diff --git a/.cloud-build/cleanup/cleanup.py b/.cloud-build/cleanup/cleanup.py index 31d828ccf..9599dbef5 100644 --- a/.cloud-build/cleanup/cleanup.py +++ b/.cloud-build/cleanup/cleanup.py @@ -5,6 +5,8 @@ from resource_cleanup_manager import ( ModelResourceCleanupManager, EndpointResourceCleanupManager, ResourceCleanupManager, + MatchingEngineIndexEndpointResourceCleanupManager, + MatchingEngineIndexResourceCleanupManager, ) rate_limit = RateLimit(max_count=25, per=60, greedy=False) @@ -40,10 +42,12 @@ if is_dry_run: print("Starting cleanup in dry run mode...") # List of all cleanup managers -managers = [ +managers: List[ResourceCleanupManager] = [ DatasetResourceCleanupManager(), EndpointResourceCleanupManager(), ModelResourceCleanupManager(), # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion. + MatchingEngineIndexEndpointResourceCleanupManager(), + MatchingEngineIndexResourceCleanupManager(), ] run_cleanup_managers(managers=managers, is_dry_run=is_dry_run) diff --git a/.cloud-build/cleanup/resource_cleanup_manager.py b/.cloud-build/cleanup/resource_cleanup_manager.py index 5992e5e96..afb67de53 100644 --- a/.cloud-build/cleanup/resource_cleanup_manager.py +++ b/.cloud-build/cleanup/resource_cleanup_manager.py @@ -109,3 +109,11 @@ class EndpointResourceCleanupManager(VertexAIResourceCleanupManager): class ModelResourceCleanupManager(VertexAIResourceCleanupManager): vertex_ai_resource = aiplatform.Model + + +class MatchingEngineIndexResourceCleanupManager(VertexAIResourceCleanupManager): + vertex_ai_resource = aiplatform.MatchingEngineIndex + + +class MatchingEngineIndexEndpointResourceCleanupManager(VertexAIResourceCleanupManager): + vertex_ai_resource = aiplatform.MatchingEngineIndexEndpoint