🐛 Ensure proper Redis client shutdown in Celery (#8237)

Author: giancarloromeo

packages/celery-library/src/celery_library/backends/_redis.py renamed to packages/celery-library/src/celery_library/backends/redis.py

@@ -12,7 +12,7 @@
     TaskMetadata,
     TaskUUID,
 )
-from servicelib.redis import RedisClientSDK
+from servicelib.redis import RedisClientSDK, handle_redis_returns_union_types
 
 from ..utils import build_task_id_prefix
 
@@ -41,18 +41,24 @@ async def create_task(
         expiry: timedelta,
     ) -> None:
         task_key = _build_key(task_id)
-        await self._redis_client_sdk.redis.hset(
-            name=task_key,
-            key=_CELERY_TASK_METADATA_KEY,
-            value=task_metadata.model_dump_json(),
-        )  # type: ignore
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=task_key,
+                key=_CELERY_TASK_METADATA_KEY,
+                value=task_metadata.model_dump_json(),
+            )
+        )
         await self._redis_client_sdk.redis.expire(
             task_key,
             expiry,
         )
 
     async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
-        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_METADATA_KEY)  # type: ignore
+        raw_result = await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hget(
+                _build_key(task_id), _CELERY_TASK_METADATA_KEY
+            )
+        )
         if not raw_result:
             return None
 
@@ -65,7 +71,11 @@ async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
             return None
 
     async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
-        raw_result = await self._redis_client_sdk.redis.hget(_build_key(task_id), _CELERY_TASK_PROGRESS_KEY)  # type: ignore
+        raw_result = await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hget(
+                _build_key(task_id), _CELERY_TASK_PROGRESS_KEY
+            )
+        )
         if not raw_result:
             return None
 
@@ -121,11 +131,13 @@ async def remove_task(self, task_id: TaskID) -> None:
         await self._redis_client_sdk.redis.delete(_build_key(task_id))
 
     async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
-        await self._redis_client_sdk.redis.hset(
-            name=_build_key(task_id),
-            key=_CELERY_TASK_PROGRESS_KEY,
-            value=report.model_dump_json(),
-        )  # type: ignore
+        await handle_redis_returns_union_types(
+            self._redis_client_sdk.redis.hset(
+                name=_build_key(task_id),
+                key=_CELERY_TASK_PROGRESS_KEY,
+                value=report.model_dump_json(),
+            )
+        )
 
     async def task_exists(self, task_id: TaskID) -> bool:
         n = await self._redis_client_sdk.redis.exists(_build_key(task_id))

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

@@ -2,13 +2,9 @@
 from typing import Any
 
 from celery import Celery  # type: ignore[import-untyped]
-from servicelib.redis import RedisClientSDK
 from settings_library.celery import CelerySettings
 from settings_library.redis import RedisDatabase
 
-from .backends._redis import RedisTaskInfoStore
-from .task_manager import CeleryTaskManager
-
 
 def _celery_configure(celery_settings: CelerySettings) -> dict[str, Any]:
     base_config = {
@@ -36,22 +32,3 @@ def create_app(settings: CelerySettings) -> Celery:
         ),
         **_celery_configure(settings),
     )
-
-
-async def create_task_manager(
-    app: Celery, settings: CelerySettings
-) -> CeleryTaskManager:
-    redis_client_sdk = RedisClientSDK(
-        settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
-            RedisDatabase.CELERY_TASKS
-        ),
-        client_name="celery_tasks",
-    )
-    await redis_client_sdk.setup()
-    # GCR please address https://github.com/ITISFoundation/osparc-simcore/issues/8159
-
-    return CeleryTaskManager(
-        app,
-        settings,
-        RedisTaskInfoStore(redis_client_sdk),
-    )

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

@@ -6,18 +6,15 @@
 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 settings_library.celery import CelerySettings
 
-from .common import create_task_manager
 from .utils import get_app_server, set_app_server
 
 _logger = logging.getLogger(__name__)
 
 
 def on_worker_init(
-    app_server: BaseAppServer,
-    celery_settings: CelerySettings,
     sender: WorkController,
+    app_server: BaseAppServer,
     **_kwargs,
 ) -> None:
     startup_complete_event = threading.Event()
@@ -26,21 +23,14 @@ def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-        async def _setup_task_manager():
-            assert sender.app  # nosec
-            assert isinstance(sender.app, Celery)  # nosec
-
-            app_server.task_manager = await create_task_manager(
-                sender.app,
-                celery_settings,
-            )
+        assert sender.app  # nosec
+        assert isinstance(sender.app, Celery)  # nosec
 
-            set_app_server(sender.app, app_server)
+        set_app_server(sender.app, app_server)
 
         app_server.event_loop = loop
 
-        loop.run_until_complete(_setup_task_manager())
-        loop.run_until_complete(app_server.lifespan(startup_complete_event))
+        loop.run_until_complete(app_server.run_until_shutdown(startup_complete_event))
 
     thread = threading.Thread(
         group=None,

packages/celery-library/tests/conftest.py

@@ -2,25 +2,30 @@
 # pylint: disable=unused-argument
 
 import datetime
+import logging
 import threading
 from collections.abc import AsyncIterator, Callable
-from functools import partial
 from typing import Any
 
 import pytest
 from celery import Celery  # type: ignore[import-untyped]
-from celery.contrib.testing.worker import TestWorkController, start_worker
+from celery.contrib.testing.worker import (
+    TestWorkController,
+    start_worker,
+)
 from celery.signals import worker_init, worker_shutdown
 from celery.worker.worker import WorkController
-from celery_library.common import create_task_manager
+from celery_library.backends.redis import RedisTaskInfoStore
 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 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.redis import RedisSettings
+from settings_library.redis import RedisDatabase, RedisSettings
 
 pytest_plugins = [
     "pytest_simcore.docker_compose",
@@ -33,11 +38,42 @@
 ]
 
 
+_logger = logging.getLogger(__name__)
+
+
 class FakeAppServer(BaseAppServer):
-    async def lifespan(self, startup_completed_event: threading.Event) -> None:
+    def __init__(self, app: Celery, settings: CelerySettings):
+        super().__init__(app)
+        self._settings = settings
+        self._task_manager: CeleryTaskManager | None = None
+
+    @property
+    def task_manager(self) -> TaskManager:
+        assert self._task_manager, "Task manager is not initialized"
+        return self._task_manager
+
+    async def run_until_shutdown(
+        self, startup_completed_event: threading.Event
+    ) -> None:
+        redis_client_sdk = RedisClientSDK(
+            self._settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
+                RedisDatabase.CELERY_TASKS
+            ),
+            client_name="pytest_celery_tasks",
+        )
+        await redis_client_sdk.setup()
+
+        self._task_manager = CeleryTaskManager(
+            self._app,
+            self._settings,
+            RedisTaskInfoStore(redis_client_sdk),
+        )
+
         startup_completed_event.set()
         await self.shutdown_event.wait()  # wait for shutdown
 
+        await redis_client_sdk.shutdown()
+
 
 @pytest.fixture
 def register_celery_tasks() -> Callable[[Celery], None]:
@@ -51,17 +87,12 @@ def _(celery_app: Celery) -> None: ...
 @pytest.fixture
 def app_environment(
     monkeypatch: pytest.MonkeyPatch,
-    redis_service: RedisSettings,
     env_devel_dict: EnvVarsDict,
 ) -> EnvVarsDict:
     return setenvs_from_dict(
         monkeypatch,
         {
             **env_devel_dict,
-            "REDIS_SECURE": redis_service.REDIS_SECURE,
-            "REDIS_HOST": redis_service.REDIS_HOST,
-            "REDIS_PORT": f"{redis_service.REDIS_PORT}",
-            "REDIS_PASSWORD": redis_service.REDIS_PASSWORD.get_secret_value(),
         },
     )
 
@@ -74,8 +105,8 @@ def celery_settings(
 
 
 @pytest.fixture
-def app_server() -> BaseAppServer:
-    return FakeAppServer(app=None)
+def app_server(celery_app: Celery, celery_settings: CelerySettings) -> BaseAppServer:
+    return FakeAppServer(app=celery_app, settings=celery_settings)
 
 
 @pytest.fixture(scope="session")
@@ -98,11 +129,10 @@ def celery_config() -> dict[str, Any]:
 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 partial(on_worker_init, app_server, celery_settings)(sender, **_kwargs)
+        return on_worker_init(sender, app_server, **_kwargs)
 
     worker_init.connect(_on_worker_init_wrapper)
     worker_shutdown.connect(on_worker_shutdown)
@@ -111,24 +141,39 @@ def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
 
     with start_worker(
         celery_app,
-        pool="threads",
         concurrency=1,
+        pool="threads",
         loglevel="info",
         perform_ping_check=False,
         queues="default",
     ) as worker:
         yield worker
 
 
+@pytest.fixture
+async def mock_celery_app(celery_config: dict[str, Any]) -> Celery:
+    return Celery(**celery_config)
+
+
 @pytest.fixture
 async def celery_task_manager(
-    celery_app: Celery,
+    mock_celery_app: Celery,
     celery_settings: CelerySettings,
-    with_celery_worker: TestWorkController,
-) -> CeleryTaskManager:
+    use_in_memory_redis: RedisSettings,
+) -> AsyncIterator[CeleryTaskManager]:
     register_celery_types()
 
-    return await create_task_manager(
-        celery_app,
-        celery_settings,
-    )
+    try:
+        redis_client_sdk = RedisClientSDK(
+            use_in_memory_redis.build_redis_dsn(RedisDatabase.CELERY_TASKS),
+            client_name="pytest_celery_tasks",
+        )
+        await redis_client_sdk.setup()
+
+        yield CeleryTaskManager(
+            mock_celery_app,
+            celery_settings,
+            RedisTaskInfoStore(redis_client_sdk),
+        )
+    finally:
+        await redis_client_sdk.shutdown()

packages/celery-library/tests/unit/test_tasks.py

@@ -12,6 +12,7 @@
 
 import pytest
 from celery import Celery, Task  # pylint: disable=no-name-in-module
+from celery.worker.worker import WorkController  # pylint: disable=no-name-in-module
 from celery_library.errors import TaskNotFoundError, TransferrableCeleryError
 from celery_library.task import register_task
 from celery_library.task_manager import CeleryTaskManager
@@ -33,6 +34,10 @@
 pytest_simcore_ops_services_selection = []
 
 
+class MyTaskFilter(TaskFilter):
+    user_id: int
+
+
 async def _fake_file_processor(
     celery_app: Celery, task_name: str, task_id: str, files: list[str]
 ) -> str:
@@ -91,8 +96,9 @@ def _(celery_app: Celery) -> None:
 
 async def test_submitting_task_calling_async_function_results_with_success_state(
     celery_task_manager: CeleryTaskManager,
+    with_celery_worker: WorkController,
 ):
-    task_filter = TaskFilter(user_id=42)
+    task_filter = MyTaskFilter(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -121,8 +127,9 @@ async def test_submitting_task_calling_async_function_results_with_success_state
 
 async def test_submitting_task_with_failure_results_with_error(
     celery_task_manager: CeleryTaskManager,
+    with_celery_worker: WorkController,
 ):
-    task_filter = TaskFilter(user_id=42)
+    task_filter = MyTaskFilter(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -149,8 +156,9 @@ async def test_submitting_task_with_failure_results_with_error(
 
 async def test_cancelling_a_running_task_aborts_and_deletes(
     celery_task_manager: CeleryTaskManager,
+    with_celery_worker: WorkController,
 ):
-    task_filter = TaskFilter(user_id=42)
+    task_filter = MyTaskFilter(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -171,8 +179,9 @@ async def test_cancelling_a_running_task_aborts_and_deletes(
 
 async def test_listing_task_uuids_contains_submitted_task(
     celery_task_manager: CeleryTaskManager,
+    with_celery_worker: WorkController,
 ):
-    task_filter = TaskFilter(user_id=42)
+    task_filter = MyTaskFilter(user_id=42)
 
     task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
@@ -190,5 +199,5 @@ async def test_listing_task_uuids_contains_submitted_task(
             tasks = await celery_task_manager.list_tasks(task_filter)
             assert any(task.uuid == task_uuid for task in tasks)
 
-        tasks = await celery_task_manager.list_tasks(task_filter)
-        assert any(task.uuid == task_uuid for task in tasks)
+    tasks = await celery_task_manager.list_tasks(task_filter)
+    assert any(task.uuid == task_uuid for task in tasks)