♻️ TasksManager uses Redis for task data

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

@@ -47,6 +47,8 @@ async def create_task_manager(
         ),
         client_name="celery_tasks",
     )
+    await redis_client_sdk.setup()
+    # GCR please address https://github.com/ITISFoundation/osparc-simcore/issues/8159
 
     return CeleryTaskManager(
         app,

packages/models-library/src/models_library/api_schemas_long_running_tasks/base.py

@@ -1,4 +1,5 @@
 import logging
+from collections.abc import Awaitable, Callable
 from typing import Annotated, TypeAlias
 
 from pydantic import BaseModel, Field, field_validator, validate_call
@@ -22,8 +23,16 @@ class TaskProgress(BaseModel):
     message: ProgressMessage = ""
     percent: ProgressPercent = 0.0
 
+    # used to propagate progress updates internally
+    _update_callback: Callable[["TaskProgress"], Awaitable[None]] | None = None
+
+    def set_update_callback(
+        self, callback: Callable[["TaskProgress"], Awaitable[None]]
+    ) -> None:
+        self._update_callback = callback
+
     @validate_call
-    def update(
+    async def update(
         self,
         *,
         message: ProgressMessage | None = None,
@@ -40,6 +49,14 @@ def update(
 
         _logger.debug("Progress update: %s", f"{self}")
 
+        if self._update_callback is not None:
+            await self._update_callback(self)
+        else:
+            _logger.warning(
+                "No update callback set for TaskProgress %s, progress will not be propagated",
+                self.task_id,
+            )
+
     @classmethod
     def create(cls, task_id: TaskId | None = None) -> "TaskProgress":
         return cls(task_id=task_id)

packages/pytest-simcore/src/pytest_simcore/long_running_tasks.py

@@ -0,0 +1,26 @@
+# pylint: disable=unused-argument
+
+import functools
+
+import pytest
+from pytest_mock import MockerFixture
+from servicelib.long_running_tasks import task
+from servicelib.long_running_tasks._store.in_memory import InMemoryStore
+
+
+def _mock_decorator_with_args(*decorator_args, **decorator_kwargs):
+    def decorator(func):
+        @functools.wraps(func)
+        async def wrapper(*func_args, **func_kwargs):
+            return await func(*func_args, **func_kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+@pytest.fixture
+def use_in_memory_long_running_tasks(mocker: MockerFixture) -> None:
+    mocker.patch.object(task, "RedisStore", InMemoryStore)
+    # for testing the exclsive is not required so it's disabled
+    mocker.patch.object(task, "exclusive", _mock_decorator_with_args)

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_manager.py

@@ -1,9 +1,11 @@
 import datetime
 
 from aiohttp import web
+from settings_library.redis import RedisSettings
 
 from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
-from ...long_running_tasks.task import TaskContext, TasksManager
+from ...long_running_tasks.models import TaskContext
+from ...long_running_tasks.task import Namespace, TasksManager
 from ._constants import APP_LONG_RUNNING_MANAGER_KEY
 from ._request import get_task_context
 
@@ -14,11 +16,15 @@ def __init__(
         app: web.Application,
         stale_task_check_interval: datetime.timedelta,
         stale_task_detect_timeout: datetime.timedelta,
+        redis_settings: RedisSettings,
+        namespace: Namespace,
     ):
         self._app = app
         self._tasks_manager = TasksManager(
             stale_task_check_interval=stale_task_check_interval,
             stale_task_detect_timeout=stale_task_detect_timeout,
+            redis_settings=redis_settings,
+            namespace=namespace,
         )
 
     @property

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_routes.py

@@ -28,7 +28,7 @@ async def list_tasks(request: web.Request) -> web.Response:
                 result_href=f"{request.app.router['get_task_result'].url_for(task_id=t.task_id)}",
                 abort_href=f"{request.app.router['cancel_and_delete_task'].url_for(task_id=t.task_id)}",
             )
-            for t in lrt_api.list_tasks(
+            for t in await lrt_api.list_tasks(
                 long_running_manager.tasks_manager,
                 long_running_manager.get_task_context(request),
             )
@@ -41,7 +41,7 @@ async def get_task_status(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
     long_running_manager = get_long_running_manager(request.app)
 
-    task_status: TaskStatus = lrt_api.get_task_status(
+    task_status: TaskStatus = await lrt_api.get_task_status(
         long_running_manager.tasks_manager,
         long_running_manager.get_task_context(request),
         path_params.task_id,

packages/service-library/src/servicelib/aiohttp/long_running_tasks/_server.py

@@ -7,15 +7,17 @@
 from aiohttp import web
 from common_library.json_serialization import json_dumps
 from pydantic import AnyHttpUrl, TypeAdapter
+from servicelib.long_running_tasks.task import Namespace
+from settings_library.redis import RedisSettings
 
 from ...aiohttp import status
 from ...long_running_tasks import lrt_api
 from ...long_running_tasks.constants import (
     DEFAULT_STALE_TASK_CHECK_INTERVAL,
     DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
-from ...long_running_tasks.models import TaskGet
-from ...long_running_tasks.task import RegisteredTaskName, TaskContext
+from ...long_running_tasks.models import TaskContext, TaskGet
+from ...long_running_tasks.task import RegisteredTaskName
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
@@ -91,9 +93,8 @@ async def start_long_running_task(
     except asyncio.CancelledError:
         # cancel the task, the client has disconnected
         if task_id:
-            long_running_manager = get_long_running_manager(request_.app)
-            await long_running_manager.tasks_manager.cancel_task(
-                task_id, with_task_context=None
+            await lrt_api.cancel_task(
+                long_running_manager.tasks_manager, task_context, task_id
             )
         raise
 
@@ -121,6 +122,8 @@ def setup(
     app: web.Application,
     *,
     router_prefix: str,
+    redis_settings: RedisSettings,
+    namespace: Namespace,
     handler_check_decorator: Callable = _no_ops_decorator,
     task_request_context_decorator: Callable = _no_task_context_decorator,
     stale_task_check_interval: datetime.timedelta = DEFAULT_STALE_TASK_CHECK_INTERVAL,
@@ -146,6 +149,8 @@ async def on_cleanup_ctx(app: web.Application) -> AsyncGenerator[None, None]:
                 app=app,
                 stale_task_check_interval=stale_task_check_interval,
                 stale_task_detect_timeout=stale_task_detect_timeout,
+                redis_settings=redis_settings,
+                namespace=namespace,
             )
         )
 

packages/service-library/src/servicelib/fastapi/long_running_tasks/_manager.py

@@ -1,9 +1,10 @@
 import datetime
 
 from fastapi import FastAPI
+from settings_library.redis import RedisSettings
 
 from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
-from ...long_running_tasks.task import TasksManager
+from ...long_running_tasks.task import Namespace, TasksManager
 
 
 class FastAPILongRunningManager(BaseLongRunningManager):
@@ -12,11 +13,15 @@ def __init__(
         app: FastAPI,
         stale_task_check_interval: datetime.timedelta,
         stale_task_detect_timeout: datetime.timedelta,
+        redis_settings: RedisSettings,
+        namespace: Namespace,
     ):
         self._app = app
         self._tasks_manager = TasksManager(
             stale_task_check_interval=stale_task_check_interval,
             stale_task_detect_timeout=stale_task_detect_timeout,
+            redis_settings=redis_settings,
+            namespace=namespace,
         )
 
     @property

packages/service-library/src/servicelib/fastapi/long_running_tasks/_routes.py

@@ -29,8 +29,8 @@ async def list_tasks(
                 request.url_for("cancel_and_delete_task", task_id=t.task_id)
             ),
         )
-        for t in lrt_api.list_tasks(
-            long_running_manager.tasks_manager, task_context=None
+        for t in await lrt_api.list_tasks(
+            long_running_manager.tasks_manager, task_context={}
         )
     ]
 
@@ -51,8 +51,8 @@ async def get_task_status(
     ],
 ) -> TaskStatus:
     assert request  # nosec
-    return lrt_api.get_task_status(
-        long_running_manager.tasks_manager, task_context=None, task_id=task_id
+    return await lrt_api.get_task_status(
+        long_running_manager.tasks_manager, task_context={}, task_id=task_id
     )
 
 
@@ -75,7 +75,7 @@ async def get_task_result(
 ) -> TaskResult | Any:
     assert request  # nosec
     return await lrt_api.get_task_result(
-        long_running_manager.tasks_manager, task_context=None, task_id=task_id
+        long_running_manager.tasks_manager, task_context={}, task_id=task_id
     )
 
 
@@ -98,5 +98,5 @@ async def cancel_and_delete_task(
 ) -> None:
     assert request  # nosec
     await lrt_api.remove_task(
-        long_running_manager.tasks_manager, task_context=None, task_id=task_id
+        long_running_manager.tasks_manager, task_context={}, task_id=task_id
     )

packages/service-library/src/servicelib/fastapi/long_running_tasks/_server.py

@@ -1,12 +1,14 @@
 import datetime
 
 from fastapi import APIRouter, FastAPI
+from settings_library.redis import RedisSettings
 
 from ...long_running_tasks.constants import (
     DEFAULT_STALE_TASK_CHECK_INTERVAL,
     DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
 from ...long_running_tasks.errors import BaseLongRunningError
+from ...long_running_tasks.task import Namespace
 from ._error_handlers import base_long_running_error_handler
 from ._manager import FastAPILongRunningManager
 from ._routes import router
@@ -16,6 +18,8 @@ def setup(
     app: FastAPI,
     *,
     router_prefix: str = "",
+    redis_settings: RedisSettings,
+    namespace: Namespace,
     stale_task_check_interval: datetime.timedelta = DEFAULT_STALE_TASK_CHECK_INTERVAL,
     stale_task_detect_timeout: datetime.timedelta = DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 ) -> None:
@@ -41,6 +45,8 @@ async def on_startup() -> None:
                 app=app,
                 stale_task_check_interval=stale_task_check_interval,
                 stale_task_detect_timeout=stale_task_detect_timeout,
+                redis_settings=redis_settings,
+                namespace=namespace,
             )
         )
         await long_running_manager.setup()

packages/service-library/src/servicelib/fastapi/redis_lifespan.py

@@ -51,6 +51,7 @@ async def redis_client_sdk_lifespan(_: FastAPI, state: State) -> AsyncIterator[S
                 redis_dsn_with_secrets,
                 client_name=redis_state.REDIS_CLIENT_NAME,
             )
+            await redis_client.setup()
 
         try:
             yield {"REDIS_CLIENT_SDK": redis_client, **called_state}

packages/service-library/src/servicelib/long_running_tasks/_store/base.py

@@ -0,0 +1,40 @@
+from abc import abstractmethod
+
+from ..models import TaskContext, TaskData, TaskId
+
+
+class BaseStore:
+
+    @abstractmethod
+    async def get_task_data(self, task_id: TaskId) -> TaskData | None:
+        """Retrieve a tracked task by its key."""
+
+    @abstractmethod
+    async def set_task_data(self, task_id: TaskId, value: TaskData) -> None:
+        """Set a tracked task with its key."""
+
+    @abstractmethod
+    async def list_tasks_data(self) -> list[TaskData]:
+        """List all tracked tasks."""
+
+    @abstractmethod
+    async def delete_task_data(self, task_id: TaskId) -> None:
+        """Delete a tracked task by its key."""
+
+    @abstractmethod
+    async def set_as_cancelled(
+        self, task_id: TaskId, with_task_context: TaskContext
+    ) -> None:
+        """Mark a tracked task as cancelled."""
+
+    @abstractmethod
+    async def get_cancelled(self) -> dict[TaskId, TaskContext]:
+        """Get cancelled tasks."""
+
+    @abstractmethod
+    async def setup(self) -> None:
+        """Setup the store, if needed."""
+
+    @abstractmethod
+    async def shutdown(self) -> None:
+        """Shutdown the store, if needed."""

packages/service-library/src/servicelib/long_running_tasks/_store/in_memory.py

@@ -0,0 +1,36 @@
+from ..models import TaskContext, TaskData, TaskId
+from .base import BaseStore
+
+
+class InMemoryStore(BaseStore):
+    def __init__(self, *args, **kwargs):
+        _ = args
+        _ = kwargs
+        self._tasks_data: dict[TaskId, TaskData] = {}
+        self._cancelled_tasks: dict[TaskId, TaskContext] = {}
+
+    async def setup(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def get_task_data(self, task_id: TaskId) -> TaskData | None:
+        return self._tasks_data.get(task_id, None)
+
+    async def set_task_data(self, task_id: TaskId, value: TaskData) -> None:
+        self._tasks_data[task_id] = value
+
+    async def list_tasks_data(self) -> list[TaskData]:
+        return list(self._tasks_data.values())
+
+    async def delete_task_data(self, task_id: TaskId) -> None:
+        self._tasks_data.pop(task_id, None)
+
+    async def set_as_cancelled(
+        self, task_id: TaskId, with_task_context: TaskContext
+    ) -> None:
+        self._cancelled_tasks[task_id] = with_task_context
+
+    async def get_cancelled(self) -> dict[TaskId, TaskContext]:
+        return self._cancelled_tasks