added model to features table

Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED

Author: Vperiodt

sdk/python/feast/ui_server.py

@@ -1,7 +1,8 @@
 import json
+import logging
 import threading
 from importlib import resources as importlib_resources
-from typing import Callable, Optional
+from typing import Callable, Dict, List, Optional
 
 import uvicorn
 from fastapi import FastAPI, Response, status
@@ -10,6 +11,8 @@
 
 import feast
 
+_logger = logging.getLogger(__name__)
+
 
 def get_app(
     store: "feast.FeatureStore",
@@ -147,6 +150,25 @@ def get_mlflow_runs(max_results: int = 50):
                 order_by=["start_time DESC"],
             )
 
+            run_id_to_models: Dict[str, List[dict]] = {}
+            try:
+                for rm in client.search_registered_models():
+                    for mv in rm.latest_versions or []:
+                        if mv.run_id:
+                            run_id_to_models.setdefault(mv.run_id, []).append(
+                                {
+                                    "model_name": rm.name,
+                                    "version": mv.version,
+                                    "stage": mv.current_stage,
+                                    "mlflow_url": (
+                                        f"{tracking_uri}/#/models/"
+                                        f"{rm.name}/versions/{mv.version}"
+                                    ),
+                                }
+                            )
+            except Exception:
+                pass
+
             result = []
             mlflow_ui_base = tracking_uri
             for run in runs:
@@ -175,6 +197,9 @@ def get_mlflow_runs(max_results: int = 50):
                             f"{mlflow_ui_base}/#/experiments/"
                             f"{run.info.experiment_id}/runs/{run.info.run_id}"
                         ),
+                        "registered_models": run_id_to_models.get(
+                            run.info.run_id, []
+                        ),
                     }
                 )
 
@@ -192,6 +217,72 @@ def get_mlflow_runs(max_results: int = 50):
                 "error": "Failed to fetch MLflow runs",
             }
 
+    @app.get("/api/mlflow-feature-models")
+    def get_mlflow_feature_models():
+        """Return a mapping of feature_ref -> registered models that use it.
+
+        Walks the MLflow Model Registry, inspects the training run for each
+        model's latest version(s), reads the ``feast.feature_refs`` tag, and
+        inverts it into a reverse index so the UI can show which registered
+        models depend on a given feature.
+        """
+        mlflow_cfg = getattr(store.config, "mlflow", None)
+        if not mlflow_cfg or not mlflow_cfg.enabled:
+            return {"feature_models": {}}
+
+        try:
+            import mlflow
+
+            tracking_uri = mlflow_cfg.tracking_uri or "http://127.0.0.1:5000"
+            client = mlflow.MlflowClient(tracking_uri=tracking_uri)
+            project_name = store.config.project
+
+            feature_models: Dict[str, List[dict]] = {}
+
+            for rm in client.search_registered_models():
+                model_name = rm.name
+                latest_versions = rm.latest_versions or []
+                for mv in latest_versions:
+                    if not mv.run_id:
+                        continue
+                    try:
+                        run = client.get_run(mv.run_id)
+                    except Exception:
+                        continue
+
+                    tags = run.data.tags
+                    if tags.get("feast.project") != project_name:
+                        continue
+
+                    refs_raw = tags.get("feast.feature_refs", "")
+                    feature_refs = [r for r in refs_raw.split(",") if r]
+
+                    model_info = {
+                        "model_name": model_name,
+                        "version": mv.version,
+                        "stage": mv.current_stage,
+                        "mlflow_url": (
+                            f"{tracking_uri}/#/models/"
+                            f"{model_name}/versions/{mv.version}"
+                        ),
+                    }
+
+                    for ref in feature_refs:
+                        feature_models.setdefault(ref, []).append(model_info)
+
+            return {"feature_models": feature_models}
+        except ImportError:
+            return {
+                "feature_models": {},
+                "error": "mlflow is not installed",
+            }
+        except Exception as e:
+            _logger.debug("Failed to fetch MLflow feature-model mapping: %s", e)
+            return {
+                "feature_models": {},
+                "error": "Failed to fetch model data",
+            }
+
     # For all other paths (such as paths that would otherwise be handled by react router), pass to React
     @app.api_route("/p/{path_name:path}", methods=["GET"])
     def catch_all():

ui/src/components/RegistryVisualization.tsx

@@ -82,6 +82,8 @@ const getNodeColor = (type: FEAST_FCO_TYPES) => {
       return "#cc0000"; // Red
     case FEAST_FCO_TYPES.mlflowRun:
       return "#0194e2"; // MLflow brand blue
+    case FEAST_FCO_TYPES.mlflowModel:
+      return "#7b2d8e"; // Purple
     default:
       return "#666666"; // Gray
   }
@@ -99,6 +101,8 @@ const getLightNodeColor = (type: FEAST_FCO_TYPES) => {
       return "#ffe6e6"; // Light red
     case FEAST_FCO_TYPES.mlflowRun:
       return "#e6f6fd"; // Light MLflow blue
+    case FEAST_FCO_TYPES.mlflowModel:
+      return "#f3e6f9"; // Light purple
     default:
       return "#f0f0f0"; // Light gray
   }
@@ -116,6 +120,8 @@ const getNodeIcon = (type: FEAST_FCO_TYPES) => {
       return "◆"; // Diamond for data source
     case FEAST_FCO_TYPES.mlflowRun:
       return "⬡"; // Hexagon for MLflow run
+    case FEAST_FCO_TYPES.mlflowModel:
+      return "⬢"; // Filled hexagon for registered model
     default:
       return "●"; // Default circle
   }
@@ -132,7 +138,11 @@ const CustomNode = ({ data }: { data: NodeData }) => {
   const hasVersion = data.versionNumber != null && data.versionNumber > 1;
 
   const handleClick = () => {
-    if (data.type === FEAST_FCO_TYPES.mlflowRun && data.metadata?.mlflow_url) {
+    if (
+      (data.type === FEAST_FCO_TYPES.mlflowRun ||
+        data.type === FEAST_FCO_TYPES.mlflowModel) &&
+      data.metadata?.mlflow_url
+    ) {
       window.open(data.metadata.mlflow_url, "_blank", "noopener,noreferrer");
       return;
     }
@@ -194,7 +204,8 @@ const CustomNode = ({ data }: { data: NodeData }) => {
             zIndex: 5,
           }}
         >
-          {data.type === FEAST_FCO_TYPES.mlflowRun
+          {data.type === FEAST_FCO_TYPES.mlflowRun ||
+          data.type === FEAST_FCO_TYPES.mlflowModel
             ? "Open in MLflow ↗"
             : "View Details"}
         </div>
@@ -412,6 +423,7 @@ const getLayoutedElements = (
     [FEAST_FCO_TYPES.featureView]: [],
     [FEAST_FCO_TYPES.featureService]: [],
     [FEAST_FCO_TYPES.mlflowRun]: [],
+    [FEAST_FCO_TYPES.mlflowModel]: [],
   };
 
   isolatedNodes.forEach((node) => {
@@ -469,6 +481,7 @@ const Legend = () => {
     { type: FEAST_FCO_TYPES.entity, label: "Entity" },
     { type: FEAST_FCO_TYPES.dataSource, label: "Data Source" },
     { type: FEAST_FCO_TYPES.mlflowRun, label: "MLflow Run" },
+    { type: FEAST_FCO_TYPES.mlflowModel, label: "Registered Model" },
   ];
 
   const isDarkMode = colorMode === "dark";
@@ -805,6 +818,52 @@ const registryToFlow = (
           });
         }
       }
+
+      if (run.registered_models && run.registered_models.length > 0) {
+        run.registered_models.forEach((model) => {
+          const modelNodeId = `model-${model.model_name}-v${model.version}`;
+          const modelExists = nodes.some((n) => n.id === modelNodeId);
+          if (!modelExists) {
+            nodes.push({
+              id: modelNodeId,
+              type: "custom",
+              data: {
+                label: `${model.model_name} v${model.version}`,
+                type: FEAST_FCO_TYPES.mlflowModel,
+                metadata: {
+                  mlflow_url: model.mlflow_url,
+                  model_name: model.model_name,
+                  version: model.version,
+                  stage: model.stage,
+                },
+              },
+              position: { x: 0, y: 0 },
+            });
+          }
+
+          edges.push({
+            id: `edge-model-${run.run_id}-${model.model_name}-v${model.version}`,
+            source: `mlflow-${run.run_id}`,
+            sourceHandle: "source",
+            target: modelNodeId,
+            targetHandle: "target",
+            animated: true,
+            style: {
+              strokeWidth: 3,
+              stroke: "#7b2d8e",
+              strokeDasharray: "10 5",
+              animation: "dataflow 2s linear infinite",
+            },
+            type: "smoothstep",
+            markerEnd: {
+              type: MarkerType.ArrowClosed,
+              width: 20,
+              height: 20,
+              color: "#7b2d8e",
+            },
+          });
+        });
+      }
     });
   }
 
@@ -823,6 +882,8 @@ const getNodePrefix = (type: FEAST_FCO_TYPES) => {
       return "ds";
     case FEAST_FCO_TYPES.mlflowRun:
       return "mlflow";
+    case FEAST_FCO_TYPES.mlflowModel:
+      return "model";
     default:
       return "unknown";
   }

ui/src/hooks/useFCOExploreSuggestions.ts

@@ -23,6 +23,7 @@ const FCO_TO_URL_NAME_MAP: Record<FEAST_FCO_TYPES, string> = {
   featureView: "/feature-view",
   featureService: "/feature-service",
   mlflowRun: "/mlflow-run",
+  mlflowModel: "/mlflow-model",
 };
 
 const createSearchLink = (

ui/src/pages/features/FeatureListPage.tsx

@@ -15,12 +15,16 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
   EuiFormRow,
+  EuiBadge,
 } from "@elastic/eui";
 import EuiCustomLink from "../../components/EuiCustomLink";
 import ExportButton from "../../components/ExportButton";
 import { useParams } from "react-router-dom";
 import useLoadRegistry from "../../queries/useLoadRegistry";
 import RegistryPathContext from "../../contexts/RegistryPathContext";
+import useLoadFeatureModels, {
+  FeatureModelInfo,
+} from "../../queries/useLoadFeatureModels";
 import { FeatureIcon } from "../../graphics/FeatureIcon";
 import { FEAST_FCO_TYPES } from "../../parsers/types";
 import {
@@ -35,6 +39,7 @@ interface Feature {
   type: string;
   project?: string;
   permissions?: any[];
+  models?: FeatureModelInfo[];
 }
 
 type FeatureColumn =
@@ -48,6 +53,7 @@ const FeatureListPage = () => {
     registryUrl,
     projectName,
   );
+  const { data: featureModelsData } = useLoadFeatureModels();
   const [searchText, setSearchText] = useState("");
   const [selectedPermissionAction, setSelectedPermissionAction] = useState("");
 
@@ -59,8 +65,11 @@ const FeatureListPage = () => {
 
   const featuresWithPermissions: Feature[] = (data?.allFeatures || []).map(
     (feature) => {
+      const featureRef = `${feature.featureView}:${feature.name}`;
       return {
         ...feature,
+        models:
+          featureModelsData?.feature_models?.[featureRef] || [],
         permissions: getEntityPermissions(
           selectedPermissionAction
             ? filterPermissionsByAction(
@@ -126,6 +135,49 @@ const FeatureListPage = () => {
       },
     },
     { name: "Type", field: "type", sortable: true },
+    {
+      name: "Models",
+      field: "models",
+      sortable: false,
+      render: (models: FeatureModelInfo[]) => {
+        if (!models || models.length === 0) {
+          return (
+            <EuiText size="xs" color="subdued">
+              --
+            </EuiText>
+          );
+        }
+        if (models.length === 1) {
+          return (
+            <EuiBadge
+              color="hollow"
+              href={models[0].mlflow_url}
+              target="_blank"
+            >
+              {models[0].model_name} v{models[0].version}
+            </EuiBadge>
+          );
+        }
+        return (
+          <EuiToolTip
+            position="top"
+            content={
+              <div>
+                {models.map((m) => (
+                  <div key={`${m.model_name}_v${m.version}`}>
+                    {m.model_name} v{m.version}
+                  </div>
+                ))}
+              </div>
+            }
+          >
+            <EuiBadge color="hollow">
+              {models.length} models
+            </EuiBadge>
+          </EuiToolTip>
+        );
+      },
+    },
     {
       name: "Permissions",
       field: "permissions",

ui/src/parsers/types.ts

@@ -4,6 +4,7 @@ enum FEAST_FCO_TYPES {
   featureView = "featureView",
   featureService = "featureService",
   mlflowRun = "mlflowRun",
+  mlflowModel = "mlflowModel",
 }
 
 export { FEAST_FCO_TYPES };