make it work

Author: sanderegg

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

@@ -23,14 +23,15 @@
 from models_library.users import UserID
 from pydantic import PositiveInt
 from servicelib.common_headers import UNDEFINED_DEFAULT_SIMCORE_USER_AGENT_VALUE
+from servicelib.logging_errors import create_troubleshootting_log_kwargs
 from servicelib.logging_utils import log_catch, log_context
 from servicelib.utils import limited_as_completed
 
 from ...core.errors import (
     ComputationalBackendNotConnectedError,
     ComputationalBackendOnDemandNotReadyError,
     ComputationalBackendTaskResultsNotReadyError,
-    TaskSchedulingError,
+    PortsValidationError,
 )
 from ...models.comp_runs import CompRunsAtDB, Iteration, RunMetadataDict
 from ...models.comp_tasks import CompTaskAtDB
@@ -59,6 +60,9 @@
 _logger = logging.getLogger(__name__)
 
 _DASK_CLIENT_RUN_REF: Final[str] = "{user_id}:{project_id}:{run_id}"
+_TASK_RETRIEVAL_ERROR_TYPE: Final[str] = "task-result-retrieval-timeout"
+_TASK_RETRIEVAL_ERROR_MSG: Final[str] = "Retrieval of task result timed-out"
+_TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY: Final[str] = "check_time"
 
 
 @asynccontextmanager
@@ -329,66 +333,121 @@ async def _process_task_result(
         _logger.debug("received %s result: %s", f"{task=}", f"{result=}")
         task_final_state = RunningState.FAILED
         simcore_platform_status = SimcorePlatformStatus.OK
-        errors: list[ErrorDict] = []
+        task_errors: list[ErrorDict] = []
+        task_completed = True
 
-        if task.job_id is not None:
-            (
-                _service_key,
-                _service_version,
-                user_id,
-                project_id,
-                node_id,
-            ) = parse_dask_job_id(task.job_id)
+        assert task.job_id  # nosec
+        (
+            _service_key,
+            _service_version,
+            user_id,
+            project_id,
+            node_id,
+        ) = parse_dask_job_id(task.job_id)
 
-            assert task.project_id == project_id  # nosec
-            assert task.node_id == node_id  # nosec
+        assert task.project_id == project_id  # nosec
+        assert task.node_id == node_id  # nosec
+        assert task.job_id  # nosec
+        log_error_context = {
+            "user_id": user_id,
+            "project_id": project_id,
+            "node_id": node_id,
+            "job_id": task.job_id,
+        }
 
+        if isinstance(result, TaskOutputData):
+            # That means the task successfully completed
             try:
-                if isinstance(result, TaskOutputData):
-                    # success!
-                    await parse_output_data(
-                        self.db_engine,
-                        task.job_id,
-                        result,
+                await parse_output_data(
+                    self.db_engine,
+                    task.job_id,
+                    result,
+                )
+                task_final_state = RunningState.SUCCESS
+            except PortsValidationError as err:
+                _logger.exception(
+                    **create_troubleshootting_log_kwargs(
+                        "Unexpected error while parsing output data, comp_tasks/comp_pipeline is not in sync with what was started",
+                        error=err,
+                        error_context=log_error_context,
+                    )
+                )
+                task_errors.extend(err.get_errors())
+                task_final_state = RunningState.FAILED
+                # NOTE: simcore platform state is still OK as the task ran fine, the issue is likely due to the service labels
+        elif isinstance(result, ComputationalBackendTaskResultsNotReadyError):
+            # Task result retrieval failed due to communication error, task will be retried
+            # so we keep it as is
+            _logger.warning(
+                **create_troubleshootting_log_kwargs(
+                    f"Retrieval of task {task.job_id} result timed-out",
+                    error=result,
+                    error_context=log_error_context,
+                    tip="This can happen if the computational backend is overloaded with requests. It will be automatically retried again.",
+                )
+            )
+
+            if task.errors:
+                for error in task.errors:
+                    if error["type"] == _TASK_RETRIEVAL_ERROR_TYPE:
+                        # already had a timeout error, let's keep it
+                        task_errors.append(error)
+                        break
+            if not task_errors:
+                # first time we have this error
+                task_errors.append(
+                    ErrorDict(
+                        loc=(f"{task.project_id}", f"{task.node_id}"),
+                        msg=f"{result}",
+                        type=_TASK_RETRIEVAL_ERROR_TYPE,
+                        ctx={
+                            _TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY: f"{arrow.utcnow()}",
+                            "user_id": user_id,
+                            "project_id": project_id,
+                            "node_id": node_id,
+                            "job_id": task.job_id,
+                        },
+                    )
+                )
+
+            task_completed = False
+        else:
+            # the task itself failed, check why
+            if isinstance(result, TaskCancelledError):
+                _logger.info(
+                    **create_troubleshootting_log_kwargs(
+                        f"Task {task.job_id} was cancelled",
+                        error=result,
+                        error_context=log_error_context,
                     )
-                    task_final_state = RunningState.SUCCESS
-                elif isinstance(result, ComputationalBackendTaskResultsNotReadyError):
-                    # Task result retrieval failed due to communication error, task will be retried
-                    # so we keep it as is
-                    assert task.job_id  # nosec
-                    return False, task.job_id
-                else:
-                    if isinstance(result, TaskCancelledError):
-                        task_final_state = RunningState.ABORTED
+                )
+                task_final_state = RunningState.ABORTED
 
-                    else:
-                        task_final_state = RunningState.FAILED
-                        errors.append(
-                            {
-                                "loc": (
-                                    f"{task.project_id}",
-                                    f"{task.node_id}",
-                                ),
-                                "msg": f"{result}",
-                                "type": "runtime",
-                            }
-                        )
-                        if isinstance(result, ComputationalBackendNotConnectedError):
-                            simcore_platform_status = SimcorePlatformStatus.BAD
-                    # we need to remove any invalid files in the storage
-                    await clean_task_output_and_log_files_if_invalid(
-                        self.db_engine, user_id, project_id, node_id
+            else:
+                _logger.info(
+                    **create_troubleshootting_log_kwargs(
+                        f"Task {task.job_id} completed with errors",
+                        error=result,
+                        error_context=log_error_context,
                     )
-            except TaskSchedulingError as err:
+                )
                 task_final_state = RunningState.FAILED
-                simcore_platform_status = SimcorePlatformStatus.BAD
-                errors = err.get_errors()
-                _logger.debug(
-                    "Unexpected failure while processing results of %s: %s",
-                    f"{task=}",
-                    f"{errors=}",
+                task_errors.append(
+                    ErrorDict(
+                        loc=(f"{task.project_id}", f"{task.node_id}"),
+                        msg=f"{result}",
+                        type="runtime",
+                    )
                 )
 
+                if isinstance(result, ComputationalBackendNotConnectedError):
+                    simcore_platform_status = SimcorePlatformStatus.BAD
+            # we need to remove any invalid files in the storage
+            await clean_task_output_and_log_files_if_invalid(
+                self.db_engine, user_id, project_id, node_id
+            )
+
+        if task_completed:
             # resource tracking
             await publish_service_resource_tracking_stopped(
                 self.rabbitmq_client,
@@ -408,17 +467,25 @@ async def _process_task_result(
                 task_final_state=task_final_state,
             )
 
-        await CompTasksRepository(self.db_engine).update_project_tasks_state(
-            task.project_id,
-            run_id,
-            [task.node_id],
-            task_final_state,
-            errors=errors,
-            optional_progress=1,
-            optional_stopped=arrow.utcnow().datetime,
-        )
-        assert task.job_id  # nosec
-        return True, task.job_id
+            await CompTasksRepository(self.db_engine).update_project_tasks_state(
+                task.project_id,
+                run_id,
+                [task.node_id],
+                task_final_state,
+                errors=task_errors,
+                optional_progress=1,
+                optional_stopped=arrow.utcnow().datetime,
+            )
+        else:
+            await CompTasksRepository(self.db_engine).update_project_tasks_state(
+                task.project_id,
+                run_id,
+                [task.node_id],
+                RunningState.STARTED,  # keep the same state as before
+                errors=task_errors,
+            )
+
+        return task_completed, task.job_id
 
     async def _task_progress_change_handler(
         self, event: tuple[UnixTimestamp, Any]

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

@@ -524,6 +524,8 @@ async def abort_computation_task(self, job_id: str) -> None:
 
     async def get_task_result(self, job_id: str) -> TaskOutputData:
         _logger.debug("getting result of %s", f"{job_id=}")
+        dask_utils.check_communication_with_scheduler_is_open(self.backend.client)
+        dask_utils.check_scheduler_status(self.backend.client)
         try:
             task_future: distributed.Future = (
                 await dask_utils.wrap_client_async_routine(

services/director-v2/tests/unit/with_dbs/comp_scheduler/test_scheduler_dask.py

@@ -11,6 +11,7 @@
 
 import asyncio
 import datetime
+import random
 from collections.abc import AsyncIterator, Awaitable, Callable
 from copy import deepcopy
 from dataclasses import dataclass
@@ -33,6 +34,7 @@
 from faker import Faker
 from fastapi.applications import FastAPI
 from models_library.computations import CollectionRunID
+from models_library.errors import ErrorDict
 from models_library.projects import ProjectAtDB, ProjectID
 from models_library.projects_nodes_io import NodeID
 from models_library.projects_state import RunningState
@@ -72,6 +74,8 @@
     BaseCompScheduler,
 )
 from simcore_service_director_v2.modules.comp_scheduler._scheduler_dask import (
+    _TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY,
+    _TASK_RETRIEVAL_ERROR_TYPE,
     DaskScheduler,
 )
 from simcore_service_director_v2.modules.comp_scheduler._utils import COMPLETED_STATES
@@ -2275,14 +2279,7 @@ async def test_run_new_pipeline_called_twice_prevents_duplicate_runs(
     )
 
 
-@pytest.mark.parametrize(
-    "exception_type",
-    [
-        ComputationalBackendTaskResultsNotReadyError,
-    ],
-)
 async def test_getting_task_result_raises_exception_does_not_fail_task_and_retries(
-    exception_type: Exception,
     with_disabled_auto_scheduling: mock.Mock,
     with_disabled_scheduler_publisher: mock.Mock,
     mocked_dask_client: mock.MagicMock,
@@ -2300,13 +2297,21 @@ async def mocked_get_tasks_status(job_ids: list[str]) -> list[RunningState]:
         return [RunningState.SUCCESS for j in job_ids]
 
     mocked_dask_client.get_tasks_status.side_effect = mocked_get_tasks_status
-    call_count = 0
 
-    async def mocked_get_task_result(_job_id: str) -> TaskOutputData:
-        nonlocal call_count
-        call_count += 1
-        if call_count > 1:
-            raise exception_type
+    computational_tasks = [
+        t for t in running_project.tasks if t.node_class is NodeClass.COMPUTATIONAL
+    ]
+    expected_timeouted_tasks = random.choices(  # noqa: S311
+        computational_tasks, k=len(computational_tasks) - 1
+    )
+    successful_tasks = [
+        t for t in computational_tasks if t not in expected_timeouted_tasks
+    ]
+
+    async def mocked_get_task_result(job_id: str) -> TaskOutputData:
+        nonlocal expected_timeouted_tasks
+        if job_id in [t.job_id for t in expected_timeouted_tasks]:
+            raise ComputationalBackendTaskResultsNotReadyError(job_id=job_id)
         return TaskOutputData.model_validate({"whatever_output": 123})
 
     mocked_dask_client.get_task_result.side_effect = mocked_get_task_result
@@ -2317,10 +2322,49 @@ async def mocked_get_task_result(_job_id: str) -> TaskOutputData:
         project_id=running_project.project.uuid,
         iteration=1,
     )
-    assert mocked_dask_client.get_task_result.call_count == len(
-        [t for t in running_project.tasks if t.node_class is NodeClass.COMPUTATIONAL]
-    )
+    assert mocked_dask_client.get_task_result.call_count == len(computational_tasks)
     mocked_dask_client.get_task_result.reset_mock()
+
+    # check the tasks in the DB, the error shall be set there and the task state is set back to STARTED
+    comp_tasks, _ = await assert_comp_tasks_and_comp_run_snapshot_tasks(
+        sqlalchemy_async_engine,
+        project_uuid=running_project.project.uuid,
+        task_ids=[t.node_id for t in expected_timeouted_tasks],
+        expected_state=RunningState.STARTED,
+        expected_progress=0,
+        run_id=running_project.runs.run_id,
+    )
+    # we should have an error in all these comp_tasks
+    retrieval_times = []
+    for t in comp_tasks:
+        assert t.errors
+        assert len(t.errors) == 1
+        error_dict = TypeAdapter(ErrorDict).validate_python(t.errors[0])
+        assert error_dict["type"] == _TASK_RETRIEVAL_ERROR_TYPE
+        assert "ctx" in error_dict
+        assert _TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY in error_dict["ctx"]
+        retrieval_times.append(
+            error_dict["ctx"][_TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY]
+        )
+
+    await assert_comp_tasks_and_comp_run_snapshot_tasks(
+        sqlalchemy_async_engine,
+        project_uuid=running_project.project.uuid,
+        task_ids=[t.node_id for t in successful_tasks],
+        expected_state=RunningState.SUCCESS,
+        expected_progress=1.0,
+        run_id=running_project.runs.run_id,
+    )
+    await assert_comp_runs(
+        sqlalchemy_async_engine,
+        expected_total=1,
+        expected_state=RunningState.STARTED,
+        where_statement=and_(
+            comp_runs.c.user_id == running_project.project.prj_owner,
+            comp_runs.c.project_uuid == f"{running_project.project.uuid}",
+        ),
+    )
+
     # calling again should not raise neither but try again
     assert running_project.project.prj_owner
     for _ in range(3):
@@ -2341,3 +2385,32 @@ async def mocked_get_task_result(_job_id: str) -> TaskOutputData:
         )
         mocked_dask_client.get_task_result.reset_mock()
         await asyncio.sleep(0.5)  # wait a bit to ensure the retry decorator has reset
+    comp_tasks, _ = await assert_comp_tasks_and_comp_run_snapshot_tasks(
+        sqlalchemy_async_engine,
+        project_uuid=running_project.project.uuid,
+        task_ids=[t.node_id for t in expected_timeouted_tasks],
+        expected_state=RunningState.STARTED,
+        expected_progress=0,
+        run_id=running_project.runs.run_id,
+    )
+    # the times shall remain the same
+    for t in comp_tasks:
+        assert t.errors
+        assert len(t.errors) == 1
+        error_dict = TypeAdapter(ErrorDict).validate_python(t.errors[0])
+        assert error_dict["type"] == _TASK_RETRIEVAL_ERROR_TYPE
+        assert "ctx" in error_dict
+        assert _TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY in error_dict["ctx"]
+        # the time shall be the same as before
+        assert (
+            error_dict["ctx"][_TASK_RETRIEVAL_ERROR_CONTEXT_TIME_KEY] in retrieval_times
+        )
+    await assert_comp_runs(
+        sqlalchemy_async_engine,
+        expected_total=1,
+        expected_state=RunningState.STARTED,
+        where_statement=and_(
+            comp_runs.c.user_id == running_project.project.prj_owner,
+            comp_runs.c.project_uuid == f"{running_project.project.uuid}",
+        ),
+    )