fix tests

Author: giancarloromeo

services/storage/src/simcore_service_storage/api/rpc/_async_jobs.py

@@ -15,7 +15,7 @@
 )
 from servicelib.rabbitmq import RPCRouter
 
-from ...modules.celery._utils import get_celery_client
+from ...modules.celery import get_celery_client
 from ...modules.celery.models import TaskStatus
 
 router = RPCRouter()

services/storage/src/simcore_service_storage/api/rpc/_data_export.py

@@ -14,7 +14,7 @@
 from ...datcore_dsm import DatCoreDataManager
 from ...dsm import get_dsm_provider
 from ...exceptions.errors import FileAccessRightError
-from ...modules.celery._utils import get_celery_client
+from ...modules.celery import get_celery_client
 from ...modules.datcore_adapter.datcore_adapter_exceptions import DatcoreAdapterError
 from ...simcore_s3_dsm import SimcoreS3DataManager
 

services/storage/src/simcore_service_storage/main.py

@@ -3,15 +3,10 @@
 import logging
 
 from servicelib.logging_utils import config_all_loggers
-from simcore_service_storage.modules.celery._utils import (
-    set_celery_app,
-    set_celery_client,
-)
+from simcore_service_storage.modules.celery import setup_celery
 
 from .core.application import create_app
 from .core.settings import ApplicationSettings
-from .modules.celery._common import create_app as create_celery_app
-from .modules.celery.client import CeleryTaskQueueClient
 
 _settings = ApplicationSettings.create_from_envs()
 
@@ -27,9 +22,6 @@
 _logger = logging.getLogger(__name__)
 
 fastapi_app = create_app(_settings)
-celery_app = create_celery_app(_settings)
-
-set_celery_app(fastapi_app, celery_app)
-set_celery_client(fastapi_app, CeleryTaskQueueClient(celery_app))
+setup_celery(fastapi_app)
 
 app = fastapi_app

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

@@ -0,0 +1,36 @@
+import logging
+from asyncio import AbstractEventLoop
+
+from fastapi import FastAPI
+from simcore_service_storage.modules.celery._common import create_app
+from simcore_service_storage.modules.celery.client import CeleryTaskQueueClient
+
+from ...core.settings import get_application_settings
+
+_logger = logging.getLogger(__name__)
+
+
+def setup_celery(app: FastAPI) -> None:
+    async def on_startup() -> None:
+        celery_settings = get_application_settings(app).STORAGE_CELERY
+        assert celery_settings  # nosec
+        celery_app = create_app(celery_settings)
+        app.state.celery_client = CeleryTaskQueueClient(celery_app)
+
+    app.add_event_handler("startup", on_startup)
+
+
+def get_celery_client(app: FastAPI) -> CeleryTaskQueueClient:
+    celery_client = app.state.celery_client
+    assert isinstance(celery_client, CeleryTaskQueueClient)
+    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

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

@@ -1,15 +1,13 @@
 import logging
 
 from celery import Celery
+from settings_library.celery import CelerySettings
 from settings_library.redis import RedisDatabase
 
-from ...core.settings import ApplicationSettings
-
 _logger = logging.getLogger(__name__)
 
 
-def create_app(settings: ApplicationSettings) -> Celery:
-    celery_settings = settings.STORAGE_CELERY
+def create_app(celery_settings: CelerySettings) -> Celery:
     assert celery_settings
 
     app = Celery(

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

@@ -22,6 +22,14 @@
 )
 _CELERY_TASK_META_PREFIX: Final[str] = "celery-task-meta-"
 _CELERY_TASK_ID_PREFIX: Final[str] = "celery"
+_CELERY_STATES_MAPPING: Final[dict[str, TaskState]] = {
+    "PENDING": TaskState.PENDING,
+    "STARTED": TaskState.PENDING,
+    "RUNNING": TaskState.RUNNING,
+    "SUCCESS": TaskState.SUCCESS,
+    "ABORTED": TaskState.ABORTED,
+    "FAILURE": TaskState.FAILURE,
+}
 
 
 def _build_context_prefix(task_context: TaskContext) -> list[str]:
@@ -80,26 +88,29 @@ def _get_progress_report(
     ) -> ProgressReport:
         task_id = _build_task_id(task_context, task_uuid)
         result = self._celery_app.AsyncResult(task_id).result
-        state = self._celery_app.AsyncResult(task_id).state
-        if result and state == TaskState.RUNNING.value:
+        state = self._get_state(task_context, task_uuid)
+        if result and state == TaskState.RUNNING:
             with contextlib.suppress(ValidationError):
                 return ProgressReport.model_validate(result)
         if state in (
-            TaskState.ABORTED.value,
-            TaskState.FAILURE.value,
-            TaskState.SUCCESS.value,
+            TaskState.ABORTED,
+            TaskState.FAILURE,
+            TaskState.SUCCESS,
         ):
             return ProgressReport(actual_value=100.0)
         return ProgressReport(actual_value=0.0)
 
+    def _get_state(self, task_context: TaskContext, task_uuid: TaskUUID) -> TaskState:
+        task_id = _build_task_id(task_context, task_uuid)
+        return _CELERY_STATES_MAPPING[self._celery_app.AsyncResult(task_id).state]
+
     @make_async()
     def get_task_status(
         self, task_context: TaskContext, task_uuid: TaskUUID
     ) -> TaskStatus:
-        task_id = _build_task_id(task_context, task_uuid)
         return TaskStatus(
             task_uuid=task_uuid,
-            task_state=self._celery_app.AsyncResult(task_id).state,
+            task_state=self._get_state(task_context, task_uuid),
             progress_report=self._get_progress_report(task_context, task_uuid),
         )
 

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

@@ -0,0 +1,27 @@
+from celery import Celery
+from fastapi import FastAPI
+
+from .worker import CeleryTaskQueueWorker
+
+_WORKER_KEY = "celery_worker"
+_FASTAPI_APP_KEY = "fastapi_app"
+
+
+def get_celery_worker(celery_app: Celery) -> CeleryTaskQueueWorker:
+    worker = celery_app.conf[_WORKER_KEY]
+    assert isinstance(worker, CeleryTaskQueueWorker)
+    return worker
+
+
+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_celery_worker(celery_app: Celery, worker: CeleryTaskQueueWorker) -> None:
+    celery_app.conf[_WORKER_KEY] = worker
+
+
+def set_fastapi_app(celery_app: Celery, fastapi_app: FastAPI) -> None:
+    celery_app.conf[_FASTAPI_APP_KEY] = fastapi_app

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

@@ -6,7 +6,7 @@
 from models_library.progress_bar import ProgressReport
 from servicelib.logging_utils import log_context
 
-from .models import TaskID, TaskState
+from .models import TaskID
 
 _logger = logging.getLogger(__name__)
 
@@ -30,6 +30,6 @@ def set_task_progress(
         ):
             self.celery_app.tasks[task_name].update_state(
                 task_id=task_id,
-                state=TaskState.RUNNING.value,
+                state="RUNNING",
                 meta=report.model_dump(mode="json"),
             )

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

@@ -3,16 +3,25 @@
 import asyncio
 import logging
 import threading
+from typing import Final
 
 from asgi_lifespan import LifespanManager
 from celery import Celery
 from celery.signals import worker_init, worker_shutdown
+from fastapi import FastAPI
 from servicelib.background_task import cancel_wait_task
 from servicelib.logging_utils import config_all_loggers
+from simcore_service_storage.modules.celery import get_event_loop, set_event_loop
+from simcore_service_storage.modules.celery.utils import (
+    CeleryTaskQueueWorker,
+    get_fastapi_app,
+    set_celery_worker,
+    set_fastapi_app,
+)
 
 from ...core.application import create_app
 from ...core.settings import ApplicationSettings
-from ._utils import create_celery_app_worker
+from ._common import create_app as create_celery_app
 
 _settings = ApplicationSettings.create_from_envs()
 
@@ -26,6 +35,8 @@
 
 _logger = logging.getLogger(__name__)
 
+_LIFESPAN_TIMEOUT: Final[int] = 10
+
 
 @worker_init.connect
 def on_worker_init(sender, **_kwargs):
@@ -39,8 +50,8 @@ def _init_fastapi():
         async def lifespan():
             async with LifespanManager(
                 fastapi_app,
-                startup_timeout=10,
-                shutdown_timeout=10,
+                startup_timeout=_LIFESPAN_TIMEOUT,
+                shutdown_timeout=_LIFESPAN_TIMEOUT,
             ):
                 try:
                     await shutdown_event.wait()
@@ -50,9 +61,10 @@ async def lifespan():
         lifespan_task = loop.create_task(lifespan())
         fastapi_app.state.lifespan_task = lifespan_task
         fastapi_app.state.shutdown_event = shutdown_event
+        set_event_loop(fastapi_app, loop)
 
-        sender.app.conf["fastapi_app"] = fastapi_app
-        sender.app.conf["loop"] = loop
+        set_fastapi_app(sender.app, fastapi_app)
+        set_celery_worker(sender.app, CeleryTaskQueueWorker(sender.app))
 
         loop.run_forever()
 
@@ -64,15 +76,17 @@ async def lifespan():
 def on_worker_shutdown(sender, **_kwargs):
     assert isinstance(sender.app, Celery)
 
-    loop = sender.app.conf["loop"]
-    fastapi_app = sender.app.conf["fastapi_app"]
+    fastapi_app = get_fastapi_app(sender.app)
+    assert isinstance(fastapi_app, FastAPI)
+    event_loop = get_event_loop(fastapi_app)
 
     async def shutdown():
         fastapi_app.state.shutdown_event.set()
 
         await cancel_wait_task(fastapi_app.state.lifespan_task, max_delay=5)
 
-    asyncio.run_coroutine_threadsafe(shutdown(), loop)
+    asyncio.run_coroutine_threadsafe(shutdown(), event_loop)
 
 
-app = create_celery_app_worker(_settings)
+assert _settings.STORAGE_CELERY
+app = create_celery_app(_settings.STORAGE_CELERY)