feat: judge cohort hardening

[tdene]

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]

/ok to test ddfa0e2

[cmunley1]

/claude review

[kajalj22]

/claude review

[claude]

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.

[claude]

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.

[tdene]

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.

[ananthsub]

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