sending events about task status

Author: sanderegg

packages/dask-task-models-library/src/dask_task_models_library/models.py

@@ -1,10 +1,10 @@
-from dataclasses import asdict, dataclass
-from typing import Any, Final, TypeAlias
+from typing import Final, Literal, TypeAlias
 
 from dask.typing import Key
 from distributed.scheduler import TaskStateState as SchedulerTaskState
 from distributed.worker_state_machine import TaskStateState as WorkerTaskState
 from models_library.projects_state import RunningState
+from pydantic import BaseModel
 
 DaskJobID: TypeAlias = str
 DaskResources: TypeAlias = dict[str, int | float]
@@ -13,37 +13,38 @@
 _SCHEDULER_TASK_STATE_TO_RUNNING_STATE: Final[
     dict[SchedulerTaskState, RunningState]
 ] = {
-    "waiting": RunningState.PENDING,
-    "no-worker": RunningState.WAITING_FOR_RESOURCES,
-    "queued": RunningState.WAITING_FOR_RESOURCES,
-    "processing": RunningState.PENDING,
-    "memory": RunningState.SUCCESS,
-    "erred": RunningState.FAILED,
-    "forgotten": RunningState.UNKNOWN,
+    "released": RunningState.NOT_STARTED,  # Known but not actively computing or in memory
+    "waiting": RunningState.PENDING,  # On track to be computed, waiting on dependencies to arrive in memory
+    "no-worker": RunningState.WAITING_FOR_RESOURCES,  # Ready to be computed, but no appropriate worker exists (for example because of resource restrictions, or because no worker is connected at all).
+    "queued": RunningState.WAITING_FOR_RESOURCES,  # Ready to be computed, but all workers are already full.
+    "processing": RunningState.PENDING,  # All dependencies are available and the task is assigned to a worker for compute (the scheduler doesn’t know whether it’s in a worker queue or actively being computed).
+    "memory": RunningState.SUCCESS,  # In memory on one or more workers
+    "erred": RunningState.FAILED,  # Task computation, or one of its dependencies, has encountered an error
+    "forgotten": RunningState.UNKNOWN,  # Task is no longer needed by any client or dependent task, so it disappears from the scheduler as well. As soon as a task reaches this state, it is immediately dereferenced from the scheduler.
 }
 
 _WORKER_TASK_STATE_TO_RUNNING_STATE: Final[dict[WorkerTaskState, RunningState]] = {
-    "cancelled": RunningState.UNKNOWN,
-    "constrained": RunningState.UNKNOWN,
-    "error": RunningState.UNKNOWN,
-    "executing": RunningState.UNKNOWN,
-    "fetch": RunningState.UNKNOWN,
-    "flight": RunningState.UNKNOWN,
-    "forgotten": RunningState.UNKNOWN,
-    "long-running": RunningState.UNKNOWN,
-    "memory": RunningState.UNKNOWN,
-    "missing": RunningState.UNKNOWN,
-    "ready": RunningState.UNKNOWN,
-    "released": RunningState.UNKNOWN,
-    "rescheduled": RunningState.UNKNOWN,
-    "resumed": RunningState.UNKNOWN,
-    "waiting": RunningState.UNKNOWN,
+    "cancelled": RunningState.ABORTED,  # The scheduler asked to forget about this task, but it’s technically impossible at the moment. See Task cancellation. The task can be found in whatever collections it was in its previous state.
+    "constrained": RunningState.PENDING,  # Like ready, but the user specified resource constraints for this task. The task can be found in the WorkerState.constrained queue.
+    "error": RunningState.FAILED,  # Task execution failed
+    "executing": RunningState.STARTED,  # The task is currently being computed on a thread. It can be found in the WorkerState.executing set and in the distributed.worker.Worker.active_threads dict.
+    "fetch": RunningState.PENDING,  # This task is in memory on one or more peer workers, but not on this worker. Its data is queued to be transferred over the network, either because it’s a dependency of a task in waiting state, or because the Active Memory Manager requested it to be replicated here. The task can be found in the WorkerState.data_needed heap.
+    "flight": RunningState.PENDING,  # The task data is currently being transferred over the network from another worker. The task can be found in the WorkerState.in_flight_tasks and WorkerState.in_flight_workers collections.
+    "forgotten": RunningState.UNKNOWN,  # The scheduler asked this worker to forget about the task, and there are neither dependents nor dependencies on the same worker.
+    "long-running": RunningState.STARTED,  # Like executing, but the user code called distributed.secede() so the task no longer counts towards the maximum number of concurrent tasks. It can be found in the WorkerState.long_running set and in the distributed.worker.Worker.active_threads dict.
+    "memory": RunningState.SUCCESS,  # Task execution completed, or the task was successfully transferred from another worker, and is now held in either WorkerState.data or WorkerState.actors.
+    "missing": RunningState.PENDING,  # Like fetch, but all peer workers that were listed by the scheduler are either unreachable or have responded they don’t actually have the task data. The worker will periodically ask the scheduler if it knows of additional replicas; when it does, the task will transition again to fetch. The task can be found in the WorkerState.missing_dep_flight set.
+    "ready": RunningState.PENDING,  # The task is ready to be computed; all of its dependencies are in memory on the current worker and it’s waiting for an available thread. The task can be found in the WorkerState.ready heap.
+    "released": RunningState.PENDING,  # Known but not actively computing or in memory. A task can stay in this state when the scheduler asked to forget it, but it has dependent tasks on the same worker.
+    "rescheduled": RunningState.PENDING,  # The task just raised the Reschedule exception. This is a transitory state, which is not stored permanently.
+    "resumed": RunningState.PENDING,  # The task was recovered from cancelled state. See Task cancellation. The task can be found in whatever collections it was in its previous state.
+    "waiting": RunningState.PENDING,  # The scheduler has added the task to the worker queue. All of its dependencies are in memory somewhere on the cluster, but not all of them are in memory on the current worker, so they need to be fetched.
 }
 
 
-@dataclass
-class TaskLifeCycleState:
-    key: Key
+class TaskLifeCycleState(BaseModel):
+    key: str
+    source: Literal["scheduler", "worker"]
     worker: str | None
     state: RunningState
 
@@ -52,10 +53,19 @@ def from_scheduler_task_state(
         cls, key: Key, worker: str | None, task_state: SchedulerTaskState
     ) -> "TaskLifeCycleState":
         return cls(
-            key=key,
+            key=f"{key!r}",
+            source="scheduler",
             worker=worker,
             state=_SCHEDULER_TASK_STATE_TO_RUNNING_STATE[task_state],
         )
 
-    def model_dump(self) -> dict[str, Any]:
-        return asdict(self)
+    @classmethod
+    def from_worker_task_state(
+        cls, key: Key, worker: str | None, task_state: WorkerTaskState
+    ) -> "TaskLifeCycleState":
+        return cls(
+            key=f"{key!r}",
+            source="worker",
+            worker=worker,
+            state=_WORKER_TASK_STATE_TO_RUNNING_STATE[task_state],
+        )

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

@@ -36,16 +36,16 @@ def transition(
         stimulus_id: str,
         **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} due to {stimulus_id=}",
+            f"Task {key!r} transition from {start} to {finish} due to {stimulus_id=}",
         ):
             assert self.scheduler  # nosec
+
             self.scheduler.log_event(
                 TASK_LIFE_CYCLE_EVENT.format(key=key),
                 TaskLifeCycleState.from_scheduler_task_state(
                     key, kwargs.get("worker"), finish
-                ).model_dump(),
+                ).model_dump(mode="json"),
             )

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

@@ -3,8 +3,10 @@
 from typing import Any
 
 from dask.typing import Key
-from distributed import Worker, WorkerPlugin
-from distributed.scheduler import TaskStateState
+from dask_task_models_library.models import TASK_LIFE_CYCLE_EVENT, TaskLifeCycleState
+from distributed import WorkerPlugin
+from distributed.worker import Worker
+from distributed.worker_state_machine import TaskStateState
 from servicelib.logging_utils import log_context
 
 _logger = logging.getLogger(__name__)
@@ -17,7 +19,7 @@ def __init__(self) -> None:
             logging.INFO,
             "TaskLifecycleWorkerPlugin init",
         ):
-            self.worker = None
+            self._worker = None
 
     def setup(self, worker: Worker) -> Awaitable[None]:
         async def _() -> None:
@@ -26,7 +28,8 @@ async def _() -> None:
                 logging.INFO,
                 "TaskLifecycleWorkerPlugin start",
             ):
-                self.worker = worker
+                assert worker  # nosec
+                self._worker = worker
 
         return _()
 
@@ -37,19 +40,15 @@ def transition(
         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}",
+            f"Task {key!r} 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,
-                },
+            assert self._worker  # nosec
+            self._worker.log_event(
+                TASK_LIFE_CYCLE_EVENT.format(key=key),
+                TaskLifeCycleState.from_worker_task_state(
+                    key, kwargs.get("worker"), finish
+                ).model_dump(mode="json"),
             )

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

@@ -1,7 +1,13 @@
+# pylint: disable=redefined-outer-name
+# pylint: disable=unused-argument
+# pylint: disable=unused-variable
+# pylint: disable=no-member
+
 import time
 
 import distributed
 import pytest
+from dask_task_models_library.models import TASK_LIFE_CYCLE_EVENT
 
 pytest_simcore_core_services_selection = [
     "rabbit",
@@ -20,7 +26,8 @@ def _some_failing_task() -> None:
 
     future = dask_client.submit(_some_task)
     assert future.result(timeout=10) == 2
-    events = dask_client.get_events(f"task-lifecycle-{future.key}")
+
+    events = dask_client.get_events(TASK_LIFE_CYCLE_EVENT.format(key=future.key))
     print("XXXX received events:")
     assert events
     assert isinstance(events, tuple)
@@ -30,7 +37,7 @@ def _some_failing_task() -> None:
     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}")
+    events = dask_client.get_events(TASK_LIFE_CYCLE_EVENT.format(key=future.key))
     print("XXXX received events:")
     assert events
     assert isinstance(events, tuple)