♻️Storage Celery integration: diverse enhancements (#7353)

Author: sanderegg

services/docker-compose.devel.yml

@@ -189,6 +189,15 @@ services:
       STORAGE_PROFILING : ${STORAGE_PROFILING}
       STORAGE_LOGLEVEL: DEBUG
 
+  sto-worker:
+    volumes:
+      - ./storage:/devel/services/storage
+      - ../packages:/devel/packages
+    environment:
+      <<: *common-environment
+      STORAGE_PROFILING : ${STORAGE_PROFILING}
+      STORAGE_LOGLEVEL: DEBUG
+
   agent:
     environment:
       <<: *common-environment

services/docker-compose.local.yml

@@ -126,6 +126,13 @@ services:
       - "8080"
       - "3003:3000"
 
+  sto-worker:
+    environment:
+      <<: *common_environment
+      STORAGE_REMOTE_DEBUGGING_PORT : 3000
+    ports:
+      - "8080"
+      - "3021:3000"
   webserver:
     environment: &webserver_environment_local
       <<: *common_environment

services/storage/docker/boot.sh

@@ -47,7 +47,7 @@ if [ "${SC_BOOT_MODE}" = "debug" ]; then
 
   exec sh -c "
     cd services/storage/src/simcore_service_storage && \
-    python -m debugpy --listen 0.0.0.0:${STORAGE_REMOTE_DEBUGGING_PORT} -m uvicorn main:the_app \
+    python -m debugpy --listen 0.0.0.0:${STORAGE_REMOTE_DEBUGGING_PORT} -m uvicorn main:app \
       --host 0.0.0.0 \
       --port ${STORAGE_PORT} \
       --reload \

services/storage/src/simcore_service_storage/_meta.py

@@ -28,4 +28,17 @@
     f"v{__version__}"
 )
 
+APP_WORKER_STARTED_BANNER_MSG = r"""
+
+  ____  _                                __        __         _
+ / ___|| |_ ___  _ __ __ _  __ _  ___    \ \      / /__  _ __| | _____ _ __
+ \___ \| __/ _ \| '__/ _` |/ _` |/ _ \____\ \ /\ / / _ \| '__| |/ / _ \ '__|
+  ___) | || (_) | | | (_| | (_| |  __/_____\ V  V / (_) | |  |   <  __/ |
+ |____/ \__\___/|_|  \__,_|\__, |\___|      \_/\_/ \___/|_|  |_|\_\___|_|
+                           |___/                                                {}
+
+""".format(
+    f"v{__version__}"
+)
+
 APP_FINISHED_BANNER_MSG = info.get_finished_banner()

services/storage/src/simcore_service_storage/core/application.py

@@ -26,6 +26,7 @@
     APP_FINISHED_BANNER_MSG,
     APP_NAME,
     APP_STARTED_BANNER_MSG,
+    APP_WORKER_STARTED_BANNER_MSG,
 )
 from ..api.rest.routes import setup_rest_api_routes
 from ..api.rpc.routes import setup_rpc_api_routes
@@ -114,7 +115,10 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
         setup_prometheus_instrumentation(app)
 
     async def _on_startup() -> None:
-        print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
+        if settings.STORAGE_WORKER_MODE:
+            print(APP_WORKER_STARTED_BANNER_MSG, flush=True)  # noqa: T201
+        else:
+            print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
 
     async def _on_shutdown() -> None:
         print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201

services/storage/src/simcore_service_storage/main.py

@@ -3,11 +3,10 @@
 import logging
 
 from servicelib.logging_utils import config_all_loggers
+from simcore_service_storage.core.application import create_app
+from simcore_service_storage.core.settings import ApplicationSettings
 from simcore_service_storage.modules.celery import setup_celery
 
-from .core.application import create_app
-from .core.settings import ApplicationSettings
-
 _settings = ApplicationSettings.create_from_envs()
 
 # SEE https://github.com/ITISFoundation/osparc-simcore/issues/3148

services/storage/src/simcore_service_storage/modules/celery/__init__.py

@@ -2,10 +2,10 @@
 from asyncio import AbstractEventLoop
 
 from fastapi import FastAPI
-from simcore_service_storage.modules.celery._common import create_app
-from simcore_service_storage.modules.celery.client import CeleryTaskQueueClient
 
 from ...core.settings import get_application_settings
+from ._common import create_app
+from .client import CeleryTaskQueueClient
 
 _logger = logging.getLogger(__name__)
 

services/storage/src/simcore_service_storage/modules/celery/_common.py

@@ -1,12 +1,12 @@
-from collections.abc import Callable
-from functools import wraps
 import logging
 import traceback
+from collections.abc import Callable
+from functools import wraps
 from typing import Any
 
-from celery import Celery, Task # type: ignore[import-untyped]
-from celery.exceptions import Ignore    # type: ignore[import-untyped]
+from celery import Celery, Task  # type: ignore[import-untyped]
 from celery.contrib.abortable import AbortableTask  # type: ignore[import-untyped]
+from celery.exceptions import Ignore  # type: ignore[import-untyped]
 from settings_library.celery import CelerySettings
 from settings_library.redis import RedisDatabase
 
@@ -39,7 +39,7 @@ def wrapper(task: Task, *args: Any, **kwargs: Any) -> Any:
         except Exception as exc:
             exc_type = type(exc).__name__
             exc_message = f"{exc}"
-            exc_traceback = traceback.format_exc().split('\n')
+            exc_traceback = traceback.format_exc().split("\n")
 
             _logger.exception(
                 "Task %s failed with exception: %s",
@@ -53,9 +53,10 @@ def wrapper(task: Task, *args: Any, **kwargs: Any) -> Any:
                     exc_type=exc_type,
                     exc_msg=exc_message,
                 ).model_dump(mode="json"),
-                traceback=exc_traceback
+                traceback=exc_traceback,
             )
-            raise Ignore from exc   # ignore doing state updates
+            raise Ignore from exc  # ignore doing state updates
+
     return wrapper
 
 

services/storage/src/simcore_service_storage/modules/celery/client.py

@@ -3,8 +3,10 @@
 from typing import Any, Final
 from uuid import uuid4
 
-from celery import Celery   # type: ignore[import-untyped]
-from celery.contrib.abortable import AbortableAsyncResult   # type: ignore[import-untyped]
+from celery import Celery  # type: ignore[import-untyped]
+from celery.contrib.abortable import (  # type: ignore[import-untyped]
+    AbortableAsyncResult,
+)
 from common_library.async_tools import make_async
 from models_library.progress_bar import ProgressReport
 from pydantic import ValidationError
@@ -47,7 +49,9 @@ def _build_task_id_prefix(task_context: TaskContext) -> str:
 
 
 def _build_task_id(task_context: TaskContext, task_uuid: TaskUUID) -> TaskID:
-    return _CELERY_TASK_ID_KEY_SEPARATOR.join([_build_task_id_prefix(task_context), f"{task_uuid}"])
+    return _CELERY_TASK_ID_KEY_SEPARATOR.join(
+        [_build_task_id_prefix(task_context), f"{task_uuid}"]
+    )
 
 
 class CeleryTaskQueueClient:
@@ -114,13 +118,13 @@ def get_task_status(
         )
 
     def _get_completed_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
-        search_key = (
-            _CELERY_TASK_META_PREFIX + _build_task_id_prefix(task_context)
-        )
+        search_key = _CELERY_TASK_META_PREFIX + _build_task_id_prefix(task_context)
         redis = self._celery_app.backend.client
         if hasattr(redis, "keys") and (keys := redis.keys(search_key + "*")):
             return {
-                TaskUUID(f"{key.decode(_CELERY_TASK_ID_KEY_ENCODING).removeprefix(search_key + _CELERY_TASK_ID_KEY_SEPARATOR)}")
+                TaskUUID(
+                    f"{key.decode(_CELERY_TASK_ID_KEY_ENCODING).removeprefix(search_key + _CELERY_TASK_ID_KEY_SEPARATOR)}"
+                )
                 for key in keys
             }
         return set()
@@ -129,15 +133,19 @@ def _get_completed_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
     def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
         all_task_ids = self._get_completed_task_uuids(task_context)
 
-        search_key = (
-            _CELERY_TASK_META_PREFIX + _build_task_id_prefix(task_context)
-        )
+        search_key = _CELERY_TASK_META_PREFIX + _build_task_id_prefix(task_context)
         for task_inspect_status in _CELERY_INSPECT_TASK_STATUSES:
             if task_ids := getattr(
                 self._celery_app.control.inspect(), task_inspect_status
             )():
                 for values in task_ids.values():
                     for value in values:
-                        all_task_ids.add(TaskUUID(value.removeprefix(search_key + _CELERY_TASK_ID_KEY_SEPARATOR)))
+                        all_task_ids.add(
+                            TaskUUID(
+                                value.removeprefix(
+                                    search_key + _CELERY_TASK_ID_KEY_SEPARATOR
+                                )
+                            )
+                        )
 
         return all_task_ids

services/storage/src/simcore_service_storage/modules/celery/signals.py

@@ -4,9 +4,10 @@
 from typing import Final
 
 from asgi_lifespan import LifespanManager
-from celery import Celery   # type: ignore[import-untyped]
+from celery import Celery  # type: ignore[import-untyped]
 from fastapi import FastAPI
 from servicelib.async_utils import cancel_wait_task
+
 from ...core.application import create_app
 from ...core.settings import ApplicationSettings
 from ...modules.celery import get_event_loop, set_event_loop