creating dask plugin

Author: sanderegg

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

@@ -1,46 +1,88 @@
-import contextlib
 import logging
-from typing import cast
 
 import distributed
 from models_library.rabbitmq_messages import RabbitMessageBase
-from servicelib.logging_utils import log_catch
 from servicelib.rabbitmq import RabbitMQClient, wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
 from .errors import ConfigurationError
 
-logger = logging.getLogger(__name__)
+_logger = logging.getLogger(__name__)
 
 
-async def on_startup(
-    worker: distributed.Worker, rabbit_settings: RabbitSettings
-) -> None:
-    worker.rabbitmq_client = None
-    settings: RabbitSettings | None = rabbit_settings
-    if not settings:
-        logger.warning("Rabbit MQ client is de-activated in the settings")
-        return
-    await wait_till_rabbitmq_responsive(settings.dsn)
-    worker.rabbitmq_client = RabbitMQClient(
-        client_name="dask-sidecar", settings=settings
-    )
+class RabbitMQPlugin(distributed.WorkerPlugin):
+    """Dask Worker Plugin for RabbitMQ integration"""
 
+    name = "rabbitmq_plugin"
+    _client: RabbitMQClient | None = None
+    _settings: RabbitSettings | None = None
 
-async def on_shutdown(worker: distributed.Worker) -> None:
-    if worker.rabbitmq_client:
-        await worker.rabbitmq_client.close()
+    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"""
+        _logger.info("Setting up RabbitMQ plugin")
+        if not self._settings:
+            _logger.warning("RabbitMQ client is de-activated (no settings provided)")
+            return
 
-def get_rabbitmq_client(worker: distributed.Worker) -> RabbitMQClient:
-    if not worker.rabbitmq_client:
-        raise ConfigurationError(
-            msg="RabbitMQ client is not available. Please check the configuration."
+        await wait_till_rabbitmq_responsive(self._settings.dsn)
+        self._client = RabbitMQClient(
+            client_name="dask-sidecar", settings=self._settings
         )
-    return cast(RabbitMQClient, worker.rabbitmq_client)
+        _logger.info("RabbitMQ client initialized successfully")
 
+    async def teardown(self, worker: distributed.Worker) -> None:
+        """Called when the worker shuts down or the plugin is removed"""
+        _logger.info("Tearing down RabbitMQ plugin")
+        if self._client:
+            await self._client.close()
+            self._client = None
+            _logger.info("RabbitMQ client closed")
 
-async def post_message(worker: distributed.Worker, message: RabbitMessageBase) -> None:
-    with log_catch(logger, reraise=False), contextlib.suppress(ConfigurationError):
-        # NOTE: if rabbitmq was not initialized the error does not need to flood the logs
-        await get_rabbitmq_client(worker).publish(message.channel_name, message)
+    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"""
+        if self._client:
+            await self._client.publish(channel_name, message)
+
+
+# async def on_startup(
+#     worker: distributed.Worker, rabbit_settings: RabbitSettings
+# ) -> None:
+#     worker.rabbitmq_client = None
+#     settings: RabbitSettings | None = rabbit_settings
+#     if not settings:
+#         __logger.warning("Rabbit MQ client is de-activated in the settings")
+#         return
+#     await wait_till_rabbitmq_responsive(settings.dsn)
+#     worker.rabbitmq_client = RabbitMQClient(
+#         client_name="dask-sidecar", settings=settings
+#     )
+
+
+# async def on_shutdown(worker: distributed.Worker) -> None:
+#     if worker.rabbitmq_client:
+#         await worker.rabbitmq_client.close()
+
+
+# def get_rabbitmq_client(worker: distributed.Worker) -> RabbitMQClient:
+#     if not worker.rabbitmq_client:
+#         raise ConfigurationError(
+#             msg="RabbitMQ client is not available. Please check the configuration."
+#         )
+#     return cast(RabbitMQClient, worker.rabbitmq_client)
+
+
+# async def post_message(worker: distributed.Worker, message: RabbitMessageBase) -> None:
+#     with log_catch(__logger, reraise=False), contextlib.suppress(ConfigurationError):
+#         # NOTE: if rabbitmq was not initialized the error does not need to flood the logs
+#         await get_rabbitmq_client(worker).publish(message.channel_name, message)

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

@@ -18,8 +18,8 @@
 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 import on_shutdown as shutdown_rabbitmq
-from .rabbitmq import on_startup as setup_rabbitmq
 from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
@@ -83,7 +83,7 @@ async def dask_setup(worker: distributed.Worker) -> None:
                 "We do have a running loop in the main thread: %s", f"{loop=}"
             )
             if settings.DASK_SIDECAR_RABBITMQ:
-                await setup_rabbitmq(worker, settings.DASK_SIDECAR_RABBITMQ)
+                await worker.plugin_add(RabbitMQPlugin(settings.DASK_SIDECAR_RABBITMQ))
 
 
 async def dask_teardown(worker: distributed.Worker) -> None:

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

@@ -1,8 +1,8 @@
-from typing import Any, Dict
+from typing import Any
 
 
 def test_sidecar_service_is_deployed_in_global_mode(
-    simcore_docker_compose: Dict[str, Any]
+    simcore_docker_compose: dict[str, Any],
 ):
     dask_sidecar_deploy_config = simcore_docker_compose["services"]["dask-sidecar"][
         "deploy"