task now revert to old state to prevent issue in heartbeat

Author: sanderegg

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

@@ -1,19 +1,27 @@
+from dataclasses import dataclass
 from typing import Literal
 
 from models_library.projects import ProjectID
 from models_library.rabbitmq_messages import RabbitMessageBase
 from models_library.users import UserID
 
 from ...models.comp_runs import Iteration
+from ...models.comp_tasks import CompTaskAtDB
 
 
 class SchedulePipelineRabbitMessage(RabbitMessageBase):
-    channel_name: Literal[
+    channel_name: Literal["simcore.services.director-v2.scheduling"] = (
         "simcore.services.director-v2.scheduling"
-    ] = "simcore.services.director-v2.scheduling"
+    )
     user_id: UserID
     project_id: ProjectID
     iteration: Iteration
 
     def routing_key(self) -> str | None:  # pylint: disable=no-self-use # abstract
         return None
+
+
+@dataclass(frozen=True, slots=True)
+class TaskStateTracker:
+    previous: CompTaskAtDB
+    current: CompTaskAtDB

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

@@ -63,6 +63,7 @@
 from ..db.repositories.comp_pipelines import CompPipelinesRepository
 from ..db.repositories.comp_runs import CompRunsRepository
 from ..db.repositories.comp_tasks import CompTasksRepository
+from ._models import TaskStateTracker
 from ._publisher import request_pipeline_scheduling
 from ._utils import (
     COMPLETED_STATES,
@@ -76,12 +77,6 @@
 _logger = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True, slots=True)
-class TaskStateTracker:
-    previous: CompTaskAtDB
-    current: CompTaskAtDB
-
-
 _MAX_WAITING_TIME_FOR_UNKNOWN_TASKS: Final[datetime.timedelta] = datetime.timedelta(
     seconds=30
 )
@@ -120,9 +115,9 @@ async def _async_cb() -> None:
 @dataclass(frozen=True, slots=True)
 class SortedTasks:
     started: list[CompTaskAtDB]
-    completed: list[CompTaskAtDB]
-    waiting: list[CompTaskAtDB]
-    potentially_lost: list[CompTaskAtDB]
+    completed: list[TaskStateTracker]
+    waiting: list[TaskStateTracker]
+    potentially_lost: list[TaskStateTracker]
 
 
 async def _triage_changed_tasks(
@@ -139,19 +134,19 @@ async def _triage_changed_tasks(
     ]
 
     completed_tasks = [
-        tracker.current
+        tracker
         for tracker in changed_tasks
         if tracker.current.state in COMPLETED_STATES
     ]
 
     waiting_for_resources_tasks = [
-        tracker.current
+        tracker
         for tracker in changed_tasks
         if tracker.current.state in WAITING_FOR_START_STATES
     ]
 
     lost_tasks = [
-        tracker.current
+        tracker
         for tracker in changed_tasks
         if (tracker.current.state is RunningState.UNKNOWN)
         and (
@@ -162,7 +157,7 @@ async def _triage_changed_tasks(
     if lost_tasks:
         _logger.warning(
             "%s are currently in unknown state. TIP: If they are running in an external cluster and it is not yet ready, that might explain it. But inform @sanderegg nevertheless!",
-            [t.node_id for t in lost_tasks],
+            [t.current.node_id for t in lost_tasks],
         )
 
     return SortedTasks(
@@ -508,16 +503,16 @@ async def _process_started_tasks(
         )
 
     async def _process_waiting_tasks(
-        self, tasks: list[CompTaskAtDB], run_id: PositiveInt
+        self, tasks: list[TaskStateTracker], run_id: PositiveInt
     ) -> None:
         comp_tasks_repo = CompTasksRepository(self.db_engine)
         await asyncio.gather(
             *(
                 comp_tasks_repo.update_project_tasks_state(
-                    t.project_id,
+                    t.current.project_id,
                     run_id,
-                    [t.node_id],
-                    t.state,
+                    [t.current.node_id],
+                    t.current.state,
                 )
                 for t in tasks
             )
@@ -608,7 +603,7 @@ async def _stop_tasks(
     async def _process_completed_tasks(
         self,
         user_id: UserID,
-        tasks: list[CompTaskAtDB],
+        tasks: list[TaskStateTracker],
         iteration: Iteration,
         comp_run: CompRunsAtDB,
     ) -> None:

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

@@ -55,6 +55,7 @@
 from ._constants import (
     MAX_CONCURRENT_PIPELINE_SCHEDULING,
 )
+from ._models import TaskStateTracker
 from ._scheduler_base import BaseCompScheduler
 from ._utils import (
     WAITING_FOR_START_STATES,
@@ -293,7 +294,7 @@ async def _stop_tasks(
     async def _process_completed_tasks(
         self,
         user_id: UserID,
-        tasks: list[CompTaskAtDB],
+        tasks: list[TaskStateTracker],
         iteration: Iteration,
         comp_run: CompRunsAtDB,
     ) -> None:
@@ -306,13 +307,20 @@ async def _process_completed_tasks(
             run_metadata=comp_run.metadata,
         ) as client:
             tasks_results = await asyncio.gather(
-                *[client.get_task_result(t.job_id or "undefined") for t in tasks],
+                *[
+                    client.get_task_result(t.current.job_id or "undefined")
+                    for t in tasks
+                ],
                 return_exceptions=True,
             )
             async for future in limited_as_completed(
                 (
                     self._process_task_result(
-                        task, result, comp_run.metadata, iteration, comp_run.run_id
+                        task,
+                        result,
+                        comp_run.metadata,
+                        iteration,
+                        comp_run.run_id,
                     )
                     for task, result in zip(tasks, tasks_results, strict=True)
                 ),
@@ -467,31 +475,30 @@ async def _handle_task_error(
 
     async def _process_task_result(
         self,
-        task: CompTaskAtDB,
+        task: TaskStateTracker,
         result: BaseException | TaskOutputData,
         run_metadata: RunMetadataDict,
         iteration: Iteration,
         run_id: PositiveInt,
     ) -> tuple[bool, str]:
         """Returns True and the job ID if the task was successfully processed and can be released from the Dask cluster."""
         _logger.debug("received %s result: %s", f"{task=}", f"{result=}")
-
-        assert task.job_id  # nosec
+        assert task.current.job_id  # nosec
         (
             _service_key,
             _service_version,
             user_id,
             project_id,
             node_id,
-        ) = parse_dask_job_id(task.job_id)
+        ) = parse_dask_job_id(task.current.job_id)
 
-        assert task.project_id == project_id  # nosec
-        assert task.node_id == node_id  # nosec
+        assert task.current.project_id == project_id  # nosec
+        assert task.current.node_id == node_id  # nosec
         log_error_context = {
             "user_id": user_id,
             "project_id": project_id,
             "node_id": node_id,
-            "job_id": task.job_id,
+            "job_id": task.current.job_id,
         }
 
         if isinstance(result, TaskOutputData):
@@ -500,7 +507,9 @@ async def _process_task_result(
                 simcore_platform_status,
                 task_errors,
                 task_completed,
-            ) = await self._handle_successful_run(task, result, log_error_context)
+            ) = await self._handle_successful_run(
+                task.current, result, log_error_context
+            )
 
         elif isinstance(result, ComputationalBackendTaskResultsNotReadyError):
             (
@@ -509,7 +518,7 @@ async def _process_task_result(
                 task_errors,
                 task_completed,
             ) = await self._handle_computational_retrieval_error(
-                task, user_id, result, log_error_context
+                task.current, user_id, result, log_error_context
             )
         elif isinstance(result, ComputationalBackendNotConnectedError):
             (
@@ -518,15 +527,15 @@ async def _process_task_result(
                 task_errors,
                 task_completed,
             ) = await self._handle_computational_backend_not_connected_error(
-                task, result, log_error_context
+                task.current, result, log_error_context
             )
         else:
             (
                 task_final_state,
                 simcore_platform_status,
                 task_errors,
                 task_completed,
-            ) = await self._handle_task_error(task, result, log_error_context)
+            ) = await self._handle_task_error(task.current, result, log_error_context)
 
             # we need to remove any invalid files in the storage
             await clean_task_output_and_log_files_if_invalid(
@@ -549,21 +558,21 @@ async def _process_task_result(
                 simcore_user_agent=run_metadata.get(
                     "simcore_user_agent", UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE
                 ),
-                task=task,
+                task=task.current,
                 task_final_state=task_final_state,
             )
 
         await CompTasksRepository(self.db_engine).update_project_tasks_state(
-            task.project_id,
+            task.current.project_id,
             run_id,
-            [task.node_id],
-            task_final_state if task_completed else RunningState.STARTED,
+            [task.current.node_id],
+            task_final_state if task_completed else task.previous.state,
             errors=task_errors,
             optional_progress=1 if task_completed else None,
             optional_stopped=arrow.utcnow().datetime if task_completed else None,
         )
 
-        return task_completed, task.job_id
+        return task_completed, task.current.job_id
 
     async def _task_progress_change_handler(
         self, event: tuple[UnixTimestamp, Any]