continue

Author: giancarloromeo

services/storage/src/simcore_service_storage/api/rest/_files.py

@@ -34,7 +34,7 @@
     StorageQueryParamsBase,
     UploadLinks,
 )
-from ...modules.celery.client import CeleryTaskQueueClient
+from ...modules.celery.client import CeleryTaskClient
 from ...modules.celery.models import TaskUUID
 from ...simcore_s3_dsm import SimcoreS3DataManager
 from .._worker_tasks._files import complete_upload_file as remote_complete_upload_file
@@ -270,7 +270,7 @@ async def abort_upload_file(
     status_code=status.HTTP_202_ACCEPTED,
 )
 async def complete_upload_file(
-    celery_client: Annotated[CeleryTaskQueueClient, Depends(get_celery_client)],
+    celery_client: Annotated[CeleryTaskClient, Depends(get_celery_client)],
     query_params: Annotated[StorageQueryParamsBase, Depends()],
     location_id: LocationID,
     file_id: StorageFileID,
@@ -324,7 +324,7 @@ async def complete_upload_file(
     response_model=Envelope[FileUploadCompleteFutureResponse],
 )
 async def is_completed_upload_file(
-    celery_client: Annotated[CeleryTaskQueueClient, Depends(get_celery_client)],
+    celery_client: Annotated[CeleryTaskClient, Depends(get_celery_client)],
     query_params: Annotated[StorageQueryParamsBase, Depends()],
     location_id: LocationID,
     file_id: StorageFileID,

services/storage/src/simcore_service_storage/api/rest/dependencies/celery.py

@@ -4,10 +4,10 @@
 from servicelib.fastapi.dependencies import get_app
 
 from ....modules.celery import get_celery_client as _get_celery_client_from_app
-from ....modules.celery.client import CeleryTaskQueueClient
+from ....modules.celery.client import CeleryTaskClient
 
 
 def get_celery_client(
     app: Annotated[FastAPI, Depends(get_app)],
-) -> CeleryTaskQueueClient:
+) -> CeleryTaskClient:
     return _get_celery_client_from_app(app)

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

@@ -97,7 +97,7 @@ async def result(
 
     if _status.task_state == TaskState.ABORTED:
         raise JobAbortedError(job_id=job_id)
-    if _status.task_state == TaskState.ERROR:
+    if _status.task_state == TaskState.FAILURE:
         # fallback exception to report
         exc_type = type(_result).__name__
         exc_msg = f"{_result}"

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

@@ -10,7 +10,7 @@
 from ._celery_types import register_celery_types
 from ._common import create_app
 from .backends._redis import RedisTaskInfoStore
-from .client import CeleryTaskQueueClient
+from .client import CeleryTaskClient
 
 _logger = logging.getLogger(__name__)
 
@@ -28,7 +28,7 @@ async def on_startup() -> None:
             client_name=f"{APP_NAME}.celery_tasks",
         )
 
-        app.state.celery_client = CeleryTaskQueueClient(
+        app.state.celery_client = CeleryTaskClient(
             celery_app,
             celery_settings,
             RedisTaskInfoStore(redis_client_sdk),
@@ -39,10 +39,10 @@ async def on_startup() -> None:
     app.add_event_handler("startup", on_startup)
 
 
-def get_celery_client(app: FastAPI) -> CeleryTaskQueueClient:
+def get_celery_client(app: FastAPI) -> CeleryTaskClient:
     assert hasattr(app.state, "celery_client")  # nosec
     celery_client = app.state.celery_client
-    assert isinstance(celery_client, CeleryTaskQueueClient)
+    assert isinstance(celery_client, CeleryTaskClient)
     return celery_client
 
 

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

@@ -14,6 +14,7 @@
 )
 from celery.exceptions import Ignore  # type: ignore[import-untyped]
 from pydantic import NonNegativeInt
+from servicelib.async_utils import cancel_wait_task
 
 from . import get_event_loop
 from .errors import encore_celery_transferrable_error
@@ -26,7 +27,8 @@
 _DEFAULT_MAX_RETRIES: Final[NonNegativeInt] = 3
 _DEFAULT_WAIT_BEFORE_RETRY: Final[timedelta] = timedelta(seconds=5)
 _DEFAULT_DONT_AUTORETRY_FOR: Final[tuple[type[Exception], ...]] = ()
-_DEFAULT_ABORT_TASK_TIMEOUT: Final[timedelta] = timedelta(seconds=0.5)
+_DEFAULT_ABORT_TASK_TIMEOUT: Final[timedelta] = timedelta(seconds=1)
+_DEFAULT_CANCEL_TASK_TIMEOUT: Final[timedelta] = timedelta(seconds=5)
 
 T = TypeVar("T")
 P = ParamSpec("P")
@@ -60,10 +62,10 @@ async def run_task(task_id: TaskID) -> R:
                     async def abort_monitor():
                         while not main_task.done():
                             if AbortableAsyncResult(task_id).is_aborted():
-                                _logger.warning(
-                                    "Task %s aborted, cancelling main task", task_id
+                                await cancel_wait_task(
+                                    main_task,
+                                    max_delay=_DEFAULT_CANCEL_TASK_TIMEOUT.total_seconds(),
                                 )
-                                main_task.cancel()
                                 return
                             await asyncio.sleep(
                                 _DEFAULT_ABORT_TASK_TIMEOUT.total_seconds()

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

@@ -1,7 +1,6 @@
-import contextlib
 import logging
 from dataclasses import dataclass
-from typing import Any, Final
+from typing import Any
 from uuid import uuid4
 
 from celery import Celery  # type: ignore[import-untyped]
@@ -10,7 +9,6 @@
 )
 from common_library.async_tools import make_async
 from models_library.progress_bar import ProgressReport
-from pydantic import ValidationError
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
@@ -27,23 +25,13 @@
 
 _logger = logging.getLogger(__name__)
 
-_CELERY_STATES_MAPPING: Final[dict[str, TaskState]] = {
-    "PENDING": TaskState.PENDING,
-    "STARTED": TaskState.PENDING,
-    "RETRY": TaskState.PENDING,
-    "RUNNING": TaskState.RUNNING,
-    "SUCCESS": TaskState.SUCCESS,
-    "ABORTED": TaskState.ABORTED,
-    "FAILURE": TaskState.ERROR,
-    "ERROR": TaskState.ERROR,
-}
 
 _MIN_PROGRESS_VALUE = 0.0
 _MAX_PROGRESS_VALUE = 1.0
 
 
 @dataclass
-class CeleryTaskQueueClient:
+class CeleryTaskClient:
     _celery_app: Celery
     _celery_settings: CelerySettings
     _task_store: TaskInfoStore
@@ -92,11 +80,6 @@ async def abort_task(self, task_context: TaskContext, task_uuid: TaskUUID) -> No
             task_id = build_task_id(task_context, task_uuid)
             await self._abort_task(task_id)
 
-    @make_async()
-    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:
@@ -114,28 +97,31 @@ async def get_task_result(
                     await self._task_store.remove(task_id)
             return result
 
-    @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)
-                return ProgressReport.model_validate(result)
+    async def _get_progress_report(
+        self, task_id: TaskID, state: TaskState
+    ) -> ProgressReport:
+        if state in (TaskState.STARTED, TaskState.RETRY):
+            progress = await self._task_store.get_progress(task_id)
+            if progress is not None:
+                return progress
         if state in (
-            TaskState.ABORTED,
-            TaskState.ERROR,
             TaskState.SUCCESS,
+            TaskState.ABORTED,
+            TaskState.FAILURE,
         ):
             return ProgressReport(
                 actual_value=_MAX_PROGRESS_VALUE, total=_MAX_PROGRESS_VALUE
             )
+
+        # task is pending
         return ProgressReport(
             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]
+        return TaskState(self._celery_app.AsyncResult(task_id).state)
 
     async def get_task_status(
         self, task_context: TaskContext, task_uuid: TaskUUID
@@ -146,11 +132,11 @@ async def get_task_status(
             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)
+            task_id = build_task_id(task_context, task_uuid)
             return TaskStatus(
                 task_uuid=task_uuid,
                 task_state=task_state,
-                progress_report=await self._get_progress_report(task_state, result),
+                progress_report=await self._get_progress_report(task_id, task_state),
             )
 
     async def get_task_uuids(self, task_context: TaskContext) -> set[TaskUUID]:

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

@@ -1,5 +1,5 @@
 from datetime import timedelta
-from enum import StrEnum, auto
+from enum import StrEnum
 from typing import Any, Final, Protocol, TypeAlias
 from uuid import UUID
 
@@ -26,11 +26,12 @@ def build_task_id(task_context: TaskContext, task_uuid: TaskUUID) -> TaskID:
 
 
 class TaskState(StrEnum):
-    PENDING = auto()
-    RUNNING = auto()
-    SUCCESS = auto()
-    ERROR = auto()
-    ABORTED = auto()
+    PENDING = "PENDING"
+    STARTED = "STARTED"
+    RETRY = "RETRY"
+    SUCCESS = "SUCCESS"
+    FAILURE = "FAILURE"
+    ABORTED = "ABORTED"
 
 
 class TasksQueue(StrEnum):
@@ -43,15 +44,15 @@ class TaskMetadata(BaseModel):
     queue: TasksQueue = TasksQueue.DEFAULT
 
 
-_TASK_DONE = {TaskState.SUCCESS, TaskState.ERROR, TaskState.ABORTED}
+_TASK_DONE = {TaskState.SUCCESS, TaskState.FAILURE, TaskState.ABORTED}
 
 
 class TaskInfoStore(Protocol):
     async def exists(self, task_id: TaskID) -> bool: ...
 
-    async def get_progress(self, task_id: TaskID) -> ProgressReport | None: ...
+    async def get_progress(self, task_id: TaskID) -> ProgressReport: ...
 
-    async def get_metadata(self, task_id: TaskID) -> TaskMetadata | None: ...
+    async def get_metadata(self, task_id: TaskID) -> TaskMetadata: ...
 
     async def get_uuids(self, task_context: TaskContext) -> set[TaskUUID]: ...
 

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

@@ -14,13 +14,14 @@
 
 from ...core.application import create_app
 from ...core.settings import ApplicationSettings
-from ...modules.celery.backends._redis import RedisTaskInfoStore, set_event_loop
-from ...modules.celery.utils import (
+from . import set_event_loop
+from .backends._redis import RedisTaskInfoStore
+from .utils import (
     get_fastapi_app,
     set_celery_worker,
     set_fastapi_app,
 )
-from ...modules.celery.worker import CeleryTaskWorker
+from .worker import CeleryTaskWorker
 
 _logger = logging.getLogger(__name__)
 

services/storage/tests/unit/test_modules_celery.py

@@ -19,7 +19,7 @@
 from servicelib.logging_utils import log_context
 from simcore_service_storage.modules.celery import get_celery_client, get_event_loop
 from simcore_service_storage.modules.celery._task import register_task
-from simcore_service_storage.modules.celery.client import CeleryTaskQueueClient
+from simcore_service_storage.modules.celery.client import CeleryTaskClient
 from simcore_service_storage.modules.celery.errors import TransferrableCeleryError
 from simcore_service_storage.modules.celery.models import (
     TaskContext,
@@ -42,7 +42,7 @@
 def celery_client(
     initialized_app: FastAPI,
     with_storage_celery_worker: CeleryTaskWorker,
-) -> CeleryTaskQueueClient:
+) -> CeleryTaskClient:
     return get_celery_client(initialized_app)
 
 
@@ -106,7 +106,7 @@ def _(celery_app: Celery) -> None:
 
 
 async def test_submitting_task_calling_async_function_results_with_success_state(
-    celery_client: CeleryTaskQueueClient,
+    celery_client: CeleryTaskClient,
 ):
     task_context = TaskContext(user_id=42)
 
@@ -134,7 +134,7 @@ async def test_submitting_task_calling_async_function_results_with_success_state
 
 
 async def test_submitting_task_with_failure_results_with_error(
-    celery_client: CeleryTaskQueueClient,
+    celery_client: CeleryTaskClient,
 ):
     task_context = TaskContext(user_id=42)
 
@@ -157,7 +157,7 @@ async def test_submitting_task_with_failure_results_with_error(
 
 
 async def test_aborting_task_results_with_aborted_state(
-    celery_client: CeleryTaskQueueClient,
+    celery_client: CeleryTaskClient,
 ):
     task_context = TaskContext(user_id=42)
 
@@ -183,7 +183,7 @@ async def test_aborting_task_results_with_aborted_state(
 
 
 async def test_listing_task_uuids_contains_submitted_task(
-    celery_client: CeleryTaskQueueClient,
+    celery_client: CeleryTaskClient,
 ):
     task_context = TaskContext(user_id=42)