Remove retry if dropped (#3998)

With the changes with database proxying this should no longer happen

Author: jmsmkn

app/grandchallenge/cases/tasks.py

@@ -7,6 +7,7 @@
 from tempfile import TemporaryDirectory
 
 from billiard.exceptions import SoftTimeLimitExceeded, TimeLimitExceeded
+from django.apps import apps
 from django.conf import settings
 from django.core.exceptions import ValidationError
 from django.core.files import File
@@ -21,10 +22,7 @@
 from grandchallenge.cases.models import Image, ImageFile, RawImageUploadSession
 from grandchallenge.components.backends.utils import safe_extract
 from grandchallenge.components.models import ComponentInterface
-from grandchallenge.components.tasks import (
-    get_model_instance,
-    lock_model_instance,
-)
+from grandchallenge.components.tasks import lock_model_instance
 from grandchallenge.core.celery import acks_late_2xlarge_task
 from grandchallenge.core.exceptions import LockNotAcquiredException
 from grandchallenge.reader_studies.models import DisplaySet
@@ -140,12 +138,12 @@ def build_images(  # noqa:C901
     )
 
     if linked_object_pk:
+        linked_model = apps.get_model(
+            app_label=linked_app_label, model_name=linked_model_name
+        )
+
         try:
-            linked_object = get_model_instance(
-                app_label=linked_app_label,
-                model_name=linked_model_name,
-                pk=linked_object_pk,
-            )
+            linked_object = linked_model.objects.get(pk=linked_object_pk)
         except (ArchiveItem.DoesNotExist, DisplaySet.DoesNotExist):
             # users can delete archive items and display sets before this task runs
             logger.info(

app/grandchallenge/components/tasks.py

@@ -717,43 +717,6 @@ def _get_image_config_file(
     return {"image_sha256": image_sha256, "config": config}
 
 
-def retry_if_dropped(func):
-    """
-    Retry a function that relies on an open database connection.
-
-    Use this decorator when you have a long running task as sometimes the db
-    connection will drop.
-    """
-
-    def wrapper(*args, **kwargs):
-        n_tries = 0
-        max_tries = 2
-        err = None
-
-        while n_tries < max_tries:
-            n_tries += 1
-
-            try:
-                return func(*args, **kwargs)
-            except OperationalError as e:
-                err = e
-
-                # This needs to be a local import
-                from django.db import connection
-
-                connection.close()
-
-        raise err
-
-    return wrapper
-
-
-@retry_if_dropped
-def get_model_instance(*, app_label, model_name, **kwargs):
-    model = apps.get_model(app_label=app_label, model_name=model_name)
-    return model.objects.get(**kwargs)
-
-
 def lock_model_instance(*, app_label, model_name, **kwargs):
     """
     Locks a model instance for update.
@@ -827,9 +790,8 @@ def execute_job(  # noqa: C901
 
     Once the job has executed it will be in the EXECUTING or FAILURE states.
     """
-    job = get_model_instance(
-        pk=job_pk, app_label=job_app_label, model_name=job_model_name
-    )
+    model = apps.get_model(app_label=job_app_label, model_name=job_model_name)
+    job = model.objects.get(pk=job_pk)
     executor = job.get_executor(backend=backend)
 
     if job.status == job.PROVISIONED:
@@ -844,15 +806,11 @@ def execute_job(  # noqa: C901
         raise PriorStepFailed(msg)
 
     try:
-        # This call is potentially very long
         executor.execute()
     except RetryStep:
         job.update_status(status=job.PROVISIONED)
         raise
     except ComponentException as e:
-        job = get_model_instance(
-            pk=job_pk, app_label=job_app_label, model_name=job_model_name
-        )
         job.update_status(
             status=job.FAILURE,
             stdout=executor.stdout,
@@ -861,19 +819,13 @@ def execute_job(  # noqa: C901
             detailed_error_message=e.message_details,
         )
     except (SoftTimeLimitExceeded, TimeLimitExceeded):
-        job = get_model_instance(
-            pk=job_pk, app_label=job_app_label, model_name=job_model_name
-        )
         job.update_status(
             status=job.FAILURE,
             stdout=executor.stdout,
             stderr=executor.stderr,
             error_message="Time limit exceeded",
         )
     except Exception:
-        job = get_model_instance(
-            pk=job_pk, app_label=job_app_label, model_name=job_model_name
-        )
         job.update_status(
             status=job.FAILURE,
             stdout=executor.stdout,
@@ -1031,9 +983,8 @@ def retry_task(
     backend: str,
 ):
     """Retries an existing task that was previously provisioned"""
-    job = get_model_instance(
-        pk=job_pk, app_label=job_app_label, model_name=job_model_name
-    )
+    model = apps.get_model(app_label=job_app_label, model_name=job_model_name)
+    job = model.objects.get(pk=job_pk)
     executor = job.get_executor(backend=backend)
 
     if job.status != job.PROVISIONED:
@@ -1062,27 +1013,24 @@ def deprovision_job(
     job_model_name: str,
     backend: str,
 ):
-    job = get_model_instance(
-        pk=job_pk, app_label=job_app_label, model_name=job_model_name
-    )
+    model = apps.get_model(app_label=job_app_label, model_name=job_model_name)
+    job = model.objects.get(pk=job_pk)
 
     executor = job.get_executor(backend=backend)
     executor.deprovision()
 
 
 @shared_task
 def start_service(*, pk: uuid.UUID, app_label: str, model_name: str):
-    session = get_model_instance(
-        pk=pk, app_label=app_label, model_name=model_name
-    )
+    model = apps.get_model(app_label=app_label, model_name=model_name)
+    session = model.objects.get(pk=pk)
     session.start()
 
 
 @shared_task
 def stop_service(*, pk: uuid.UUID, app_label: str, model_name: str):
-    session = get_model_instance(
-        pk=pk, app_label=app_label, model_name=model_name
-    )
+    model = apps.get_model(app_label=app_label, model_name=model_name)
+    session = model.objects.get(pk=pk)
     session.stop()
 
 
@@ -1235,60 +1183,52 @@ def preload_interactive_algorithms():
 
 
 @acks_late_micro_short_task
+@transaction.atomic
 def add_image_to_component_interface_value(
     *, component_interface_value_pk, upload_session_pk
 ):
-    with transaction.atomic():
-        session = RawImageUploadSession.objects.get(pk=upload_session_pk)
+    from grandchallenge.components.models import ComponentInterfaceValue
 
-        if session.image_set.count() != 1:
-            session.status = RawImageUploadSession.FAILURE
-            session.error_message = (
-                "Image imports should result in a single image"
-            )
-            session.save()
-            return
+    session = RawImageUploadSession.objects.get(pk=upload_session_pk)
 
-        civ = get_model_instance(
-            pk=component_interface_value_pk,
-            app_label="components",
-            model_name="componentinterfacevalue",
-        )
+    if session.image_set.count() != 1:
+        session.status = RawImageUploadSession.FAILURE
+        session.error_message = "Image imports should result in a single image"
+        session.save()
+        return
 
-        civ.image = session.image_set.get()
-        civ.full_clean()
-        civ.save()
+    civ = ComponentInterfaceValue.objects.get(pk=component_interface_value_pk)
+
+    civ.image = session.image_set.get()
+    civ.full_clean()
+    civ.save()
 
-        civ.image.update_viewer_groups_permissions()
+    civ.image.update_viewer_groups_permissions()
 
 
 @acks_late_2xlarge_task
+@transaction.atomic
 def civ_value_to_file(*, civ_pk):
-    with transaction.atomic():
-        civ = get_model_instance(
-            pk=civ_pk,
-            app_label="components",
-            model_name="componentinterfacevalue",
-        )
+    from grandchallenge.components.models import ComponentInterfaceValue
 
-        if civ.value is None:
-            raise RuntimeError("CIV value is None")
+    civ = ComponentInterfaceValue.objects.get(pk=civ_pk)
 
-        civ.file = ContentFile(
-            json.dumps(civ.value).encode("utf-8"),
-            name=Path(civ.interface.relative_path).name,
-        )
-        civ.value = None
-        civ.save()
+    if civ.value is None:
+        raise RuntimeError("CIV value is None")
+
+    civ.file = ContentFile(
+        json.dumps(civ.value).encode("utf-8"),
+        name=Path(civ.interface.relative_path).name,
+    )
+    civ.value = None
+    civ.save()
 
 
 @acks_late_2xlarge_task
 def validate_voxel_values(*, civ_pk):
-    civ = get_model_instance(
-        pk=civ_pk,
-        app_label="components",
-        model_name="componentinterfacevalue",
-    )
+    from grandchallenge.components.models import ComponentInterfaceValue
+
+    civ = ComponentInterfaceValue.objects.get(pk=civ_pk)
 
     first_file = civ.image.files.first()
     if (