Skip to content

feat(rollout): expose per-rollout enrichment to run_examples callers - #5023

Draft
pthombre wants to merge 1 commit into
mainfrom
pthombre/rollout-enrichment
Draft

pthombre wants to merge 1 commit into
mainfrom
pthombre/rollout-enrichment

Conversation

@pthombre

@pthombre pthombre commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

What

run_examples returns the raw /run result, so a training framework that dispatches through it (NeMo-RL) never gets what ng_collect_rollouts records for each rollout: the model server's captured calls, the normalized ng_trajectory, and ng_perf. This PR exposes those collector steps and fixes two gaps that showed up once a framework uses them.

  • enrich_run_result(row, result, *, capture_dirs, observability_enabled, include_payloads=True, rollout_latency_ms=None, delete_capture=False) (public, in rollout_collection.py). It runs the collector's per-rollout steps in place:

    1. wait for capture quiescence;
    2. merge_model_call_capture_into_record, including the Claude Code compaction association;
    3. _attach_trajectory_record;
    4. _attach_ng_perf.

    It copies the row's rollout identity onto the result first, because the capture key is read from the record. Reward and harness output are never touched. delete_capture removes the rollout's capture file after the merge. ng_collect_rollouts now calls this function, so there is one implementation and no behavior change.

  • run_examples(..., enrich_results=False, include_model_call_payloads=True, delete_model_call_capture=False): an opt-in way to have run_examples apply the function to each result. With the default off, the documented "exactly what /run returned" contract is unchanged.

  • Anthropic Messages turns. Turns built from captured calls (_turn_content) handled only Responses and Chat Completions, so every Claude Code turn had answer=None. Anthropic calls are now mapped onto Responses items with AnthropicConverter, which covers system in the question, thinking as reasoning and tool_use as function_call. If the converter rejects a block type, the turn falls back to the raw blocks.

  • Fix: _trajectory_identity raised KeyError whenever a row had no _ng_task_index, even when it carried task_id. Python evaluates the default argument of next(...) eagerly. run_examples rows need not have a task index, so this silently replaced the projection with a trajectory_projection_failed gap. The task id now falls back to the explicit rollout id.

Why

NeMo-RL is adding a debug log of every Gym rollout (NVIDIA-NeMo/RL PR, linked below). The only complete record of what a rollout did is the model server's capture: every policy call, subagents, summarizers, and judge calls made through ServerClient. RL reaches Gym through run_examples, so without this the log holds only each agent's self-report, and for many agents that is just the final answer.

No issue filed. The motivation and design are in the RL PR and its plan.

Validation

  • Unit tests (tests/unit_tests/test_rollout_collection.py). The new classes are TestEnrichRunResult: merge, metadata-only plus deletion, identity-only without capture, a missing task index, and run_examples opt-in. Also TestAnthropicTurns: tool_use + thinking, and the unsupported-block fallback. Run inside the 20261002 nemo-rl nightly-gym image with a baked Gym venv:
    • this branch (on 6f295d7): -k "EnrichRunResult or AnthropicTurns or TurnsFromModelCalls or capture or trajectory" → 47 passed.
    • the same change ported to 17abf2d, the image's Gym, without the Anthropic part since that Gym has no turn projection → 32 passed.
  • ruff check and ruff format --check on the changed files are clean. pre-commit run --all-files was not run.

Rollout evidence

Nemotron 3.5 Nano RLVR single-controller smoke: 10 nodes, the amplified-dolphin blend across 11–12 agents including LLM-judge environments. It ran the 17abf2d port of this change through NeMo-RL's actor, calling enrich_run_result per result, with observability_enabled:

token capture off token capture on
rollouts logged 415 438
with merged model calls (bodies included) 415 (563 calls, 1–5 per rollout) 438 (590 calls)
enrich errors / capture gaps 0 / 0 0 / 0
capture files left for already-merged rollouts 0 0

How NeMo-RL uses it

NVIDIA-NeMo/RL#4890 calls enrich_run_result directly in its NemoGym actor, on a shallow copy of each /run result, rather than run_examples(enrich_results=True). That keeps the large call bodies in the actor's log, and the result sent on to the trainer is unchanged. The function only replaces top-level keys, so a shallow copy is safe. run_examples(enrich_results=True) remains for callers that want the enriched result itself.

RL also clears each rollout's stale capture file before dispatch with clear_model_call_captures_for_rollouts. Unlike the collector, RL can re-dispatch a row under the same rollout id.

Compatibility

All new behavior is additive and opt-in. ng_collect_rollouts output is unchanged except for two cases:

  • Anthropic-dialect turns now have answers.
  • A row without a task index now gets a trajectory instead of a projection-failure gap.

🤖 Generated with Claude Code

ng_collect_rollouts merges each rollout's captured model calls, projects
ng_trajectory and assembles ng_perf, but run_examples returns the raw /run
result, so training frameworks that dispatch through it never see the
model-server record of the rollout. Move those steps into a public
enrich_run_result() used by the collector, and let run_examples apply it
with enrich_results=True (optionally metadata-only, optionally deleting
each rollout's capture file after the merge).

Trajectory turns built from captured calls now also understand the
Anthropic Messages dialect (Claude Code), mapping it onto Responses items;
previously those turns had answer=None.

Also name a trajectory by its explicit rollout id when the row carries no task index (run_examples rows need not): _trajectory_identity evaluated next()'s default eagerly and raised KeyError even for rows with a task_id.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Oct 23, 2026

Copy link
Copy Markdown

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.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added area:evaluation Evaluation, rollout collection, reward profiling, comparison, and exporters feature New capabilities, enhancements, or enablement work labels Oct 23, 2026

@marcromeyn marcromeyn left a comment

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.

Thanks, one shared implementation for ng_collect_rollouts and run_examples is the right base; I'd like to build the RL-path records (FEP-1907) on it. A few points, reviewed at 0b1ee91b:

  1. Rebase onto main. This head predates #5081 (merged Oct 26), which adds record_omitted_rollouts: rows marked _ng_omitted_rollout for a try killed before it produced a result and for a row the dispatch budget never started, with readers taking a rollout's latest executed try. The no_result branch here sits right next to that code, so worth checking after the rebase that enrich_run_result keeps the same distinction (interrupted vs never started), since a framework calling it will see those results too.
  2. Task id fallback. When a row has neither a task id field nor _ng_task_index, the trajectory's task_id becomes the rollout id. That makes every repeat of a task look like a different task to anything grouping by task. I'd leave task_id unset in that case (unknown rather than guessed), or have the caller stamp the task index.
  3. delete_capture with include_payloads=False. The docstring says that without payloads "the bodies stay in the capture store", but delete_capture then removes that store, so the bodies are gone. Either refuse the combination or say in the docstring that deleting drops them.
  4. Identity copy. Copying the row's rollout id, task, repeat and attempt index only when the result lacks them looks right. Once the capture-identity design lands (agent instance, model_ref, weight_version), those belong in the same copy; no change needed now.
  5. Small: the Anthropic fallback swallows every converter error silently. A debug log with the error type would help when turns lose their normalization.

Happy to look at the tests in more detail after the rebase.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:evaluation Evaluation, rollout collection, reward profiling, comparison, and exporters feature New capabilities, enhancements, or enablement work

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants