diff --git a/arkindex/documents/fixtures/data.json b/arkindex/documents/fixtures/data.json
index 0ef2f121170e2db74c72bc3e874e9b6b114c49ae..efd6ebcd8817dca400f85166fbe46631cc9d803d 100644
--- a/arkindex/documents/fixtures/data.json
+++ b/arkindex/documents/fixtures/data.json
@@ -309,7 +309,7 @@
},
"state": "available",
"gpu_usage": "required",
- "model_usage": false,
+ "model_usage": "disabled",
"docker_image": "a93301c3-bdac-4d06-ba22-b7066decd282",
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
@@ -328,7 +328,7 @@
},
"state": "created",
"gpu_usage": "disabled",
- "model_usage": false,
+ "model_usage": "disabled",
"docker_image": null,
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
@@ -345,7 +345,7 @@
"configuration": {},
"state": "available",
"gpu_usage": "disabled",
- "model_usage": false,
+ "model_usage": "disabled",
"docker_image": "a93301c3-bdac-4d06-ba22-b7066decd282",
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
@@ -364,7 +364,7 @@
},
"state": "available",
"gpu_usage": "disabled",
- "model_usage": false,
+ "model_usage": "disabled",
"docker_image": "a93301c3-bdac-4d06-ba22-b7066decd282",
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
@@ -383,7 +383,7 @@
},
"state": "available",
"gpu_usage": "disabled",
- "model_usage": true,
+ "model_usage": "required",
"docker_image": "a93301c3-bdac-4d06-ba22-b7066decd282",
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
@@ -402,7 +402,7 @@
},
"state": "available",
"gpu_usage": "disabled",
- "model_usage": false,
+ "model_usage": "disabled",
"docker_image": "a93301c3-bdac-4d06-ba22-b7066decd282",
"docker_image_iid": null,
"created": "2020-02-02T01:23:45.678Z",
diff --git a/arkindex/documents/management/commands/bootstrap.py b/arkindex/documents/management/commands/bootstrap.py
index 40235e163826c475f3cc8ae770e88c01a4c96545..23dda2d5e2b8e70cd10a90720ff506411d4b290f 100644
--- a/arkindex/documents/management/commands/bootstrap.py
+++ b/arkindex/documents/management/commands/bootstrap.py
@@ -10,14 +10,7 @@ from rest_framework.authtoken.models import Token
from arkindex.images.models import ImageServer
from arkindex.ponos.models import Farm
-from arkindex.process.models import (
- Repository,
- Worker,
- WorkerType,
- WorkerVersion,
- WorkerVersionGPUUsage,
- WorkerVersionState,
-)
+from arkindex.process.models import FeatureUsage, Repository, Worker, WorkerType, WorkerVersion, WorkerVersionState
from arkindex.users.models import User
# Constants used in architecture project
@@ -210,7 +203,7 @@ class Command(BaseCommand):
'id': IMPORT_WORKER_VERSION_ID,
'configuration': {},
'state': WorkerVersionState.Created,
- 'gpu_usage': WorkerVersionGPUUsage.Disabled,
+ 'gpu_usage': FeatureUsage.Disabled,
'docker_image': None,
'docker_image_iid': None,
}
@@ -243,7 +236,7 @@ class Command(BaseCommand):
# Ensure it has the right attributes
version.configuration = {}
version.state = WorkerVersionState.Created
- version.gpu_usage = WorkerVersionGPUUsage.Disabled
+ version.gpu_usage = FeatureUsage.Disabled
version.docker_image = None
version.docker_image_iid = None
version.save()
diff --git a/arkindex/documents/management/commands/build_fixtures.py b/arkindex/documents/management/commands/build_fixtures.py
index 6727f672d7c447eaa640c92a8edbc5843ea5ce3d..e7775bea98f0aeb1e8e326bd7dcf7fb9e7771aa7 100644
--- a/arkindex/documents/management/commands/build_fixtures.py
+++ b/arkindex/documents/management/commands/build_fixtures.py
@@ -10,13 +10,13 @@ from arkindex.documents.models import Corpus, Element, MetaData, MetaType
from arkindex.images.models import Image, ImageServer
from arkindex.ponos.models import Farm, State
from arkindex.process.models import (
+ FeatureUsage,
Process,
ProcessMode,
Worker,
WorkerRun,
WorkerType,
WorkerVersion,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.project.tools import fake_now
@@ -140,7 +140,7 @@ class Command(BaseCommand):
revision=revision,
configuration={'test': 42},
state=WorkerVersionState.Available,
- model_usage=False,
+ model_usage=FeatureUsage.Disabled,
docker_image=docker_image
)
dla_worker = WorkerVersion.objects.create(
@@ -152,7 +152,7 @@ class Command(BaseCommand):
revision=revision,
configuration={'test': 42},
state=WorkerVersionState.Available,
- model_usage=False,
+ model_usage=FeatureUsage.Disabled,
docker_image=docker_image
)
@@ -165,7 +165,7 @@ class Command(BaseCommand):
revision=revision,
configuration={},
state=WorkerVersionState.Available,
- model_usage=False,
+ model_usage=FeatureUsage.Disabled,
docker_image=docker_image,
)
@@ -178,9 +178,9 @@ class Command(BaseCommand):
revision=revision,
configuration={'test': 42},
state=WorkerVersionState.Available,
- model_usage=False,
+ model_usage=FeatureUsage.Disabled,
docker_image=docker_image,
- gpu_usage=WorkerVersionGPUUsage.Required
+ gpu_usage=FeatureUsage.Required
)
# Create a generic worker and its version that uses a ML Model
@@ -193,8 +193,8 @@ class Command(BaseCommand):
revision=revision,
configuration={'test': 42},
state=WorkerVersionState.Available,
- gpu_usage=WorkerVersionGPUUsage.Disabled,
- model_usage=True,
+ gpu_usage=FeatureUsage.Disabled,
+ model_usage=FeatureUsage.Required,
docker_image=docker_image
)
diff --git a/arkindex/process/builder.py b/arkindex/process/builder.py
index 76458142ce7590f0e6909debe990c8989110a28a..4d882ae4f8a346bcd091aae7d347a2a33eb78e7e 100644
--- a/arkindex/process/builder.py
+++ b/arkindex/process/builder.py
@@ -162,11 +162,11 @@ class ProcessBuilder(object):
@prefetch_worker_runs
def validate_gpu_requirement(self):
- from arkindex.process.models import WorkerVersionGPUUsage
+ from arkindex.process.models import FeatureUsage
if (
not self.process.use_gpu
- and any(run.version.gpu_usage == WorkerVersionGPUUsage.Required for run in self.process.worker_runs.all())
+ and any(run.version.gpu_usage == FeatureUsage.Required for run in self.process.worker_runs.all())
):
raise ValidationError("Some worker versions require a GPU and the `use_gpu` option is disabled.")
diff --git a/arkindex/process/migrations/0001_initial.py b/arkindex/process/migrations/0001_initial.py
index 43fbdce9a47026fd4593729ce0504bf8f46c564e..4a0b2688941211f69672706a0feff07b791a4d39 100644
--- a/arkindex/process/migrations/0001_initial.py
+++ b/arkindex/process/migrations/0001_initial.py
@@ -151,7 +151,7 @@ class Migration(migrations.Migration):
('id', models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)),
('configuration', models.JSONField()),
('state', enumfields.fields.EnumField(default='created', enum=arkindex.process.models.WorkerVersionState, max_length=10)),
- ('gpu_usage', enumfields.fields.EnumField(blank=True, default='disabled', enum=arkindex.process.models.WorkerVersionGPUUsage, max_length=10)),
+ ('gpu_usage', enumfields.fields.EnumField(blank=True, default='disabled', enum=arkindex.process.models.FeatureUsage, max_length=10)),
('model_usage', models.BooleanField(default=False)),
('docker_image_iid', models.CharField(blank=True, max_length=80, null=True)),
('corpora', models.ManyToManyField(related_name='worker_versions', through='process.CorpusWorkerVersion', to='documents.corpus')),
diff --git a/arkindex/process/migrations/0023_alter_workerversion_model_usage.py b/arkindex/process/migrations/0023_alter_workerversion_model_usage.py
new file mode 100644
index 0000000000000000000000000000000000000000..adead68c2bdd5d46560495607599b42cebe4d0a6
--- /dev/null
+++ b/arkindex/process/migrations/0023_alter_workerversion_model_usage.py
@@ -0,0 +1,47 @@
+# Generated by Django 4.1.7 on 2023-11-23 10:17
+
+import enumfields.fields
+from django.db import migrations
+
+import arkindex.process.models
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ('process', '0022_worker_description'),
+ ]
+
+ operations = [
+ migrations.RunSQL(
+ [
+ # This uses a `USING` expression to tell PostgreSQL how to perform the data migration
+ """
+ ALTER TABLE process_workerversion
+ ALTER COLUMN model_usage
+ SET DATA TYPE varchar(10)
+ USING CASE
+ WHEN TRUE THEN 'required'
+ ELSE 'disabled'
+ END
+ """
+ ],
+ reverse_sql=[
+ # When going backwards, we will set `model_usage=True` only for versions where models are required.
+ # `supported` versions will be disabled again, just like before we implemented this feature.
+ """
+ ALTER TABLE process_workerversion
+ ALTER COLUMN model_usage
+ SET DATA TYPE boolean
+ USING model_usage = 'required'
+ """
+ ],
+ state_operations=[
+ migrations.AlterField(
+ model_name='workerversion',
+ name='model_usage',
+ field=enumfields.fields.EnumField(blank=True, default='disabled', enum=arkindex.process.models.FeatureUsage, max_length=10),
+ ),
+ ],
+ ),
+ ]
diff --git a/arkindex/process/models.py b/arkindex/process/models.py
index 7044efa16dad41c7564d319507dfd718daca96c4..5b9ae5361478d572752c276ee570823d259730ff 100644
--- a/arkindex/process/models.py
+++ b/arkindex/process/models.py
@@ -699,7 +699,7 @@ class WorkerVersionState(Enum):
Error = 'error'
-class WorkerVersionGPUUsage(Enum):
+class FeatureUsage(Enum):
Disabled = 'disabled'
Supported = 'supported'
Required = 'required'
@@ -722,8 +722,8 @@ class WorkerVersion(models.Model):
)
configuration = models.JSONField()
state = EnumField(WorkerVersionState, default=WorkerVersionState.Created)
- gpu_usage = EnumField(WorkerVersionGPUUsage, default=WorkerVersionGPUUsage.Disabled, blank=True)
- model_usage = models.BooleanField(default=False)
+ gpu_usage = EnumField(FeatureUsage, default=FeatureUsage.Disabled, blank=True)
+ model_usage = EnumField(FeatureUsage, default=FeatureUsage.Disabled, blank=True)
# Foreign key towards the ponos Artifact holding the image content,
# for workers imported via the Git imports
@@ -954,7 +954,7 @@ class WorkerRun(models.Model):
token=token,
process=process,
extra_files=extra_files,
- requires_gpu=process.use_gpu and self.version.gpu_usage in (WorkerVersionGPUUsage.Required, WorkerVersionGPUUsage.Supported)
+ requires_gpu=process.use_gpu and self.version.gpu_usage in (FeatureUsage.Required, FeatureUsage.Supported)
)
return task, parents
diff --git a/arkindex/process/serializers/imports.py b/arkindex/process/serializers/imports.py
index 6fd8f7c56b9ecca0882b9d09e72d1a795d78b1d0..c58a03cf297ede95615fc96ea3eee3d105eb76fd 100644
--- a/arkindex/process/serializers/imports.py
+++ b/arkindex/process/serializers/imports.py
@@ -11,10 +11,10 @@ from arkindex.ponos.serializers import FarmSerializer, TaskLightSerializer
from arkindex.process.models import (
ActivityState,
DataFile,
+ FeatureUsage,
Process,
ProcessMode,
WorkerRun,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.process.serializers.git import RevisionSerializer
@@ -435,17 +435,17 @@ class StartProcessSerializer(serializers.Serializer):
# Use process.worker_runs.all() to access the (prefetched) worker_runs to avoid new SQL queries
# The related worker versions have also been prefetched
if len(list(self.instance.worker_runs.all())) > 0:
- if validated_data.get('use_gpu') and (not any(item.version.gpu_usage != WorkerVersionGPUUsage.Disabled for item in self.instance.worker_runs.all())):
+ if validated_data.get('use_gpu') and (not any(item.version.gpu_usage != FeatureUsage.Disabled for item in self.instance.worker_runs.all())):
errors['use_gpu'].append('The process is configured to use a GPU, but does not include any workers that support GPU usage.')
- # Check if a worker run has no model version but version.model_usage = True
+ # Check if a worker run has no model version but version.model_usage = Required
missing_model_versions = []
unavailable_versions = []
has_unavailable_model_versions = False
missing_required_configurations = []
for worker_run in self.instance.worker_runs.all():
- if worker_run.version.model_usage and worker_run.model_version_id is None:
+ if worker_run.version.model_usage == FeatureUsage.Required and worker_run.model_version_id is None:
missing_model_versions.append(worker_run.version.worker.name)
if (
worker_run.version.state != WorkerVersionState.Available
diff --git a/arkindex/process/serializers/training.py b/arkindex/process/serializers/training.py
index e50c1e9863a3712727ffd58aa2c631db2ce1a0b7..44ce215e2c1f21d397656c915b9d1fb795f27547 100644
--- a/arkindex/process/serializers/training.py
+++ b/arkindex/process/serializers/training.py
@@ -7,13 +7,13 @@ from rest_framework.exceptions import PermissionDenied, ValidationError
from arkindex.documents.models import Corpus, Element
from arkindex.ponos.models import Farm, Task
from arkindex.process.models import (
+ FeatureUsage,
Process,
ProcessDataset,
ProcessMode,
WorkerConfiguration,
WorkerRun,
WorkerVersion,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.process.utils import get_default_farm
@@ -197,9 +197,9 @@ class StartTrainingSerializer(serializers.ModelSerializer, WorkerACLMixin, Train
raise ValidationError({'__all__': ['Train, validation and test folders must be different.']})
# Ensure the worker version is compatible with the GPU mode
- if data['worker_version'].gpu_usage == WorkerVersionGPUUsage.Disabled and data["use_gpu"]:
+ if data['worker_version'].gpu_usage == FeatureUsage.Disabled and data["use_gpu"]:
raise ValidationError({'__all__': ['This worker version does not support GPU mode.']})
- elif data['worker_version'].gpu_usage == WorkerVersionGPUUsage.Required and not data["use_gpu"]:
+ elif data['worker_version'].gpu_usage == FeatureUsage.Required and not data["use_gpu"]:
raise ValidationError({'__all__': ['This worker version does only support GPU mode.']})
# Ensure the worker configuration match the given worker
diff --git a/arkindex/process/serializers/worker_runs.py b/arkindex/process/serializers/worker_runs.py
index 72ce2e3c2ad1fb8983645bdff419dd493da38b55..0de887642a0dd490722dfc365a8c23dd15fb8ef7 100644
--- a/arkindex/process/serializers/worker_runs.py
+++ b/arkindex/process/serializers/worker_runs.py
@@ -6,6 +6,7 @@ from rest_framework import serializers
from rest_framework.exceptions import ValidationError
from arkindex.process.models import (
+ FeatureUsage,
GitRef,
Process,
ProcessMode,
@@ -166,7 +167,7 @@ class WorkerRunSerializer(WorkerACLMixin, serializers.ModelSerializer):
if model_version.state != ModelVersionState.Available:
errors['model_version_id'].append('This ModelVersion is not in an Available state.')
- if not worker_version.model_usage:
+ if worker_version.model_usage == FeatureUsage.Disabled:
errors['model_version_id'].append('This worker version does not support models.')
existing_worker_run = process.worker_runs.filter(version=worker_version, model_version_id=model_version, configuration=configuration)
diff --git a/arkindex/process/serializers/workers.py b/arkindex/process/serializers/workers.py
index 916f1cb46871836e0f0e5c282767ce2445fb7f50..7ef377036bb0103e889a249b6cd3a73edc99d61a 100644
--- a/arkindex/process/serializers/workers.py
+++ b/arkindex/process/serializers/workers.py
@@ -15,6 +15,7 @@ from arkindex.ponos.models import Task
from arkindex.ponos.utils import get_process_from_task_auth
from arkindex.process.models import (
CorpusWorkerVersion,
+ FeatureUsage,
GitRef,
Process,
ProcessMode,
@@ -26,7 +27,6 @@ from arkindex.process.models import (
WorkerConfiguration,
WorkerType,
WorkerVersion,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.process.serializers.git import GitRefSerializer, RevisionWithRefsSerializer
@@ -220,8 +220,8 @@ class WorkerVersionSerializer(serializers.ModelSerializer):
state = EnumField(WorkerVersionState, required=False)
worker = WorkerLightSerializer(read_only=True)
revision = RevisionWithRefsSerializer(required=False, read_only=True, allow_null=True)
- gpu_usage = EnumField(WorkerVersionGPUUsage, required=False, default=WorkerVersionGPUUsage.Disabled)
- model_usage = serializers.BooleanField(required=False, default=False)
+ gpu_usage = EnumField(FeatureUsage, required=False, default=FeatureUsage.Disabled)
+ model_usage = EnumField(FeatureUsage, required=False, default=FeatureUsage.Disabled)
# Ensure worker version configuration JSON body is an object
configuration = serializers.DictField()
@@ -610,8 +610,8 @@ class DockerWorkerVersionSerializer(serializers.ModelSerializer):
required=False,
help_text="Optional list of types and names to create GitRefs",
)
- gpu_usage = EnumField(WorkerVersionGPUUsage, required=False, default=WorkerVersionGPUUsage.Disabled)
- model_usage = serializers.BooleanField(default=False)
+ gpu_usage = EnumField(FeatureUsage, required=False, default=FeatureUsage.Disabled)
+ model_usage = EnumField(FeatureUsage, required=False, default=FeatureUsage.Disabled)
class Meta:
model = WorkerVersion
diff --git a/arkindex/process/tests/test_create_training_process.py b/arkindex/process/tests/test_create_training_process.py
index 6f6257642adde13b36bc3b2efec7ce38e01bdb67..e129c15053d77e23f391b78d767576974baa3034 100644
--- a/arkindex/process/tests/test_create_training_process.py
+++ b/arkindex/process/tests/test_create_training_process.py
@@ -8,12 +8,12 @@ from rest_framework import status
from arkindex.ponos.models import Artifact, Farm
from arkindex.ponos.models import State as PonosState
from arkindex.process.models import (
+ FeatureUsage,
Process,
ProcessMode,
WorkerConfiguration,
WorkerRun,
WorkerVersion,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.project.tests import FixtureTestCase
@@ -31,7 +31,7 @@ class TestCreateTrainingProcess(FixtureTestCase):
cls.test_folder = cls.corpus.elements.create(type=cls.train_folder.type, name='Volume 3')
cls.training_worker_version = WorkerVersion.objects.get(worker__slug='worker-gpu')
cls.training_worker_version.worker.memberships.create(user=cls.user, level=Role.Admin.value)
- cls.training_worker_version.gpu_usage = WorkerVersionGPUUsage.Supported
+ cls.training_worker_version.gpu_usage = FeatureUsage.Supported
cls.training_worker_version.docker_image_iid = 1337
cls.training_worker_version.docker_image = Artifact.objects.first()
cls.training_worker_version.save()
@@ -223,7 +223,7 @@ class TestCreateTrainingProcess(FixtureTestCase):
An error is raised in case the GPU mode is incompatible with the worker
"""
self.client.force_login(self.user)
- self.training_worker_version.gpu_usage = WorkerVersionGPUUsage.Required
+ self.training_worker_version.gpu_usage = FeatureUsage.Required
self.training_worker_version.save()
with self.assertNumQueries(14):
response = self.client.post(reverse('api:process-training'), {
@@ -240,7 +240,7 @@ class TestCreateTrainingProcess(FixtureTestCase):
An error is raised in case the GPU mode is incompatible with the worker
"""
self.client.force_login(self.user)
- self.training_worker_version.gpu_usage = WorkerVersionGPUUsage.Disabled
+ self.training_worker_version.gpu_usage = FeatureUsage.Disabled
self.training_worker_version.save()
with self.assertNumQueries(14):
response = self.client.post(reverse('api:process-training'), {
diff --git a/arkindex/process/tests/test_docker_worker_version.py b/arkindex/process/tests/test_docker_worker_version.py
index 9342f1f9f952ffbe66a0b913c3d5328de858a611..c743d169c404755592f718b9fa00fa0cb5031154 100644
--- a/arkindex/process/tests/test_docker_worker_version.py
+++ b/arkindex/process/tests/test_docker_worker_version.py
@@ -2,7 +2,7 @@ from django.urls import reverse
from rest_framework import status
from arkindex.ponos.models import Farm
-from arkindex.process.models import GitRefType, ProcessMode, Repository, Worker, WorkerType, WorkerVersionGPUUsage
+from arkindex.process.models import FeatureUsage, GitRefType, ProcessMode, Repository, Worker, WorkerType
from arkindex.project.tests import FixtureAPITestCase
from arkindex.users.models import Role, Scope
@@ -128,8 +128,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
self.assertDictEqual(response.json(), {
'docker_image_iid': ['Not a valid string.'],
- 'gpu_usage': ['Value is not of type WorkerVersionGPUUsage'],
- 'model_usage': ['Must be a valid boolean.'],
+ 'gpu_usage': ['Value is not of type FeatureUsage'],
+ 'model_usage': ['Value is not of type FeatureUsage'],
'repository_url': ['Not a valid string.'],
'revision_author': ['Not a valid string.'],
'revision_hash': ['This field may not be blank.'],
@@ -190,8 +190,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
'docker_image': None,
'docker_image_iid': 'a_docker_image',
'docker_image_name': new_version.docker_image_name,
- 'gpu_usage': WorkerVersionGPUUsage.Disabled.value,
- 'model_usage': False,
+ 'gpu_usage': FeatureUsage.Disabled.value,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(new_revision.id),
'author': 'default',
@@ -232,8 +232,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
{'type': 'branch', 'name': 'master'},
{'type': 'tag', 'name': '2.0'},
],
- 'gpu_usage': WorkerVersionGPUUsage.Required.value,
- 'model_usage': True,
+ 'gpu_usage': FeatureUsage.Required.value,
+ 'model_usage': FeatureUsage.Supported.value,
},
format='json',
)
@@ -251,8 +251,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
'docker_image': None,
'docker_image_iid': 'e' * 512,
'docker_image_name': new_version.docker_image_name,
- 'gpu_usage': WorkerVersionGPUUsage.Required.value,
- 'model_usage': True,
+ 'gpu_usage': FeatureUsage.Required.value,
+ 'model_usage': FeatureUsage.Supported.value,
'revision': {
'id': str(new_revision.id),
'author': 'Iwan Roberts',
@@ -311,8 +311,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
{'type': 'branch', 'name': 'master'},
{'type': 'tag', 'name': '2.0'},
],
- 'gpu_usage': WorkerVersionGPUUsage.Required.value,
- 'model_usage': True,
+ 'gpu_usage': FeatureUsage.Required.value,
+ 'model_usage': FeatureUsage.Required.value,
},
format='json',
)
@@ -351,8 +351,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
'docker_image': None,
'docker_image_iid': 'a_docker_image',
'docker_image_name': new_version.docker_image_name,
- 'gpu_usage': WorkerVersionGPUUsage.Required.value,
- 'model_usage': True,
+ 'gpu_usage': FeatureUsage.Required.value,
+ 'model_usage': FeatureUsage.Required.value,
'revision': {
'id': str(new_revision.id),
'author': 'Iwan Roberts',
@@ -412,8 +412,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
'docker_image': str(self.version.docker_image_id),
'docker_image_iid': None,
'docker_image_name': self.version.docker_image_name,
- 'gpu_usage': WorkerVersionGPUUsage.Disabled.value,
- 'model_usage': False,
+ 'gpu_usage': FeatureUsage.Disabled.value,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version.revision.id),
'author': 'Test user',
@@ -479,8 +479,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
'docker_image': None,
'docker_image_iid': 'docker_image_42',
'docker_image_name': new_version.docker_image_name,
- 'gpu_usage': WorkerVersionGPUUsage.Disabled.value,
- 'model_usage': False,
+ 'gpu_usage': FeatureUsage.Disabled.value,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(new_revision.id),
'author': 'default',
@@ -514,8 +514,8 @@ class TestDockerWorkerVersion(FixtureAPITestCase):
self.assertEqual(new_version.configuration, {'key': 'value'})
self.assertEqual(new_version.docker_image_id, None)
self.assertEqual(new_version.docker_image_iid, 'docker_image_42')
- self.assertEqual(new_version.gpu_usage, WorkerVersionGPUUsage.Disabled)
- self.assertEqual(new_version.model_usage, False)
+ self.assertEqual(new_version.gpu_usage, FeatureUsage.Disabled)
+ self.assertEqual(new_version.model_usage, FeatureUsage.Disabled)
self.assertEqual(new_worker_type.slug, 'docker')
self.assertEqual(new_worker_type.display_name, 'Docker')
# User is granted an admin role on the repository
diff --git a/arkindex/process/tests/test_user_workerruns.py b/arkindex/process/tests/test_user_workerruns.py
index 124c6c4cc5b2009039d983a4c7d7fb55dab99f06..554b77208cb28fe5f0e1782d57471f39cdf6c7a3 100644
--- a/arkindex/process/tests/test_user_workerruns.py
+++ b/arkindex/process/tests/test_user_workerruns.py
@@ -6,6 +6,7 @@ from django.utils import timezone
from rest_framework import status
from arkindex.process.models import (
+ FeatureUsage,
Process,
ProcessMode,
WorkerConfiguration,
@@ -99,7 +100,7 @@ class TestUserWorkerRuns(FixtureAPITestCase):
'docker_image_name': f'my_repo.fake/workers/worker/reco:{str(self.version_1.id)}',
'gpu_usage': 'disabled',
'id': str(self.version_1.id),
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'author': 'Test user',
'commit_url': 'http://my_repo.fake/workers/worker/commit/1337',
@@ -147,7 +148,7 @@ class TestUserWorkerRuns(FixtureAPITestCase):
'docker_image_name': None,
'gpu_usage': 'disabled',
'id': str(self.custom_version.id),
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': None,
'state': 'created',
'version': 1,
@@ -229,7 +230,7 @@ class TestUserWorkerRuns(FixtureAPITestCase):
'docker_image_name': None,
'gpu_usage': 'disabled',
'id': str(self.other_version.id),
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': None,
'state': 'created',
'version': 2,
@@ -406,7 +407,7 @@ class TestUserWorkerRuns(FixtureAPITestCase):
'docker_image_name': None,
'gpu_usage': 'disabled',
'id': str(self.other_version.id),
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': None,
'state': 'created',
'version': 2,
diff --git a/arkindex/process/tests/test_workerruns.py b/arkindex/process/tests/test_workerruns.py
index 3b14e3405f7848e30ed26cb697e7af4556526275..14f6a5ca39aacf641c48cd591c8e383dd556832e 100644
--- a/arkindex/process/tests/test_workerruns.py
+++ b/arkindex/process/tests/test_workerruns.py
@@ -9,6 +9,7 @@ from rest_framework.exceptions import ValidationError
from arkindex.ponos.models import Agent, Artifact, Farm, State
from arkindex.process.models import (
+ FeatureUsage,
GitRefType,
ProcessMode,
Revision,
@@ -135,7 +136,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -235,7 +236,7 @@ class TestWorkerRuns(FixtureAPITestCase):
def test_create_unique(self):
self.client.force_login(self.user)
- self.version_1.model_usage = True
+ self.version_1.model_usage = FeatureUsage.Required
self.version_1.save()
cases = [
(None, None),
@@ -384,7 +385,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -481,7 +482,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -590,7 +591,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{version.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(version.revision.id),
'author': 'Teklia Bot',
@@ -675,7 +676,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -739,7 +740,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -793,7 +794,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': None,
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': None,
'version': 1,
'created': self.version_1.created.isoformat().replace('+00:00', 'Z'),
@@ -911,7 +912,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -985,7 +986,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -1234,7 +1235,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -1300,7 +1301,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -1374,7 +1375,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -1450,7 +1451,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=False
+ model_usage=FeatureUsage.Disabled
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -1486,7 +1487,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -1523,7 +1524,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -1570,7 +1571,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -1624,7 +1625,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run = self.process_1.worker_runs.create(
version=version,
@@ -1662,7 +1663,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Supported
)
run = self.process_1.worker_runs.create(
version=version_with_model,
@@ -1729,7 +1730,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{version_with_model.id}',
'gpu_usage': 'disabled',
- 'model_usage': True,
+ 'model_usage': FeatureUsage.Supported.value,
'revision': {
'id': str(rev.id),
'author': 'bob',
@@ -1768,7 +1769,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run = self.process_1.worker_runs.create(
version=version_with_model,
@@ -1832,7 +1833,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{version_with_model.id}',
'gpu_usage': 'disabled',
- 'model_usage': True,
+ 'model_usage': FeatureUsage.Required.value,
'revision': {
'id': str(rev_2.id),
'author': 'bob',
@@ -1911,7 +1912,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -1936,7 +1937,7 @@ class TestWorkerRuns(FixtureAPITestCase):
def test_update_unique(self):
self.client.force_login(self.user)
- self.version_1.model_usage = True
+ self.version_1.model_usage = FeatureUsage.Required
self.version_1.save()
cases = [
(None, None),
@@ -2142,7 +2143,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -2208,7 +2209,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -2279,7 +2280,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -2354,7 +2355,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=False
+ model_usage=FeatureUsage.Disabled
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -2389,7 +2390,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -2425,7 +2426,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -2465,7 +2466,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run_2 = self.process_1.worker_runs.create(
version=version_no_model,
@@ -2518,7 +2519,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run = self.process_1.worker_runs.create(
version=version,
@@ -2556,7 +2557,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run = self.process_1.worker_runs.create(
version=version_with_model,
@@ -2621,7 +2622,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{version_with_model.id}',
'gpu_usage': 'disabled',
- 'model_usage': True,
+ 'model_usage': FeatureUsage.Required.value,
'revision': {
'id': str(rev.id),
'author': 'bob',
@@ -2661,7 +2662,7 @@ class TestWorkerRuns(FixtureAPITestCase):
worker=self.worker_1,
revision=rev_2,
configuration={"test": "test2"},
- model_usage=True
+ model_usage=FeatureUsage.Required
)
run = self.process_1.worker_runs.create(
version=version_with_model,
@@ -2721,7 +2722,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{version_with_model.id}',
'gpu_usage': 'disabled',
- 'model_usage': True,
+ 'model_usage': FeatureUsage.Required.value,
'revision': {
'id': str(rev_2.id),
'author': 'bob',
@@ -2799,7 +2800,7 @@ class TestWorkerRuns(FixtureAPITestCase):
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'revision': {
'id': str(self.version_1.revision.id),
'author': 'Test user',
@@ -2824,7 +2825,7 @@ class TestWorkerRuns(FixtureAPITestCase):
def test_partial_update_unique(self):
self.client.force_login(self.user)
- self.version_1.model_usage = True
+ self.version_1.model_usage = FeatureUsage.Required
self.version_1.save()
cases = [
(None, None),
diff --git a/arkindex/process/tests/test_workers.py b/arkindex/process/tests/test_workers.py
index 01c36db284f24f3ee2638a79b754d5149a6751fb..a020684b2e92d1a5e72c164c746bfe39eca410fb 100644
--- a/arkindex/process/tests/test_workers.py
+++ b/arkindex/process/tests/test_workers.py
@@ -6,6 +6,7 @@ from rest_framework import status
from arkindex.ponos.models import Farm
from arkindex.process.models import (
CorpusWorkerVersion,
+ FeatureUsage,
GitRefType,
Process,
ProcessMode,
@@ -15,7 +16,6 @@ from arkindex.process.models import (
WorkerConfiguration,
WorkerType,
WorkerVersion,
- WorkerVersionGPUUsage,
WorkerVersionState,
)
from arkindex.project.tests import FixtureAPITestCase
@@ -77,7 +77,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
def setUp(self):
super().setUp()
self.version_1 = WorkerVersion.objects.get(worker=self.worker_reco)
- self.version_1.gpu_usage = WorkerVersionGPUUsage.Disabled
+ self.version_1.gpu_usage = FeatureUsage.Disabled
self.version_2 = self.worker_dla.versions.get()
self.version_custom = self.worker_custom.versions.get()
@@ -578,7 +578,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image_name': f'my_repo.fake/workers/worker/reco:{last_version.id}',
'state': 'error',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'worker': {
'id': str(self.worker_reco.id),
'name': self.worker_reco.name,
@@ -636,7 +636,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
worker=worker_2,
revision=self.rev,
configuration={"test": "test2"},
- gpu_usage=WorkerVersionGPUUsage.Disabled
+ gpu_usage=FeatureUsage.Disabled
)
with self.assertNumQueries(14):
@@ -776,7 +776,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
self.client.force_login(self.user)
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_custom.id)}),
- data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
@@ -789,7 +789,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
self.client.force_login(self.user)
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_reco.id)}),
- data={'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
@@ -802,7 +802,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
self.client.force_login(self.user)
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_custom.id)}),
- data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
@@ -818,7 +818,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
with self.assertNumQueries(4):
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_custom.id)}),
- data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
HTTP_AUTHORIZATION=f'Ponos {self.task.token}',
)
@@ -832,7 +832,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
with self.assertNumQueries(6):
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_custom.id)}),
- data={'configuration': {"test": "val"}, 'model_usage': True},
+ data={'configuration': {"test": "val"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@@ -843,7 +843,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
with self.assertNumQueries(10):
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_dla.id)}),
- data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
@@ -862,7 +862,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
with self.assertNumQueries(9):
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_custom.id)}),
- data={'configuration': {"test": "val"}, 'model_usage': True},
+ data={'configuration': {"test": "val"}, 'model_usage': FeatureUsage.Required.value},
format='json',
)
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
@@ -875,8 +875,8 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image': None,
'docker_image_iid': None,
'docker_image_name': None,
- 'gpu_usage': WorkerVersionGPUUsage.Disabled.value,
- 'model_usage': True,
+ 'gpu_usage': FeatureUsage.Disabled.value,
+ 'model_usage': FeatureUsage.Required.value,
'revision': None,
'state': WorkerVersionState.Created.value,
'version': 42,
@@ -905,7 +905,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
def test_create_version(self):
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_reco.id)}),
- data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True},
+ data={'revision_id': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': FeatureUsage.Required.value},
format='json',
HTTP_AUTHORIZATION=f'Ponos {self.task.token}',
)
@@ -917,7 +917,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
self.assertEqual(data['revision']['id'], str(self.rev2.id))
self.assertEqual(data['state'], 'created')
self.assertEqual(data['gpu_usage'], 'disabled')
- self.assertEqual(data['model_usage'], True)
+ self.assertEqual(data['model_usage'], FeatureUsage.Required.value)
def test_create_version_wrong_gpu_usage(self):
response = self.client.post(
@@ -1463,7 +1463,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'state': 'available',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'worker': {
'id': str(self.worker_reco.id),
'name': self.worker_reco.name,
@@ -1489,7 +1489,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
'state': 'available',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'worker': {
'id': str(self.worker_reco.id),
'name': self.worker_reco.name,
@@ -1508,7 +1508,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image': '12341234-1234-1234-1234-123412341234',
'state': 'error',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@@ -1523,7 +1523,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image': '12341234-1234-1234-1234-123412341234',
'state': 'error',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
@@ -1544,7 +1544,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'docker_image': str(docker_image.id),
'state': 'error',
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@@ -1699,7 +1699,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
worker=self.worker_reco,
revision=self.rev2,
configuration={"test": "test1"},
- gpu_usage=WorkerVersionGPUUsage.Disabled
+ gpu_usage=FeatureUsage.Disabled
)
self.assertEqual(
version_2.docker_image_name,
@@ -1715,7 +1715,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
worker=self.worker_reco,
revision=self.rev2,
configuration={"test": "test1"},
- gpu_usage=WorkerVersionGPUUsage.Disabled
+ gpu_usage=FeatureUsage.Disabled
)
self.assertEqual(version_2.slug, 'reco_012345')
@@ -1763,7 +1763,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_2.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/dla:{self.version_2.id}',
@@ -1794,7 +1794,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_1.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
@@ -1856,7 +1856,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_2.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/dla:{self.version_2.id}',
@@ -1887,7 +1887,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_1.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
@@ -1946,7 +1946,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_2.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/dla:{self.version_2.id}',
@@ -1992,7 +1992,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_1.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
@@ -2057,7 +2057,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': 'disabled',
- 'model_usage': False,
+ 'model_usage': FeatureUsage.Disabled.value,
'docker_image': str(self.version_1.docker_image_id),
'docker_image_iid': None,
'docker_image_name': f'my_repo.fake/workers/worker/reco:{self.version_1.id}',
diff --git a/arkindex/project/settings.py b/arkindex/project/settings.py
index 7339c279353647f3f2c4f12977705a64dc34f997..c17c2471e469659d631c268b2a69c59af7da618b 100644
--- a/arkindex/project/settings.py
+++ b/arkindex/project/settings.py
@@ -254,6 +254,7 @@ SPECTACULAR_SETTINGS = {
'ModelVersionState': 'arkindex.training.models.ModelVersionState',
'MetricMode': 'arkindex.training.models.MetricMode',
'DatasetState': 'arkindex.training.models.DatasetState',
+ 'FeatureUsage': 'arkindex.process.models.FeatureUsage',
},
'TAGS': [
{'name': 'classifications'},