removed daskstate and upgrade test using running state

Author: sanderegg

services/director-v2/src/simcore_service_director_v2/models/dask_subsystem.py

@@ -157,7 +157,7 @@ async def _get_tasks_status(
                 use_on_demand_clusters=comp_run.use_on_demand_clusters,
                 run_metadata=comp_run.metadata,
             ) as client:
-                return await client.get_tasks_status2([f"{t.job_id}" for t in tasks])
+                return await client.get_tasks_status([f"{t.job_id}" for t in tasks])
             #     tasks_statuses = await client.get_tasks_status(
             #         [f"{t.job_id}" for t in tasks]
             #     )

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

@@ -17,7 +17,6 @@
 from http.client import HTTPException
 from typing import Any, Final, cast
 
-import dask.typing
 import distributed
 from aiohttp import ClientResponseError
 from common_library.json_serialization import json_dumps
@@ -427,7 +426,7 @@ async def send_computation_tasks(
 
         return list_of_node_id_to_job_id
 
-    async def get_tasks_status2(self, job_ids: Iterable[str]) -> list[RunningState]:
+    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
         )
@@ -436,6 +435,7 @@ async def get_tasks_status2(self, job_ids: Iterable[str]) -> list[RunningState]:
 
         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)
@@ -478,74 +478,6 @@ async def _get_job_id_status(job_id: str) -> RunningState:
 
         return await asyncio.gather(*(_get_job_id_status(job_id) for job_id in job_ids))
 
-    async def get_tasks_status(self, job_ids: list[str]) -> list[DaskClientTaskState]:
-        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)
-
-        # try to get the task from the scheduler
-        def _get_pipeline_statuses(
-            dask_scheduler: distributed.Scheduler,
-        ) -> dict[dask.typing.Key, DaskSchedulerTaskState | None]:
-            statuses: dict[dask.typing.Key, DaskSchedulerTaskState | None] = (
-                dask_scheduler.get_task_status(keys=job_ids)
-            )
-            return statuses
-
-        task_statuses: dict[dask.typing.Key, DaskSchedulerTaskState | None] = (
-            await self.backend.client.run_on_scheduler(_get_pipeline_statuses)
-        )
-        assert isinstance(task_statuses, dict)  # nosec
-
-        _logger.debug("found dask task statuses: %s", f"{task_statuses=}")
-
-        running_states: list[DaskClientTaskState] = []
-        for job_id in job_ids:
-            dask_status = cast(
-                DaskSchedulerTaskState | None, task_statuses.get(job_id, "lost")
-            )
-            if dask_status == "erred":
-                try:
-                    # find out if this was a cancellation
-                    var = distributed.Variable(job_id, client=self.backend.client)
-                    future: distributed.Future = await var.get(
-                        timeout=_DASK_DEFAULT_TIMEOUT_S
-                    )
-                    exception = await future.exception(timeout=_DASK_DEFAULT_TIMEOUT_S)
-                    assert isinstance(exception, Exception)  # nosec
-
-                    if isinstance(exception, TaskCancelledError):
-                        running_states.append(DaskClientTaskState.ABORTED)
-                    else:
-                        assert exception  # nosec
-                        _logger.warning(
-                            "Task  %s completed in error:\n%s\nTrace:\n%s",
-                            job_id,
-                            exception,
-                            "".join(traceback.format_exception(exception)),
-                        )
-                        running_states.append(DaskClientTaskState.ERRED)
-                except TimeoutError:
-                    _logger.warning(
-                        "Task  %s could not be retrieved from dask-scheduler, it is lost\n"
-                        "TIP:If the task was unpublished this can happen, or if the dask-scheduler was restarted.",
-                        job_id,
-                    )
-                    running_states.append(DaskClientTaskState.LOST)
-
-            elif dask_status is None:
-                running_states.append(DaskClientTaskState.LOST)
-            else:
-                running_states.append(
-                    _DASK_TASK_STATUS_DASK_CLIENT_TASK_STATE_MAP.get(
-                        dask_status, DaskClientTaskState.LOST
-                    )
-                )
-
-        return running_states
-
     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
         # a sidecar can be cancelled that way.

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

@@ -43,6 +43,7 @@
 from models_library.docker import to_simcore_runtime_docker_label_key
 from models_library.projects import ProjectID
 from models_library.projects_nodes_io import NodeID
+from models_library.projects_state import RunningState
 from models_library.resource_tracker import HardwareInfo
 from models_library.services_types import ServiceRunID
 from models_library.users import UserID
@@ -60,7 +61,6 @@
 )
 from simcore_service_director_v2.models.comp_runs import RunMetadataDict
 from simcore_service_director_v2.models.comp_tasks import Image
-from simcore_service_director_v2.models.dask_subsystem import DaskClientTaskState
 from simcore_service_director_v2.modules.dask_client import DaskClient, TaskHandlers
 from tenacity.asyncio import AsyncRetrying
 from tenacity.retry import retry_if_exception_type
@@ -90,7 +90,7 @@ async def _assert_wait_for_cb_call(mocked_fct, timeout: int | None = None):
 async def _assert_wait_for_task_status(
     job_id: str,
     dask_client: DaskClient,
-    expected_status: DaskClientTaskState,
+    expected_status: RunningState,
     timeout: int | None = None,  # noqa: ASYNC109
 ):
     async for attempt in AsyncRetrying(
@@ -105,11 +105,11 @@ async def _assert_wait_for_task_status(
                 f"Attempt={attempt.retry_state.attempt_number}"
             )
             got = (await dask_client.get_tasks_status([job_id]))[0]
-            assert isinstance(got, DaskClientTaskState)
+            assert isinstance(got, RunningState)
             print(f"{got=} vs {expected_status=}")
-            if got is DaskClientTaskState.ERRED and expected_status not in [
-                DaskClientTaskState.ERRED,
-                DaskClientTaskState.LOST,
+            if got is RunningState.FAILED and expected_status not in [
+                RunningState.FAILED,
+                RunningState.UNKNOWN,
             ]:
                 try:
                     # we can fail fast here
@@ -516,7 +516,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task.job_id,
         dask_client,
-        expected_status=DaskClientTaskState.PENDING_OR_STARTED,
+        expected_status=RunningState.STARTED,
     )
 
     # using the event we let the remote fct continue
@@ -530,7 +530,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task.job_id,
         dask_client,
-        expected_status=DaskClientTaskState.SUCCESS,
+        expected_status=RunningState.SUCCESS,
     )
 
     # check the results
@@ -544,7 +544,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task.job_id,
         dask_client,
-        expected_status=DaskClientTaskState.LOST,
+        expected_status=RunningState.UNKNOWN,
         timeout=60,
     )
 
@@ -609,7 +609,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        expected_status=DaskClientTaskState.SUCCESS,
+        expected_status=RunningState.SUCCESS,
     )
     assert published_computation_task[0].node_id in image_params.fake_tasks
     # creating a new future shows that it is not done????
@@ -704,7 +704,7 @@ def fake_remote_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        DaskClientTaskState.PENDING_OR_STARTED,
+        RunningState.STARTED,
     )
 
     # we wait to be sure the remote fct is started
@@ -721,7 +721,7 @@ def fake_remote_fct(
 
     await _assert_wait_for_cb_call(mocked_user_completed_cb)
     await _assert_wait_for_task_status(
-        published_computation_task[0].job_id, dask_client, DaskClientTaskState.ABORTED
+        published_computation_task[0].job_id, dask_client, RunningState.ABORTED
     )
 
     # getting the results should throw the cancellation error
@@ -732,9 +732,12 @@ def fake_remote_fct(
     await dask_client.release_task_result(published_computation_task[0].job_id)
     # NOTE: this change of status takes a very long time to happen and is not relied upon so we skip it since it
     # makes the test fail a lot for no gain (it's kept here in case it ever becomes an issue)
-    # await _assert_wait_for_task_status(
-    #     job_id, dask_client, RunningState.UNKNOWN, timeout=120
-    # )
+    await _assert_wait_for_task_status(
+        published_computation_task[0].job_id,
+        dask_client,
+        RunningState.UNKNOWN,
+        timeout=120,
+    )
 
 
 async def test_failed_task_returns_exceptions(
@@ -784,7 +787,7 @@ def fake_failing_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        expected_status=DaskClientTaskState.ERRED,
+        expected_status=RunningState.FAILED,
     )
     with pytest.raises(
         ValueError,
@@ -1047,7 +1050,7 @@ def fake_remote_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        DaskClientTaskState.PENDING_OR_STARTED,
+        RunningState.STARTED,
     )
 
     # let the remote fct run through now
@@ -1057,7 +1060,7 @@ def fake_remote_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        DaskClientTaskState.ERRED if fail_remote_fct else DaskClientTaskState.SUCCESS,
+        RunningState.FAILED if fail_remote_fct else RunningState.SUCCESS,
     )
     # release the task results
     await dask_client.release_task_result(published_computation_task[0].job_id)
@@ -1069,7 +1072,7 @@ def fake_remote_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        DaskClientTaskState.LOST,
+        RunningState.UNKNOWN,
         timeout=60,
     )
 
@@ -1209,7 +1212,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        expected_status=DaskClientTaskState.PENDING_OR_STARTED,
+        expected_status=RunningState.STARTED,
     )
 
     # check we have one worker using the resources
@@ -1240,7 +1243,7 @@ def fake_sidecar_fct(
     await _assert_wait_for_task_status(
         published_computation_task[0].job_id,
         dask_client,
-        expected_status=DaskClientTaskState.SUCCESS,
+        expected_status=RunningState.SUCCESS,
     )
 
     # check the resources are released