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

[andre-merzky]

Summary

• Wrap V3 function-task targets in a module-level helper (_v3_function_wrapper) that redirects sys.stdout/sys.stderr to per-rank files under self._work_dir and writes the Python traceback to the stderr file before re-raising.
• _deliver_batch globs the per-rank files on failure and folds their contents into the augmented exception's message, preserving the original exception type so callers re-raising via task['exception'] see the real cause.
• Gated to function tasks only — executable tasks (capture_stdio shell-redirect path) are unaffected.

Closes #50.

Why

Without this change, a function-task rank that exits non-zero leaves the awaiter with only RuntimeError: Error(s) in user-provided code resulted in exit codes: [1] / Dragon Error Code: DRAGON_USER_CODE_ERROR and no information about why — see the issue for the full repro. The fix surfaces the actual rank-side Python traceback under a --- worker output --- section in the augmented exception's message:

RuntimeError: Error(s) in user-provided code resulted in exit codes: [1]
Dragon Error Code: DRAGON_USER_CODE_ERROR
--- worker output ---
[rank pid716750]
Traceback (most recent call last):
  File ".../dragon.py", line 1695, in _v3_function_wrapper
    return user_func(*args, **kwargs)
  File "/tmp/repro.py", line 31, in bad_function
    return x["nonexistent_key"]
KeyError: 'nonexistent_key'
[rank pid716751]
...

Implementation notes

• _v3_function_wrapper is module-level (picklable) and uses functools.partial to bake in stdout_path/stderr_path/user_func before dragon's batch pickles the ProcessTemplate.
• Rank label prefers MPI runtime env vars (PMI_RANK / OMPI_COMM_WORLD_RANK / PALS_RANKID / SLURM_PROCID) and falls back to pid<N> for non-MPI ProcessGroup ranks so distinct ranks land in distinct files.
• _deliver_batch uses task_info['script_path'] (set only for the executable redirect path) to distinguish glob-prefix paths (function wrap) from literal file paths (exec redirect), so executable tasks keep their existing behaviour.
• The augmented exception only contains per-rank captured tracebacks — dragon's batch-level tb (the internal _do_task_impl → grp.join() → _manage_exception trace) is intentionally not included; str(result) already carries the user-facing dragon error message.

Test plan

• Verified locally (single-node, dragon batch with two ranks) using a minimal repro: one task that succeeds, one task that raises KeyError from inside process_templates=[(2, {})].
  • Successful task returns normally; per-rank capture files are present but ignored on the success path.
  • Failed task surfaces the KeyError traceback for both ranks under --- worker output --- and the awaiter's RuntimeError carries the augmented message.
• Existing rhapsody unit/integration tests untouched and still pass — please re-run CI.
• Verify on a multi-node HPC environment (left for reviewer / next test cycle).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces two new execution backends: EdgeExecutionBackend, which delegates task execution to remote RADICAL Edge nodes with support for submission batching and Python version compatibility checks, and NoopExecutionBackend for performance benchmarking. It also enhances the Dragon V3 backend by implementing per-rank stdout/stderr redirection and traceback capture for function tasks. Feedback highlights several critical issues in the Dragon V3 implementation, including potential output loss, blocking I/O on the event loop, and resource leaks. Additionally, inconsistencies in the Edge backend's default configuration and interface mismatches in the No-op backend were identified.

[gemini-code-assist]

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

[gemini-code-assist]

The rank detection logic in the wrapper should include DRAGON_RANK. This is the standard environment variable used by Dragon to identify ranks within a process group. Including it ensures consistency with the V1 and V2 backends and provides a reliable fallback for non-MPI environments.

Suggested change

	rank = (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()}")
	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()}")

[gemini-code-assist]

There is an inconsistency between the docstring and the implementation for the default batch_window. The docstring states it is 0.05 seconds, but _DEFAULT_BATCH_WINDOW is defined as 0.25 seconds. These should be aligned to avoid confusing users.

[gemini-code-assist]

The default remote backend is set to ['dragon_v3'], but the DragonExecutionBackendV3 class defaults its name to "dragon". Unless the remote session explicitly overrides the backend name, this default will cause a 'Backend not found' error. It is recommended to change the default to ['dragon'] or ensure it matches the standard registration name.

[gemini-code-assist]

The call to _get_rhapsody_handle() is synchronous and performs network I/O. Since it is called within an async method (_async_init), it will block the event loop. This should be wrapped in asyncio.to_thread to maintain responsiveness.

Suggested change

	self._rh = self._get_rhapsody_handle()
	self._rh = await asyncio.to_thread(self._get_rhapsody_handle)

[gemini-code-assist]

The submit_tasks method signature and return type do not match the BaseBackend interface, which defines it as returning None. Maintaining interface consistency is important for backend interoperability.

Suggested change

	async def submit_tasks(self, tasks: list[dict[str, Any]]) -> list[asyncio.Task]:
	if self._backend_state != BackendMainStates.RUNNING:
	self._backend_state = BackendMainStates.RUNNING
	submitted = []
	for task in tasks:
	task.update({
	"return_value": True,
	"stdout": "",
	"stderr": "",
	"exit_code": 0,
	})
	self.tasks[task["uid"]] = task
	future = asyncio.create_task(self._complete(task))
	submitted.append(future)
	return submitted
	async def submit_tasks(self, tasks: list[dict[str, Any]]) -> None:
	if self._backend_state != BackendMainStates.RUNNING:
	self._backend_state = BackendMainStates.RUNNING
	for task in tasks:
	task.update({
	"return_value": True,
	"stdout": "",
	"stderr": "",
	"exit_code": 0,
	})
	self.tasks[task["uid"]] = task
	asyncio.create_task(self._complete(task))

[gemini-code-assist]

The task_state_cb method signature does not match the BaseBackend interface, which expects (self, task: dict, state: str). This will cause a TypeError if the method is ever invoked.

Suggested change

	def task_state_cb(self):
	def task_state_cb(self, task: dict, state: str):

[andre-merzky]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces the EdgeExecutionBackend for remote task execution via RADICAL Edge and the NoopExecutionBackend for benchmarking, alongside updates to the Dragon V3 backend to capture per-rank stdout/stderr and tracebacks. The review feedback highlights several critical robustness improvements, such as ensuring deserialized exceptions are instances of BaseException to prevent silent task successes, avoiding file descriptor leaks using with statements, handling non-UTF-8 characters when draining rank files, and wrapping batch flushes and shutdown sequences in try...except blocks to prevent tasks from hanging or leaking resources during failures.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

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}"
	)

[gemini-code-assist]

In _apply_task_update, if the task failed (state == 'FAILED') or if exception/error is present in body, the exception might be a string or dict because it was serialized over the network. If it is not a BaseException instance, session.py will not recognize it as an exception and will silently treat the failed task as successful! Converting string/dict exceptions to RuntimeError ensures failures are correctly propagated to the client.

Suggested change

	# Decode base64-encoded return values (bytes results)
	if body.get('_return_value_encoding') == 'base64':
	import base64
	body['return_value'] = base64.b64decode(body['return_value'])
	del body['_return_value_encoding']
	# Update local task dict with remote results
	for key in ('state', 'stdout', 'stderr', 'exit_code',
	'return_value', 'exception', 'traceback', 'error'):
	if key in body:
	task[key] = body[key]
	# Decode base64-encoded return values (bytes results)
	if body.get('_return_value_encoding') == 'base64':
	import base64
	body['return_value'] = base64.b64decode(body['return_value'])
	del body['_return_value_encoding']
	# Ensure exception is a BaseException if the task failed, so session.py raises it
	if body.get('state') == 'FAILED' or 'exception' in body or 'error' in body:
	exc = body.get('exception') or body.get('error') or "Unknown remote error"
	if not isinstance(exc, BaseException):
	body['exception'] = RuntimeError(str(exc))
	# Update local task dict with remote results
	for key in ('state', 'stdout', 'stderr', 'exit_code',
	'return_value', 'exception', 'traceback', 'error'):
	if key in body:
	task[key] = body[key]

[gemini-code-assist]

In _flush_batch, if self._rh.submit_tasks raises an exception (e.g., due to a network error or remote service failure), the exception is unhandled because _locked_flush runs as a background task. The tasks in the batch are left in limbo and the client hangs indefinitely. Wrapping the submission in a try...except block and failing all tasks in the batch ensures that failures are correctly reported to the user.

    async def _flush_batch(self):
        """Send all buffered tasks in one bulk request.

        Must be called while holding ``_batch_lock``.
        """
        if not self._batch_buffer:
            return

        batch = self._batch_buffer
        self._batch_buffer = []

        if self._flush_handle is not None:
            self._flush_handle.cancel()
            self._flush_handle = None

        try:
            self._check_python_compat(batch)
            self.logger.debug("Flushing batch of %d tasks", len(batch))

            prof = self._prof
            if prof:
                for t in batch:
                    prof.prof('task_batch_flush', uid=t.get('uid', '?'))

            await asyncio.to_thread(self._rh.submit_tasks, batch)
        except Exception as exc:
            self.logger.error("Failed to submit batch of %d tasks: %s", len(batch), exc)
            for t in batch:
                uid = t.get('uid')
                if uid in self._tasks:
                    task = self._tasks[uid]
                    task['exception'] = exc
                    task['state'] = 'FAILED'
                    self._callback_func(task, 'FAILED')
            raise

[gemini-code-assist]

In _v3_function_wrapper, if opening the second file (stderr_path) fails, the first file (stdout_path) remains open and is leaked. Using a with statement to manage both files ensures they are always closed properly and prevents resource leaks.

Suggested change

	out = open(f"{stdout_path}.{rank}", "w")
	err = open(f"{stderr_path}.{rank}", "w")
	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
	out.close()
	err.close()
	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

[gemini-code-assist]

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 = ""

[gemini-code-assist]

In shutdown, if await self._locked_flush() raises an exception, the rest of the shutdown sequence is skipped, leaking the _rh and _bc clients. Wrapping it in a try...except block ensures that cleanup always runs.

    async def shutdown(self) -> None:
        """Flush pending tasks, close session and BridgeClient."""
        try:
            await self._locked_flush()
        except Exception as e:
            self.logger.warning("Failed to flush pending tasks during shutdown: %s", e)

        self._backend_state = BackendMainStates.SHUTDOWN

        if self._rh:
            try:
                self._rh.close()
            except Exception as e:
                self.logger.warning("Failed to close session: %s", e)
            self._rh = None

        if self._bc:
            try:
                self._bc.close()
            except Exception as e:
                self.logger.warning("Failed to close bridge client: %s", e)
            self._bc = None

        if self._prof:
            self._prof.close()

        self.logger.info("Edge execution backend shutdown complete")