diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 4eb4caa..ea4f18a 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -298,6 +298,9 @@ def _deliver_batch(self, completions: list) -> None: continue task_desc = task_info["description"] if raised: + # surface the remote traceback in the endpoint log -- the + # consumer side only ever sees the exception message + self.logger.error(f"task {uid} failed: {result}\n{tb or '(no traceback captured)'}") task_desc["exception"] = result task_desc["stderr"] = stderr if stderr else (tb if tb else str(result)) task_desc["stdout"] = stdout @@ -498,7 +501,13 @@ def get_task_states_map(self): async def cancel_task(self, uid: str) -> bool: if uid not in self._task_registry: - raise ValueError(f"Task {uid} not found") + # Idempotent, like the concurrent backend: a cancel can race a + # task that already finished or was never dispatched (ROSE + # cancels its losing candidate branch this way). Raising here + # breaks the caller's teardown and turns a clean branch-cancel + # into a DependencyFailureError downstream. + self.logger.debug(f"cancel_task: {uid} not tracked (already done?)") + return False batch_task = self._task_registry[uid]["batch_task"] loop = asyncio.get_running_loop() diff --git a/src/rhapsody/backends/execution/orbit.py b/src/rhapsody/backends/execution/orbit.py index 0a8f3ba..e0bd67d 100644 --- a/src/rhapsody/backends/execution/orbit.py +++ b/src/rhapsody/backends/execution/orbit.py @@ -670,7 +670,12 @@ async def cancel_task(self, uid: str) -> bool: if not buffered: await asyncio.to_thread(self._rh.cancel_task, uid) - task = self._tasks[uid] + # the task may have completed under the in-flight cancel -- its + # terminal callback then already fired via the remote notification + # and it left `_tasks`; nothing remains to cancel + task = self._tasks.get(uid) + if task is None: + return False task["state"] = "CANCELED" self._fire_callback(task, "CANCELED") return True diff --git a/tests/unit/test_backend_dragon_unit.py b/tests/unit/test_backend_dragon_unit.py new file mode 100644 index 0000000..705b2bd --- /dev/null +++ b/tests/unit/test_backend_dragon_unit.py @@ -0,0 +1,72 @@ +"""Dragon-backend behavior that must hold WITHOUT a dragon install. + +The full dragon test module skips when the runtime is absent, so the +contracts fixed for the remote DT demo are pinned here on a bare +instance: cancel of an untracked task is idempotent, and a failed +delivery logs the remote traceback endpoint-side. +""" + +import logging + +import pytest + +from rhapsody.backends.execution.dragon import DragonExecutionBackend + + +def bare_backend(): + """The slice of the backend the methods under test touch.""" + + backend = object.__new__(DragonExecutionBackend) + backend.logger = logging.getLogger("test.dragon.unit") + backend._task_registry = {} + backend._cancelled_tasks = set() + + delivered = [] + backend._callback_func = lambda task, state: delivered.append((dict(task), state)) + + return backend, delivered + + +@pytest.mark.asyncio +async def test_cancel_of_an_untracked_task_is_a_noop(): + """A cancel can race a task that already finished or was never + dispatched -- ROSE cancels its losing candidate branch exactly this + way. Raising here broke the caller's teardown and turned a clean + branch-cancel into a DependencyFailureError on the surviving + pipeline.""" + + backend, delivered = bare_backend() + + assert await backend.cancel_task("task.000000") is False + assert delivered == [] + + +def test_a_failed_delivery_logs_the_remote_traceback(caplog): + """The worker captures the traceback and the consumer side only ever sees the exception message + -- the endpoint log is where it must surface.""" + + backend, delivered = bare_backend() + backend._task_registry["task.000001"] = {"description": {"uid": "task.000001"}} + + with caplog.at_level(logging.ERROR, logger="test.dragon.unit"): + backend._deliver_batch( + [ + ( + "task.000001", + "boom happened", + "Traceback (most recent call last):\n ...", + True, + "", + "", + ), + ] + ) + + assert len(delivered) == 1 + task, state = delivered[0] + assert state == "FAILED" + assert task["exception"] == "boom happened" + assert "Traceback" in task["stderr"] + + assert "task.000001 failed: boom happened" in caplog.text + assert "Traceback (most recent call last):" in caplog.text diff --git a/tests/unit/test_backend_execution_orbit.py b/tests/unit/test_backend_execution_orbit.py index 9b21927..2424327 100644 --- a/tests/unit/test_backend_execution_orbit.py +++ b/tests/unit/test_backend_execution_orbit.py @@ -210,6 +210,27 @@ async def test_cancel_unknown_task(): assert result is False +@pytest.mark.asyncio +async def test_cancel_racing_a_completion_is_a_noop(): + """The task completes while the remote cancel is in flight: its + terminal callback already fired via the notification and it left the + registry -- the cancel must answer False, not KeyError (observed as + teardown noise on every twin_close under load).""" + + backend = await _init_backend() + backend._tasks["t.001"] = {"uid": "t.001", "state": "RUNNING"} + cb = MagicMock() + backend.register_callback(cb) + + def completes_meanwhile(uid): + backend._tasks.pop(uid, None) + + backend._mock_rh.cancel_task.side_effect = completes_meanwhile + + assert await backend.cancel_task("t.001") is False + cb.assert_not_called() + + # --------------------------------------------------------------------------- # cancel_all_tasks # ---------------------------------------------------------------------------