✨ Add Celery routing queues (#7471)

giancarloromeo · sanderegg · web-flow · commit 08ce0fe6d867 · 2025-04-09T23:19:16.000+02:00
Co-authored-by: sanderegg &lt;35365065+sanderegg@users.noreply.github.com&gt;
diff --git a/packages/pytest-simcore/src/pytest_simcore/simcore_services.py b/packages/pytest-simcore/src/pytest_simcore/simcore_services.py
@@ -6,9 +6,9 @@
 import json
 import logging
 import warnings
+from collections.abc import Iterator
 from dataclasses import dataclass
 from io import StringIO
-from typing import Iterator
 
 import aiohttp
 import pytest
@@ -38,6 +38,7 @@
     "traefik",
     "whoami",
     "sto-worker",
+    "sto-worker-cpu-bound",
 }
 # TODO: unify healthcheck policies see  https://github.com/ITISFoundation/osparc-simcore/pull/2281
 SERVICE_PUBLISHED_PORT = {}
diff --git a/packages/settings-library/src/settings_library/celery.py b/packages/settings-library/src/settings_library/celery.py
@@ -22,6 +22,12 @@ class CelerySettings(BaseCustomSettings):
             description="Time after which task results will be deleted (default to seconds, or see https://pydantic-docs.helpmanual.io/usage/types/#datetime-types for string formating)."
         ),
     ] = timedelta(days=7)
+    CELERY_EPHEMERAL_RESULT_EXPIRES: Annotated[
+        timedelta,
+        Field(
+            description="Time after which ephemeral task results will be deleted (default to seconds, or see https://pydantic-docs.helpmanual.io/usage/types/#datetime-types for string formating)."
+        ),
+    ] = timedelta(hours=1)
     CELERY_RESULT_PERSISTENT: Annotated[
         bool,
         Field(
diff --git a/services/docker-compose.devel.yml b/services/docker-compose.devel.yml
@@ -215,6 +215,16 @@ services:
       STORAGE_PROFILING : ${STORAGE_PROFILING}
       STORAGE_LOGLEVEL: DEBUG
 
+  sto-worker-cpu-bound:
+    volumes:
+      - ./storage:/devel/services/storage
+      - ../packages:/devel/packages
+      - ${HOST_UV_CACHE_DIR}:/home/scu/.cache/uv
+    environment:
+      <<: *common-environment
+      STORAGE_PROFILING : ${STORAGE_PROFILING}
+      STORAGE_LOGLEVEL: DEBUG
+
   agent:
     environment:
       <<: *common-environment
diff --git a/services/docker-compose.local.yml b/services/docker-compose.local.yml
@@ -133,6 +133,14 @@ services:
     ports:
       - "8080"
       - "3021:3000"
+
+  sto-worker-cpu-bound:
+    environment:
+      <<: *common_environment
+      STORAGE_REMOTE_DEBUGGING_PORT : 3000
+    ports:
+      - "8080"
+      - "3022:3000"
   webserver:
     environment: &webserver_environment_local
       <<: *common_environment
diff --git a/services/docker-compose.yml b/services/docker-compose.yml
@@ -1193,10 +1193,21 @@ services:
     image: ${DOCKER_REGISTRY:-itisfoundation}/storage:${DOCKER_IMAGE_TAG:-master-github-latest}
     init: true
     hostname: "sto-worker-{{.Node.Hostname}}-{{.Task.Slot}}"
+    environment:
+      <<: *storage_environment
+      STORAGE_WORKER_MODE: "true"
+      CELERY_CONCURRENCY: 100
+    networks: *storage_networks
+
+  sto-worker-cpu-bound:
+    image: ${DOCKER_REGISTRY:-itisfoundation}/storage:${DOCKER_IMAGE_TAG:-master-github-latest}
+    init: true
+    hostname: "sto-worker-cpu-bound-{{.Node.Hostname}}-{{.Task.Slot}}"
     environment:
       <<: *storage_environment
       STORAGE_WORKER_MODE: "true"
       CELERY_CONCURRENCY: 1
+      CELERY_QUEUES: "cpu-bound"
     networks: *storage_networks
 
   rabbit:
diff --git a/services/storage/docker/boot.sh b/services/storage/docker/boot.sh
@@ -54,13 +54,15 @@ if [ "${STORAGE_WORKER_MODE}" = "true" ]; then
       --app=simcore_service_storage.modules.celery.worker_main:app \
       worker --pool=threads \
       --loglevel="${SERVER_LOG_LEVEL}" \
-      --concurrency="${CELERY_CONCURRENCY}"
+      --concurrency="${CELERY_CONCURRENCY}" \
+      --queues="${CELERY_QUEUES:-default}"
   else
     exec celery \
       --app=simcore_service_storage.modules.celery.worker_main:app \
       worker --pool=threads \
       --loglevel="${SERVER_LOG_LEVEL}" \
-      --concurrency="${CELERY_CONCURRENCY}"
+      --concurrency="${CELERY_CONCURRENCY}" \
+      --queues="${CELERY_QUEUES:-default}"
   fi
 else
   if [ "${SC_BOOT_MODE}" = "debug" ]; then
diff --git a/services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py b/services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py
@@ -21,7 +21,10 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ...modules.celery import get_celery_client
-from ...modules.celery.errors import decode_celery_transferrable_error
+from ...modules.celery.errors import (
+    TransferrableCeleryError,
+    decode_celery_transferrable_error,
+)
 from ...modules.celery.models import TaskState
 
 _logger = logging.getLogger(__name__)
@@ -102,6 +105,7 @@ async def result(
         # try to recover the original error
         exception = None
         with log_catch(_logger, reraise=False):
+            assert isinstance(_result, TransferrableCeleryError)  # nosec
             exception = decode_celery_transferrable_error(_result)
             exc_type = type(exception).__name__
             exc_msg = f"{exception}"
diff --git a/services/storage/src/simcore_service_storage/api/rpc/_simcore_s3.py b/services/storage/src/simcore_service_storage/api/rpc/_simcore_s3.py
@@ -8,6 +8,7 @@
 from servicelib.rabbitmq import RPCRouter
 
 from ...modules.celery import get_celery_client
+from ...modules.celery.models import TaskMetadata, TasksQueue
 from .._worker_tasks._simcore_s3 import deep_copy_files_from_project, export_data
 
 router = RPCRouter()
@@ -36,6 +37,10 @@ async def start_export_data(
     task_uuid = await get_celery_client(app).send_task(
         export_data.__name__,
         task_context=job_id_data.model_dump(),
+        task_metadata=TaskMetadata(
+            ephemeral=False,
+            queue=TasksQueue.CPU_BOUND,
+        ),
         user_id=job_id_data.user_id,
         paths_to_export=paths_to_export,
     )
diff --git a/services/storage/src/simcore_service_storage/modules/celery/__init__.py b/services/storage/src/simcore_service_storage/modules/celery/__init__.py
@@ -9,7 +9,7 @@
 from ...core.settings import get_application_settings
 from ._celery_types import register_celery_types
 from ._common import create_app
-from .backends._redis import RedisTaskStore
+from .backends._redis import RedisTaskMetadataStore
 from .client import CeleryTaskQueueClient
 
 _logger = logging.getLogger(__name__)
@@ -29,7 +29,9 @@ async def on_startup() -> None:
         )
 
         app.state.celery_client = CeleryTaskQueueClient(
-            celery_app, RedisTaskStore(redis_client_sdk)
+            celery_app,
+            celery_settings,
+            RedisTaskMetadataStore(redis_client_sdk),
         )
 
         register_celery_types()
diff --git a/services/storage/src/simcore_service_storage/modules/celery/_common.py b/services/storage/src/simcore_service_storage/modules/celery/_common.py
@@ -15,6 +15,7 @@ def _celery_configure(celery_settings: CelerySettings) -> dict[str, Any]:
         "result_expires": celery_settings.CELERY_RESULT_EXPIRES,
         "result_extended": True,
         "result_serializer": "json",
+        "task_default_queue": "default",
         "task_send_sent_event": True,
         "task_track_started": True,
         "worker_send_task_events": True,
diff --git a/services/storage/src/simcore_service_storage/modules/celery/backends/_redis.py b/services/storage/src/simcore_service_storage/modules/celery/backends/_redis.py
@@ -1,21 +1,43 @@
+import logging
+from datetime import timedelta
 from typing import Final
 
+from celery.result import AsyncResult  # type: ignore[import-untyped]
 from servicelib.redis._client import RedisClientSDK
 
-from ..models import TaskContext, TaskData, TaskID, TaskUUID, build_task_id_prefix
+from ..models import TaskContext, TaskID, TaskMetadata, TaskUUID, build_task_id_prefix
 
-_CELERY_TASK_META_PREFIX: Final[str] = "celery-task-meta-"
+_CELERY_TASK_METADATA_PREFIX: Final[str] = "celery-task-metadata-"
+_CELERY_TASK_ID_KEY_ENCODING = "utf-8"
 _CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
 _CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 10000
-_CELERY_TASK_ID_KEY_ENCODING = "utf-8"
+
+_logger = logging.getLogger(__name__)
 
 
-class RedisTaskStore:
+def _build_key(task_id: TaskID) -> str:
+    return _CELERY_TASK_METADATA_PREFIX + task_id
+
+
+class RedisTaskMetadataStore:
     def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
         self._redis_client_sdk = redis_client_sdk
 
-    async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
-        search_key = build_task_id_prefix(task_context) + _CELERY_TASK_ID_KEY_SEPARATOR
+    async def exists(self, task_id: TaskID) -> bool:
+        n = await self._redis_client_sdk.redis.exists(_build_key(task_id))
+        assert isinstance(n, int)  # nosec
+        return n > 0
+
+    async def get(self, task_id: TaskID) -> TaskMetadata | None:
+        result = await self._redis_client_sdk.redis.get(_build_key(task_id))
+        return TaskMetadata.model_validate_json(result) if result else None
+
+    async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
+        search_key = (
+            _CELERY_TASK_METADATA_PREFIX
+            + build_task_id_prefix(task_context)
+            + _CELERY_TASK_ID_KEY_SEPARATOR
+        )
         keys = set()
         async for key in self._redis_client_sdk.redis.scan_iter(
             match=search_key + "*", count=_CELERY_TASK_SCAN_COUNT_PER_BATCH
@@ -29,13 +51,15 @@ async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
             keys.add(TaskUUID(_key.removeprefix(search_key)))
         return keys
 
-    async def task_exists(self, task_id: TaskID) -> bool:
-        n = await self._redis_client_sdk.redis.exists(task_id)
-        assert isinstance(n, int)  # nosec
-        return n > 0
+    async def remove(self, task_id: TaskID) -> None:
+        await self._redis_client_sdk.redis.delete(_build_key(task_id))
+        AsyncResult(task_id).forget()
 
-    async def set_task(self, task_id: TaskID, task_data: TaskData) -> None:
+    async def set(
+        self, task_id: TaskID, task_metadata: TaskMetadata, expiry: timedelta
+    ) -> None:
         await self._redis_client_sdk.redis.set(
-            task_id,
-            task_data.model_dump_json(),
+            _build_key(task_id),
+            task_metadata.model_dump_json(),
+            ex=expiry,
         )
diff --git a/services/storage/src/simcore_service_storage/modules/celery/client.py b/services/storage/src/simcore_service_storage/modules/celery/client.py
@@ -12,13 +12,15 @@
 from models_library.progress_bar import ProgressReport
 from pydantic import ValidationError
 from servicelib.logging_utils import log_context
+from settings_library.celery import CelerySettings
 
 from .models import (
     TaskContext,
-    TaskData,
+    TaskID,
+    TaskMetadata,
+    TaskMetadataStore,
     TaskState,
     TaskStatus,
-    TaskStore,
     TaskUUID,
     build_task_id,
 )
@@ -43,10 +45,16 @@
 @dataclass
 class CeleryTaskQueueClient:
     _celery_app: Celery
-    _task_store: TaskStore
+    _celery_settings: CelerySettings
+    _task_store: TaskMetadataStore
 
     async def send_task(
-        self, task_name: str, *, task_context: TaskContext, **task_params
+        self,
+        task_name: str,
+        *,
+        task_context: TaskContext,
+        task_metadata: TaskMetadata | None = None,
+        **task_params,
     ) -> TaskUUID:
         with log_context(
             _logger,
@@ -55,39 +63,59 @@ async def send_task(
         ):
             task_uuid = uuid4()
             task_id = build_task_id(task_context, task_uuid)
-            self._celery_app.send_task(task_name, task_id=task_id, kwargs=task_params)
-            await self._task_store.set_task(
-                task_id, TaskData(status=TaskState.PENDING.name)
+            task_metadata = task_metadata or TaskMetadata()
+            self._celery_app.send_task(
+                task_name,
+                task_id=task_id,
+                kwargs=task_params,
+                queue=task_metadata.queue.value,
             )
+
+            expiry = (
+                self._celery_settings.CELERY_EPHEMERAL_RESULT_EXPIRES
+                if task_metadata.ephemeral
+                else self._celery_settings.CELERY_RESULT_EXPIRES
+            )
+            await self._task_store.set(task_id, task_metadata, expiry=expiry)
             return task_uuid
 
-    @staticmethod
     @make_async()
-    def abort_task(task_context: TaskContext, task_uuid: TaskUUID) -> None:
+    def _abort_task(self, task_id: TaskID) -> None:
+        AbortableAsyncResult(task_id, app=self._celery_app).abort()
+
+    async def abort_task(self, task_context: TaskContext, task_uuid: TaskUUID) -> None:
         with log_context(
             _logger,
             logging.DEBUG,
             msg=f"Abort task: {task_context=} {task_uuid=}",
         ):
             task_id = build_task_id(task_context, task_uuid)
-            AbortableAsyncResult(task_id).abort()
+            return await self._abort_task(task_id)
 
     @make_async()
-    def get_task_result(self, task_context: TaskContext, task_uuid: TaskUUID) -> Any:
+    def _get_result(self, task_context: TaskContext, task_uuid: TaskUUID) -> Any:
+        task_id = build_task_id(task_context, task_uuid)
+        return self._celery_app.AsyncResult(task_id).result
+
+    async def get_task_result(
+        self, task_context: TaskContext, task_uuid: TaskUUID
+    ) -> Any:
         with log_context(
             _logger,
             logging.DEBUG,
             msg=f"Get task result: {task_context=} {task_uuid=}",
         ):
             task_id = build_task_id(task_context, task_uuid)
-            return self._celery_app.AsyncResult(task_id).result
+            async_result = self._celery_app.AsyncResult(task_id)
+            result = async_result.result
+            if async_result.ready():
+                task_metadata = await self._task_store.get(task_id)
+                if task_metadata is not None and task_metadata.ephemeral:
+                    await self._task_store.remove(task_id)
+            return result
 
-    def _get_progress_report(
-        self, task_context: TaskContext, task_uuid: TaskUUID
-    ) -> ProgressReport:
-        task_id = build_task_id(task_context, task_uuid)
-        result = self._celery_app.AsyncResult(task_id).result
-        state = self._get_state(task_context, task_uuid)
+    @staticmethod
+    async def _get_progress_report(state, result) -> ProgressReport:
         if result and state == TaskState.RUNNING:
             with contextlib.suppress(ValidationError):
                 # avoids exception if result is not a ProgressReport (or overwritten by a Celery's state update)
@@ -104,23 +132,25 @@ def _get_progress_report(
             actual_value=_MIN_PROGRESS_VALUE, total=_MAX_PROGRESS_VALUE
         )
 
+    @make_async()
     def _get_state(self, task_context: TaskContext, task_uuid: TaskUUID) -> TaskState:
         task_id = build_task_id(task_context, task_uuid)
         return _CELERY_STATES_MAPPING[self._celery_app.AsyncResult(task_id).state]
 
-    @make_async()
-    def get_task_status(
+    async def get_task_status(
         self, task_context: TaskContext, task_uuid: TaskUUID
     ) -> TaskStatus:
         with log_context(
             _logger,
             logging.DEBUG,
             msg=f"Getting task status: {task_context=} {task_uuid=}",
         ):
+            task_state = await self._get_state(task_context, task_uuid)
+            result = await self._get_result(task_context, task_uuid)
             return TaskStatus(
                 task_uuid=task_uuid,
-                task_state=self._get_state(task_context, task_uuid),
-                progress_report=self._get_progress_report(task_context, task_uuid),
+                task_state=task_state,
+                progress_report=await self._get_progress_report(task_state, result),
             )
 
     async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
@@ -129,4 +159,4 @@ async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:
             logging.DEBUG,
             msg=f"Getting task uuids: {task_context=}",
         ):
-            return await self._task_store.get_task_uuids(task_context)
+            return await self._task_store.get_uuids(task_context)
diff --git a/services/storage/src/simcore_service_storage/modules/celery/models.py b/services/storage/src/simcore_service_storage/modules/celery/models.py
diff --git a/services/storage/tests/conftest.py b/services/storage/tests/conftest.py
diff --git a/services/web/server/tests/integration/conftest.py b/services/web/server/tests/integration/conftest.py
