improve error handling

Author: giancarloromeo

services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py

@@ -56,7 +56,7 @@ async def get_result(
     assert job_id  # nosec
     assert job_id_data  # nosec
 
-    result = await get_celery_client(app).get_result(
+    result = await get_celery_client(app).get_task_result(
         task_context=job_id_data.model_dump(),
         task_uuid=job_id,
     )

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

@@ -83,15 +83,16 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     setup_client_session(app)
 
     setup_rabbitmq(app)
-    setup_rpc_api_routes(app)
+    if not settings.STORAGE_WORKER_MODE:
+        setup_rpc_api_routes(app)
     setup_rest_api_long_running_tasks_for_uploads(app)
     setup_rest_api_routes(app, API_VTAG)
     set_exception_handlers(app)
 
     setup_redis(app)
 
     setup_dsm(app)
-    if settings.STORAGE_CLEANER_INTERVAL_S:
+    if settings.STORAGE_CLEANER_INTERVAL_S and not settings.STORAGE_WORKER_MODE:
         setup_dsm_cleaner(app)
 
     if settings.STORAGE_PROFILING:

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

@@ -1,9 +1,16 @@
+from collections.abc import Callable
+from functools import wraps
 import logging
+import traceback
 
-from celery import Celery
+from celery import Celery, Task
+from celery.exceptions import Ignore
+from celery.contrib.abortable import AbortableTask
 from settings_library.celery import CelerySettings
 from settings_library.redis import RedisDatabase
 
+from .models import TaskError
+
 _logger = logging.getLogger(__name__)
 
 
@@ -18,5 +25,36 @@ def create_app(celery_settings: CelerySettings) -> Celery:
     )
     app.conf.result_expires = celery_settings.CELERY_RESULT_EXPIRES
     app.conf.result_extended = True  # original args are included in the results
+    app.conf.result_serializer = "json"
     app.conf.task_track_started = True
     return app
+
+
+def error_handling(func: Callable):
+    @wraps(func)
+    def wrapper(task: Task, *args, **kwargs):
+        try:
+            return func(task, *args, **kwargs)
+        except Exception as exc:
+            exc_type = type(exc).__name__
+            exc_message = f"{exc}"
+            exc_traceback = traceback.format_exc().split('\n')
+
+            task.update_state(
+                state="ERROR",
+                meta=TaskError(
+                    exc_type=exc_type,
+                    exc_msg=exc_message,
+                ).model_dump(mode="json"),
+                traceback=exc_traceback
+            )
+            raise Ignore from exc
+    return wrapper
+
+
+def define_task(app: Celery, fn: Callable, task_name: str | None = None):
+    app.task(
+        name=task_name or fn.__name__,
+        bind=True,
+        base=AbortableTask,
+    )(error_handling(fn))

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

@@ -1,16 +1,16 @@
 import contextlib
 import logging
-from typing import Any, Final
+from typing import Any, Final, Type
 from uuid import uuid4
 
 from celery import Celery
 from celery.contrib.abortable import AbortableAsyncResult
 from common_library.async_tools import make_async
 from models_library.progress_bar import ProgressReport
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
 from servicelib.logging_utils import log_context
 
-from .models import TaskContext, TaskID, TaskState, TaskStatus, TaskUUID
+from .models import TaskContext, TaskError, TaskID, TaskResult, TaskState, TaskStatus, TaskUUID
 
 _logger = logging.getLogger(__name__)
 
@@ -28,20 +28,23 @@
     "RUNNING": TaskState.RUNNING,
     "SUCCESS": TaskState.SUCCESS,
     "ABORTED": TaskState.ABORTED,
-    "FAILURE": TaskState.FAILURE,
+    "FAILURE": TaskState.ERROR,
+    "ERROR": TaskState.ERROR,
 }
+_CELERY_TASK_ID_KEY_SEPARATOR: Final[str] = ":"
+_CELERY_TASK_ID_KEY_ENCODING = "utf-8"
 
 
 def _build_context_prefix(task_context: TaskContext) -> list[str]:
     return [f"{task_context[key]}" for key in sorted(task_context)]
 
 
 def _build_task_id_prefix(task_context: TaskContext) -> str:
-    return ":".join(_build_context_prefix(task_context))
+    return _CELERY_TASK_ID_KEY_SEPARATOR.join(_build_context_prefix(task_context))
 
 
 def _build_task_id(task_context: TaskContext, task_uuid: TaskUUID) -> TaskID:
-    return ":".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:
@@ -76,9 +79,11 @@ def abort_task(  # pylint: disable=R6301
         AbortableAsyncResult(task_id).abort()
 
     @make_async()
-    def get_result(self, task_context: TaskContext, task_uuid: TaskUUID) -> Any:
+    def get_task_result(self, task_context: TaskContext, task_uuid: TaskUUID) -> TaskResult:
         task_id = _build_task_id(task_context, task_uuid)
-        return self._celery_app.AsyncResult(task_id).result
+        return TypeAdapter(TaskResult).validate_python(
+            self._celery_app.AsyncResult(task_id).result
+        )
 
     def _get_progress_report(
         self, task_context: TaskContext, task_uuid: TaskUUID
@@ -91,7 +96,7 @@ def _get_progress_report(
                 return ProgressReport.model_validate(result)
         if state in (
             TaskState.ABORTED,
-            TaskState.FAILURE,
+            TaskState.ERROR,
             TaskState.SUCCESS,
         ):
             return ProgressReport(actual_value=100.0)
@@ -113,12 +118,12 @@ 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) + "*"
+            _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)):
+        if hasattr(redis, "keys") and (keys := redis.keys(search_key + "*")):
             return {
-                TaskUUID(f"{key}".removeprefix(_CELERY_TASK_META_PREFIX))
+                TaskUUID(f"{key.decode(_CELERY_TASK_ID_KEY_ENCODING).removeprefix(search_key + _CELERY_TASK_ID_KEY_SEPARATOR)}")
                 for key in keys
             }
         return set()

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

@@ -1,9 +1,9 @@
 from enum import StrEnum, auto
-from typing import Any, Final, Self, TypeAlias
+from typing import Annotated, Any, Final, Self, TypeAlias
 from uuid import UUID
 
 from models_library.progress_bar import ProgressReport
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, Field, model_validator
 
 TaskContext: TypeAlias = dict[str, Any]
 TaskID: TypeAlias = str
@@ -17,11 +17,11 @@ class TaskState(StrEnum):
     PENDING = auto()
     RUNNING = auto()
     SUCCESS = auto()
-    FAILURE = auto()
+    ERROR = auto()
     ABORTED = auto()
 
 
-_TASK_DONE = {TaskState.SUCCESS, TaskState.FAILURE, TaskState.ABORTED}
+_TASK_DONE = {TaskState.SUCCESS, TaskState.ERROR, TaskState.ABORTED}
 
 
 class TaskStatus(BaseModel):
@@ -42,11 +42,21 @@ def _check_consistency(self) -> Self:
             TaskState.RUNNING: _MIN_PROGRESS <= value <= _MAX_PROGRESS,
             TaskState.SUCCESS: value == _MAX_PROGRESS,
             TaskState.ABORTED: value == _MAX_PROGRESS,
-            TaskState.FAILURE: value == _MAX_PROGRESS,
+            TaskState.ERROR: value == _MAX_PROGRESS,
         }
 
         if not valid_states.get(self.task_state, True):
             msg = f"Inconsistent progress actual value for state={self.task_state}: {value}"
             raise ValueError(msg)
 
         return self
+
+
+class TaskError(BaseModel):
+    exc_type: str
+    exc_msg: str
+
+
+TaskResult: TypeAlias = Annotated[
+    TaskError | Any, Field(union_mode="left_to_right")
+]

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

@@ -2,7 +2,6 @@
 
 import logging
 
-from celery.contrib.abortable import AbortableTask
 from celery.signals import worker_init, worker_shutdown
 from servicelib.logging_utils import config_all_loggers
 from simcore_service_storage.modules.celery.signals import (
@@ -11,8 +10,8 @@
 )
 
 from ...core.settings import ApplicationSettings
-from ._common import create_app as create_celery_app
-from .tasks import export_data
+from ._common import create_app as create_celery_app, define_task
+from .tasks import export_data, export_data_with_error
 
 _settings = ApplicationSettings.create_from_envs()
 
@@ -30,4 +29,6 @@
 app = create_celery_app(_settings.STORAGE_CELERY)
 worker_init.connect(on_worker_init)
 worker_shutdown.connect(on_worker_shutdown)
-app.task(name="export_data", bind=True, base=AbortableTask)(export_data)
+
+define_task(app, export_data)
+define_task(app, export_data_with_error)

services/storage/tests/unit/modules/celery/test_celery.py

@@ -4,14 +4,16 @@
 from collections.abc import Callable
 from random import randint
 
+from pydantic import ValidationError
 import pytest
 from celery import Celery, Task
 from celery.contrib.abortable import AbortableTask
 from models_library.progress_bar import ProgressReport
 from servicelib.logging_utils import log_context
 from simcore_service_storage.modules.celery import get_event_loop
+from simcore_service_storage.modules.celery._common import define_task
 from simcore_service_storage.modules.celery.client import CeleryTaskQueueClient
-from simcore_service_storage.modules.celery.models import TaskContext, TaskState
+from simcore_service_storage.modules.celery.models import TaskContext, TaskError, TaskState
 from simcore_service_storage.modules.celery.utils import (
     get_celery_worker,
     get_fastapi_app,
@@ -69,11 +71,9 @@ def dreamer_task(task: AbortableTask) -> list[int]:
 @pytest.fixture
 def register_celery_tasks() -> Callable[[Celery], None]:
     def _(celery_app: Celery) -> None:
-        celery_app.task(name="sync_archive", bind=True)(sync_archive)
-        celery_app.task(name="failure_task", bind=True)(failure_task)
-        celery_app.task(name="dreamer_task", base=AbortableTask, bind=True)(
-            dreamer_task
-        )
+        define_task(celery_app, sync_archive)
+        define_task(celery_app, failure_task)
+        define_task(celery_app, dreamer_task)
 
     return _
 
@@ -99,6 +99,9 @@ async def test_sumitting_task_calling_async_function_results_with_success_state(
             progress = await celery_client.get_task_status(task_context, task_uuid)
             assert progress.task_state == TaskState.SUCCESS
 
+    assert (
+        await celery_client.get_task_result(task_context, task_uuid)
+    ) == "archive.zip"
     assert (
         await celery_client.get_task_status(task_context, task_uuid)
     ).task_state == TaskState.SUCCESS
@@ -113,20 +116,19 @@ async def test_submitting_task_with_failure_results_with_error(
     task_uuid = await celery_client.send_task("failure_task", task_context=task_context)
 
     for attempt in Retrying(
-        retry=retry_if_exception_type(AssertionError),
+        retry=retry_if_exception_type((AssertionError, ValidationError)),
         wait=wait_fixed(1),
         stop=stop_after_delay(30),
     ):
         with attempt:
-            result = await celery_client.get_result(task_context, task_uuid)
-            assert isinstance(result, ValueError)
+            result = await celery_client.get_task_result(task_context, task_uuid)
+            assert isinstance(result, TaskError)
 
     assert (
         await celery_client.get_task_status(task_context, task_uuid)
-    ).task_state == TaskState.FAILURE
-    result = await celery_client.get_result(task_context, task_uuid)
-    assert isinstance(result, ValueError)
-    assert f"{result}" == "my error here"
+    ).task_state == TaskState.ERROR
+    result = await celery_client.get_task_result(task_context, task_uuid)
+    assert f"{result.exc_msg}" == "my error here"
 
 
 @pytest.mark.usefixtures("celery_worker")