diff --git a/arkindex/dataimport/admin.py b/arkindex/dataimport/admin.py
index d2cc69cee0886bc1010803cfac29438bcbf63afb..365a7072af6bdee7fa955c855e0f2c1e2196bc05 100644
--- a/arkindex/dataimport/admin.py
+++ b/arkindex/dataimport/admin.py
@@ -115,7 +115,7 @@ class WorkerTypeAdmin(admin.ModelAdmin):
class WorkerVersionAdmin(admin.ModelAdmin):
list_display = ('id', 'worker', 'revision')
list_filter = ('worker', )
- fields = ('id', 'worker', 'revision', 'configuration')
+ fields = ('id', 'worker', 'revision', 'configuration', 'model_usage')
readonly_fields = ('id', )
raw_id_fields = ('docker_image', 'revision')
diff --git a/arkindex/dataimport/api.py b/arkindex/dataimport/api.py
index dad6d83f7d84f022e16171a2ac5b1e6614b55aec..e6b4b4c088321700b2b40a28e4a3f173be58bc4a 100644
--- a/arkindex/dataimport/api.py
+++ b/arkindex/dataimport/api.py
@@ -963,6 +963,7 @@ class WorkerVersionList(WorkerACLMixin, ListCreateAPIView):
defaults={
'configuration': serializer.validated_data['configuration'],
'gpu_usage': serializer.validated_data['gpu_usage'],
+ 'model_usage': serializer.validated_data['model_usage'],
}
)
diff --git a/arkindex/dataimport/migrations/0047_workerversion_model_usage.py b/arkindex/dataimport/migrations/0047_workerversion_model_usage.py
new file mode 100644
index 0000000000000000000000000000000000000000..9cc56f42526245db073481f0c7f1a6d3b1c01dc7
--- /dev/null
+++ b/arkindex/dataimport/migrations/0047_workerversion_model_usage.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.0.2 on 2022-05-03 15:12
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ('dataimport', '0046_workertype_alter_worker_type'),
+ ]
+
+ operations = [
+ migrations.AddField(
+ model_name='workerversion',
+ name='model_usage',
+ field=models.BooleanField(default=False),
+ ),
+ ]
diff --git a/arkindex/dataimport/models.py b/arkindex/dataimport/models.py
index 176a71215728701c57ecfb0773b57e62e9411a9c..85aefc4b8f7584c7c787e54bb5de24c892f41a5b 100644
--- a/arkindex/dataimport/models.py
+++ b/arkindex/dataimport/models.py
@@ -576,6 +576,7 @@ 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)
# Foreign key towards the ponos Artifact holding the image content
docker_image = models.ForeignKey(Artifact, on_delete=models.SET_NULL, null=True, blank=True)
diff --git a/arkindex/dataimport/serializers/workers.py b/arkindex/dataimport/serializers/workers.py
index 16a6246035b92e9f0290951811db6c225c3836cc..d277a6897386fa676efbb011bceb67465ee52b6f 100644
--- a/arkindex/dataimport/serializers/workers.py
+++ b/arkindex/dataimport/serializers/workers.py
@@ -129,6 +129,7 @@ class WorkerVersionSerializer(serializers.ModelSerializer):
revision = serializers.UUIDField()
element_count = serializers.IntegerField(read_only=True)
gpu_usage = EnumField(WorkerVersionGPUUsage, required=False, default=WorkerVersionGPUUsage.Disabled)
+ model_usage = serializers.BooleanField(required=False, default=False)
# Serialize worker with its basic information
class Meta:
@@ -142,6 +143,7 @@ class WorkerVersionSerializer(serializers.ModelSerializer):
'docker_image_name',
'state',
'gpu_usage',
+ 'model_usage',
'element_count',
'worker'
)
diff --git a/arkindex/dataimport/tests/test_workers.py b/arkindex/dataimport/tests/test_workers.py
index cf0c34cf0f53955966ab7cd23c50d18b4094906c..2cda8dc6d90effddeeaf23932b0990ef83449238 100644
--- a/arkindex/dataimport/tests/test_workers.py
+++ b/arkindex/dataimport/tests/test_workers.py
@@ -441,6 +441,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,
'worker': {
'id': str(self.worker_1.id),
'name': self.worker_1.name,
@@ -654,15 +655,17 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
self.client.force_login(self.internal_user)
response = self.client.post(
reverse('api:worker-versions', kwargs={'pk': str(self.worker_1.id)}),
- data={'revision': str(self.rev2.id), 'configuration': {"test": "test2"}}, format='json'
+ data={'revision': str(self.rev2.id), 'configuration': {"test": "test2"}, 'model_usage': True}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+
data = response.json()
self.assertNotEqual(data['id'], str(self.version_1.id))
self.assertEqual(data['configuration'], {"test": "test2"})
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)
def test_no_user_configuration_ok(self):
self.client.force_login(self.internal_user)
@@ -991,6 +994,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,
'worker': {
'id': str(self.worker_1.id),
'name': self.worker_1.name,
@@ -1014,6 +1018,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,
'worker': {
'id': str(self.worker_1.id),
'name': self.worker_1.name,
@@ -1029,7 +1034,8 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'configuration': {"test": "test2"},
'docker_image': '12341234-1234-1234-1234-123412341234',
'state': 'error',
- 'gpu_usage': 'disabled'
+ 'gpu_usage': 'disabled',
+ 'model_usage': False,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@@ -1043,7 +1049,8 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'configuration': {"test": "test2"},
'docker_image': '12341234-1234-1234-1234-123412341234',
'state': 'error',
- 'gpu_usage': 'disabled'
+ 'gpu_usage': 'disabled',
+ 'model_usage': False,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
@@ -1063,7 +1070,8 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'configuration': {"test": "test2"},
'docker_image': str(docker_image.id),
'state': 'error',
- 'gpu_usage': 'disabled'
+ 'gpu_usage': 'disabled',
+ 'model_usage': False,
}, format='json'
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@@ -1250,6 +1258,7 @@ class TestWorkersWorkerVersions(FixtureAPITestCase):
'refs': []
},
'gpu_usage': version.gpu_usage.value,
+ 'model_usage': version.model_usage,
'docker_image': str(version.docker_image_id),
'docker_image_iid': None,
'docker_image_name': version.docker_image_name,
diff --git a/arkindex/sql_validation/indexer_prefetch.sql b/arkindex/sql_validation/indexer_prefetch.sql
index a3615d5d6b1a89eee429bb7507386236c942e76a..bd1d0779661fd967cf627e8250d20c418070d590 100644
--- a/arkindex/sql_validation/indexer_prefetch.sql
+++ b/arkindex/sql_validation/indexer_prefetch.sql
@@ -52,6 +52,7 @@ SELECT "dataimport_workerversion"."id",
"dataimport_workerversion"."configuration",
"dataimport_workerversion"."state",
"dataimport_workerversion"."gpu_usage",
+ "dataimport_workerversion"."model_usage",
"dataimport_workerversion"."docker_image_id",
"dataimport_workerversion"."docker_image_iid"
FROM "dataimport_workerversion"
@@ -118,6 +119,7 @@ SELECT "dataimport_workerversion"."id",
"dataimport_workerversion"."configuration",
"dataimport_workerversion"."state",
"dataimport_workerversion"."gpu_usage",
+ "dataimport_workerversion"."model_usage",
"dataimport_workerversion"."docker_image_id",
"dataimport_workerversion"."docker_image_iid"
FROM "dataimport_workerversion"