fix(fully_async): per-task timeout + exception-safe requeue in rollout worker

[Shi-Dong]

Summary

Fixes two compounding bugs in AsyncRolloutWorker.continuous_worker_loop
(examples/fully_async/fully_async_rollout.py) that together caused rollout-
level soft-deadlock once a single trial went pathological. Observed live on
GLM-4.7-Flash async agentic RL Run 4 rollout 3, where 11 long-tail SWE-bench
trials (agent step counts 60-148 vs. normal 20-40) held 11/32 concurrency
slots for ~7.5 hours and would have done so indefinitely. Full forensic
write-up lives in the memory-bank entity "Run 4 rollout-3 stall — full
forensic diagnosis".

Before

1. No per-task timeout. asyncio.create_task(generate_and_rm_group(...))
   was never wrapped in asyncio.wait_for. A group-task awaiting a stuck
   HTTP call held its slot forever. Line 73 while len(active_tasks) < max_concurrent_tasks could never dispatch new groups once enough slots
   were pathological.
2. Silent group-loss. generate_and_rm_group uses asyncio.gather
   without return_exceptions=True, so one sub-task raise propagates. The
   worker's lines 67-69 task.result() / except Exception: print swallowed
   it and did not call data_buffer.add_samples — the group vanished and
   the outer collector's while len(data) < 32 never recovered.

After

Both paths funnel through a new module-level helper _run_group_with_timeout
that wraps the per-group coroutine with asyncio.wait_for and a full
exception guard. Its contract: always returns the group, never raises.
On timeout or any exception, every sample in the group is marked
Sample.Status.ABORTED and returned normally, so the existing aborted-
requeue path in generate_rollout_async (re-adds to data_buffer via
add_samples) handles recovery. CancelledError is still propagated (after
marking) so worker shutdown works correctly.

On why return_exceptions=True in generate_and_rm_group was not chosen

Two reasons:

• generate_and_rm_group is on a broader code path than fully-async and
  changing its gather semantics risks regressions in consumers that rely on
  exceptions escaping (e.g. abort() at sglang_rollout.py:334 orchestrates
  its own cancellation).
• The wrapper approach also fixes (1) — the timeout — with the same ~50 new
  lines. A single choke point is cleaner than changes in two files.

CLI

New arg --rollout-group-timeout-s (default 1800.0 seconds). Existing
runs that omit it get the 1800s default, which is strictly better than the
prior infinite behavior. Set to None (or <= 0) to disable the time
bound; exception safety still applies.

Test plan

• Unit tests in tests/fast/rollout/test_fully_async_rollout_timeout.py:
  • Normal completion — samples flow through unchanged.
  • Timeout — samples marked ABORTED.
  • Direct exception raised by coroutine — group returned as ABORTED,
    not silently dropped.
  • Sub-task exception propagating out of asyncio.gather (exact
    production failure mode) — group returned as ABORTED.
  • timeout_s=None disables the bound but keeps exception safety.
  • timeout_s=0.0 disables the bound.
  • CancelledError propagates after samples are marked.
• Smoke-tested all 7 scenarios locally against the edited file (full
  miles dep stack is heavy; the pytest file uses the same logic).
• Reviewer: run pytest tests/fast/rollout/test_fully_async_rollout_timeout.py
  in a full miles env before merge.
• Reviewer: eyeball the CLI-arg default and confirm 1800s matches your
  expectation for your longest-running rollout shape.

Notes

• No behavior change for happy-path runs that never hit the timeout.
• No protocol changes with the agent server or SGLang.
• Do not merge without review; merge is human-only.

[gemini-code-assist]

Code Review

This pull request introduces a per-task timeout and exception-handling wrapper for the fully async rollout worker to prevent hung tasks from occupying concurrency slots and to ensure groups are not silently dropped on failure. Feedback was provided regarding the handling of asyncio.CancelledError, which currently propagates and could cause the main worker loop to terminate unexpectedly because it is not caught by the Exception block. Additionally, a correction was suggested for the command-line argument help text, as the float type constraint prevents users from passing None directly via the CLI.

[gemini-code-assist]

Re-raising asyncio.CancelledError here will cause the task.result() call in the continuous_worker_loop (line 169) and the task_done_callback (line 202) to raise the exception. Since CancelledError inherits from BaseException in Python 3.8+, it is not caught by the except Exception block (line 170), which will lead to the unexpected termination of the worker loop. Furthermore, if the exception propagates, the group is never put into the output_queue, causing the collector to lose those samples and potentially hang while waiting for the target batch size. To maintain the "exception-safe" contract and prevent group loss, it is recommended to catch the cancellation, mark the samples as ABORTED, and return the group normally. When resetting these samples for retry, ensure all metadata fields (e.g., remove_sample, train_metadata) are cleared to their default state to prevent carrying over stale data. Note that you should also update the corresponding test case test_cancellation_propagates_after_marking in tests/fast/rollout/test_fully_async_rollout_timeout.py to reflect this change.

    except asyncio.CancelledError:
        logger.warning(
            "Group task was cancelled; marking %d samples ABORTED for requeue",
            len(group),
        )
        for s in group:
            s.status = Sample.Status.ABORTED
            s.remove_sample = False
            s.train_metadata = {}
        return group

References

1. When implementing a function to reset Sample objects for retry, ensure all generated outputs and metadata fields (e.g., remove_sample, train_metadata) are cleared to their default state to prevent carrying over stale data from previous attempts.

[gemini-code-assist]

The help text suggests setting the value to None to disable the timeout. However, because the argument type is set to float, argparse will raise a ValueError if a user attempts to pass the string "None" via the command line. It is clearer to instruct users to use a non-positive value (e.g., 0) to disable the time bound.

Suggested change

	"concurrency slot forever. Set to None (or <= 0) to disable."
	"concurrency slot forever. Set to <= 0 to disable."