feat(token-id-capture): Claude Code external-harness example - #2128
Closed
ananthsub wants to merge 9 commits into
Closed
feat(token-id-capture): Claude Code external-harness example#2128ananthsub wants to merge 9 commits into
ananthsub wants to merge 9 commits into
Conversation
|
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. |
ananthsub
changed the title
This was referenced Jul 23, 2026
ananthsub
force-pushed
the
ananthsub/tokidcap/example
branch
from
July 29, 2026 06:55
a91b5aa to
a85399f
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 06:55
4a07f8d to
774339d
Compare
An external agent harness returns no token ids, and the ids do not survive to the client for a streamed response or one translated to Anthropic Messages. The model server holds the assembled response with the ids on it for a moment before either of those happens, which is the one point that covers every dialect. nemo_gym/token_id_capture/ records one TokenEntry per correlated model call: prompt ids, generated ids, one log prob per generated token, and the assistant text and tool calls, since a trainer reads the text for its own penalties. Records go to <rollout_id>.tokens.jsonl, per-file flock and fsync, awaited so a record is durable before the call returns. Nothing is added to the client response. TokenSink and TokenSource are the write and read seams. Gym owns the record shape and the capture code; a framework supplies the implementation and runs it where its tokens are produced, so sink placement is a deployment choice rather than a fork in the design. The package is a leaf -- no fastapi, ray, uvicorn, aiohttp or torch -- so a framework's inference worker can import it; a subprocess test enforces that. TokenEntry also carries optional parent_call_id, cum_len and digest. cum_len and digest are stamped at capture; parent_call_id stays null until the model server can resolve a parent. A capture failure is logged and also writes a <rollout_id>.tokens.incomplete marker, so a rollout that lost a call is distinguishable from a complete one. Capture stays best-effort: a bad payload must not break the harness run. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The builder is a pure function over a rollout's TokenEntry records. per_request emits one sequence per call. prefix_merging chains calls by the token-prefix relationship, parenting each call to the earlier call whose prompt-plus- generation is the longest prefix of this call's prompt, so an append-only multi-turn rollout becomes one chain. A prompt that extends nothing starts a new root, which is what a compacted or rewritten context looks like. Both are order-independent. Loss masks follow provenance: generated tokens are trained with their captured log probs, and anything re-fed into a prompt is not. The projection re-emits contiguous Responses items carrying both content and token ids, which is what NeMo-RL's postprocess already consumes. Main-chain selection is by generated-token mass across all roots, not by the first root. Entries are processed in increasing prompt length, so the first root is whichever root has the shortest prompt; a rollout's own first call is large because of the harness system prompt, so an auxiliary short-prompt call would be selected instead and the rollout dropped at delivery without an error. Recorded parent links are used when present and verified by digest rather than trusted, falling back to prefix inference on mismatch. A retry of the final call is reported unresolved rather than tie-broken, since nothing can say which generation the client received. Chain count, quarantined fraction and delivered fraction are returned rather than discarded, and a malformed capture returns an unbuilt result instead of raising into the caller's loop. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…med records Rollout collection replaces response.output with the merged, contiguous items for agents that opted into capture, so NeMo-RL reads response.output the same way for native and external-harness rollouts. Native agents are excluded by the per-agent opt-in: they already return exact ids inline, and a rebuild could differ from what the model server returned. Retention runs in both directions. Consumed records are deleted once folded into response.output (NG_KEEP_TOKCAP retains them), and stale records are cleared before dispatch. Both are needed because rollout ids are deterministic and the store appends, so a rerun would otherwise stitch a previous attempt's calls together with this one's. The build's counts ride the record under _ng_token_capture. Without them a rollout that trained on one of five calls is indistinguishable from one that trained on all five. A rollout captured incompletely, or whose final call was retried ambiguously, is marked for masking and warns. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Claude Code makes model calls that are not part of the rollout: it generates a conversation title and probes quota. They reach the model server on the same rollout-prefixed URL and get captured, and because they are genuine policy output -- real token ids, real log probs -- nothing downstream can tell they do not belong. Training on them optimizes the policy to write conversation titles under the rollout's reward. The record now keeps what the *harness* asked for (requested_model, has_tools), read off the parsed request body at the handler rather than by touching the body again in middleware. That is the signal, because a harness asks for a small model for these calls even though the server serves one model. Classification uses two signals and needs no harness-specific code in the core: an optional explicit pattern list for deployments that know their harness, and self-calibration -- whichever model generated the most tokens in a rollout is the policy model, and calls asking for a different one are side calls. Records written before this field existed carry an empty requested_model and are all kept, so nothing changes for them. Excluded calls are reported (side_calls_excluded) rather than silently dropped, and a rollout whose calls were *all* side calls is masked instead of yielding an empty trajectory. This is the second half of the title-call problem. The first was structural: a short side call became the main chain and the real rollout was dropped, fixed by selecting on generated-token mass. Even with the right chain selected, the side call would still have been stitched in. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The builder infers lineage after the fact, by finding the earlier call whose token sequence is the longest prefix of this call's prompt. That cannot do two things. It cannot resolve a retry: capture records a call once the response is assembled, including one the client never received, so a retry leaves two records with the same prompt and different generations, and both are equally valid children. And it cannot run before the call, which is where supplying the engine an exact prefix has to happen. This resolves the parent at request time using only what the harness already sends. A harness must echo the conversation to continue it, so the assistant turns in a request are the ones we produced; hashing them in order identifies the call that produced the last one. Nothing is added to the wire and nothing depends on the harness preserving a field we invented. Tool-call arguments are canonicalized (sorted-key JSON) for comparison only, because harnesses re-serialize them between turns -- compact one turn, pretty the next. Without that every tool-using turn would miss. The model's original argument string is what stays in the record. The index is a map keyed by call, not a running cursor, so it is a tree: two sub-agents branching from one parent both resolve to that parent and both would get the same prefix. A cursor would hand the second branch a prefix containing the first branch's generation, which the splice would apply without complaint. Entries are added and never mutated, so concurrent sub-agents cannot corrupt each other's lineage. Three cases deliberately return no parent rather than a guess: a new or rewritten conversation, an ambiguous match (two recorded calls with byte-identical output), and an evicted rollout. Each falls back to prefix inference, which is what happens today. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 07:12
774339d to
7342d4e
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/example
branch
from
July 29, 2026 07:12
a85399f to
16d8b00
Compare
This was referenced Jul 29, 2026
Without this the engine builds every prompt by re-rendering the conversation through the chat template. Re-tokenizing an assistant turn can produce a different split than the one sampled; a tool-call parser truncates the turn at the tool call and discards anything the model generated after it; and for a reasoning model the template drops earlier thinking entirely. Any of these means the new prompt does not extend the previous prompt-plus-generation, the builder cannot chain, and only the first call is trained on. Supplying the parent's cumulative ids as required_prefix_token_ids makes NeMo-RL's splice keep them verbatim and append only the newly rendered tail. The harness does not send this field and could not. Gym constructs the outbound engine request itself, so it is injected in _preprocess_chat_completion_create_params next to the sampling pin. The inbound request supplies identity; the outbound request supplies payload. The field also has to ride the separate /tokenize call that recovers prompt_token_ids. Generation applies the supplied prefix, so without it there the recorded prompt is a plain re-render that does not extend the previous call -- the chain looks broken even though the engine generated from the right tokens, and the record is what training consumes. NeMo-RL accepts the field on that endpoint for this reason. The splice applies whatever it is given without checking that it belongs to this conversation, so supply fires only on a unique, verified parent and otherwise forwards the request untouched. Each record carries prefix_supplied and the server tracks a supplied/eligible ratio, so a run can be audited afterwards rather than inferred from whether chains happen to be contiguous. Measured on a 2-step GRPO smoke, Qwen2.5-3B with a tool-calling task: 1/8 contiguous parent-child links with supply off, 6/6 with it on, and generation KL error 0.054 -> 0.023. Off by default; it needs a backend that honours the field. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
The route that serves a rollout's raw training tokens is registered on the same app the harness calls to generate. That is acceptable inside a trusted cluster. It is not acceptable once the harness runs in a sandbox whose only egress is this server, because the harness could read its own training data -- or another rollout's. token_id_capture_read_token requires a bearer token on the route, compared in constant time. When it is unset the route stays open and warns once, so existing deployments keep working and the gap is visible rather than silent. This is the one piece of sandboxing that has to land before the sandbox work, not with it. Everything else about capture is unaffected by sandboxing: capture happens in the model server, outside the sandbox, and the sandbox only ever sees text. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Add a framework-agnostic sampling pin to vllm_model: sampling_overrides forces temperature/top_p on every chat request, read from generic policy_generation_* keys with on-policy defaults. This keeps an external harness's rollouts on-policy for training without Gym holding any framework-specific knowledge — the training framework publishes its sampling into the generic keys. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Turns on token_id_capture for the claude_code_agent configs and restricts the harness to the Bash tool. Claude Code's system prompt is dominated by tool definitions and each allowed tool adds its schema to every call, so restricting the tool set lowers the fixed prompt cost; reasoning_gym only needs Bash, which its own system prompt already says. Also adds a last-quoted-string fallback to reasoning_gym answer extraction, so a conversational reply whose final answer is quoted is scored instead of an intermediate value pulled from the reasoning text. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 09:19
7342d4e to
f507b6f
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/example
branch
from
July 29, 2026 09:19
16d8b00 to
3860daf
Compare
ananthsub
force-pushed
the
ananthsub/tokidcap/sampling-pin
branch
from
July 29, 2026 12:38
f507b6f to
190f25c
Compare
Contributor
Author
|
Closing: six of the ten files here were scratch Python that the rollouts wrote during experiments, not intended content. The one useful change — |
adil-a
added a commit
that referenced
this pull request
Jul 30, 2026
…gent services (#2163) Part of the external-agent-integration epic #1396. Supersedes the collector-side `agent_url` approach proposed in #2006 (closed unmerged in favor of this PR). ## What An agent server, `responses_api_agents/remote_agent/`, that drives an agent service running **outside Gym's process tree** (your own repo, your own infrastructure) through Gym's tool loop. The remote service implements one endpoint **compliant with the OpenAI `/v1/responses` contract**, and the two servers compose as Responses-speaking agents: `POST {agent_base_url}/v1/responses` — receives the conversation so far (the row's `responses_create_params` with accumulated output and tool results appended to `input`) and returns a Responses API object: unpaired `function_call` items to ask Gym to execute environment tools, paired call+output items as records of its own internal tools (passed through untouched), or a final assistant message to finish. Gym owns everything else: it seeds the session and holds its cookies, executes the asked-for tools against the resources server, appends results and calls the service again, validates each reply against the Responses API schema (failing the rollout on mismatch), verifies, and returns the verify response from `/run`. **The resources server is never exposed to the service** — tool execution, session cookies, and `verifier_metadata` all stay inside Gym. The service's own cookies are round-tripped per call so it can keep per-rollout state; the only network direction is Gym → service. To the collector it is a normal named agent — resume, aggregation, and profiling work unchanged. The loop is `simple_agent`'s `responses()` with two marked in-loop divergences — the model hop is a hardened POST to `agent_base_url`, and only *unpaired* function_calls are executed (a call the service already answered itself is its own record, not an ask) — plus the never-raise `/run` contract that turns transport/validation errors into failure rows, and one deliberate cookie difference: the service's own cookies are round-tripped to it but stay out of the outgoing Set-Cookie (they are its private session, not Gym's). `run()` mirrors `simple_agent`'s seed → self-post `/v1/responses` → verify. ## Why a server instead of collector-side `agent_url` (#2006) - The token-capture/training stack (#2124–#2128, unmerged) gates participation by the agent's **name in Gym's config**; a url-dispatched agent has no config entry and is excluded by construction. A named RemoteAgent is compatible by design — note this PR contains no capture wiring; a wiring example is a follow-up once that stack lands. - Failure handling, validation, timeouts, and session plumbing live in the agent-server layer — this server's `/run` implementation plus the shared SimpleServer/ServerClient plumbing — instead of being reimplemented inside the shared collection loop. - `/run` responsibilities stay in the agent-server layer; the collector stays agent-agnostic. - Same pattern as every other harness integration (claude_code_agent, codex, …). ## Failure contract Failures never raise out of `/run`: remote endpoint down (3 total connection attempts, `ClientOSError`/`ServerDisconnectedError` only; timeouts and all other errors fail after a single attempt), per-call and whole-rollout timeouts, malformed/interrupted replies, seed/verify errors, and internal bugs all become reward-0 sentinel verify-responses (`_ng_failure_class="remote_agent_error"`) that rollout collection routes to the failures sidecar and retries on resume — non-terminal failures only, up to `NEMO_GYM_MAX_ROLLOUT_ATTEMPTS` (default 3). An invalid Responses object from the service is terminal (a schema bug will not fix itself on retry); the terminal flag crosses the HTTP self-post boundary by exception name. Reused rollout/failures JSONL rows carrying stale result keys are sanitized rather than crashing or leaking routing flags. Tool-level errors are NOT rollout failures: an unknown tool name or malformed arguments come back to the service as that call's `function_call_output`, matching `simple_agent`'s semantics. ## Testing - 51 offline tests (mocked ServerClient + mocked aiohttp seam; the `/v1/responses` self-post is routed into the real `responses()` with the exception middleware emulated): config validation, loop mechanics (multi-turn tool execution, paired-call pass-through, unknown-tool feedback, malformed-arguments feedback, max_steps, service-cookie round-trip, usage accumulation), every transport failure mode → sentinel, terminal classification across the route boundary, semaphore bounds incl. release-on-failure, run-wallclock-after-semaphore semantics, aggregate proxy + bound, route-level serialization (HTTP 200, never 500). - **Stateful E2E in-suite**: the real `example_session_state_mgmt` counter server in-process; the service returns unpaired tool asks, GYM executes them on the seeded session, reward 1.0 through the real verifier — and the test asserts the service was fed the counter value Gym read back. - **Collector round-trip in-suite**: real `RolloutCollectionHelper.run_from_config` driving this agent — successes to the main JSONL, sentinel rows to the failures sidecar. - **Live E2E, off-host, real agent**: a containerized service (own network namespace) wrapping the **Claude CLI** (opus via an internal gateway) as the decision-maker, `agent_base_url` pointing at the container's bridge IP; the counter resources server stayed on loopback with a random port — unreachable from the container by construction. **5/5 rollouts reward 1.0**, trajectories reading `function_call → function_call_output → … → message`, real token usage accumulated across loop turns (~8.3k mean/rollout). The only network path was Gym → container. ## Follow-ups - run-mode (`/run`-exposing self-scoring services) as a config mode on this server - token-capture wiring example once #2124–#2128 land (the service would route its model calls through Gym's model server at the per-rollout prefixed URL) - per review discussion: opt-in bounded-retry/timeout kwargs on core `request()` (then this server's transport collapses onto core), and two additive core helpers with existing adopters — a failure-row builder (stirrup/pinchbench hand-roll the same) and an aggregate-metrics proxy helper (eight agents hand-roll it) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: adil-a <adil.asif2000@hotmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
OlegSudakov
pushed a commit
to OlegSudakov/Gym
that referenced
this pull request
Aug 7, 2026
…gent services (NVIDIA-NeMo#2163) Part of the external-agent-integration epic NVIDIA-NeMo#1396. Supersedes the collector-side `agent_url` approach proposed in NVIDIA-NeMo#2006 (closed unmerged in favor of this PR). ## What An agent server, `responses_api_agents/remote_agent/`, that drives an agent service running **outside Gym's process tree** (your own repo, your own infrastructure) through Gym's tool loop. The remote service implements one endpoint **compliant with the OpenAI `/v1/responses` contract**, and the two servers compose as Responses-speaking agents: `POST {agent_base_url}/v1/responses` — receives the conversation so far (the row's `responses_create_params` with accumulated output and tool results appended to `input`) and returns a Responses API object: unpaired `function_call` items to ask Gym to execute environment tools, paired call+output items as records of its own internal tools (passed through untouched), or a final assistant message to finish. Gym owns everything else: it seeds the session and holds its cookies, executes the asked-for tools against the resources server, appends results and calls the service again, validates each reply against the Responses API schema (failing the rollout on mismatch), verifies, and returns the verify response from `/run`. **The resources server is never exposed to the service** — tool execution, session cookies, and `verifier_metadata` all stay inside Gym. The service's own cookies are round-tripped per call so it can keep per-rollout state; the only network direction is Gym → service. To the collector it is a normal named agent — resume, aggregation, and profiling work unchanged. The loop is `simple_agent`'s `responses()` with two marked in-loop divergences — the model hop is a hardened POST to `agent_base_url`, and only *unpaired* function_calls are executed (a call the service already answered itself is its own record, not an ask) — plus the never-raise `/run` contract that turns transport/validation errors into failure rows, and one deliberate cookie difference: the service's own cookies are round-tripped to it but stay out of the outgoing Set-Cookie (they are its private session, not Gym's). `run()` mirrors `simple_agent`'s seed → self-post `/v1/responses` → verify. ## Why a server instead of collector-side `agent_url` (NVIDIA-NeMo#2006) - The token-capture/training stack (NVIDIA-NeMo#2124–NVIDIA-NeMo#2128, unmerged) gates participation by the agent's **name in Gym's config**; a url-dispatched agent has no config entry and is excluded by construction. A named RemoteAgent is compatible by design — note this PR contains no capture wiring; a wiring example is a follow-up once that stack lands. - Failure handling, validation, timeouts, and session plumbing live in the agent-server layer — this server's `/run` implementation plus the shared SimpleServer/ServerClient plumbing — instead of being reimplemented inside the shared collection loop. - `/run` responsibilities stay in the agent-server layer; the collector stays agent-agnostic. - Same pattern as every other harness integration (claude_code_agent, codex, …). ## Failure contract Failures never raise out of `/run`: remote endpoint down (3 total connection attempts, `ClientOSError`/`ServerDisconnectedError` only; timeouts and all other errors fail after a single attempt), per-call and whole-rollout timeouts, malformed/interrupted replies, seed/verify errors, and internal bugs all become reward-0 sentinel verify-responses (`_ng_failure_class="remote_agent_error"`) that rollout collection routes to the failures sidecar and retries on resume — non-terminal failures only, up to `NEMO_GYM_MAX_ROLLOUT_ATTEMPTS` (default 3). An invalid Responses object from the service is terminal (a schema bug will not fix itself on retry); the terminal flag crosses the HTTP self-post boundary by exception name. Reused rollout/failures JSONL rows carrying stale result keys are sanitized rather than crashing or leaking routing flags. Tool-level errors are NOT rollout failures: an unknown tool name or malformed arguments come back to the service as that call's `function_call_output`, matching `simple_agent`'s semantics. ## Testing - 51 offline tests (mocked ServerClient + mocked aiohttp seam; the `/v1/responses` self-post is routed into the real `responses()` with the exception middleware emulated): config validation, loop mechanics (multi-turn tool execution, paired-call pass-through, unknown-tool feedback, malformed-arguments feedback, max_steps, service-cookie round-trip, usage accumulation), every transport failure mode → sentinel, terminal classification across the route boundary, semaphore bounds incl. release-on-failure, run-wallclock-after-semaphore semantics, aggregate proxy + bound, route-level serialization (HTTP 200, never 500). - **Stateful E2E in-suite**: the real `example_session_state_mgmt` counter server in-process; the service returns unpaired tool asks, GYM executes them on the seeded session, reward 1.0 through the real verifier — and the test asserts the service was fed the counter value Gym read back. - **Collector round-trip in-suite**: real `RolloutCollectionHelper.run_from_config` driving this agent — successes to the main JSONL, sentinel rows to the failures sidecar. - **Live E2E, off-host, real agent**: a containerized service (own network namespace) wrapping the **Claude CLI** (opus via an internal gateway) as the decision-maker, `agent_base_url` pointing at the container's bridge IP; the counter resources server stayed on loopback with a random port — unreachable from the container by construction. **5/5 rollouts reward 1.0**, trajectories reading `function_call → function_call_output → … → message`, real token usage accumulated across loop turns (~8.3k mean/rollout). The only network path was Gym → container. ## Follow-ups - run-mode (`/run`-exposing self-scoring services) as a config mode on this server - token-capture wiring example once NVIDIA-NeMo#2124–NVIDIA-NeMo#2128 land (the service would route its model calls through Gym's model server at the per-rollout prefixed URL) - per review discussion: opt-in bounded-retry/timeout kwargs on core `request()` (then this server's transport collapses onto core), and two additive core helpers with existing adopters — a failure-row builder (stirrup/pinchbench hand-roll the same) and an aggregate-metrics proxy helper (eight agents hand-roll it) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: adil-a <adil.asif2000@hotmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Claude Code external-agent example
Top PR in the stack (base: #2127).
A runnable end-to-end example that exercises the whole stack, plus a scoring fix the example needs.
What it adds
claude_code_agentconfig, pointed at the Gym model server, with capture turned on. This is the reference setup: the agent calls the model server, the server records the token ids, and the rollout is reassembled for training.reasoning_gymanswer extraction: when the model's output has no<answer>tag and no\boxed{}, use the last quoted string as the answer. A conversational agent like Claude Code usually states its final answer as the last quoted phrase; scoring the whole reasoning text instead can pick up an intermediate guess and mark a correct answer wrong.Why it's here
This is the concrete target for the stack: reasoning_gym with Claude Code, capture on, producing rollouts that carry token ids and can be trained on. It's also what motivates the scoring fix — without it, correct answers scored zero and the reward didn't move.
Tests
Verified in the 2-GPU Megatron GRPO run: reward moved and the run stayed on-policy (generation KL ~0.04). The reasoning_gym change is a small fallback in
_extract_answer_from_response.Stack
One commit per PR, each based on the previous branch (bottom of the stack targets
main):All nine are drafts. #2183 supersedes #2127, which could not be retargeted after
the stack grew (GitHub forbids changing the base of a PR that is part of a stack).