Add Model Run ID and Version Number to JobTable (dev to staging merge)

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "pandera~=0.13",
     "mlflow~=2.15.0",
     "cachetools",
+    "types-cachetools",
 ]
 
 [project.urls]

src/webapp/database.py

@@ -483,7 +483,10 @@ class ModelTable(Base):
     )
     inst: Mapped["InstTable"] = relationship(back_populates="models")
 
-    jobs: Mapped[Set["JobTable"]] = relationship(back_populates="model")
+    jobs: Mapped[Set["JobTable"]] = relationship(
+        back_populates="model",
+        passive_deletes=True,
+    )
 
     name: Mapped[str] = mapped_column(String(VAR_CHAR_STANDARD_LENGTH), nullable=False)
     # What configuration of schemas are allowed (list of maps e.g. [PDP Course : 1 + PDP Cohort : 1, X_schema :1 + Y_schema: 2])
@@ -551,8 +554,11 @@ class JobTable(Base):
         String(VAR_CHAR_STANDARD_LENGTH), nullable=True
     )
     completed: Mapped[bool] = mapped_column(nullable=True)
-    framework: Mapped[str | None] = mapped_column(
-        String(VAR_CHAR_STANDARD_LENGTH), nullable=False, default="sklearn"
+    model_version: Mapped[str | None] = mapped_column(
+        String(VAR_CHAR_STANDARD_LENGTH), nullable=True
+    )
+    model_run_id: Mapped[str | None] = mapped_column(
+        String(VAR_CHAR_STANDARD_LENGTH), nullable=True
     )
 
 

src/webapp/databricks.py

@@ -50,7 +50,6 @@ class DatabricksInferenceRunRequest(BaseModel):
     # Note that the following should be the filepath.
     filepath_to_type: dict[str, list[SchemaType]]
     model_name: str
-    model_type: str
     # The email where notifications will get sent.
     email: str
     gcp_external_bucket_name: str
@@ -89,10 +88,10 @@ def _sha256_json(obj: Any) -> str:
     ).hexdigest()
 
 
-L1_RESP_CACHE_TTL = int("120")  # seconds
-L1_VER_CACHE_TTL = int("60")  # seconds
-L1_RESP_CACHE = TTLCache(maxsize=128, ttl=L1_RESP_CACHE_TTL)
-L1_VER_CACHE = TTLCache(maxsize=256, ttl=L1_VER_CACHE_TTL)
+L1_RESP_CACHE_TTL = int("600")  # seconds
+L1_VER_CACHE_TTL = int("3600")  # seconds
+L1_RESP_CACHE: Any = TTLCache(maxsize=128, ttl=L1_RESP_CACHE_TTL)
+L1_VER_CACHE: Any = TTLCache(maxsize=256, ttl=L1_VER_CACHE_TTL)
 _L1_LOCK = threading.RLock()
 
 
@@ -252,7 +251,6 @@ def run_pdp_inference(
                     ],  # is this value the same PER environ? dev/staging/prod
                     "gcp_bucket_name": req.gcp_external_bucket_name,
                     "model_name": req.model_name,
-                    "model_type": req.model_type,
                     "notification_email": req.email,
                 },
             )
@@ -334,7 +332,7 @@ def fetch_table_data(
         inst_name: str,
         table_name: str,
         warehouse_id: str,
-    ) -> List[Dict[str, Any]]:
+    ) -> Any:
         """
         Execute SELECT * via Databricks SQL Statement Execution API using EXTERNAL_LINKS.
         Blocks server-side for up to 30s; if not SUCCEEDED, raises. Downloads presigned
@@ -367,9 +365,9 @@ def fetch_table_data(
 
             if not ver_resp.status or ver_resp.status.state != StatementState.SUCCEEDED:
                 raise TimeoutError("DESCRIBE HISTORY did not finish within 30s")
-            cols = [c.name for c in ver_resp.manifest.schema.columns]
+            cols = [c.name for c in ver_resp.manifest.schema.columns]  # type: ignore
             idx = {n: i for i, n in enumerate(cols)}
-            rows = ver_resp.result.data_array or []
+            rows = ver_resp.result.data_array or []  # type: ignore
             if not rows or "version" not in idx:
                 raise ValueError("DESCRIBE HISTORY returned no version")
             table_version = str(rows[0][idx["version"]])
@@ -433,13 +431,13 @@ def fetch_table_data(
             resp.manifest and resp.manifest.schema and resp.manifest.schema.columns
         ):
             raise ValueError("Schema/columns missing (EXTERNAL_LINKS).")
-        cols: List[str] = []
+        cols: List[str] = []  # type: ignore
         for c in resp.manifest.schema.columns:
             if c.name is None:
                 raise ValueError("Encountered a column without a name.")
             cols.append(c.name)
 
-        records: List[Dict[str, Any]] = []
+        records: Any = []
 
         # Helper: consume one chunk-like object (first result or subsequent chunk)
         def _consume_chunk(chunk_obj: Any) -> int | None:
@@ -505,6 +503,62 @@ def _consume_chunk(chunk_obj: Any) -> int | None:
                 pass
         return records
 
+    def fetch_model_version(
+        self, catalog_name: str, inst_name: str, model_name: str
+    ) -> Any:
+        schema = databricksify_inst_name(inst_name)
+        model_name_path = f"{catalog_name}.{schema}_gold.{model_name}"
+
+        try:
+            w = WorkspaceClient(
+                host=databricks_vars["DATABRICKS_HOST_URL"],
+                google_service_account=gcs_vars["GCP_SERVICE_ACCOUNT_EMAIL"],
+            )
+        except Exception as e:
+            LOGGER.exception(
+                "Failed to create Databricks WorkspaceClient with host: %s and service account: %s",
+                databricks_vars["DATABRICKS_HOST_URL"],
+                gcs_vars["GCP_SERVICE_ACCOUNT_EMAIL"],
+            )
+            raise ValueError(f"setup_new_inst(): Workspace client creation failed: {e}")
+
+        model_versions: Any = list(
+            w.model_versions.list(
+                full_name=model_name_path,
+            )
+        )
+
+        if not model_versions:
+            raise ValueError(f"No versions found for model: {model_name_path}")
+
+        latest_version = max(model_versions, key=lambda v: int(v.version))
+
+        return latest_version
+
+    def delete_model(self, catalog_name: str, inst_name: str, model_name: str) -> None:
+        schema = databricksify_inst_name(inst_name)
+        model_name_path = f"{catalog_name}.{schema}_gold.{model_name}"
+
+        try:
+            w = WorkspaceClient(
+                host=databricks_vars["DATABRICKS_HOST_URL"],
+                google_service_account=gcs_vars["GCP_SERVICE_ACCOUNT_EMAIL"],
+            )
+        except Exception as e:
+            LOGGER.exception(
+                "Failed to create Databricks WorkspaceClient with host: %s and service account: %s",
+                databricks_vars["DATABRICKS_HOST_URL"],
+                gcs_vars["GCP_SERVICE_ACCOUNT_EMAIL"],
+            )
+            raise ValueError(f"setup_new_inst(): Workspace client creation failed: {e}")
+
+        try:
+            w.registered_models.delete(full_name=model_name_path)
+            LOGGER.info("Deleted registration model: %s", model_name_path)
+        except Exception:
+            LOGGER.exception("Failed to delete registered model: %s", model_name_path)
+            raise
+
     def get_key_for_file(
         self, mapping: Dict[str, Any], file_name: str
     ) -> Optional[str]: