using long_running_manager for aiohttp

Author: Andrei Neagu

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

@@ -18,12 +18,12 @@ class TaskResult(BaseModel):
     error: Any | None
 
 
-class TaskGetWithoutHref(BaseModel):
+class TaskBase(BaseModel):
     task_id: TaskId
     task_name: str
 
 
-class TaskGet(TaskGetWithoutHref):
+class TaskGet(TaskBase):
     status_href: str
     result_href: str
     abort_href: str

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

@@ -3,9 +3,9 @@
 from pydantic import PositiveFloat
 
 MINUTE: Final[PositiveFloat] = 60
-APP_LONG_RUNNING_TASKS_MANAGER_KEY: Final[
-    str
-] = f"{__name__ }.long_running_tasks.tasks_manager"
-RQT_LONG_RUNNING_TASKS_CONTEXT_KEY: Final[
-    str
-] = f"{__name__}.long_running_tasks.context"
+APP_LONG_RUNNING_MANAGER_KEY: Final[str] = (
+    f"{__name__ }.long_running_tasks.tasks_manager"
+)
+RQT_LONG_RUNNING_TASKS_CONTEXT_KEY: Final[str] = (
+    f"{__name__}.long_running_tasks.context"
+)

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

@@ -1,11 +1,10 @@
-from aiohttp import web
+from typing import Any
 
-from ...long_running_tasks.task import TasksManager
-from ._constants import APP_LONG_RUNNING_TASKS_MANAGER_KEY
+from aiohttp import web
 
-# NOTE: figure out how to remove these and expose them differently if possible
+from ._constants import RQT_LONG_RUNNING_TASKS_CONTEXT_KEY
 
 
-def get_tasks_manager(app: web.Application) -> TasksManager:
-    output: TasksManager = app[APP_LONG_RUNNING_TASKS_MANAGER_KEY]
+def get_task_context(request: web.Request) -> dict[str, Any]:
+    output: dict[str, Any] = request[RQT_LONG_RUNNING_TASKS_CONTEXT_KEY]
     return output

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

@@ -0,0 +1,41 @@
+import datetime
+
+from aiohttp import web
+
+from ...long_running_tasks.base_long_running_manager import BaseLongRunningManager
+from ...long_running_tasks.task import TaskContext, TasksManager
+from ._constants import APP_LONG_RUNNING_MANAGER_KEY
+from ._dependencies import get_task_context
+
+
+class AiohttpLongRunningManager(BaseLongRunningManager):
+    def __init__(
+        self,
+        app: web.Application,
+        stale_task_check_interval: datetime.timedelta,
+        stale_task_detect_timeout: datetime.timedelta,
+    ):
+        self._app = app
+        self._tasks_manager = TasksManager(
+            stale_task_check_interval=stale_task_check_interval,
+            stale_task_detect_timeout=stale_task_detect_timeout,
+        )
+
+    @property
+    def tasks_manager(self) -> TasksManager:
+        return self._tasks_manager
+
+    async def setup(self) -> None:
+        await self._tasks_manager.setup()
+
+    async def teardown(self) -> None:
+        await self._tasks_manager.teardown()
+
+    @staticmethod
+    def get_task_context(request: web.Request) -> TaskContext:
+        return get_task_context(request)
+
+
+def get_long_running_manager(app: web.Application) -> AiohttpLongRunningManager:
+    output: AiohttpLongRunningManager = app[APP_LONG_RUNNING_MANAGER_KEY]
+    return output

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

@@ -3,30 +3,23 @@
 from aiohttp import web
 from pydantic import BaseModel
 from servicelib.aiohttp import status
-from servicelib.aiohttp.rest_responses import create_data_response
 
 from ...long_running_tasks import http_endpoint_responses
 from ...long_running_tasks.models import TaskGet, TaskId, TaskStatus
 from ..requests_validation import parse_request_path_parameters_as
-from ._constants import RQT_LONG_RUNNING_TASKS_CONTEXT_KEY
-from ._dependencies import get_tasks_manager
+from ..rest_responses import create_data_response
+from ._manager import get_long_running_manager
 
 routes = web.RouteTableDef()
 
 
-def get_task_context(request: web.Request) -> dict[str, Any]:
-    output: dict[str, Any] = request[RQT_LONG_RUNNING_TASKS_CONTEXT_KEY]
-    return output
-
-
 class _PathParam(BaseModel):
     task_id: TaskId
 
 
 @routes.get("", name="list_tasks")
 async def list_tasks(request: web.Request) -> web.Response:
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
     return create_data_response(
         [
             TaskGet(
@@ -36,41 +29,48 @@ 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 http_endpoint_responses.list_tasks(tasks_manager, task_context)
+            for t in http_endpoint_responses.list_tasks(
+                long_running_manager.tasks_manager,
+                long_running_manager.get_task_context(request),
+            )
         ]
     )
 
 
 @routes.get("/{task_id}", name="get_task_status")
 async def get_task_status(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
 
     task_status: TaskStatus = http_endpoint_responses.get_task_status(
-        tasks_manager, task_context, path_params.task_id
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
     )
     return create_data_response(task_status)
 
 
 @routes.get("/{task_id}/result", name="get_task_result")
 async def get_task_result(request: web.Request) -> web.Response | Any:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
 
     # NOTE: this might raise an exception that will be catched by the _error_handlers
     return await http_endpoint_responses.get_task_result(
-        tasks_manager, task_context, path_params.task_id
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
     )
 
 
 @routes.delete("/{task_id}", name="cancel_and_delete_task")
 async def cancel_and_delete_task(request: web.Request) -> web.Response:
     path_params = parse_request_path_parameters_as(_PathParam, request)
-    tasks_manager = get_tasks_manager(request.app)
-    task_context = get_task_context(request)
+    long_running_manager = get_long_running_manager(request.app)
+
     await http_endpoint_responses.remove_task(
-        tasks_manager, task_context, path_params.task_id
+        long_running_manager.tasks_manager,
+        long_running_manager.get_task_context(request),
+        path_params.task_id,
     )
     return web.json_response(status=status.HTTP_204_NO_CONTENT)

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

@@ -1,6 +1,5 @@
 import asyncio
 import datetime
-import logging
 from collections.abc import AsyncGenerator, Callable
 from functools import wraps
 from typing import Any
@@ -15,22 +14,15 @@
     DEFAULT_STALE_TASK_DETECT_TIMEOUT,
 )
 from ...long_running_tasks.models import TaskGet
-from ...long_running_tasks.task import (
-    TaskContext,
-    TaskProtocol,
-    TasksManager,
-    start_task,
-)
+from ...long_running_tasks.task import TaskContext, TaskProtocol, start_task
 from ..typing_extension import Handler
 from . import _routes
 from ._constants import (
-    APP_LONG_RUNNING_TASKS_MANAGER_KEY,
+    APP_LONG_RUNNING_MANAGER_KEY,
     RQT_LONG_RUNNING_TASKS_CONTEXT_KEY,
 )
-from ._dependencies import get_tasks_manager
 from ._error_handlers import base_long_running_error_handler
-
-_logger = logging.getLogger(__name__)
+from ._manager import AiohttpLongRunningManager, get_long_running_manager
 
 
 def no_ops_decorator(handler: Handler):
@@ -59,12 +51,12 @@ async def start_long_running_task(
     task_context: TaskContext,
     **task_kwargs: Any,
 ) -> web.Response:
-    task_manager = get_tasks_manager(request_.app)
+    long_running_manager = get_long_running_manager(request_.app)
     task_name = _create_task_name_from_request(request_)
     task_id = None
     try:
         task_id = start_task(
-            task_manager,
+            long_running_manager.tasks_manager,
             task_,
             fire_and_forget=fire_and_forget,
             task_context=task_context,
@@ -99,8 +91,10 @@ async def start_long_running_task(
     except asyncio.CancelledError:
         # cancel the task, the client has disconnected
         if task_id:
-            task_manager = get_tasks_manager(request_.app)
-            await task_manager.cancel_task(task_id, with_task_context=None)
+            long_running_manager = get_long_running_manager(request_.app)
+            await long_running_manager.tasks_manager.cancel_task(
+                task_id, with_task_context=None
+            )
         raise
 
 
@@ -147,19 +141,20 @@ async def on_cleanup_ctx(app: web.Application) -> AsyncGenerator[None, None]:
         app.middlewares.append(base_long_running_error_handler)
 
         # add components to state
-        app[APP_LONG_RUNNING_TASKS_MANAGER_KEY] = long_running_task_manager = (
-            TasksManager(
+        app[APP_LONG_RUNNING_MANAGER_KEY] = long_running_manager = (
+            AiohttpLongRunningManager(
+                app=app,
                 stale_task_check_interval=stale_task_check_interval,
                 stale_task_detect_timeout=stale_task_detect_timeout,
             )
         )
 
-        await long_running_task_manager.setup()
+        await long_running_manager.setup()
 
         yield
 
         # cleanup
-        await long_running_task_manager.teardown()
+        await long_running_manager.teardown()
 
     # add routing (done at setup-time)
     _wrap_and_add_routes(

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

@@ -6,13 +6,11 @@
 running task.
 """
 
-from ._dependencies import get_tasks_manager
-from ._routes import get_task_context
+from ._manager import get_long_running_manager
 from ._server import setup, start_long_running_task
 
 __all__: tuple[str, ...] = (
-    "get_task_context",
-    "get_tasks_manager",
+    "get_long_running_manager",
     "setup",
     "start_long_running_task",
 )

packages/service-library/src/servicelib/long_running_tasks/http_endpoint_responses.py

@@ -1,19 +1,16 @@
 from typing import Any
 
-from .models import TaskGetWithoutHref, TaskId, TaskStatus
+from .models import TaskBase, TaskId, TaskStatus
 from .task import TaskContext, TasksManager, TrackedTask
 
 
 def list_tasks(
     tasks_manager: TasksManager, task_context: TaskContext | None
-) -> list[TaskGetWithoutHref]:
+) -> list[TaskBase]:
     tracked_tasks: list[TrackedTask] = tasks_manager.list_tasks(
         with_task_context=task_context
     )
-    return [
-        TaskGetWithoutHref(task_id=t.task_id, task_name=t.task_name)
-        for t in tracked_tasks
-    ]
+    return [TaskBase(task_id=t.task_id, task_name=t.task_name) for t in tracked_tasks]
 
 
 def get_task_status(

packages/service-library/src/servicelib/long_running_tasks/models.py

@@ -12,8 +12,8 @@
     TaskProgress,
 )
 from models_library.api_schemas_long_running_tasks.tasks import (
+    TaskBase,
     TaskGet,
-    TaskGetWithoutHref,
     TaskResult,
     TaskStatus,
 )
@@ -79,7 +79,7 @@ async def result(self) -> Any:
     "ProgressMessage",
     "ProgressPercent",
     "TaskGet",
-    "TaskGetWithoutHref",
+    "TaskBase",
     "TaskId",
     "TaskProgress",
     "TaskResult",

services/web/server/src/simcore_service_webserver/tasks/_rest.py

@@ -22,8 +22,7 @@
 from pydantic import BaseModel
 from servicelib.aiohttp import status
 from servicelib.aiohttp.long_running_tasks.server import (
-    get_task_context,
-    get_tasks_manager,
+    get_long_running_manager,
 )
 from servicelib.aiohttp.requests_validation import (
     parse_request_path_parameters_as,
@@ -57,10 +56,10 @@
 @handle_export_data_exceptions
 @webserver_request_context_decorator
 async def get_async_jobs(request: web.Request) -> web.Response:
-    inprocess_task_manager = get_tasks_manager(request.app)
-    inprocess_task_context = get_task_context(request)
+    inprocess_long_running_manager = get_long_running_manager(request.app)
     inprocess_tracked_tasks = http_endpoint_responses.list_tasks(
-        inprocess_task_manager, inprocess_task_context
+        inprocess_long_running_manager.tasks_manager,
+        inprocess_long_running_manager.get_task_context(request),
     )
 
     _req_ctx = AuthenticatedRequestContext.model_validate(request)