Skip to content

Commit 55c2f18

Browse files
committed
perf: Cache feature view resolution in get_online_features to reduce per-request overhead
Signed-off-by: ntkathole <[email protected]>
1 parent 1941d4d commit 55c2f18

5 files changed

Lines changed: 307 additions & 43 deletions

File tree

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -919,6 +919,16 @@ def refresh(self, project: Optional[str] = None):
919919
"""Refreshes the state of the registry cache by fetching the registry state from the remote registry store."""
920920
raise NotImplementedError
921921

922+
def is_cache_valid(self) -> bool:
923+
"""Check whether the registry's local cache is still within its TTL.
924+
925+
Returns True if cached data can be used without a refresh.
926+
Subclasses that support caching should override this.
927+
Registries without caching always return False (every read goes
928+
to the backing store).
929+
"""
930+
return False
931+
922932
# Lineage operations
923933
def get_registry_lineage(
924934
self,

‎sdk/python/feast/infra/registry/caching_registry.py‎

Lines changed: 18 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -442,32 +442,26 @@ def refresh(self, project: Optional[str] = None):
442442
except Exception as e:
443443
logger.debug(f"Error while refreshing registry: {e}", exc_info=True)
444444

445+
def is_cache_valid(self) -> bool:
446+
if (
447+
self.cached_registry_proto is None
448+
or self.cached_registry_proto == RegistryProto()
449+
):
450+
return False
451+
if (
452+
not hasattr(self, "cached_registry_proto_created")
453+
or self.cached_registry_proto_created is None
454+
):
455+
return False
456+
if self.cached_registry_proto_ttl.total_seconds() > 0 and _utc_now() > (
457+
self.cached_registry_proto_created + self.cached_registry_proto_ttl
458+
):
459+
return False
460+
return True
461+
445462
def _refresh_cached_registry_if_necessary(self):
446463
if self.cache_mode == "sync":
447-
448-
def is_cache_expired():
449-
if (
450-
self.cached_registry_proto is None
451-
or self.cached_registry_proto == RegistryProto()
452-
):
453-
return True
454-
455-
# Cache is expired if creation time is None
456-
if (
457-
not hasattr(self, "cached_registry_proto_created")
458-
or self.cached_registry_proto_created is None
459-
):
460-
return True
461-
462-
# Cache is expired if TTL > 0 and current time exceeds creation + TTL
463-
if self.cached_registry_proto_ttl.total_seconds() > 0 and _utc_now() > (
464-
self.cached_registry_proto_created + self.cached_registry_proto_ttl
465-
):
466-
return True
467-
468-
return False
469-
470-
if is_cache_expired():
464+
if not self.is_cache_valid():
471465
if not self._refresh_lock.acquire(blocking=False):
472466
logger.debug(
473467
"Skipping refresh if lock is already held by another thread"

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1287,6 +1287,17 @@ def refresh(self, project: Optional[str] = None):
12871287
"""Refreshes the state of the registry cache by fetching the registry state from the remote registry store."""
12881288
self._get_registry_proto(project=project, allow_cache=False)
12891289

1290+
def is_cache_valid(self) -> bool:
1291+
if self.cached_registry_proto_created is None:
1292+
return False
1293+
if (
1294+
self.cached_registry_proto_ttl.total_seconds() > 0
1295+
and _utc_now()
1296+
> self.cached_registry_proto_created + self.cached_registry_proto_ttl
1297+
):
1298+
return False
1299+
return True
1300+
12901301
def teardown(self):
12911302
"""Tears down (removes) the registry."""
12921303
self._registry_store.teardown()

‎sdk/python/feast/utils.py‎

Lines changed: 97 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import copy
22
import itertools
3+
import logging
34
import os
5+
import threading
46
import typing
57
import warnings
68
from collections import Counter, defaultdict
@@ -1380,6 +1382,96 @@ def _get_online_request_context(
13801382
)
13811383

13821384

1385+
_feature_resolution_cache: Dict[tuple, tuple] = {}
1386+
_feature_resolution_cache_lock = threading.Lock()
1387+
_feature_resolution_registry_ts: Optional[datetime] = None
1388+
1389+
_logger = logging.getLogger(__name__)
1390+
1391+
1392+
def _get_cached_request_context(
1393+
registry,
1394+
project: str,
1395+
features: Union[List[str], "FeatureService"],
1396+
full_feature_names: bool,
1397+
):
1398+
"""Return the output of _get_online_request_context plus resolved ODFV
1399+
entity join keys, using a cache that is invalidated whenever the
1400+
registry refreshes or its cache TTL expires."""
1401+
from feast.feature_service import FeatureService as _FS
1402+
1403+
global _feature_resolution_cache, _feature_resolution_registry_ts
1404+
1405+
registry_ts = getattr(registry, "cached_registry_proto_created", None)
1406+
1407+
if isinstance(features, _FS):
1408+
features_key: tuple = ("__fs__", features.name)
1409+
else:
1410+
features_key = tuple(features)
1411+
1412+
cache_key = (features_key, project, full_feature_names)
1413+
1414+
is_cache_valid = getattr(registry, "is_cache_valid", None)
1415+
registry_cache_valid = is_cache_valid() if callable(is_cache_valid) else False
1416+
1417+
if registry_ts is not None and registry_cache_valid:
1418+
with _feature_resolution_cache_lock:
1419+
if registry_ts != _feature_resolution_registry_ts:
1420+
_feature_resolution_cache.clear()
1421+
_feature_resolution_registry_ts = registry_ts
1422+
_logger.debug("Feature resolution cache cleared (registry refreshed)")
1423+
else:
1424+
cached = _feature_resolution_cache.get(cache_key)
1425+
if cached is not None:
1426+
return cached
1427+
1428+
ctx = _get_online_request_context(registry, project, features, full_feature_names)
1429+
1430+
(
1431+
feature_refs,
1432+
requested_on_demand_feature_views,
1433+
entity_name_to_join_key_map,
1434+
entity_type_map,
1435+
join_keys_set,
1436+
grouped_refs,
1437+
requested_result_row_names,
1438+
needed_request_data,
1439+
entityless_case,
1440+
) = ctx
1441+
1442+
odfv_join_keys: set = set()
1443+
for on_demand_feature_view in requested_on_demand_feature_views:
1444+
entities_for_odfv = getattr(on_demand_feature_view, "entities", [])
1445+
if len(entities_for_odfv) > 0 and isinstance(entities_for_odfv[0], str):
1446+
entities_for_odfv = [
1447+
registry.get_entity(entity_name, project, allow_cache=True)
1448+
for entity_name in entities_for_odfv
1449+
]
1450+
odfv_join_keys.update(entities_for_odfv)
1451+
1452+
result = (
1453+
feature_refs,
1454+
requested_on_demand_feature_views,
1455+
entity_name_to_join_key_map,
1456+
entity_type_map,
1457+
join_keys_set | odfv_join_keys,
1458+
grouped_refs,
1459+
frozenset(requested_result_row_names),
1460+
needed_request_data,
1461+
entityless_case,
1462+
)
1463+
1464+
registry_ts_after = getattr(registry, "cached_registry_proto_created", None)
1465+
if registry_ts_after is not None:
1466+
with _feature_resolution_cache_lock:
1467+
if registry_ts_after != _feature_resolution_registry_ts:
1468+
_feature_resolution_cache.clear()
1469+
_feature_resolution_registry_ts = registry_ts_after
1470+
_feature_resolution_cache[cache_key] = result
1471+
1472+
return result
1473+
1474+
13831475
def _prepare_entities_to_read_from_online_store(
13841476
registry,
13851477
project,
@@ -1399,10 +1491,13 @@ def _prepare_entities_to_read_from_online_store(
13991491
entity_type_map,
14001492
join_keys_set,
14011493
grouped_refs,
1402-
requested_result_row_names,
1494+
requested_result_row_names_frozen,
14031495
needed_request_data,
14041496
entityless_case,
1405-
) = _get_online_request_context(registry, project, features, full_feature_names)
1497+
) = _get_cached_request_context(registry, project, features, full_feature_names)
1498+
1499+
# Mutable copy — downstream code adds join keys to this set.
1500+
requested_result_row_names = set(requested_result_row_names_frozen)
14061501

14071502
# Extract Sequence from RepeatedValue Protobuf.
14081503
entity_value_lists: Dict[str, Union[List[Any], List[ValueProto]]] = {
@@ -1424,23 +1519,6 @@ def _prepare_entities_to_read_from_online_store(
14241519

14251520
num_rows = _validate_entity_values(entity_proto_values)
14261521

1427-
odfv_entities: List[Entity] = []
1428-
request_source_keys: List[str] = []
1429-
for on_demand_feature_view in requested_on_demand_feature_views:
1430-
entities_for_odfv = getattr(on_demand_feature_view, "entities", [])
1431-
if len(entities_for_odfv) > 0 and isinstance(entities_for_odfv[0], str):
1432-
entities_for_odfv = [
1433-
registry.get_entity(entity_name, project, allow_cache=True)
1434-
for entity_name in entities_for_odfv
1435-
]
1436-
odfv_entities.extend(entities_for_odfv)
1437-
for source in on_demand_feature_view.source_request_sources:
1438-
source_schema = on_demand_feature_view.source_request_sources[source].schema
1439-
for column in source_schema:
1440-
request_source_keys.append(column.name)
1441-
1442-
join_keys_set.update(set(odfv_entities))
1443-
14441522
join_key_values: Dict[str, List[ValueProto]] = {}
14451523
request_data_features: Dict[str, List[ValueProto]] = {}
14461524
# Entity rows may be either entities or request data.

0 commit comments

Comments
 (0)