added specific events to be more effective

Author: sanderegg

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

@@ -10,6 +10,7 @@
 DaskResources: TypeAlias = dict[str, int | float]
 
 TASK_LIFE_CYCLE_EVENT: Final[str] = "task-lifecycle-{key}"
+TASK_RUNNING_PROGRESS_EVENT: Final[str] = "task-progress-{key}"
 _SCHEDULER_TASK_STATE_TO_RUNNING_STATE: Final[
     dict[SchedulerTaskState, RunningState]
 ] = {

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

@@ -13,6 +13,7 @@
 )
 from dask_task_models_library.container_tasks.io import TaskCancelEventName
 from dask_task_models_library.container_tasks.protocol import TaskOwner
+from dask_task_models_library.models import TASK_RUNNING_PROGRESS_EVENT
 from distributed.worker import get_worker
 from distributed.worker_state_machine import TaskState
 from models_library.progress_bar import ProgressReport
@@ -178,4 +179,10 @@ def publish_event(
         log_catch(_logger, reraise=False),
         log_context(_logger, logging.DEBUG, msg=f"publishing {event=}"),
     ):
-        worker.log_event(TaskProgressEvent.topic_name(), event.model_dump_json())
+        worker.log_event(
+            [
+                TaskProgressEvent.topic_name(),
+                TASK_RUNNING_PROGRESS_EVENT.format(key=event.job_id),
+            ],
+            event.model_dump_json(),
+        )

services/director-v2/src/simcore_service_director_v2/modules/dask_client.py

@@ -40,6 +40,7 @@
 from dask_task_models_library.container_tasks.utils import generate_dask_job_id
 from dask_task_models_library.models import (
     TASK_LIFE_CYCLE_EVENT,
+    TASK_RUNNING_PROGRESS_EVENT,
     DaskJobID,
     DaskResources,
     TaskLifeCycleState,
@@ -407,30 +408,25 @@ async def send_computation_tasks(
 
     async def get_tasks_progress(
         self, job_ids: list[str]
-    ) -> tuple[TaskProgressEvent | None, ...]:
+    ) -> list[TaskProgressEvent | None]:
         dask_utils.check_scheduler_is_still_the_same(
             self.backend.scheduler_id, self.backend.client
         )
         dask_utils.check_communication_with_scheduler_is_open(self.backend.client)
         dask_utils.check_scheduler_status(self.backend.client)
 
-        dask_events = await self.backend.client.get_events(
-            TaskProgressEvent.topic_name()
-        )
-
-        if not dask_events:
-            return tuple([None] * len(job_ids))
-        last_task_progress = []
-        for job_id in job_ids:
-            progress_event = None
-            for dask_event in reversed(dask_events):
-                parsed_event = TaskProgressEvent.model_validate_json(dask_event[1])
-                if parsed_event.job_id == job_id:
-                    progress_event = parsed_event
-                    break
-            last_task_progress.append(progress_event)
+        async def _get_task_progress(job_id: str) -> TaskProgressEvent | None:
+            dask_events: tuple[tuple[UnixTimestamp, str], ...] = (
+                await self.backend.client.get_events(
+                    TASK_RUNNING_PROGRESS_EVENT.format(key=job_id)
+                )
+            )
+            if not dask_events:
+                return None
+            # we are interested in the last event
+            return TaskProgressEvent.model_validate_json(dask_events[-1][1])
 
-        return tuple(last_task_progress)
+        return await asyncio.gather(*(_get_task_progress(job_id) for job_id in job_ids))
 
     async def get_tasks_status(self, job_ids: Iterable[str]) -> list[RunningState]:
         dask_utils.check_scheduler_is_still_the_same(
@@ -439,7 +435,7 @@ async def get_tasks_status(self, job_ids: Iterable[str]) -> list[RunningState]:
         dask_utils.check_communication_with_scheduler_is_open(self.backend.client)
         dask_utils.check_scheduler_status(self.backend.client)
 
-        async def _get_job_id_status(job_id: str) -> RunningState:
+        async def _get_task_state(job_id: str) -> RunningState:
             dask_events: tuple[tuple[UnixTimestamp, str], ...] = (
                 await self.backend.client.get_events(
                     TASK_LIFE_CYCLE_EVENT.format(key=job_id)
@@ -480,7 +476,7 @@ async def _get_job_id_status(job_id: str) -> RunningState:
 
             return parsed_event.state
 
-        return await asyncio.gather(*(_get_job_id_status(job_id) for job_id in job_ids))
+        return await asyncio.gather(*(_get_task_state(job_id) for job_id in job_ids))
 
     async def abort_computation_task(self, job_id: str) -> None:
         # Dask future may be cancelled, but only a future that was not already taken by