add progress sending via rabbitmq

Author: Andrei Neagu

packages/models-library/src/models_library/osparc_jobs.py

@@ -0,0 +1,3 @@
+from typing import TypeAlias
+
+OsparcJobId: TypeAlias = str

packages/models-library/src/models_library/rabbitmq_messages.py

@@ -6,6 +6,7 @@
 from typing import Any, Literal, TypeAlias
 
 import arrow
+from models_library.osparc_jobs import OsparcJobId
 from pydantic import BaseModel, Field
 
 from .products import ProductName
@@ -123,6 +124,8 @@ def routing_key(self) -> str | None:
 
 
 class ProgressRabbitMessageWorkerJob(ProgressMessageMixin, WorkerJobMessageBase):
+    osparc_job_id: OsparcJobId
+
     def routing_key(self) -> str | None:
         return f"{self.user_id}.worker_job"
 

services/storage/src/simcore_service_storage/api/celery/_export_task.py

@@ -1,22 +1,22 @@
 from asyncio import AbstractEventLoop, get_event_loop, run_coroutine_threadsafe
-from typing import TypeAlias, cast
+from typing import cast
 from uuid import uuid4
 
 from fastapi import FastAPI
+from models_library.osparc_jobs import OsparcJobId
 from models_library.projects_nodes_io import StorageFileID
 from models_library.users import UserID
 from servicelib.progress_bar import ProgressReport, ReportCB
 
 from ...dsm import SimcoreS3DataManager, get_dsm_provider
+from ...modules.rabbitmq import post_task_progress_message
 from ._tqdm_utils import get_export_progress, set_absolute_progress
 
-OsparcJobID: TypeAlias = str
-
 
 async def _async_export(
     app: FastAPI,
     *,
-    osparc_job_id: OsparcJobID,
+    osparc_job_id: OsparcJobId,
     user_id: UserID,
     paths_to_export: list[StorageFileID],
     progress_cb: ReportCB,
@@ -26,9 +26,7 @@ async def _async_export(
         get_dsm_provider(app).get(SimcoreS3DataManager.get_location_id()),
     )
 
-    with get_export_progress(
-        total=1, description=f"export job: {osparc_job_id}"
-    ) as pbar:
+    with get_export_progress(total=1, description=f"{osparc_job_id}") as pbar:
 
         def _sync_progress_updates(report: ProgressReport) -> None:
             set_absolute_progress(pbar, current_progress=report.actual_value)
@@ -38,7 +36,7 @@ def _sync_progress_updates(report: ProgressReport) -> None:
 
         async def _progress_cb(report: ProgressReport) -> None:
             await get_event_loop().run_in_executor(None, _sync_progress_updates, report)
-            # TODO: propagate progress via socketio -> also include `osparc_job_id`
+            await post_task_progress_message(app, user_id, osparc_job_id, report)
 
         return await dsm.create_s3_export(
             user_id, paths_to_export, progress_cb=_progress_cb
@@ -55,10 +53,11 @@ def fake_celery_export(  # TODO: replace this with the  Celery interface
     # The Celery task should propagate this information to the FE togehter with the progress (nmaybe as metadata?).
     # When progress is published via socket.io `osparc_job_id` will be used by the FE to associate the value
     # to the proper elment.
-    osparc_job_id = f"some_random_id_{uuid4()}"
+    osparc_job_id = f"{user_id}_fake_celery_export_{uuid4()}"
 
     def _progress_cb(report: ProgressReport) -> None:
-        _ = report  # TODO: propagate to celery task status
+        # TODO: propagate progress via CeleryWorkerInterface
+        _ = report
 
     return run_coroutine_threadsafe(
         _async_export(

services/storage/src/simcore_service_storage/modules/rabbitmq.py

@@ -1,8 +1,17 @@
 import logging
+from contextlib import suppress
 from typing import cast
 
 from fastapi import FastAPI
-from servicelib.logging_utils import log_context
+from models_library.osparc_jobs import OsparcJobId
+from models_library.progress_bar import ProgressReport
+from models_library.rabbitmq_messages import (
+    ProgressRabbitMessageWorkerJob,
+    ProgressType,
+    RabbitMessageBase,
+)
+from models_library.users import UserID
+from servicelib.logging_utils import log_catch, log_context
 from servicelib.rabbitmq import (
     RabbitMQClient,
     RabbitMQRPCClient,
@@ -52,7 +61,7 @@ async def on_shutdown() -> None:
 
 
 def get_rabbitmq_client(app: FastAPI) -> RabbitMQClient:
-    if not app.state.rabbitmq_client:
+    if not hasattr(app.state, "rabbitmq_client"):
         raise ConfigurationError(
             msg="RabbitMQ client is not available. Please check the configuration."
         )
@@ -62,3 +71,22 @@ def get_rabbitmq_client(app: FastAPI) -> RabbitMQClient:
 def get_rabbitmq_rpc_server(app: FastAPI) -> RabbitMQRPCClient:
     assert app.state.rabbitmq_rpc_server  # nosec
     return cast(RabbitMQRPCClient, app.state.rabbitmq_rpc_server)
+
+
+async def post_message(app: FastAPI, message: RabbitMessageBase) -> None:
+    with log_catch(_logger, reraise=False), suppress(ConfigurationError):
+        # NOTE: if rabbitmq was not initialized the error does not need to flood the logs
+        await get_rabbitmq_client(app).publish(message.channel_name, message)
+
+
+async def post_task_progress_message(
+    app: FastAPI, user_id: UserID, osparc_job_id: OsparcJobId, report: ProgressReport
+) -> None:
+    with log_catch(_logger, reraise=False):
+        message = ProgressRabbitMessageWorkerJob.model_construct(
+            user_id=user_id,
+            osparc_job_id=osparc_job_id,
+            progress_type=ProgressType.WORKER_JOB_EXPORTING,
+            report=report,
+        )
+        await post_message(app, message)

services/storage/tests/unit/test_api_celery__export_task.py

@@ -0,0 +1,38 @@
+from asyncio import AbstractEventLoop
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import FastAPI
+from models_library.projects_nodes_io import StorageFileID
+from pydantic import TypeAdapter
+from pytest_mock import MockerFixture
+from simcore_service_storage.api.celery._export_task import fake_celery_export
+
+pytest_simcore_core_services_selection = ["postgres"]
+pytest_simcore_ops_services_selection = ["adminer"]
+
+
+@pytest.fixture
+def mock_rabbitmq_post_message(mocker: MockerFixture) -> AsyncMock:
+    return mocker.patch("simcore_service_storage.modules.rabbitmq.post_message")
+
+
+async def test_celery_export_task(
+    event_loop: AbstractEventLoop,
+    initialized_app: FastAPI,
+    mock_rabbitmq_post_message: AsyncMock,
+):
+    # TODO: rewrite this test with the Celery job subscription tests from GCR
+
+    file_id = await event_loop.run_in_executor(
+        None, fake_celery_export, event_loop, initialized_app, 1, []
+    )
+
+    # check progress values, since there are no files to read progress goes from 0 to 1
+    assert mock_rabbitmq_post_message.call_count == 2
+    progress_values = [
+        x.args[1].report.actual_value for x in mock_rabbitmq_post_message.call_args_list
+    ]
+    assert progress_values == [0, 1]
+
+    assert TypeAdapter(StorageFileID).validate_python(file_id)