Skip to content

feat: judge cohort hardening - #2385

Closed
tdene wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
tdene:tde/genrm_cohort_timeout
Closed

tdene wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
tdene:tde/genrm_cohort_timeout

Conversation

@tdene

@tdene tdene commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

A judge model will judge an entire GRPO group (labeled as cohort in the Gym code) at once.

This means that it must wait for all related episodes to finish.

Frequently, all related episodes are unable to finish. Commonly there will be at least 1 eternal hang within each cohort.

This PR introduces a timeout argument that causes the judge to stop waiting beyond a certain point.
If, at that point, there is a partial cohort (more than 1 episode) available, the judge will judge the partial cohort.
If not, the cohort will fail.
Regardless, this presents infinite hangs on the cohorts.

@tdene

tdene commented Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test ddfa0e2

@github-actions github-actions Bot added the sla:triage-overdue Review assignment is over the one-business-day SLA label Aug 7, 2026
@cmunley1

cmunley1 commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

/claude review

1 similar comment
@kajalj22

kajalj22 commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

/claude review

@claude

claude Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

SHIP WITH CARE

This fixes a real deadlock: previously a cohort whose peer sub-request died upstream of verify() would leave every waiter awaiting a future that never resolves — an indefinite hang on the training/eval path. The new design is sound. Every waiter runs wait_for(shield(future), timeout); the buffer holds all waiters' futures; whoever pops the buffer under _cohort_lock (either the cohort-completing arrival, or the first waiter to time out) resolves every future synchronously in _resolve_cohort (no await inside it), so late waiters that pop None still find their future already resolved and collect it. I traced the four handled cases — full cohort, timeout-with-partial, timeout-vs-final race (slow_final), and poisoned metadata — and found no double-resolution (the lock-guarded pop hands the non-None buffer to exactly one coroutine) and no remaining hang path. Tests cover all four. aggregate_scores runs outside the lock. Good.

One thing to settle before this drives a real run, raised inline:

  • RISK (app.py:285): On timeout with ≥2 arrivals, the partial cohort is scored for real and returns genuine relative rewards — miscalibrated vs. a full cohort, since rewards are cohort-relative — distinguished only by failure_reason="cohort_timeout". Nothing in this diff consumes that new field. If the downstream trainer/collector doesn't gate on failure_reason, timed-out partial cohorts silently feed skewed reward to RLHF. Confirm the consumer drops/down-weights non-null failure_reason, or return default_score for partial cohorts instead. Still strictly better than the prior deadlock — this is a "make the gating deliberate" flag, not a merge blocker.

  • NOTE (app.py:67): cohort state and the lock are process-global; _cohort_lock now serializes the (short) pop/claim critical section across all cohorts. Fine, just noting.

No async-HTTP, ray.get, or public-API concerns. Default cohort_timeout_s=1800.0 lives on the config field rather than exemplar YAML; consistent with the file's other numeric defaults, so not flagging.


reward = await future
return BaseVerifyResponse(
self._resolve_cohort(stale_buf, *(stale_jit or ([], [])), cause="cohort_timeout")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

RISK — partial-cohort timeout produces real (skewed) relative rewards, gated only by a field nobody may read.

WHAT: On timeout with ≥2 arrivals, _resolve_cohort(..., cause="cohort_timeout") aggregates the partial cohort and returns genuine per-rollout rewards, differing only from a full cohort in that failure_reason="cohort_timeout" is set on the response. aggregate_scores computes rewards relative to the cohort, so a 3/4 cohort yields a different reward distribution than the intended 4/4.

BLAST RADIUS: These rewards feed the training signal. If the downstream consumer (rollout collection / trainer) does not inspect GenRMCompareVerifyResponse.failure_reason, timed-out partial cohorts silently contribute miscalibrated relative rewards to RLHF — the exact silent-scoring-corruption failure mode. failure_reason is a brand-new field; nothing in this diff consumes it.

FIX: Confirm the training/eval consumer explicitly drops or down-weights rows where failure_reason is non-null (at minimum cohort_timeout/aggregation_failed/no_comparisons). If no consumer gates on it, the safer default here is to return default_score for timed-out partial cohorts rather than a partial relative reward. This is still strictly better than the prior infinite deadlock — flagging so the gating is a deliberate decision, not implicit.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thank you Claude, I appreciate the review.

This is intended.

If this is a concern to the maintainers of this repository, I will start a conversation with the training team and discuss this topic further.

@github-actions github-actions Bot removed the sla:triage-overdue Review assignment is over the one-business-day SLA label Aug 7, 2026
@tdene
tdene force-pushed the tde/genrm_cohort_timeout branch from ddfa0e2 to 4c7cc9b Compare August 18, 2026 10:59
@tdene
tdene force-pushed the tde/genrm_cohort_timeout branch from 4c7cc9b to c01dd24 Compare August 21, 2026 15:16
@tdene
tdene force-pushed the tde/genrm_cohort_timeout branch from c01dd24 to bb858c2 Compare August 25, 2026 08:56
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@ananthsub

Copy link
Copy Markdown
Contributor

#3181 now owns incomplete GenRM cohort behavior. This PR can proceed independently on timeout and cleanup, but its final policy must reconcile expected, arrived, missing, retried, and late members and must not turn partial membership into an unmarked ordinary reward. Persistence integration follows #2135/#3179; please link the PR to #3181 and resolve the blocked checks.

pull Bot pushed a commit to Stars1233/Gym that referenced this pull request Sep 23, 2026
)

This follows merged NVIDIA-NeMo#2814, which established group, attempt, and member
identities. It addresses NVIDIA-NeMo#3181's immediate timeout and failure-cleanup
requirements.

- **While collecting answers:** enforce a finite deadline. If a member
never arrives, fail the group and release waiting requests without
publishing partial rewards.
- **Once all answers are available:** retry transient judge HTTP
failures, including interrupted response bodies, within a bounded budget
using the same answers.
- **If judging ultimately fails:** preserve each generated answer and
return an explicitly masked failure. The zero reward is a placeholder,
excluded from Gym evaluation metrics by default.
- **On failure, cancellation, or shutdown:** release waiters and clean
up pending work.

### Recovery and compatibility

A judge failure returns HTTP 200 with `mask_sample: true`,
`failure_kind: judge_failed`, a bounded `failure_reason`, and the
existing failure-sidecar tags. The shared judge failsafe also preserves
`instance_config.mask_sample: true` for older RL consumers. This shared
response fix applies to every environment using `judge_failsafe`.

The existing evaluation opt-in `count_failure_classes_as_zero:
[judge_failed]` still includes selected failures as zeros. Only the
metric-input copy is unmasked; saved failure records and training masks
remain unchanged. Online and offline aggregation apply the same policy,
including when every rollout fails. With the opt-in, an
all-`judge_failed` run reports `mean/reward: 0.0`; without eligible
counted failures or successful results, collection still raises. The
main rollout JSONL never receives synthetic successes.

**A masked HTTP 200 result does not automatically make RL retry the
group.** It is a completed, unusable result for the current RL caller.
If the caller chooses replacement:

- **Explicit-ID groups:** advance the shared group-attempt number for
every member and submit a complete replacement group.
- **Legacy groups without an ID:** failed groups retain a failure
record. Reusing their task/prompt key returns failure with migration
guidance; recovery requires a fresh explicit group ID for the complete
group. This prevents delayed old members from joining a replacement.
Successful legacy groups still allow sequential new evaluations.

The judge-request default is now **1,800 seconds**, matching the overall
evaluation budget. This allows substantially more room for long
generations and queueing than the previous 300-second default;
deployments must still size both limits for their workload.

**Required RL coordination:** before NeMo RL adopts this Gym revision,
coordinate with [RL NVIDIA-NeMo#4160](NVIDIA-NeMo/RL#4160)
/ [NVIDIA-NeMo#4061](NVIDIA-NeMo/RL#4061). Reading the
mask alone is insufficient: TransferQueue must preserve it in the
training loss mask, and unavailable judge rewards must be excluded from
group statistics. NVIDIA-NeMo#4160 includes these fixes with `masked_reward_policy:
exclude` as the default. Top-level field ingestion remains separate
migration work, so the nested compatibility field remains. Before
updating RL's Gym revision, validate a real failed Gym response through
RL's loss mask and group statistics. Gym can merge first if RL keeps its
existing Gym revision until those fixes are available.

GenRM-specific usage and recovery documentation is now in the
[resources-server
README](https://github.com/NVIDIA-NeMo/Gym/blob/yuhez/3181-genrm-cohort-lifecycle/resources_servers/genrm_compare/README.md#genrm-comparison-groups),
rather than a new Fern site page.

Closes NVIDIA-NeMo#3181.

<details>
<summary>Validation</summary>

Current code: `6a7f30b707aee62c46478d757005f3fcc9fabe0d`, incorporating
main at `d4c0f86c7`. The latest fix is `059adb409`.

- **5,110 core, GenRM, and OSWorld dependency-policy tests passed** on
Python 3.13.14: `python -m pytest tests/unit_tests
resources_servers/genrm_compare/tests
responses_api_agents/osworld_agent/tests/test_dependency_policy.py -m
'not sandbox' -q --tb=short`. There were 103 skips, 673 sandbox tests
deselected, and 184 passing subtests. Sandbox tests were not run
locally.
- The all-failure regression checks one and four masked judge failures,
zero counting enabled/disabled, an unrelated counted class, and deferred
aggregation. Four cases fail on the preceding code and pass with this
fix; default and wrong-class error behavior remains intact. Existing
mixed-success tests remain covered.
- Actual TCP collector → SimpleAgent → GenRM tests check both collection
timeout and judge outage, with zero counting on/off. Online and offline
metrics match, the failure sidecar remains byte-for-byte unchanged, and
the main/merged rollout JSONL stays empty.
- Final all-files pre-commit, whitespace checks, and DCO sign-off
passed. [Exact-head
CI](https://github.com/NVIDIA-NeMo/Gym/actions/runs/35867697700) passed,
including core and all eight server shards. The preceding run passed
core and seven server shards but failed OSWorld dependency resolution;
merged NVIDIA-NeMo#3649 fixes that NumPy metadata conflict and is included in this
head. [Previous-head
CI](https://github.com/NVIDIA-NeMo/Gym/actions/runs/35681811360) passed
on `be7e6353d`, including core and all eight server shards.
- No new model run for the latest collector aggregation fix: it changes
metric eligibility, not policy generation or judging. The earlier
real-model evidence below is from `be7e6353d`.
- **Real-model functional validation passed:** Slurm job **19084624**,
exit 0 in **14m58s**, on `be7e6353d`. Gym's native CLI/HTTP path used
Qwen3-0.6B policy and Qwen3-Nemotron-235B-A22B-GenRM-2603 on eight
H100s. Eight policy generations and four completed real judge calls were
captured.
- The run preserved four masked failure answers during an injected judge
outage, then scored a caller-coordinated complete replacement attempt.
Four HTTP 503s and four interrupted HTTP response bodies recovered
within the retry budget without extra policy generation. A second
collector resume made no further model calls. Failed legacy groups
rejected delayed members, explicit zero-count aggregation preserved
saved masks, and rewards exactly matched the captured judgments. Saved
artifacts were independently inspected.
- The model run enabled `CUDA_LAUNCH_BLOCKING=1` for diagnostics after
an earlier CUDA failure. That failure did not recur; this does not
establish its root cause. The run's functional assertions passed before
the separate launcher teardown issue described below.
- No full RL training, checkpoint restart, or comparative throughput
claim.

</details>

<details>
<summary>Limits and follow-ups</summary>

The caller coordinates complete replacement attempts. There is no
partial scoring, automatic group scheduling, durable answer reuse, or
change to generation/judge overlap. Legacy identities cannot isolate
overlapping runs or provide completed replay. Explicit-ID replay and
failure records are process-local and bounded by retention; eviction or
restart is not a recovery protocol.

Ordinary individual reverification is unsupported. The existing
`judge_failed_only` guard bypass does not create groups or advance
attempts: retained failures remain failures until the caller supplies a
complete replacement with the required identity.

Active-group memory limits, retention optimization, caller-minted group
identities, durable outcomes, richer telemetry, and improvements to
retry backoff remain follow-up work. The inherited judge-failed-only
mode-guard bypass needs a separate compatibility change; documenting it
does not make individual GenRM reverification supported.

A Ray teardown abort was previously reproduced on both main and an
earlier PR head; its native cause remains undiagnosed. In the passing
job 19084624, the Gym launcher again exited with SIGABRT after the
functional assertions completed. The harness stopped the remaining owned
children and recorded no survivors; the separate policy server exited
normally and Slurm released the allocation. This validates the GenRM
failure/recovery behavior, not a clean Ray launcher shutdown.

</details>

### Related PRs: what is covered and what could follow

Built on NVIDIA-NeMo#2814 (Anish Mahishi), with deadline and cleanup requirements
from NVIDIA-NeMo#2385 (Teodor-Dumitru Ene) and NVIDIA-NeMo#2903 (Jiacheng Xu). Contributor
credit remains in commit trailers.

**NVIDIA-NeMo#2385:** the timeout requirement is covered. Its partial-scoring
policy is not adopted: if A/B/C arrive but D is missing, this PR fails
the incomplete group. Whole-group failure is allowed by NVIDIA-NeMo#3181. Close
NVIDIA-NeMo#2385 as superseded if that policy is accepted; otherwise retain/rework
partial scoring as separate follow-up work with explicit missing-member
and training/metric semantics.

**NVIDIA-NeMo#2903:** together, NVIDIA-NeMo#2814 and this PR cover request deduplication,
duplicate-disconnect safety, and failure cleanup. Three policy
differences remain:

- NVIDIA-NeMo#2903 permits a fresh group under the same identity after failure.
Here, replacement requires a newer shared explicit attempt, or a fresh
explicit ID when migrating a failed legacy group.
- NVIDIA-NeMo#2903 keeps completed rewards for the server's lifetime, including
legacy groups. Here, completed replay is bounded and restricted to
identical responses in explicit-ID groups; legacy groups have no
completed replay.
- NVIDIA-NeMo#2903 starts comparing available answer pairs while other answers are
still generating. Here, we preserve NVIDIA-NeMo#2814's wait-for-all schedule. Each
judge call needs only two answers, so early judging remains a possible
optimization; it does not require partial reward publication, and its
performance benefit has not been benchmarked.

After this PR merges, maintainers can close NVIDIA-NeMo#2903 as superseded if these
choices are accepted, tracking desired early judging or broader
replay/retry behavior separately. Alternatively, rework the existing PR
against current main to focus on the remaining scope. Closing either
older PR means its core problem is addressed and remaining choices are
accepted or tracked, not that every proposed behavior was merged.

---------

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Co-authored-by: Anish Mahishi <amahishi@nvidia.com>
Co-authored-by: Jiacheng Xu <jiachengx@nvidia.com>
Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

This branch was previously deployed

1 inactive deployment
public — fc8a43b3 Deployed Sep 1, 2026 by copy-pr-bot[bot] via release / finalize / notify #2383
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants