diff --git a/.cloud-build/cleanup/cleanup.py b/.cloud-build/cleanup/cleanup.py index 0430ac0a2..e82de544b 100644 --- a/.cloud-build/cleanup/cleanup.py +++ b/.cloud-build/cleanup/cleanup.py @@ -16,6 +16,7 @@ from resource_cleanup_manager import ( ResourceCleanupManager, MatchingEngineIndexEndpointResourceCleanupManager, MatchingEngineIndexResourceCleanupManager, + FeatureStoreLegacyCleanupManager, FeatureStoreCleanupManager, PipelineJobCleanupManager, TrainingJobCleanupManager, @@ -62,6 +63,7 @@ managers: List[ResourceCleanupManager] = [ ModelResourceCleanupManager(), # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion. MatchingEngineIndexEndpointResourceCleanupManager(), MatchingEngineIndexResourceCleanupManager(), + FeatureStoreLegacyCleanupManager(), FeatureStoreCleanupManager(), PipelineJobCleanupManager(), TrainingJobCleanupManager(), diff --git a/.cloud-build/cleanup/resource_cleanup_manager.py b/.cloud-build/cleanup/resource_cleanup_manager.py index 70f3d5ccc..477fdc14a 100644 --- a/.cloud-build/cleanup/resource_cleanup_manager.py +++ b/.cloud-build/cleanup/resource_cleanup_manager.py @@ -11,6 +11,7 @@ import abc from typing import Any, Type from google.cloud import aiplatform +from google.cloud import aiplatform_v1beta1 as v1beta1 from google.cloud.aiplatform import base from google.cloud import storage from proto.datetime_helpers import DatetimeWithNanoseconds @@ -129,19 +130,35 @@ class MatchingEngineIndexEndpointResourceCleanupManager(VertexAIResourceCleanupM resource.undeploy_all() resource.delete(force=True) -class FeatureStoreCleanupManager(VertexAIResourceCleanupManager): +class FeatureStoreLegacyCleanupManager(VertexAIResourceCleanupManager): # TODO: only deleting legacy # not deleting ingestions jobs + # ingest_from_xxx methods do not return a job ID, there is no list command, aka no python way to delete # not deleting batch serving jobs + # batch_serve_to_xxx methods do not return a job ID, there is no list command, aka no python way to delete vertex_ai_resource = aiplatform.Featurestore - # for FS 2.0 - # TODO: use _v1beta1, and gapic clients - # delete features, feature groups, feature views, feature online stores - def resource_name(self, resource: Any) -> str: return resource.name + +class FeatureStoreCleanupManager(VertexAIResourceCleanupManager): + # for FS 2.0 + # TODO: use _v1beta1, and gapic clients + # delete features, feature groups, feature views, feature online stores + vertex_ai_resource = v1beta1.FeatureOnlineStore + + def resource_name(self, resource: Any) -> str: + return resource.name + + def type_name(self) -> str: + return "FeatureOnlineStore" + + # TODO: FeatureOnlineStore has no list method + def list(self) -> Any: + return [] + + class PipelineJobCleanupManager(VertexAIResourceCleanupManager): vertex_ai_resource = aiplatform.PipelineJob