Skip to content

backends.dragon V3: capture per-rank stderr for function tasks #51

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andre-merzky wants to merge 1 commit into
base: main
Choose a base branch
from
Open

backends.dragon V3: capture per-rank stderr for function tasks #51

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 132 additions & 9 deletions src/rhapsody/backends/execution/dragon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import asyncio
import functools
import glob
import logging
import os
import shlex
Expand Down Expand Up @@ -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:
Expand All @@ -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,
}

Expand Down Expand Up @@ -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 ``<stderr_path>.<rank>`` 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
Expand Down Expand Up @@ -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:
Expand All @@ -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 ``<path>.<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 = ""
Comment on lines +3299 to +3306

@gemini-code-assist gemini-code-assist Bot Jun 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In _drain_rank_files, reading the file with the default system encoding can raise a UnicodeDecodeError if the rank's output contains non-UTF-8 binary data (e.g., from a C library or progress bar). This would cause the entire output of that rank to be discarded. Specifying encoding="utf-8" and errors="replace" makes the file reading much more robust.

Suggested change
for f in sorted(glob.glob(f"{path}.*")):
try:
with open(f) as fh:
content = fh.read()
except Exception:
content = ""
for f in sorted(glob.glob(f"{path}.*")):
try:
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
Expand All @@ -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}"
)
Comment on lines +3351 to +3363

@gemini-code-assist gemini-code-assist Bot Jun 30, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In _deliver_batch, cls = type(result) is used to instantiate an augmented exception. If result is not a BaseException subclass (for example, if it is a string due to serialization fallbacks across nodes), cls(...) will succeed (e.g., str(...)) but augmented will be a string. Since session.py only raises exceptions if they are instances of BaseException, a string exception will be bypassed, causing the failed task to be silently treated as successful! Adding an issubclass(cls, BaseException) check prevents this critical issue.

Suggested change
try:
cls = type(result)
augmented = cls(f"{result}\n--- worker output ---\n{diagnostic}")
except Exception:
augmented = RuntimeError(
f"{result}\n--- worker output ---\n{diagnostic}"
)
try:
cls = type(result)
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")
Comment on lines 3340 to 3378

@gemini-code-assist gemini-code-assist Bot May 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current implementation of _deliver_batch for V3 function tasks has several issues:

  1. Lost Output: Since _v3_function_wrapper redirects sys.stdout and sys.stderr to files, Dragon's default capture mechanism (the stdout and stderr variables at line 3266) will return empty strings. However, this method only reads the stderr files on failure (lines 3281-3292) and never reads the stdout files. Consequently, all stdout is lost for function tasks, and all output is lost for successful function tasks.
  2. Blocking I/O: The use of glob.glob and fh.read() inside _deliver_batch (which runs on the event loop via call_soon_threadsafe) will block the loop. For jobs with many ranks or large output files, this can cause significant latency or freeze the application.
  3. Resource Leak: The per-rank output files created in the session's work directory are never deleted, which may lead to excessive disk usage over time.

Consider reading both stdout and stderr files for all function tasks and deleting them after they are processed. To avoid blocking the event loop, perform the file I/O in a separate thread (e.g., using asyncio.to_thread or moving the logic to the _monitor_loop thread).


async def submit_tasks(self, tasks: list[dict]) -> None:
Expand Down Expand Up @@ -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}
Expand Down
38 changes: 30 additions & 8 deletions tests/unit/test_backend_execution_dragon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"] == ""
Expand All @@ -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"
Expand All @@ -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"]
Expand All @@ -689,19 +689,37 @@ 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"
task_desc = {"uid": uid}
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
Expand Down Expand Up @@ -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)
Expand All @@ -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"
Expand All @@ -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)
Loading