Compare commits

...
2 changed files with 13 additions and 1 deletions
+5 -1
View File
@@ -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)
@@ -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