Skip to content

Commit 4e6a3f7

Browse files
felixwang9817felixwang9817
authored
Registry teardown (#1718)
* Teardown Registry during FeatureStore teardown Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Modify CLI teardown command to delegate to FeatureStore teardown method Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Ensure tests teardown FeatureStores Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Ensure tests teardown Registries Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Remove unnecessary return statements Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Move import Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Clarify teardown comment Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 13645c1 commit 4e6a3f7

7 files changed

Lines changed: 87 additions & 29 deletions

File tree

‎sdk/python/feast/feature_store.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -267,10 +267,7 @@ def teardown(self):
267267
entities = self.list_entities()
268268

269269
self._get_provider().teardown_infra(self.project, tables, entities)
270-
for feature_view in feature_views:
271-
self.delete_feature_view(feature_view.name)
272-
for feature_table in feature_tables:
273-
self._registry.delete_feature_table(feature_table.name, self.project)
270+
self._registry.teardown()
274271

275272
@log_exceptions_and_usage
276273
def get_historical_features(

‎sdk/python/feast/registry.py‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,6 @@ def __init__(self, registry_path: str, repo_path: Path, cache_ttl: timedelta):
7373
f"Registry path {registry_path} has unsupported scheme {uri.scheme}. Supported schemes are file and gs."
7474
)
7575
self.cached_registry_proto_ttl = cache_ttl
76-
return
7776

7877
def _initialize_registry(self):
7978
"""Explicitly initializes the registry with an empty proto."""
@@ -109,7 +108,6 @@ def apply_entity(self, entity: Entity, project: str, commit: bool = True):
109108
self.cached_registry_proto.entities.append(entity_proto)
110109
if commit:
111110
self.commit()
112-
return
113111

114112
def list_entities(self, project: str, allow_cache: bool = False) -> List[Entity]:
115113
"""
@@ -396,6 +394,10 @@ def refresh(self):
396394
"""Refreshes the state of the registry cache by fetching the registry state from the remote registry store."""
397395
self._get_registry_proto(allow_cache=False)
398396

397+
def teardown(self):
398+
"""Tears down (removes) the registry."""
399+
self._registry_store.teardown()
400+
399401
def _prepare_registry_for_changes(self):
400402
"""Prepares the Registry for changes by refreshing the cache if necessary."""
401403
try:
@@ -469,6 +471,13 @@ def update_registry_proto(self, registry_proto: RegistryProto):
469471
"""
470472
pass
471473

474+
@abstractmethod
475+
def teardown(self):
476+
"""
477+
Tear down all resources.
478+
"""
479+
pass
480+
472481

473482
class LocalRegistryStore(RegistryStore):
474483
def __init__(self, repo_path: Path, registry_path_string: str):
@@ -489,15 +498,21 @@ def get_registry_proto(self):
489498

490499
def update_registry_proto(self, registry_proto: RegistryProto):
491500
self._write_registry(registry_proto)
492-
return
501+
502+
def teardown(self):
503+
try:
504+
self._filepath.unlink()
505+
except FileNotFoundError:
506+
# If the file deletion fails with FileNotFoundError, the file has already
507+
# been deleted.
508+
pass
493509

494510
def _write_registry(self, registry_proto: RegistryProto):
495511
registry_proto.version_id = str(uuid.uuid4())
496512
registry_proto.last_updated.FromDatetime(datetime.utcnow())
497513
file_dir = self._filepath.parent
498514
file_dir.mkdir(exist_ok=True)
499515
self._filepath.write_bytes(registry_proto.SerializeToString())
500-
return
501516

502517

503518
class GCSRegistryStore(RegistryStore):
@@ -513,7 +528,6 @@ def __init__(self, uri: str):
513528
self._uri = urlparse(uri)
514529
self._bucket = self._uri.hostname
515530
self._blob = self._uri.path.lstrip("/")
516-
return
517531

518532
def get_registry_proto(self):
519533
from google.cloud import storage
@@ -540,7 +554,16 @@ def get_registry_proto(self):
540554

541555
def update_registry_proto(self, registry_proto: RegistryProto):
542556
self._write_registry(registry_proto)
543-
return
557+
558+
def teardown(self):
559+
from google.cloud.exceptions import NotFound
560+
561+
gs_bucket = self.gcs_client.get_bucket(self._bucket)
562+
try:
563+
gs_bucket.delete_blob(self._blob)
564+
except NotFound:
565+
# If the blob deletion fails with NotFound, it has already been deleted.
566+
pass
544567

545568
def _write_registry(self, registry_proto: RegistryProto):
546569
registry_proto.version_id = str(uuid.uuid4())
@@ -552,7 +575,6 @@ def _write_registry(self, registry_proto: RegistryProto):
552575
file_obj.write(registry_proto.SerializeToString())
553576
file_obj.seek(0)
554577
blob.upload_from_file(file_obj)
555-
return
556578

557579

558580
class S3RegistryStore(RegistryStore):
@@ -605,7 +627,9 @@ def get_registry_proto(self):
605627

606628
def update_registry_proto(self, registry_proto: RegistryProto):
607629
self._write_registry(registry_proto)
608-
return
630+
631+
def teardown(self):
632+
self.s3_client.Object(self._bucket, self._key).delete()
609633

610634
def _write_registry(self, registry_proto: RegistryProto):
611635
registry_proto.version_id = str(uuid.uuid4())

‎sdk/python/feast/repo_operations.py‎

Lines changed: 4 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from click.exceptions import BadParameter
1313

1414
from feast import Entity, FeatureTable
15+
from feast.feature_store import FeatureStore
1516
from feast.feature_view import FeatureView
1617
from feast.inference import (
1718
update_data_sources_with_inferred_event_timestamp_col,
@@ -242,23 +243,9 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
242243

243244
@log_exceptions_and_usage
244245
def teardown(repo_config: RepoConfig, repo_path: Path):
245-
registry_config = repo_config.get_registry_config()
246-
registry = Registry(
247-
registry_path=registry_config.path,
248-
repo_path=repo_path,
249-
cache_ttl=timedelta(seconds=registry_config.cache_ttl_seconds),
250-
)
251-
project = repo_config.project
252-
registry_tables: List[Union[FeatureTable, FeatureView]] = []
253-
registry_tables.extend(registry.list_feature_tables(project=project))
254-
registry_tables.extend(registry.list_feature_views(project=project))
255-
256-
registry_entities: List[Entity] = registry.list_entities(project=project)
257-
258-
infra_provider = get_provider(repo_config, repo_path)
259-
infra_provider.teardown_infra(
260-
project, tables=registry_tables, entities=registry_entities
261-
)
246+
# Cannot pass in both repo_path and repo_config to FeatureStore.
247+
feature_store = FeatureStore(repo_path=repo_path, config=None)
248+
feature_store.teardown()
262249

263250

264251
@log_exceptions_and_usage

‎sdk/python/tests/integration/materialization/test_offline_online_store_consistency.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,8 @@ def prep_redshift_fs_and_fv(
181181

182182
yield fs, fv
183183

184+
fs.teardown()
185+
184186
# Clean up the uploaded Redshift table
185187
aws_utils.execute_redshift_statement(
186188
client,

‎sdk/python/tests/integration/offline_store/test_historical_retrieval.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,8 @@ def test_historical_features_from_parquet_sources(
344344
).reset_index(drop=True),
345345
)
346346

347+
store.teardown()
348+
347349

348350
@pytest.mark.integration
349351
@pytest.mark.parametrize(
@@ -596,6 +598,8 @@ def test_historical_features_from_bigquery_sources(
596598
actual_df_from_df_entities, table_from_df_entities.to_pandas()
597599
)
598600

601+
store.teardown()
602+
599603

600604
@pytest.mark.integration
601605
def test_timestamp_bound_inference_from_entity_df_using_bigquery():

‎sdk/python/tests/integration/registration/test_feature_store.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ def test_apply_entity_success(test_feature_store):
113113
and entity.labels["team"] == "matchmaking"
114114
)
115115

116+
test_feature_store.teardown()
117+
116118

117119
@pytest.mark.integration
118120
@pytest.mark.parametrize(
@@ -154,6 +156,8 @@ def test_apply_entity_integration(test_feature_store):
154156
and entity.labels["team"] == "matchmaking"
155157
)
156158

159+
test_feature_store.teardown()
160+
157161

158162
@pytest.mark.parametrize(
159163
"test_feature_store", [lazy_fixture("feature_store_with_local_registry")],
@@ -202,6 +206,8 @@ def test_apply_feature_view_success(test_feature_store):
202206
and feature_views[0].entities[0] == "fs1_my_entity_1"
203207
)
204208

209+
test_feature_store.teardown()
210+
205211

206212
@pytest.mark.integration
207213
@pytest.mark.parametrize(
@@ -266,6 +272,8 @@ def test_feature_view_inference_success(test_feature_store, dataframe_source):
266272
== actual_bq_using_query_arg_source
267273
)
268274

275+
test_feature_store.teardown()
276+
269277

270278
@pytest.mark.integration
271279
@pytest.mark.parametrize(
@@ -337,6 +345,8 @@ def test_apply_feature_view_integration(test_feature_store):
337345
feature_views = test_feature_store.list_feature_views()
338346
assert len(feature_views) == 0
339347

348+
test_feature_store.teardown()
349+
340350

341351
@pytest.mark.parametrize(
342352
"test_feature_store", [lazy_fixture("feature_store_with_local_registry")],
@@ -398,6 +408,8 @@ def test_apply_object_and_read(test_feature_store):
398408
assert fv2 != fv1_actual
399409
assert e2 != e1_actual
400410

411+
test_feature_store.teardown()
412+
401413

402414
def test_apply_remote_repo():
403415
fd, registry_path = mkstemp()
@@ -466,3 +478,5 @@ def test_reapply_feature_view_success(test_feature_store, dataframe_source):
466478
# Check Feature View
467479
fv_stored = test_feature_store.get_feature_view(fv1.name)
468480
assert len(fv_stored.materialization_intervals) == 0
481+
482+
test_feature_store.teardown()

‎sdk/python/tests/integration/registration/test_registry.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,12 @@ def test_apply_entity_success(test_registry):
9696
and entity.labels["team"] == "matchmaking"
9797
)
9898

99+
test_registry.teardown()
100+
101+
# Will try to reload registry, which will fail because the file has been deleted
102+
with pytest.raises(FileNotFoundError):
103+
test_registry._get_registry_proto()
104+
99105

100106
@pytest.mark.integration
101107
@pytest.mark.parametrize(
@@ -135,6 +141,12 @@ def test_apply_entity_integration(test_registry):
135141
and entity.labels["team"] == "matchmaking"
136142
)
137143

144+
test_registry.teardown()
145+
146+
# Will try to reload registry, which will fail because the file has been deleted
147+
with pytest.raises(FileNotFoundError):
148+
test_registry._get_registry_proto()
149+
138150

139151
@pytest.mark.parametrize(
140152
"test_registry", [lazy_fixture("local_registry")],
@@ -203,6 +215,12 @@ def test_apply_feature_view_success(test_registry):
203215
feature_views = test_registry.list_feature_views(project)
204216
assert len(feature_views) == 0
205217

218+
test_registry.teardown()
219+
220+
# Will try to reload registry, which will fail because the file has been deleted
221+
with pytest.raises(FileNotFoundError):
222+
test_registry._get_registry_proto()
223+
206224

207225
@pytest.mark.integration
208226
@pytest.mark.parametrize(
@@ -272,6 +290,12 @@ def test_apply_feature_view_integration(test_registry):
272290
feature_views = test_registry.list_feature_views(project)
273291
assert len(feature_views) == 0
274292

293+
test_registry.teardown()
294+
295+
# Will try to reload registry, which will fail because the file has been deleted
296+
with pytest.raises(FileNotFoundError):
297+
test_registry._get_registry_proto()
298+
275299

276300
def test_commit():
277301
fd, registry_path = mkstemp()
@@ -345,3 +369,9 @@ def test_commit():
345369
and "team" in entity.labels
346370
and entity.labels["team"] == "matchmaking"
347371
)
372+
373+
test_registry.teardown()
374+
375+
# Will try to reload registry, which will fail because the file has been deleted
376+
with pytest.raises(FileNotFoundError):
377+
test_registry._get_registry_proto()

0 commit comments

Comments
 (0)