Make ModelVersion size field PositiveBigIntegerField

arkindex/training/migrations/0010_alter_modelversion_size.py

+# 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),
+        ),
+    ]

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)

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