dragon+orbit backends: idempotent cancels; dragon failure tracebacks

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()

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

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

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
 # ---------------------------------------------------------------------------