From 699c51b3dadf180675889c18f779736f32c05a2b Mon Sep 17 00:00:00 2001
From: mlbonhomme <bonhomme@teklia.com>
Date: Wed, 6 Nov 2024 11:41:18 +0100
Subject: [PATCH] Make ModelVersion size field PositiveBigIntegerField
---
.../0010_alter_modelversion_size.py | 18 ++++++++
arkindex/training/models.py | 2 +-
arkindex/training/tests/test_model_api.py | 46 +++++++++++++++++++
3 files changed, 65 insertions(+), 1 deletion(-)
create mode 100644 arkindex/training/migrations/0010_alter_modelversion_size.py
diff --git a/arkindex/training/migrations/0010_alter_modelversion_size.py b/arkindex/training/migrations/0010_alter_modelversion_size.py
new file mode 100644
index 0000000000..ce11bc782f
--- /dev/null
+++ b/arkindex/training/migrations/0010_alter_modelversion_size.py
@@ -0,0 +1,18 @@
+# Generated by Django 5.0.8 on 2024-11-06 10:28
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ("training", "0009_remove_metrics"),
+ ]
+
+ operations = [
+ migrations.AlterField(
+ model_name="modelversion",
+ name="size",
+ field=models.PositiveBigIntegerField(blank=True, help_text="file size in bytes", null=True),
+ ),
+ ]
diff --git a/arkindex/training/models.py b/arkindex/training/models.py
index 4df86f5ad1..1e08f99d9f 100644
--- a/arkindex/training/models.py
+++ b/arkindex/training/models.py
@@ -105,7 +105,7 @@ class ModelVersion(S3FileMixin, IndexableModel):
)
# Size of the archive
- size = models.PositiveIntegerField(null=True, blank=True, help_text="file size in bytes")
+ size = models.PositiveBigIntegerField(null=True, blank=True, help_text="file size in bytes")
# Store dictionary of parameters given by the ML developer
configuration = models.JSONField(blank=True, default=dict)
diff --git a/arkindex/training/tests/test_model_api.py b/arkindex/training/tests/test_model_api.py
index 0c7afe0b71..114abe1b97 100644
--- a/arkindex/training/tests/test_model_api.py
+++ b/arkindex/training/tests/test_model_api.py
@@ -837,6 +837,52 @@ class TestModelAPI(FixtureAPITestCase):
}
)
+ @patch("arkindex.project.aws.s3.meta.client.generate_presigned_url")
+ def test_create_huge_model_version(self, s3_presigned_url_mock):
+ """
+ Creates a new model version that is very big
+ """
+ self.client.force_login(self.user1)
+ s3_presigned_url_mock.return_value = "http://s3/upload_put_url"
+ fake_now = timezone.now()
+ # To mock the creation date
+ with patch("django.utils.timezone.now") as mock_now:
+ mock_now.return_value = fake_now
+ with self.assertNumQueries(6):
+ response = self.client.post(
+ reverse("api:model-versions", kwargs={"pk": str(self.model1.id)}),
+ {
+ "tag": "TAG",
+ "size": 9223372036854775806,
+ "hash": "d" * 32,
+ "archive_hash": "e" * 32,
+ "description": "description",
+ "configuration": {"hello": "this is me"},
+ },
+ )
+ self.assertEqual(response.status_code, status.HTTP_201_CREATED)
+ data = response.json()
+ self.assertIn("id", data)
+ df = ModelVersion.objects.get(id=data["id"])
+
+ self.assertDictEqual(
+ data,
+ {
+ "id": str(df.id),
+ "model_id": str(self.model1.id),
+ "parent": None,
+ "description": "description",
+ "state": ModelVersionState.Created.value,
+ "configuration": {"hello": "this is me"},
+ "tag": "TAG",
+ "size": 9223372036854775806,
+ "hash": "d" * 32,
+ "created": fake_now.isoformat().replace("+00:00", "Z"),
+ "s3_url": None,
+ "s3_put_url": s3_presigned_url_mock.return_value
+ }
+ )
+
def test_create_model_version_unique_tag_model(self):
"""
Raises 400 when creating a model version that already exists, same model_id and tag
--
GitLab