fix: celery task manager setup

Author: giancarloromeo

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

@@ -1,18 +1,48 @@
-from celery_library.common import create_app, create_task_manager
+import logging
+
+from celery_library.backends._redis import RedisTaskInfoStore
+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
+from servicelib.logging_utils import log_context
+from servicelib.redis import RedisClientSDK
 from settings_library.celery import CelerySettings
+from settings_library.redis import RedisDatabase
 
 from ..celery_worker.worker_tasks.tasks import pydantic_types_to_register
 
+_logger = logging.getLogger(__name__)
+
 
-def setup_task_manager(app: FastAPI, celery_settings: CelerySettings) -> None:
+def setup_task_manager(app: FastAPI, settings: CelerySettings) -> None:
     async def on_startup() -> None:
-        app.state.task_manager = await create_task_manager(
-            create_app(celery_settings), celery_settings
-        )
+        with log_context(_logger, logging.INFO, "Setting up Celery"):
+            redis_client_sdk = RedisClientSDK(
+                settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
+                    RedisDatabase.CELERY_TASKS
+                ),
+                client_name="api_server_celery_tasks",
+            )
+            app.state.celery_tasks_redis_client_sdk = redis_client_sdk
+            await redis_client_sdk.setup()
+
+            app.state.task_manager = CeleryTaskManager(
+                create_app(settings),
+                settings,
+                RedisTaskInfoStore(redis_client_sdk),
+            )
+
+            register_celery_types()
+            register_pydantic_types(*pydantic_types_to_register)
 
-        register_celery_types()
-        register_pydantic_types(*pydantic_types_to_register)
+    async def on_shutdown() -> None:
+        with log_context(_logger, logging.INFO, "Shutting down Celery"):
+            redis_client_sdk: RedisClientSDK | None = (
+                app.state.celery_tasks_redis_client_sdk
+            )
+            if redis_client_sdk:
+                await redis_client_sdk.shutdown()
 
     app.add_event_handler("startup", on_startup)
+    app.add_event_handler("shutdown", on_shutdown)

services/storage/src/simcore_service_storage/core/application.py

@@ -36,7 +36,7 @@
 from ..dsm import setup_dsm
 from ..dsm_cleaner import setup_dsm_cleaner
 from ..exceptions.handlers import set_exception_handlers
-from ..modules.celery import setup_celery
+from ..modules.celery import setup_task_manager
 from ..modules.db import setup_db
 from ..modules.rabbitmq import setup as setup_rabbitmq
 from ..modules.redis import setup as setup_redis
@@ -71,7 +71,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:  # noqa: C901
     setup_client_session(app, tracing_settings=settings.STORAGE_TRACING)
 
     if settings.STORAGE_CELERY:
-        setup_celery(app, settings=settings.STORAGE_CELERY)
+        setup_task_manager(app, settings=settings.STORAGE_CELERY)
 
         if not settings.STORAGE_WORKER_MODE:
             setup_rabbitmq(app)

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

@@ -19,14 +19,14 @@
 _logger = logging.getLogger(__name__)
 
 
-def setup_celery(app: FastAPI, settings: CelerySettings) -> None:
+def setup_task_manager(app: FastAPI, settings: CelerySettings) -> None:
     async def on_startup() -> None:
         with log_context(_logger, logging.INFO, "Setting up Celery"):
             redis_client_sdk = RedisClientSDK(
                 settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
                     RedisDatabase.CELERY_TASKS
                 ),
-                client_name="celery_tasks",
+                client_name="storage_celery_tasks",
             )
             app.state.celery_tasks_redis_client_sdk = redis_client_sdk
             await redis_client_sdk.setup()