added life cycle plugins

Author: sanderegg

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

@@ -0,0 +1,55 @@
+import logging
+from collections.abc import Awaitable
+from typing import Any
+
+from dask.typing import Key
+from distributed import Worker, WorkerPlugin
+from distributed.scheduler import TaskStateState
+from servicelib.logging_utils import log_context
+
+_logger = logging.getLogger(__name__)
+
+
+class TaskLifecycleWorkerPlugin(WorkerPlugin):
+    def __init__(self) -> None:
+        with log_context(
+            _logger,
+            logging.INFO,
+            "TaskLifecycleWorkerPlugin init",
+        ):
+            self.worker = None
+
+    def setup(self, worker: Worker) -> Awaitable[None]:
+        async def _() -> None:
+            with log_context(
+                _logger,
+                logging.INFO,
+                "TaskLifecycleWorkerPlugin start",
+            ):
+                self.worker = worker
+
+        return _()
+
+    def transition(
+        self,
+        key: Key,
+        start: TaskStateState,
+        finish: TaskStateState,
+        **kwargs: Any,
+    ):
+        # Start state: one of released, waiting, processing, memory, error
+        with log_context(
+            _logger,
+            logging.INFO,
+            f"Task {key} transition from {start} to {finish}",
+        ):
+            assert self.worker  # nosec
+            self.worker.log_event(
+                f"task-lifecycle-{key}",
+                {
+                    "key": key,
+                    "worker": kwargs.get("worker"),
+                    "start": start,
+                    "finish": finish,
+                },
+            )

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

@@ -18,6 +18,9 @@
 from .computational_sidecar.core import ComputationalSidecar
 from .rabbitmq_plugin import RabbitMQPlugin
 from .settings import ApplicationSettings
+from .task_life_cycle_worker_plugin import (
+    TaskLifecycleWorkerPlugin,
+)
 from .utils.dask import (
     TaskPublisher,
     get_current_task_resources,
@@ -78,6 +81,11 @@ async def dask_setup(worker: distributed.Worker) -> None:
             except Exception:
                 await worker.close(reason="failed to add RabbitMQ plugin")
                 raise
+        try:
+            await worker.plugin_add(TaskLifecycleWorkerPlugin(), catch_errors=False)
+        except Exception:
+            await worker.close(reason="failed to add TaskLifecycleWorkerPlugin")
+            raise
 
         print_dask_sidecar_banner()
 

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

@@ -3,7 +3,12 @@
 # pylint: disable=unused-variable
 # pylint: disable=no-member
 
+import asyncio
+from unittest import mock
+
 import distributed
+import pytest
+from pytest_mock import MockerFixture
 
 # Selection of core and tool services started in this swarm fixture (integration)
 pytest_simcore_core_services_selection = [
@@ -14,3 +19,19 @@
 
 
 def test_rabbitmq_plugin_initializes(dask_client: distributed.Client): ...
+
+
+@pytest.fixture
+def erroring_rabbitmq_plugin(mocker: MockerFixture) -> mock.Mock:
+    return mocker.patch(
+        "simcore_service_dask_sidecar.rabbitmq_plugin.RabbitMQPlugin",
+        autospec=True,
+        side_effect=RuntimeError("Pytest: RabbitMQ plugin initialization failed"),
+    )
+
+
+async def test_dask_worker_closes_if_plugin_fails_on_start(
+    erroring_rabbitmq_plugin: mock.Mock,
+    local_cluster: distributed.LocalCluster,
+):
+    await asyncio.sleep(10)

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

@@ -1,9 +1,38 @@
+import time
+
 import distributed
+import pytest
 
 pytest_simcore_core_services_selection = [
     "rabbit",
 ]
 
 
 def test_scheduler(dask_client: distributed.Client) -> None:
-    assert True
+    def _some_task() -> int:
+        time.sleep(1)
+        return 2
+
+    def _some_failing_task() -> None:
+        time.sleep(1)
+        msg = "Some error"
+        raise RuntimeError(msg)
+
+    future = dask_client.submit(_some_task)
+    assert future.result(timeout=10) == 2
+    events = dask_client.get_events(f"task-lifecycle-{future.key}")
+    print("XXXX received events:")
+    assert events
+    assert isinstance(events, tuple)
+    for event in events:
+        print(f"\t{event}")
+
+    future = dask_client.submit(_some_failing_task)
+    with pytest.raises(RuntimeError):
+        future.result(timeout=10)
+    events = dask_client.get_events(f"task-lifecycle-{future.key}")
+    print("XXXX received events:")
+    assert events
+    assert isinstance(events, tuple)
+    for event in events:
+        print(f"\t{event}")