diff --git a/arkindex/process/tests/test_workerruns.py b/arkindex/process/tests/test_workerruns.py
deleted file mode 100644
index 40ad68fcdeb66fe6477998646c1e3d51ad717598..0000000000000000000000000000000000000000
--- a/arkindex/process/tests/test_workerruns.py
+++ /dev/null
@@ -1,2969 +0,0 @@
-import uuid
-from datetime import datetime, timezone
-from unittest.mock import call, patch
-
-from django.db import transaction
-from django.urls import reverse
-from rest_framework import status
-from rest_framework.exceptions import ValidationError
-
-from arkindex.ponos.models import Agent, Farm, State
-from arkindex.process.models import (
- FeatureUsage,
- ProcessMode,
- Worker,
- WorkerRun,
- WorkerVersion,
- WorkerVersionState,
-)
-from arkindex.project.tests import FixtureAPITestCase
-from arkindex.training.models import Model, ModelVersion, ModelVersionState
-from arkindex.users.models import Role
-
-ENV = {
- "ARKINDEX_PROCESS_ID": "12345"
-}
-
-
-class TestWorkerRuns(FixtureAPITestCase):
- """
- Test worker runs endpoints and methods
- """
-
- @classmethod
- def setUpTestData(cls):
- super().setUpTestData()
- cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
- cls.farm = Farm.objects.first()
- cls.process_1 = cls.corpus.processes.create(
- creator=cls.user,
- mode=ProcessMode.Workers,
- farm=cls.farm,
- )
- cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
- cls.worker_1 = cls.version_1.worker
- cls.worker_custom = Worker.objects.get(slug="custom")
- cls.version_2 = WorkerVersion.objects.get(worker__slug="dla")
- cls.version_custom = cls.worker_custom.versions.get()
- cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
- cls.run_custom = cls.local_process.worker_runs.get(version=cls.version_custom)
- cls.configuration_1 = cls.worker_1.configurations.create(name="My config", configuration={"key": "value"})
- worker_version = WorkerVersion.objects.exclude(worker=cls.version_1.worker).first()
- cls.configuration_2 = worker_version.worker.configurations.create(name="Config")
- cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
- # Add an execution access right on the worker
- cls.worker_1.memberships.create(user=cls.user, level=Role.Contributor.value)
-
- # Model and Model version setup
- cls.model_1 = Model.objects.create(name="My model")
- cls.model_1.memberships.create(user=cls.user, level=Role.Contributor.value)
- cls.model_version_1 = ModelVersion.objects.create(
- model=cls.model_1,
- state=ModelVersionState.Available,
- size=8,
- hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
- archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
- )
-
- cls.model_2 = Model.objects.create(name="Their model")
- cls.model_2.memberships.create(user=cls.user, level=Role.Guest.value)
- cls.model_version_2 = cls.model_2.versions.create(
- state=ModelVersionState.Available,
- tag="blah",
- size=8,
- hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
- archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
- )
-
- cls.model_3 = Model.objects.create(name="Our model", public=True)
- cls.model_version_3 = cls.model_3.versions.create(
- state=ModelVersionState.Available,
- tag="blah",
- size=8,
- hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
- archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
- )
-
- cls.agent = Agent.objects.create(
- farm=cls.farm,
- hostname="claude",
- cpu_cores=42,
- cpu_frequency=1e15,
- ram_total=99e9,
- last_ping=datetime.now(timezone.utc),
- )
- # Add custom attributes to make the agent usable as an authenticated user
- cls.agent.is_agent = True
- cls.agent.is_anonymous = False
-
- def test_list_requires_login(self):
- with self.assertNumQueries(0):
- response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- def test_list_no_execution_right(self):
- """
- Worker runs attached to a process can be listed even if the user has no execution rights to workers
- This is due to the fact a user can see a process running on a corpus they have access
- """
- self.worker_1.memberships.all().delete()
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json()["count"], 1)
-
- def test_list(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- data = response.json()
- self.assertEqual(data["results"], [{
- "id": str(self.run_1.id),
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- }])
-
- def test_list_filter_process(self):
- run_2 = self.process_2.worker_runs.create(
- version=self.version_1,
- parents=[],
- )
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}))
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- data = response.json()
- self.assertEqual(data["count"], 1)
- self.assertEqual(data["results"][0]["id"], str(run_2.id))
-
- def test_create_requires_login(self):
- with self.assertNumQueries(0):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- self.assertDictEqual(response.json(), {"detail": "Authentication credentials were not provided."})
-
- @patch("arkindex.users.utils.get_max_level", return_value=Role.Guest.value)
- def test_create_no_execution_right(self, max_level_mock):
- """
- An execution access on the target worker version is required to create a worker run
- """
- self.worker_1.memberships.update(level=Role.Guest.value)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {"worker_version_id": ["You do not have an execution access to this worker."]})
-
- self.assertListEqual(max_level_mock.call_args_list, [
- call(self.user, self.worker_1)
- ])
-
- def test_create_invalid_version(self):
- self.client.force_login(self.user)
- version_id = uuid.uuid4()
-
- with self.assertNumQueries(4):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": version_id},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "worker_version_id": [f'Invalid pk "{version_id}" - object does not exist.'],
- })
-
- def test_create_unique(self):
- self.client.force_login(self.user)
- self.version_1.model_usage = FeatureUsage.Required
- self.version_1.save()
- cases = [
- (None, None),
- (None, self.configuration_1),
- (self.model_version_1, None),
- (self.model_version_1, self.configuration_1),
- ]
- for model_version, configuration in cases:
- with self.subTest(model_version=model_version, configuration=configuration):
- self.run_1.model_version = model_version
- self.run_1.configuration = configuration
- self.run_1.save()
-
- # Having a model version or a configuration adds one query for each
- query_count = 5 + bool(model_version) + bool(configuration)
-
- with self.assertNumQueries(query_count):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}),
- data={
- "worker_version_id": str(self.version_1.id),
- "model_version_id": str(model_version.id) if model_version else None,
- "configuration_id": str(configuration.id) if configuration else None,
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
- })
-
- def test_create_unavailable_version(self):
- self.version_1.state = WorkerVersionState.Error
- self.version_1.save()
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {"worker_version_id": ["This WorkerVersion is not in an Available state."]})
-
- def test_create_archived_worker(self):
- self.worker_1.archived = datetime.now(timezone.utc)
- self.worker_1.save()
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {"worker_version_id": ["This WorkerVersion is part of an archived worker."]})
-
- def test_create_archived_model(self):
- self.model_1.archived = datetime.now(timezone.utc)
- self.model_1.save()
- self.version_1.model_usage = FeatureUsage.Supported
- self.version_1.save()
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "model_version_id": str(self.model_version_1.id)},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {"model_version_id": ["This ModelVersion is part of an archived model."]})
-
- def test_create_invalid_process_id(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": uuid.uuid4()}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- self.assertEqual(response.json(), {"detail": "No Process matches the given query."})
-
- def test_create_invalid_process_mode(self):
- self.client.force_login(self.user)
-
- for mode in set(ProcessMode) - {ProcessMode.Workers, ProcessMode.Dataset, ProcessMode.Local}:
- with self.subTest(mode=mode):
- self.process_2.mode = mode
- self.process_2.save()
-
- with self.assertNumQueries(5):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- {"worker_version_id": str(self.version_1.id)},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
- })
-
- def test_create_non_corpus_process_mode(self):
- process = self.user.processes.get(mode=ProcessMode.Local)
- self.client.force_login(self.user)
- with self.assertNumQueries(3):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(process.id)}),
- {"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_create_process_already_started(self):
- process = self.corpus.processes.create(
- creator=self.user,
- mode=ProcessMode.Workers,
- farm=self.farm,
- )
- process.run()
-
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(process.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
- })
-
- def test_create(self):
- self.client.force_login(self.user)
-
- for mode in (ProcessMode.Workers, ProcessMode.Dataset):
- self.process_2.worker_runs.filter(version=self.version_1).delete()
-
- with self.subTest(mode=mode):
- self.process_2.mode = mode
- self.process_2.save()
-
- with self.assertNumQueries(6):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- {"worker_version_id": str(self.version_1.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-
- data = response.json()
- pk = data.pop("id")
- self.assertNotEqual(pk, self.run_1.id)
- self.assertDictEqual(data, {
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.process_2.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": mode.value,
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
- run = WorkerRun.objects.get(pk=pk)
- # Check generated summary
- self.assertEqual(run.summary, "Worker Recognizer @ version 1")
-
- def test_create_empty(self):
- """
- The worker_version_id is required to create a worker run
- """
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertDictEqual(response.json(), {
- "worker_version_id": ["This field is required."],
- })
-
- def test_create_configuration(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(7):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={
- "worker_version_id": str(self.version_1.id),
- "parents": [],
- "configuration_id": str(self.configuration_1.id)
- },
- )
- self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-
- data = response.json()
- pk = data.pop("id")
- self.assertNotEqual(pk, self.run_1.id)
- self.assertDictEqual(data, {
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": {
- "id": str(self.configuration_1.id),
- "archived": False,
- "configuration": {"key": "value"},
- "name": "My config"
- },
- "process": {
- "id": str(self.process_2.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
- })
- run = WorkerRun.objects.get(pk=pk)
- # Check generated summary
- self.assertEqual(run.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
-
- def test_create_invalid_configuration(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data={"worker_version_id": str(self.version_1.id), "parents": [], "configuration_id": str(self.configuration_2.id)},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
-
- def test_summary(self):
- self.client.force_login(self.user)
- test_version = self.worker_1.versions.create(
- version=2,
- state=WorkerVersionState.Available,
- docker_image_iid="registry.gitlab.com/dead-sea-scrolls:004",
- model_usage=FeatureUsage.Supported
- )
-
- cases = [
- ("eva-01", "revision_url", self.model_version_1, self.configuration_1, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]}) with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
- (None, "revision_url", self.model_version_1, self.configuration_1, f"Worker Recognizer @ {str(test_version.id)[:6]} with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
- ("eva-01", "revision_url", None, self.configuration_1, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]}) using configuration 'My config'"),
- (None, "revision_url", self.model_version_1, None, f"Worker Recognizer @ {str(test_version.id)[:6]} with model My model @ {str(self.model_version_1.id)[:6]}"),
- ("eva-01", "revision_url", None, None, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]})"),
- (None, "revision_url", None, None, f"Worker Recognizer @ {str(test_version.id)[:6]}"),
- ("eva-01", None, self.model_version_1, self.configuration_1, f"Worker Recognizer @ eva-01 (version 2) with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
- (None, None, self.model_version_1, self.configuration_1, f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
- ("eva-01", None, None, self.configuration_1, "Worker Recognizer @ eva-01 (version 2) using configuration 'My config'"),
- (None, None, self.model_version_1, None, f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]}"),
- ("eva-01", None, None, None, "Worker Recognizer @ eva-01 (version 2)"),
- (None, None, None, None, "Worker Recognizer @ version 2"),
- ]
-
- for tag, revision_url, model_version, config, expected_summary in cases:
- with self.subTest(tag=tag, model_version=model_version, config=config), transaction.atomic():
- # Clear the process of worker runs
- self.process_2.worker_runs.all().delete()
-
- num_queries = 6
-
- test_version.version = None
- test_version.tag = tag
- test_version.revision_url = revision_url
- if not revision_url:
- test_version.version = 2
- test_version.save()
-
- payload = {
- "worker_version_id": str(test_version.id),
- "parents": [],
- }
- if model_version:
- payload["model_version_id"] = model_version.id
- # If there is a model version, it adds a query
- num_queries += 1
- if config:
- payload["configuration_id"] = config.id
- # If there is a configuration, it adds a query
- num_queries += 1
-
- with self.assertNumQueries(num_queries):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
- data=payload,
- )
- if response.status_code == 400:
- self.assertDictEqual(response.json(), {})
- self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-
- created_run = WorkerRun.objects.get(pk=response.json()["id"])
- self.assertEqual(created_run.summary, expected_summary)
-
- def test_retrieve_requires_login(self):
- with self.assertNumQueries(0):
- response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}))
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- def test_retrieve_no_worker_execution_right(self):
- """
- A user can retrieve any worker run if they have an admin access to the process project
- """
- self.worker_1.memberships.update(level=Role.Guest.value)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(4):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertDictEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_retrieve_invalid_id(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"})
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_retrieve(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(4):
- response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}))
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertDictEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_retrieve_custom(self):
- self.client.force_login(self.user)
- # Update process attributes to ensure they do not cause extra requests
- page = self.corpus.elements.get(name="Volume 1, page 1r")
- self.local_process.element = page
- self.local_process.element_type = page.type
- self.local_process.folder_type = page.type
- self.local_process.prefix = "a"
- self.local_process.save()
- with self.assertNumQueries(5):
- response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_custom.id)}))
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertDictEqual(response.json(), {
- "id": str(self.run_custom.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.local_process.id),
- "activity_state": "disabled",
- "corpus": None,
- "chunks": 1,
- "mode": "local",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": {
- "id": str(page.id),
- "type": "page",
- "corpus": {
- "id": str(self.corpus.id),
- "name": "Unit Tests",
- "public": True
- },
- "name": "Volume 1, page 1r",
- "rotation_angle": 0,
- "mirrored": False,
- "thumbnail_url": None,
- "zone": {
- "id": str(page.id),
- "image": {
- "id": str(page.image.id),
- "height": 1000,
- "path": "img1",
- "s3_url": None,
- "server": {
- "display_name": "Test Server",
- "max_height": None,
- "max_width": None,
- "url": "http://server"
- },
- "status": "unchecked",
- "url": "http://server/img1",
- "width": 1000
- },
- "polygon": [[int(x), int(y)] for x, y in page.polygon],
- "url": "http://server/img1/0,0,1000,1000/full/0/default.jpg",
- },
- },
- "element_type": "page",
- "folder_type": "page",
- "prefix": "a",
- },
- "worker_version": {
- "id": str(self.version_custom.id),
- "configuration": {"custom": "value"},
- "docker_image_iid": None,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "created",
- "worker": {
- "id": str(self.worker_custom.id),
- "name": "Custom worker",
- "slug": "custom",
- "type": "custom",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "summary": "Worker Custom worker @ version 1",
- "use_gpu": False,
- })
-
- def test_retrieve_local(self):
- """
- A user can retrieve a run on their own local process
- """
- run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertDictEqual(response.json(), {
- "id": str(run.id),
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.local_process.id),
- "activity_state": "disabled",
- "corpus": None,
- "chunks": 1,
- "mode": "local",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_retrieve_local_only_current_user(self):
- """
- A user cannot retrieve a run on another user's local process
- """
- run = WorkerRun.objects.filter(process__creator=self.superuser, process__mode=ProcessMode.Local).first()
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_retrieve_agent(self):
- """
- A Ponos agent can retrieve a WorkerRun on a process where it has some assigned tasks
- """
- self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
-
- # Agent auth is not implemented in CE
- self.client.force_authenticate(user=self.agent)
- with self.assertNumQueries(3):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "mode": "workers",
- "chunks": 1,
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_retrieve_agent_unassigned(self):
- """
- A Ponos agent cannot retrieve a WorkerRun on a process where it does not have any assigned tasks
- """
- self.process_1.tasks.create(run=0, depth=0, slug="something", agent=None)
-
- # Agent auth is not implemented in CE
- self.client.force_authenticate(user=self.agent)
- with self.assertNumQueries(1):
- response = self.client.get(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_update_requires_nothing(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={},
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- def test_update_requires_login(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
-
- with self.assertNumQueries(0):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- @patch("arkindex.project.mixins.get_max_level", return_value=Role.Contributor.value)
- def test_update_no_project_admin_right(self, get_max_level_mock):
- """
- A user cannot update a worker run if they have no admin access on its process project
- """
- self.corpus.memberships.filter(user=self.user).update(level=Role.Contributor.value)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": []
- }
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- self.assertEqual(response.json(), {"detail": "You do not have an admin access to this process."})
-
- self.assertEqual(get_max_level_mock.call_count, 1)
- self.assertEqual(get_max_level_mock.call_args, call(self.user, self.corpus))
-
- def test_update_invalid_process_mode(self):
- """
- A user cannot update a worker run on a local process
- """
- run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "parents": []
- }
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
- })
-
- def test_update_invalid_id(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_update_nonexistent_parent(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": ["12341234-1234-1234-1234-123412341234"],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), [
- f"Can't add or update WorkerRun {self.run_1.id} because parents field isn't properly defined. It can be either because"
- " one or several UUIDs don't refer to existing WorkerRuns or either because listed WorkerRuns doesn't belong to the"
- " same Process than this WorkerRun."
- ])
-
- def test_update_duplicate_parents(self):
- self.client.force_login(self.user)
- run_2 = self.process_1.worker_runs.create(version=self.version_2)
-
- with self.assertNumQueries(4):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [
- str(run_2.id),
- str(run_2.id),
- ],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "parents": ["The parents of a WorkerRun must be unique."],
- })
-
- def test_duplicate_parents_signal(self):
- """
- Duplicates in WorkerRun.parents are also detected outside of the WorkerRun APIs
- """
- run_2 = self.process_1.worker_runs.create(
- version=self.version_2,
- parents=[],
- )
-
- run_2.parents = [self.run_1.id, self.run_1.id]
- with self.assertRaisesRegex(ValidationError, f"Can't add or update WorkerRun {run_2.id} because it has duplicate parents."):
- run_2.parents = [self.run_1.id, self.run_1.id]
- run_2.save()
-
- def test_update_process_id(self):
- """
- Process field cannot be updated
- """
- self.client.force_login(self.user)
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "process_id": str(self.process_2.id),
- "parents": [],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
- self.run_1.refresh_from_db()
- self.assertEqual(self.run_1.process.id, self.process_1.id)
-
- def test_update_worker_version_id(self):
- """
- Version field cannot be updated
- """
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "worker_version_id": str(self.version_2.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- },
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
- self.run_1.refresh_from_db()
- self.assertNotEqual(self.run_1.version_id, self.version_2.id)
-
- def test_update_configuration(self):
- self.client.force_login(self.user)
- self.assertEqual(self.run_1.configuration, None)
- # Check generated summary, before updating, it should not be that verbose
- self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1")
-
- with self.assertNumQueries(6):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
- data={
- "parents": [],
- "configuration_id": str(self.configuration_1.id)
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.run_1.refresh_from_db()
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": {
- "id": str(self.configuration_1.id),
- "archived": False,
- "configuration": {"key": "value"},
- "name": "My config"
- },
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
- })
- self.assertEqual(self.run_1.configuration.id, self.configuration_1.id)
- # Check generated summary, after the update, the configuration should be displayed as well
- self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
-
- def test_update_invalid_configuration(self):
- self.client.force_login(self.user)
- self.assertEqual(self.run_1.configuration, None)
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
- data={"parents": [], "configuration_id": str(self.configuration_2.id)},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
-
- def test_update_process_already_started(self):
- """
- Update dependencies of a worker run is not possible once the process is started
- """
- self.process_1.run()
- self.assertTrue(self.process_1.tasks.exists())
- self.client.force_login(self.user)
-
- with self.assertNumQueries(4):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
- })
-
- def test_update_model_version_not_allowed(self):
- """
- The model_version UUID is not allowed when the related version doesn't allow model_usage
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Disabled
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This worker version does not support models."]
- })
-
- def test_update_unknown_model_version(self):
- """
- Cannot use a model version id that doesn't exist
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
- random_model_version_uuid = str(uuid.uuid4())
-
- with self.assertNumQueries(4):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": random_model_version_uuid,
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{random_model_version_uuid}" - object does not exist.']
- })
-
- @patch("arkindex.users.managers.BaseACLManager.filter_rights", return_value=ModelVersion.objects.none())
- def test_update_model_version_no_access(self, filter_rights_mock):
- """
- Cannot update a worker run with a model_version UUID, when you don't have access to the model version
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- # Create a model version, the user has no access to
- model_no_access = Model.objects.create(name="Secret model")
- model_version_no_access = ModelVersion.objects.create(
- model=model_no_access,
- state=ModelVersionState.Available,
- size=8,
- hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
- archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
- )
-
- with self.assertNumQueries(3):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(model_version_no_access.id),
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{model_version_no_access.id}" - object does not exist.'],
- })
-
- self.assertListEqual(filter_rights_mock.call_args_list, [
- call(self.user, Model, Role.Guest.value),
- call(self.user, Model, Role.Contributor.value),
- ])
-
- @patch("arkindex.users.managers.BaseACLManager.filter_rights")
- def test_update_model_version_guest(self, filter_rights_mock):
- """
- Cannot update a worker run with a model_version when you only have guest access to the model,
- and the model version has no tag or is not available
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- def filter_rights(user, model, level):
- """
- The filter_rights mock needs to return nothing when called for contributor access,
- and the models we will test on when called for guest access
- """
- if level == Role.Guest.value:
- return Model.objects.filter(id__in=(self.model_2.id, self.model_3.id))
- return Model.objects.none()
-
- filter_rights_mock.side_effect = filter_rights
-
- cases = [
- # On a model with a membership giving guest access
- (self.model_version_2, None, ModelVersionState.Created),
- (self.model_version_2, None, ModelVersionState.Error),
- (self.model_version_2, None, ModelVersionState.Available),
- (self.model_version_2, "blah", ModelVersionState.Created),
- (self.model_version_2, "blah", ModelVersionState.Error),
- # On a public model with no membership
- (self.model_version_3, None, ModelVersionState.Created),
- (self.model_version_3, None, ModelVersionState.Error),
- (self.model_version_3, None, ModelVersionState.Available),
- (self.model_version_3, "blah", ModelVersionState.Created),
- (self.model_version_3, "blah", ModelVersionState.Error),
- ]
-
- for model_version, tag, state in cases:
- filter_rights_mock.reset_mock()
- with self.subTest(model_version=model_version, tag=tag, state=state):
- model_version.tag = tag
- model_version.state = state
- model_version.save()
-
- with self.assertNumQueries(4):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(model_version.id),
- "parents": [],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{model_version.id}" - object does not exist.'],
- })
-
- self.assertListEqual(filter_rights_mock.call_args_list, [
- call(self.user, Model, Role.Guest.value),
- call(self.user, Model, Role.Contributor.value),
- ])
-
- def test_update_model_version_unavailable(self):
- self.client.force_login(self.user)
- version = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(
- version=version,
- parents=[],
- )
- self.model_version_1.state = ModelVersionState.Error
- self.model_version_1.save()
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This ModelVersion is not in an Available state."]
- })
-
- def test_update_model_archived(self):
- self.client.force_login(self.user)
- version = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(version=version)
- self.model_1.archived = datetime.now(timezone.utc)
- self.model_1.save()
-
- with self.assertNumQueries(5):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This ModelVersion is part of an archived model."],
- })
-
- def test_update_model_version_id(self):
- """
- Update the worker run by adding a model_version with a worker version that supports it
- """
- self.client.force_login(self.user)
- version_with_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Supported
- )
- run = self.process_1.worker_runs.create(
- version=version_with_model,
- parents=[],
- )
- self.assertEqual(run.model_version, None)
- # Check generated summary, before updating, there should be only information about the worker version
- self.assertEqual(run.summary, "Worker Recognizer @ version 2")
-
- model_versions = [
- # Version on a model with contributor access
- self.model_version_1,
- # Available version with tag on a model with guest access
- self.model_version_2,
- # Available version with tag on a public model
- self.model_version_3,
- ]
- for model_version in model_versions:
- with self.subTest(model_version=model_version):
- with self.assertNumQueries(6):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(model_version.id),
- "parents": [],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- run.refresh_from_db()
- self.assertEqual(response.json(), {
- "id": str(run.id),
- "configuration": None,
- "model_version": {
- "id": str(model_version.id),
- "configuration": {},
- "model": {
- "id": str(model_version.model.id),
- "name": model_version.model.name
- },
- "size": 8,
- "state": "available",
- "tag": model_version.tag,
- },
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(version_with_model.id),
- "configuration": {"test": "test2"},
- "docker_image_iid": None,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Supported.value,
- "revision_url": None,
- "version": version_with_model.version,
- "tag": None,
- "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
- "state": "created",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}",
- })
- self.assertEqual(run.model_version_id, model_version.id)
- self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}")
-
- def test_update_configuration_and_model_version(self):
- """
- Update the worker run by adding both a model_version and a worker configuration
- """
- self.client.force_login(self.user)
- version_with_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(
- version=version_with_model,
- parents=[],
- )
- self.assertIsNone(run.model_version)
- self.assertIsNone(run.configuration)
- # Check generated summary, before updating, there should be only information about the worker version
- self.assertEqual(run.summary, "Worker Recognizer @ version 2")
-
- with self.assertNumQueries(7):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- "configuration_id": str(self.configuration_1.id),
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- run.refresh_from_db()
- self.assertEqual(response.json(), {
- "id": str(run.id),
- "configuration": {
- "id": str(self.configuration_1.id),
- "archived": False,
- "configuration": {"key": "value"},
- "name": "My config"
- },
- "model_version": {
- "id": str(self.model_version_1.id),
- "configuration": {},
- "model": {
- "id": str(self.model_1.id),
- "name": "My model"
- },
- "size": 8,
- "state": "available",
- "tag": None
- },
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(version_with_model.id),
- "configuration": {"test": "test2"},
- "docker_image_iid": None,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Required.value,
- "revision_url": None,
- "version": version_with_model.version,
- "tag": None,
- "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
- "state": "created",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'",
- })
- self.assertEqual(run.model_version_id, self.model_version_1.id)
- # Check generated summary, after updating, there should be information about the model loaded
- self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {self.model_version_1.model.name} @ {str(self.model_version_1.id)[:6]} using configuration '{self.configuration_1.name}'")
-
- def test_update(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
- self.client.force_login(self.user)
-
- with self.assertNumQueries(7):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.run_1.refresh_from_db()
- self.assertDictEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [str(run_2.id)],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_update_unique(self):
- self.client.force_login(self.user)
- self.version_1.model_usage = FeatureUsage.Required
- self.version_1.save()
- cases = [
- (None, None),
- (None, self.configuration_1),
- (self.model_version_1, None),
- (self.model_version_1, self.configuration_1),
- ]
- for model_version, configuration in cases:
- with self.subTest(model_version=model_version, configuration=configuration):
- # Erase any previous failures
- self.process_1.worker_runs.exclude(id=self.run_1.id).delete()
-
- self.run_1.model_version = model_version
- self.run_1.configuration = configuration
- self.run_1.save()
-
- # Ensure the other run has different values before updating to avoid conflicts
- run_2 = self.process_1.worker_runs.create(
- version=self.version_1,
- model_version=None if model_version else self.model_version_1,
- configuration=None if configuration else self.configuration_1,
- )
-
- # Having a model version or a configuration adds one query for each
- query_count = 4 + bool(model_version) + bool(configuration)
-
- with self.assertNumQueries(query_count):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- # Update the second worker run to the first worker run's values to cause a conflict
- "model_version_id": str(model_version.id) if model_version else None,
- "configuration_id": str(configuration.id) if configuration else None,
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
- })
-
- def test_update_agent(self):
- """
- Ponos agents cannot update WorkerRuns, even when they can access them
- """
- self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
-
- # Agent auth is not implemented in CE
- self.client.force_authenticate(user=self.agent)
- with self.assertNumQueries(1):
- response = self.client.put(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
- self.assertEqual(response.json(), {
- "detail": "You do not have an admin access to this process.",
- })
-
- def test_partial_update_requires_login(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
-
- with self.assertNumQueries(0):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- @patch("arkindex.project.mixins.get_max_level", return_value=Role.Contributor.value)
- def test_partial_update_no_project_admin_right(self, get_max_level_mock):
- """
- A user cannot update a worker run if they have no admin access on its process project
- """
- self.corpus.memberships.filter(user=self.user).update(level=Role.Contributor.value)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(3):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- self.assertEqual(response.json(), {"detail": "You do not have an admin access to this process."})
-
- self.assertEqual(get_max_level_mock.call_count, 1)
- self.assertEqual(get_max_level_mock.call_args, call(self.user, self.corpus))
-
- def test_partial_update_invalid_id(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
-
- self.client.force_login(self.user)
- with self.assertNumQueries(3):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_partial_update_local(self):
- """
- A user cannot update a worker run on a local process
- """
- run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "parents": []
- }
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
- })
-
- def test_partial_update_nonexistent_parent(self):
- self.client.force_login(self.user)
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": ["12341234-1234-1234-1234-123412341234"],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), [
- f"Can't add or update WorkerRun {self.run_1.id} because parents field isn't properly defined. It can be either because"
- " one or several UUIDs don't refer to existing WorkerRuns or either because listed WorkerRuns doesn't belong to the"
- " same Process than this WorkerRun."
- ])
-
- def test_partial_update_process_id(self):
- """
- Process field cannot be updated
- """
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "process_id": str(self.process_2.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
- self.run_1.refresh_from_db()
- self.assertEqual(self.run_1.process.id, self.process_1.id)
-
- def test_partial_update_worker_version_id(self):
- """
- Version field cannot be updated
- """
- self.client.force_login(self.user)
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "worker_version_id": str(self.version_2.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
- self.run_1.refresh_from_db()
- self.assertNotEqual(self.run_1.version_id, self.version_2.id)
-
- def test_partial_update_configuration(self):
- self.client.force_login(self.user)
- self.assertEqual(self.run_1.configuration, None)
- self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1")
-
- with self.assertNumQueries(6):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
- data={
- "configuration_id": str(self.configuration_1.id)
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.run_1.refresh_from_db()
- self.assertEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": {
- "id": str(self.configuration_1.id),
- "archived": False,
- "configuration": {"key": "value"},
- "name": "My config"
- },
- "model_version": None,
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
- })
- self.assertEqual(self.run_1.configuration.id, self.configuration_1.id)
- self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
-
- def test_partial_update_invalid_configuration(self):
- self.client.force_login(self.user)
- self.assertEqual(self.run_1.configuration, None)
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
- data={"configuration_id": str(self.configuration_2.id)},
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
-
- def test_partial_update_process_already_started(self):
- """
- Update dependencies of a worker run is not possible once the process is started
- """
- self.process_1.run()
- self.assertTrue(self.process_1.tasks.exists())
- self.client.force_login(self.user)
-
- with self.assertNumQueries(4):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
- })
-
- def test_partial_update_model_version_not_allowed(self):
- """
- The model_version UUID is not allowed when the related version doesn't allow model_usage
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Disabled
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This worker version does not support models."]
- })
-
- def test_partial_update_unknown_model_version(self):
- """
- Cannot use a model version id that doesn't exist
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
- random_model_version_uuid = str(uuid.uuid4())
-
- with self.assertNumQueries(4):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": random_model_version_uuid,
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{random_model_version_uuid}" - object does not exist.']
- })
-
- @patch("arkindex.users.managers.BaseACLManager.filter_rights", return_value=ModelVersion.objects.none())
- def test_partial_update_model_version_no_access(self, filter_rights_mock):
- """
- Cannot update a worker run with a model_version UUID, when you don't have access to the model version
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- # Create a model version, the user has no access to
- model_no_access = Model.objects.create(name="Secret model")
- model_version_no_access = ModelVersion.objects.create(model=model_no_access, state=ModelVersionState.Available, size=8, hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
-
- with self.assertNumQueries(3):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(model_version_no_access.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{model_version_no_access.id}" - object does not exist.'],
- })
-
- self.assertListEqual(filter_rights_mock.call_args_list, [
- call(self.user, Model, Role.Guest.value),
- call(self.user, Model, Role.Contributor.value),
- ])
-
- @patch("arkindex.users.managers.BaseACLManager.filter_rights")
- def test_partial_update_model_version_guest(self, filter_rights_mock):
- """
- Cannot update a worker run with a model_version when you only have guest access to the model,
- and the model version has no tag or is not available
- """
- self.client.force_login(self.user)
- version_no_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_no_model,
- parents=[],
- )
-
- def filter_rights(user, model, level):
- """
- The filter_rights mock needs to return nothing when called for contributor access,
- and the models we will test on when called for guest access
- """
- if level == Role.Guest.value:
- return Model.objects.filter(id__in=(self.model_2.id, self.model_3.id))
- return Model.objects.none()
-
- filter_rights_mock.side_effect = filter_rights
-
- cases = [
- # On a model with a membership giving guest access
- (self.model_version_2, None, ModelVersionState.Created),
- (self.model_version_2, None, ModelVersionState.Error),
- (self.model_version_2, None, ModelVersionState.Available),
- (self.model_version_2, "blah", ModelVersionState.Created),
- (self.model_version_2, "blah", ModelVersionState.Error),
- # On a public model with no membership
- (self.model_version_3, None, ModelVersionState.Created),
- (self.model_version_3, None, ModelVersionState.Error),
- (self.model_version_3, None, ModelVersionState.Available),
- (self.model_version_3, "blah", ModelVersionState.Created),
- (self.model_version_3, "blah", ModelVersionState.Error),
- ]
-
- for model_version, tag, state in cases:
- filter_rights_mock.reset_mock()
- with self.subTest(model_version=model_version, tag=tag, state=state):
- model_version.tag = tag
- model_version.state = state
- model_version.save()
-
- with self.assertNumQueries(4):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- "model_version_id": str(model_version.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": [f'Invalid pk "{model_version.id}" - object does not exist.'],
- })
-
- self.assertListEqual(filter_rights_mock.call_args_list, [
- call(self.user, Model, Role.Guest.value),
- call(self.user, Model, Role.Contributor.value),
- ])
-
- def test_partial_update_model_version_unavailable(self):
- self.client.force_login(self.user)
- version = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(
- version=version,
- parents=[],
- )
- self.model_version_1.state = ModelVersionState.Error
- self.model_version_1.save()
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- "parents": []
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This ModelVersion is not in an Available state."],
- })
-
- def test_partial_update_model_archived(self):
- self.client.force_login(self.user)
- version = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(version=version)
- self.model_1.archived = datetime.now(timezone.utc)
- self.model_1.save()
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "model_version_id": ["This ModelVersion is part of an archived model."],
- })
-
- def test_partial_update_model_version(self):
- """
- Update the worker run by adding a model_version with a worker version that supports it
- """
- self.client.force_login(self.user)
- version_with_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(
- version=version_with_model,
- parents=[],
- )
- self.assertIsNone(run.model_version_id)
- self.assertEqual(run.summary, "Worker Recognizer @ version 2")
-
- model_versions = [
- # Version on a model with contributor access
- self.model_version_1,
- # Available version with tag on a model with guest access
- self.model_version_2,
- # Available version with tag on a public model
- self.model_version_3,
- ]
- for model_version in model_versions:
- with self.subTest(model_version=model_version):
- with self.assertNumQueries(6):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(model_version.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- run.refresh_from_db()
- self.assertEqual(response.json(), {
- "id": str(run.id),
- "configuration": None,
- "model_version": {
- "id": str(model_version.id),
- "configuration": {},
- "model": {
- "id": str(model_version.model.id),
- "name": model_version.model.name
- },
- "size": 8,
- "state": "available",
- "tag": model_version.tag,
- },
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(version_with_model.id),
- "configuration": {"test": "test2"},
- "docker_image_iid": None,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Required.value,
- "revision_url": None,
- "version": version_with_model.version,
- "tag": None,
- "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
- "state": "created",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}",
- })
- self.assertEqual(run.model_version_id, model_version.id)
- self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}")
-
- def test_partial_update_model_version_with_configuration(self):
- """
- Updating the worker run by adding a model_version with a worker version
- doesn't erase previously loaded worker configuration
- """
- self.client.force_login(self.user)
- version_with_model = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"},
- model_usage=FeatureUsage.Required
- )
- run = self.process_1.worker_runs.create(
- version=version_with_model,
- parents=[],
- configuration=self.configuration_1
- )
- self.assertEqual(run.model_version_id, None)
-
- with self.assertNumQueries(6):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={
- "model_version_id": str(self.model_version_1.id),
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- run.refresh_from_db()
- self.assertEqual(response.json(), {
- "id": str(run.id),
- "configuration": {
- "archived": False,
- "configuration": {"key": "value"},
- "id": str(self.configuration_1.id),
- "name": "My config"
- },
- "model_version": {
- "id": str(self.model_version_1.id),
- "configuration": {},
- "model": {
- "id": str(self.model_1.id),
- "name": "My model"
- },
- "size": 8,
- "state": "available",
- "tag": None
- },
- "parents": [],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(version_with_model.id),
- "configuration": {"test": "test2"},
- "docker_image_iid": None,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Required.value,
- "revision_url": None,
- "version": version_with_model.version,
- "tag": None,
- "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
- "state": "created",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'",
- })
- self.assertEqual(run.model_version_id, self.model_version_1.id)
- self.assertEqual(run.configuration_id, self.configuration_1.id)
-
- def test_partial_update(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[],
- )
- self.client.force_login(self.user)
-
- with self.assertNumQueries(7):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- data={
- "parents": [str(run_2.id)],
- },
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- self.run_1.refresh_from_db()
- self.assertDictEqual(response.json(), {
- "id": str(self.run_1.id),
- "configuration": None,
- "model_version": None,
- "parents": [str(run_2.id)],
- "process": {
- "id": str(self.process_1.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- },
- "worker_version": {
- "id": str(self.version_1.id),
- "configuration": {"test": 42},
- "docker_image_iid": self.version_1.docker_image_iid,
- "gpu_usage": "disabled",
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 1,
- "tag": None,
- "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker_1.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "use_gpu": False,
- "summary": "Worker Recognizer @ version 1",
- })
-
- def test_partial_update_unique(self):
- self.client.force_login(self.user)
- self.version_1.model_usage = FeatureUsage.Required
- self.version_1.save()
- cases = [
- (None, None),
- (None, self.configuration_1),
- (self.model_version_1, None),
- (self.model_version_1, self.configuration_1),
- ]
- for model_version, configuration in cases:
- with self.subTest(model_version=model_version, configuration=configuration):
- # Erase any previous failures
- self.process_1.worker_runs.exclude(id=self.run_1.id).delete()
-
- self.run_1.model_version = model_version
- self.run_1.configuration = configuration
- self.run_1.save()
-
- # Ensure the other run has different values before updating to avoid conflicts
- run_2 = self.process_1.worker_runs.create(
- version=self.version_1,
- model_version=None if model_version else self.model_version_1,
- configuration=None if configuration else self.configuration_1,
- )
-
- # Having a model version or a configuration adds one query for each
- query_count = 4 + bool(model_version) + bool(configuration)
-
- with self.assertNumQueries(query_count):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
- data={
- # Update the second worker run to the first worker run's values to cause a conflict
- "model_version_id": str(model_version.id) if model_version else None,
- "configuration_id": str(configuration.id) if configuration else None,
- },
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), {
- "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
- })
-
- def test_partial_update_agent(self):
- """
- Ponos agents cannot update WorkerRuns, even when they can access them
- """
- self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
-
- # Agent auth is not implemented in CE
- self.client.force_authenticate(user=self.agent)
- with self.assertNumQueries(1):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
- self.assertEqual(response.json(), {
- "detail": "You do not have an admin access to this process.",
- })
-
- def test_delete_requires_login(self):
- with self.assertNumQueries(0):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
-
- def test_delete_invalid_id(self):
- self.client.force_login(self.user)
- with self.assertNumQueries(3):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"})
- )
- self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
-
- def test_delete_no_worker_execution_right(self):
- """
- A user can delete a worker run with no right on its worker but an admin access to the process project
- """
- self.worker_1.memberships.update(level=Role.Guest.value)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
-
- with self.assertRaises(WorkerRun.DoesNotExist):
- self.run_1.refresh_from_db()
-
- def test_delete_local(self):
- """
- A user cannot delete a worker run on a local process
- """
- run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
- self.client.force_login(self.user)
-
- with self.assertNumQueries(4):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), ["WorkerRuns can only be deleted from Workers or Dataset processes."])
-
- def test_delete(self):
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
-
- with self.assertRaises(WorkerRun.DoesNotExist):
- self.run_1.refresh_from_db()
-
- def test_delete_with_children(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- version_3 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=3,
- configuration={"test": "test3"}
- )
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[self.run_1.id],
- )
- run_3 = self.process_1.worker_runs.create(
- version=version_3,
- parents=[self.run_1.id, run_2.id],
- )
-
- self.assertTrue(self.run_1.id in run_2.parents)
- self.assertTrue(self.run_1.id in run_3.parents)
- self.client.force_login(self.user)
-
- with self.assertNumQueries(6):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
-
- with self.assertRaises(WorkerRun.DoesNotExist):
- self.run_1.refresh_from_db()
- run_2.refresh_from_db()
- run_3.refresh_from_db()
- self.assertEqual(run_2.parents, [])
- self.assertEqual(run_3.parents, [run_2.id])
-
- def test_delete_started_process(self):
- """
- A user shouldn't be able to delete the worker run of a started process
- """
- self.client.force_login(self.user)
- self.process_1.run()
- self.process_1.tasks.update(state=State.Running)
-
- with self.assertNumQueries(3):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
-
- self.assertEqual(response.json(), ["WorkerRuns cannot be deleted from a process that has already started."])
-
- def test_delete_agent(self):
- """
- Ponos agents cannot delete WorkerRuns, even when they can access them
- """
- self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
-
- # Agent auth is not implemented in CE
- self.client.force_authenticate(user=self.agent)
- with self.assertNumQueries(1):
- response = self.client.delete(
- reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
- )
- self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
- self.assertEqual(response.json(), {
- "detail": "You do not have an admin access to this process.",
- })
-
- def test_build_task_no_parent(self):
- task, parent_slugs = self.run_1.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json")
-
- self.assertEqual(task.slug, f"reco_{str(self.run_1.id)[0:6]}")
- self.assertEqual(task.image, self.version_1.docker_image_iid)
- self.assertEqual(task.command, None)
- self.assertEqual(task.shm_size, None)
- self.assertEqual(parent_slugs, ["import"])
- self.assertEqual(task.env, {
- "ARKINDEX_PROCESS_ID": "12345",
- "ARKINDEX_TASK_TOKEN": str(task.token),
- "TASK_ELEMENTS": "/data/import/elements.json",
- "ARKINDEX_WORKER_RUN_ID": str(self.run_1.id),
- })
-
- def test_build_task_with_chunk(self):
- task, parent_slugs = self.run_1.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json", chunk=4)
-
- self.assertEqual(task.slug, f"reco_{str(self.run_1.id)[0:6]}_4")
- self.assertEqual(task.image, self.version_1.docker_image_iid)
- self.assertEqual(task.command, None)
- self.assertEqual(task.shm_size, None)
- self.assertEqual(parent_slugs, ["import"])
- self.assertEqual(task.env, {
- "ARKINDEX_PROCESS_ID": "12345",
- "ARKINDEX_TASK_TOKEN": str(task.token),
- "ARKINDEX_TASK_CHUNK": "4",
- "TASK_ELEMENTS": "/data/import/elements.json",
- "ARKINDEX_WORKER_RUN_ID": str(self.run_1.id),
- })
-
- def test_build_task_with_parent(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- version_2.docker_image_iid = "evaunit:latest"
- version_2.state = WorkerVersionState.Available
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[self.run_1.id],
- )
-
- task, parent_slugs = run_2.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json")
-
- self.assertEqual(task.slug, f"reco_{str(run_2.id)[0:6]}")
- self.assertEqual(task.image, version_2.docker_image_iid)
- self.assertEqual(task.command, None)
- self.assertEqual(task.shm_size, None)
- self.assertEqual(parent_slugs, [f"reco_{str(self.run_1.id)[0:6]}"])
- self.assertEqual(task.env, {
- "ARKINDEX_PROCESS_ID": "12345",
- "ARKINDEX_TASK_TOKEN": str(task.token),
- "TASK_ELEMENTS": "/data/import/elements.json",
- "ARKINDEX_WORKER_RUN_ID": str(run_2.id),
- })
-
- def test_build_task_with_parent_and_chunk(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- version_2.docker_image_iid = "evaunit:latest"
- version_2.state = WorkerVersionState.Available
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[self.run_1.id],
- )
-
- task, parent_slugs = run_2.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json", chunk=4)
-
- self.assertEqual(task.slug, f"reco_{str(run_2.id)[0:6]}_4")
- self.assertEqual(task.image, version_2.docker_image_iid)
- self.assertEqual(task.command, None)
- self.assertEqual(task.shm_size, None)
- self.assertEqual(parent_slugs, [f"reco_{str(self.run_1.id)[0:6]}_4"])
- self.assertEqual(task.env, {
- "ARKINDEX_PROCESS_ID": "12345",
- "ARKINDEX_TASK_TOKEN": str(task.token),
- "ARKINDEX_TASK_CHUNK": "4",
- "TASK_ELEMENTS": "/data/import/elements.json",
- "ARKINDEX_WORKER_RUN_ID": str(run_2.id),
- })
-
- def test_build_task_shm_size(self):
- self.version_1.configuration = {
- "docker": {
- "shm_size": 505,
- }
- }
- task, parent_slugs = self.run_1.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json")
-
- self.assertEqual(task.slug, f"reco_{str(self.run_1.id)[0:6]}")
- self.assertEqual(task.image, self.version_1.docker_image_iid)
- self.assertEqual(task.command, None)
- self.assertEqual(task.shm_size, 505)
- self.assertEqual(parent_slugs, ["import"])
- self.assertEqual(task.env, {
- "ARKINDEX_PROCESS_ID": "12345",
- "ARKINDEX_TASK_TOKEN": str(task.token),
- "TASK_ELEMENTS": "/data/import/elements.json",
- "ARKINDEX_WORKER_RUN_ID": str(self.run_1.id),
- })
-
- def test_build_task_unavailable_version(self):
- version_2 = WorkerVersion.objects.create(
- worker=self.worker_1,
- version=2,
- configuration={"test": "test2"}
- )
- version_2.docker_image_iid = "evaunit:latest"
- self.assertEqual(version_2.state, WorkerVersionState.Created)
- run_2 = self.process_1.worker_runs.create(
- version=version_2,
- parents=[self.run_1.id],
- )
-
- with self.assertRaisesRegex(
- AssertionError,
- f"Worker Version {version_2.id} is not available and cannot be used to build a task."
- ):
- run_2.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json")
-
- def test_build_task_unavailable_model_version(self):
- self.model_version_1.state = ModelVersionState.Created
- self.model_version_1.save()
- self.run_1.model_version = self.model_version_1
- self.run_1.save()
- with self.assertRaisesRegex(
- AssertionError,
- f"ModelVersion {self.model_version_1.id} is not available and cannot be used to build a task."
- ):
- self.run_1.build_task(self.process_1, ENV.copy(), "import", "/data/import/elements.json")
diff --git a/arkindex/process/tests/test_workerruns_use_gpu.py b/arkindex/process/tests/test_workerruns_use_gpu.py
deleted file mode 100644
index 79704d3cf80653b5210d34b9c889286eb76a8c54..0000000000000000000000000000000000000000
--- a/arkindex/process/tests/test_workerruns_use_gpu.py
+++ /dev/null
@@ -1,266 +0,0 @@
-from django.urls import reverse
-from rest_framework import status
-
-from arkindex.ponos.models import Farm
-from arkindex.process.models import FeatureUsage, ProcessMode, Worker, WorkerRun, WorkerVersion, WorkerVersionState
-from arkindex.project.tests import FixtureAPITestCase
-
-
-class TestWorkerRunsGPU(FixtureAPITestCase):
- @classmethod
- def setUpTestData(cls):
- super().setUpTestData()
- cls.farm = Farm.objects.first()
- cls.test_process = cls.corpus.processes.create(
- creator=cls.user,
- mode=ProcessMode.Workers,
- farm=cls.farm,
- )
- cls.worker = Worker.objects.get(slug="reco")
- cls.version_gpu_required = WorkerVersion.objects.create(
- worker=cls.worker,
- configuration={"pilot": "ikari shinji"},
- state=WorkerVersionState.Available,
- model_usage=FeatureUsage.Disabled,
- docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
- gpu_usage=FeatureUsage.Required,
- version=2
- )
- cls.version_gpu_supported = WorkerVersion.objects.create(
- worker=cls.worker,
- configuration={"pilot": "ayanami rei"},
- state=WorkerVersionState.Available,
- model_usage=FeatureUsage.Disabled,
- docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
- gpu_usage=FeatureUsage.Supported,
- version=3
- )
- cls.version_gpu_disabled = WorkerVersion.objects.create(
- worker=cls.worker,
- configuration={"pilot": "soryu asuka langley"},
- state=WorkerVersionState.Available,
- model_usage=FeatureUsage.Disabled,
- docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
- gpu_usage=FeatureUsage.Disabled,
- version=4
- )
-
- def test_create_gpu_usage_auto(self):
- """
- use_gpu on the worker run is set automatically depending on the worker version's gpu_usage
- """
- self.client.force_login(self.user)
-
- cases = [
- (self.version_gpu_disabled, False),
- (self.version_gpu_supported, False),
- (self.version_gpu_required, True)
- ]
-
- for worker_version, use_gpu in cases:
- with self.subTest(worker_version=worker_version, use_gpu=use_gpu):
- with self.assertNumQueries(6):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.test_process.id)}),
- {"worker_version_id": str(worker_version.id), "parents": []},
- )
- self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-
- data = response.json()
- pk = data.pop("id")
- self.assertDictEqual(data, {
- "worker_version": {
- "id": str(worker_version.id),
- "configuration": worker_version.configuration,
- "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
- "gpu_usage": worker_version.gpu_usage.value,
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": worker_version.version,
- "tag": None,
- "created": worker_version.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.test_process.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "summary": f"Worker Recognizer @ version {worker_version.version}",
- "use_gpu": use_gpu
- })
- run = WorkerRun.objects.get(pk=pk)
- self.assertEqual(run.use_gpu, use_gpu)
-
- def test_create_use_gpu_ignored(self):
- """
- if "use_gpu" is being sent when creating a worker run, it is ignored
- """
- self.client.force_login(self.user)
- with self.assertNumQueries(6):
- response = self.client.post(
- reverse("api:worker-run-list", kwargs={"pk": str(self.test_process.id)}),
- {"worker_version_id": str(self.version_gpu_required.id), "parents": [], "use_gpu": False},
- )
- self.assertEqual(response.status_code, status.HTTP_201_CREATED)
-
- data = response.json()
- pk = data.pop("id")
- self.assertDictEqual(data, {
- "worker_version": {
- "id": str(self.version_gpu_required.id),
- "configuration": {"pilot": "ikari shinji"},
- "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
- "gpu_usage": FeatureUsage.Required.value,
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": 2,
- "tag": None,
- "created": self.version_gpu_required.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.test_process.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "summary": "Worker Recognizer @ version 2",
- "use_gpu": True
- })
- run = WorkerRun.objects.get(pk=pk)
- self.assertEqual(run.use_gpu, True)
-
- def test_update_use_gpu_errors(self):
- self.client.force_login(self.user)
-
- cases = [
- (self.version_gpu_disabled, True, "This worker version does not support GPU usage."),
- (self.version_gpu_required, False, "This worker version requires GPU usage.")
- ]
-
- for worker_version, use_gpu, error_message in cases:
- with self.subTest(worker_version=worker_version, use_gpu=use_gpu, error_message=error_message):
- run = WorkerRun.objects.create(
- process=self.test_process,
- version=worker_version,
- parents=[]
- )
- self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
- with self.assertNumQueries(3):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={"use_gpu": use_gpu}
- )
- self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
- self.assertDictEqual(response.json(), {"use_gpu": [error_message]})
-
- run.refresh_from_db()
- self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
-
- def test_update_use_gpu(self):
- self.client.force_login(self.user)
-
- cases = [
- (self.version_gpu_disabled, False),
- (self.version_gpu_required, True),
- (self.version_gpu_supported, False),
- (self.version_gpu_supported, True)
- ]
-
- for worker_version, use_gpu in cases:
- with self.subTest(worker_version=worker_version, use_gpu=use_gpu):
- self.test_process.worker_runs.all().delete()
- run = WorkerRun.objects.create(
- process=self.test_process,
- version=worker_version,
- parents=[]
- )
- self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
-
- with self.assertNumQueries(5):
- response = self.client.patch(
- reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
- data={"use_gpu": use_gpu}
- )
- self.assertEqual(response.status_code, status.HTTP_200_OK)
-
- data = response.json()
- self.assertDictEqual(data, {
- "id": str(run.id),
- "worker_version": {
- "id": str(worker_version.id),
- "configuration": worker_version.configuration,
- "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
- "gpu_usage": worker_version.gpu_usage.value,
- "model_usage": FeatureUsage.Disabled.value,
- "revision_url": None,
- "version": worker_version.version,
- "tag": None,
- "created": worker_version.created.isoformat().replace("+00:00", "Z"),
- "state": "available",
- "worker": {
- "id": str(self.worker.id),
- "name": "Recognizer",
- "slug": "reco",
- "type": "recognizer",
- "description": "",
- "repository_url": None,
- "archived": False,
- }
- },
- "parents": [],
- "model_version": None,
- "configuration": None,
- "process": {
- "id": str(self.test_process.id),
- "activity_state": "disabled",
- "corpus": str(self.corpus.id),
- "element": None,
- "element_type": None,
- "folder_type": None,
- "prefix": None,
- "chunks": 1,
- "mode": "workers",
- "name": None,
- "state": "unscheduled",
- "use_cache": False,
- },
- "summary": f"Worker Recognizer @ version {worker_version.version}",
- "use_gpu": use_gpu
- })
- run.refresh_from_db()
- self.assertEqual(run.use_gpu, use_gpu)
diff --git a/arkindex/process/tests/worker_runs/__init__.py b/arkindex/process/tests/worker_runs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/arkindex/process/tests/worker_runs/test_build_task.py b/arkindex/process/tests/worker_runs/test_build_task.py
new file mode 100644
index 0000000000000000000000000000000000000000..51c8d672036b28ae34b19aced225b83778c1487a
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_build_task.py
@@ -0,0 +1,175 @@
+from arkindex.ponos.models import Farm
+from arkindex.process.models import ProcessMode, WorkerVersion, WorkerVersionState
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.training.models import Model, ModelVersion, ModelVersionState
+from arkindex.users.models import Role
+
+ENV = {
+ "ARKINDEX_PROCESS_ID": "12345"
+}
+
+
+class TestWorkerRunsBuildTask(FixtureAPITestCase):
+ """
+ Test task creation from worker runs
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.farm = Farm.objects.first()
+ cls.process = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker = cls.version.worker
+ cls.worker_run = cls.process.worker_runs.create(version=cls.version, parents=[])
+
+ # Model and Model version setup
+ cls.model_1 = Model.objects.create(name="My model")
+ cls.model_1.memberships.create(user=cls.user, level=Role.Contributor.value)
+ cls.model_version = ModelVersion.objects.create(
+ model=cls.model_1,
+ state=ModelVersionState.Available,
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ def test_build_task_no_parent(self):
+ task, parent_slugs = self.worker_run.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json")
+
+ self.assertEqual(task.slug, f"reco_{str(self.worker_run.id)[0:6]}")
+ self.assertEqual(task.image, self.version.docker_image_iid)
+ self.assertEqual(task.command, None)
+ self.assertEqual(task.shm_size, None)
+ self.assertEqual(parent_slugs, ["import"])
+ self.assertEqual(task.env, {
+ "ARKINDEX_PROCESS_ID": "12345",
+ "ARKINDEX_TASK_TOKEN": str(task.token),
+ "TASK_ELEMENTS": "/data/import/elements.json",
+ "ARKINDEX_WORKER_RUN_ID": str(self.worker_run.id),
+ })
+
+ def test_build_task_with_chunk(self):
+ task, parent_slugs = self.worker_run.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json", chunk=4)
+
+ self.assertEqual(task.slug, f"reco_{str(self.worker_run.id)[0:6]}_4")
+ self.assertEqual(task.image, self.version.docker_image_iid)
+ self.assertEqual(task.command, None)
+ self.assertEqual(task.shm_size, None)
+ self.assertEqual(parent_slugs, ["import"])
+ self.assertEqual(task.env, {
+ "ARKINDEX_PROCESS_ID": "12345",
+ "ARKINDEX_TASK_TOKEN": str(task.token),
+ "ARKINDEX_TASK_CHUNK": "4",
+ "TASK_ELEMENTS": "/data/import/elements.json",
+ "ARKINDEX_WORKER_RUN_ID": str(self.worker_run.id),
+ })
+
+ def test_build_task_with_parent(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ version_2.docker_image_iid = "evaunit:latest"
+ version_2.state = WorkerVersionState.Available
+ run_2 = self.process.worker_runs.create(
+ version=version_2,
+ parents=[self.worker_run.id],
+ )
+
+ task, parent_slugs = run_2.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json")
+
+ self.assertEqual(task.slug, f"reco_{str(run_2.id)[0:6]}")
+ self.assertEqual(task.image, version_2.docker_image_iid)
+ self.assertEqual(task.command, None)
+ self.assertEqual(task.shm_size, None)
+ self.assertEqual(parent_slugs, [f"reco_{str(self.worker_run.id)[0:6]}"])
+ self.assertEqual(task.env, {
+ "ARKINDEX_PROCESS_ID": "12345",
+ "ARKINDEX_TASK_TOKEN": str(task.token),
+ "TASK_ELEMENTS": "/data/import/elements.json",
+ "ARKINDEX_WORKER_RUN_ID": str(run_2.id),
+ })
+
+ def test_build_task_with_parent_and_chunk(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ version_2.docker_image_iid = "evaunit:latest"
+ version_2.state = WorkerVersionState.Available
+ run_2 = self.process.worker_runs.create(
+ version=version_2,
+ parents=[self.worker_run.id],
+ )
+
+ task, parent_slugs = run_2.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json", chunk=4)
+
+ self.assertEqual(task.slug, f"reco_{str(run_2.id)[0:6]}_4")
+ self.assertEqual(task.image, version_2.docker_image_iid)
+ self.assertEqual(task.command, None)
+ self.assertEqual(task.shm_size, None)
+ self.assertEqual(parent_slugs, [f"reco_{str(self.worker_run.id)[0:6]}_4"])
+ self.assertEqual(task.env, {
+ "ARKINDEX_PROCESS_ID": "12345",
+ "ARKINDEX_TASK_TOKEN": str(task.token),
+ "ARKINDEX_TASK_CHUNK": "4",
+ "TASK_ELEMENTS": "/data/import/elements.json",
+ "ARKINDEX_WORKER_RUN_ID": str(run_2.id),
+ })
+
+ def test_build_task_shm_size(self):
+ self.version.configuration = {
+ "docker": {
+ "shm_size": 505,
+ }
+ }
+ task, parent_slugs = self.worker_run.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json")
+
+ self.assertEqual(task.slug, f"reco_{str(self.worker_run.id)[0:6]}")
+ self.assertEqual(task.image, self.version.docker_image_iid)
+ self.assertEqual(task.command, None)
+ self.assertEqual(task.shm_size, 505)
+ self.assertEqual(parent_slugs, ["import"])
+ self.assertEqual(task.env, {
+ "ARKINDEX_PROCESS_ID": "12345",
+ "ARKINDEX_TASK_TOKEN": str(task.token),
+ "TASK_ELEMENTS": "/data/import/elements.json",
+ "ARKINDEX_WORKER_RUN_ID": str(self.worker_run.id),
+ })
+
+ def test_build_task_unavailable_version(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ version_2.docker_image_iid = "evaunit:latest"
+ self.assertEqual(version_2.state, WorkerVersionState.Created)
+ run_2 = self.process.worker_runs.create(
+ version=version_2,
+ parents=[self.worker_run.id],
+ )
+
+ with self.assertRaisesRegex(
+ AssertionError,
+ f"Worker Version {version_2.id} is not available and cannot be used to build a task."
+ ):
+ run_2.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json")
+
+ def test_build_task_unavailable_model_version(self):
+ self.model_version.state = ModelVersionState.Created
+ self.model_version.save()
+ self.worker_run.model_version = self.model_version
+ self.worker_run.save()
+ with self.assertRaisesRegex(
+ AssertionError,
+ f"ModelVersion {self.model_version.id} is not available and cannot be used to build a task."
+ ):
+ self.worker_run.build_task(self.process, ENV.copy(), "import", "/data/import/elements.json")
diff --git a/arkindex/process/tests/worker_runs/test_create.py b/arkindex/process/tests/worker_runs/test_create.py
new file mode 100644
index 0000000000000000000000000000000000000000..0cf94a07eb15d8e4a55c9a35e056c5a2575b04f5
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_create.py
@@ -0,0 +1,603 @@
+import uuid
+from datetime import datetime, timezone
+from unittest.mock import call, patch
+
+from django.db import transaction
+from django.urls import reverse
+from rest_framework import status
+
+from arkindex.ponos.models import Farm
+from arkindex.process.models import (
+ FeatureUsage,
+ ProcessMode,
+ WorkerRun,
+ WorkerVersion,
+ WorkerVersionState,
+)
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.training.models import Model, ModelVersion, ModelVersionState
+from arkindex.users.models import Role
+
+
+class TestWorkerRunsCreate(FixtureAPITestCase):
+ """
+ Test worker runs create endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+ cls.configuration_1 = cls.worker_1.configurations.create(name="My config", configuration={"key": "value"})
+ worker_version = WorkerVersion.objects.exclude(worker=cls.version_1.worker).first()
+ cls.configuration_2 = worker_version.worker.configurations.create(name="Config")
+ cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
+
+ # Model and Model version setup
+ cls.model_1 = Model.objects.create(name="My model")
+ cls.model_1.memberships.create(user=cls.user, level=Role.Contributor.value)
+ cls.model_version_1 = ModelVersion.objects.create(
+ model=cls.model_1,
+ state=ModelVersionState.Available,
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+ cls.version_gpu_required = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "ikari shinji"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Required,
+ version=2
+ )
+ cls.version_gpu_supported = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "ayanami rei"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Supported,
+ version=3
+ )
+ cls.version_gpu_disabled = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "soryu asuka langley"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Disabled,
+ version=4
+ )
+
+ def test_create_requires_login(self):
+ with self.assertNumQueries(0):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ self.assertDictEqual(response.json(), {"detail": "Authentication credentials were not provided."})
+
+ @patch("arkindex.users.utils.get_max_level", return_value=Role.Guest.value)
+ def test_create_no_execution_right(self, max_level_mock):
+ """
+ An execution access on the target worker version is required to create a worker run
+ """
+ self.worker_1.memberships.update(level=Role.Guest.value)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {"worker_version_id": ["You do not have an execution access to this worker."]})
+
+ self.assertListEqual(max_level_mock.call_args_list, [
+ call(self.user, self.worker_1)
+ ])
+
+ def test_create_invalid_version(self):
+ self.client.force_login(self.user)
+ version_id = uuid.uuid4()
+
+ with self.assertNumQueries(4):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": version_id},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "worker_version_id": [f'Invalid pk "{version_id}" - object does not exist.'],
+ })
+
+ def test_create_unique(self):
+ self.client.force_login(self.user)
+ self.version_1.model_usage = FeatureUsage.Required
+ self.version_1.save()
+ cases = [
+ (None, None),
+ (None, self.configuration_1),
+ (self.model_version_1, None),
+ (self.model_version_1, self.configuration_1),
+ ]
+ for model_version, configuration in cases:
+ with self.subTest(model_version=model_version, configuration=configuration):
+ self.run_1.model_version = model_version
+ self.run_1.configuration = configuration
+ self.run_1.save()
+
+ # Having a model version or a configuration adds one query for each
+ query_count = 5 + bool(model_version) + bool(configuration)
+
+ with self.assertNumQueries(query_count):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}),
+ data={
+ "worker_version_id": str(self.version_1.id),
+ "model_version_id": str(model_version.id) if model_version else None,
+ "configuration_id": str(configuration.id) if configuration else None,
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
+ })
+
+ def test_create_unavailable_version(self):
+ self.version_1.state = WorkerVersionState.Error
+ self.version_1.save()
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {"worker_version_id": ["This WorkerVersion is not in an Available state."]})
+
+ def test_create_archived_worker(self):
+ self.worker_1.archived = datetime.now(timezone.utc)
+ self.worker_1.save()
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {"worker_version_id": ["This WorkerVersion is part of an archived worker."]})
+
+ def test_create_archived_model(self):
+ self.model_1.archived = datetime.now(timezone.utc)
+ self.model_1.save()
+ self.version_1.model_usage = FeatureUsage.Supported
+ self.version_1.save()
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "model_version_id": str(self.model_version_1.id)},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {"model_version_id": ["This ModelVersion is part of an archived model."]})
+
+ def test_create_invalid_process_id(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": uuid.uuid4()}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ self.assertEqual(response.json(), {"detail": "No Process matches the given query."})
+
+ def test_create_invalid_process_mode(self):
+ self.client.force_login(self.user)
+
+ for mode in set(ProcessMode) - {ProcessMode.Workers, ProcessMode.Dataset, ProcessMode.Local}:
+ with self.subTest(mode=mode):
+ self.process_2.mode = mode
+ self.process_2.save()
+
+ with self.assertNumQueries(5):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ {"worker_version_id": str(self.version_1.id)},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
+ })
+
+ def test_create_non_corpus_process_mode(self):
+ process = self.user.processes.get(mode=ProcessMode.Local)
+ self.client.force_login(self.user)
+ with self.assertNumQueries(3):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(process.id)}),
+ {"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_create_process_already_started(self):
+ process = self.corpus.processes.create(
+ creator=self.user,
+ mode=ProcessMode.Workers,
+ farm=self.farm,
+ )
+ process.run()
+
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(process.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
+ })
+
+ def test_create(self):
+ self.client.force_login(self.user)
+
+ for mode in (ProcessMode.Workers, ProcessMode.Dataset):
+ self.process_2.worker_runs.filter(version=self.version_1).delete()
+
+ with self.subTest(mode=mode):
+ self.process_2.mode = mode
+ self.process_2.save()
+
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ {"worker_version_id": str(self.version_1.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+ data = response.json()
+ pk = data.pop("id")
+ self.assertNotEqual(pk, self.run_1.id)
+ self.assertDictEqual(data, {
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_2.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": mode.value,
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+ run = WorkerRun.objects.get(pk=pk)
+ # Check generated summary
+ self.assertEqual(run.summary, "Worker Recognizer @ version 1")
+
+ def test_create_empty(self):
+ """
+ The worker_version_id is required to create a worker run
+ """
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertDictEqual(response.json(), {
+ "worker_version_id": ["This field is required."],
+ })
+
+ def test_create_configuration(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(7):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={
+ "worker_version_id": str(self.version_1.id),
+ "parents": [],
+ "configuration_id": str(self.configuration_1.id)
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+ data = response.json()
+ pk = data.pop("id")
+ self.assertNotEqual(pk, self.run_1.id)
+ self.assertDictEqual(data, {
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": {
+ "id": str(self.configuration_1.id),
+ "archived": False,
+ "configuration": {"key": "value"},
+ "name": "My config"
+ },
+ "process": {
+ "id": str(self.process_2.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
+ })
+ run = WorkerRun.objects.get(pk=pk)
+ # Check generated summary
+ self.assertEqual(run.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
+
+ def test_create_invalid_configuration(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data={"worker_version_id": str(self.version_1.id), "parents": [], "configuration_id": str(self.configuration_2.id)},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
+
+ def test_create_summary(self):
+ self.client.force_login(self.user)
+ test_version = self.worker_1.versions.create(
+ version=5,
+ state=WorkerVersionState.Available,
+ docker_image_iid="registry.gitlab.com/dead-sea-scrolls:004",
+ model_usage=FeatureUsage.Supported
+ )
+
+ cases = [
+ ("eva-01", "revision_url", self.model_version_1, self.configuration_1, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]}) with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
+ (None, "revision_url", self.model_version_1, self.configuration_1, f"Worker Recognizer @ {str(test_version.id)[:6]} with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
+ ("eva-01", "revision_url", None, self.configuration_1, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]}) using configuration 'My config'"),
+ (None, "revision_url", self.model_version_1, None, f"Worker Recognizer @ {str(test_version.id)[:6]} with model My model @ {str(self.model_version_1.id)[:6]}"),
+ ("eva-01", "revision_url", None, None, f"Worker Recognizer @ eva-01 ({str(test_version.id)[:6]})"),
+ (None, "revision_url", None, None, f"Worker Recognizer @ {str(test_version.id)[:6]}"),
+ ("eva-01", None, self.model_version_1, self.configuration_1, f"Worker Recognizer @ eva-01 (version 5) with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
+ (None, None, self.model_version_1, self.configuration_1, f"Worker Recognizer @ version 5 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'"),
+ ("eva-01", None, None, self.configuration_1, "Worker Recognizer @ eva-01 (version 5) using configuration 'My config'"),
+ (None, None, self.model_version_1, None, f"Worker Recognizer @ version 5 with model My model @ {str(self.model_version_1.id)[:6]}"),
+ ("eva-01", None, None, None, "Worker Recognizer @ eva-01 (version 5)"),
+ (None, None, None, None, "Worker Recognizer @ version 5"),
+ ]
+
+ for tag, revision_url, model_version, config, expected_summary in cases:
+ with self.subTest(tag=tag, model_version=model_version, config=config), transaction.atomic():
+ # Clear the process of worker runs
+ self.process_2.worker_runs.all().delete()
+
+ num_queries = 6
+
+ test_version.version = None
+ test_version.tag = tag
+ test_version.revision_url = revision_url
+ if not revision_url:
+ test_version.version = 5
+ test_version.save()
+
+ payload = {
+ "worker_version_id": str(test_version.id),
+ "parents": [],
+ }
+ if model_version:
+ payload["model_version_id"] = model_version.id
+ # If there is a model version, it adds a query
+ num_queries += 1
+ if config:
+ payload["configuration_id"] = config.id
+ # If there is a configuration, it adds a query
+ num_queries += 1
+
+ with self.assertNumQueries(num_queries):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}),
+ data=payload,
+ )
+ if response.status_code == 400:
+ self.assertDictEqual(response.json(), {})
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+ created_run = WorkerRun.objects.get(pk=response.json()["id"])
+ self.assertEqual(created_run.summary, expected_summary)
+
+ def test_create_gpu_usage_auto(self):
+ """
+ use_gpu on the worker run is set automatically depending on the worker version's gpu_usage
+ """
+ self.client.force_login(self.user)
+
+ cases = [
+ (self.version_gpu_disabled, False),
+ (self.version_gpu_supported, False),
+ (self.version_gpu_required, True)
+ ]
+
+ for worker_version, use_gpu in cases:
+ with self.subTest(worker_version=worker_version, use_gpu=use_gpu):
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}),
+ {"worker_version_id": str(worker_version.id), "parents": []},
+ )
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+ data = response.json()
+ pk = data.pop("id")
+ self.assertDictEqual(data, {
+ "worker_version": {
+ "id": str(worker_version.id),
+ "configuration": worker_version.configuration,
+ "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
+ "gpu_usage": worker_version.gpu_usage.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": worker_version.version,
+ "tag": None,
+ "created": worker_version.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "summary": f"Worker Recognizer @ version {worker_version.version}",
+ "use_gpu": use_gpu
+ })
+ run = WorkerRun.objects.get(pk=pk)
+ self.assertEqual(run.use_gpu, use_gpu)
+
+ def test_create_use_gpu_ignored(self):
+ """
+ if "use_gpu" is being sent when creating a worker run, it is ignored
+ """
+ self.client.force_login(self.user)
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}),
+ {"worker_version_id": str(self.version_gpu_required.id), "parents": [], "use_gpu": False},
+ )
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
+ data = response.json()
+ pk = data.pop("id")
+ self.assertDictEqual(data, {
+ "worker_version": {
+ "id": str(self.version_gpu_required.id),
+ "configuration": {"pilot": "ikari shinji"},
+ "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
+ "gpu_usage": FeatureUsage.Required.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 2,
+ "tag": None,
+ "created": self.version_gpu_required.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "summary": "Worker Recognizer @ version 2",
+ "use_gpu": True
+ })
+ run = WorkerRun.objects.get(pk=pk)
+ self.assertEqual(run.use_gpu, True)
diff --git a/arkindex/process/tests/worker_runs/test_delete.py b/arkindex/process/tests/worker_runs/test_delete.py
new file mode 100644
index 0000000000000000000000000000000000000000..a288d868afd7736e671c3cf75c13202bc062ef70
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_delete.py
@@ -0,0 +1,170 @@
+from datetime import datetime, timezone
+
+from django.urls import reverse
+from rest_framework import status
+
+from arkindex.ponos.models import Agent, Farm, State
+from arkindex.process.models import ProcessMode, WorkerRun, WorkerVersion
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.users.models import Role
+
+
+class TestWorkerRunsDelete(FixtureAPITestCase):
+ """
+ Test worker runs delete endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+ cls.version_2 = WorkerVersion.objects.get(worker__slug="dla")
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+
+ cls.agent = Agent.objects.create(
+ farm=cls.farm,
+ hostname="claude",
+ cpu_cores=42,
+ cpu_frequency=1e15,
+ ram_total=99e9,
+ last_ping=datetime.now(timezone.utc),
+ )
+ # Add custom attributes to make the agent usable as an authenticated user
+ cls.agent.is_agent = True
+ cls.agent.is_anonymous = False
+
+ def test_delete_requires_login(self):
+ with self.assertNumQueries(0):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ def test_delete_invalid_id(self):
+ self.client.force_login(self.user)
+ with self.assertNumQueries(3):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"})
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_delete_no_worker_execution_right(self):
+ """
+ A user can delete a worker run with no right on its worker but an admin access to the process project
+ """
+ self.worker_1.memberships.update(level=Role.Guest.value)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
+
+ with self.assertRaises(WorkerRun.DoesNotExist):
+ self.run_1.refresh_from_db()
+
+ def test_delete_local(self):
+ """
+ A user cannot delete a worker run on a local process
+ """
+ run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(4):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), ["WorkerRuns can only be deleted from Workers or Dataset processes."])
+
+ def test_delete(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
+
+ with self.assertRaises(WorkerRun.DoesNotExist):
+ self.run_1.refresh_from_db()
+
+ def test_delete_with_children(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ version_3 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=3,
+ configuration={"test": "test3"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[self.run_1.id],
+ )
+ run_3 = self.process_1.worker_runs.create(
+ version=version_3,
+ parents=[self.run_1.id, run_2.id],
+ )
+
+ self.assertTrue(self.run_1.id in run_2.parents)
+ self.assertTrue(self.run_1.id in run_3.parents)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
+
+ with self.assertRaises(WorkerRun.DoesNotExist):
+ self.run_1.refresh_from_db()
+ run_2.refresh_from_db()
+ run_3.refresh_from_db()
+ self.assertEqual(run_2.parents, [])
+ self.assertEqual(run_3.parents, [run_2.id])
+
+ def test_delete_started_process(self):
+ """
+ A user shouldn't be able to delete the worker run of a started process
+ """
+ self.client.force_login(self.user)
+ self.process_1.run()
+ self.process_1.tasks.update(state=State.Running)
+
+ with self.assertNumQueries(3):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), ["WorkerRuns cannot be deleted from a process that has already started."])
+
+ def test_delete_agent(self):
+ """
+ Ponos agents cannot delete WorkerRuns, even when they can access them
+ """
+ self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
+
+ # Agent auth is not implemented in CE
+ self.client.force_authenticate(user=self.agent)
+ with self.assertNumQueries(1):
+ response = self.client.delete(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+ self.assertEqual(response.json(), {
+ "detail": "You do not have an admin access to this process.",
+ })
diff --git a/arkindex/process/tests/worker_runs/test_list.py b/arkindex/process/tests/worker_runs/test_list.py
new file mode 100644
index 0000000000000000000000000000000000000000..63123a908d33f5c725d11be5f6fd3994356b6b81
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_list.py
@@ -0,0 +1,108 @@
+from django.urls import reverse
+from rest_framework import status
+
+from arkindex.ponos.models import Farm
+from arkindex.process.models import FeatureUsage, ProcessMode, WorkerVersion
+from arkindex.project.tests import FixtureAPITestCase
+
+
+class TestWorkerRunsList(FixtureAPITestCase):
+ """
+ Test worker runs list endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+ cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
+
+ def test_list_requires_login(self):
+ with self.assertNumQueries(0):
+ response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ def test_list_no_execution_right(self):
+ """
+ Worker runs attached to a process can be listed even if the user has no execution rights to workers
+ This is due to the fact a user can see a process running on a corpus they have access
+ """
+ self.worker_1.memberships.all().delete()
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json()["count"], 1)
+
+ def test_list(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_1.id)}))
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ data = response.json()
+ self.assertEqual(data["results"], [{
+ "id": str(self.run_1.id),
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ }])
+
+ def test_list_filter_process(self):
+ run_2 = self.process_2.worker_runs.create(
+ version=self.version_1,
+ parents=[],
+ )
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(6):
+ response = self.client.get(reverse("api:worker-run-list", kwargs={"pk": str(self.process_2.id)}))
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ data = response.json()
+ self.assertEqual(data["count"], 1)
+ self.assertEqual(data["results"][0]["id"], str(run_2.id))
diff --git a/arkindex/process/tests/worker_runs/test_partial_update.py b/arkindex/process/tests/worker_runs/test_partial_update.py
new file mode 100644
index 0000000000000000000000000000000000000000..43b484a6a5634da378cbe41d20f369784875b316
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_partial_update.py
@@ -0,0 +1,1058 @@
+import uuid
+from datetime import datetime, timezone
+from unittest.mock import call, patch
+
+from django.urls import reverse
+from rest_framework import status
+
+from arkindex.ponos.models import Agent, Farm
+from arkindex.process.models import FeatureUsage, ProcessMode, WorkerRun, WorkerVersion, WorkerVersionState
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.training.models import Model, ModelVersion, ModelVersionState
+from arkindex.users.models import Role
+
+
+class TestWorkerRunsPartialUpdate(FixtureAPITestCase):
+ """
+ Test worker runs partial update endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+ cls.version_2 = WorkerVersion.objects.get(worker__slug="dla")
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+ cls.configuration_1 = cls.worker_1.configurations.create(name="My config", configuration={"key": "value"})
+ worker_version = WorkerVersion.objects.exclude(worker=cls.version_1.worker).first()
+ cls.configuration_2 = worker_version.worker.configurations.create(name="Config")
+ cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
+ cls.version_gpu_required = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "ikari shinji"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Required,
+ version=3
+ )
+ cls.version_gpu_supported = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "ayanami rei"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Supported,
+ version=4
+ )
+ cls.version_gpu_disabled = WorkerVersion.objects.create(
+ worker=cls.worker_1,
+ configuration={"pilot": "soryu asuka langley"},
+ state=WorkerVersionState.Available,
+ model_usage=FeatureUsage.Disabled,
+ docker_image_iid="registry.nerv.co.jp/neon-genesis:evangelion",
+ gpu_usage=FeatureUsage.Disabled,
+ version=5
+ )
+
+ # Model and Model version setup
+ cls.model_1 = Model.objects.create(name="My model")
+ cls.model_1.memberships.create(user=cls.user, level=Role.Contributor.value)
+ cls.model_version_1 = ModelVersion.objects.create(
+ model=cls.model_1,
+ state=ModelVersionState.Available,
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.model_2 = Model.objects.create(name="Their model")
+ cls.model_2.memberships.create(user=cls.user, level=Role.Guest.value)
+ cls.model_version_2 = cls.model_2.versions.create(
+ state=ModelVersionState.Available,
+ tag="blah",
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.model_3 = Model.objects.create(name="Our model", public=True)
+ cls.model_version_3 = cls.model_3.versions.create(
+ state=ModelVersionState.Available,
+ tag="blah",
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.agent = Agent.objects.create(
+ farm=cls.farm,
+ hostname="claude",
+ cpu_cores=42,
+ cpu_frequency=1e15,
+ ram_total=99e9,
+ last_ping=datetime.now(timezone.utc),
+ )
+ # Add custom attributes to make the agent usable as an authenticated user
+ cls.agent.is_agent = True
+ cls.agent.is_anonymous = False
+
+ def test_partial_update_requires_login(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+
+ with self.assertNumQueries(0):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ @patch("arkindex.project.mixins.get_max_level", return_value=Role.Contributor.value)
+ def test_partial_update_no_project_admin_right(self, get_max_level_mock):
+ """
+ A user cannot update a worker run if they have no admin access on its process project
+ """
+ self.corpus.memberships.filter(user=self.user).update(level=Role.Contributor.value)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ self.assertEqual(response.json(), {"detail": "You do not have an admin access to this process."})
+
+ self.assertEqual(get_max_level_mock.call_count, 1)
+ self.assertEqual(get_max_level_mock.call_args, call(self.user, self.corpus))
+
+ def test_partial_update_invalid_id(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+
+ self.client.force_login(self.user)
+ with self.assertNumQueries(3):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_partial_update_local(self):
+ """
+ A user cannot update a worker run on a local process
+ """
+ run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "parents": []
+ }
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
+ })
+
+ def test_partial_update_nonexistent_parent(self):
+ self.client.force_login(self.user)
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": ["12341234-1234-1234-1234-123412341234"],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), [
+ f"Can't add or update WorkerRun {self.run_1.id} because parents field isn't properly defined. It can be either because"
+ " one or several UUIDs don't refer to existing WorkerRuns or either because listed WorkerRuns doesn't belong to the"
+ " same Process than this WorkerRun."
+ ])
+
+ def test_partial_update_process_id(self):
+ """
+ Process field cannot be updated
+ """
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "process_id": str(self.process_2.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+ self.run_1.refresh_from_db()
+ self.assertEqual(self.run_1.process.id, self.process_1.id)
+
+ def test_partial_update_worker_version_id(self):
+ """
+ Version field cannot be updated
+ """
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "worker_version_id": str(self.version_2.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+ self.run_1.refresh_from_db()
+ self.assertNotEqual(self.run_1.version_id, self.version_2.id)
+
+ def test_partial_update_configuration(self):
+ self.client.force_login(self.user)
+ self.assertEqual(self.run_1.configuration, None)
+ self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1")
+
+ with self.assertNumQueries(6):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
+ data={
+ "configuration_id": str(self.configuration_1.id)
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.run_1.refresh_from_db()
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": {
+ "id": str(self.configuration_1.id),
+ "archived": False,
+ "configuration": {"key": "value"},
+ "name": "My config"
+ },
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
+ })
+ self.assertEqual(self.run_1.configuration.id, self.configuration_1.id)
+ self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
+
+ def test_partial_update_invalid_configuration(self):
+ self.client.force_login(self.user)
+ self.assertEqual(self.run_1.configuration, None)
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
+ data={"configuration_id": str(self.configuration_2.id)},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
+
+ def test_partial_update_process_already_started(self):
+ """
+ Update dependencies of a worker run is not possible once the process is started
+ """
+ self.process_1.run()
+ self.assertTrue(self.process_1.tasks.exists())
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(4):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
+ })
+
+ def test_partial_update_model_version_not_allowed(self):
+ """
+ The model_version UUID is not allowed when the related version doesn't allow model_usage
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Disabled
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This worker version does not support models."]
+ })
+
+ def test_partial_update_unknown_model_version(self):
+ """
+ Cannot use a model version id that doesn't exist
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+ random_model_version_uuid = str(uuid.uuid4())
+
+ with self.assertNumQueries(4):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": random_model_version_uuid,
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{random_model_version_uuid}" - object does not exist.']
+ })
+
+ @patch("arkindex.users.managers.BaseACLManager.filter_rights", return_value=ModelVersion.objects.none())
+ def test_partial_update_model_version_no_access(self, filter_rights_mock):
+ """
+ Cannot update a worker run with a model_version UUID, when you don't have access to the model version
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ # Create a model version, the user has no access to
+ model_no_access = Model.objects.create(name="Secret model")
+ model_version_no_access = ModelVersion.objects.create(model=model_no_access, state=ModelVersionState.Available, size=8, hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
+
+ with self.assertNumQueries(3):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(model_version_no_access.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{model_version_no_access.id}" - object does not exist.'],
+ })
+
+ self.assertListEqual(filter_rights_mock.call_args_list, [
+ call(self.user, Model, Role.Guest.value),
+ call(self.user, Model, Role.Contributor.value),
+ ])
+
+ @patch("arkindex.users.managers.BaseACLManager.filter_rights")
+ def test_partial_update_model_version_guest(self, filter_rights_mock):
+ """
+ Cannot update a worker run with a model_version when you only have guest access to the model,
+ and the model version has no tag or is not available
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ def filter_rights(user, model, level):
+ """
+ The filter_rights mock needs to return nothing when called for contributor access,
+ and the models we will test on when called for guest access
+ """
+ if level == Role.Guest.value:
+ return Model.objects.filter(id__in=(self.model_2.id, self.model_3.id))
+ return Model.objects.none()
+
+ filter_rights_mock.side_effect = filter_rights
+
+ cases = [
+ # On a model with a membership giving guest access
+ (self.model_version_2, None, ModelVersionState.Created),
+ (self.model_version_2, None, ModelVersionState.Error),
+ (self.model_version_2, None, ModelVersionState.Available),
+ (self.model_version_2, "blah", ModelVersionState.Created),
+ (self.model_version_2, "blah", ModelVersionState.Error),
+ # On a public model with no membership
+ (self.model_version_3, None, ModelVersionState.Created),
+ (self.model_version_3, None, ModelVersionState.Error),
+ (self.model_version_3, None, ModelVersionState.Available),
+ (self.model_version_3, "blah", ModelVersionState.Created),
+ (self.model_version_3, "blah", ModelVersionState.Error),
+ ]
+
+ for model_version, tag, state in cases:
+ filter_rights_mock.reset_mock()
+ with self.subTest(model_version=model_version, tag=tag, state=state):
+ model_version.tag = tag
+ model_version.state = state
+ model_version.save()
+
+ with self.assertNumQueries(4):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(model_version.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{model_version.id}" - object does not exist.'],
+ })
+
+ self.assertListEqual(filter_rights_mock.call_args_list, [
+ call(self.user, Model, Role.Guest.value),
+ call(self.user, Model, Role.Contributor.value),
+ ])
+
+ def test_partial_update_model_version_unavailable(self):
+ self.client.force_login(self.user)
+ version = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(
+ version=version,
+ parents=[],
+ )
+ self.model_version_1.state = ModelVersionState.Error
+ self.model_version_1.save()
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This ModelVersion is not in an Available state."],
+ })
+
+ def test_partial_update_model_archived(self):
+ self.client.force_login(self.user)
+ version = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(version=version)
+ self.model_1.archived = datetime.now(timezone.utc)
+ self.model_1.save()
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This ModelVersion is part of an archived model."],
+ })
+
+ def test_partial_update_model_version(self):
+ """
+ Update the worker run by adding a model_version with a worker version that supports it
+ """
+ self.client.force_login(self.user)
+ version_with_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(
+ version=version_with_model,
+ parents=[],
+ )
+ self.assertIsNone(run.model_version_id)
+ self.assertEqual(run.summary, "Worker Recognizer @ version 2")
+
+ model_versions = [
+ # Version on a model with contributor access
+ self.model_version_1,
+ # Available version with tag on a model with guest access
+ self.model_version_2,
+ # Available version with tag on a public model
+ self.model_version_3,
+ ]
+ for model_version in model_versions:
+ with self.subTest(model_version=model_version):
+ with self.assertNumQueries(6):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(model_version.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ run.refresh_from_db()
+ self.assertEqual(response.json(), {
+ "id": str(run.id),
+ "configuration": None,
+ "model_version": {
+ "id": str(model_version.id),
+ "configuration": {},
+ "model": {
+ "id": str(model_version.model.id),
+ "name": model_version.model.name
+ },
+ "size": 8,
+ "state": "available",
+ "tag": model_version.tag,
+ },
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(version_with_model.id),
+ "configuration": {"test": "test2"},
+ "docker_image_iid": None,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Required.value,
+ "revision_url": None,
+ "version": version_with_model.version,
+ "tag": None,
+ "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
+ "state": "created",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}",
+ })
+ self.assertEqual(run.model_version_id, model_version.id)
+ self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}")
+
+ def test_partial_update_model_version_with_configuration(self):
+ """
+ Updating the worker run by adding a model_version with a worker version
+ doesn't erase previously loaded worker configuration
+ """
+ self.client.force_login(self.user)
+ version_with_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(
+ version=version_with_model,
+ parents=[],
+ configuration=self.configuration_1
+ )
+ self.assertEqual(run.model_version_id, None)
+
+ with self.assertNumQueries(6):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ run.refresh_from_db()
+ self.assertEqual(response.json(), {
+ "id": str(run.id),
+ "configuration": {
+ "archived": False,
+ "configuration": {"key": "value"},
+ "id": str(self.configuration_1.id),
+ "name": "My config"
+ },
+ "model_version": {
+ "id": str(self.model_version_1.id),
+ "configuration": {},
+ "model": {
+ "id": str(self.model_1.id),
+ "name": "My model"
+ },
+ "size": 8,
+ "state": "available",
+ "tag": None
+ },
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(version_with_model.id),
+ "configuration": {"test": "test2"},
+ "docker_image_iid": None,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Required.value,
+ "revision_url": None,
+ "version": version_with_model.version,
+ "tag": None,
+ "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
+ "state": "created",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'",
+ })
+ self.assertEqual(run.model_version_id, self.model_version_1.id)
+ self.assertEqual(run.configuration_id, self.configuration_1.id)
+
+ def test_partial_update(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(7):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.run_1.refresh_from_db()
+ self.assertDictEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [str(run_2.id)],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_partial_update_unique(self):
+ self.client.force_login(self.user)
+ self.version_1.model_usage = FeatureUsage.Required
+ self.version_1.save()
+ cases = [
+ (None, None),
+ (None, self.configuration_1),
+ (self.model_version_1, None),
+ (self.model_version_1, self.configuration_1),
+ ]
+ for model_version, configuration in cases:
+ with self.subTest(model_version=model_version, configuration=configuration):
+ # Erase any previous failures
+ self.process_1.worker_runs.exclude(id=self.run_1.id).delete()
+
+ self.run_1.model_version = model_version
+ self.run_1.configuration = configuration
+ self.run_1.save()
+
+ # Ensure the other run has different values before updating to avoid conflicts
+ run_2 = self.process_1.worker_runs.create(
+ version=self.version_1,
+ model_version=None if model_version else self.model_version_1,
+ configuration=None if configuration else self.configuration_1,
+ )
+
+ # Having a model version or a configuration adds one query for each
+ query_count = 4 + bool(model_version) + bool(configuration)
+
+ with self.assertNumQueries(query_count):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ # Update the second worker run to the first worker run's values to cause a conflict
+ "model_version_id": str(model_version.id) if model_version else None,
+ "configuration_id": str(configuration.id) if configuration else None,
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
+ })
+
+ def test_partial_update_agent(self):
+ """
+ Ponos agents cannot update WorkerRuns, even when they can access them
+ """
+ self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
+
+ # Agent auth is not implemented in CE
+ self.client.force_authenticate(user=self.agent)
+ with self.assertNumQueries(1):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+ self.assertEqual(response.json(), {
+ "detail": "You do not have an admin access to this process.",
+ })
+
+ def test_update_use_gpu_errors(self):
+ self.client.force_login(self.user)
+
+ cases = [
+ (self.version_gpu_disabled, True, "This worker version does not support GPU usage."),
+ (self.version_gpu_required, False, "This worker version requires GPU usage.")
+ ]
+
+ for worker_version, use_gpu, error_message in cases:
+ with self.subTest(worker_version=worker_version, use_gpu=use_gpu, error_message=error_message):
+ run = WorkerRun.objects.create(
+ process=self.process_1,
+ version=worker_version,
+ parents=[]
+ )
+ self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
+ with self.assertNumQueries(3):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={"use_gpu": use_gpu}
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+ self.assertDictEqual(response.json(), {"use_gpu": [error_message]})
+
+ run.refresh_from_db()
+ self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
+
+ def test_update_use_gpu(self):
+ self.client.force_login(self.user)
+
+ cases = [
+ (self.version_gpu_disabled, False),
+ (self.version_gpu_required, True),
+ (self.version_gpu_supported, False),
+ (self.version_gpu_supported, True)
+ ]
+
+ for worker_version, use_gpu in cases:
+ with self.subTest(worker_version=worker_version, use_gpu=use_gpu):
+ self.process_1.worker_runs.all().delete()
+ run = WorkerRun.objects.create(
+ process=self.process_1,
+ version=worker_version,
+ parents=[]
+ )
+ self.assertEqual(run.use_gpu, True if worker_version.gpu_usage == FeatureUsage.Required else False)
+
+ with self.assertNumQueries(5):
+ response = self.client.patch(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={"use_gpu": use_gpu}
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ data = response.json()
+ self.assertDictEqual(data, {
+ "id": str(run.id),
+ "worker_version": {
+ "id": str(worker_version.id),
+ "configuration": worker_version.configuration,
+ "docker_image_iid": "registry.nerv.co.jp/neon-genesis:evangelion",
+ "gpu_usage": worker_version.gpu_usage.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": worker_version.version,
+ "tag": None,
+ "created": worker_version.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ },
+ "summary": f"Worker Recognizer @ version {worker_version.version}",
+ "use_gpu": use_gpu
+ })
+ run.refresh_from_db()
+ self.assertEqual(run.use_gpu, use_gpu)
diff --git a/arkindex/process/tests/worker_runs/test_retrieve.py b/arkindex/process/tests/worker_runs/test_retrieve.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c2471ee10aa682c89da96dbb4e429af0d564047
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_retrieve.py
@@ -0,0 +1,405 @@
+from datetime import datetime, timezone
+
+from django.urls import reverse
+from rest_framework import status
+
+from arkindex.ponos.models import Agent, Farm
+from arkindex.process.models import (
+ FeatureUsage,
+ ProcessMode,
+ Worker,
+ WorkerRun,
+ WorkerVersion,
+)
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.users.models import Role
+
+
+class TestWorkerRunsretrieve(FixtureAPITestCase):
+ """
+ Test worker runs retrieve endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+
+ cls.worker_custom = Worker.objects.get(slug="custom")
+ cls.version_custom = cls.worker_custom.versions.get()
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+ cls.run_custom = cls.local_process.worker_runs.get(version=cls.version_custom)
+ cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
+
+ cls.agent = Agent.objects.create(
+ farm=cls.farm,
+ hostname="claude",
+ cpu_cores=42,
+ cpu_frequency=1e15,
+ ram_total=99e9,
+ last_ping=datetime.now(timezone.utc),
+ )
+ # Add custom attributes to make the agent usable as an authenticated user
+ cls.agent.is_agent = True
+ cls.agent.is_anonymous = False
+
+ def test_retrieve_requires_login(self):
+ with self.assertNumQueries(0):
+ response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}))
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ def test_retrieve_no_worker_execution_right(self):
+ """
+ A user can retrieve any worker run if they have an admin access to the process project
+ """
+ self.worker_1.memberships.update(level=Role.Guest.value)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(4):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)})
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertDictEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_retrieve_invalid_id(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"})
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_retrieve(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(4):
+ response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}))
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertDictEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_retrieve_custom(self):
+ self.client.force_login(self.user)
+ # Update process attributes to ensure they do not cause extra requests
+ page = self.corpus.elements.get(name="Volume 1, page 1r")
+ self.local_process.element = page
+ self.local_process.element_type = page.type
+ self.local_process.folder_type = page.type
+ self.local_process.prefix = "a"
+ self.local_process.save()
+ with self.assertNumQueries(5):
+ response = self.client.get(reverse("api:worker-run-details", kwargs={"pk": str(self.run_custom.id)}))
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertDictEqual(response.json(), {
+ "id": str(self.run_custom.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.local_process.id),
+ "activity_state": "disabled",
+ "corpus": None,
+ "chunks": 1,
+ "mode": "local",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": {
+ "id": str(page.id),
+ "type": "page",
+ "corpus": {
+ "id": str(self.corpus.id),
+ "name": "Unit Tests",
+ "public": True
+ },
+ "name": "Volume 1, page 1r",
+ "rotation_angle": 0,
+ "mirrored": False,
+ "thumbnail_url": None,
+ "zone": {
+ "id": str(page.id),
+ "image": {
+ "id": str(page.image.id),
+ "height": 1000,
+ "path": "img1",
+ "s3_url": None,
+ "server": {
+ "display_name": "Test Server",
+ "max_height": None,
+ "max_width": None,
+ "url": "http://server"
+ },
+ "status": "unchecked",
+ "url": "http://server/img1",
+ "width": 1000
+ },
+ "polygon": [[int(x), int(y)] for x, y in page.polygon],
+ "url": "http://server/img1/0,0,1000,1000/full/0/default.jpg",
+ },
+ },
+ "element_type": "page",
+ "folder_type": "page",
+ "prefix": "a",
+ },
+ "worker_version": {
+ "id": str(self.version_custom.id),
+ "configuration": {"custom": "value"},
+ "docker_image_iid": None,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "created",
+ "worker": {
+ "id": str(self.worker_custom.id),
+ "name": "Custom worker",
+ "slug": "custom",
+ "type": "custom",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "summary": "Worker Custom worker @ version 1",
+ "use_gpu": False,
+ })
+
+ def test_retrieve_local(self):
+ """
+ A user can retrieve a run on their own local process
+ """
+ run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertDictEqual(response.json(), {
+ "id": str(run.id),
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.local_process.id),
+ "activity_state": "disabled",
+ "corpus": None,
+ "chunks": 1,
+ "mode": "local",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_retrieve_local_only_current_user(self):
+ """
+ A user cannot retrieve a run on another user's local process
+ """
+ run = WorkerRun.objects.filter(process__creator=self.superuser, process__mode=ProcessMode.Local).first()
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_retrieve_agent(self):
+ """
+ A Ponos agent can retrieve a WorkerRun on a process where it has some assigned tasks
+ """
+ self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
+
+ # Agent auth is not implemented in CE
+ self.client.force_authenticate(user=self.agent)
+ with self.assertNumQueries(3):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "parents": [],
+ "model_version": None,
+ "configuration": None,
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "mode": "workers",
+ "chunks": 1,
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_retrieve_agent_unassigned(self):
+ """
+ A Ponos agent cannot retrieve a WorkerRun on a process where it does not have any assigned tasks
+ """
+ self.process_1.tasks.create(run=0, depth=0, slug="something", agent=None)
+
+ # Agent auth is not implemented in CE
+ self.client.force_authenticate(user=self.agent)
+ with self.assertNumQueries(1):
+ response = self.client.get(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
diff --git a/arkindex/process/tests/worker_runs/test_update.py b/arkindex/process/tests/worker_runs/test_update.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4e38afa695846dd92b96d5dd732c26dbbbb9a04
--- /dev/null
+++ b/arkindex/process/tests/worker_runs/test_update.py
@@ -0,0 +1,999 @@
+import uuid
+from datetime import datetime, timezone
+from unittest.mock import call, patch
+
+from django.urls import reverse
+from rest_framework import status
+from rest_framework.exceptions import ValidationError
+
+from arkindex.ponos.models import Agent, Farm
+from arkindex.process.models import FeatureUsage, ProcessMode, WorkerVersion
+from arkindex.project.tests import FixtureAPITestCase
+from arkindex.training.models import Model, ModelVersion, ModelVersionState
+from arkindex.users.models import Role
+
+
+class TestWorkerRunsUpdate(FixtureAPITestCase):
+ """
+ Test worker runs update endpoint
+ """
+
+ @classmethod
+ def setUpTestData(cls):
+ super().setUpTestData()
+ cls.local_process = cls.user.processes.get(mode=ProcessMode.Local)
+ cls.farm = Farm.objects.first()
+ cls.process_1 = cls.corpus.processes.create(
+ creator=cls.user,
+ mode=ProcessMode.Workers,
+ farm=cls.farm,
+ )
+ cls.version_1 = WorkerVersion.objects.get(worker__slug="reco")
+ cls.worker_1 = cls.version_1.worker
+ cls.version_2 = WorkerVersion.objects.get(worker__slug="dla")
+ cls.run_1 = cls.process_1.worker_runs.create(version=cls.version_1, parents=[])
+ cls.configuration_1 = cls.worker_1.configurations.create(name="My config", configuration={"key": "value"})
+ worker_version = WorkerVersion.objects.exclude(worker=cls.version_1.worker).first()
+ cls.configuration_2 = worker_version.worker.configurations.create(name="Config")
+ cls.process_2 = cls.corpus.processes.create(creator=cls.user, mode=ProcessMode.Workers)
+
+ # Model and Model version setup
+ cls.model_1 = Model.objects.create(name="My model")
+ cls.model_1.memberships.create(user=cls.user, level=Role.Contributor.value)
+ cls.model_version_1 = ModelVersion.objects.create(
+ model=cls.model_1,
+ state=ModelVersionState.Available,
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.model_2 = Model.objects.create(name="Their model")
+ cls.model_2.memberships.create(user=cls.user, level=Role.Guest.value)
+ cls.model_version_2 = cls.model_2.versions.create(
+ state=ModelVersionState.Available,
+ tag="blah",
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.model_3 = Model.objects.create(name="Our model", public=True)
+ cls.model_version_3 = cls.model_3.versions.create(
+ state=ModelVersionState.Available,
+ tag="blah",
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ cls.agent = Agent.objects.create(
+ farm=cls.farm,
+ hostname="claude",
+ cpu_cores=42,
+ cpu_frequency=1e15,
+ ram_total=99e9,
+ last_ping=datetime.now(timezone.utc),
+ )
+ # Add custom attributes to make the agent usable as an authenticated user
+ cls.agent.is_agent = True
+ cls.agent.is_anonymous = False
+
+ def test_update_requires_nothing(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={},
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ def test_update_requires_login(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+
+ with self.assertNumQueries(0):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ @patch("arkindex.project.mixins.get_max_level", return_value=Role.Contributor.value)
+ def test_update_no_project_admin_right(self, get_max_level_mock):
+ """
+ A user cannot update a worker run if they have no admin access on its process project
+ """
+ self.corpus.memberships.filter(user=self.user).update(level=Role.Contributor.value)
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": []
+ }
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+ self.assertEqual(response.json(), {"detail": "You do not have an admin access to this process."})
+
+ self.assertEqual(get_max_level_mock.call_count, 1)
+ self.assertEqual(get_max_level_mock.call_args, call(self.user, self.corpus))
+
+ def test_update_invalid_process_mode(self):
+ """
+ A user cannot update a worker run on a local process
+ """
+ run = self.local_process.worker_runs.create(version=self.version_1, parents=[])
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "parents": []
+ }
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns can only be created or updated on Workers or Dataset processes."],
+ })
+
+ def test_update_invalid_id(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(3):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": "12341234-1234-1234-1234-123412341234"}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
+ def test_update_nonexistent_parent(self):
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": ["12341234-1234-1234-1234-123412341234"],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), [
+ f"Can't add or update WorkerRun {self.run_1.id} because parents field isn't properly defined. It can be either because"
+ " one or several UUIDs don't refer to existing WorkerRuns or either because listed WorkerRuns doesn't belong to the"
+ " same Process than this WorkerRun."
+ ])
+
+ def test_update_duplicate_parents(self):
+ self.client.force_login(self.user)
+ run_2 = self.process_1.worker_runs.create(version=self.version_2)
+
+ with self.assertNumQueries(4):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [
+ str(run_2.id),
+ str(run_2.id),
+ ],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "parents": ["The parents of a WorkerRun must be unique."],
+ })
+
+ def test_duplicate_parents_signal(self):
+ """
+ Duplicates in WorkerRun.parents are also detected outside of the WorkerRun APIs
+ """
+ run_2 = self.process_1.worker_runs.create(
+ version=self.version_2,
+ parents=[],
+ )
+
+ run_2.parents = [self.run_1.id, self.run_1.id]
+ with self.assertRaisesRegex(ValidationError, f"Can't add or update WorkerRun {run_2.id} because it has duplicate parents."):
+ run_2.parents = [self.run_1.id, self.run_1.id]
+ run_2.save()
+
+ def test_update_process_id(self):
+ """
+ Process field cannot be updated
+ """
+ self.client.force_login(self.user)
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "process_id": str(self.process_2.id),
+ "parents": [],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+ self.run_1.refresh_from_db()
+ self.assertEqual(self.run_1.process.id, self.process_1.id)
+
+ def test_update_worker_version_id(self):
+ """
+ Version field cannot be updated
+ """
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "worker_version_id": str(self.version_2.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ },
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+ self.run_1.refresh_from_db()
+ self.assertNotEqual(self.run_1.version_id, self.version_2.id)
+
+ def test_update_configuration(self):
+ self.client.force_login(self.user)
+ self.assertEqual(self.run_1.configuration, None)
+ # Check generated summary, before updating, it should not be that verbose
+ self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1")
+
+ with self.assertNumQueries(6):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
+ data={
+ "parents": [],
+ "configuration_id": str(self.configuration_1.id)
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.run_1.refresh_from_db()
+
+ self.assertEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": {
+ "id": str(self.configuration_1.id),
+ "archived": False,
+ "configuration": {"key": "value"},
+ "name": "My config"
+ },
+ "model_version": None,
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1 using configuration 'My config'",
+ })
+ self.assertEqual(self.run_1.configuration.id, self.configuration_1.id)
+ # Check generated summary, after the update, the configuration should be displayed as well
+ self.assertEqual(self.run_1.summary, "Worker Recognizer @ version 1 using configuration 'My config'")
+
+ def test_update_invalid_configuration(self):
+ self.client.force_login(self.user)
+ self.assertEqual(self.run_1.configuration, None)
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": self.run_1.id}),
+ data={"parents": [], "configuration_id": str(self.configuration_2.id)},
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertDictEqual(response.json(), {"configuration_id": ["The configuration must be part of the same worker."]})
+
+ def test_update_process_already_started(self):
+ """
+ Update dependencies of a worker run is not possible once the process is started
+ """
+ self.process_1.run()
+ self.assertTrue(self.process_1.tasks.exists())
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(4):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "process_id": ["WorkerRuns cannot be added or updated on processes that have already started."],
+ })
+
+ def test_update_model_version_not_allowed(self):
+ """
+ The model_version UUID is not allowed when the related version doesn't allow model_usage
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Disabled
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This worker version does not support models."]
+ })
+
+ def test_update_unknown_model_version(self):
+ """
+ Cannot use a model version id that doesn't exist
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+ random_model_version_uuid = str(uuid.uuid4())
+
+ with self.assertNumQueries(4):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": random_model_version_uuid,
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{random_model_version_uuid}" - object does not exist.']
+ })
+
+ @patch("arkindex.users.managers.BaseACLManager.filter_rights", return_value=ModelVersion.objects.none())
+ def test_update_model_version_no_access(self, filter_rights_mock):
+ """
+ Cannot update a worker run with a model_version UUID, when you don't have access to the model version
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ # Create a model version, the user has no access to
+ model_no_access = Model.objects.create(name="Secret model")
+ model_version_no_access = ModelVersion.objects.create(
+ model=model_no_access,
+ state=ModelVersionState.Available,
+ size=8,
+ hash="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+ archive_hash="bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
+ )
+
+ with self.assertNumQueries(3):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(model_version_no_access.id),
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{model_version_no_access.id}" - object does not exist.'],
+ })
+
+ self.assertListEqual(filter_rights_mock.call_args_list, [
+ call(self.user, Model, Role.Guest.value),
+ call(self.user, Model, Role.Contributor.value),
+ ])
+
+ @patch("arkindex.users.managers.BaseACLManager.filter_rights")
+ def test_update_model_version_guest(self, filter_rights_mock):
+ """
+ Cannot update a worker run with a model_version when you only have guest access to the model,
+ and the model version has no tag or is not available
+ """
+ self.client.force_login(self.user)
+ version_no_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_no_model,
+ parents=[],
+ )
+
+ def filter_rights(user, model, level):
+ """
+ The filter_rights mock needs to return nothing when called for contributor access,
+ and the models we will test on when called for guest access
+ """
+ if level == Role.Guest.value:
+ return Model.objects.filter(id__in=(self.model_2.id, self.model_3.id))
+ return Model.objects.none()
+
+ filter_rights_mock.side_effect = filter_rights
+
+ cases = [
+ # On a model with a membership giving guest access
+ (self.model_version_2, None, ModelVersionState.Created),
+ (self.model_version_2, None, ModelVersionState.Error),
+ (self.model_version_2, None, ModelVersionState.Available),
+ (self.model_version_2, "blah", ModelVersionState.Created),
+ (self.model_version_2, "blah", ModelVersionState.Error),
+ # On a public model with no membership
+ (self.model_version_3, None, ModelVersionState.Created),
+ (self.model_version_3, None, ModelVersionState.Error),
+ (self.model_version_3, None, ModelVersionState.Available),
+ (self.model_version_3, "blah", ModelVersionState.Created),
+ (self.model_version_3, "blah", ModelVersionState.Error),
+ ]
+
+ for model_version, tag, state in cases:
+ filter_rights_mock.reset_mock()
+ with self.subTest(model_version=model_version, tag=tag, state=state):
+ model_version.tag = tag
+ model_version.state = state
+ model_version.save()
+
+ with self.assertNumQueries(4):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ "model_version_id": str(model_version.id),
+ "parents": [],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": [f'Invalid pk "{model_version.id}" - object does not exist.'],
+ })
+
+ self.assertListEqual(filter_rights_mock.call_args_list, [
+ call(self.user, Model, Role.Guest.value),
+ call(self.user, Model, Role.Contributor.value),
+ ])
+
+ def test_update_model_version_unavailable(self):
+ self.client.force_login(self.user)
+ version = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(
+ version=version,
+ parents=[],
+ )
+ self.model_version_1.state = ModelVersionState.Error
+ self.model_version_1.save()
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This ModelVersion is not in an Available state."]
+ })
+
+ def test_update_model_archived(self):
+ self.client.force_login(self.user)
+ version = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(version=version)
+ self.model_1.archived = datetime.now(timezone.utc)
+ self.model_1.save()
+
+ with self.assertNumQueries(5):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "model_version_id": ["This ModelVersion is part of an archived model."],
+ })
+
+ def test_update_model_version_id(self):
+ """
+ Update the worker run by adding a model_version with a worker version that supports it
+ """
+ self.client.force_login(self.user)
+ version_with_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Supported
+ )
+ run = self.process_1.worker_runs.create(
+ version=version_with_model,
+ parents=[],
+ )
+ self.assertEqual(run.model_version, None)
+ # Check generated summary, before updating, there should be only information about the worker version
+ self.assertEqual(run.summary, "Worker Recognizer @ version 2")
+
+ model_versions = [
+ # Version on a model with contributor access
+ self.model_version_1,
+ # Available version with tag on a model with guest access
+ self.model_version_2,
+ # Available version with tag on a public model
+ self.model_version_3,
+ ]
+ for model_version in model_versions:
+ with self.subTest(model_version=model_version):
+ with self.assertNumQueries(6):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(model_version.id),
+ "parents": [],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ run.refresh_from_db()
+ self.assertEqual(response.json(), {
+ "id": str(run.id),
+ "configuration": None,
+ "model_version": {
+ "id": str(model_version.id),
+ "configuration": {},
+ "model": {
+ "id": str(model_version.model.id),
+ "name": model_version.model.name
+ },
+ "size": 8,
+ "state": "available",
+ "tag": model_version.tag,
+ },
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(version_with_model.id),
+ "configuration": {"test": "test2"},
+ "docker_image_iid": None,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Supported.value,
+ "revision_url": None,
+ "version": version_with_model.version,
+ "tag": None,
+ "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
+ "state": "created",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}",
+ })
+ self.assertEqual(run.model_version_id, model_version.id)
+ self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {model_version.model.name} @ {str(model_version.id)[:6]}")
+
+ def test_update_configuration_and_model_version(self):
+ """
+ Update the worker run by adding both a model_version and a worker configuration
+ """
+ self.client.force_login(self.user)
+ version_with_model = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"},
+ model_usage=FeatureUsage.Required
+ )
+ run = self.process_1.worker_runs.create(
+ version=version_with_model,
+ parents=[],
+ )
+ self.assertIsNone(run.model_version)
+ self.assertIsNone(run.configuration)
+ # Check generated summary, before updating, there should be only information about the worker version
+ self.assertEqual(run.summary, "Worker Recognizer @ version 2")
+
+ with self.assertNumQueries(7):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run.id)}),
+ data={
+ "model_version_id": str(self.model_version_1.id),
+ "configuration_id": str(self.configuration_1.id),
+ "parents": []
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ run.refresh_from_db()
+ self.assertEqual(response.json(), {
+ "id": str(run.id),
+ "configuration": {
+ "id": str(self.configuration_1.id),
+ "archived": False,
+ "configuration": {"key": "value"},
+ "name": "My config"
+ },
+ "model_version": {
+ "id": str(self.model_version_1.id),
+ "configuration": {},
+ "model": {
+ "id": str(self.model_1.id),
+ "name": "My model"
+ },
+ "size": 8,
+ "state": "available",
+ "tag": None
+ },
+ "parents": [],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(version_with_model.id),
+ "configuration": {"test": "test2"},
+ "docker_image_iid": None,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Required.value,
+ "revision_url": None,
+ "version": version_with_model.version,
+ "tag": None,
+ "created": version_with_model.created.isoformat().replace("+00:00", "Z"),
+ "state": "created",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": f"Worker Recognizer @ version 2 with model My model @ {str(self.model_version_1.id)[:6]} using configuration 'My config'",
+ })
+ self.assertEqual(run.model_version_id, self.model_version_1.id)
+ # Check generated summary, after updating, there should be information about the model loaded
+ self.assertEqual(run.summary, f"Worker Recognizer @ version 2 with model {self.model_version_1.model.name} @ {str(self.model_version_1.id)[:6]} using configuration '{self.configuration_1.name}'")
+
+ def test_update(self):
+ version_2 = WorkerVersion.objects.create(
+ worker=self.worker_1,
+ version=2,
+ configuration={"test": "test2"}
+ )
+ run_2 = self.process_1.worker_runs.create(
+ version=version_2,
+ parents=[],
+ )
+ self.client.force_login(self.user)
+
+ with self.assertNumQueries(7):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ data={
+ "parents": [str(run_2.id)],
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_200_OK)
+
+ self.run_1.refresh_from_db()
+ self.assertDictEqual(response.json(), {
+ "id": str(self.run_1.id),
+ "configuration": None,
+ "model_version": None,
+ "parents": [str(run_2.id)],
+ "process": {
+ "id": str(self.process_1.id),
+ "activity_state": "disabled",
+ "corpus": str(self.corpus.id),
+ "chunks": 1,
+ "mode": "workers",
+ "name": None,
+ "state": "unscheduled",
+ "use_cache": False,
+ "element": None,
+ "element_type": None,
+ "folder_type": None,
+ "prefix": None,
+ },
+ "worker_version": {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "gpu_usage": "disabled",
+ "model_usage": FeatureUsage.Disabled.value,
+ "revision_url": None,
+ "version": 1,
+ "tag": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ "state": "available",
+ "worker": {
+ "id": str(self.worker_1.id),
+ "name": "Recognizer",
+ "slug": "reco",
+ "type": "recognizer",
+ "description": "",
+ "repository_url": None,
+ "archived": False,
+ }
+ },
+ "use_gpu": False,
+ "summary": "Worker Recognizer @ version 1",
+ })
+
+ def test_update_unique(self):
+ self.client.force_login(self.user)
+ self.version_1.model_usage = FeatureUsage.Required
+ self.version_1.save()
+ cases = [
+ (None, None),
+ (None, self.configuration_1),
+ (self.model_version_1, None),
+ (self.model_version_1, self.configuration_1),
+ ]
+ for model_version, configuration in cases:
+ with self.subTest(model_version=model_version, configuration=configuration):
+ # Erase any previous failures
+ self.process_1.worker_runs.exclude(id=self.run_1.id).delete()
+
+ self.run_1.model_version = model_version
+ self.run_1.configuration = configuration
+ self.run_1.save()
+
+ # Ensure the other run has different values before updating to avoid conflicts
+ run_2 = self.process_1.worker_runs.create(
+ version=self.version_1,
+ model_version=None if model_version else self.model_version_1,
+ configuration=None if configuration else self.configuration_1,
+ )
+
+ # Having a model version or a configuration adds one query for each
+ query_count = 4 + bool(model_version) + bool(configuration)
+
+ with self.assertNumQueries(query_count):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(run_2.id)}),
+ data={
+ # Update the second worker run to the first worker run's values to cause a conflict
+ "model_version_id": str(model_version.id) if model_version else None,
+ "configuration_id": str(configuration.id) if configuration else None,
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+ self.assertEqual(response.json(), {
+ "__all__": ["A WorkerRun already exists on this process with the selected worker version, model version and configuration."],
+ })
+
+ def test_update_agent(self):
+ """
+ Ponos agents cannot update WorkerRuns, even when they can access them
+ """
+ self.process_1.tasks.create(run=0, depth=0, slug="something", agent=self.agent)
+
+ # Agent auth is not implemented in CE
+ self.client.force_authenticate(user=self.agent)
+ with self.assertNumQueries(1):
+ response = self.client.put(
+ reverse("api:worker-run-details", kwargs={"pk": str(self.run_1.id)}),
+ )
+ self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+ self.assertEqual(response.json(), {
+ "detail": "You do not have an admin access to this process.",
+ })