fix: Get tracking metadata for running models

Include model info and metadata retrieved from the tracking server when
fetching the list of running models, if the 'verbose' query parameter is
set to 'true' and the running instance holds a label with the associated
model URI. The API was previously attempting to fetch information about
a registered model using the URI which would always fail (the function
expects a model name, not a URI).

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

cogstack_model_gateway/common/tracking.py

@@ -1,6 +1,7 @@
 import logging
 
 import mlflow
+import mlflow.models
 from mlflow import MlflowClient, MlflowException
 from mlflow.entities import Run, RunStatus
 
@@ -76,6 +77,7 @@ class TrackingClient:
     def __init__(self, tracking_uri: str = None):
         self.tracking_uri = tracking_uri or mlflow.get_tracking_uri()
         self._mlflow_client = MlflowClient(self.tracking_uri)
+        mlflow.set_tracking_uri(self.tracking_uri)
 
     def get_task(self, tracking_id: str) -> TrackingTask:
         """Get a task by its tracking ID."""
@@ -119,3 +121,20 @@ def get_model_uri(self, tracking_id: str) -> str:
         except Exception as e:
             log.error(f"Failed to get model URI for task with tracking ID '{tracking_id}': {e}")
             return None
+
+    def get_model_metadata(self, model_uri: str) -> dict:
+        """Get metadata for a model URI."""
+        try:
+            model_info = mlflow.models.get_model_info(model_uri)
+            return {
+                "uuid": model_info.model_uuid,
+                "signature": model_info.signature.to_dict() if model_info.signature else {},
+                "flavors": model_info.flavors,
+                "run_id": model_info.run_id,
+                "artifact_path": model_info.artifact_path,
+                "utc_time_created": model_info.utc_time_created,
+                "mlflow_version": model_info.mlflow_version,
+            }
+        except MlflowException as e:
+            log.error(f"Failed to get model metadata for model URI '{model_uri}': {e}")
+            return None

cogstack_model_gateway/gateway/core/models.py

@@ -1,9 +1,7 @@
 import os
 
 import docker
-import mlflow
 from docker.models.containers import Container
-from mlflow.entities.model_registry import RegisteredModel
 
 from cogstack_model_gateway.common.containers import (
     IS_MODEL_LABEL,
@@ -37,14 +35,6 @@ def get_running_models() -> list[dict]:
     ]
 
 
-def get_model_meta(model_uri: str) -> RegisteredModel:
-    try:
-        client = mlflow.tracking.MlflowClient()
-        return client.get_registered_model(model_uri)
-    except mlflow.exceptions.MlflowException:
-        return None
-
-
 def run_model_container(model_name: str, model_uri: str, ttl: int):
     client = docker.from_env()
     cms_project = os.getenv(CMS_PROJECT_ENV_VAR)

cogstack_model_gateway/gateway/routers/models.py

@@ -4,19 +4,15 @@
 
 import requests
 from docker.errors import DockerException
-from fastapi import APIRouter, Body, Depends, Header, HTTPException, Request
+from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query, Request
 from starlette.datastructures import UploadFile as StarletteUploadFile
 
 from cogstack_model_gateway.common.config import Config, get_config
 from cogstack_model_gateway.common.object_store import ObjectStoreManager
 from cogstack_model_gateway.common.queue import QueueManager
 from cogstack_model_gateway.common.tasks import Status, TaskManager
 from cogstack_model_gateway.common.tracking import TrackingClient
-from cogstack_model_gateway.gateway.core.models import (
-    get_model_meta,
-    get_running_models,
-    run_model_container,
-)
+from cogstack_model_gateway.gateway.core.models import get_running_models, run_model_container
 from cogstack_model_gateway.gateway.core.priority import calculate_task_priority
 from cogstack_model_gateway.gateway.routers.utils import (
     get_content_type,
@@ -103,11 +99,15 @@
 
 
 @router.get("/models/", response_model=list[dict], tags=["models"])
-async def get_models():
+async def get_models(
+    verbose: Annotated[
+        bool | None, Query(description="Include model metadata from the tracking server")
+    ] = False,
+):
     models = get_running_models()
     for model in models:
-        if model["uri"]:
-            if model_info := get_model_meta(model["uri"]):
+        if model["uri"] and verbose:
+            if model_info := TrackingClient().get_model_metadata(model["uri"]):
                 model["info"] = model_info
     return models