renaming and testing

Author: sanderegg

services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq.py

@@ -0,0 +1,63 @@
+import logging
+
+import distributed
+from models_library.rabbitmq_messages import RabbitMessageBase
+from servicelib.logging_utils import log_catch, log_context
+from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive
+from settings_library.rabbit import RabbitSettings
+
+from .errors import ConfigurationError
+
+_logger = logging.getLogger(__name__)
+
+
+class RabbitMQPlugin(distributed.WorkerPlugin):
+    """Dask Worker Plugin for RabbitMQ integration"""
+
+    name = "rabbitmq_plugin"
+    _client: RabbitMQClient | None = None
+    _settings: RabbitSettings | None = None
+
+    def __init__(self, settings: RabbitSettings):
+        self._settings = settings
+
+    async def setup(self, worker: distributed.Worker) -> None:
+        """Called when the plugin is attached to a worker"""
+        if not self._settings:
+            _logger.warning("RabbitMQ client is de-activated (no settings provided)")
+            return
+
+        with log_context(
+            _logger,
+            logging.INFO,
+            f"RabbitMQ client initialization for worker {worker.address}",
+        ):
+            await wait_till_rabbitmq_responsive(self._settings.dsn)
+            self._client = RabbitMQClient(
+                client_name="dask-sidecar", settings=self._settings
+            )
+
+    async def teardown(self, worker: distributed.Worker) -> None:
+        """Called when the worker shuts down or the plugin is removed"""
+        with log_context(
+            _logger,
+            logging.INFO,
+            f"RabbitMQ client teardown for worker {worker.address}",
+        ):
+            if self._client:
+                await self._client.close()
+                self._client = None
+
+    def get_client(self) -> RabbitMQClient:
+        """Returns the RabbitMQ client or raises an error if not available"""
+        if not self._client:
+            raise ConfigurationError(
+                msg="RabbitMQ client is not available. Please check the configuration."
+            )
+        return self._client
+
+    async def publish(self, *, channel_name: str, message: RabbitMessageBase) -> None:
+        """Publishes a message to the specified channel"""
+        with log_catch(_logger, reraise=False):
+            if self._client:
+                await self._client.publish(channel_name, message)

services/dask-sidecar/src/simcore_service_dask_sidecar/rabbitmq_plugin.py

@@ -18,7 +18,7 @@
 from ._meta import print_dask_sidecar_banner
 from .computational_sidecar.core import ComputationalSidecar
 from .dask_utils import TaskPublisher, get_current_task_resources, monitor_task_abortion
-from .rabbitmq import RabbitMQPlugin
+from .rabbitmq_plugin import RabbitMQPlugin
 from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
@@ -41,7 +41,7 @@ def __init__(self, worker: distributed.Worker):
 
     def exit_gracefully(self, *_args):
         tasks = asyncio.all_tasks()
-        dask_worker_logger.warning(
+        _logger.warning(
             "Application shutdown detected!\n %s",
             pformat([t.get_name() for t in tasks]),
         )
@@ -69,24 +69,22 @@ async def dask_setup(worker: distributed.Worker) -> None:
         tracing_settings=None,  # no tracing for dask sidecar
     )
 
-    with log_context(dask_worker_logger, logging.INFO, "Launch dask worker"):
-        dask_worker_logger.info("app settings: %s", settings.model_dump_json(indent=1))
+    with log_context(_logger, logging.INFO, "Launch dask worker"):
+        _logger.info("app settings: %s", settings.model_dump_json(indent=1))
 
         print_dask_sidecar_banner()
 
         if threading.current_thread() is threading.main_thread():
             GracefulKiller(worker)
 
             loop = asyncio.get_event_loop()
-            dask_worker_logger.info(
-                "We do have a running loop in the main thread: %s", f"{loop=}"
-            )
+            _logger.info("We do have a running loop in the main thread: %s", f"{loop=}")
             if settings.DASK_SIDECAR_RABBITMQ:
                 await worker.plugin_add(RabbitMQPlugin(settings.DASK_SIDECAR_RABBITMQ))
 
 
-async def dask_teardown(_worker: distributed.Worker) -> None:
-    with log_context(dask_worker_logger, logging.INFO, "tear down dask worker"):
+async def dask_teardown(worker: distributed.Worker) -> None:
+    with log_context(_logger, logging.INFO, f"tear down dask {worker.address}"):
         ...
 
 

services/dask-sidecar/src/simcore_service_dask_sidecar/tasks.py

@@ -37,6 +37,7 @@
     "pytest_simcore.docker_swarm",
     "pytest_simcore.environment_configs",
     "pytest_simcore.faker_users_data",
+    "pytest_simcore.rabbit_service",
     "pytest_simcore.repository_paths",
 ]
 
@@ -231,7 +232,7 @@ def creator() -> AnyUrl:
         open_file = fsspec.open(f"{new_remote_file}", mode="wt", **s3_storage_kwargs)
         with open_file as fp:
             fp.write(  # type: ignore
-                f"This is the file contents of file #'{(len(list_of_created_files)+1):03}'\n"
+                f"This is the file contents of file #'{(len(list_of_created_files) + 1):03}'\n"
             )
             for s in faker.sentences(5):
                 fp.write(f"{s}\n")  # type: ignore

services/dask-sidecar/tests/unit/conftest.py

@@ -0,0 +1,31 @@
+import distributed
+import pytest
+from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
+from pytest_simcore.helpers.typing_env import EnvVarsDict
+from settings_library.rabbit import RabbitSettings
+
+# Selection of core and tool services started in this swarm fixture (integration)
+pytest_simcore_core_services_selection = [
+    "rabbit",
+]
+
+pytest_simcore_ops_services_selection = []
+
+
+@pytest.fixture
+def app_environment(
+    app_environment: EnvVarsDict,
+    monkeypatch: pytest.MonkeyPatch,
+    rabbit_service: RabbitSettings,
+) -> EnvVarsDict:
+    # configured as worker
+    envs = setenvs_from_dict(
+        monkeypatch,
+        {
+            "DASK_WORKER_RABBITMQ": rabbit_service.model_dump_json(),
+        },
+    )
+    return app_environment | envs
+
+
+def test_rabbitmq_plugin_initializes(dask_client: distributed.Client): ...