now we also test the scheduler preload

Author: sanderegg

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

@@ -22,6 +22,6 @@ async def dask_setup(scheduler: distributed.Scheduler) -> None:
         print_dask_scheduler_banner()
 
 
-async def dask_teardown(_worker: distributed.Worker) -> None:
-    with log_context(_logger, logging.INFO, "Tear down dask scheduler"):
+async def dask_teardown(scheduler: distributed.Scheduler) -> None:
+    with log_context(_logger, logging.INFO, f"Tear down dask {scheduler.address}"):
         ...

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

@@ -0,0 +1,35 @@
+from typing import Any
+
+from dask.typing import Key
+from distributed import SchedulerPlugin
+from distributed.scheduler import TaskStateState
+
+
+class SchedulerLifecyclePlugin(SchedulerPlugin):
+    def __init__(self) -> None:
+        self.scheduler = None
+
+    def add_task(self, key, **kwargs):
+        """Task published to cluster"""
+        self.scheduler.log_event(
+            "task-published",
+            {"key": key, "timestamp": time.time(), "client": kwargs.get("client")},
+        )
+
+    def transition(
+        self,
+        key: Key,
+        start: TaskStateState,
+        finish: TaskStateState,
+        **kwargs: Any,
+    ):
+        """State transitions"""
+        if finish in ("waiting", "processing", "memory", "erred"):
+            self.scheduler.log_event(
+                f"task-{finish}",
+                {
+                    "key": key,
+                    "worker": kwargs.get("worker"),
+                    "duration": time.time() - self.scheduler.tasks[key].start_time,
+                },
+            )

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

@@ -117,6 +117,7 @@ def local_cluster(app_environment: EnvVarsDict) -> Iterator[distributed.LocalClu
     with distributed.LocalCluster(
         worker_class=distributed.Worker,
         resources={"CPU": 10, "GPU": 10},
+        scheduler_kwargs={"preload": "simcore_service_dask_sidecar.scheduler"},
         preload="simcore_service_dask_sidecar.worker",
     ) as cluster:
         assert cluster

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

@@ -0,0 +1,9 @@
+import distributed
+
+pytest_simcore_core_services_selection = [
+    "rabbit",
+]
+
+
+def test_scheduler(dask_client: distributed.Client) -> None:
+    assert True