🐛 Use Celery prefork pool for CPU-bound tasks

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

@@ -2,32 +2,19 @@
 import logging
 import threading
 
-from celery import Celery  # type: ignore[import-untyped]
-from celery.worker.worker import WorkController  # type: ignore[import-untyped]
 from servicelib.celery.app_server import BaseAppServer
 from servicelib.logging_utils import log_context
 
-from .utils import get_app_server, set_app_server
-
 _logger = logging.getLogger(__name__)
 
 
-def on_worker_init(
-    sender: WorkController,
-    app_server: BaseAppServer,
-    **_kwargs,
-) -> None:
+def on_worker_init(app_server: BaseAppServer, **_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)
 
-        assert sender.app  # nosec
-        assert isinstance(sender.app, Celery)  # nosec
-
-        set_app_server(sender.app, app_server)
-
         app_server.event_loop = loop
 
         loop.run_until_complete(app_server.run_until_shutdown(startup_complete_event))
@@ -44,9 +31,6 @@ def _init(startup_complete_event: threading.Event) -> None:
     startup_complete_event.wait()
 
 
-def on_worker_shutdown(sender, **_kwargs) -> None:
+def on_worker_shutdown(app_server: BaseAppServer, **_kwargs) -> None:
     with log_context(_logger, logging.INFO, "Worker shutdown"):
-        assert isinstance(sender.app, Celery)
-        app_server = get_app_server(sender.app)
-
         app_server.shutdown_event.set()

packages/celery-library/src/celery_library/task.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/worker/app.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/worker/signals.py

@@ -0,0 +1,59 @@
+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 CelerySettings
+
+from .app_server import get_app_server, set_app_server
+
+
+def register_worker_signals(
+    app: Celery,
+    settings: CelerySettings,
+    app_server_factory: Callable[[], BaseAppServer],
+) -> None:
+    def _worker_init_wrapper(**_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()
+
+    def _worker_shutdown_wrapper(**_kwargs) -> None:
+        get_app_server(app).shutdown_event.set()
+
+    match settings.CELERY_POOL:
+        case "prefork":
+            worker_process_init.connect(_worker_init_wrapper, weak=False)
+            worker_process_shutdown.connect(_worker_shutdown_wrapper, weak=False)
+        case _:
+            worker_init.connect(_worker_init_wrapper, weak=False)
+            worker_shutdown.connect(_worker_shutdown_wrapper, weak=False)

packages/celery-library/tests/conftest.py

@@ -14,11 +14,11 @@
     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.app_server import set_app_server
 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
@@ -131,11 +131,15 @@ async def with_celery_worker(
     app_server: BaseAppServer,
     register_celery_tasks: Callable[[Celery], None],
 ) -> AsyncIterator[TestWorkController]:
-    def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
-        return on_worker_init(sender, app_server, **_kwargs)
+    def _on_worker_init_wrapper(**kwargs):
+        set_app_server(celery_app, app_server)
+        return on_worker_init(app_server, **kwargs)
+
+    def _on_worker_shutdown_wrapper(**kwargs):
+        return on_worker_shutdown(app_server, **kwargs)
 
     worker_init.connect(_on_worker_init_wrapper)
-    worker_shutdown.connect(on_worker_shutdown)
+    worker_shutdown.connect(_on_worker_shutdown_wrapper)
 
     register_celery_tasks(celery_app)
 

packages/celery-library/tests/unit/test_task_manager.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/service-library/src/servicelib/fastapi/celery/app_server.py

@@ -30,6 +30,6 @@ async def run_until_shutdown(
             startup_timeout=None,  # waits for full app initialization (DB migrations, etc.)
             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

packages/settings-library/src/settings_library/celery.py

@@ -1,5 +1,5 @@
 from datetime import timedelta
-from typing import Annotated
+from typing import Annotated, Literal
 
 from pydantic import Field
 from pydantic_settings import SettingsConfigDict
@@ -35,6 +35,13 @@ class CelerySettings(BaseCustomSettings):
         ),
     ] = True
 
+    CELERY_POOL: Annotated[
+        Literal["prefork", "eventlet", "gevent", "solo", "threads"],
+        Field(
+            description="Type of pool to use. One of: prefork, eventlet, gevent, solo, threads. See https://docs.celeryq.dev/en/stable/userguide/concurrency/index.html for details.",
+        ),
+    ] = "prefork"
+
     model_config = SettingsConfigDict(
         json_schema_extra={
             "examples": [

services/api-server/docker/boot.sh

@@ -48,18 +48,16 @@ if [ "${API_SERVER_WORKER_MODE}" = "true" ]; then
       --recursive \
       -- \
       celery \
-      --app=boot_celery_worker:app \
-      --workdir=services/api-server/docker \
-      worker --pool=threads \
+      --app=simcore_service_api_server.modules.celery.worker.main:app \
+      worker --pool="${CELERY_POOL}" \
       --loglevel="${API_SERVER_LOGLEVEL}" \
       --concurrency="${CELERY_CONCURRENCY}" \
       --hostname="${API_SERVER_WORKER_NAME}" \
       --queues="${CELERY_QUEUES:-default}"
   else
     exec celery \
-      --app=boot_celery_worker:app \
-      --workdir=services/api-server/docker \
-      worker --pool=threads \
+      --app=simcore_service_api_server.modules.celery.worker.main:app \
+      worker --pool="${CELERY_POOL}" \
       --loglevel="${API_SERVER_LOGLEVEL}" \
       --concurrency="${CELERY_CONCURRENCY}" \
       --hostname="${API_SERVER_WORKER_NAME}" \

services/api-server/src/simcore_service_api_server/clients/celery_task_manager.py

@@ -1,7 +1,7 @@
 import logging
 
+from celery_library.app import create_app
 from celery_library.backends.redis import RedisTaskStore
-from celery_library.common import create_app
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.types import register_celery_types, register_pydantic_types
 from fastapi import FastAPI

services/api-server/src/simcore_service_api_server/celery_worker/worker_tasks/functions_tasks.py → services/api-server/src/simcore_service_api_server/modules/celery/worker/_functions_tasks.py

@@ -1,17 +1,17 @@
 from celery import (  # type: ignore[import-untyped] # pylint: disable=no-name-in-module
     Task,
 )
-from celery_library.utils import get_app_server  # pylint: disable=no-name-in-module
+from celery_library.worker.app_server import get_app_server
 from fastapi import FastAPI
 from models_library.functions import RegisteredFunction, RegisteredFunctionJob
 from models_library.projects import ProjectID
 from models_library.projects_nodes_io import NodeID
 from servicelib.celery.models import TaskKey
 from simcore_service_api_server._service_function_jobs import FunctionJobService
 
-from ...api.dependencies.authentication import Identity
-from ...api.dependencies.rabbitmq import get_rabbitmq_rpc_client
-from ...api.dependencies.services import (
+from ....api.dependencies.authentication import Identity
+from ....api.dependencies.rabbitmq import get_rabbitmq_rpc_client
+from ....api.dependencies.services import (
     get_catalog_service,
     get_directorv2_service,
     get_function_job_service,
@@ -20,13 +20,16 @@
     get_solver_service,
     get_storage_service,
 )
-from ...api.dependencies.webserver_http import get_session_cookie, get_webserver_session
-from ...api.dependencies.webserver_rpc import get_wb_api_rpc_client
-from ...models.api_resources import JobLinks
-from ...models.domain.functions import PreRegisteredFunctionJobData
-from ...models.schemas.jobs import JobPricingSpecification
-from ...services_http.director_v2 import DirectorV2Api
-from ...services_http.storage import StorageApi
+from ....api.dependencies.webserver_http import (
+    get_session_cookie,
+    get_webserver_session,
+)
+from ....api.dependencies.webserver_rpc import get_wb_api_rpc_client
+from ....models.api_resources import JobLinks
+from ....models.domain.functions import PreRegisteredFunctionJobData
+from ....models.schemas.jobs import JobPricingSpecification
+from ....services_http.director_v2 import DirectorV2Api
+from ....services_http.storage import StorageApi
 
 
 async def _assemble_function_job_service(

services/api-server/src/simcore_service_api_server/modules/celery/worker/main.py

@@ -0,0 +1,39 @@
+from celery_library.worker.app import create_worker_app
+from servicelib.fastapi.celery.app_server import FastAPIAppServer
+from servicelib.logging_utils import setup_loggers
+from servicelib.tracing import TracingConfig
+
+from ....core.application import create_app
+from ....core.settings import ApplicationSettings
+from .tasks import register_worker_tasks
+
+_settings = ApplicationSettings.create_from_envs()
+_tracing_settings = _settings.API_SERVER_TRACING
+_tracing_config = TracingConfig.create(
+    tracing_settings=_tracing_settings,
+    service_name="api-server-celery-worker",
+)
+
+
+def get_app():
+    setup_loggers(
+        log_format_local_dev_enabled=_settings.API_SERVER_LOG_FORMAT_LOCAL_DEV_ENABLED,
+        logger_filter_mapping=_settings.API_SERVER_LOG_FILTER_MAPPING,
+        tracing_config=_tracing_config,
+        log_base_level=_settings.log_level,
+        noisy_loggers=None,
+    )
+
+    def _app_server_factory() -> FastAPIAppServer:
+        fastapi_app = create_app(_settings, tracing_config=_tracing_config)
+        return FastAPIAppServer(app=fastapi_app)
+
+    assert _settings.API_SERVER_CELERY  # nosec
+    return create_worker_app(
+        _settings.API_SERVER_CELERY,
+        register_worker_tasks,
+        _app_server_factory,
+    )
+
+
+app = get_app()

services/api-server/src/simcore_service_api_server/celery_worker/worker_tasks/tasks.py → services/api-server/src/simcore_service_api_server/modules/celery/worker/tasks.py

@@ -7,13 +7,13 @@
 from celery_library.types import register_celery_types, register_pydantic_types
 from servicelib.logging_utils import log_context
 
-from ...models.domain.celery_models import pydantic_types_to_register
-from .functions_tasks import run_function
+from ....models.domain.celery_models import pydantic_types_to_register
+from ._functions_tasks import run_function
 
 _logger = logging.getLogger(__name__)
 
 
-def setup_worker_tasks(app: Celery) -> None:
+def register_worker_tasks(app: Celery) -> None:
     register_celery_types()
     register_pydantic_types(*pydantic_types_to_register)