feat(token-id-capture): deliver rebuilt trajectories safely

[ananthsub]

Attaches a safely rebuilt token trajectory to each participating rollout and retires capture evidence only after durable handoff.

Snapshot and retirement flow

sequenceDiagram
    participant C as Rollout collector or framework
    participant S as TokenSource
    participant B as Trajectory builder
    participant D as Durable downstream

    C->>C: select static agent capture or all_agents
    C->>S: freeze(rollout_id)
    S-->>C: TokenCaptureSnapshot(entries, incomplete, snapshot_id, version)
    C->>B: build frozen snapshot
    alt incomplete, missing, ambiguous, or unsafe
        B-->>C: mask_sample=true and diagnostics
        Note over C,S: Retain frozen evidence
    else one safe trainable chain
        B-->>C: rebuilt response and metrics
        C->>D: write rebuilt rollout
        D-->>C: durable acknowledgement
        C->>S: drop(rollout_id, snapshot_id, version)
        S-->>C: conditional result - frozen tombstone remains
    end

Loading

Retirement is outside reconstruction. A delayed consumer cannot delete a newer attempt, and a late writer cannot recreate a retired attempt.

Summary

• Clears and finalizes records only for agents selected by static token_id_capture configuration or all_agents; tokenless nonparticipants remain ordinary evaluation rollouts.
• Freezes the source before reconstruction so entries and incomplete state come from one atomic snapshot.
• Masks missing, malformed, incomplete, ambiguous, multi-root, multi-chain, and empty-token builds.
• Separates finalization from retirement and conditionally drops only the matching snapshot_id and version after durability.
• Fsyncs Gym's JSONL result before retirement and retains failed or masked evidence.
• Uses an installed caller-owned source without closing it, or constructs and closes Gym's default file source.
• Fails before rollout dispatch when selected rows require rebuilding and the collector process has no source.

Depends on #2125. Documented in #2341.

[copy-pr-bot]

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

Contributors can view more details about this message here.

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

[cmunley1]

left a question above, but i think its unlikely in normal cases and we could follow up unless you want to reply there now.

[ananthsub]

/claude review

[claude]

SHIP — no reliability concerns.

Reviewed the token-capture delivery step and its wiring into run_rollouts. The correctness-critical properties hold:

• Durable-before-delete ordering is correct. finalize_rollout_token_capture mutates result["response"]["output"] before orjson.dumps(result), so the persisted record carries the rebuilt trajectory; os.fsync(results_file.fileno()) runs before retire_rollout_token_capture, so source records are dropped only after the rebuilt trajectory is durable. No crash-window data loss.
• Reward/scorer untouched — only response.output is replaced; reward and components are left alone (asserted in tests).
• Never-raises holds — trajectories_from_source and source.drop are both wrapped; a failing transport masks one rollout, not the batch.
• Retirement is safe against deterministic-id reuse — drop is conditioned on seal_id+version; a stale build cannot delete a newer attempt's records. Masked/failed builds are retained as evidence, not retired.
• Ownership/lifecycle correct — owned_token_source is closed only when Gym constructed it (configured source or file store), never a process-installed source; second finalize call is a no-op.
• Async hygiene clean — no httpx, no ray.get, all coroutines awaited; no undeclared imports; all symbols exported from the package __init__.

Test coverage is real and exercises masking, no-capture-key, nothing-recorded, mixed native/external batches, retention, and stale-record clearing.

NOTE (non-blocking, author's call): finalize_rollout_token_capture is invoked for every result including no_persist (kill_shaped) and failure_class rows that are never persisted, so a rebuild + possible masking warning is computed and discarded for those. Harmless — no reward/output impact and results aren't written — just minor wasted work and log noise on failure-heavy runs.

[ananthsub]

/claude review

[claude]

Review incomplete: the PR revision changed during review. The captured base (970603d) and head (84755ea) that this review analyzed no longer match the live PR, which now points at base b9746c7 and head 21e7ab0. The final pre-publish revision check therefore failed, so inline findings are withheld to avoid anchoring comments to a stale commit. Please re-run the review against the current head.

[ananthsub]

/claude review

[cmunley1]

With all_agents=True, if capture records are written, this early return skips the snapshot retirement https://github.com/NVIDIA-NeMo/Gym/pull/2126/changes#diff-4ed3026c5b1fe1967854548a34299c33ba36c8fd55c7cd2249324108dd0a550cR915, so those redundant records may accumulate. Should we avoid capture for agents who provide the token ids or still retire their captured snapshot?