fix(async-grpo): fail fast on generation worker errors

[jthomson04]

Summary

Async GRPO can otherwise leave the trainer waiting for replay-buffer data when rollout batch workers fail in the background. This change makes those failures visible while preserving recovery from transient failures:

• Count consecutive failed _run_rollout_batch_worker invocations and reset the streak after a successful batch.
• Preserve failure timing, traceback logging, target release, and worker-thread cleanup.
• Store the first fatal diagnostic after failures exceed max_generation_failures, including the backend, generation weight, target weight, count, threshold, original exception, and worker traceback.
• Raise a fresh RuntimeError from check_health() and check health during initial buffer fill and once per training-loop iteration.
• Give AsyncGRPOConfig.max_generation_failures the documented Pydantic default of 0, while keeping explicit YAML values in runnable async-GRPO configs and the reference snapshot.
• Wake paused collection after a tolerated worker failure so released targets can be gap-filled, and re-arm the pause event on every limit check.
• Exclude errors raised only because normal dataloader exhaustion stopped the collector, while preserving reservation release and worker-thread cleanup.
• Set the four Nemotron-3-Super stage recipes, two Qwen SWE recipes, and seven Nemotron-3-Ultra recipes to tolerate 3 consecutive failures, aborting on the fourth.
• Document the setting in both copyable async-GRPO guide examples.

Superseded original fix

This PR originally targeted incomplete per-target prompt-group accounting. That original bugfix is no longer part of this PR: a separate batched reservation and gap-filling fix landed on main and superseded it.

The rebase therefore drops obsolete commit 96ee36e72, preserves the target-accounting behavior now on main, and contains no replay-buffer or target-accounting delta.

Failure semantics

max_generation_failures: 0 is fail-fast. Positive values tolerate that many consecutive failed batch-worker invocations. A successful batch resets the streak. Once the threshold is exceeded, the first fatal diagnostic remains sticky and subsequent check_health() calls raise fresh exceptions carrying that diagnostic.

Validation

• Parameterized worker-health coverage for healthy checks, immediate failure, positive tolerance, success reset, threshold crossing, sticky first-fatal diagnostics, fresh exceptions, traceback context, and reservation release.
• Regression coverage for tolerated-failure wake-up at max_trajectory_age_steps: 1 and for normal shutdown errors not affecting worker health.
• Async trainer coverage for unconditional per-step health propagation and cleanup.
• Pydantic default-model and router-replay compatibility after the config migration on main.
• Normal async trainer max-step path.
• Config-reference validation, including the current GRPO contract for all seven Nemotron-3-Ultra recipes.
• Focused and broader unit validation: 93 tests passed.
• Ruff check and format validation on affected Python files.
• git diff --check.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[jthomson04]

/ok to test 24e0667

[yuki-97]

@jthomson04 thanks for fixing this, overall lgtm. left some comments.

[jthomson04]

/ok to test 105aa4d

[jthomson04]

/ok to test 333c87c

[yuki-97]

thanks @jthomson04 , lgtm!
@terrykong @mehraakash could you help to take a review as well?

[terrykong]

@jthomson04 IIUC, this now allows backfill for failing prompts instead of stalling/hanging. Please correct me if I'm wrong. If my understanding is correct, I feel like we should actually be erroring out (surface the error) b/c this kind of behavior seems like it can allow a lot of issues to be swept under the rug and ignored

@ananthsub to opine on gym surfacing errors and how this could happen

[jthomson04]

Pushed an update to set a threshold for errors we tolerate. It loudly logs errors, and throws an exception after a certain amount of errors have been reached. By default, we retain the prior behavior of considering any generation error as fatal.

[yuki-97]

thanks @jthomson04 , lgtm. could you help to rebase main and fix the conflict?

@terrykong @mehraakash could you help to take a review as well?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jthomson04]

/ok to test 3fc9465

[jthomson04]

/ok to test c68a039

[jthomson04]

/ok to test 52f5a72

[yuki-97]

thanks @jthomson04 , LGTM

[jthomson04]

/ok to test a168378

[jthomson04]

/ok to test 777ba10

[yuki-97]

/ok to test 30c98ad

[yuki-97]

/ok to test 5244f6e

[yuki-97]

@terrykong @yfw could you help to take a review as well?

[jthomson04]

/ok to test 109036f

[jthomson04]

/ok to test b5c7dee

[terrykong]

Reviewed with a team of five agents plus an adversarial pass; the main finding has a runnable repro attached to the first comment (verbatim-extracted methods, no GPU/ray/torch, deterministic 3/3).

The fast-fail is correctly wired for max_generation_failures: 0, and the bare raise at grpo.py:5104 is load-bearing rather than incidental — without it check_health()'s RuntimeError would be caught and swallowed by that same block and the feature would be a no-op. It also fixes a real pre-existing bug where a crashed async run exited 0. Docs, exemplar and the reference config are updated in lockstep.

The one thing I'd hold on: at max_trajectory_age_steps: 1 — which all 13 recipes this PR sets to max_generation_failures: 3 resolve to — once a batch worker exhausts its in-worker retries, the #2651 gap-fill that would top the short target back up never executes. The collection loop parks while the worker is in flight and a released slot can't wake it, so trajectories_needed is never recomputed and the step never completes. 0 aborts cleanly on such a failure; 3 hangs silently. Two-line fix, verified by the repro.

Two smaller independent items: a config default that contradicts its own documentation, and a miscount that turns normal dataloader exhaustion into a fatal generation failure.

Generated by Claude Code

[terrykong]

1 action item, 1 follow-up. PR-introduced.

AI-1 — after a worker gives up, the gap-fill that would top the target back up never runs, so the step never completes

Simplified to one thread, this is the collector:

for batch in dataloader:
    if no_target_slot_free():                 # slots = next max_trajectory_age_steps steps
        wait_for_event()                      # <-- PARKS. only set_weight_version() wakes it
    target = reserve_slot()
    needed = 16 - already_buffered(target)    # <-- the gap-fill, every iteration
    spawn_thread(worker, batch[:needed], target)   # returns in ms; worker runs for minutes

def worker(groups, target):
    for attempt in range(4 if gym else 1):    # retries the SAME prompts
        ... buffer each group as it completes ...
        if not missing:
            fail_count = 0; release_slot(target); return
    fail_count += 1                           # ONE increment for the whole batch
    if fail_count > max_generation_failures:
        fatal = "abort"                       # -> trainer's check_health() raises
    release_slot(target)                      # <-- frees the slot, never wakes the loop

Normally the gap-fill handles this. Say 15 of 16 groups buffered and the 16th failed all 4 attempts (max_attempts = 1 + _MAX_NEMO_GYM_STREAM_RETRIES, gym only — native gets 1). Groups are buffered individually as they stream (_enqueue_rollout_group), so the target is simply short by one. On the next loop iteration:

needed = get_trajectories_needed(target=7)   -> 1
groups = next_batch[:1]                      -> fresh prompts, not the failed one
"🎯 Reserved target weight 7 for gap-filling (need 1/16 more trajectories)"

That's the #2651 machinery and it is correct.

It never gets to run. _process_batch returns in ms; the loop's next iteration finds the slot still reserved — at max_trajectory_age_steps: 1 (L61) it is the only slot (_calculate_target_weights) — so _should_pause_for_generation_limits → True → .clear() → parks on .wait(). Minutes later the worker gives up and _release_target frees the slot but does not set _generation_limit_cleared — set only at :114 and in set_weight_version, whose callers are grpo.py:4110 (once, pre-fill) and grpo.py:4687 (post-refit). So needed is never recomputed, the target stays at 15/16, has_complete_batch stays False, and the trainer blocks. The escape at grpo.py:4272-4285 can't fire either: it needs not running, and a parked loop never reaches its finally.

Measured (repro below; verbatim methods, no GPU, deterministic 3/3):

A) generation dead                            batches  count  need  outcome
   current code, mgf=3  <- the 13 recipes           2      2     4  HUNG
   current code, mgf=0  [control]                   2      2     1  ABORTED
   fix A only: .set()                             200      3     4  HUNG
   fix A+B: .set() + .clear() un-gated              6      4     4  ABORTED

B) ONE worker exhausts its retries          steps done
   baseline: no failure at all                    8 / 8   ok
   one exhausted worker                           2       STRANDED
   one exhausted worker + fix A+B                 8 / 8   ok

Net for this recipe: 0 latches fatal before the park matters and aborts cleanly; 3 latches nothing and hangs. The tolerance setting is strictly worse than the default here.

Action — two changes; the first alone is not enough:

A. In _run_rollout_batch_worker's except, after the increment: self._generation_limit_cleared.set().

B. Move .clear() at :324 out of the if self._last_limit_warning_version != self.current_weight_version: guard at :309. That guard logs once per weight version, but .clear() is nested in it while .wait() is not — so after the first warning the event can never be re-armed and the pause degenerates into a spin (row "fix A only": 200 dataloader batches burned, zero GPU work).

Please don't put the set() in _release_target — it also runs on the success path, where parking is correct.

A startup validator on max_generation_failures < max_trajectory_age_steps + 1 isn't a substitute: at max_age: 1 it would reject every non-zero value, deleting the feature for the recipes that want it.

Caveats. (a) The park needs the worker to fail slower than one loop iteration (~tens of ms; several blocking Ray round-trips) — anything with a network or subprocess round trip qualifies. (b) Event check-then-clear still has a lost-wakeup window; it never fired here, but a Condition with a predicate loop would be airtight if you prefer.

Follow-up

A regression test for a worker exhausting its retries at max_trajectory_age_steps: 1. No current test could catch this — they all drive _run_rollout_batch_worker directly and never exercise the loop's pause path.

repro (save as repro.py, run from repo root, ~60 s)

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///
"""PR #2368 repro. No GPU/ray/torch. Run from the repo root:  ./repro.py

The six load-bearing methods are extracted VERBATIM from trajectory_collector.py via
ast and exec'd, so the control flow under test is the real code. Modeled only:
  ray.get/.remote -> pass-through | replay buffer -> per-target group counts
  _process_batch  -> reserve + spawn                  (mirrors :390-462)
  worker epilogue -> increment/compare/latch/release  (mirrors :854-895)

Scenario A: generation is dead      -> can the counter ever reach the threshold?
Scenario B: ONE worker exhausts its retries, all others succeed
                                    -> does the gap-fill top the short target back up?
"""
import ast, sys, threading, time, io, contextlib
from contextlib import contextmanager
from pathlib import Path

SRC = Path(sys.argv[1] if len(sys.argv) > 1
           else "nemo_rl/algorithms/async_utils/trajectory_collector.py")
WANT = ["_calculate_target_weights", "_get_next_target_for_generation", "set_weight_version",
        "_should_pause_for_generation_limits", "_collection_loop", "_release_target"]


def extract(path):
    txt = path.read_text(); lines = txt.splitlines()
    cls = next(n for n in ast.walk(ast.parse(txt))
               if isinstance(n, ast.ClassDef) and n.name == "AsyncTrajectoryCollector")
    out = {}
    for n in cls.body:
        if isinstance(n, ast.FunctionDef) and n.name in WANT:
            b = "\n".join(lines[n.lineno - 1:n.end_lineno])
            out[n.name] = "\n".join(l[4:] if l.startswith("    ") else l for l in b.splitlines())
    assert not set(WANT) - set(out), sorted(set(WANT) - set(out))
    return out


class _Ref:
    def __init__(s, f): s._f = f
    def remote(s, *a, **k): return s._f(*a, **k)


class _Ray:
    @staticmethod
    def get(v): return v


class Buf:
    def __init__(s, n): s.n, s.counts, s.last = n, {}, -1
    @property
    def get_last_target_weight_already_generated(s): return _Ref(lambda: s.last)
    @property
    def get_trajectories_needed(s):
        return _Ref(lambda t, num, m=None: max(0, s.n - s.counts.get(t, 0)))
    def complete(s, t): return s.counts.get(t, 0) >= s.n
    def fill(s, t): s.counts[t] = s.n


@contextmanager
def _null(*a, **k): yield


class _T:
    def time(s, *a, **k): return _null()
    def record(s, *a, **k): pass


class Cfg:
    class grpo:
        num_prompts_per_step = 4
        class async_grpo: max_trajectory_age_steps = 1


class C:
    def __init__(s, limit, mode, fix, fail_target=2, delay=0.15):
        s.master_config, s.replay_buffer = Cfg, Buf(Cfg.grpo.num_prompts_per_step)
        s.running, s.dataloader = True, None
        s.data_exhausted = s.collection_failed = False
        s.initial_weight_version = s.current_weight_version = 0
        s._last_limit_warning_version = -1
        s._generating_targets, s._inflight_threads = set(), set()
        s._generation_check_lock = threading.Lock(); s._threads_lock = threading.Lock()
        s._generation_limit_cleared = threading.Event(); s._generation_limit_cleared.set()
        s._manual_pause_cleared = threading.Event(); s._manual_pause_cleared.set()
        s._refit_pause_cleared = threading.Event(); s._refit_pause_cleared.set()
        s._efficiency_timer = _T()
        s._failure_lock = threading.Lock()
        s._failure_count, s._fatal_error_message = 0, None
        s._max_generation_failures = limit
        s.mode, s.fix, s.fail_target, s.delay = mode, fix, fail_target, delay
        s.spawned = s.batches = 0
        s.failed_once = False

    def _process_batch(s, batch):                       # mirrors :390-462
        s.batches += 1
        t = s._get_next_target_for_generation(s.current_weight_version)
        if t is None: return
        s.spawned += 1
        threading.Thread(target=s._worker, args=(t,), daemon=True).start()

    def _worker(s, target):                             # mirrors :854-895
        time.sleep(s.delay)
        fails = s.mode == "all_fail" or (
            s.mode == "one" and target == s.fail_target and not s.failed_once)
        if fails:
            if s.mode != "all_fail": s.failed_once = True
            with s._failure_lock:
                s._failure_count += 1
                if (s._failure_count > s._max_generation_failures
                        and not s._fatal_error_message):
                    s._fatal_error_message = (
                        f"aborting: {s._failure_count} failure(s) exceeded "
                        f"max_generation_failures={s._max_generation_failures}")
        else:
            with s._failure_lock: s._failure_count = 0   # :855 success resets
            s.replay_buffer.fill(target)
        s._release_target(target)                        # verbatim :647
        if s.fix in ("set-only", "set+ungate"):
            s._generation_limit_cleared.set()

    def check_health(s):
        with s._failure_lock: m = s._fatal_error_message
        if m: raise RuntimeError(m)


def install(methods, ungate):
    import typing
    ns = {"ray": _Ray, "_threading": threading, "Optional": typing.Optional,
          "Any": typing.Any, "StatefulDataLoader": object, "BatchedDataDict": dict,
          "DatumSpec": dict}
    for name, src in methods.items():
        if name == "_collection_loop" and ungate:
            # move .clear() OUT of the `log once per weight version` guard
            src = src.replace("                    self._generation_limit_cleared.clear()", "", 1)
            src = src.replace("                # Efficiently wait for generation limits",
                              "                self._generation_limit_cleared.clear()\n"
                              "                # Efficiently wait for generation limits", 1)
        exec(compile("from __future__ import annotations\n" + src, f"<{name}>", "exec"), ns)
        setattr(C, name, ns[name])


def run(limit, mode, fix, steps=8, settle=2.5, train_time=0.35, n_batches=200):
    install(METHODS, ungate=(fix == "set+ungate"))
    c = C(limit, mode, fix); c.dataloader = iter(range(n_batches))
    with contextlib.redirect_stdout(io.StringIO()):
        threading.Thread(target=c._collection_loop, daemon=True).start()
        fatal, done = False, 0
        for _ in range(steps):
            end, advanced = time.time() + settle, False
            while time.time() < end:
                try: c.check_health()
                except RuntimeError: fatal = True; break
                if c.replay_buffer.complete(c.current_weight_version):
                    time.sleep(train_time)
                    c.current_weight_version += 1
                    c.set_weight_version(c.current_weight_version)
                    done += 1; advanced = True; break
                time.sleep(0.02)
            if fatal or not advanced: break
    return dict(count=c._failure_count, need=limit + 1, batches=c.batches,
                steps=done, fatal=fatal)


if __name__ == "__main__":
    METHODS = extract(SRC)
    print(f"\nverbatim from {SRC.name}: {', '.join(sorted(METHODS))}")

    print("\nA) generation dead  (max_age=1) — can the counter reach the threshold?")
    print(f"   {'variant':46} {'batches':>8} {'count':>6} {'need':>5}  outcome")
    for lbl, lim, fix in (("current code, mgf=3  <- the 13 recipes", 3, "none"),
                          ("current code, mgf=0  [control]        ", 0, "none"),
                          ("fix A: .set() only                    ", 3, "set-only"),
                          ("fix A+B: .set() + .clear() un-gated    ", 3, "set+ungate")):
        r = run(lim, "all_fail", fix)
        print(f"   {lbl:46} {r['batches']:>8} {r['count']:>6} {r['need']:>5}  "
              f"{'ABORTED' if r['fatal'] else 'HUNG'}")

    print("\nB) ONE worker exhausts its retries, all others succeed (max_age=1, mgf=3)")
    print(f"   {'variant':46} {'steps done':>11}  outcome")
    res = {}
    for lbl, mode, fix in (("baseline: no failure at all           ", "none", "none"),
                           ("one exhausted worker                  ", "one", "none"),
                           ("one exhausted worker + fix A+B        ", "one", "set+ungate")):
        r = run(3, mode, fix); res[lbl.strip()] = r
        print(f"   {lbl:46} {r['steps']:>11}  "
              f"{'ok' if r['steps'] >= 8 else 'STRANDED at step %d' % r['steps']}")

    ok = (res["one exhausted worker"]["steps"] < 8
          and res["one exhausted worker + fix A+B"]["steps"] >= 8)
    print("\n" + ("REPRODUCED" if ok else "did not reproduce"))
    sys.exit(0 if ok else 1)

[jthomson04]

Addressed in 87f55d4. A failed batch worker now sets _generation_limit_cleared after releasing its failure state so a paused collector can re-evaluate and gap-fill the released target, and the collection loop clears the event on every limit pause rather than only in the log-once branch. I added a regression using the real collection loop and worker path with max_trajectory_age_steps=1; it verifies the tolerated failure wakes the loop, releases the reservation, and starts gap filling without latching a fatal error.

[terrykong]

fixed in 87f55d4

[terrykong]

1 action item. PR-introduced, non-blocking and independent of the other finding.

AI-1 — the documented default contradicts the schema default

examples/configs/grpo_math_1B.yaml:57 added by this PR says # 0 (default) = fail on the first worker exception, but the schema default here is None, which raises. There is no configuration in which the documented default is the actual default.

Underneath that, T | None = None is how this repo expresses a genuinely optional field, while this one is mandatory in practice, enforced by hand-written checks in two consumers (grpo.py:3920 and the Ray actor constructor at trajectory_collector.py:134). .claude/skills/config-conventions/SKILL.md asks that a default "live in exactly one place" and, for BaseModel, "on the BaseModel field as a Python value". Every other field on AsyncGRPOConfig carries a concrete default.

It also introduces an unannounced compat break: an out-of-repo config with async_grpo.enabled: true that doesn't inherit from an updated exemplar now dies at startup with KeyError: 'grpo.async_grpo.max_generation_failures is required for async GRPO', under a fix: commit with no migration note.

Action: give it a real default and delete both raise KeyError blocks. Zero behavior change for every in-repo config — all already resolve to 0 or 3.

Suggested change

	max_generation_failures: int | None = None
	# Generation-worker failures tolerated before the AsyncTrajectoryCollector
	# aborts the run. A successful batch worker resets the count.
	# 0 = the first worker exception is fatal.
	max_generation_failures: int = 0

If you'd rather keep required-when-enabled semantics, the matching pattern is 28 lines below — RewardPenaltyConfig._require_unwanted_token_ids_when_penalized uses a pydantic @model_validator(mode="after") raising ValueError. Production loads go through full validation so a validator fires there; tests use model_construct, which skips validators, so nothing breaks. Minor either way: KeyError is the wrong type for a present-but-None value, and Python quote-wraps KeyError messages so the traceback reads like a key name.

[jthomson04]

Addressed in 87f55d4. AsyncGRPOConfig.max_generation_failures is now int = 0, and the manual None/KeyError guards were removed from both the trainer and collector. The nested-default test now asserts the Pydantic default is 0; the config-reference and current GRPO contract validation also pass.

[terrykong]

fixed in 87f55d4

[terrykong]

1 action item. PR-introduced, independent of the other finding. Confirmed with a repro (below).

AI-1 — normal dataloader exhaustion is counted as a generation failure

This isn't a rare race — it's the default end-of-data path. The loop spawns a worker for the last batch, immediately exhausts the dataloader, and _collection_loop's finally sets self.running = False while that worker is still generating. The worker then hits raise RuntimeError("Trajectory collection stopped before enqueue completed") — note it drops the group without even attempting the add — which surfaces as batch_error from _collect_rollout_batch and lands in this counting block as an ordinary generation failure.

With max_generation_failures: 0 — what every non-agentic config resolves to — that latches a fatal on the first such worker.

Repro output (_collection_loop and _enqueue_rollout_group extracted verbatim; no GPU/ray/torch):

control  (running=True)          : buffered OK
after exhaustion                 : running=False, data_exhausted=True
in-flight worker enqueue         : RuntimeError('Trajectory collection stopped before enqueue completed')
_failure_count                   : 1 (max_generation_failures=0)

What the user sees at the next check_health():
   AsyncTrajectoryCollector aborting: 1 batch-worker failure(s) exceeded
   max_generation_failures=0. Last failure in native batch worker ...
   RuntimeError('Trajectory collection stopped before enqueue completed')

What they should have seen (grpo.py:4272-4285):
   Trajectory collector stopped: dataloader exhausted ...
   Increase data.train.max_num_epochs or use a larger dataset.

Cost: the run was ending anyway, so this isn't a corrupted-training issue — but the user gets an error blaming generation for what is actually "your dataset ran out", and the job dies at the next check_health() instead of training the steps still sitting in the buffer.

Action: don't count shutdown-induced enqueue aborts. Either guard the counting block with if self.running:, or raise a distinct sentinel from _enqueue_rollout_group and re-raise it past the counter.

repro (save as repro_exhaust.py, run from repo root)

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///
"""PR #2368: does NORMAL dataloader exhaustion get counted as a generation failure?

Run from the repo root:  ./repro_exhaust.py

_collection_loop and _enqueue_rollout_group are extracted VERBATIM from
trajectory_collector.py via ast and exec'd, so the running=False -> raise chain is
the real code. Modeled: ray.get/.remote pass-through, replay buffer, and the
worker's except block (mirrors :857-879).

Sequence under test:
  1. the dataloader runs out -> _collection_loop's `finally` sets self.running = False   (:349)
  2. a still-in-flight worker calls _enqueue_rollout_group -> `while self.running:` is
     False -> raise RuntimeError("Trajectory collection stopped before enqueue completed")
  3. that lands in the batch-worker except block -> _failure_count += 1                   (:866)
  4. with max_generation_failures=0 that latches a fatal, and check_health() raises it
     instead of the intended "dataloader exhausted / increase max_num_epochs" message
"""
import ast, sys, asyncio, threading, time, io, contextlib
from contextlib import contextmanager
from pathlib import Path

SRC = Path(sys.argv[1] if len(sys.argv) > 1
           else "nemo_rl/algorithms/async_utils/trajectory_collector.py")
WANT = ["_collection_loop", "_enqueue_rollout_group", "_should_pause_for_generation_limits",
        "_calculate_target_weights", "_get_next_target_for_generation", "_release_target"]


def extract(path):
    txt = path.read_text(); lines = txt.splitlines()
    cls = next(n for n in ast.walk(ast.parse(txt))
               if isinstance(n, ast.ClassDef) and n.name == "AsyncTrajectoryCollector")
    out = {}
    for n in cls.body:
        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n.name in WANT:
            b = "\n".join(lines[n.lineno - 1:n.end_lineno])
            out[n.name] = "\n".join(l[4:] if l.startswith("    ") else l for l in b.splitlines())
    assert not set(WANT) - set(out), sorted(set(WANT) - set(out))
    return out


class _Ref:
    def __init__(s, f): s._f = f
    def remote(s, *a, **k): return s._f(*a, **k)


class _AsyncRef:
    """replay_buffer.add.remote(...) is awaited in _enqueue_rollout_group."""
    def __init__(s, f): s._f = f
    def remote(s, *a, **k):
        async def go(): return s._f(*a, **k)
        return go()


class _Ray:
    @staticmethod
    def get(v): return v


class Buf:
    def __init__(s, n): s.n, s.counts, s.last = n, {}, -1
    @property
    def get_last_target_weight_already_generated(s): return _Ref(lambda: s.last)
    @property
    def get_trajectories_needed(s):
        return _Ref(lambda t, num, m=None: max(0, s.n - s.counts.get(t, 0)))
    @property
    def add(s):
        def _add(group, gwv, twv):
            s.counts[twv] = s.counts.get(twv, 0) + 1
            return "success"
        return _AsyncRef(_add)


class _FakeBatch(dict):
    def to(s, dev): return s


class _Result:
    def __init__(s):
        s.final_batch = _FakeBatch()
        s.rollout_metrics = {}
        s.group_index = 0
        s.task_index = None


@contextmanager
def _null(*a, **k): yield


class _T:
    def time(s, *a, **k): return _null()
    def record(s, *a, **k): pass


class Cfg:
    class grpo:
        num_prompts_per_step = 4
        class async_grpo: max_trajectory_age_steps = 1


class C:
    def __init__(s, limit):
        s.master_config, s.replay_buffer = Cfg, Buf(4)
        s.running, s.dataloader = True, None
        s.data_exhausted = s.collection_failed = False
        s.initial_weight_version = s.current_weight_version = 0
        s._last_limit_warning_version = -1
        s._generating_targets, s._inflight_threads = set(), set()
        s._generation_check_lock = threading.Lock(); s._threads_lock = threading.Lock()
        s._generation_limit_cleared = threading.Event(); s._generation_limit_cleared.set()
        s._manual_pause_cleared = threading.Event(); s._manual_pause_cleared.set()
        s._refit_pause_cleared = threading.Event(); s._refit_pause_cleared.set()
        s._efficiency_timer = _T()
        s._has_distillation_teachers = False
        s.tokenizer = None
        s._failure_lock = threading.Lock()
        s._failure_count, s._fatal_error_message = 0, None
        s._max_generation_failures = limit

    def _process_batch(s, batch):
        pass  # not under test here

    # mirrors _run_rollout_batch_worker's except block, :857-879
    def count_worker_failure(s, error, backend="native", gwv=0, twv=0):
        with s._failure_lock:
            s._failure_count += 1
            n = s._failure_count
            fatal = n > s._max_generation_failures
            if fatal and s._fatal_error_message is None:
                s._fatal_error_message = (
                    f"AsyncTrajectoryCollector aborting: {n} batch-worker failure(s) "
                    f"exceeded max_generation_failures={s._max_generation_failures}. "
                    f"Last failure in {backend} batch worker for generation_weight={gwv}, "
                    f"target_weight={twv}: {error!r}")

    def check_health(s):
        with s._failure_lock: m = s._fatal_error_message
        if m: raise RuntimeError(m)


def install(methods):
    import typing
    ns = {"ray": _Ray, "_threading": threading, "asyncio": asyncio, "time": time,
          "Optional": typing.Optional, "Any": typing.Any, "StatefulDataLoader": object,
          "BatchedDataDict": dict, "DatumSpec": dict,
          "NEMO_GYM_TASK_INDEX_KEY": "_ng_task_index",
          "_REPLAY_BUFFER_MAX_BACKOFF_SECONDS": 0.5}
    for name, src in methods.items():
        exec(compile("from __future__ import annotations\n" + src, f"<{name}>", "exec"), ns)
        setattr(C, name, ns[name])


def enqueue(c):
    """Run the verbatim _enqueue_rollout_group exactly as a worker would."""
    return asyncio.run(c._enqueue_rollout_group(
        rollout_result=_Result(), generation_weight_version=0, target_weight_version=1,
        expected_prompt_groups=4, buffered_group_indices=set(),
        collection_started_at=time.perf_counter()))


if __name__ == "__main__":
    install(extract(SRC))
    print(f"\nverbatim: _collection_loop, _enqueue_rollout_group (+4 helpers) "
          f"from {SRC.name}\n")

    # --- control: collector still running -> the group buffers fine ---
    c = C(limit=0); c.dataloader = iter([])
    with contextlib.redirect_stdout(io.StringIO()):
        try:
            enqueue(c); control = "buffered OK"
        except Exception as e:
            control = f"raised {e!r}"
    print(f"control  (running=True)          : {control}")

    # --- the real sequence: let the dataloader exhaust, THEN enqueue ---
    c = C(limit=0); c.dataloader = iter(range(2))
    with contextlib.redirect_stdout(io.StringIO()):
        c._collection_loop()                    # VERBATIM: exhausts, finally sets running=False
    print(f"after exhaustion                 : running={c.running}, "
          f"data_exhausted={c.data_exhausted}")

    with contextlib.redirect_stdout(io.StringIO()):
        try:
            enqueue(c); outcome = None
        except Exception as e:
            outcome = e
    print(f"in-flight worker enqueue         : {outcome!r}")

    if outcome is None:
        print("\nNOT reproduced: enqueue did not raise"); sys.exit(1)

    c.count_worker_failure(outcome)             # mirrors :857-879
    print(f"_failure_count                   : {c._failure_count} "
          f"(max_generation_failures={c._max_generation_failures})")

    try:
        c.check_health(); print("\nNOT reproduced: check_health stayed silent"); sys.exit(1)
    except RuntimeError as e:
        print("\nWhat the user sees at the next check_health():\n")
        print("   " + str(e)[:300])
        print("\nWhat they should have seen (grpo.py:4272-4285):\n")
        print("   RuntimeError: Trajectory collector stopped: dataloader exhausted ...")
        print("   The dataset ran out of data before training could start.")
        print("   Increase data.train.max_num_epochs or use a larger dataset.")
        print("\nREPRODUCED: normal exhaustion is reported as a generation failure")
        sys.exit(0)

[jthomson04]

Addressed in 87f55d4. _run_rollout_batch_worker now skips failure accounting when the collector has already stopped, while the existing finally block still releases the target and removes the worker thread. The new shutdown regression simulates the enqueue-side error after normal exhaustion and verifies the failure count/fatal diagnostic remain untouched and cleanup completes.

[terrykong]

fixed in 87f55d4

[jthomson04]

/ok to test 87f55d4

[jthomson04]

Addressed in 87f55d4. The collector now tracks consecutive failures, preserves the first formatted diagnostic, raises fresh health exceptions, wakes gap filling after tolerated failures, and excludes shutdown-induced enqueue errors from failure accounting. AsyncGRPOConfig.max_generation_failures now defaults to 0, long-running agentic recipes explicitly tolerate 3, the trainer checks health every iteration, and the copyable docs/config references are updated. Validation passed: 93 broader unit tests, 4 focused post-format worker regressions, Ruff lint/format, and git diff --check.

[jthomson04]

0 is an aggressive default for these production agentic recipes specifically.

Today a failed batch worker is not fatal by construction: _release_target frees the reservation and the next _process_batch re-picks the same target and gap-fills it (trajectory_collector.py:403-421). A one-off failure costs one wasted batch, not the run.

With this change that same one-off arms _fatal_error, and the very next pass through the normal generation-bound wait in grpo.py:4188 ("⏳ Buffer empty or not enough groups to form a full step, waiting..." → check_health() → sleep(0.5) → continue) kills the job. That branch is routine in async GRPO, not exceptional — so the effective semantics are "any single rollout failure ends the run."

The four nemotron-3-super/stage* configs and the SWE recipes are exactly the multi-day agentic runs with in_flight_weight_updates: true where an occasional env/tool/backend exception is expected. Worth setting a nonzero value here even if the exemplar default stays 0.

[jthomson04]

The counter is process-lifetime, so max_generation_failures: 5 means "5 failures over the entire run", not "5 in a row". On a 10k-step run any nonzero budget is exhausted early and the setting collapses back to fail-fast — which makes it hard to use for its stated purpose ("transient generation errors ... acceptable to drop").

Suggest resetting on success so the threshold counts consecutive failures:

try:
    await self._collect_rollout_batch(...)
    with self._failure_lock:
        self._failure_count = 0
except Exception as error:
    ...

and updating the AsyncGRPOConfig comment accordingly. That makes the knob discriminate transient flakes from a genuinely broken backend, which is the distinction this PR is drawing.

[jthomson04]

{error!r} gives the driver the exception type and args but no stack — the real traceback only exists in the collector actor's stdout. Since check_health() is what terminates the job, the driver-side failure should be self-contained: fold failure_traceback into the message, or set err.__cause__ = error when storing it.

Related: check_health does raise err on the same stored instance every call, which appends frames to err.__traceback__ each time. Storing the message and constructing a fresh RuntimeError per call avoids the growth and is no more code.

[jthomson04]

Making this required is the right convention for a legacy TypedDict config, but two follow-ups are missing:

1. docs/guides/async-grpo.md:39-45 and :49-72 both show async_grpo blocks users are meant to copy, and neither includes max_generation_failures. Someone following the guide gets a bare KeyError: 'max_generation_failures' raised inside the AsyncTrajectoryCollector actor constructor — after the virtual cluster, policy, and generation workers are already up. Same for any hand-rolled config that doesn't inherit from an exemplar.

2. Reading the key driver-side in async_grpo_train, next to where max_trajectory_age_steps is read, would turn that late actor-init KeyError into an early readable message.

[jthomson04]

Both check_health() call sites are inside wait loops, so if the buffer happens to stay full a stored fatal error sits idle indefinitely. One call at the top of the while step < master_config.grpo["max_num_steps"] loop costs one RPC per step and makes the guarantee hold unconditionally rather than depending on the starvation branch being hit. Minor given how often that branch does run, but it removes the coincidence.

[jthomson04]

Nicely constructed — the [0, 2] parametrization covers both immediate-fatal and tolerance, and it asserts the diagnostic fields, reservation release, and stickiness across two calls. Three gaps worth closing:

• No assertion that check_health() is a no-op with zero failures — that's the path taken on essentially every real call.
• The if is_fatal and self._fatal_error is None guard is untested: nothing verifies a second post-threshold failure leaves the first error intact while still incrementing the count.
• Nothing exercises the grpo.py wiring. MockCollector.check_health returns None, so async_grpo_train never sees a raise. A case where check_health.remote() raises and the test asserts async_grpo_train propagates it would lock in the behavior this PR exists to provide.

[jthomson04]

Addressed in 87f55d4. A failed batch worker now sets _generation_limit_cleared after releasing its failure state so a paused collector can re-evaluate and gap-fill the released target, and the collection loop clears the event on every limit pause rather than only in the log-once branch. I added a regression using the real collection loop and worker path with max_trajectory_age_steps=1; it verifies the tolerated failure wakes the loop, releases the reservation, and starts gap filling without latching a fatal error.

[jthomson04]

Addressed in 87f55d4. AsyncGRPOConfig.max_generation_failures is now int = 0, and the manual None/KeyError guards were removed from both the trainer and collector. The nested-default test now asserts the Pydantic default is 0; the config-reference and current GRPO contract validation also pass.

[jthomson04]

Addressed in 87f55d4. _run_rollout_batch_worker now skips failure accounting when the collector has already stopped, while the existing finally block still releases the target and removes the worker thread. The new shutdown regression simulates the enqueue-side error after normal exhaustion and verifies the failure count/fatal diagnostic remain untouched and cleanup completes.

[jthomson04]

/ok to test 07be815