Add Model Run ID and Version Number to JobTable

pyproject.toml

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

src/webapp/database.py

@@ -551,8 +551,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]:

src/webapp/routers/models.py

@@ -5,7 +5,7 @@
 import jsonpickle
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import BaseModel
-from sqlalchemy import and_
+from sqlalchemy import and_, update, or_
 from sqlalchemy.orm import Session
 from sqlalchemy.future import select
 from ..databricks import DatabricksControl, DatabricksInferenceRunRequest
@@ -33,6 +33,7 @@
 import traceback
 import logging
 from ..gcsdbutils import update_db_from_bucket
+from ..config import env_vars
 
 from ..gcsutil import StorageControl
 
@@ -310,6 +311,50 @@ def read_inst_model(
     }
 
 
+@router.delete("/{inst_id}/models/{model_name}")
+def delete_model(
+    inst_id: str,
+    model_name: str,
+    current_user: Annotated[BaseUser, Depends(get_current_active_user)],
+    sql_session: Annotated[Session, Depends(get_session)],
+) -> Any:
+    transformed_model_name = str(decode_url_piece(model_name)).strip()
+    has_access_to_inst_or_err(inst_id, current_user)
+    model_owner_and_higher_or_err(current_user, "modify batch")
+
+    local_session.set(sql_session)
+    sess = local_session.get()
+
+    model_list = sess.execute(
+        select(ModelTable).where(
+            ModelTable.name == transformed_model_name,
+            ModelTable.inst_id == str_to_uuid(inst_id),
+        )
+    ).scalar_one_or_none()
+    if model_list is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Model not found."
+        )
+
+    # 2) Optionally Delete models from databricks itself
+    # TODO: Add databricks deletion functionality
+
+    try:
+        sess.delete(model_list)
+        sess.commit()
+    except Exception as e:
+        sess.rollback()
+        raise HTTPException(
+            status_code=500, detail=f"DB batch delete failed after file cleanup: {e}"
+        )
+
+    return {
+        "inst_id": inst_id,
+        "model_name": transformed_model_name,
+        "status": "Model deleted",
+    }
+
+
 @router.get("/{inst_id}/models/{model_name}/runs", response_model=list[RunInfo])
 def read_inst_model_outputs(
     inst_id: str,
@@ -364,6 +409,8 @@ def read_inst_model_outputs(
                 "inst_id": uuid_to_str(query_result[0][0].inst_id),
                 "m_name": query_result[0][0].name,
                 "run_id": elem.id,
+                "model_run_id": elem.model_run_id,
+                "model_version": elem.model_version,
                 "created_by": uuid_to_str(elem.created_by),
                 "triggered_at": elem.triggered_at,
                 "batch_name": elem.batch_name,
@@ -555,7 +602,6 @@ def trigger_inference_run(
         gcp_external_bucket_name=get_external_bucket_name(inst_id),
         # The institution email to which pipeline success/failure notifications will get sent.
         email=cast(str, current_user.email),
-        model_type=query_result[0][0].framework,
     )
     try:
         res = databricks_control.run_pdp_inference(db_req)
@@ -567,14 +613,20 @@ def trigger_inference_run(
             detail=f"Databricks run_pdp_inference error. Error = {str(e)}",
         ) from e
     triggered_timestamp = datetime.now()
+    latest_model_version = databricks_control.fetch_model_version(
+        catalog_name=str(env_vars["CATALOG_NAME"]),
+        inst_name=inst_result[0][0].name,
+        model_name=model_name,
+    )
     job = JobTable(
         id=res.job_run_id,
         triggered_at=triggered_timestamp,
         created_by=str_to_uuid(current_user.user_id),
         batch_name=req.batch_name,
         model_id=query_result[0][0].id,
         output_valid=False,
-        framework=query_result[0][0].framework,
+        model_version=latest_model_version.version,
+        model_run_id=latest_model_version.run_id,
     )
     local_session.get().add(job)
     return {
@@ -585,5 +637,130 @@ def trigger_inference_run(
         "triggered_at": triggered_timestamp,
         "batch_name": req.batch_name,
         "output_valid": False,
-        "framework": query_result[0][0].framework,
+        "model_version": latest_model_version.version,
+        "model_run_id": latest_model_version.run_id,
+    }
+
+
+@router.get("/{inst_id}/models/{model_name}/get-model-versions")
+def get_model_versions(
+    inst_id: str,
+    model_name: str,
+    current_user: Annotated[BaseUser, Depends(get_current_active_user)],
+    sql_session: Annotated[Session, Depends(get_session)],
+    databricks_control: Annotated[DatabricksControl, Depends(DatabricksControl)],
+) -> Any:
+    transformed_model_name = str(decode_url_piece(model_name)).strip()
+    has_access_to_inst_or_err(inst_id, current_user)
+
+    local_session.set(sql_session)
+    query_result = (
+        local_session.get()
+        .execute(select(InstTable).where(InstTable.id == str_to_uuid(inst_id)))
+        .all()
+    )
+    if not query_result or len(query_result) == 0:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Institution not found.",
+        )
+    if len(query_result) > 1:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Institution duplicates found.",
+        )
+
+    print(f"Initial model name = {model_name}")
+    print(f"Converted model name {transformed_model_name}")
+
+    latest_model_version = databricks_control.fetch_model_version(
+        catalog_name=str(env_vars["CATALOG_NAME"]),
+        inst_name=f"{query_result[0][0].name}",
+        model_name=transformed_model_name,
+    )
+
+    return latest_model_version
+
+
+@router.post("/{inst_id}/models/{model_name}/backfill-model-runs")
+def backfill_model_runs(
+    inst_id: str,
+    model_name: str,
+    current_user: Annotated[BaseUser, Depends(get_current_active_user)],
+    sql_session: Annotated[Session, Depends(get_session)],
+    databricks_control: Annotated[DatabricksControl, Depends(DatabricksControl)],
+) -> Any:
+    """Backfills missing model run metadata and returns the latest model version info.
+
+    Temporary endpoint to populate model_run_id and model_version on existing jobs for this model.
+    Use only when backfilling historical job runs, not for regular operation.
+    """
+    model_name = str(decode_url_piece(model_name)).strip()
+    has_access_to_inst_or_err(inst_id, current_user)
+
+    # Load institution
+    local_session.set(sql_session)
+    inst_row = (
+        local_session.get()
+        .execute(select(InstTable).where(InstTable.id == str_to_uuid(inst_id)))
+        .all()
+    )
+
+    model_id = (
+        local_session.get()
+        .execute(
+            select(ModelTable).where(
+                and_(
+                    ModelTable.inst_id == str_to_uuid(inst_id),
+                    ModelTable.name == model_name,
+                )
+            )
+        )
+        .all()
+    )
+
+    if not inst_row or len(inst_row) == 0:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Institution not found.",
+        )
+    if len(inst_row) > 1:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Institution duplicates found.",
+        )
+
+    latest_mv = databricks_control.fetch_model_version(
+        catalog_name=str(env_vars["CATALOG_NAME"]),
+        inst_name=f"{inst_row[0][0].name}",
+        model_name=model_name,
+    )
+
+    mv_version = str(latest_mv.version)
+    mv_run_id = str(latest_mv.run_id)
+
+    # UPDATE existing jobs for this model (only those missing values)
+    stmt = (
+        update(JobTable)
+        .where(JobTable.model_id == model_id[0][0].id)
+        .where(
+            or_(
+                JobTable.model_run_id.is_(None),
+                JobTable.model_run_id == "",
+                JobTable.model_version.is_(None),
+                JobTable.model_version == "",
+            )
+        )
+        .values(model_run_id=mv_run_id, model_version=mv_version)
+    )
+    result = local_session.get().execute(stmt)
+    updated_count = result.rowcount or 0  # type: ignore
+    local_session.get().commit()
+
+    return {
+        "inst_id": str(inst_id),
+        "model_id": str(model_id[0][0].id),
+        "model_name": model_name,
+        "latest_model_version": {"version": mv_version, "run_id": mv_run_id},
+        "updated_count": updated_count,
     }