feat: Added Model Metadata support in Registry

Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>

Author: ntkathole

README.md

@@ -9,7 +9,7 @@
 
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/feast)](https://pypi.org/project/feast/)
 [![GitHub contributors](https://img.shields.io/github/contributors/feast-dev/feast)](https://github.com/feast-dev/feast/graphs/contributors)
-[![unit-tests](https://github.com/feast-dev/feast/actions/workflows/unit_tests.yml/badge.svg?branch=master)](https://github.com/feast-dev/feast/actions/workflows/unit_tests.yml)
+[![unit-tests](https://github.com/feast-dev/feast/actions/workflows/unit_tests.yml/badge.svg?branch=master&event=pull_request)](https://github.com/feast-dev/feast/actions/workflows/unit_tests.yml)
 [![integration-tests-and-build](https://github.com/feast-dev/feast/actions/workflows/master_only.yml/badge.svg?branch=master&event=push)](https://github.com/feast-dev/feast/actions/workflows/master_only.yml)
 [![java-integration-tests](https://github.com/feast-dev/feast/actions/workflows/java_master_only.yml/badge.svg?branch=master&event=push)](https://github.com/feast-dev/feast/actions/workflows/java_master_only.yml)
 [![linter](https://github.com/feast-dev/feast/actions/workflows/linter.yml/badge.svg?branch=master&event=push)](https://github.com/feast-dev/feast/actions/workflows/linter.yml)
@@ -21,6 +21,8 @@
 ## Join us on Slack!
 👋👋👋 [Come say hi on Slack!](https://communityinviter.com/apps/feastopensource/feast-the-open-source-feature-store)
 
+[Check out our DeepWiki!](https://deepwiki.com/feast-dev/feast)
+
 ## Overview
 
 Feast (**Fea**ture **St**ore) is an open source feature store for machine learning. Feast is the fastest path to manage existing infrastructure to productionize analytic data for model training and online inference.
@@ -257,4 +259,4 @@ Thanks goes to these incredible people:
 
 <a href="https://github.com/feast-dev/feast/graphs/contributors">
   <img src="https://contrib.rocks/image?repo=feast-dev/feast" />
-</a>
+</a>

protos/feast/core/Model.proto

@@ -0,0 +1,34 @@
+syntax = "proto3";
+package feast.core;
+
+option go_package = "github.com/feast-dev/feast/go/protos/feast/core";
+option java_outer_classname = "ModelMetadataProto";
+option java_package = "feast.proto.core";
+
+import "google/protobuf/timestamp.proto";
+
+message ModelMetadata {
+  // Required: Unique model name
+  string name = 1;
+
+  // Optional: Associated feature view name (if used)
+  repeated string feature_view = 2;
+
+  // Optional: Associated feature service name (if used)
+  repeated string feature_service = 3;
+
+  // Optional: Direct list of fully qualified feature references
+  repeated string features = 4;
+
+  // Project this model belongs to
+  string project = 5;
+
+  // Metadata (e.g., model purpose, version, tags)
+  map<string, string> tags = 6;
+
+  // Timestamp when model was trained
+  google.protobuf.Timestamp training_timestamp = 7;
+
+  // Optional: Description of the model
+  string description = 8;
+}

protos/feast/core/Registry.proto

@@ -34,8 +34,9 @@ import "feast/core/ValidationProfile.proto";
 import "google/protobuf/timestamp.proto";
 import "feast/core/Permission.proto";
 import "feast/core/Project.proto";
+import "feast/core/Model.proto";
 
-// Next id: 18
+// Next id: 19
 message Registry {
     repeated Entity entities = 1;
     repeated FeatureTable feature_tables = 2;
@@ -55,6 +56,7 @@ message Registry {
     google.protobuf.Timestamp last_updated = 5;
     repeated Permission permissions = 16;
     repeated Project projects = 17;
+    repeated ModelMetadata model_metadata = 18; 
 }
 
 message ProjectMetadata {

sdk/python/feast/diff/registry_diff.py

@@ -18,6 +18,7 @@
     FeatureService as FeatureServiceProto,
 )
 from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
+from feast.protos.feast.core.Model_pb2 import ModelMetadata as ModelMetadataProto
 from feast.protos.feast.core.OnDemandFeatureView_pb2 import (
     OnDemandFeatureView as OnDemandFeatureViewProto,
 )
@@ -133,8 +134,10 @@ def diff_registry_objects(
     current_spec: FeastObjectSpecProto
     new_spec: FeastObjectSpecProto
     if isinstance(
-        current_proto, (DataSourceProto, ValidationReferenceProto)
-    ) or isinstance(new_proto, (DataSourceProto, ValidationReferenceProto)):
+        current_proto, (DataSourceProto, ValidationReferenceProto, ModelMetadataProto)
+    ) or isinstance(
+        new_proto, (DataSourceProto, ValidationReferenceProto, ModelMetadataProto)
+    ):
         assert type(current_proto) == type(new_proto)
         current_spec = cast(DataSourceProto, current_proto)
         new_spec = cast(DataSourceProto, new_proto)

sdk/python/feast/errors.py

@@ -522,3 +522,8 @@ def grpc_status_code(self) -> "GrpcStatusCode":
 
     def http_status_code(self) -> int:
         return HttpStatusCode.HTTP_403_FORBIDDEN
+
+
+class ModelObjectNotFoundException(Exception):
+    def __init__(self, name: str, project: str):
+        super().__init__(f"Model {name} not found in project {project}")

sdk/python/feast/feast_object.py

@@ -8,6 +8,7 @@
 from .entity import Entity
 from .feature_service import FeatureService
 from .feature_view import FeatureView
+from .model import ModelMetadata
 from .on_demand_feature_view import OnDemandFeatureView
 from .permissions.permission import Permission
 from .protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
@@ -37,6 +38,7 @@
     ValidationReference,
     SavedDataset,
     Permission,
+    ModelMetadata,
 ]
 
 FeastObjectSpecProto = Union[

sdk/python/feast/feature_store.py

@@ -71,6 +71,7 @@
 from feast.infra.registry.base_registry import BaseRegistry
 from feast.infra.registry.registry import Registry
 from feast.infra.registry.sql import SqlRegistry
+from feast.model import ModelMetadata
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.online_response import OnlineResponse
 from feast.permissions.permission import Permission
@@ -809,6 +810,7 @@ def apply(
             FeatureService,
             ValidationReference,
             Permission,
+            ModelMetadata,
             List[FeastObject],
         ],
         objects_to_delete: Optional[List[FeastObject]] = None,
@@ -924,6 +926,10 @@ def apply(
 
         data_sources_to_update = list(data_sources_set_to_update)
 
+        model_metadata_to_update: List[ModelMetadata] = [
+            ob for ob in objects if isinstance(ob, ModelMetadata)
+        ]
+
         # Handle all entityless feature views by using DUMMY_ENTITY as a placeholder entity.
         entities_to_update.append(DUMMY_ENTITY)
 
@@ -963,11 +969,16 @@ def apply(
             self._registry.apply_permission(
                 permission, project=self.project, commit=False
             )
+        for model_metadata in model_metadata_to_update:
+            self._registry.apply_model(
+                model_metadata, project=self.project, commit=False
+            )
 
         entities_to_delete = []
         views_to_delete = []
         sfvs_to_delete = []
         permissions_to_delete = []
+        model_metadata_to_delete = []
         if not partial:
             # Delete all registry objects that should not exist.
             entities_to_delete = [
@@ -999,6 +1010,9 @@ def apply(
             permissions_to_delete = [
                 ob for ob in objects_to_delete if isinstance(ob, Permission)
             ]
+            model_metadata_to_delete = [
+                ob for ob in objects_to_delete if isinstance(ob, ModelMetadata)
+            ]
 
             for data_source in data_sources_to_delete:
                 self._registry.delete_data_source(
@@ -1032,6 +1046,10 @@ def apply(
                 self._registry.delete_permission(
                     permission.name, project=self.project, commit=False
                 )
+            for model_metadata in model_metadata_to_delete:
+                self._registry.delete_model(
+                    model_metadata.name, project=self.project, commit=False
+                )
 
         tables_to_delete: List[FeatureView] = (
             views_to_delete + sfvs_to_delete if not partial else []  # type: ignore
@@ -2531,6 +2549,55 @@ def list_saved_datasets(
             self.project, allow_cache=allow_cache, tags=tags
         )
 
+    def get_model(self, name: str, allow_cache: bool = False) -> ModelMetadata:
+        """
+        Retrieves a specific model by ID.
+
+        Args:
+            name: The name of the model.
+            allow_cache: Whether to allow returning from cache.
+
+        Returns:
+            The ModelMetadata for the specified model.
+
+        Raises:
+            ModelObjectNotFoundException: If the model does not exist.
+        """
+        return self._registry.get_model(
+            name=name, project=self.project, allow_cache=allow_cache
+        )
+
+    def list_models(
+        self,
+        allow_cache: bool = False,
+        tags: Optional[dict[str, str]] = None,
+    ) -> List[ModelMetadata]:
+        """
+        List models stored in the registry.
+
+        Args:
+            allow_cache: Whether to allow returning from a cached registry.
+            tags: Optional dictionary to filter models by tags.
+
+        Returns:
+            A list of ModelMetadata objects.
+        """
+        return self._registry.list_models(
+            allow_cache=allow_cache,
+            project=self.project,
+            tags=tags,
+        )
+
+    def delete_model(self, name: str) -> None:
+        """
+        Deletes a model from the registry.
+
+        Args:
+            name: The name of the model to delete.
+            feast_project: The project the model belongs to (defaults to current project).
+        """
+        self._registry.delete_model(name=name, project=self.project)
+
     async def initialize(self) -> None:
         """Initialize long-lived clients and/or resources needed for accessing datastores"""
         await self._get_provider().initialize(self.config)

sdk/python/feast/infra/registry/base_registry.py

@@ -27,6 +27,7 @@
 from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
 from feast.infra.infra_object import Infra
+from feast.model import ModelMetadata
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.permissions.permission import Permission
 from feast.project import Project
@@ -768,6 +769,71 @@ def list_projects(
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def apply_model(
+        self, model_metadata: ModelMetadata, project: str, commit: bool = True
+    ) -> None:
+        """
+        Register or update model metadata in the registry.
+
+        Args:
+            model_metadata: The model metadata object to apply.
+            project: Feast project that this entity belongs to
+            commit: Whether the change should be persisted immediately
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_model(
+        self, name: str, project: str, allow_cache: bool = False
+    ) -> ModelMetadata:
+        """
+        Retrieve a model from the registry by model ID.
+
+        Args:
+            name: The name of the model to retrieve.
+            allow_cache: Whether to allow returning model from cached registry.
+
+        Returns:
+            ModelMetadata
+
+        Raises:
+            ModelObjectNotFoundException: If the model is not found.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def list_models(
+        self,
+        project: str,
+        allow_cache: bool = False,
+        tags: Optional[Dict[str, str]] = None,
+    ) -> List[ModelMetadata]:
+        """
+        List all models registered in the registry.
+
+        Args:
+            allow_cache: Whether to allow returning models from cached registry.
+            project: Project name to filter models.
+            tags: Tags to filter models.
+
+        Returns:
+            List of ModelMetadata objects.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_model(self, name: str, project: str, commit: bool = True) -> None:
+        """
+        Delete a model from the registry.
+
+        Args:
+            name: The name of the model to delete.
+            project: The project the model belongs to.
+            commit: Whether the change should be persisted immediately
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def proto(self) -> RegistryProto:
         """

sdk/python/feast/infra/registry/proto_registry_utils.py

@@ -10,13 +10,15 @@
     EntityNotFoundException,
     FeatureServiceNotFoundException,
     FeatureViewNotFoundException,
+    ModelObjectNotFoundException,
     PermissionObjectNotFoundException,
     ProjectObjectNotFoundException,
     SavedDatasetNotFound,
     ValidationReferenceNotFound,
 )
 from feast.feature_service import FeatureService
 from feast.feature_view import FeatureView
+from feast.model import ModelMetadata
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.permissions.permission import Permission
 from feast.project import Project
@@ -393,3 +395,48 @@ def get_project(registry_proto: RegistryProto, name: str) -> Project:
         if projects_proto.spec.name == name:
             return Project.from_proto(projects_proto)
     raise ProjectObjectNotFoundException(name=name)
+
+
+@registry_proto_cache_with_tags
+def list_models(
+    registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
+) -> List[ModelMetadata]:
+    models = []
+    registry_proto_models = getattr(registry_proto, "model_metadata", [])
+    for model_proto in registry_proto_models:
+        if model_proto.project == project and utils.has_all_tags(
+            model_proto.tags, tags
+        ):
+            models.append(ModelMetadata.from_proto(model_proto))
+    return models
+
+
+def get_model(registry_proto: RegistryProto, name: str, project: str) -> ModelMetadata:
+    registry_proto_models = getattr(registry_proto, "model_metadata", [])
+    for model_proto in registry_proto_models:
+        if model_proto.project == project and model_proto.name == name:
+            return ModelMetadata.from_proto(model_proto)
+    raise ModelObjectNotFoundException(name=name, project=project)
+
+
+def apply_model(registry_proto: RegistryProto, model_metadata: ModelMetadata) -> None:
+    registry_proto_models = getattr(registry_proto, "model_metadata", [])
+    for i, existing_proto in enumerate(registry_proto_models):
+        if (
+            existing_proto.name == model_metadata.name
+            and existing_proto.project == model_metadata.project
+        ):
+            registry_proto.model_metadata[i].CopyFrom(model_metadata.to_proto())
+            return
+
+    registry_proto.model_metadata.append(model_metadata.to_proto())
+
+
+def delete_model(registry_proto: RegistryProto, name: str, project: str) -> None:
+    registry_proto_models = getattr(registry_proto, "model_metadata", [])
+    for i, model_proto in enumerate(registry_proto_models):
+        if model_proto.name == name and model_proto.project == project:
+            del registry_proto.model_metadata[i]
+            return
+
+    raise ModelObjectNotFoundException(name=name, project=project)