From e5b83d7d438cfb2f2a3398d1a1cf64876230e7a7 Mon Sep 17 00:00:00 2001 From: amickan Date: Fri, 24 May 2024 14:27:15 +0200 Subject: [PATCH] Add test for model import task --- ...rithmmodeluserobjectpermission_and_more.py | 5 ++- app/grandchallenge/algorithms/models.py | 14 +++++++ app/grandchallenge/algorithms/tasks.py | 22 +++++++--- app/tests/algorithms_tests/test_tasks.py | 40 +++++++++++++------ 4 files changed, 60 insertions(+), 21 deletions(-) diff --git a/app/grandchallenge/algorithms/migrations/0049_algorithmmodel_algorithmmodeluserobjectpermission_and_more.py b/app/grandchallenge/algorithms/migrations/0049_algorithmmodel_algorithmmodeluserobjectpermission_and_more.py index 8c7f49095a..60a75427a2 100644 --- a/app/grandchallenge/algorithms/migrations/0049_algorithmmodel_algorithmmodeluserobjectpermission_and_more.py +++ b/app/grandchallenge/algorithms/migrations/0049_algorithmmodel_algorithmmodeluserobjectpermission_and_more.py @@ -1,4 +1,4 @@ -# Generated by Django 4.2.13 on 2024-05-23 11:17 +# Generated by Django 4.2.13 on 2024-05-24 10:50 import uuid @@ -15,9 +15,9 @@ class Migration(migrations.Migration): dependencies = [ - migrations.swappable_dependency(settings.AUTH_USER_MODEL), ("uploads", "0006_userupload_mimetype"), ("auth", "0012_alter_user_first_name_max_length"), + migrations.swappable_dependency(settings.AUTH_USER_MODEL), ("algorithms", "0048_job_detailed_error_message"), ] @@ -52,6 +52,7 @@ class Migration(migrations.Migration): default=0, ), ), + ("status", models.TextField(editable=False)), ( "model", models.FileField( diff --git a/app/grandchallenge/algorithms/models.py b/app/grandchallenge/algorithms/models.py index f2e8d2cead..694b4b00ab 100644 --- a/app/grandchallenge/algorithms/models.py +++ b/app/grandchallenge/algorithms/models.py @@ -1,6 +1,7 @@ import logging from datetime import datetime, timedelta +import boto3 from actstream.actions import follow, is_following from actstream.models import Follow from django.conf import settings @@ -949,6 +950,7 @@ class AlgorithmModel(UUIDModel): default=ImportStatusChoices.INITIALIZED, db_index=True, ) + status = models.TextField(editable=False) user_upload = models.ForeignKey( UserUpload, blank=True, @@ -1012,6 +1014,18 @@ def mark_desired_version(self, peer_models=None): models.append(self) self.__class__.objects.bulk_update(models, ["is_desired_version"]) + def delete_model_file(self): + if not self.import_status == ImportStatusChoices.FAILED: + raise RuntimeError("Cannot delete model from completed upload.") + + s3_client = boto3.client( + "s3", + endpoint_url=settings.AWS_S3_ENDPOINT_URL, + ) + s3_client.delete_object( + Bucket=self.model.storage.bucket_name, Key=self.model.name + ) + class AlgorithmModelUserObjectPermission(UserObjectPermissionBase): content_object = models.ForeignKey( diff --git a/app/grandchallenge/algorithms/tasks.py b/app/grandchallenge/algorithms/tasks.py index 891ad23211..3bb2fb2e24 100644 --- a/app/grandchallenge/algorithms/tasks.py +++ b/app/grandchallenge/algorithms/tasks.py @@ -484,7 +484,7 @@ def set_credits_per_job(): @transaction.atomic() -@shared_task(**settings.CELERY_TASK_DECORATOR_KWARGS["acks-late-2xlarge"]) +@shared_task(**settings.CELERY_TASK_DECORATOR_KWARGS["acks-late-micro-short"]) def assign_algorithm_model_from_upload(*, algorithm_model_pk, retries=0): from grandchallenge.algorithms.models import AlgorithmModel @@ -512,18 +512,28 @@ def assign_algorithm_model_from_upload(*, algorithm_model_pk, retries=0): ) return - # catch errors with uploading? current_model.user_upload.copy_object(to_field=current_model.model) - # retrieve sha256 and check if it's unique, error out if not - current_model.sha256 = get_object_sha256(current_model.model) + + sha256 = get_object_sha256(current_model.model) + if AlgorithmModel.objects.filter(sha256=sha256).exists(): + current_model.import_status = ImportStatusChoices.FAILED + current_model.status = ( + "Algorithm model with this sha256 already exists." + ) + current_model.save() + current_model.user_upload.delete() + current_model.delete_model_file() + return + + current_model.sha256 = sha256 current_model.size_in_storage = current_model.model.size current_model.import_status = ImportStatusChoices.COMPLETED current_model.save() + current_model.user_upload.delete() # mark as desired version and pass locked peer models directly since else - # mark_desired_version will try to lock the peer models a second time, - # which will fail + # mark_desired_version will fail trying to access the locked models current_model.mark_desired_version(peer_models=peer_models) diff --git a/app/tests/algorithms_tests/test_tasks.py b/app/tests/algorithms_tests/test_tasks.py index e73452ce2b..9a0538e2f7 100644 --- a/app/tests/algorithms_tests/test_tasks.py +++ b/app/tests/algorithms_tests/test_tasks.py @@ -9,7 +9,7 @@ from django.core.files.base import ContentFile, File from requests import put -from grandchallenge.algorithms.models import AlgorithmModel, Job +from grandchallenge.algorithms.models import Job from grandchallenge.algorithms.tasks import ( assign_algorithm_model_from_upload, create_algorithm_jobs, @@ -35,7 +35,6 @@ AlgorithmJobFactory, AlgorithmModelFactory, ) -from tests.cases_tests import RESOURCE_PATH from tests.cases_tests.factories import RawImageUploadSessionFactory from tests.components_tests.factories import ( ComponentInterfaceFactory, @@ -772,29 +771,44 @@ def test_setting_credits_per_job( assert alg.credits_per_job == test["credits"] -@pytest.mark.django_db(transaction=True) -def test_assign_algorithm_model_from_upload( - algorithm_io_image, settings, django_capture_on_commit_callbacks -): +@pytest.mark.django_db() +def test_assign_algorithm_model_from_upload(settings): # Override the celery settings settings.task_eager_propagates = (True,) settings.task_always_eager = (True,) user = UserFactory() alg = AlgorithmFactory() - alg.add_editor(user) upload = create_upload_from_file( - creator=user, file_path=RESOURCE_PATH / "test.zip" + creator=user, + file_path=Path(__file__).parent / "resources" / "model.tar.gz", ) model = AlgorithmModelFactory( algorithm=alg, creator=user, user_upload=upload ) assert model.is_desired_version is False - with django_capture_on_commit_callbacks(): - assign_algorithm_model_from_upload( - algorithm_model_pk=model.pk, - ) - model = AlgorithmModel.objects.get(pk=model.pk) + assign_algorithm_model_from_upload( + algorithm_model_pk=model.pk, + ) + model.refresh_from_db() assert model.is_desired_version assert model.import_status == ImportStatusChoices.COMPLETED + + upload2 = create_upload_from_file( + creator=user, + file_path=Path(__file__).parent / "resources" / "model.tar.gz", + ) + model2 = AlgorithmModelFactory( + algorithm=alg, creator=user, user_upload=upload2 + ) + assign_algorithm_model_from_upload( + algorithm_model_pk=model2.pk, + ) + model2.refresh_from_db() + assert not model2.is_desired_version + assert model2.import_status == ImportStatusChoices.FAILED + assert model2.status == "Algorithm model with this sha256 already exists." + assert not model2.user_upload + with pytest.raises(FileNotFoundError): + model2.model.file