Fixed cleanup scripts

.cloud-build/cleanup/cleanup.py

@@ -1,9 +1,13 @@
 from typing import List
+from ratemate import RateLimit
+from resource_cleanup_manager import (
+    DatasetResourceCleanupManager,
+    ModelResourceCleanupManager,
+    EndpointResourceCleanupManager,
+    ResourceCleanupManager,
+)
 
-from resource_cleanup_manager import (DatasetResourceCleanupManager,
-                                      EndpointResourceCleanupManager,
-                                      ModelResourceCleanupManager,
-                                      ResourceCleanupManager)
+rate_limit = RateLimit(max_count=25, per=60, greedy=False)
 
 
 def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: bool):
@@ -14,17 +18,18 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
         resources = manager.list()
         print(f"Found {len(resources)} {type_name}'s")
         for resource in resources:
-            if not manager.is_deletable(resource):
-                continue
+            try:
+                if not manager.is_deletable(resource):
+                    continue
 
-            if is_dry_run:
-                resource_name = manager.resource_name(resource)
-                print(f"Will delete '{type_name}': {resource_name}")
-            else:
-                try:
+                if is_dry_run:
+                    resource_name = manager.resource_name(resource)
+                    print(f"Will delete '{type_name}': {resource_name}")
+                else:
+                    rate_limit.wait()  # wait before deleting
                     manager.delete(resource)
-                except Exception as exception:
-                    print(exception)
+            except Exception as exception:
+                print(exception)
 
         print("")
 
@@ -38,7 +43,7 @@ if is_dry_run:
 managers = [
     DatasetResourceCleanupManager(),
     EndpointResourceCleanupManager(),
-    ModelResourceCleanupManager(),
+    ModelResourceCleanupManager(),  # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion.
 ]
 
 run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)

.cloud-build/cleanup/resource_cleanup_manager.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any
+from typing import Any, Type
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
@@ -41,7 +41,7 @@ class ResourceCleanupManager(abc.ABC):
         # Check that it wasn't created too recently, to prevent race conditions
         if time_difference <= RESOURCE_UPDATE_BUFFER_IN_SECONDS:
             print(
-                f"Skipping '{resource}' due update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
+                f"Skipping '{resource}' due to update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
             )
             return False
 
@@ -51,7 +51,7 @@ class ResourceCleanupManager(abc.ABC):
 class VertexAIResourceCleanupManager(ResourceCleanupManager):
     @property
     @abc.abstractmethod
-    def vertex_ai_resource(self) -> base.VertexAiResourceNounWithFutureManager:
+    def vertex_ai_resource(self) -> Type[base.VertexAiResourceNounWithFutureManager]:
         pass
 
     @property
@@ -61,7 +61,9 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
     def list(self) -> Any:
         return self.vertex_ai_resource.list()
 
-    def resource_name(self, resource: Any) -> str:
+    def resource_name(
+        self, resource: Type[base.VertexAiResourceNounWithFutureManager]
+    ) -> str:
         return resource.display_name
 
     def delete(self, resource):
@@ -75,12 +77,33 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
 
 class DatasetResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.datasets._Dataset
+    dataset_types = [
+        aiplatform.ImageDataset,
+        aiplatform.TabularDataset,
+        aiplatform.TextDataset,
+        aiplatform.TimeSeriesDataset,
+        aiplatform.VideoDataset,
+    ]
+
+    def list(self) -> Any:
+        return [
+            dataset
+            for dataset_type in self.dataset_types
+            for dataset in dataset_type.list()
+        ]
 
 
 class EndpointResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.Endpoint
 
     def delete(self, resource):
+        # TODO: Remove this once https://github.com/googleapis/python-aiplatform/issues/1441 is fixed
+        resource._sync_gca_resource()
+        for deployed_model_id in [
+            models.id for models in resource._gca_resource.deployed_models
+        ]:
+            resource._undeploy(deployed_model_id=deployed_model_id)
+
         resource.delete(force=True)