Skip to content

Commit 147adad

Browse files
jiwidishuchu
jiwidi
and
shuchu
authored
fix: Key registry proto caches on the proto, not its id() (#7085)
* fix: Key registry proto caches on the proto, not its id() The proto_registry_utils caches keyed on id(registry_proto) and version_id. SqlRegistry leaves version_id empty. After a registry proto is freed, a new proto can get the same id, and the cache then returns objects from the old registry. For example, write_to_online_store validated a vector against a feature view from an earlier store and rejected valid rows. The caches now keep the proto and compare it by identity. Signed-off-by: jiwidi <[email protected]> * style: Fix the import block spacing in proto_registry_utils Signed-off-by: jiwidi <[email protected]> * test: Make the registry cache id-reuse test deterministic The test depended on CPython giving the second proto the freed address of the first. It now patches id() in the module so both live protos share one id, and it covers all three cache decorators. Signed-off-by: jiwidi <[email protected]> --------- Signed-off-by: jiwidi <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent 8ef6e75 commit 147adad

2 files changed

Lines changed: 75 additions & 16 deletions

File tree

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

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -28,28 +28,34 @@
2828
from feast.saved_dataset import SavedDataset, ValidationReference
2929
from feast.stream_feature_view import StreamFeatureView
3030

31+
# The caches hold the registry proto itself and compare by identity. An id() key
32+
# is unsafe: once a proto is freed, a new proto can reuse its id, and registries
33+
# that leave version_id empty (SqlRegistry) then read another registry's objects.
34+
3135

3236
def registry_proto_cache(func):
37+
cache_proto = None
3338
cache_key = None
3439
cache_value = None
3540

3641
@wraps(func)
3742
def wrapper(registry_proto: RegistryProto, project: str):
38-
nonlocal cache_key, cache_value
43+
nonlocal cache_proto, cache_key, cache_value
3944

40-
key = tuple([id(registry_proto), registry_proto.version_id, project])
45+
key = (registry_proto.version_id, project)
4146

42-
if key == cache_key:
47+
if registry_proto is cache_proto and key == cache_key:
4348
return cache_value
4449
else:
4550
cache_value = func(registry_proto, project)
46-
cache_key = key
51+
cache_proto, cache_key = registry_proto, key
4752
return cache_value
4853

4954
return wrapper
5055

5156

5257
def registry_proto_cache_with_tags(func):
58+
cache_proto = None
5359
cache_key = None
5460
cache_value = None
5561

@@ -60,42 +66,43 @@ def wrapper(
6066
tags: Optional[dict[str, str]],
6167
**kwargs,
6268
):
63-
nonlocal cache_key, cache_value
69+
nonlocal cache_proto, cache_key, cache_value
6470

6571
kwargs_key = tuple(sorted(kwargs.items())) if kwargs else ()
6672
# Snapshot the filter so caller mutations cannot change a cached key.
6773
tags_key = frozenset(tags.items()) if tags is not None else None
6874
key = (
69-
id(registry_proto),
7075
registry_proto.version_id,
7176
project,
7277
tags_key,
7378
kwargs_key,
7479
)
7580

76-
if key == cache_key:
81+
if registry_proto is cache_proto and key == cache_key:
7782
return cache_value
7883
else:
7984
cache_value = func(registry_proto, project, tags, **kwargs)
80-
cache_key = key
85+
cache_proto, cache_key = registry_proto, key
8186
return cache_value
8287

8388
return wrapper
8489

8590

8691
def registry_proto_cache_name(func):
87-
cache_proto_key = None
92+
cache_proto = None
93+
cache_version = None
8894
cache = {}
8995

9096
@wraps(func)
9197
def wrapper(registry_proto: RegistryProto, name: str, project: str):
92-
nonlocal cache_proto_key, cache
93-
94-
proto_key = (id(registry_proto), registry_proto.version_id)
98+
nonlocal cache_proto, cache_version, cache
9599

96-
if proto_key != cache_proto_key:
100+
if (
101+
registry_proto is not cache_proto
102+
or registry_proto.version_id != cache_version
103+
):
97104
cache = {}
98-
cache_proto_key = proto_key
105+
cache_proto, cache_version = registry_proto, registry_proto.version_id
99106

100107
key = (project, name)
101108
if key in cache:

‎sdk/python/tests/unit/infra/registry/test_proto_registry_utils.py‎

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,16 @@
1+
from typing import Callable
12
from unittest.mock import patch
23

3-
from feast import Entity
4-
from feast.infra.registry.proto_registry_utils import list_entities, list_saved_datasets
4+
import pytest
5+
6+
from feast import Entity, FeatureView, FileSource
7+
from feast.infra.registry import proto_registry_utils
8+
from feast.infra.registry.proto_registry_utils import (
9+
get_any_feature_view,
10+
list_entities,
11+
list_project_metadata,
12+
list_saved_datasets,
13+
)
514
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
615
from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto
716
from feast.protos.feast.core.SavedDataset_pb2 import (
@@ -35,6 +44,49 @@ def test_equivalent_tag_filters_reuse_cached_entities() -> None:
3544
from_proto.assert_called_once()
3645

3746

47+
def _registry(marker: str) -> RegistryProto:
48+
# SqlRegistry leaves version_id empty, so only the proto object tells them apart.
49+
proto = RegistryProto()
50+
proto.project_metadata.add(project="test_project", project_uuid=marker)
51+
entity = Entity(name="driver", value_type=ValueType.INT64, description=marker)
52+
proto.entities.append(entity.to_proto())
53+
proto.entities[0].spec.project = "test_project"
54+
view = FeatureView(
55+
name="driver_stats",
56+
source=FileSource(path="driver.parquet"),
57+
description=marker,
58+
)
59+
proto.feature_views.append(view.to_proto())
60+
proto.feature_views[0].spec.project = "test_project"
61+
return proto
62+
63+
64+
@pytest.mark.parametrize(
65+
"read",
66+
[
67+
lambda proto: list_project_metadata(proto, "test_project")[0].project_uuid,
68+
lambda proto: list_entities(proto, "test_project", None)[0].description,
69+
lambda proto: (
70+
get_any_feature_view(proto, "driver_stats", "test_project").description
71+
),
72+
],
73+
ids=[
74+
"registry_proto_cache",
75+
"registry_proto_cache_with_tags",
76+
"registry_proto_cache_name",
77+
],
78+
)
79+
def test_registries_that_share_an_id_are_not_served_from_each_others_cache(
80+
monkeypatch: pytest.MonkeyPatch, read: Callable[[RegistryProto], str]
81+
) -> None:
82+
# Give both protos the same id(), as CPython does when a freed address is reused.
83+
monkeypatch.setattr(proto_registry_utils, "id", lambda _proto: 0, raising=False)
84+
first, second = _registry("old"), _registry("new")
85+
86+
assert read(first) == "old"
87+
assert read(second) == "new"
88+
89+
3890
class TestRegistryProto:
3991
"""Test class for proto_registry_utils functions"""
4092

0 commit comments

Comments
 (0)