Merge branch 'master' into feature/search-files

Author: odeimaiz

packages/celery-library/src/celery_library/common.py renamed to packages/celery-library/src/celery_library/app.py

@@ -13,7 +13,7 @@
 from servicelib.celery.models import TaskKey
 
 from .errors import encode_celery_transferrable_error
-from .utils import get_app_server
+from .worker.app_server import get_app_server
 
 _logger = logging.getLogger(__name__)
 

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

@@ -0,0 +1,20 @@
+from collections.abc import Callable
+
+from celery import Celery  # type: ignore[import-untyped]
+from servicelib.celery.app_server import BaseAppServer
+from settings_library.celery import CelerySettings
+
+from ..app import create_app
+from .signals import register_worker_signals
+
+
+def create_worker_app(
+    settings: CelerySettings,
+    register_worker_tasks_cb: Callable[[Celery], None],
+    app_server_factory_cb: Callable[[], BaseAppServer],
+) -> Celery:
+    app = create_app(settings)
+    register_worker_tasks_cb(app)
+    register_worker_signals(app, settings, app_server_factory_cb)
+
+    return app

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

@@ -0,0 +1,71 @@
+import asyncio
+import threading
+from collections.abc import Callable
+
+from celery import Celery  # type: ignore[import-untyped]
+from celery.signals import (  # type: ignore[import-untyped]
+    worker_init,
+    worker_process_init,
+    worker_process_shutdown,
+    worker_shutdown,
+)
+from servicelib.celery.app_server import BaseAppServer
+from settings_library.celery import CeleryPoolType, CelerySettings
+
+from .app_server import get_app_server, set_app_server
+
+
+def _worker_init_wrapper(
+    app: Celery, app_server_factory: Callable[[], BaseAppServer]
+) -> Callable[..., None]:
+    def _worker_init_handler(**_kwargs) -> None:
+        startup_complete_event = threading.Event()
+
+        def _init(startup_complete_event: threading.Event) -> None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            app_server = app_server_factory()
+            app_server.event_loop = loop
+
+            set_app_server(app, app_server)
+
+            loop.run_until_complete(
+                app_server.run_until_shutdown(startup_complete_event)
+            )
+
+        thread = threading.Thread(
+            group=None,
+            target=_init,
+            name="app_server_init",
+            args=(startup_complete_event,),
+            daemon=True,
+        )
+        thread.start()
+
+        startup_complete_event.wait()
+
+    return _worker_init_handler
+
+
+def _worker_shutdown_wrapper(app: Celery) -> Callable[..., None]:
+    def _worker_shutdown_handler(**_kwargs) -> None:
+        get_app_server(app).shutdown_event.set()
+
+    return _worker_shutdown_handler
+
+
+def register_worker_signals(
+    app: Celery,
+    settings: CelerySettings,
+    app_server_factory: Callable[[], BaseAppServer],
+) -> None:
+    match settings.CELERY_POOL:
+        case CeleryPoolType.PREFORK:
+            worker_process_init.connect(
+                _worker_init_wrapper(app, app_server_factory), weak=False
+            )
+            worker_process_shutdown.connect(_worker_shutdown_wrapper(app), weak=False)
+        case _:
+            worker_init.connect(_worker_init_wrapper(app, app_server_factory), weak=False)
+            worker_shutdown.connect(_worker_shutdown_wrapper(app), weak=False)

services/api-server/src/simcore_service_api_server/celery_worker/__init__.py renamed to packages/celery-library/src/celery_library/worker/__init__.py

@@ -14,17 +14,16 @@
     start_worker,
 )
 from celery.signals import worker_init, worker_shutdown
-from celery.worker.worker import WorkController
 from celery_library.backends.redis import RedisTaskStore
-from celery_library.signals import on_worker_init, on_worker_shutdown
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.types import register_celery_types
+from celery_library.worker.signals import _worker_init_wrapper, _worker_shutdown_wrapper
 from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
 from pytest_simcore.helpers.typing_env import EnvVarsDict
 from servicelib.celery.app_server import BaseAppServer
 from servicelib.celery.task_manager import TaskManager
 from servicelib.redis import RedisClientSDK
-from settings_library.celery import CelerySettings
+from settings_library.celery import CeleryPoolType, CelerySettings
 from settings_library.redis import RedisDatabase, RedisSettings
 
 pytest_plugins = [
@@ -104,11 +103,6 @@ def celery_settings(
     return CelerySettings.create_from_envs()
 
 
-@pytest.fixture
-def app_server(celery_app: Celery, celery_settings: CelerySettings) -> BaseAppServer:
-    return FakeAppServer(app=celery_app, settings=celery_settings)
-
-
 @pytest.fixture(scope="session")
 def celery_config() -> dict[str, Any]:
     return {
@@ -128,21 +122,25 @@ def celery_config() -> dict[str, Any]:
 @pytest.fixture
 async def with_celery_worker(
     celery_app: Celery,
-    app_server: BaseAppServer,
+    celery_settings: CelerySettings,
     register_celery_tasks: Callable[[Celery], None],
 ) -> AsyncIterator[TestWorkController]:
-    def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
-        return on_worker_init(sender, app_server, **_kwargs)
 
-    worker_init.connect(_on_worker_init_wrapper)
-    worker_shutdown.connect(on_worker_shutdown)
+    def _app_server_factory() -> BaseAppServer:
+        return FakeAppServer(app=celery_app, settings=celery_settings)
+
+    # NOTE: explicitly connect the signals in tests
+    worker_init.connect(
+        _worker_init_wrapper(celery_app, _app_server_factory), weak=False
+    )
+    worker_shutdown.connect(_worker_shutdown_wrapper(celery_app), weak=False)
 
     register_celery_tasks(celery_app)
 
     with start_worker(
         celery_app,
         concurrency=1,
-        pool="threads",
+        pool=CeleryPoolType.THREADS,
         loglevel="info",
         perform_ping_check=False,
         queues="default",

packages/celery-library/src/celery_library/worker/app.py

@@ -16,7 +16,7 @@
 from celery_library.errors import TaskNotFoundError, TransferrableCeleryError
 from celery_library.task import register_task
 from celery_library.task_manager import CeleryTaskManager
-from celery_library.utils import get_app_server
+from celery_library.worker.app_server import get_app_server
 from common_library.errors_classes import OsparcErrorMixin
 from faker import Faker
 from models_library.progress_bar import ProgressReport

packages/celery-library/src/celery_library/utils.py renamed to packages/celery-library/src/celery_library/worker/app_server.py

@@ -9,6 +9,7 @@
 from ...celery.app_server import BaseAppServer
 from ...celery.task_manager import TaskManager
 
+_STARTUP_TIMEOUT: Final[float] = datetime.timedelta(minutes=5).total_seconds()
 _SHUTDOWN_TIMEOUT: Final[float] = datetime.timedelta(seconds=10).total_seconds()
 
 _logger = logging.getLogger(__name__)
@@ -27,9 +28,10 @@ async def run_until_shutdown(
     ) -> None:
         async with LifespanManager(
             self.app,
-            startup_timeout=None,  # waits for full app initialization (DB migrations, etc.)
+            startup_timeout=_STARTUP_TIMEOUT,
             shutdown_timeout=_SHUTDOWN_TIMEOUT,
         ):
-            _logger.info("fastapi app initialized")
+            _logger.info("FastAPI initialized: %s", self.app)
             startup_completed_event.set()
             await self.shutdown_event.wait()  # NOTE: wait here until shutdown is requested
+            _logger.info("FastAPI shutdown completed: %s", self.app)