Skip to content

Commit 2918073

Browse files
committed
feat: Added Model Metadata support in Registry
Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent 648c53d commit 2918073

19 files changed

+601
-14
lines changed

README.md

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
[![PyPI - Downloads](https://img.shields.io/pypi/dm/feast)](https://pypi.org/project/feast/)
1111
[![GitHub contributors](https://img.shields.io/github/contributors/feast-dev/feast)](https://github.com/feast-dev/feast/graphs/contributors)
12-
[![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)
12+
[![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)
1313
[![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)
1414
[![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)
1515
[![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 @@
2121
## Join us on Slack!
2222
👋👋👋 [Come say hi on Slack!](https://communityinviter.com/apps/feastopensource/feast-the-open-source-feature-store)
2323

24+
[Check out our DeepWiki!](https://deepwiki.com/feast-dev/feast)
25+
2426
## Overview
2527

2628
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:
257259

258260
<a href="https://github.com/feast-dev/feast/graphs/contributors">
259261
<img src="https://contrib.rocks/image?repo=feast-dev/feast" />
260-
</a>
262+
</a>

protos/feast/core/Model.proto

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
syntax = "proto3";
2+
package feast.core;
3+
4+
option go_package = "github.com/feast-dev/feast/go/protos/feast/core";
5+
option java_outer_classname = "ModelMetadataProto";
6+
option java_package = "feast.proto.core";
7+
8+
import "google/protobuf/timestamp.proto";
9+
10+
message ModelMetadata {
11+
// Required: Unique model name
12+
string name = 1;
13+
14+
// Optional: Associated feature view name (if used)
15+
repeated string feature_view = 2;
16+
17+
// Optional: Associated feature service name (if used)
18+
repeated string feature_service = 3;
19+
20+
// Optional: Direct list of fully qualified feature references
21+
repeated string features = 4;
22+
23+
// Project this model belongs to
24+
string project = 5;
25+
26+
// Metadata (e.g., model purpose, version, tags)
27+
map<string, string> tags = 6;
28+
29+
// Timestamp when model was trained
30+
google.protobuf.Timestamp training_timestamp = 7;
31+
32+
// Optional: Description of the model
33+
string description = 8;
34+
}

protos/feast/core/Registry.proto

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,9 @@ import "feast/core/ValidationProfile.proto";
3434
import "google/protobuf/timestamp.proto";
3535
import "feast/core/Permission.proto";
3636
import "feast/core/Project.proto";
37+
import "feast/core/Model.proto";
3738

38-
// Next id: 18
39+
// Next id: 19
3940
message Registry {
4041
repeated Entity entities = 1;
4142
repeated FeatureTable feature_tables = 2;
@@ -55,6 +56,7 @@ message Registry {
5556
google.protobuf.Timestamp last_updated = 5;
5657
repeated Permission permissions = 16;
5758
repeated Project projects = 17;
59+
repeated ModelMetadata model_metadata = 18;
5860
}
5961

6062
message ProjectMetadata {

sdk/python/feast/diff/registry_diff.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
FeatureService as FeatureServiceProto,
1919
)
2020
from feast.protos.feast.core.FeatureView_pb2 import FeatureView as FeatureViewProto
21+
from feast.protos.feast.core.Model_pb2 import ModelMetadata as ModelMetadataProto
2122
from feast.protos.feast.core.OnDemandFeatureView_pb2 import (
2223
OnDemandFeatureView as OnDemandFeatureViewProto,
2324
)
@@ -133,8 +134,10 @@ def diff_registry_objects(
133134
current_spec: FeastObjectSpecProto
134135
new_spec: FeastObjectSpecProto
135136
if isinstance(
136-
current_proto, (DataSourceProto, ValidationReferenceProto)
137-
) or isinstance(new_proto, (DataSourceProto, ValidationReferenceProto)):
137+
current_proto, (DataSourceProto, ValidationReferenceProto, ModelMetadataProto)
138+
) or isinstance(
139+
new_proto, (DataSourceProto, ValidationReferenceProto, ModelMetadataProto)
140+
):
138141
assert type(current_proto) == type(new_proto)
139142
current_spec = cast(DataSourceProto, current_proto)
140143
new_spec = cast(DataSourceProto, new_proto)

sdk/python/feast/errors.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,3 +522,8 @@ def grpc_status_code(self) -> "GrpcStatusCode":
522522

523523
def http_status_code(self) -> int:
524524
return HttpStatusCode.HTTP_403_FORBIDDEN
525+
526+
527+
class ModelObjectNotFoundException(Exception):
528+
def __init__(self, name: str, project: str):
529+
super().__init__(f"Model {name} not found in project {project}")

sdk/python/feast/feast_object.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from .entity import Entity
99
from .feature_service import FeatureService
1010
from .feature_view import FeatureView
11+
from .model import ModelMetadata
1112
from .on_demand_feature_view import OnDemandFeatureView
1213
from .permissions.permission import Permission
1314
from .protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
@@ -37,6 +38,7 @@
3738
ValidationReference,
3839
SavedDataset,
3940
Permission,
41+
ModelMetadata,
4042
]
4143

4244
FeastObjectSpecProto = Union[

sdk/python/feast/feature_store.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
from feast.infra.registry.base_registry import BaseRegistry
7272
from feast.infra.registry.registry import Registry
7373
from feast.infra.registry.sql import SqlRegistry
74+
from feast.model import ModelMetadata
7475
from feast.on_demand_feature_view import OnDemandFeatureView
7576
from feast.online_response import OnlineResponse
7677
from feast.permissions.permission import Permission
@@ -809,6 +810,7 @@ def apply(
809810
FeatureService,
810811
ValidationReference,
811812
Permission,
813+
ModelMetadata,
812814
List[FeastObject],
813815
],
814816
objects_to_delete: Optional[List[FeastObject]] = None,
@@ -924,6 +926,10 @@ def apply(
924926

925927
data_sources_to_update = list(data_sources_set_to_update)
926928

929+
model_metadata_to_update: List[ModelMetadata] = [
930+
ob for ob in objects if isinstance(ob, ModelMetadata)
931+
]
932+
927933
# Handle all entityless feature views by using DUMMY_ENTITY as a placeholder entity.
928934
entities_to_update.append(DUMMY_ENTITY)
929935

@@ -963,11 +969,16 @@ def apply(
963969
self._registry.apply_permission(
964970
permission, project=self.project, commit=False
965971
)
972+
for model_metadata in model_metadata_to_update:
973+
self._registry.apply_model(
974+
model_metadata, project=self.project, commit=False
975+
)
966976

967977
entities_to_delete = []
968978
views_to_delete = []
969979
sfvs_to_delete = []
970980
permissions_to_delete = []
981+
model_metadata_to_delete = []
971982
if not partial:
972983
# Delete all registry objects that should not exist.
973984
entities_to_delete = [
@@ -999,6 +1010,9 @@ def apply(
9991010
permissions_to_delete = [
10001011
ob for ob in objects_to_delete if isinstance(ob, Permission)
10011012
]
1013+
model_metadata_to_delete = [
1014+
ob for ob in objects_to_delete if isinstance(ob, ModelMetadata)
1015+
]
10021016

10031017
for data_source in data_sources_to_delete:
10041018
self._registry.delete_data_source(
@@ -1032,6 +1046,10 @@ def apply(
10321046
self._registry.delete_permission(
10331047
permission.name, project=self.project, commit=False
10341048
)
1049+
for model_metadata in model_metadata_to_delete:
1050+
self._registry.delete_model(
1051+
model_metadata.name, project=self.project, commit=False
1052+
)
10351053

10361054
tables_to_delete: List[FeatureView] = (
10371055
views_to_delete + sfvs_to_delete if not partial else [] # type: ignore
@@ -2531,6 +2549,55 @@ def list_saved_datasets(
25312549
self.project, allow_cache=allow_cache, tags=tags
25322550
)
25332551

2552+
def get_model(self, name: str, allow_cache: bool = False) -> ModelMetadata:
2553+
"""
2554+
Retrieves a specific model by ID.
2555+
2556+
Args:
2557+
name: The name of the model.
2558+
allow_cache: Whether to allow returning from cache.
2559+
2560+
Returns:
2561+
The ModelMetadata for the specified model.
2562+
2563+
Raises:
2564+
ModelObjectNotFoundException: If the model does not exist.
2565+
"""
2566+
return self._registry.get_model(
2567+
name=name, project=self.project, allow_cache=allow_cache
2568+
)
2569+
2570+
def list_models(
2571+
self,
2572+
allow_cache: bool = False,
2573+
tags: Optional[dict[str, str]] = None,
2574+
) -> List[ModelMetadata]:
2575+
"""
2576+
List models stored in the registry.
2577+
2578+
Args:
2579+
allow_cache: Whether to allow returning from a cached registry.
2580+
tags: Optional dictionary to filter models by tags.
2581+
2582+
Returns:
2583+
A list of ModelMetadata objects.
2584+
"""
2585+
return self._registry.list_models(
2586+
allow_cache=allow_cache,
2587+
project=self.project,
2588+
tags=tags,
2589+
)
2590+
2591+
def delete_model(self, name: str) -> None:
2592+
"""
2593+
Deletes a model from the registry.
2594+
2595+
Args:
2596+
name: The name of the model to delete.
2597+
feast_project: The project the model belongs to (defaults to current project).
2598+
"""
2599+
self._registry.delete_model(name=name, project=self.project)
2600+
25342601
async def initialize(self) -> None:
25352602
"""Initialize long-lived clients and/or resources needed for accessing datastores"""
25362603
await self._get_provider().initialize(self.config)

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from feast.feature_service import FeatureService
2828
from feast.feature_view import FeatureView
2929
from feast.infra.infra_object import Infra
30+
from feast.model import ModelMetadata
3031
from feast.on_demand_feature_view import OnDemandFeatureView
3132
from feast.permissions.permission import Permission
3233
from feast.project import Project
@@ -768,6 +769,71 @@ def list_projects(
768769
"""
769770
raise NotImplementedError
770771

772+
@abstractmethod
773+
def apply_model(
774+
self, model_metadata: ModelMetadata, project: str, commit: bool = True
775+
) -> None:
776+
"""
777+
Register or update model metadata in the registry.
778+
779+
Args:
780+
model_metadata: The model metadata object to apply.
781+
project: Feast project that this entity belongs to
782+
commit: Whether the change should be persisted immediately
783+
"""
784+
raise NotImplementedError
785+
786+
@abstractmethod
787+
def get_model(
788+
self, name: str, project: str, allow_cache: bool = False
789+
) -> ModelMetadata:
790+
"""
791+
Retrieve a model from the registry by model ID.
792+
793+
Args:
794+
name: The name of the model to retrieve.
795+
allow_cache: Whether to allow returning model from cached registry.
796+
797+
Returns:
798+
ModelMetadata
799+
800+
Raises:
801+
ModelObjectNotFoundException: If the model is not found.
802+
"""
803+
raise NotImplementedError
804+
805+
@abstractmethod
806+
def list_models(
807+
self,
808+
project: str,
809+
allow_cache: bool = False,
810+
tags: Optional[Dict[str, str]] = None,
811+
) -> List[ModelMetadata]:
812+
"""
813+
List all models registered in the registry.
814+
815+
Args:
816+
allow_cache: Whether to allow returning models from cached registry.
817+
project: Project name to filter models.
818+
tags: Tags to filter models.
819+
820+
Returns:
821+
List of ModelMetadata objects.
822+
"""
823+
raise NotImplementedError
824+
825+
@abstractmethod
826+
def delete_model(self, name: str, project: str, commit: bool = True) -> None:
827+
"""
828+
Delete a model from the registry.
829+
830+
Args:
831+
name: The name of the model to delete.
832+
project: The project the model belongs to.
833+
commit: Whether the change should be persisted immediately
834+
"""
835+
raise NotImplementedError
836+
771837
@abstractmethod
772838
def proto(self) -> RegistryProto:
773839
"""

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,15 @@
1010
EntityNotFoundException,
1111
FeatureServiceNotFoundException,
1212
FeatureViewNotFoundException,
13+
ModelObjectNotFoundException,
1314
PermissionObjectNotFoundException,
1415
ProjectObjectNotFoundException,
1516
SavedDatasetNotFound,
1617
ValidationReferenceNotFound,
1718
)
1819
from feast.feature_service import FeatureService
1920
from feast.feature_view import FeatureView
21+
from feast.model import ModelMetadata
2022
from feast.on_demand_feature_view import OnDemandFeatureView
2123
from feast.permissions.permission import Permission
2224
from feast.project import Project
@@ -393,3 +395,48 @@ def get_project(registry_proto: RegistryProto, name: str) -> Project:
393395
if projects_proto.spec.name == name:
394396
return Project.from_proto(projects_proto)
395397
raise ProjectObjectNotFoundException(name=name)
398+
399+
400+
@registry_proto_cache_with_tags
401+
def list_models(
402+
registry_proto: RegistryProto, project: str, tags: Optional[dict[str, str]]
403+
) -> List[ModelMetadata]:
404+
models = []
405+
registry_proto_models = getattr(registry_proto, "model_metadata", [])
406+
for model_proto in registry_proto_models:
407+
if model_proto.project == project and utils.has_all_tags(
408+
model_proto.tags, tags
409+
):
410+
models.append(ModelMetadata.from_proto(model_proto))
411+
return models
412+
413+
414+
def get_model(registry_proto: RegistryProto, name: str, project: str) -> ModelMetadata:
415+
registry_proto_models = getattr(registry_proto, "model_metadata", [])
416+
for model_proto in registry_proto_models:
417+
if model_proto.project == project and model_proto.name == name:
418+
return ModelMetadata.from_proto(model_proto)
419+
raise ModelObjectNotFoundException(name=name, project=project)
420+
421+
422+
def apply_model(registry_proto: RegistryProto, model_metadata: ModelMetadata) -> None:
423+
registry_proto_models = getattr(registry_proto, "model_metadata", [])
424+
for i, existing_proto in enumerate(registry_proto_models):
425+
if (
426+
existing_proto.name == model_metadata.name
427+
and existing_proto.project == model_metadata.project
428+
):
429+
registry_proto.model_metadata[i].CopyFrom(model_metadata.to_proto())
430+
return
431+
432+
registry_proto.model_metadata.append(model_metadata.to_proto())
433+
434+
435+
def delete_model(registry_proto: RegistryProto, name: str, project: str) -> None:
436+
registry_proto_models = getattr(registry_proto, "model_metadata", [])
437+
for i, model_proto in enumerate(registry_proto_models):
438+
if model_proto.name == name and model_proto.project == project:
439+
del registry_proto.model_metadata[i]
440+
return
441+
442+
raise ModelObjectNotFoundException(name=name, project=project)

0 commit comments

Comments
 (0)