diff --git a/src/rhapsody/backends/execution/dragon.py b/src/rhapsody/backends/execution/dragon.py index 9098034..e3ac46c 100644 --- a/src/rhapsody/backends/execution/dragon.py +++ b/src/rhapsody/backends/execution/dragon.py @@ -1,4 +1,6 @@ import asyncio +import functools +import glob import logging import os import shlex @@ -363,6 +365,7 @@ def _aggregate_function_results( "exit_code": result.exit_code, "return_value": result.return_value, "exception": result.exception, + "traceback": result.traceback, "success": result.success, } else: @@ -389,6 +392,10 @@ def _aggregate_function_results( "exception": None if all_successful else "; ".join(str(r.exception) for r in results if not r.success), + "traceback": None + if all_successful + else "\n".join(r.traceback for r in results + if not r.success and r.traceback), "success": all_successful, } @@ -1656,6 +1663,38 @@ class TaskStateMapperV3: terminal_states = {DONE, FAILED, CANCELED} +def _v3_function_wrapper(user_func, stdout_path, stderr_path, *args, **kwargs): + """Run *user_func* with sys.stdout/stderr redirected to per-rank files. + + Dragon batch surfaces only a dragon-side ``DragonUserCodeError`` when a + ProcessGroup rank exits non-zero — the rank's own Python traceback is + dropped. Persisting it to ``.`` before re-raising + lets ``_deliver_batch`` read it post-mortem. + """ + import os + import sys + import traceback + + rank = (os.environ.get("DRAGON_RANK") + or os.environ.get("PMI_RANK") + or os.environ.get("OMPI_COMM_WORLD_RANK") + or os.environ.get("PALS_RANKID") + or os.environ.get("SLURM_PROCID") + or f"pid{os.getpid()}") + + with open(f"{stdout_path}.{rank}", "w") as out, \ + open(f"{stderr_path}.{rank}", "w") as err: + old_out, old_err = sys.stdout, sys.stderr + sys.stdout, sys.stderr = out, err + try: + return user_func(*args, **kwargs) + except BaseException: + traceback.print_exc(file=err) + raise + finally: + sys.stdout, sys.stderr = old_out, old_err + + # ============================================================================ # Main Backend Classes # V1 Integrates with Dragon HPC Native API @@ -3218,9 +3257,19 @@ def _monitor_loop(self): )[tuid] except KeyError: continue - self._monitored_batches.pop(tuid, None) - completed.append((uid, result, tb, raised, stdout, stderr)) + # Drain the per-rank function stdout/stderr files here, in + # the monitor thread, so the file I/O never blocks the + # asyncio loop that runs `_deliver_batch`. Executable tasks + # redirect via a shell script (`script_path` set) and have no + # per-rank files to drain. + fn_stdout = fn_stderr = None + info = self._task_registry.get(uid) + if info and not info.get("script_path"): + fn_stdout = self._drain_rank_files(info.get("stdout_path")) + fn_stderr = self._drain_rank_files(info.get("stderr_path")) + completed.append((uid, result, tb, raised, stdout, stderr, + fn_stdout, fn_stderr)) # One cross-thread wakeup for the entire sweep batch if completed: @@ -3235,14 +3284,46 @@ def _monitor_loop(self): self.logger.debug("Dragon batch monitor loop stopped") + @staticmethod + def _drain_rank_files(path: str) -> str: + """Read and unlink all per-rank ``.`` files written by + :func:`_v3_function_wrapper`, returning their concatenated content. + + Runs in the monitor thread (never the event loop), reads both the + success and failure output of every rank, and removes the files + afterwards so they do not accumulate in the work dir. + """ + if not path: + return "" + chunks = [] + for f in sorted(glob.glob(f"{path}.*")): + try: + # Worker output may contain non-UTF-8 bytes (C libs, progress + # bars); replace rather than drop the rank's whole output. + with open(f, encoding="utf-8", errors="replace") as fh: + content = fh.read() + except Exception: + content = "" + try: + os.unlink(f) + except OSError: + pass + if content: + rank = f.rsplit(".", 1)[-1] + chunks.append(f"[rank {rank}]\n{content}") + return "\n".join(chunks) + def _deliver_batch(self, completions: list) -> None: """Deliver a batch of completed tasks. Runs on the asyncio event loop (via call_soon_threadsafe). Called once per monitor sweep with all tasks that completed in that sweep, reducing cross- - thread wakeups from O(tasks) to O(sweeps). + thread wakeups from O(tasks) to O(sweeps). Per-rank function output has + already been drained off-loop by the monitor thread and arrives as + ``fn_stdout`` / ``fn_stderr`` strings, so this method does no file I/O. """ - for uid, result, tb, raised, stdout, stderr in completions: + for (uid, result, tb, raised, stdout, stderr, + fn_stdout, fn_stderr) in completions: task_info = self._task_registry.pop(uid, None) if not task_info: continue @@ -3252,15 +3333,48 @@ def _deliver_batch(self, completions: list) -> None: task_desc = task_info["description"] stdout_path = task_info.get("stdout_path") stderr_path = task_info.get("stderr_path") + # Executable redirect writes one literal file per path (script_path + # set), so callers receive the path. Function tasks stream to + # per-rank files already folded into fn_stdout / fn_stderr. + is_exec_redirect = bool(task_info.get("script_path")) if raised: - task_desc["exception"] = result - task_desc["stderr"] = stderr_path if stderr_path else (tb if tb else str(result)) - task_desc["stdout"] = stdout_path or stdout or "" + if is_exec_redirect: + task_desc["exception"] = result + task_desc["stderr"] = stderr_path or (tb if tb else str(result)) + task_desc["stdout"] = stdout_path or stdout or "" + else: + diagnostic = fn_stderr or "" + if diagnostic: + # Fold the rank's own traceback into the exception so a + # re-raise surfaces the real cause, not just Dragon's + # generic DragonUserCodeError. + try: + cls = type(result) + # A non-BaseException result (e.g. a string from a + # serialization fallback) would build fine but never + # be raised by session.py, silently masking the + # failure as success — coerce it to RuntimeError. + if not issubclass(cls, BaseException): + raise TypeError("result is not a BaseException") + augmented = cls(f"{result}\n--- worker output ---\n{diagnostic}") + except Exception: + augmented = RuntimeError( + f"{result}\n--- worker output ---\n{diagnostic}" + ) + task_desc["exception"] = augmented + else: + task_desc["exception"] = result + task_desc["stderr"] = diagnostic or stderr or (tb if tb else str(result)) + task_desc["stdout"] = fn_stdout or stdout or "" self._callback_func(task_desc, "FAILED") else: task_desc["return_value"] = result - task_desc["stdout"] = stdout_path or stdout or "" - task_desc["stderr"] = stderr_path or stderr or "" + if is_exec_redirect: + task_desc["stdout"] = stdout_path or stdout or "" + task_desc["stderr"] = stderr_path or stderr or "" + else: + task_desc["stdout"] = fn_stdout or stdout or "" + task_desc["stderr"] = fn_stderr or stderr or "" self._callback_func(task_desc, "DONE") async def submit_tasks(self, tasks: list[dict]) -> None: @@ -3401,6 +3515,15 @@ def target(*a, **kw): target = "/bin/bash" task_args = (script_path,) + # Wrap every function target so each rank persists its own + # stdout/stderr; see _v3_function_wrapper for why. + if is_function: + stdout_path = os.path.join(self._work_dir, f"{uid}.stdout") + stderr_path = os.path.join(self._work_dir, f"{uid}.stderr") + target = functools.partial( + _v3_function_wrapper, target, stdout_path, stderr_path + ) + def _build_process_template_kwargs(template_cfg: dict[str, Any]) -> dict[str, Any]: """Build ProcessTemplate kwargs: stdout_pipe as default; user config overrides.""" return {"stdout": stdout_pipe, **template_cfg, "args": task_args, "kwargs": task_kwargs} diff --git a/tests/unit/test_backend_execution_dragon.py b/tests/unit/test_backend_execution_dragon.py index 5aa8551..b4a1d84 100644 --- a/tests/unit/test_backend_execution_dragon.py +++ b/tests/unit/test_backend_execution_dragon.py @@ -644,7 +644,7 @@ def test_v3_deliver_batch_success_stores_value_and_fires_done(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 42, None, False, None, None)]) + backend_v3._deliver_batch([(uid, 42, None, False, None, None, None, None)]) assert task_desc["return_value"] == 42 assert task_desc["stdout"] == "" @@ -659,7 +659,7 @@ def test_v3_deliver_batch_propagates_stdout_stderr(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n")]) + backend_v3._deliver_batch([(uid, "ok", None, False, "hello\n", "warn\n", None, None)]) assert task_desc["stdout"] == "hello\n" assert task_desc["stderr"] == "warn\n" @@ -673,7 +673,7 @@ def test_v3_deliver_batch_failure_stores_exc_and_fires_failed(backend_v3): exc = RuntimeError("something went wrong") # raised=True, tb=None: stderr falls back to str(exc) - backend_v3._deliver_batch([(uid, exc, None, True, None, None)]) + backend_v3._deliver_batch([(uid, exc, None, True, None, None, None, None)]) assert task_desc["exception"] is exc assert "something went wrong" in task_desc["stderr"] @@ -689,11 +689,29 @@ def test_v3_deliver_batch_prefers_traceback_over_str_exc(backend_v3): exc = RuntimeError("boom") tb = "Traceback (most recent call last):\n File ...\nRuntimeError: boom" - backend_v3._deliver_batch([(uid, exc, tb, True, None, None)]) + backend_v3._deliver_batch([(uid, exc, tb, True, None, None, None, None)]) assert task_desc["stderr"] == tb +def test_v3_deliver_batch_non_exception_result_coerced_to_baseexception(backend_v3): + """A non-BaseException failure result must still surface as a raisable + exception, otherwise session.py would silently treat the failure as DONE.""" + uid = "task.unit-failed-str" + task_desc = {"uid": uid} + backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} + + # result is a plain string (e.g. a cross-node serialization fallback) and a + # per-rank diagnostic is present, so the augmentation path runs. + backend_v3._deliver_batch( + [(uid, "remote failure", None, True, None, None, None, "[rank 0]\ntb")] + ) + + assert isinstance(task_desc["exception"], BaseException) + assert "remote failure" in str(task_desc["exception"]) + backend_v3._callback_func.assert_called_once_with(task_desc, "FAILED") + + def test_v3_cancelled_task_skips_callback(backend_v3): """_deliver_batch is a no-op for UIDs in _cancelled_tasks.""" uid = "task.unit-cancelled" @@ -701,7 +719,7 @@ def test_v3_cancelled_task_skips_callback(backend_v3): backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} backend_v3._cancelled_tasks.add(uid) - backend_v3._deliver_batch([(uid, 99, None, False, None, None)]) + backend_v3._deliver_batch([(uid, 99, None, False, None, None, None, None)]) backend_v3._callback_func.assert_not_called() assert "return_value" not in task_desc @@ -1036,7 +1054,7 @@ def test_v3_deliver_batch_empty_dragon_stdout_written_as_empty_string(backend_v3 task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_v3._deliver_batch([(uid, 0, None, False, "", "", None, None)]) assert isinstance(task_desc["stdout"], str) assert isinstance(task_desc["stderr"], str) @@ -1053,9 +1071,13 @@ def test_v3_deliver_batch_stdout_path_takes_priority_over_empty_stdout(backend_v "description": task_desc, "stdout_path": "/work/uid.stdout", "stderr_path": "/work/uid.stderr", + # capture_stdio redirect writes a wrapper script alongside the literal + # stdout/stderr files; script_path marks this as an exec redirect (vs a + # function-wrap glob prefix), which is what makes the path take priority. + "script_path": "/work/uid.sh", } - backend_v3._deliver_batch([(uid, 0, None, False, "", "")]) + backend_v3._deliver_batch([(uid, 0, None, False, "", "", None, None)]) assert task_desc["stdout"] == "/work/uid.stdout" assert task_desc["stderr"] == "/work/uid.stderr" @@ -1068,6 +1090,6 @@ def test_v3_deliver_batch_failed_stdout_is_string(backend_v3): task_desc = {"uid": uid} backend_v3._task_registry[uid] = {"uid": uid, "description": task_desc} - backend_v3._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "")]) + backend_v3._deliver_batch([(uid, RuntimeError("boom"), None, True, "", "", None, None)]) assert isinstance(task_desc["stdout"], str)