✨ Add Celery task manager to Web Server ⚠️ (#8436)

On behalf of @giancarloromeo

Author: giancarloromeo

api/specs/web-server/_long_running_tasks.py

@@ -12,27 +12,24 @@
 from servicelib.aiohttp.long_running_tasks._routes import _PathParam
 from servicelib.long_running_tasks.models import TaskGet, TaskStatus
 from simcore_service_webserver._meta import API_VTAG
-from simcore_service_webserver.tasks._exception_handlers import (
-    _TO_HTTP_ERROR_MAP as export_data_http_error_map,
+from simcore_service_webserver.tasks._controller._rest_exceptions import (
+    _TO_HTTP_ERROR_MAP,
 )
 
 router = APIRouter(
     prefix=f"/{API_VTAG}",
     tags=[
         "long-running-tasks",
     ],
+    responses={
+        i.status_code: {"model": EnvelopedError} for i in _TO_HTTP_ERROR_MAP.values()
+    },
 )
 
-_export_data_responses: dict[int | str, dict[str, Any]] = {
-    i.status_code: {"model": EnvelopedError}
-    for i in export_data_http_error_map.values()
-}
-
 
 @router.get(
     "/tasks",
     response_model=Envelope[list[TaskGet]],
-    responses=_export_data_responses,
 )
 def get_async_jobs():
     """Lists all long running tasks"""
@@ -41,7 +38,6 @@ def get_async_jobs():
 @router.get(
     "/tasks/{task_id}",
     response_model=Envelope[TaskStatus],
-    responses=_export_data_responses,
 )
 def get_async_job_status(
     _path_params: Annotated[_PathParam, Depends()],
@@ -51,7 +47,6 @@ def get_async_job_status(
 
 @router.delete(
     "/tasks/{task_id}",
-    responses=_export_data_responses,
     status_code=status.HTTP_204_NO_CONTENT,
 )
 def cancel_async_job(
@@ -63,7 +58,6 @@ def cancel_async_job(
 @router.get(
     "/tasks/{task_id}/result",
     response_model=Any,
-    responses=_export_data_responses,
 )
 def get_async_job_result(
     _path_params: Annotated[_PathParam, Depends()],

api/specs/web-server/_storage.py

@@ -4,7 +4,7 @@
 # pylint: disable=too-many-arguments
 
 
-from typing import Annotated, Any, TypeAlias
+from typing import Annotated, TypeAlias
 
 from fastapi import APIRouter, Depends, Query, status
 from models_library.api_schemas_long_running_tasks.tasks import (
@@ -35,8 +35,8 @@
 from servicelib.fastapi.rest_pagination import CustomizedPathsCursorPage
 from simcore_service_webserver._meta import API_VTAG
 from simcore_service_webserver.storage.schemas import DatasetMetaData, FileMetaData
-from simcore_service_webserver.tasks._exception_handlers import (
-    _TO_HTTP_ERROR_MAP as export_data_http_error_map,
+from simcore_service_webserver.tasks._controller._rest_exceptions import (
+    _TO_HTTP_ERROR_MAP,
 )
 
 router = APIRouter(
@@ -220,19 +220,14 @@ async def is_completed_upload_file(
     """Returns state of upload completion"""
 
 
-# data export
-_export_data_responses: dict[int | str, dict[str, Any]] = {
-    i.status_code: {"model": EnvelopedError}
-    for i in export_data_http_error_map.values()
-}
-
-
 @router.post(
     "/storage/locations/{location_id}/export-data",
     response_model=Envelope[TaskGet],
     name="export_data",
     description="Export data",
-    responses=_export_data_responses,
+    responses={
+        i.status_code: {"model": EnvelopedError} for i in _TO_HTTP_ERROR_MAP.values()
+    },
 )
 async def export_data(export_data: DataExportPost, location_id: LocationID):
     """Trigger data export. Returns async job id for getting status and results"""

packages/celery-library/src/celery_library/backends/redis.py

@@ -1,5 +1,6 @@
 import contextlib
 import logging
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import TYPE_CHECKING, Final
 
@@ -10,35 +11,36 @@
     ExecutionMetadata,
     OwnerMetadata,
     Task,
-    TaskInfoStore,
     TaskKey,
+    TaskStore,
 )
 from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
-_CELERY_TASK_INFO_PREFIX: Final[str] = "celery-task-info-"
+_CELERY_TASK_PREFIX: Final[str] = "celery-task-"
 _CELERY_TASK_ID_KEY_ENCODING = "utf-8"
 _CELERY_TASK_SCAN_COUNT_PER_BATCH: Final[int] = 1000
 _CELERY_TASK_METADATA_KEY: Final[str] = "metadata"
 _CELERY_TASK_PROGRESS_KEY: Final[str] = "progress"
 
+
 _logger = logging.getLogger(__name__)
 
 
-def _build_key(task_key: TaskKey) -> str:
-    return _CELERY_TASK_INFO_PREFIX + task_key
+def _build_redis_task_key(task_key: TaskKey) -> str:
+    return _CELERY_TASK_PREFIX + task_key
 
 
-class RedisTaskInfoStore:
-    def __init__(self, redis_client_sdk: RedisClientSDK) -> None:
-        self._redis_client_sdk = redis_client_sdk
+@dataclass(frozen=True)
+class RedisTaskStore:
+    _redis_client_sdk: RedisClientSDK
 
     async def create_task(
         self,
         task_key: TaskKey,
         execution_metadata: ExecutionMetadata,
         expiry: timedelta,
     ) -> None:
-        redis_key = _build_key(task_key)
+        redis_key = _build_redis_task_key(task_key)
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
                 name=redis_key,
@@ -54,7 +56,8 @@ async def create_task(
     async def get_task_metadata(self, task_key: TaskKey) -> ExecutionMetadata | None:
         raw_result = await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hget(
-                _build_key(task_key), _CELERY_TASK_METADATA_KEY
+                _build_redis_task_key(task_key),
+                _CELERY_TASK_METADATA_KEY,
             )
         )
         if not raw_result:
@@ -73,7 +76,8 @@ async def get_task_metadata(self, task_key: TaskKey) -> ExecutionMetadata | None
     async def get_task_progress(self, task_key: TaskKey) -> ProgressReport | None:
         raw_result = await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hget(
-                _build_key(task_key), _CELERY_TASK_PROGRESS_KEY
+                _build_redis_task_key(task_key),
+                _CELERY_TASK_PROGRESS_KEY,
             )
         )
         if not raw_result:
@@ -90,7 +94,7 @@ async def get_task_progress(self, task_key: TaskKey) -> ProgressReport | None:
             return None
 
     async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
-        search_key = _CELERY_TASK_INFO_PREFIX + owner_metadata.model_dump_task_key(
+        search_key = _CELERY_TASK_PREFIX + owner_metadata.model_dump_task_key(
             task_uuid=WILDCARD
         )
 
@@ -127,24 +131,28 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         return tasks
 
     async def remove_task(self, task_key: TaskKey) -> None:
-        await self._redis_client_sdk.redis.delete(_build_key(task_key))
+        await self._redis_client_sdk.redis.delete(
+            _build_redis_task_key(task_key),
+        )
 
     async def set_task_progress(
         self, task_key: TaskKey, report: ProgressReport
     ) -> None:
         await handle_redis_returns_union_types(
             self._redis_client_sdk.redis.hset(
-                name=_build_key(task_key),
+                name=_build_redis_task_key(task_key),
                 key=_CELERY_TASK_PROGRESS_KEY,
                 value=report.model_dump_json(),
             )
         )
 
     async def task_exists(self, task_key: TaskKey) -> bool:
-        n = await self._redis_client_sdk.redis.exists(_build_key(task_key))
+        n = await self._redis_client_sdk.redis.exists(
+            _build_redis_task_key(task_key),
+        )
         assert isinstance(n, int)  # nosec
         return n > 0
 
 
 if TYPE_CHECKING:
-    _: type[TaskInfoStore] = RedisTaskInfoStore
+    _: type[TaskStore] = RedisTaskStore

packages/celery-library/src/celery_library/errors.py

@@ -1,6 +1,8 @@
 import base64
 import pickle
+from functools import wraps
 
+from celery.exceptions import CeleryError  # type: ignore[import-untyped]
 from common_library.errors_classes import OsparcErrorMixin
 
 
@@ -32,3 +34,18 @@ class TaskSubmissionError(OsparcErrorMixin, Exception):
 
 class TaskNotFoundError(OsparcErrorMixin, Exception):
     msg_template = "Task with uuid '{task_uuid}' and owner_metadata '{owner_metadata}' was not found"
+
+
+class TaskManagerError(OsparcErrorMixin, Exception):
+    msg_template = "An internal error occurred"
+
+
+def handle_celery_errors(func):
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except CeleryError as exc:
+            raise TaskManagerError from exc
+
+    return wrapper

packages/celery-library/src/celery_library/task_manager.py

@@ -12,17 +12,17 @@
     ExecutionMetadata,
     OwnerMetadata,
     Task,
-    TaskInfoStore,
     TaskKey,
     TaskState,
     TaskStatus,
+    TaskStore,
     TaskUUID,
 )
 from servicelib.celery.task_manager import TaskManager
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
-from .errors import TaskNotFoundError, TaskSubmissionError
+from .errors import TaskNotFoundError, TaskSubmissionError, handle_celery_errors
 
 _logger = logging.getLogger(__name__)
 
@@ -35,8 +35,9 @@
 class CeleryTaskManager:
     _celery_app: Celery
     _celery_settings: CelerySettings
-    _task_info_store: TaskInfoStore
+    _task_info_store: TaskStore
 
+    @handle_celery_errors
     async def submit_task(
         self,
         execution_metadata: ExecutionMetadata,
@@ -85,6 +86,7 @@ async def submit_task(
 
             return task_uuid
 
+    @handle_celery_errors
     async def cancel_task(
         self, owner_metadata: OwnerMetadata, task_uuid: TaskUUID
     ) -> None:
@@ -109,6 +111,7 @@ async def task_exists(self, task_key: TaskKey) -> bool:
     def _forget_task(self, task_key: TaskKey) -> None:
         self._celery_app.AsyncResult(task_key).forget()
 
+    @handle_celery_errors
     async def get_task_result(
         self, owner_metadata: OwnerMetadata, task_uuid: TaskUUID
     ) -> Any:
@@ -154,6 +157,7 @@ async def _get_task_progress_report(
     def _get_task_celery_state(self, task_key: TaskKey) -> TaskState:
         return TaskState(self._celery_app.AsyncResult(task_key).state)
 
+    @handle_celery_errors
     async def get_task_status(
         self, owner_metadata: OwnerMetadata, task_uuid: TaskUUID
     ) -> TaskStatus:
@@ -177,6 +181,7 @@ async def get_task_status(
                 ),
             )
 
+    @handle_celery_errors
     async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         with log_context(
             _logger,
@@ -185,6 +190,7 @@ async def list_tasks(self, owner_metadata: OwnerMetadata) -> list[Task]:
         ):
             return await self._task_info_store.list_tasks(owner_metadata)
 
+    @handle_celery_errors
     async def set_task_progress(
         self, task_key: TaskKey, report: ProgressReport
     ) -> None:

packages/celery-library/tests/conftest.py

@@ -15,7 +15,7 @@
 )
 from celery.signals import worker_init, worker_shutdown
 from celery.worker.worker import WorkController
-from celery_library.backends.redis import RedisTaskInfoStore
+from celery_library.backends.redis import RedisTaskStore
 from celery_library.signals import on_worker_init, on_worker_shutdown
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.types import register_celery_types
@@ -66,7 +66,7 @@ async def run_until_shutdown(
         self._task_manager = CeleryTaskManager(
             self._app,
             self._settings,
-            RedisTaskInfoStore(redis_client_sdk),
+            RedisTaskStore(redis_client_sdk),
         )
 
         startup_completed_event.set()
@@ -156,11 +156,11 @@ async def mock_celery_app(celery_config: dict[str, Any]) -> Celery:
 
 
 @pytest.fixture
-async def celery_task_manager(
+async def task_manager(
     mock_celery_app: Celery,
     celery_settings: CelerySettings,
     use_in_memory_redis: RedisSettings,
-) -> AsyncIterator[CeleryTaskManager]:
+) -> AsyncIterator[TaskManager]:
     register_celery_types()
 
     try:
@@ -173,7 +173,7 @@ async def celery_task_manager(
         yield CeleryTaskManager(
             mock_celery_app,
             celery_settings,
-            RedisTaskInfoStore(redis_client_sdk),
+            RedisTaskStore(redis_client_sdk),
         )
     finally:
         await redis_client_sdk.shutdown()

packages/celery-library/tests/unit/test_async_jobs.py

@@ -138,13 +138,13 @@ async def async_job(task: Task, task_key: TaskKey, action: Action, payload: Any)
 
 @pytest.fixture
 async def register_rpc_routes(
-    async_jobs_rabbitmq_rpc_client: RabbitMQRPCClient, celery_task_manager: TaskManager
+    async_jobs_rabbitmq_rpc_client: RabbitMQRPCClient, task_manager: TaskManager
 ) -> None:
     await async_jobs_rabbitmq_rpc_client.register_router(
-        _async_jobs.router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=celery_task_manager
+        _async_jobs.router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=task_manager
     )
     await async_jobs_rabbitmq_rpc_client.register_router(
-        router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=celery_task_manager
+        router, ASYNC_JOBS_RPC_NAMESPACE, task_manager=task_manager
     )