Skip to content

Commit 709c709

Browse files
authored
feat: Apply cache to load proto registry for performance (#3702)
* 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>
1 parent 774ed33 commit 709c709

File tree

1 file changed

+32
-1
lines changed

1 file changed

+32
-1
lines changed

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

+32-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import uuid
2+
from functools import wraps
23
from typing import List, Optional
34

45
from feast import usage
@@ -23,6 +24,26 @@
2324
from feast.stream_feature_view import StreamFeatureView
2425

2526

27+
def registry_proto_cache(func):
28+
cache_key = None
29+
cache_value = None
30+
31+
@wraps(func)
32+
def wrapper(registry_proto: RegistryProto, project: str):
33+
nonlocal cache_key, cache_value
34+
35+
key = tuple([id(registry_proto), registry_proto.version_id, project])
36+
37+
if key == cache_key:
38+
return cache_value
39+
else:
40+
cache_value = func(registry_proto, project)
41+
cache_key = key
42+
return cache_value
43+
44+
return wrapper
45+
46+
2647
def init_project_metadata(cached_registry_proto: RegistryProto, project: str):
2748
new_project_uuid = f"{uuid.uuid4()}"
2849
usage.set_current_project_uuid(new_project_uuid)
@@ -137,8 +158,9 @@ def get_validation_reference(
137158
raise ValidationReferenceNotFound(name, project=project)
138159

139160

161+
@registry_proto_cache
140162
def list_feature_services(
141-
registry_proto: RegistryProto, project: str, allow_cache: bool = False
163+
registry_proto: RegistryProto, project: str
142164
) -> List[FeatureService]:
143165
feature_services = []
144166
for feature_service_proto in registry_proto.feature_services:
@@ -147,6 +169,7 @@ def list_feature_services(
147169
return feature_services
148170

149171

172+
@registry_proto_cache
150173
def list_feature_views(
151174
registry_proto: RegistryProto, project: str
152175
) -> List[FeatureView]:
@@ -157,6 +180,7 @@ def list_feature_views(
157180
return feature_views
158181

159182

183+
@registry_proto_cache
160184
def list_request_feature_views(
161185
registry_proto: RegistryProto, project: str
162186
) -> List[RequestFeatureView]:
@@ -169,6 +193,7 @@ def list_request_feature_views(
169193
return feature_views
170194

171195

196+
@registry_proto_cache
172197
def list_stream_feature_views(
173198
registry_proto: RegistryProto, project: str
174199
) -> List[StreamFeatureView]:
@@ -181,6 +206,7 @@ def list_stream_feature_views(
181206
return stream_feature_views
182207

183208

209+
@registry_proto_cache
184210
def list_on_demand_feature_views(
185211
registry_proto: RegistryProto, project: str
186212
) -> List[OnDemandFeatureView]:
@@ -193,6 +219,7 @@ def list_on_demand_feature_views(
193219
return on_demand_feature_views
194220

195221

222+
@registry_proto_cache
196223
def list_entities(registry_proto: RegistryProto, project: str) -> List[Entity]:
197224
entities = []
198225
for entity_proto in registry_proto.entities:
@@ -201,6 +228,7 @@ def list_entities(registry_proto: RegistryProto, project: str) -> List[Entity]:
201228
return entities
202229

203230

231+
@registry_proto_cache
204232
def list_data_sources(registry_proto: RegistryProto, project: str) -> List[DataSource]:
205233
data_sources = []
206234
for data_source_proto in registry_proto.data_sources:
@@ -209,6 +237,7 @@ def list_data_sources(registry_proto: RegistryProto, project: str) -> List[DataS
209237
return data_sources
210238

211239

240+
@registry_proto_cache
212241
def list_saved_datasets(
213242
registry_proto: RegistryProto, project: str
214243
) -> List[SavedDataset]:
@@ -219,6 +248,7 @@ def list_saved_datasets(
219248
return saved_datasets
220249

221250

251+
@registry_proto_cache
222252
def list_validation_references(
223253
registry_proto: RegistryProto, project: str
224254
) -> List[ValidationReference]:
@@ -231,6 +261,7 @@ def list_validation_references(
231261
return validation_references
232262

233263

264+
@registry_proto_cache
234265
def list_project_metadata(
235266
registry_proto: RegistryProto, project: str
236267
) -> List[ProjectMetadata]:

0 commit comments

Comments
 (0)