Skip to content

Commit

Permalink
feat: Apply cache to load proto registry for performance (#3702)
Browse files Browse the repository at this point in the history
* fix: Remove unused parameter

Signed-off-by: Jiwon Park <bakjeeone@hotmail.com>

* feat: Apply cache to load proto registry for performance

Signed-off-by: Jiwon Park <bakjeeone@hotmail.com>

* fix: Fix update key when cache missed

Signed-off-by: phil-park <bakjeeone@hotmail.com>

---------

Signed-off-by: Jiwon Park <bakjeeone@hotmail.com>
Signed-off-by: phil-park <bakjeeone@hotmail.com>
  • Loading branch information
phil-park authored Sep 12, 2023
1 parent 774ed33 commit 709c709
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion sdk/python/feast/infra/registry/proto_registry_utils.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import uuid
from functools import wraps
from typing import List, Optional

from feast import usage
Expand All @@ -23,6 +24,26 @@
from feast.stream_feature_view import StreamFeatureView


def registry_proto_cache(func):
cache_key = None
cache_value = None

@wraps(func)
def wrapper(registry_proto: RegistryProto, project: str):
nonlocal cache_key, cache_value

key = tuple([id(registry_proto), registry_proto.version_id, project])

if key == cache_key:
return cache_value
else:
cache_value = func(registry_proto, project)
cache_key = key
return cache_value

return wrapper


def init_project_metadata(cached_registry_proto: RegistryProto, project: str):
new_project_uuid = f"{uuid.uuid4()}"
usage.set_current_project_uuid(new_project_uuid)
Expand Down Expand Up @@ -137,8 +158,9 @@ def get_validation_reference(
raise ValidationReferenceNotFound(name, project=project)


@registry_proto_cache
def list_feature_services(
registry_proto: RegistryProto, project: str, allow_cache: bool = False
registry_proto: RegistryProto, project: str
) -> List[FeatureService]:
feature_services = []
for feature_service_proto in registry_proto.feature_services:
Expand All @@ -147,6 +169,7 @@ def list_feature_services(
return feature_services


@registry_proto_cache
def list_feature_views(
registry_proto: RegistryProto, project: str
) -> List[FeatureView]:
Expand All @@ -157,6 +180,7 @@ def list_feature_views(
return feature_views


@registry_proto_cache
def list_request_feature_views(
registry_proto: RegistryProto, project: str
) -> List[RequestFeatureView]:
Expand All @@ -169,6 +193,7 @@ def list_request_feature_views(
return feature_views


@registry_proto_cache
def list_stream_feature_views(
registry_proto: RegistryProto, project: str
) -> List[StreamFeatureView]:
Expand All @@ -181,6 +206,7 @@ def list_stream_feature_views(
return stream_feature_views


@registry_proto_cache
def list_on_demand_feature_views(
registry_proto: RegistryProto, project: str
) -> List[OnDemandFeatureView]:
Expand All @@ -193,6 +219,7 @@ def list_on_demand_feature_views(
return on_demand_feature_views


@registry_proto_cache
def list_entities(registry_proto: RegistryProto, project: str) -> List[Entity]:
entities = []
for entity_proto in registry_proto.entities:
Expand All @@ -201,6 +228,7 @@ def list_entities(registry_proto: RegistryProto, project: str) -> List[Entity]:
return entities


@registry_proto_cache
def list_data_sources(registry_proto: RegistryProto, project: str) -> List[DataSource]:
data_sources = []
for data_source_proto in registry_proto.data_sources:
Expand All @@ -209,6 +237,7 @@ def list_data_sources(registry_proto: RegistryProto, project: str) -> List[DataS
return data_sources


@registry_proto_cache
def list_saved_datasets(
registry_proto: RegistryProto, project: str
) -> List[SavedDataset]:
Expand All @@ -219,6 +248,7 @@ def list_saved_datasets(
return saved_datasets


@registry_proto_cache
def list_validation_references(
registry_proto: RegistryProto, project: str
) -> List[ValidationReference]:
Expand All @@ -231,6 +261,7 @@ def list_validation_references(
return validation_references


@registry_proto_cache
def list_project_metadata(
registry_proto: RegistryProto, project: str
) -> List[ProjectMetadata]:
Expand Down

0 comments on commit 709c709

Please sign in to comment.