feat: [TMX-393] update inference tables enablement for serving endpoi…

…nts (#29) * change codeowners * change inference tables enablement to UC from databricks * upgrade version * change the update of the inference table to be in `update_compute_config` function instead --------- Co-authored-by: Shu Ming Peh <>
hipagesgroup · Nov 6, 2023 · 9e95100 · 9e95100
1 parent d850f41
commit 9e95100
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -1,2 +1,2 @@
 # code owners of repos
-* @hipagesgroup/data-platform @hipagesgroup/lead-claim-experience
+* @hipagesgroup/data-platform @wattychops @mithunra @stimsoni9 @shumingpeh @hipagesgroup/tradie-marketplace-experience
diff --git a/hip_data_ml_utils/mlflow_databricks/mlflow_serve.py b/hip_data_ml_utils/mlflow_databricks/mlflow_serve.py
@@ -32,7 +32,7 @@ def enable_endpoint(
         token of the databricks workspace
     dbfs_table_path: str
         dbfs file path for inference logging table
-    model_version: inte
+    model_version: int
         registered model version of specified stage tag
     request_time_out: int
         timeout for the request
@@ -62,7 +62,6 @@ def enable_endpoint(
                 }
             ]
         },
-        "inference_table_config": {"dbfs_destination_path": dbfs_table_path},
     }
 
     response = requests.post(
@@ -174,6 +173,7 @@ def update_compute_config(
     workload_size_id: str,
     scale_to_zero_enabled: str,
     model_version: int,
+    catalog_name: str,
     request_time_out: int = 60,
 ) -> int:
     """
@@ -197,6 +197,8 @@ def update_compute_config(
         "false" to disable scaling to zero
     model_version: int
         registered model version of specified stage tag
+    catalog_name: str
+        name of catalog
     request_time_out: int
         timeout for the request
 
@@ -220,6 +222,11 @@ def update_compute_config(
         "traffic_config": {
             "routes": [{"served_model_name": model_name, "traffic_percentage": 100}]
         },
+        "auto_capture_config": {
+            "catalog_name": catalog_name,
+            "schema_name": "ml_features",
+            "table_name_prefix": model_name,
+        },
     }
 
     compute_config_response = requests.put(

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hip_data_ml_utils"
-version = "1.1.4"
+version = "1.2.4"
 description = "Common Python tools and utilities for Hipages ML work"
 authors = ["Hipages Data Team <datascience@hipagesgroup.com.au>"]
 license = "MIT"

diff --git a/tests/hip_data_ml_utils/mlflow_databricks/test_mlflow_serve.py b/tests/hip_data_ml_utils/mlflow_databricks/test_mlflow_serve.py
@@ -212,6 +212,7 @@ def test_update_compute_config(self, dummy_url, request_time_out: int = 30) -> N
             workload_size_id="test_workload_size_id",
             scale_to_zero_enabled="test_scale_to_zero_enabled",
             model_version=151,
+            catalog_name="lakehouse_development",
         )
 
         assert r.json()["registered_model_name"] == "test_model"  # noqa: S101
@@ -250,6 +251,7 @@ def test_update_compute_config_error(
             workload_size_id="test_workload_size_id",
             scale_to_zero_enabled="test_scale_to_zero_enabled",
             model_version=151,
+            catalog_name="lakehouse_development",
         )
         assert r.json()["registered_model_name"] == "test_model"  # noqa: S101
         assert update_compute_config_response == 1  # noqa: S101