feat: wrap app server

Author: giancarloromeo

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

@@ -1,5 +1,4 @@
 import logging
-from asyncio import AbstractEventLoop
 
 from fastapi import FastAPI
 from settings_library.celery import CelerySettings
@@ -27,13 +26,3 @@ def get_celery_client(app: FastAPI) -> CeleryTaskManager:
     celery_client = app.state.celery_client
     assert isinstance(celery_client, CeleryTaskManager)
     return celery_client
-
-
-def get_event_loop(app: FastAPI) -> AbstractEventLoop:
-    event_loop = app.state.event_loop
-    assert isinstance(event_loop, AbstractEventLoop)
-    return event_loop
-
-
-def set_event_loop(app: FastAPI, event_loop: AbstractEventLoop) -> None:
-    app.state.event_loop = event_loop

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

@@ -1,50 +1,39 @@
 import asyncio
-import datetime
 import logging
 import threading
-from collections.abc import Callable
-from typing import Final
 
-from asgi_lifespan import LifespanManager
 from celery import Celery  # type: ignore[import-untyped]
 from celery.worker.worker import WorkController  # type: ignore[import-untyped]
-from fastapi import FastAPI
+from servicelib.base_app_server import BaseAppServer
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
-from . import set_event_loop
 from .common import create_task_manager
 from .utils import (
-    get_fastapi_app,
-    set_fastapi_app,
+    get_app_server,
+    set_app_server,
     set_task_manager,
 )
 
 _logger = logging.getLogger(__name__)
 
-_SHUTDOWN_TIMEOUT: Final[float] = datetime.timedelta(seconds=10).total_seconds()
-_STARTUP_TIMEOUT: Final[float] = datetime.timedelta(minutes=1).total_seconds()
-
 
 def on_worker_init(
-    app_factory: Callable[[], FastAPI],
+    app_server: BaseAppServer,
     celery_settings: CelerySettings,
     sender: WorkController,
     **_kwargs,
 ) -> None:
-    startup_complete_event = threading.Event()
-
-    def _init(startup_complete_event: threading.Event) -> None:
+    def _init() -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
-        shutdown_event = asyncio.Event()
 
-        fastapi_app = app_factory()
-        assert isinstance(fastapi_app, FastAPI)  # nosec
+        app_server.event_loop = loop
 
         async def setup_task_manager():
             assert sender.app  # nosec
             assert isinstance(sender.app, Celery)  # nosec
+
             set_task_manager(
                 sender.app,
                 create_task_manager(
@@ -53,45 +42,22 @@ async def setup_task_manager():
                 ),
             )
 
-        async def fastapi_lifespan(
-            startup_complete_event: threading.Event, shutdown_event: asyncio.Event
-        ) -> None:
-            async with LifespanManager(
-                fastapi_app,
-                startup_timeout=_STARTUP_TIMEOUT,
-                shutdown_timeout=_SHUTDOWN_TIMEOUT,
-            ):
-                try:
-                    _logger.info("fastapi APP started!")
-                    startup_complete_event.set()
-                    await shutdown_event.wait()
-                except asyncio.CancelledError:
-                    _logger.warning("Lifespan task cancelled")
-
-        fastapi_app.state.shutdown_event = shutdown_event
-        set_event_loop(fastapi_app, loop)
-
-        set_fastapi_app(sender.app, fastapi_app)
+        set_app_server(sender.app, app_server)
         loop.run_until_complete(setup_task_manager())
-        loop.run_until_complete(
-            fastapi_lifespan(startup_complete_event, shutdown_event)
-        )
+        loop.run_until_complete(app_server.startup())
 
     thread = threading.Thread(
         group=None,
         target=_init,
-        name="fastapi_app",
-        args=(startup_complete_event,),
+        name="app_server_init",
         daemon=True,
     )
     thread.start()
-    # ensure the fastapi app is ready before going on
-    startup_complete_event.wait(_STARTUP_TIMEOUT * 1.1)
 
 
 def on_worker_shutdown(sender, **_kwargs) -> None:
-    with log_context(_logger, logging.INFO, "Worker Shuts-down"):
+    with log_context(_logger, logging.INFO, "Worker shutdown"):
         assert isinstance(sender.app, Celery)
-        fastapi_app = get_fastapi_app(sender.app)
-        assert isinstance(fastapi_app, FastAPI)
-        fastapi_app.state.shutdown_event.set()
+        app_server = get_app_server(sender.app)
+
+        app_server.event_loop.run_until_complete(app_server.shutdown())

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

@@ -15,10 +15,9 @@
 from pydantic import NonNegativeInt
 from servicelib.async_utils import cancel_wait_task
 
-from . import get_event_loop
 from .errors import encode_celery_transferrable_error
 from .models import TaskID
-from .utils import get_fastapi_app
+from .utils import get_app_server
 
 _logger = logging.getLogger(__name__)
 
@@ -48,7 +47,7 @@ def decorator(
     ) -> Callable[Concatenate[AbortableTask, P], R]:
         @wraps(coro)
         def wrapper(task: AbortableTask, *args: P.args, **kwargs: P.kwargs) -> R:
-            fastapi_app = get_fastapi_app(app)
+            app_server = get_app_server(app)
             # NOTE: task.request is a thread local object, so we need to pass the id explicitly
             assert task.request.id is not None  # nosec
 
@@ -90,7 +89,7 @@ async def abort_monitor():
 
             return asyncio.run_coroutine_threadsafe(
                 run_task(task.request.id),
-                get_event_loop(fastapi_app),
+                app_server.event_loop,
             ).result()
 
         return wrapper

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

@@ -1,27 +1,27 @@
 from celery import Celery  # type: ignore[import-untyped]
-from fastapi import FastAPI
+from servicelib.base_app_server import BaseAppServer
 
 from .task_manager import CeleryTaskManager
 
-_WORKER_KEY = "celery_worker"
-_FASTAPI_APP_KEY = "fastapi_app"
+_APP_SERVER_KEY = "app_server"
+_TASK_MANAGER_KEY = "task_manager"
 
 
-def set_task_manager(celery_app: Celery, worker: CeleryTaskManager) -> None:
-    celery_app.conf[_WORKER_KEY] = worker
+def get_app_server(app: Celery) -> BaseAppServer:
+    app_server = app.conf[_APP_SERVER_KEY]
+    assert isinstance(app_server, BaseAppServer)
+    return app_server
+
+
+def set_app_server(app: Celery, app_server: BaseAppServer) -> None:
+    app.conf[_APP_SERVER_KEY] = app_server
 
 
 def get_task_manager(celery_app: Celery) -> CeleryTaskManager:
-    worker = celery_app.conf[_WORKER_KEY]
+    worker = celery_app.conf[_TASK_MANAGER_KEY]
     assert isinstance(worker, CeleryTaskManager)
     return worker
 
 
-def set_fastapi_app(celery_app: Celery, fastapi_app: FastAPI) -> None:
-    celery_app.conf[_FASTAPI_APP_KEY] = fastapi_app
-
-
-def get_fastapi_app(celery_app: Celery) -> FastAPI:
-    fastapi_app = celery_app.conf[_FASTAPI_APP_KEY]
-    assert isinstance(fastapi_app, FastAPI)
-    return fastapi_app
+def set_task_manager(celery_app: Celery, worker: CeleryTaskManager) -> None:
+    celery_app.conf[_TASK_MANAGER_KEY] = worker

packages/service-library/src/servicelib/base_app_server.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from asyncio import AbstractEventLoop
+from contextlib import suppress
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    with suppress(ImportError):
+        from fastapi import FastAPI
+    with suppress(ImportError):
+        from aiohttp.web import Application
+
+
+class BaseAppServer(ABC):
+    @property
+    def fastapi_app(self) -> "FastAPI":
+        raise NotImplementedError
+
+    @property
+    def aiohttp_app(self) -> "Application":
+        raise NotImplementedError
+
+    @abstractmethod
+    async def startup(self):
+        pass
+
+    @property
+    def event_loop(self) -> AbstractEventLoop:
+        return self._event_loop
+
+    @event_loop.setter
+    def event_loop(self, loop: AbstractEventLoop) -> None:
+        self._event_loop = loop
+
+    @abstractmethod
+    async def shutdown(self):
+        pass

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

@@ -0,0 +1,38 @@
+import asyncio
+from datetime import timedelta
+from typing import Final
+
+from asgi_lifespan import LifespanManager
+from fastapi import FastAPI
+
+from ..base_app_server import BaseAppServer
+
+_SHUTDOWN_TIMEOUT: Final[float] = timedelta(seconds=10).total_seconds()
+_STARTUP_TIMEOUT: Final[float] = timedelta(minutes=1).total_seconds()
+
+
+class FastAPIAppServer(BaseAppServer):
+    def __init__(self, app: FastAPI):
+        self._app = app
+        self._lifespan_manager = None
+        self._shutdown_event = asyncio.Event()
+
+    @property
+    def fastapi_app(self) -> FastAPI:
+        assert isinstance(self._app, FastAPI)  # nosec
+        return self._app
+
+    async def startup(self):
+        self._lifespan_manager = LifespanManager(
+            self.fastapi_app,
+            startup_timeout=_STARTUP_TIMEOUT,
+            shutdown_timeout=_SHUTDOWN_TIMEOUT,
+        )
+        await self._lifespan_manager.__aenter__()
+        await self._shutdown_event.wait()
+
+    async def shutdown(self):
+        self._shutdown_event.set()
+        if self._lifespan_manager is None:
+            return
+        await self._lifespan_manager.__aexit__(None, None, None)

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

@@ -2,7 +2,7 @@
 
 from celery import Task  # type: ignore[import-untyped]
 from celery_library.models import TaskID
-from celery_library.utils import get_fastapi_app
+from celery_library.utils import get_app_server
 from models_library.api_schemas_storage.storage_schemas import (
     FileUploadCompletionBody,
 )
@@ -30,7 +30,7 @@ async def complete_upload_file(
         logging.INFO,
         msg=f"completing upload of file {user_id=}, {location_id=}, {file_id=}",
     ):
-        dsm = get_dsm_provider(get_fastapi_app(task.app)).get(location_id)
+        dsm = get_dsm_provider(get_app_server(task.app).fastapi_app).get(location_id)
         # NOTE: completing a multipart upload on AWS can take up to several minutes
         # if it returns slow we return a 202 - Accepted, the client will have to check later
         # for completeness

services/storage/src/simcore_service_storage/api/_worker_tasks/_paths.py

@@ -3,7 +3,7 @@
 
 from celery import Task  # type: ignore[import-untyped]
 from celery_library.models import TaskID
-from celery_library.utils import get_fastapi_app
+from celery_library.utils import get_app_server
 from models_library.projects_nodes_io import LocationID, StorageFileID
 from models_library.users import UserID
 from pydantic import ByteSize, TypeAdapter
@@ -25,7 +25,7 @@ async def compute_path_size(
         logging.INFO,
         msg=f"computing path size {user_id=}, {location_id=}, {path=}",
     ):
-        dsm = get_dsm_provider(get_fastapi_app(task.app)).get(location_id)
+        dsm = get_dsm_provider(get_app_server(task.app).fastapi_app).get(location_id)
         return await dsm.compute_path_size(user_id, path=Path(path))
 
 
@@ -42,7 +42,7 @@ async def delete_paths(
         logging.INFO,
         msg=f"delete {paths=} in {location_id=} for {user_id=}",
     ):
-        dsm = get_dsm_provider(get_fastapi_app(task.app)).get(location_id)
+        dsm = get_dsm_provider(get_app_server(task.app).fastapi_app).get(location_id)
         files_ids: set[StorageFileID] = {
             TypeAdapter(StorageFileID).validate_python(f"{path}") for path in paths
         }

services/storage/src/simcore_service_storage/api/_worker_tasks/_simcore_s3.py

@@ -5,7 +5,7 @@
 from aws_library.s3._models import S3ObjectKey
 from celery import Task  # type: ignore[import-untyped]
 from celery_library.models import TaskID
-from celery_library.utils import get_fastapi_app, get_task_manager
+from celery_library.utils import get_app_server, get_task_manager
 from models_library.api_schemas_storage.storage_schemas import FoldersBody
 from models_library.api_schemas_webserver.storage import PathToExport
 from models_library.progress_bar import ProgressReport
@@ -40,7 +40,7 @@ async def deep_copy_files_from_project(
         logging.INFO,
         msg=f"copying {body.source['uuid']} -> {body.destination['uuid']} with {task.request.id}",
     ):
-        dsm = get_dsm_provider(get_fastapi_app(task.app)).get(
+        dsm = get_dsm_provider(get_app_server(task.app).fastapi_app).get(
             SimcoreS3DataManager.get_location_id()
         )
         assert isinstance(dsm, SimcoreS3DataManager)  # nosec
@@ -75,7 +75,7 @@ async def export_data(
         logging.INFO,
         f"'{task_id}' export data (for {user_id=}) fom selection: {paths_to_export}",
     ):
-        dsm = get_dsm_provider(get_fastapi_app(task.app)).get(
+        dsm = get_dsm_provider(get_app_server(task.app).fastapi_app).get(
             SimcoreS3DataManager.get_location_id()
         )
         assert isinstance(dsm, SimcoreS3DataManager)  # nosec

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

@@ -9,6 +9,7 @@
     on_worker_init,
     on_worker_shutdown,
 )
+from servicelib.fastapi.app_server import FastAPIAppServer
 from servicelib.logging_utils import config_all_loggers
 from simcore_service_storage.api._worker_tasks.tasks import setup_worker_tasks
 
@@ -28,11 +29,12 @@
 
 assert _settings.STORAGE_CELERY  # nosec
 app = create_celery_app(_settings.STORAGE_CELERY)
-app_factory = partial(create_app, _settings)
+
+app_server = FastAPIAppServer(app=create_app(_settings))
 
 
 def worker_init_wrapper(sender, **_kwargs):
-    return partial(on_worker_init, app_factory, _settings.STORAGE_CELERY)(
+    return partial(on_worker_init, app_server, _settings.STORAGE_CELERY)(
         sender, **_kwargs
     )