fix: load remote image once

CHANGELOG.md

@@ -39,6 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Convert asset uuid to str in `FailedAssetLogsViewSet` [#804](https://github.com/Substra/substra-backend/pull/804)
+- Organisations which are not function owner loads function image only once ([#833](https://github.com/Substra/substra-backend/pull/833/files))
 
 ## [0.42.2](https://github.com/Substra/substra-backend/releases/tag/0.42.2) 2023-10-18
 

backend/substrapp/compute_tasks/execute.py

@@ -14,6 +14,7 @@
 
 from substrapp.compute_tasks import compute_task as task_utils
 from substrapp.compute_tasks import errors as compute_task_errors
+from substrapp.compute_tasks import image_builder
 from substrapp.compute_tasks import utils
 from substrapp.compute_tasks.command import get_exec_command
 from substrapp.compute_tasks.command import get_exec_command_args
@@ -27,6 +28,8 @@
 from substrapp.compute_tasks.volumes import get_worker_subtuple_pvc_name
 from substrapp.docker_registry import get_container_image_name
 from substrapp.docker_registry import get_entrypoint
+from substrapp.exceptions import OrganizationError
+from substrapp.exceptions import OrganizationHttpError
 from substrapp.exceptions import PodReadinessTimeoutError
 from substrapp.kubernetes_utils import delete_pod
 from substrapp.kubernetes_utils import execute
@@ -35,6 +38,7 @@
 from substrapp.kubernetes_utils import wait_for_pod_readiness
 from substrapp.models import ImageEntrypoint
 from substrapp.orchestrator import get_orchestrator_client
+from substrapp.utils import get_owner
 from substrapp.utils import timeit
 
 logger = structlog.get_logger(__name__)
@@ -51,17 +55,24 @@ def execute_compute_task(ctx: Context) -> None:
     env = get_environment(ctx)
     image = get_container_image_name(container_image_tag)
 
-    # save entrypoint to DB
-    entrypoint = get_entrypoint(container_image_tag)
-    ImageEntrypoint.objects.get_or_create(
-        archive_checksum=ctx.function.archive_address.checksum, entrypoint_json=entrypoint
-    )
-
     k8s_client = _get_k8s_client()
 
     should_create_pod = not pod_exists_by_label_selector(k8s_client, compute_pod.label_selector)
 
     if should_create_pod:
+        if get_owner() != ctx.function.owner:
+            try:
+                image_builder.load_remote_function_image(ctx.function, channel_name)
+            except OrganizationHttpError as e:
+                raise compute_task_errors.CeleryNoRetryError() from e
+            except OrganizationError as e:
+                raise compute_task_errors.CeleryRetryError() from e
+
+        # save entrypoint to DB
+        entrypoint = get_entrypoint(container_image_tag)
+        ImageEntrypoint.objects.get_or_create(
+            archive_checksum=ctx.function.archive_address.checksum, entrypoint_json=entrypoint
+        )
         volume_mounts, volumes = get_volumes(ctx)
 
         with get_orchestrator_client(channel_name) as client:

backend/substrapp/tasks/tasks_compute_task.py

@@ -29,7 +29,6 @@
 from substrapp.clients import organization as organization_client
 from substrapp.compute_tasks import compute_task as task_utils
 from substrapp.compute_tasks import errors as compute_task_errors
-from substrapp.compute_tasks import image_builder
 from substrapp.compute_tasks.asset_buffer import add_assets_to_taskdir
 from substrapp.compute_tasks.asset_buffer import add_task_assets_to_buffer
 from substrapp.compute_tasks.asset_buffer import clear_assets_buffer
@@ -48,13 +47,11 @@
 from substrapp.compute_tasks.lock import MAX_TASK_DURATION
 from substrapp.compute_tasks.lock import acquire_compute_plan_lock
 from substrapp.compute_tasks.outputs import OutputSaver
-from substrapp.exceptions import OrganizationError
 from substrapp.exceptions import OrganizationHttpError
 from substrapp.lock_local import lock_resource
 from substrapp.orchestrator import get_orchestrator_client
 from substrapp.tasks.task import ComputeTask
 from substrapp.utils import Timer
-from substrapp.utils import get_owner
 from substrapp.utils import list_dir
 from substrapp.utils import retry
 from substrapp.utils.url import TASK_PROFILING_BASE_URL
@@ -198,14 +195,6 @@ def _run(
         # start build_image timer
         timer.start()
 
-        if get_owner() != ctx.function.owner:
-            try:
-                image_builder.load_remote_function_image(ctx.function, channel_name)
-            except OrganizationHttpError as e:
-                raise compute_task_errors.CeleryNoRetryError() from e
-            except OrganizationError as e:
-                raise compute_task_errors.CeleryRetryError() from e
-
         # stop build_image timer
         _create_task_profiling_step(channel_name, task.key, ComputeTaskSteps.BUILD_IMAGE, timer.stop())
 

backend/substrapp/tests/tasks/test_compute_task.py

@@ -29,7 +29,6 @@ def test_compute_task_exception(mocker: MockerFixture):
     mock_init_task_dirs = mocker.patch("substrapp.tasks.tasks_compute_task.init_task_dirs")
     mock_add_asset_to_buffer = mocker.patch("substrapp.tasks.tasks_compute_task.add_task_assets_to_buffer")
     mock_add_asset_to_task_dir = mocker.patch("substrapp.tasks.tasks_compute_task.add_assets_to_taskdir")
-    mock_load_remote_function_image = mocker.patch("substrapp.compute_tasks.image_builder.load_remote_function_image")
     mock_execute_compute_task = mocker.patch("substrapp.tasks.tasks_compute_task.execute_compute_task")
     saver = mocker.MagicMock()
     mock_output_saver = mocker.patch("substrapp.tasks.tasks_compute_task.OutputSaver", return_value=saver)
@@ -64,7 +63,6 @@ class FakeDirectories:
     mock_init_task_dirs.assert_called_once()
     mock_add_asset_to_buffer.assert_called_once()
     mock_add_asset_to_task_dir.assert_called_once()
-    mock_load_remote_function_image.assert_called_once()
     mock_execute_compute_task.assert_called_once()
     saver.save_outputs.assert_called_once()
     mock_output_saver.assert_called_once()