Cleanup archived workers and models

arkindex/documents/management/commands/cleanup.py

@@ -14,10 +14,21 @@ from rq.utils import as_text
 from arkindex.documents.models import CorpusExport, CorpusExportState, Element
 from arkindex.images.models import Image, ImageServer
 from arkindex.ponos.models import Artifact, Task
-from arkindex.process.models import DataFile, GitRef, GitRefType, Process, WorkerVersion, WorkerVersionState
+from arkindex.process.models import (
+    CorpusWorkerVersion,
+    DataFile,
+    GitRef,
+    GitRefType,
+    Process,
+    Worker,
+    WorkerActivity,
+    WorkerRun,
+    WorkerVersion,
+    WorkerVersionState,
+)
 from arkindex.project.aws import s3
 from arkindex.project.rq_overrides import Job
-from arkindex.training.models import ModelVersion
+from arkindex.training.models import Model, ModelVersion
 from redis.exceptions import ConnectionError
 
 # Ponos artifacts use the path: <task id>/<path>
@@ -32,6 +43,9 @@ class Command(BaseCommand):
     help = "Clean up old corpus exports, trashed DataFiles, expired processes and S3 buckets"
 
     def handle(self, *args, **options):
+        # Cleaning up workers could free some artifacts, so clean them before artifacts
+        self.cleanup_archived_workers()
+
         self.cleanup_artifacts()
 
         self.cleanup_expired_processes()
@@ -48,6 +62,8 @@ class Command(BaseCommand):
 
         self.cleanup_ponos_logs()
 
+        self.cleanup_archived_models()
+
         self.cleanup_unlinked_model_versions()
 
         self.cleanup_rq_user_registries()
@@ -294,6 +310,71 @@ class Command(BaseCommand):
 
         self.stdout.write(self.style.SUCCESS("Successfully cleaned up orphaned Ponos logs."))
 
+    def cleanup_archived_workers(self):
+        """
+        Remove Worker instances that have been archived for longer than the configured worker cleanup delay
+        and that are not being used in any worker result.
+        """
+        self.stdout.write("Removing archived workers…")
+
+        workers = Worker.objects.filter(archived__lte=timezone.now() - timedelta(days=settings.WORKER_CLEANUP_DELAY))
+        skipped, deleted = 0, 0
+        for worker in workers.iterator():
+            # There are both foreign keys for worker versions and worker runs on worker results.
+            # Some old results might only have a worker version ID, but when a worker run ID is set,
+            # the worker version ID is deduced from it, so we only have to check on the version.
+            if worker.versions.all().in_use():
+                skipped += 1
+                continue
+
+            # Skip any workers whose WorkerConfigurations are in use.
+            # This should never happen since we already filter on the WorkerVersions,
+            # but that could lead to deleting worker results when we didn't want to.
+            if WorkerRun.objects.filter(configuration__worker=worker).in_use():
+                self.stdout.write(self.style.WARNING(
+                    f"Worker {worker.name} ({worker.id}) does not have any worker versions used by worker results, "
+                    "but some of its worker configurations are in use."
+                ))
+                continue
+
+            self.stdout.write(f"Removing worker {worker.name} ({worker.id})")
+            worker.delete()
+            deleted += 1
+
+        if skipped:
+            self.stdout.write(f"Skipping {skipped} archived workers that have worker versions or configurations used in worker results.")
+        self.stdout.write(self.style.SUCCESS(f"Successfully cleaned up {deleted} archived workers."))
+
+    def cleanup_archived_models(self):
+        """
+        Remove Model instances that have been archived for longer than the configured model cleanup delay
+        and that are not being used in any worker result.
+        """
+        self.stdout.write("Removing archived models…")
+        models = Model.objects.filter(archived__lte=timezone.now() - timedelta(days=settings.MODEL_CLEANUP_DELAY))
+
+        skipped, deleted = 0, 0
+        for model in models.iterator():
+            if WorkerRun.objects.filter(model_version__model=model).in_use():
+                skipped += 1
+                continue
+
+            self.stdout.write(f"Removing model {model.name} ({model.id})")
+
+            # Remove CorpusWorkerVersions and WorkerActivities first
+            # Those normally use SET_NULL, but this can cause the unique constraints to complain
+            # if there already are rows with a model version set to None.
+            WorkerActivity.objects.filter(model_version__model=model).delete()
+            CorpusWorkerVersion.objects.filter(model_version__model=model).delete()
+
+            model.delete()
+
+            deleted += 1
+
+        if skipped:
+            self.stdout.write(f"Skipping {skipped} archived models that have model versions used in worker results.")
+        self.stdout.write(self.style.SUCCESS(f"Successfully cleaned up {deleted} archived models."))
+
     def cleanup_unlinked_model_versions(self):
         self.stdout.write("Removing orphaned model versions archives…")
         bucket = s3.Bucket(settings.AWS_TRAINING_BUCKET)

arkindex/process/managers.py

@@ -130,6 +130,32 @@ class CorpusWorkerVersionManager(Manager):
             ], ignore_conflicts=True)
 
 
+class WorkerResultSourceQuerySet(QuerySet):
+
+    @property
+    def worker_result_relations(self):
+        from arkindex.process.models import WorkerRun
+        return (
+            field
+            for field in self.model._meta.get_fields()
+            if isinstance(field, ManyToOneRel)
+            # Ignore reverse links to processes tasks
+            and field.related_model is not Task
+            # Ignore the WorkerRun→WorkerVersion link
+            and field.related_model is not WorkerRun
+        )
+
+    def in_use(self):
+        """
+        Check if any data is linked to this worker run or worker version
+        """
+        ids = self.values_list("id", flat=True)
+        return any(
+            relation.related_model.objects.filter(**{f"{relation.field.name}__in": ids}).exists()
+            for relation in self.worker_result_relations
+        )
+
+
 class WorkerVersionManager(Manager):
 
     @cached_property
@@ -145,6 +171,9 @@ class WorkerVersionManager(Manager):
             .get(id=settings.IMPORTS_WORKER_VERSION)
         )
 
+    def get_queryset(self):
+        return WorkerResultSourceQuerySet(self.model, using=self._db)
+
 
 class WorkerManager(BaseACLManager):
 
@@ -172,25 +201,7 @@ class WorkerManager(BaseACLManager):
         ).distinct()
 
 
-class WorkerRunQuerySet(QuerySet):
-    def worker_results_models(self):
-        return [
-            field.related_model
-            for field in self.model._meta.get_fields()
-            if isinstance(field, ManyToOneRel)
-            # Ignore reverse links to processes tasks
-            and field.related_model is not Task
-        ]
-
-    def in_use(self):
-        """
-        Check if any data is linked to worker runs
-        """
-        ids = self.values_list("id", flat=True)
-        for field in self.worker_results_models():
-            if field.objects.filter(worker_run__in=ids).exists():
-                return True
-        return False
+class WorkerRunQuerySet(WorkerResultSourceQuerySet):
 
     def set_has_results(self):
         """
@@ -204,9 +215,9 @@ class WorkerRunQuerySet(QuerySet):
     def has_results_expression(self):
         return reduce(operator.or_, [
             Exists(
-                model.objects.filter(worker_run_id=OuterRef("pk"))
+                relation.related_model.objects.filter(worker_run_id=OuterRef("pk"))
             )
-            for model in self.worker_results_models()
+            for relation in self.worker_result_relations
         ])
 
 

arkindex/project/config.py

@@ -226,4 +226,8 @@ def get_settings_parser(base_dir):
     ingest_parser.add_option("extra_buckets", type=str, many=True, default=[])
     ingest_parser.add_option("prefix_by_bucket_name", type=bool, default=True)
 
+    cleanup_parser = parser.add_subparser("cleanup", default={})
+    cleanup_parser.add_option("worker_delay", type=int, default=30)
+    cleanup_parser.add_option("model_delay", type=int, default=30)
+
     return parser

arkindex/project/settings.py

@@ -566,6 +566,8 @@ FRONTEND_DOORBELL = conf["doorbell"]
 # Banner settings
 FRONTEND_BANNER = conf["banner"]
 
+WORKER_CLEANUP_DELAY = conf["cleanup"]["worker_delay"]
+MODEL_CLEANUP_DELAY = conf["cleanup"]["model_delay"]
 
 # Optional default group to add new users to when they complete email verification
 SIGNUP_DEFAULT_GROUP = conf["signup_default_group"]

arkindex/project/tests/config_samples/defaults.yaml

@@ -7,6 +7,9 @@ cache:
   path: null
   type: null
   url: null
+cleanup:
+  model_delay: 30
+  worker_delay: 30
 cors:
   origin_whitelist:
   - http://localhost:8080

arkindex/project/tests/config_samples/errors.yaml

@@ -4,6 +4,9 @@ banner:
   style: big
 cache:
   type: redis
+cleanup:
+  model_delay: forever
+  worker_delay: immediately
 cors:
   origin_whitelist: france
   suffixes: 1

arkindex/project/tests/config_samples/expected_errors.yaml

@@ -2,6 +2,9 @@ banner:
   style: "'big' is not a valid BannerStyle"
 cache:
   url: cache.url is required for a Redis cache
+cleanup:
+  model_delay: "invalid literal for int() with base 10: 'forever'"
+  worker_delay: "invalid literal for int() with base 10: 'immediately'"
 cors:
   suffixes: "'int' object is not iterable"
 csrf:

arkindex/project/tests/config_samples/override.yaml

@@ -8,6 +8,9 @@ cache:
   path: /
   type: filesystem
   url: http://aaa
+cleanup:
+  model_delay: 9998
+  worker_delay: 9999
 cors:
   origin_whitelist:
   - localtoast:1337