diff --git a/arkindex/process/serializers/workers.py b/arkindex/process/serializers/workers.py index bd2b5dd0dc6dbd8beff5ef267572aecd4e4af06c..4561ea6f7a8ee62120d96ab09474247fefd992a9 100644 --- a/arkindex/process/serializers/workers.py +++ b/arkindex/process/serializers/workers.py @@ -29,6 +29,8 @@ from arkindex.users.models import Role from docker.auth import resolve_repository_name from docker.errors import InvalidRepository from docker.utils import parse_repository_tag +from packaging.version import InvalidVersion +from packaging.version import parse as parse_version class ExistingWorkerException(Exception): @@ -371,6 +373,17 @@ class WorkerVersionCreateSerializer(WorkerVersionSerializer): last_version = worker.versions.using("default").aggregate(last_version=Max("version"))["last_version"] data["version"] = last_version + 1 if last_version else 1 + # Try to autodetect the tag from the docker image + if "tag" not in data and "docker_image_iid" in data: + _, tag = parse_repository_tag(data["docker_image_iid"]) + if tag is not None: + try: + parse_version(tag) + except InvalidVersion: + pass + else: + data["tag"] = tag + data["worker_id"] = worker.id return data diff --git a/arkindex/process/tests/worker_versions/test_create.py b/arkindex/process/tests/worker_versions/test_create.py index 054907a1bd353c7e738ecfdf718511a22914dd2b..eac2805562506fdeb27f6773bc62d085e5743418 100644 --- a/arkindex/process/tests/worker_versions/test_create.py +++ b/arkindex/process/tests/worker_versions/test_create.py @@ -660,3 +660,90 @@ class TestWorkerVersionCreate(FixtureAPITestCase): }, } }) + + def test_docker_image_iid(self): + with self.assertNumQueries(6): + response = self.client.post( + reverse("api:worker-versions", kwargs={"pk": str(self.worker_reco.id)}), + data={ + "docker_image_iid": "registry.dev.test/workers/test:0001.0.0rc1", + "revision_url": "https://dev.test/some/commit/aaaa", + "configuration": {"test": 42}, + "state": "available", + }, + HTTP_AUTHORIZATION=f"Ponos {self.task.token}", + ) + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + + data = response.json() + worker_version = WorkerVersion.objects.get(id=data["id"]) + self.assertDictEqual(data, { + "branch": None, + "configuration": {"test": 42}, + "created": worker_version.created.isoformat().replace("+00:00", "Z"), + "docker_image_iid": "registry.dev.test/workers/test:0001.0.0rc1", + "gpu_usage": "disabled", + "id": str(worker_version.id), + "model_usage": "disabled", + "revision_url": "https://dev.test/some/commit/aaaa", + "state": "available", + # Tag is automatically detected from the docker image + "tag": "0001.0.0rc1", + "version": None, + "worker": { + "archived": False, + "cost_1k_elements": "0.250", + "cost_cpu_hour": "1.312", + "cost_gpu_hour": "0.000", + "description": "", + "id": str(worker_version.worker_id), + "name": "Recognizer", + "repository_url": None, + "slug": "reco", + "type": "recognizer", + } + }) + + def test_docker_image_iid_tag_override(self): + with self.assertNumQueries(6): + response = self.client.post( + reverse("api:worker-versions", kwargs={"pk": str(self.worker_reco.id)}), + data={ + "docker_image_iid": "registry.dev.test/workers/test:0001.0.0rc1", + "tag": None, + "revision_url": "https://dev.test/some/commit/aaaa", + "configuration": {"test": 42}, + "state": "available", + }, + HTTP_AUTHORIZATION=f"Ponos {self.task.token}", + ) + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + + data = response.json() + worker_version = WorkerVersion.objects.get(id=data["id"]) + self.assertDictEqual(data, { + "branch": None, + "configuration": {"test": 42}, + "created": worker_version.created.isoformat().replace("+00:00", "Z"), + "docker_image_iid": "registry.dev.test/workers/test:0001.0.0rc1", + "gpu_usage": "disabled", + "id": str(worker_version.id), + "model_usage": "disabled", + "revision_url": "https://dev.test/some/commit/aaaa", + "state": "available", + # Tag is not automatically detected from the docker image + "tag": None, + "version": None, + "worker": { + "archived": False, + "cost_1k_elements": "0.250", + "cost_cpu_hour": "1.312", + "cost_gpu_hour": "0.000", + "description": "", + "id": str(worker_version.worker_id), + "name": "Recognizer", + "repository_url": None, + "slug": "reco", + "type": "recognizer", + } + })