Cleanup unused worker types in cleanup command

arkindex/documents/management/commands/cleanup.py

@@ -7,7 +7,7 @@ import django_rq
 from botocore.exceptions import ClientError
 from django.conf import settings
 from django.core.management.base import BaseCommand
-from django.db.models import Exists, Max, OuterRef
+from django.db.models import Count, Exists, Max, OuterRef
 from django.utils import timezone
 from rq.utils import as_text
 
@@ -21,6 +21,7 @@ from arkindex.process.models import (
     Worker,
     WorkerActivity,
     WorkerRun,
+    WorkerType,
 )
 from arkindex.project.aws import s3
 from arkindex.project.rq_overrides import Job
@@ -64,6 +65,8 @@ class Command(BaseCommand):
 
         self.cleanup_rq_user_registries()
 
+        self.cleanup_worker_types()
+
     def cleanup_artifacts(self):
         """
         Remove all Ponos artifacts that are not tied to a Process
@@ -424,3 +427,16 @@ class Command(BaseCommand):
             user_count += 1
 
         self.stdout.write(self.style.SUCCESS(f"Successfully cleaned up {deleted_jobs} deleted jobs from {user_count} RQ user registries."))
+
+    def cleanup_worker_types(self):
+        """
+        Delete unused worker types (no workers with this type)
+        """
+        self.stdout.write("Removing unused worker types…")
+        unused_types = (
+            WorkerType.objects.annotate(workers_count=Count("workers"))
+            .filter(workers_count=0)
+        )
+        count = unused_types.count()
+        unused_types.delete()
+        self.stdout.write(self.style.SUCCESS(f"Successfully removed {count} unused worker types."))

arkindex/documents/tests/commands/test_cleanup.py

@@ -105,6 +105,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -149,6 +151,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -195,6 +199,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -248,6 +254,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -303,6 +311,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -367,6 +377,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -443,6 +455,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -551,6 +565,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -625,6 +641,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -703,6 +721,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -730,7 +750,7 @@ class TestCleanupCommand(FixtureTestCase):
             image_no_element_old_2 = Image.objects.create(path="path/pathpathpath/img", width=12, height=12, server=img_server)
         image_no_element_new = Image.objects.create(path="path/pathpath/img", width=12, height=12, server=img_server)
 
-        with self.assertNumQueries(19):
+        with self.assertNumQueries(21):
             self.assertEqual(
                 self.cleanup(),
                 dedent(
@@ -764,6 +784,8 @@ class TestCleanupCommand(FixtureTestCase):
                     Successfully cleaned up orphaned model versions archives.
                     Cleaning up deleted jobs from RQ user registries…
                     Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                    Removing unused worker types…
+                    Successfully removed 0 unused worker types.
                     """
                 ).strip()
             )
@@ -844,6 +866,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -926,6 +950,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -1007,6 +1033,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -1083,6 +1111,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 3 deleted jobs from 2 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -1190,6 +1220,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 1 unused worker types.
                 """
             ).strip()
         )
@@ -1261,6 +1293,8 @@ class TestCleanupCommand(FixtureTestCase):
                 Successfully cleaned up orphaned model versions archives.
                 Cleaning up deleted jobs from RQ user registries…
                 Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 0 unused worker types.
                 """
             ).strip()
         )
@@ -1270,3 +1304,53 @@ class TestCleanupCommand(FixtureTestCase):
         used_model.refresh_from_db()
         non_archived.refresh_from_db()
         recently_archived.refresh_from_db()
+
+
+    def test_cleanup_worker_types(self, s3_mock, rq_mock):
+        for worker_type in WorkerType.objects.all():
+            self.assertTrue(Worker.objects.filter(type=worker_type).exists())
+        unused_type = WorkerType.objects.create(slug="unused", display_name="A worker type that no worker uses")
+        self.assertFalse(Worker.objects.filter(type=unused_type).exists())
+        self.assertEqual(WorkerType.objects.count(), 7)
+
+        self.assertEqual(
+            self.cleanup(),
+            dedent(
+                """
+                Removing archived workers…
+                Successfully cleaned up 0 archived workers.
+                Removing orphaned Ponos artifacts…
+                Successfully cleaned up orphaned Ponos artifacts.
+                Removing 0 artifacts of expired processes from S3…
+                Removing logs for 0 tasks of expired processes from S3…
+                Removing 0 artifacts of expired processes…
+                Removing 0 tasks of expired processes…
+                Successfully cleaned up expired processes.
+                Removing 0 old corpus exports from S3…
+                Removing 0 old corpus exports…
+                Successfully cleaned up old corpus exports.
+                Removing orphaned corpus exports…
+                Successfully cleaned up orphaned corpus exports.
+                Deleting 0 DataFiles marked as trashed from S3 and the database…
+                Successfully cleaned up DataFiles marked as trashed.
+                Removing orphan images…
+                Successfully cleaned up orphan images.
+                Removing orphaned local images…
+                Successfully cleaned up orphaned local images.
+                Removing orphaned Ponos logs…
+                Successfully cleaned up orphaned Ponos logs.
+                Removing archived models…
+                Successfully cleaned up 0 archived models.
+                Removing orphaned model versions archives…
+                Successfully cleaned up orphaned model versions archives.
+                Cleaning up deleted jobs from RQ user registries…
+                Successfully cleaned up 0 deleted jobs from 0 RQ user registries.
+                Removing unused worker types…
+                Successfully removed 1 unused worker types.
+                """
+            ).strip()
+        )
+
+        with self.assertRaises(WorkerType.DoesNotExist):
+            unused_type.refresh_from_db()
+        self.assertEqual(WorkerType.objects.count(), 6)

arkindex/process/migrations/0044_alter_worker_type.py

+# Generated by Django 5.0.8 on 2024-10-10 12:10
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("process", "0043_workerversion_branch"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="worker",
+            name="type",
+            field=models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, related_name="workers", to="process.workertype"),
+        ),
+    ]

arkindex/process/models.py

@@ -525,7 +525,7 @@ class Worker(models.Model):
     id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
     name = models.CharField(max_length=100)
     slug = models.CharField(max_length=100)
-    type = models.ForeignKey("process.WorkerType", on_delete=models.PROTECT, related_name="type")
+    type = models.ForeignKey("process.WorkerType", on_delete=models.PROTECT, related_name="workers")
     description = models.TextField(default="", blank=True)
     memberships = GenericRelation("users.Right", "content_id")