manage to poll task progress now

Author: sanderegg

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_base.py

@@ -116,11 +116,11 @@ class SortedTasks:
 
 
 async def _triage_changed_tasks(
-    changed_tasks: list[tuple[_Previous, _Current]],
+    changed_tasks_or_executing: list[tuple[_Previous, _Current]],
 ) -> SortedTasks:
     started_tasks = [
         current
-        for previous, current in changed_tasks
+        for previous, current in changed_tasks_or_executing
         if current.state in RUNNING_STATES
         or (
             previous.state in WAITING_FOR_START_STATES
@@ -130,17 +130,21 @@ async def _triage_changed_tasks(
 
     # NOTE: some tasks can be both started and completed since we might have the time they were running
     completed_tasks = [
-        current for _, current in changed_tasks if current.state in COMPLETED_STATES
+        current
+        for _, current in changed_tasks_or_executing
+        if current.state in COMPLETED_STATES
     ]
 
     waiting_for_resources_tasks = [
         current
-        for previous, current in changed_tasks
+        for previous, current in changed_tasks_or_executing
         if current.state in WAITING_FOR_START_STATES
     ]
 
     lost_or_momentarily_lost_tasks = [
-        current for _, current in changed_tasks if current.state is RunningState.UNKNOWN
+        current
+        for _, current in changed_tasks_or_executing
+        if current.state is RunningState.UNKNOWN
     ]
     if lost_or_momentarily_lost_tasks:
         _logger.warning(
@@ -321,21 +325,30 @@ async def _get_changed_tasks_from_backend(
         user_id: UserID,
         processing_tasks: list[CompTaskAtDB],
         comp_run: CompRunsAtDB,
-    ) -> list[tuple[_Previous, _Current]]:
+    ) -> tuple[list[tuple[_Previous, _Current]], list[CompTaskAtDB]]:
         tasks_backend_status = await self._get_tasks_status(
             user_id, processing_tasks, comp_run
         )
 
-        return [
-            (
-                task,
-                task.model_copy(update={"state": backend_state}),
-            )
-            for task, backend_state in zip(
-                processing_tasks, tasks_backend_status, strict=True
-            )
-            if task.state is not backend_state
-        ]
+        return (
+            [
+                (
+                    task,
+                    task.model_copy(update={"state": backend_state}),
+                )
+                for task, backend_state in zip(
+                    processing_tasks, tasks_backend_status, strict=True
+                )
+                if task.state is not backend_state
+            ],
+            [
+                task
+                for task, backend_state in zip(
+                    processing_tasks, tasks_backend_status, strict=True
+                )
+                if task.state is backend_state is RunningState.STARTED
+            ],
+        )
 
     async def _process_started_tasks(
         self,
@@ -476,7 +489,10 @@ async def _update_states_from_comp_backend(
             return
 
         # get the tasks which state actually changed since last check
-        tasks_with_changed_states = await self._get_changed_tasks_from_backend(
+        (
+            tasks_with_changed_states,
+            executing_tasks,
+        ) = await self._get_changed_tasks_from_backend(
             user_id, tasks_inprocess, comp_run
         )
         # NOTE: typical states a task goes through
@@ -511,6 +527,9 @@ async def _update_states_from_comp_backend(
         if sorted_tasks.waiting:
             await self._process_waiting_tasks(sorted_tasks.waiting)
 
+        if executing_tasks:
+            await self._process_executing_tasks(user_id, executing_tasks, comp_run)
+
     @abstractmethod
     async def _start_tasks(
         self,
@@ -545,6 +564,15 @@ async def _process_completed_tasks(
     ) -> None:
         """process tasks from the 3rd party backend"""
 
+    @abstractmethod
+    async def _process_executing_tasks(
+        self,
+        user_id: UserID,
+        tasks: list[CompTaskAtDB],
+        comp_run: CompRunsAtDB,
+    ) -> None:
+        """process executing tasks from the 3rd party backend"""
+
     async def apply(
         self,
         *,

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_dask.py

@@ -37,6 +37,7 @@
 )
 from ...utils.dask_client_utils import TaskHandlers, UnixTimestamp
 from ...utils.rabbitmq import (
+    publish_service_progress,
     publish_service_resource_tracking_stopped,
     publish_service_stopped_metrics,
 )
@@ -150,6 +151,62 @@ async def _get_tasks_status(
             _logger.info("The on demand computational backend is not ready yet...")
             return [RunningState.WAITING_FOR_CLUSTER] * len(tasks)
 
+    async def _process_executing_tasks(
+        self,
+        user_id: UserID,
+        tasks: list[CompTaskAtDB],
+        comp_run: CompRunsAtDB,
+    ) -> None:
+        task_progresses = []
+        try:
+            async with _cluster_dask_client(
+                user_id,
+                self,
+                use_on_demand_clusters=comp_run.use_on_demand_clusters,
+                run_metadata=comp_run.metadata,
+            ) as client:
+                task_progresses = await client.get_tasks_progress(
+                    [f"{t.job_id}" for t in tasks],
+                )
+            for task_progress_event in task_progresses:
+                if task_progress_event:
+                    await CompTasksRepository(
+                        self.db_engine
+                    ).update_project_task_progress(
+                        task_progress_event.task_owner.project_id,
+                        task_progress_event.task_owner.node_id,
+                        task_progress_event.progress,
+                    )
+
+        except ComputationalBackendOnDemandNotReadyError:
+            _logger.info("The on demand computational backend is not ready yet...")
+
+        comp_tasks_repo = CompTasksRepository(self.db_engine)
+        await asyncio.gather(
+            *(
+                comp_tasks_repo.update_project_task_progress(
+                    t.task_owner.project_id,
+                    t.task_owner.node_id,
+                    t.progress,
+                )
+                for t in task_progresses
+                if t
+            )
+        )
+        await asyncio.gather(
+            *(
+                publish_service_progress(
+                    self.rabbitmq_client,
+                    user_id=t.task_owner.user_id,
+                    project_id=t.task_owner.project_id,
+                    node_id=t.task_owner.node_id,
+                    progress=t.progress,
+                )
+                for t in task_progresses
+                if t
+            )
+        )
+
     async def _stop_tasks(
         self, user_id: UserID, tasks: list[CompTaskAtDB], comp_run: CompRunsAtDB
     ) -> None:

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

@@ -405,6 +405,33 @@ async def send_computation_tasks(
 
         return list_of_node_id_to_job_id
 
+    async def get_tasks_progress(
+        self, job_ids: list[str]
+    ) -> tuple[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)
+
+        return tuple(last_task_progress)
+
     async def get_tasks_status(self, job_ids: Iterable[str]) -> list[RunningState]:
         dask_utils.check_scheduler_is_still_the_same(
             self.backend.scheduler_id, self.backend.client
@@ -413,8 +440,6 @@ async def get_tasks_status(self, job_ids: Iterable[str]) -> list[RunningState]:
         dask_utils.check_scheduler_status(self.backend.client)
 
         async def _get_job_id_status(job_id: str) -> RunningState:
-            # TODO: maybe we should define an event just for that, instead of multiple calls
-            # but the max length by default is 1000. We should test it
             dask_events: tuple[tuple[UnixTimestamp, str], ...] = (
                 await self.backend.client.get_events(
                     TASK_LIFE_CYCLE_EVENT.format(key=job_id)