From 767c48834bc96f92c71752169000495487659d63 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:31:30 -0700 Subject: [PATCH 1/9] feat(token-id-capture): capture exact token ids per model call 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 .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 .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 --- nemo_gym/base_responses_api_model.py | 109 +++-- nemo_gym/global_config.py | 3 + nemo_gym/token_id_capture/__init__.py | 81 ++++ nemo_gym/token_id_capture/config.py | 55 +++ nemo_gym/token_id_capture/protocols.py | 93 +++++ nemo_gym/token_id_capture/reader.py | 61 +++ nemo_gym/token_id_capture/records.py | 180 +++++++++ nemo_gym/token_id_capture/routes.py | 55 +++ nemo_gym/token_id_capture/sink.py | 142 +++++++ nemo_gym/token_id_capture/source.py | 58 +++ nemo_gym/token_id_capture/store.py | 144 +++++++ tests/unit_tests/test_token_id_capture.py | 472 ++++++++++++++++++++++ 12 files changed, 1430 insertions(+), 23 deletions(-) create mode 100644 nemo_gym/token_id_capture/__init__.py create mode 100644 nemo_gym/token_id_capture/config.py create mode 100644 nemo_gym/token_id_capture/protocols.py create mode 100644 nemo_gym/token_id_capture/reader.py create mode 100644 nemo_gym/token_id_capture/records.py create mode 100644 nemo_gym/token_id_capture/routes.py create mode 100644 nemo_gym/token_id_capture/sink.py create mode 100644 nemo_gym/token_id_capture/source.py create mode 100644 nemo_gym/token_id_capture/store.py create mode 100644 tests/unit_tests/test_token_id_capture.py diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 99426db471..25a2191a12 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -67,6 +67,16 @@ BaseServer, SimpleServer, ) +from nemo_gym.token_id_capture import ( + CaptureContext, + capture_tokens, + reset_token_sink, + set_token_sink, +) + +# The read route and its store factory need Gym's server stack, so they are not +# re-exported from the leaf package (see nemo_gym/token_id_capture/__init__.py). +from nemo_gym.token_id_capture.routes import install_token_capture_routes, make_token_store logger = logging.getLogger(__name__) @@ -90,7 +100,12 @@ def setup_webserver(self) -> FastAPI: self.setup_session_middleware(app) capture_config = ModelCallCaptureConfig.model_validate(self.server_client.global_config_dict) - install_model_call_capture(app, capture_config, model_server_name=self.config.name) + install_model_call_capture( + app, + capture_config, + model_server_name=self.config.name, + global_config_dict=self.server_client.global_config_dict, + ) app.post("/v1/chat/completions")(self.chat_completions_dispatch) @@ -192,8 +207,11 @@ async def _invoke_chat_completions( # only `body`. Dispatch on whichever this server declares so the shared dispatch works for # all of them. if "request" in inspect.signature(self.chat_completions).parameters: - return await self.chat_completions(request=request, body=params) - return await self.chat_completions(body=params) + completion = await self.chat_completions(request=request, body=params) + else: + completion = await self.chat_completions(body=params) + await capture_tokens(completion) + return completion async def messages(self, request: Request, body: dict = Body()): """Default Anthropic Messages <-> Responses mapping shared by every Gym model server. @@ -206,6 +224,7 @@ async def messages(self, request: Request, body: dict = Body()): """ params = _ANTHROPIC_CONVERTER.anthropic_request_to_responses(body) response = await self._invoke_responses(request, params) + # Capture here: the Anthropic response returned below has already dropped token ids. model_name = body.get("model") or response.model anthropic_response = _ANTHROPIC_CONVERTER.responses_to_anthropic_response(response, model=model_name) if body.get("stream"): @@ -222,8 +241,14 @@ async def _invoke_responses( # `body`. Dispatch on whichever this server declares so the default messages() works for # all of them. if "request" in inspect.signature(self.responses).parameters: - return await self.responses(request=request, body=params) - return await self.responses(body=params) + response = await self.responses(request=request, body=params) + else: + response = await self.responses(body=params) + # Capture here rather than at the route: the streaming dispatch returns a StreamingResponse + # and the Anthropic mapping drops the token fields, so this is the last point where the + # assembled response still carries them, for every dialect. + await capture_tokens(response) + return response def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming: @@ -1003,10 +1028,19 @@ class _CaptureMiddleware: prefix and forwards only. """ - def __init__(self, app: Any, *, store: Optional[CaptureStore], model_server_name: Optional[str]) -> None: + def __init__( + self, + app: Any, + *, + store: Optional[CaptureStore], + model_server_name: Optional[str], + token_store: Any = None, + ) -> None: self._app = app self._store = store self._model_server_name = model_server_name + # When set, correlated+observed calls also record training tokens via a per-request sink. + self._token_store = token_store async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: if scope.get("type") != "http": @@ -1021,24 +1055,36 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None path = prefix_match.group("rest") scope = {**scope, "path": path, "raw_path": path.encode("utf-8")} - # Capture disabled: the prefix is already stripped (routing preserved), so just forward. - if self._store is None: - await self._app(scope, receive, send) - return + dialect = _OBSERVED_PATHS.get(path) - # Only explicitly correlated model calls are captured. An unprefixed call is forwarded - # unchanged rather than being mixed with unrelated calls under a shared fallback key. - if rollout_from_path is None: + # Nothing to capture: neither store is active, the call isn't correlated to a rollout, or the + # path isn't an observed model endpoint. The prefix is already stripped, so just forward. + # An unprefixed call is forwarded rather than mixed with unrelated calls under a shared key. + if (self._store is None and self._token_store is None) or rollout_from_path is None or dialect is None: await self._app(scope, receive, send) return - dialect = _OBSERVED_PATHS.get(path) - if dialect is None: - await self._app(scope, receive, send) # not observed (or a stripped non-/v1 path) - return - rollout_id = rollout_from_path model_call_id = uuid4().hex + + # Hand the model server a per-request token sink keyed to this call. It records token ids + # from its complete response (the middleware can't -- token ids are dropped on the SSE wire). + sink_token = None + if self._token_store is not None: + sink_token = set_token_sink( + CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, store=self._token_store) + ) + + # Training-token capture only: no evaluation record, so skip the response buffering entirely + # and just forward with the sink live. + if self._store is None: + try: + await self._app(scope, receive, send) + finally: + if sink_token is not None: + reset_token_sink(sink_token) + return + request_body = bytearray() async def _receive() -> dict[str, Any]: @@ -1120,6 +1166,10 @@ async def _flush_deferred_response() -> None: finally: await _flush_deferred_response() raise + finally: + # The sink is only needed while the model server produces the response. + if sink_token is not None: + reset_token_sink(sink_token) completed_at = time.time() latency_ms = (time.perf_counter() - start) * 1000.0 @@ -1187,22 +1237,35 @@ def _parse_and_record() -> None: def install_model_call_capture( - app: Any, config: ModelCallCaptureConfig, *, model_server_name: Optional[str] = None + app: Any, + config: ModelCallCaptureConfig, + *, + model_server_name: Optional[str] = None, + global_config_dict: Any = None, ) -> None: """Install model-call capture middleware. Always installed so the ``/ng-rollout/`` correlation prefix is stripped before routing regardless of whether capture is enabled (otherwise a default ``gym eval`` would 404 on every - prefixed model call). When capture is enabled the middleware additionally records each observed - call's request + response into a rollout-keyed CaptureStore while forwarding bytes downstream - unchanged (non-terminal SSE chunks are forwarded as they arrive; the terminal event follows the - durable capture write). + prefixed model call). When evaluation capture is enabled the middleware additionally records each + observed call's request + response into a rollout-keyed CaptureStore while forwarding bytes + downstream unchanged (non-terminal SSE chunks are forwarded as they arrive; the terminal event + follows the durable capture write). + + Training-token capture is a separate, independently-gated concern that reuses the same + correlation point: when enabled, the middleware hands the model server a per-request token sink + (keyed by the same rollout id and model_call_id) and the server records token ids from its + complete response. The read route is registered only when that capture is enabled. """ + token_store = make_token_store(global_config_dict) if global_config_dict is not None else None app.add_middleware( _CaptureMiddleware, store=make_capture_store(config), model_server_name=model_server_name, + token_store=token_store, ) + if token_store is not None: + install_token_capture_routes(app, token_store) # --- Run-level capture helpers (rollout-collection side) --- diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 206984628d..62f2a9e9e1 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -89,6 +89,9 @@ QUERY_KEY_NAME = "query" OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" +TOKEN_ID_CAPTURE_ENABLED_KEY_NAME = "token_id_capture_enabled" +# Per-agent opt-in (on an agent's config block) for participating in training token capture. +TOKEN_ID_CAPTURE_KEY_NAME = "token_id_capture" COMPONENT_NAME_KEY_NAME = "component_name" NEMO_GYM_RESERVED_TOP_LEVEL_KEYS = [ CONFIG_PATHS_KEY_NAME, diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py new file mode 100644 index 0000000000..f7cbdd925f --- /dev/null +++ b/nemo_gym/token_id_capture/__init__.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Training-token capture: produce, store, read, and source ``TokenEntry`` records. + +This is the per-model-call training data path, kept separate from evaluation +capture. The capture middleware sets a per-request capture context; the model +server records a ``TokenEntry`` from its complete response; a trainer reads a +rollout's entries through a ``TokenSource``. + +**This package is a leaf.** Importing it must not pull in fastapi, ray, uvicorn, +aiohttp, requests, or torch, because a training framework's inference worker +imports the record, the protocols, and the capture core to write into its own +data plane (see ``protocols.py``). The HTTP read route and the HTTP reader do +need Gym's server stack, so they are deliberately *not* re-exported here -- +import ``nemo_gym.token_id_capture.routes`` / ``.reader`` directly from server +code. ``tests/unit_tests/test_token_id_capture.py::test_package_is_dependency_free_leaf`` +enforces this. +""" + +from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.protocols import ( + TokenSink, + TokenSource, + install_token_sink, + installed_token_sink, +) +from nemo_gym.token_id_capture.records import ( + TOKEN_FIELDS, + TokenEntry, + compute_digest, + cumulative_tokens, + encode_token_ids, + extract_token_fields, + stamp_lineage, +) +from nemo_gym.token_id_capture.sink import ( + CaptureContext, + TokenCaptureContext, + capture_tokens, + reset_token_sink, + set_token_sink, +) +from nemo_gym.token_id_capture.source import CaptureTokenSource +from nemo_gym.token_id_capture.store import TokenCaptureStore, validate_rollout_id + + +__all__ = [ + "TokenIdCaptureConfig", + "TokenEntry", + "TOKEN_FIELDS", + "extract_token_fields", + "compute_digest", + "encode_token_ids", + "cumulative_tokens", + "stamp_lineage", + "TokenCaptureStore", + "validate_rollout_id", + "TokenSink", + "TokenSource", + "install_token_sink", + "installed_token_sink", + "CaptureContext", + "TokenCaptureContext", + "set_token_sink", + "reset_token_sink", + "capture_tokens", + "CaptureTokenSource", +] diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py new file mode 100644 index 0000000000..8c587822a7 --- /dev/null +++ b/nemo_gym/token_id_capture/config.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run-wide switch for training-token capture. + +This is a separate switch from evaluation capture (``observability_enabled``). +Evaluation capture records a compact request/response summary; training-token +capture records token ids and log probabilities for RL. A run can enable either, +both, or neither. When no dedicated directory is given, tokens are written +alongside the eval capture files in ``model_call_capture_dir``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Optional + +from pydantic import BaseModel, ConfigDict, model_validator + + +class TokenIdCaptureConfig(BaseModel): + model_config = ConfigDict(extra="ignore") + + token_id_capture_enabled: bool = False + token_id_capture_dir: Optional[Path] = None + # Shared fallback directory (also used by evaluation capture). + model_call_capture_dir: Optional[Path] = None + + @model_validator(mode="after") + def _validate(self) -> "TokenIdCaptureConfig": + if not self.token_id_capture_enabled: + return self + directory = self.resolved_dir() + if directory is None: + raise ValueError( + "token_id_capture_dir (or model_call_capture_dir) is required when token_id_capture_enabled=true" + ) + if not directory.is_absolute(): + raise ValueError("training-token capture directory must be an absolute path") + return self + + def resolved_dir(self) -> Optional[Path]: + return self.token_id_capture_dir or self.model_call_capture_dir diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py new file mode 100644 index 0000000000..74d1c268a9 --- /dev/null +++ b/nemo_gym/token_id_capture/protocols.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The write and read seams for captured training tokens. + +Gym owns the record shape, these two protocols, and the code that builds a +record. A training framework supplies the implementation and runs it wherever +its tokens are produced. Neither side imports the other's transport. + +This makes the *placement* of the write a deployment choice rather than a fork +in the design: + +- Gym owns serving (today): install the sink in the model server, which already + holds the assembled response. No extra hop. +- A framework owns the inference worker: install the sink there, so bulk token + arrays go straight to the framework's data plane instead of riding back + through Gym's HTTP response. + +Same capture code both times. This module must stay dependency-free (no +fastapi, ray, torch, aiohttp) so a framework's worker can import it without +pulling in Gym's server stack; ``tests/unit_tests/test_token_id_capture.py`` +enforces that. +""" + +from __future__ import annotations + +from typing import Optional, Protocol, runtime_checkable + +from nemo_gym.token_id_capture.records import TokenEntry + + +@runtime_checkable +class TokenSink(Protocol): + """Where captured records go. Implemented by Gym's file store, or by a + framework over its own transport.""" + + async def put(self, entry: TokenEntry) -> None: + """Append one record. + + MUST be durable on return: a later ``tokens_for`` for the same rollout + has to see it. Delete-on-consume and post-rollout reads are only correct + because of this. + + May raise. The caller counts the failure and marks the rollout, and + never fails the model call because of it. + """ + ... + + +@runtime_checkable +class TokenSource(Protocol): + """Where a trajectory builder reads records from.""" + + async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: + """All records for a rollout, in any order. + + Order carries no meaning: calls run concurrently and may be served by + different workers. The builder recovers structure from the records + themselves (parent links, or token-prefix relationships). + """ + ... + + async def drop(self, rollout_id: str) -> None: + """Retire a rollout's records once they have been consumed.""" + ... + + +# Installed once at process startup by whoever owns the process: Gym's model +# server, or a framework's inference worker. The capture path reads it when a +# request-scoped context does not carry an explicit sink. +_INSTALLED_SINK: Optional[TokenSink] = None + + +def install_token_sink(sink: Optional[TokenSink]) -> None: + """Set (or clear, with ``None``) the process-wide default sink.""" + global _INSTALLED_SINK + _INSTALLED_SINK = sink + + +def installed_token_sink() -> Optional[TokenSink]: + return _INSTALLED_SINK diff --git a/nemo_gym/token_id_capture/reader.py b/nemo_gym/token_id_capture/reader.py new file mode 100644 index 0000000000..638540c52b --- /dev/null +++ b/nemo_gym/token_id_capture/reader.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Readers that return a rollout's ``TokenEntry`` records. + +Two readers implement the same async ``read`` method: + +* ``LocalTokenReader`` reads the store's files directly. Use it when the reader + runs in the same process (or box) as the model server that wrote them. +* ``HttpTokenReader`` reads over the model server's ``/ng-capture/tokens`` + route. Use it when the trainer is not co-located with the store -- the common + case once serving and training run on different nodes. It goes through Gym's + shared aiohttp client, so reading at high rollout concurrency never stalls the + event loop. +""" + +from __future__ import annotations + +from typing import Protocol + +from nemo_gym.server_utils import raise_for_status, request +from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +class TokenReader(Protocol): + async def read(self, rollout_id: str) -> list[TokenEntry]: ... + + +class LocalTokenReader: + def __init__(self, store: TokenCaptureStore) -> None: + self._store = store + + async def read(self, rollout_id: str) -> list[TokenEntry]: + return self._store.read_entries(rollout_id) + + +class HttpTokenReader: + def __init__(self, base_url: str, api_key: str = "dummy_key") -> None: + self._base_url = base_url.rstrip("/") + self._headers = {"Authorization": f"Bearer {api_key}"} + + async def read(self, rollout_id: str) -> list[TokenEntry]: + # Timeouts and retries are governed by Gym's shared aiohttp client. + url = f"{self._base_url}/ng-capture/tokens/{rollout_id}" + response = await request("GET", url, headers=self._headers) + await raise_for_status(response) + text = await response.text() + return [TokenEntry.model_validate_json(line) for line in text.splitlines() if line.strip()] diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py new file mode 100644 index 0000000000..abd1eb39f8 --- /dev/null +++ b/nemo_gym/token_id_capture/records.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The training-token record and how to pull it off a served response. + +A ``TokenEntry`` holds only what a trainer needs from one model call: the exact +prompt token ids the engine ran on, the generated token ids, and one log +probability per generated token. It is deliberately separate from the model-call +capture record used for evaluation (``ModelCallRecord``): the eval record is a +compact request/response summary and never carries token ids, while a +``TokenEntry`` is large and read only when building training data. Keeping them +apart lets eval reads skip the token payloads and lets training token ids move +to a different store later without touching the eval schema. + +Both records for the same model call share a ``model_call_id``, so training can +join a ``TokenEntry`` to its ``ModelCallRecord`` when it needs the eval context. +""" + +from __future__ import annotations + +import hashlib +import struct +from typing import Any, Optional + +from pydantic import BaseModel, ConfigDict + + +# The fields the model server attaches to a served response when token-id return +# is on. ``routed_experts`` is present only for MoE backends that report it. +TOKEN_FIELDS = ("prompt_token_ids", "generation_token_ids", "generation_log_probs", "routed_experts") + +# Bumped if the digest encoding below ever changes, so a stale digest fails to +# verify instead of silently comparing equal. +DIGEST_VERSION = 1 +_DIGEST_DOMAIN = b"nemo-gym-tokens" +_EMPTY_DIGEST = hashlib.sha256(_DIGEST_DOMAIN).hexdigest() + + +def encode_token_ids(token_ids: list[int]) -> bytes: + """Stable, length-delimited, big-endian encoding of a token sequence. + + Structurally the same scheme NeMo-RL's rollout writer uses for its staging + hashes, so the two sides can verify the same bytes the same way. + """ + encoded = bytearray(struct.pack(">BQ", DIGEST_VERSION, len(token_ids))) + for token_id in token_ids: + if token_id < 0: + raise ValueError(f"token ids must be non-negative, got {token_id}") + encoded.extend(struct.pack(">Q", token_id)) + return bytes(encoded) + + +def compute_digest(token_ids: list[int]) -> str: + """Digest of an exact token sequence. + + Used to verify a claimed ``parent_call_id`` in O(1) instead of comparing + whole arrays, and to detect a stale or interleaved store (a rerun that + appended onto a previous attempt's records fails the check rather than + merging silently). + """ + if not token_ids: + return _EMPTY_DIGEST + return hashlib.sha256(_DIGEST_DOMAIN + encode_token_ids(token_ids)).hexdigest() + + +class TokenEntry(BaseModel): + """One model call's captured record: the content-bearing output items (assistant + text, tool calls) together with the token fields, keyed to its rollout and to the + ``model_call_id`` the capture middleware minted for the call. + + ``output_items`` holds the served response's output items with their content, so a + trainer can read the text (e.g. NeMo-RL's invalid-tool-call / malformed-thinking + penalties) — token ids alone are not sufficient. The top-level token arrays are the + same fields carried on the generated item, kept here for the builder's chaining. + """ + + model_config = ConfigDict(extra="allow") + + rollout_id: str + model_call_id: str + model: str = "" + prompt_token_ids: list[int] + generation_token_ids: list[int] + generation_log_probs: list[float] + routed_experts: Optional[Any] = None + # The served response's output items (Responses shape), content preserved. + output_items: list[dict] = [] + # Non-semantic; a cheap diagnostic for retry/sibling-branch cases. + created_at: float = 0.0 + + # --- Lineage. Optional: null when the model server could not identify the + # parent, in which case the builder infers it from token prefixes instead. + # + # When present these make the chain exact rather than inferred. Two calls + # can share a prompt and differ only in their generation (a harness retry, + # since capture records a response the client may never have received); + # prefix inference has to guess between them, whereas the next call's parent + # link names the one the harness actually kept. They also turn the builder's + # O(N^2) prefix scan into an O(1) lookup, and let a claimed parent be + # verified rather than trusted. + parent_call_id: Optional[str] = None + # len(prompt_token_ids) + len(generation_token_ids) for THIS call: the + # length of the prefix a child of this call must start with. + cum_len: Optional[int] = None + # compute_digest(prompt_token_ids + generation_token_ids). + digest: Optional[str] = None + + +def cumulative_tokens(entry: TokenEntry) -> list[int]: + """The full sequence a child of this call must start with.""" + return list(entry.prompt_token_ids) + list(entry.generation_token_ids) + + +def stamp_lineage(entry: TokenEntry, parent_call_id: Optional[str]) -> TokenEntry: + """Fill ``cum_len`` and ``digest`` (always) and ``parent_call_id`` (when known). + + ``cum_len``/``digest`` describe this call and are always computable; the + parent link is only known when the model server resolved one. + """ + cumulative = cumulative_tokens(entry) + entry.cum_len = len(cumulative) + entry.digest = compute_digest(cumulative) + entry.parent_call_id = parent_call_id + return entry + + +def response_to_output_items(payload: dict) -> list[dict]: + """Normalize a served response to a list of content-bearing Responses output items. + + Responses payloads already carry ``output``. Chat payloads carry + ``choices[*].message``; the assistant message is wrapped as a single Responses + ``message`` item so the training record is dialect-uniform. + """ + output = payload.get("output") + if isinstance(output, list) and output: + return [item for item in output if isinstance(item, dict)] + items: list[dict] = [] + for choice in payload.get("choices") or []: + message = (choice or {}).get("message") or {} + if not isinstance(message, dict): + continue + item = dict(message) + item.setdefault("type", "message") + item.setdefault("role", "assistant") + items.append(item) + return items + + +def extract_token_fields(response_json: dict) -> Optional[dict]: + """Pull the token-id fields off a served response, or ``None`` if absent. + + Handles both shapes a Gym model server can return: a Responses-style + ``output`` list (the fields ride the last output item that carries them) and + a chat-completions ``choices[*].message``. Returns ``None`` when no item + carries token ids (e.g. token-id return is off, or an empty completion). + """ + candidates: list[dict] = [] + for item in response_json.get("output") or []: + if isinstance(item, dict) and item.get("generation_token_ids") is not None: + candidates.append(item) + for choice in response_json.get("choices") or []: + message = (choice or {}).get("message") or {} + if isinstance(message, dict) and message.get("generation_token_ids") is not None: + candidates.append(message) + if not candidates: + return None + source = candidates[-1] + return {field: source.get(field) for field in TOKEN_FIELDS} diff --git a/nemo_gym/token_id_capture/routes.py b/nemo_gym/token_id_capture/routes.py new file mode 100644 index 0000000000..dddcc400ce --- /dev/null +++ b/nemo_gym/token_id_capture/routes.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The model server's read route for captured training tokens. + +The route is registered on a model server only when training-token capture is +enabled, so a default run exposes nothing. It lets a non-co-located trainer pull +a rollout's ``TokenEntry`` records over HTTP (see ``HttpTokenReader``) instead of +requiring shared-filesystem access to the store. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Optional + +from fastapi import APIRouter, HTTPException, Response + +from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +def make_token_store(global_config_dict: Any) -> Optional[TokenCaptureStore]: + """Build the training-token store, or ``None`` when capture is disabled.""" + config = TokenIdCaptureConfig.model_validate(global_config_dict) + if not config.token_id_capture_enabled: + return None + return TokenCaptureStore(config.resolved_dir()) + + +def install_token_capture_routes(app: Any, store: TokenCaptureStore) -> None: + router = APIRouter() + + @router.get("/ng-capture/tokens/{rollout_id}") + async def get_tokens(rollout_id: str) -> Response: + try: + entries = await asyncio.to_thread(store.read_entries, rollout_id) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + body = "\n".join(entry.model_dump_json() for entry in entries) + return Response(content=body, media_type="application/x-ndjson") + + app.include_router(router) diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py new file mode 100644 index 0000000000..f4978145a2 --- /dev/null +++ b/nemo_gym/token_id_capture/sink.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Served-layer token capture for one model call. + +Token ids are dropped on the wire for streaming responses (Anthropic +``/v1/messages``, OpenAI chat SSE), so the capture middleware -- which only sees +the streamed bytes -- cannot record them. But the model server holds the +complete response WITH token ids for a moment, just before it synthesizes the +SSE stream. The middleware therefore hands the model server a per-request "token +sink" through a request-scoped ContextVar; the server calls ``capture_tokens`` +on its complete response and the sink writes a ``TokenEntry``. + +The sink carries the ``model_call_id`` the middleware minted for the same call, +so a captured ``TokenEntry`` joins its ``ModelCallRecord``. Only the middleware +sets a sink (for rollout-correlated, observed calls), so ordinary untagged +traffic captures nothing. +""" + +from __future__ import annotations + +import logging +import time +from contextvars import ContextVar, Token +from dataclasses import dataclass +from typing import Any, Optional + +from nemo_gym.token_id_capture.protocols import TokenSink +from nemo_gym.token_id_capture.records import ( + TokenEntry, + extract_token_fields, + response_to_output_items, + stamp_lineage, +) +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +logger = logging.getLogger(__name__) + + +@dataclass +class CaptureContext: + """What the capture middleware hands the model server for one call: which + rollout and call this is, and where the record goes. + + ``store`` is Gym's file store today; anything satisfying ``TokenSink`` works, + which is how a training framework redirects the write to its own data plane + without changing the capture code. + """ + + rollout_id: str + model_call_id: str + store: TokenCaptureStore + model: str = "" + + @property + def sink(self) -> TokenSink: + return self.store + + +# Kept so existing call sites and tests that refer to the request-scoped context +# by its old name keep working; ``TokenSink`` is now the write protocol. +TokenCaptureContext = CaptureContext + + +_TOKEN_SINK: ContextVar[Optional[CaptureContext]] = ContextVar("nemo_gym_token_sink", default=None) + + +def set_token_sink(sink: CaptureContext) -> Token: + return _TOKEN_SINK.set(sink) + + +def reset_token_sink(token: Token) -> None: + _TOKEN_SINK.reset(token) + + +async def capture_tokens(response: Any, parent_call_id: Optional[str] = None) -> None: + """Record a ``TokenEntry`` from a complete model response when a sink is set. + + ``response`` is a served response as a pydantic model or dict. No-op when no + sink is active (untagged traffic) or the response carries no token ids. The + write is offloaded and awaited, so the entry is durable before the model call + returns -- a post-rollout reader always sees it, with no background writer to + drain. + """ + sink = _TOKEN_SINK.get() + if sink is None: + return + if hasattr(response, "model_dump"): + payload = response.model_dump() + elif isinstance(response, dict): + payload = response + else: + return + info = extract_token_fields(payload) + if info is None: + return + try: + entry = TokenEntry( + rollout_id=sink.rollout_id, + model_call_id=sink.model_call_id, + model=sink.model or str(payload.get("model") or ""), + prompt_token_ids=info.get("prompt_token_ids") or [], + generation_token_ids=info.get("generation_token_ids") or [], + generation_log_probs=info.get("generation_log_probs") or [], + routed_experts=info.get("routed_experts"), + # Keep the content (assistant text, tool calls) so the trajectory the trainer + # reads is not token-only -- text-based penalties need it. + output_items=response_to_output_items(payload), + created_at=time.time(), + ) + # cum_len/digest describe this call and are always computable; the parent + # link is filled only when the model server resolved one. + stamp_lineage(entry, parent_call_id) + await sink.sink.put(entry) + except Exception: + # Capture is best-effort per call: a bad token payload must never fail the + # model call and break the harness's run. But a rollout that lost a call + # must not look identical to a complete one, so mark it -- delivery reads + # the marker and masks the sample rather than training on a hole. + logger.warning( + "Training-token capture failed for model call %s of rollout %s.", + sink.model_call_id, + sink.rollout_id, + exc_info=True, + ) + try: + sink.store.mark_incomplete(sink.rollout_id, sink.model_call_id) + except Exception: + logger.warning("Could not mark rollout %s incomplete.", sink.rollout_id, exc_info=True) diff --git a/nemo_gym/token_id_capture/source.py b/nemo_gym/token_id_capture/source.py new file mode 100644 index 0000000000..fd9fb4243d --- /dev/null +++ b/nemo_gym/token_id_capture/source.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The seam a trajectory builder reads training tokens through. + +``TokenSource`` is the only interface a trajectory builder needs: given a +rollout id, return its ``TokenEntry`` records. This decouples "where the tokens +came from" from "how the trajectory is assembled". ``CaptureTokenSource`` is the +implementation backed by Gym's capture store (through a local or HTTP reader). +Alternative sources -- e.g. records staged by a training framework's own +transport -- can implement the same protocol without changing the builder. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from nemo_gym.token_id_capture.protocols import TokenSource +from nemo_gym.token_id_capture.records import TokenEntry + + +if TYPE_CHECKING: # pragma: no cover - the HTTP reader pulls in Gym's server stack + from nemo_gym.token_id_capture.reader import TokenReader + + +__all__ = ["TokenSource", "CaptureTokenSource"] + + +class CaptureTokenSource: + """A ``TokenSource`` backed by Gym's capture store via a ``TokenReader``.""" + + def __init__(self, reader: "TokenReader") -> None: + self._reader = reader + + async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: + return await self._reader.read(rollout_id) + + async def drop(self, rollout_id: str) -> None: + """Delete-on-consume, when the underlying reader can retire records. + + The HTTP reader is read-only against a remote store, so this is a no-op + there and the owning process retires the rollout itself. + """ + drop = getattr(self._reader, "drop", None) + if drop is not None: + await drop(rollout_id) diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py new file mode 100644 index 0000000000..9f31b73797 --- /dev/null +++ b/nemo_gym/token_id_capture/store.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Append-only, rollout-keyed store for training ``TokenEntry`` records. + +One file per rollout (``.tokens.jsonl``), separate from the +evaluation capture file (``.capture.jsonl``) so token payloads never +bloat eval reads. Each write fsyncs and holds a per-file ``flock`` (which +excludes other threads and worker processes writing the *same* rollout file), +because a killed box must not lose a rollout's training tokens. + +Concurrency is per file, not global: there is deliberately no process-wide lock. +Every model call appends to its own rollout's file, so a global lock would +serialize all of them behind one fsync -- on a shared/network filesystem that +collapses throughput to ~1/fsync-latency regardless of core count. The per-file +flock keeps concurrent writers to one rollout correct while letting writes to +different rollouts proceed in parallel. +""" + +from __future__ import annotations + +import asyncio +import fcntl +import os +from pathlib import Path + +import orjson + +from nemo_gym.token_id_capture.records import TokenEntry + + +def validate_rollout_id(rollout_id: str) -> str: + """Reject anything that could escape the store directory or index a bad file.""" + if not rollout_id or any(not (char.isascii() and (char.isalnum() or char in "._-")) for char in rollout_id): + raise ValueError(f"Invalid rollout id: {rollout_id!r}") + return rollout_id + + +class TokenCaptureStore: + """Durable, rollout-keyed JSONL sink for ``TokenEntry`` records.""" + + def __init__(self, root: str | Path) -> None: + self._root = Path(root) + self._root.mkdir(parents=True, exist_ok=True) + + @property + def root(self) -> Path: + return self._root + + def path_for(self, rollout_id: str) -> Path: + return self._root / f"{validate_rollout_id(rollout_id)}.tokens.jsonl" + + def incomplete_path_for(self, rollout_id: str) -> Path: + """Sentinel marking that at least one call of this rollout failed to capture.""" + return self._root / f"{validate_rollout_id(rollout_id)}.tokens.incomplete" + + def mark_incomplete(self, rollout_id: str, reason: str = "") -> None: + """Record that a call was lost. + + Capture is best-effort per call -- a bad payload must never break the + harness -- but a rollout that captured 9 of 10 calls must not be + indistinguishable from a complete one. The marker is a file rather than + an in-process counter because the writer (model server) and the reader + (rollout collection, or the trainer) are different processes. + """ + try: + with self.incomplete_path_for(rollout_id).open("a") as handle: + handle.write(f"{reason}\n") + except OSError: + # Never let bookkeeping about a failed capture cause another failure. + pass + + def is_incomplete(self, rollout_id: str) -> bool: + return self.incomplete_path_for(rollout_id).exists() + + def append(self, entry: TokenEntry) -> None: + """Append one entry and fsync. Blocking file IO -- callers on the event + loop must offload it (e.g. ``asyncio.to_thread``).""" + line = orjson.dumps(entry.model_dump(), option=orjson.OPT_APPEND_NEWLINE) + path = self.path_for(entry.rollout_id) + with path.open("ab") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX) + try: + handle.write(line) + handle.flush() + os.fsync(handle.fileno()) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + # --- TokenSink / TokenSource. The file store is Gym's default implementation + # of both seams; a framework swaps in its own without touching capture. + + async def put(self, entry: TokenEntry) -> None: + """``TokenSink``: durable on return. The blocking append is offloaded so + it does not sit on the event loop, and awaited so a reader after the + rollout never races a partial file.""" + await asyncio.to_thread(self.append, entry) + + async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: + """``TokenSource``.""" + return await asyncio.to_thread(self.read_entries, rollout_id) + + async def drop(self, rollout_id: str) -> None: + """``TokenSource``: delete-on-consume.""" + await asyncio.to_thread(self.delete, rollout_id) + + def delete(self, rollout_id: str) -> None: + """Remove a rollout's records and its incomplete marker. + + Records are large (hundreds of KB per rollout) and the append opens in + "ab" mode, so leaving a consumed file behind both grows the directory + without bound and lets a rerun that reuses the id append onto stale + records. + """ + self.path_for(rollout_id).unlink(missing_ok=True) + self.incomplete_path_for(rollout_id).unlink(missing_ok=True) + + def read_entries(self, rollout_id: str) -> list[TokenEntry]: + path = self.path_for(rollout_id) + if not path.exists(): + return [] + entries: list[TokenEntry] = [] + with path.open("rb") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH) + try: + for line in handle: + stripped = line.strip() + if stripped: + entries.append(TokenEntry.model_validate(orjson.loads(stripped))) + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + return entries diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py new file mode 100644 index 0000000000..d021d490c3 --- /dev/null +++ b/tests/unit_tests/test_token_id_capture.py @@ -0,0 +1,472 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Training-token capture: schema, store, readers, source, and the served path. + +The served-path tests build a real ``SimpleResponsesAPIModel`` so the full chain runs: +the capture middleware mints a ``model_call_id`` and sets a per-request token sink, the +model server records a ``TokenEntry`` from its complete response, and the entry is read +back through the store, the HTTP route, and a ``TokenSource``. +""" + +import asyncio +import subprocess +import sys +from time import time +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from fastapi import Body, Request +from fastapi.testclient import TestClient + +from nemo_gym.base_responses_api_model import ( + BaseResponsesAPIModelConfig, + CaptureStore, + SimpleResponsesAPIModel, + read_model_call_records, +) +from nemo_gym.openai_utils import ( + NeMoGymChatCompletion, + NeMoGymChatCompletionCreateParamsNonStreaming, + NeMoGymResponse, + NeMoGymResponseCreateParamsNonStreaming, +) +from nemo_gym.server_utils import ServerClient +from nemo_gym.token_id_capture import ( + CaptureTokenSource, + TokenCaptureStore, + TokenEntry, + TokenIdCaptureConfig, + compute_digest, + cumulative_tokens, + extract_token_fields, + stamp_lineage, +) +from nemo_gym.token_id_capture import reader as reader_module +from nemo_gym.token_id_capture.reader import HttpTokenReader, LocalTokenReader +from nemo_gym.token_id_capture.routes import make_token_store + + +PTOKS = [1, 2, 3] +GTOKS = [4, 5] +LPS = [-0.1, -0.2] + + +# --- schema / extractor ------------------------------------------------------- + + +def test_extract_token_fields_responses_shape(): + payload = { + "output": [ + {"type": "message", "prompt_token_ids": PTOKS, "generation_token_ids": GTOKS, "generation_log_probs": LPS} + ] + } + assert extract_token_fields(payload) == { + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + "routed_experts": None, + } + + +def test_extract_token_fields_chat_shape(): + payload = { + "choices": [ + {"message": {"prompt_token_ids": [1], "generation_token_ids": [7], "generation_log_probs": [-0.3]}} + ] + } + got = extract_token_fields(payload) + assert got["generation_token_ids"] == [7] and got["prompt_token_ids"] == [1] + + +def test_extract_token_fields_absent_returns_none(): + assert extract_token_fields({"output": [{"type": "message"}]}) is None + assert extract_token_fields({}) is None + + +# --- store -------------------------------------------------------------------- + + +def test_token_store_round_trip(tmp_path): + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="t0-r0", + model_call_id="abc", + model="m", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + store.append(entry) + store.append(entry.model_copy(update={"model_call_id": "def"})) + read = store.read_entries("t0-r0") + assert [e.model_call_id for e in read] == ["abc", "def"] + assert read[0].prompt_token_ids == PTOKS + assert store.read_entries("missing") == [] + + +@pytest.mark.parametrize("bad", ["", "a/b", "../x", "a b"]) +def test_token_store_rejects_unsafe_rollout_ids(tmp_path, bad): + with pytest.raises(ValueError): + TokenCaptureStore(tmp_path).path_for(bad) + + +# --- config ------------------------------------------------------------------- + + +def test_config_disabled_needs_no_dir(): + cfg = TokenIdCaptureConfig.model_validate({}) + assert cfg.token_id_capture_enabled is False + assert make_token_store({}) is None + + +def test_config_enabled_requires_absolute_dir(tmp_path): + with pytest.raises(ValueError): + TokenIdCaptureConfig(token_id_capture_enabled=True) + with pytest.raises(ValueError): + TokenIdCaptureConfig(token_id_capture_enabled=True, token_id_capture_dir="relative/dir") + cfg = TokenIdCaptureConfig(token_id_capture_enabled=True, token_id_capture_dir=str(tmp_path)) + assert cfg.resolved_dir() == tmp_path + + +def test_config_falls_back_to_model_call_capture_dir(tmp_path): + cfg = TokenIdCaptureConfig(token_id_capture_enabled=True, model_call_capture_dir=str(tmp_path)) + assert cfg.resolved_dir() == tmp_path + + +# --- source / readers --------------------------------------------------------- + + +def test_capture_token_source_over_local_reader(tmp_path): + store = TokenCaptureStore(tmp_path) + store.append( + TokenEntry( + rollout_id="r", + model_call_id="c", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + ) + source = CaptureTokenSource(LocalTokenReader(store)) + entries = asyncio.run(source.tokens_for("r")) + assert len(entries) == 1 and entries[0].generation_token_ids == GTOKS + + +def test_http_token_reader_parses_ndjson(monkeypatch): + entry = TokenEntry( + rollout_id="r", model_call_id="c", prompt_token_ids=PTOKS, generation_token_ids=GTOKS, generation_log_probs=LPS + ) + body = entry.model_dump_json() + "\n" + + class _FakeResp: + async def text(self): + return body + + async def _fake_request(method, url, **kwargs): + assert method == "GET" and url.endswith("/ng-capture/tokens/r") + return _FakeResp() + + async def _fake_raise(_resp): + return None + + monkeypatch.setattr(reader_module, "request", _fake_request) + monkeypatch.setattr(reader_module, "raise_for_status", _fake_raise) + entries = asyncio.run(HttpTokenReader("http://model:9000").read("r")) + assert len(entries) == 1 and entries[0].model_call_id == "c" + + +# --- served path (full model server) ----------------------------------------- + + +def _training_response(text: str, model: str = "downstream-model") -> NeMoGymResponse: + return NeMoGymResponse( + id=f"resp_{uuid4().hex}", + created_at=int(time()), + model=model, + object="response", + output=[ + { + "type": "message", + "id": f"msg_{uuid4().hex}", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + } + ], + tool_choice="auto", + parallel_tool_calls=True, + tools=[], + ) + + +def _training_chat_completion(model: str = "downstream-model") -> NeMoGymChatCompletion: + return NeMoGymChatCompletion.model_validate( + { + "id": f"chatcmpl_{uuid4().hex}", + "created": int(time()), + "model": model, + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "hi", + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + }, + } + ], + } + ) + + +class _CapturingModel(SimpleResponsesAPIModel): + config: BaseResponsesAPIModelConfig + model_config = {"arbitrary_types_allowed": True} + + async def responses( + self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body() + ) -> NeMoGymResponse: + return _training_response("hi from responses") + + async def chat_completions( + self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body() + ) -> NeMoGymChatCompletion: + return _training_chat_completion() + + +def _server(global_config_dict) -> SimpleResponsesAPIModel: + return _CapturingModel( + config=BaseResponsesAPIModelConfig(host="0.0.0.0", port=8099, entrypoint="", name="srv"), + server_client=MagicMock(spec=ServerClient, global_config_dict=global_config_dict), + ) + + +def _both_enabled(tmp_path) -> dict: + return { + "observability_enabled": True, + "model_call_capture_dir": str(tmp_path), + "token_id_capture_enabled": True, + "token_id_capture_dir": str(tmp_path), + } + + +def test_responses_call_captures_tokens_joined_to_eval_record(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post("/ng-rollout/task0-roll0/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + + tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll0") + assert len(tokens) == 1 + assert tokens[0].generation_token_ids == GTOKS and tokens[0].prompt_token_ids == PTOKS + + records = read_model_call_records(CaptureStore(tmp_path), "task0-roll0") + assert len(records) == 1 + # The training entry joins its eval record by the middleware-minted model_call_id. + assert tokens[0].model_call_id == records[0].model_call_id + + +def test_captured_entry_carries_content(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/task0-rollC/v1/responses", json={"input": "hi"}) + tokens = TokenCaptureStore(tmp_path).read_entries("task0-rollC") + assert len(tokens) == 1 + # Not token-only: the captured record carries the content-bearing output items. + assert tokens[0].output_items + text = tokens[0].output_items[-1]["content"][0]["text"] + assert text == "hi from responses" + + +def test_messages_call_captures_tokens(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post( + "/ng-rollout/task0-roll1/v1/messages", + json={"model": "claude-x", "max_tokens": 16, "messages": [{"role": "user", "content": "hello"}]}, + ) + assert resp.status_code == 200 + # The Anthropic response on the wire never carries token ids. + assert "generation_token_ids" not in resp.text + tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll1") + assert len(tokens) == 1 and tokens[0].generation_token_ids == GTOKS + + +def test_chat_completions_call_captures_tokens(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post( + "/ng-rollout/task0-roll2/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]} + ) + assert resp.status_code == 200 + tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll2") + assert len(tokens) == 1 and tokens[0].generation_token_ids == GTOKS + + +def test_tokens_captured_even_when_eval_capture_disabled(tmp_path): + config = {"token_id_capture_enabled": True, "token_id_capture_dir": str(tmp_path)} + client = TestClient(_server(config).setup_webserver()) + resp = client.post("/ng-rollout/task1-roll0/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + assert len(TokenCaptureStore(tmp_path).read_entries("task1-roll0")) == 1 + # No eval capture file was written. + assert read_model_call_records(CaptureStore(tmp_path), "task1-roll0") == [] + + +def test_uncorrelated_call_captures_nothing(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post("/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + # No rollout prefix -> nothing recorded, no file created. + assert list(tmp_path.glob("*.tokens.jsonl")) == [] + + +def test_http_route_returns_tokens_and_404_when_disabled(tmp_path): + # Enabled: the route serves the captured entries as ndjson. + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/task2-roll0/v1/responses", json={"input": "hi"}) + got = client.get("/ng-capture/tokens/task2-roll0") + assert got.status_code == 200 + lines = [line for line in got.text.splitlines() if line.strip()] + assert len(lines) == 1 + parsed = TokenEntry.model_validate_json(lines[0]) + assert parsed.generation_token_ids == GTOKS + + # Disabled: the route is not registered. + disabled = TestClient(_server({}).setup_webserver()) + assert disabled.get("/ng-capture/tokens/task2-roll0").status_code == 404 + + +def test_package_is_dependency_free_leaf(): + """``nemo_gym.token_id_capture`` must import without Gym's server stack. + + A training framework's inference worker imports the record, the protocols, + and the capture core so it can write into its own data plane (see + ``protocols.py``). If the package drags in ray/fastapi/uvicorn, that is not + possible. Run in a subprocess so this test is unaffected by whatever the + rest of the suite has already imported. + """ + heavy = ("ray", "fastapi", "uvicorn", "aiohttp", "requests", "torch") + program = ( + f"import sys; import nemo_gym.token_id_capture; print(','.join(m for m in {heavy!r} if m in sys.modules))" + ) + proc = subprocess.run([sys.executable, "-c", program], capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + assert proc.stdout.strip() == "", f"leaf package pulled in: {proc.stdout.strip()}" + + +def test_streamed_messages_capture_tokens_absent_from_the_stream(tmp_path): + """The Claude Code shape: streamed /v1/messages. + + Token ids exist only on the assembled response, before it is converted to + Anthropic and split into SSE. This is the case the whole design turns on, so + it is asserted end to end rather than only through the non-streamed path. + """ + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + with client.stream( + "POST", + "/ng-rollout/stream0-roll0/v1/messages", + json={ + "model": "claude-x", + "max_tokens": 16, + "stream": True, + "messages": [{"role": "user", "content": "hello"}], + }, + ) as resp: + assert resp.status_code == 200 + body = "".join(resp.iter_text()) + # Nothing on the wire carries token ids. + assert "generation_token_ids" not in body + assert "prompt_token_ids" not in body + # ...yet the record is complete. + entries = TokenCaptureStore(tmp_path).read_entries("stream0-roll0") + assert len(entries) == 1 + assert entries[0].generation_token_ids == GTOKS + assert entries[0].prompt_token_ids == PTOKS + assert entries[0].output_items, "content must be captured alongside the tokens" + + +def test_capture_stamps_cum_len_and_digest(tmp_path): + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/lineage0-roll0/v1/responses", json={"input": "hi"}) + (entry,) = TokenCaptureStore(tmp_path).read_entries("lineage0-roll0") + assert entry.cum_len == len(PTOKS) + len(GTOKS) + assert entry.digest == compute_digest(PTOKS + GTOKS) + # No parent index yet, so the link is absent and the builder infers instead. + assert entry.parent_call_id is None + + +def test_digest_round_trip_and_stamp_lineage(): + entry = TokenEntry( + rollout_id="r", + model_call_id="c", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.5], + ) + stamp_lineage(entry, "parent-1") + assert cumulative_tokens(entry) == [1, 2, 3] + assert entry.cum_len == 3 + assert entry.parent_call_id == "parent-1" + assert entry.digest == compute_digest([1, 2, 3]) + # Distinct sequences must not collide, and the empty sequence is well defined. + assert compute_digest([1, 2, 3]) != compute_digest([1, 2, 4]) + assert compute_digest([]) == compute_digest([]) + with pytest.raises(ValueError): + compute_digest([-1]) + + +def test_capture_failure_marks_the_rollout_incomplete(tmp_path, monkeypatch): + """A lost call must not leave the rollout looking complete. + + Capture stays best-effort so a bad payload cannot break the harness's run, + but delivery has to be able to tell "10 of 10 captured" from "9 of 10". + """ + store = TokenCaptureStore(tmp_path) + + async def boom(self, entry): + raise RuntimeError("sink is down") + + monkeypatch.setattr(TokenCaptureStore, "put", boom) + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + # The model call still succeeds. + resp = client.post("/ng-rollout/fail0-roll0/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + assert store.read_entries("fail0-roll0") == [] + assert store.is_incomplete("fail0-roll0") + + +def test_delete_removes_records_and_marker(tmp_path): + store = TokenCaptureStore(tmp_path) + store.append( + TokenEntry( + rollout_id="gone-0", + model_call_id="c", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + ) + store.mark_incomplete("gone-0", "c") + assert store.path_for("gone-0").exists() and store.is_incomplete("gone-0") + store.delete("gone-0") + assert not store.path_for("gone-0").exists() + assert not store.is_incomplete("gone-0") + # Idempotent: consuming a rollout twice must not raise. + store.delete("gone-0") From 8e163583d69d2403b0f70aa7e00af182c8d4ee6e Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:32:13 -0700 Subject: [PATCH 2/9] feat(token-id-capture): build trajectories from captured token records 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 --- nemo_gym/token_id_capture/__init__.py | 35 +- nemo_gym/token_id_capture/builder.py | 446 ++++++++++++++++++ nemo_gym/token_id_capture/consumer.py | 167 +++++++ .../data/rg_ext_train_metrics.json | 47 ++ .../claude_code_agent/predicate.sh | 0 tests/unit_tests/test_trajectory_builder.py | 406 ++++++++++++++++ 6 files changed, 1096 insertions(+), 5 deletions(-) create mode 100644 nemo_gym/token_id_capture/builder.py create mode 100644 nemo_gym/token_id_capture/consumer.py create mode 100644 resources_servers/reasoning_gym/data/rg_ext_train_metrics.json create mode 100755 responses_api_agents/claude_code_agent/predicate.sh create mode 100644 tests/unit_tests/test_trajectory_builder.py diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index f7cbdd925f..2c3ed56793 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -16,9 +16,9 @@ """Training-token capture: produce, store, read, and source ``TokenEntry`` records. This is the per-model-call training data path, kept separate from evaluation -capture. The capture middleware sets a per-request capture context; the model -server records a ``TokenEntry`` from its complete response; a trainer reads a -rollout's entries through a ``TokenSource``. +capture. The capture middleware sets a per-request token sink; the model server +records a ``TokenEntry`` from its complete response; a trainer reads a rollout's +entries through a ``TokenSource`` and stitches them into a trajectory. **This package is a leaf.** Importing it must not pull in fastapi, ray, uvicorn, aiohttp, requests, or torch, because a training framework's inference worker @@ -26,11 +26,25 @@ data plane (see ``protocols.py``). The HTTP read route and the HTTP reader do need Gym's server stack, so they are deliberately *not* re-exported here -- import ``nemo_gym.token_id_capture.routes`` / ``.reader`` directly from server -code. ``tests/unit_tests/test_token_id_capture.py::test_package_is_dependency_free_leaf`` -enforces this. +code. """ +from nemo_gym.token_id_capture.builder import ( + BuildOutput, + Chain, + Trajectory, + assert_nemo_rl_contiguity, + build_trajectories, + per_request, + prefix_merging, + project_main_chain_response, +) from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.consumer import ( + token_id_capture_dirs_from_config, + trajectories_for_rollout, + trajectories_from_source, +) from nemo_gym.token_id_capture.protocols import ( TokenSink, TokenSource, @@ -78,4 +92,15 @@ "reset_token_sink", "capture_tokens", "CaptureTokenSource", + "build_trajectories", + "per_request", + "prefix_merging", + "project_main_chain_response", + "assert_nemo_rl_contiguity", + "Trajectory", + "Chain", + "BuildOutput", + "trajectories_for_rollout", + "trajectories_from_source", + "token_id_capture_dirs_from_config", ] diff --git a/nemo_gym/token_id_capture/builder.py b/nemo_gym/token_id_capture/builder.py new file mode 100644 index 0000000000..e960fa6e78 --- /dev/null +++ b/nemo_gym/token_id_capture/builder.py @@ -0,0 +1,446 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Turn a rollout's captured token records into trainable trajectories. + +The builder is a pure function over a list of ``TokenEntry`` records (whatever a +``TokenSource`` returns). It has two strategies: + + per_request assumes nothing about how the calls relate; every call becomes + its own training sequence. Always valid. + prefix_merging chains calls by the token-prefix relationship: each call is + parented to the earlier call whose full token sequence (prompt + plus generation) is the longest prefix of this call's prompt. + This rebuilds a multi-turn, append-only rollout into one chain. + A prompt that no longer extends any earlier call starts a new + root (a compacted or rewritten context). Two candidate parents + with identical sequences are ambiguous, so that subtree is + quarantined rather than guessed. + +Both are order-independent: they do not depend on arrival order or any sequence +number. ``prefix_merging`` processes entries by increasing prompt length, which +is derived from the tokens themselves (a parent's prompt is shorter than its +child's), so concurrent or out-of-order capture yields the same result. + +Loss masks follow provenance: tokens the policy generated are marked 1 (with +their captured log probabilities), and everything re-fed into a prompt (history, +tool output, tokens added between calls) is marked 0. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Optional + +from pydantic import BaseModel, Field + +from nemo_gym.token_id_capture.records import TokenEntry, compute_digest + + +@dataclass +class ChainLink: + entry: TokenEntry + interstitial: list[int] # prompt tokens added since the parent (tool output, new user turn); mask 0 + + +@dataclass +class Chain: + chain_id: str + links: list[ChainLink] = field(default_factory=list) + root_prompt: list[int] = field(default_factory=list) + + def flatten(self) -> tuple[list[int], list[int], list[Optional[float]], list[tuple[int, int, str]]]: + """Expand the chain into (token_ids, loss_mask, log_probs, spans). + + Generated tokens get mask 1 and their log probabilities; prompt/interstitial + tokens get mask 0 and no log probability. ``spans`` records, per generated + segment, ``(start, end, model_call_id)`` so a trainer can attach per-call + metadata (e.g. a weight version under async) to the tokens of each call.""" + ids: list[int] = list(self.root_prompt) + mask: list[int] = [0] * len(ids) + lps: list[Optional[float]] = [None] * len(ids) + spans: list[tuple[int, int, str]] = [] + for link in self.links: + ids += link.interstitial + mask += [0] * len(link.interstitial) + lps += [None] * len(link.interstitial) + gen = link.entry.generation_token_ids + glp = link.entry.generation_log_probs + if len(glp) != len(gen): + raise ValueError( + f"log-prob/token length mismatch on {link.entry.model_call_id}: {len(glp)} vs {len(gen)}" + ) + start = len(ids) + ids += gen + mask += [1] * len(gen) + lps += list(glp) + spans.append((start, len(ids), link.entry.model_call_id)) + return ids, mask, lps, spans + + +@dataclass +class BuildOutput: + chains: list[Chain] + quarantined: list[str] = field(default_factory=list) # model_call_ids + notes: dict = field(default_factory=dict) + + +class Trajectory(BaseModel): + """One trainable sequence: a flat token stream with a per-token loss mask and + the behavior-policy log probabilities at the generated positions.""" + + rollout_id: str + chain_id: str + token_ids: list[int] + loss_mask: list[int] + log_probs: list[Optional[float]] + # Per-call provenance: (start, end, model_call_id) for each generated span, so async training can + # attach a per-call weight version to the right tokens. + spans: list[tuple[int, int, str]] = Field(default_factory=list) + # The scalar reward drives single-objective training (GRPO). reward_components carries the + # named per-objective scores for multi-objective training (GDPO); it is None for single-reward + # environments. Both are per-response (per-rollout), copied from the verifier result; neither is + # an engine fact, so neither appears on TokenEntry. + reward: float = 0.0 + reward_components: Optional[dict[str, float]] = None + provenance: dict = Field(default_factory=dict) + + +def per_request(entries: list[TokenEntry]) -> BuildOutput: + ordered = sorted(entries, key=lambda e: (len(e.prompt_token_ids), e.model_call_id)) + chains = [ + Chain(chain_id=f"req-{i}", root_prompt=list(e.prompt_token_ids), links=[ChainLink(entry=e, interstitial=[])]) + for i, e in enumerate(ordered) + ] + return BuildOutput(chains=chains, notes={"builder": "per_request"}) + + +def _is_prefix(a: list[int], b: list[int]) -> bool: + return len(a) <= len(b) and b[: len(a)] == a + + +@dataclass(eq=False) # identity-based, so nodes are hashable for set membership +class _Node: + entry: TokenEntry + cumulative: list[int] # prompt + generation for this call + parent: Optional["_Node"] = None + children: list["_Node"] = field(default_factory=list) + quarantined: bool = False + + +def _resolve_parent( + node: "_Node", + by_call_id: dict[str, "_Node"], + candidates: list["_Node"], +) -> tuple[Optional["_Node"], bool, Optional[str]]: + """Find this call's parent, preferring the recorded link over inference. + + Returns ``(parent, ambiguous, note)``. ``note`` names why the recorded link + was not used, so a run can report how often it fell back instead of silently + behaving differently. + + The recorded link is verified rather than trusted: the child's prompt must + actually start with the parent's cumulative sequence, checked by digest so + it costs a hash of a prefix instead of comparing whole arrays. A stale store + -- a rerun that appended onto a previous attempt's records -- fails here and + falls back rather than merging two attempts into one trajectory. + """ + prompt = list(node.entry.prompt_token_ids) + claimed = node.entry.parent_call_id + if claimed is not None: + parent = by_call_id.get(claimed) + if parent is None: + return _infer_parent(prompt, candidates) + ("parent_call_id_missing",) + cum_len = parent.entry.cum_len + if cum_len is None: + cum_len = len(parent.cumulative) + if cum_len <= len(prompt) and compute_digest(prompt[:cum_len]) == ( + parent.entry.digest or compute_digest(parent.cumulative) + ): + return parent, False, None + return _infer_parent(prompt, candidates) + ("parent_digest_mismatch",) + return _infer_parent(prompt, candidates) + (None,) + + +def _infer_parent(prompt: list[int], candidates: list["_Node"]) -> tuple[Optional["_Node"], bool]: + """Fallback when no verified parent link is recorded: the earlier call whose + cumulative sequence is the longest prefix of this prompt. + + Two candidates with identical cumulative sequences are indistinguishable, so + the subtree is quarantined rather than guessed. + """ + matches = [n for n in candidates if _is_prefix(n.cumulative, prompt)] + if not matches: + return None, False + best_len = max(len(n.cumulative) for n in matches) + best = [n for n in matches if len(n.cumulative) == best_len] + return best[0], len(best) > 1 + + +def prefix_merging(entries: list[TokenEntry]) -> BuildOutput: + # Increasing prompt length is an order derived from the tokens: a parent's + # cumulative sequence is a prefix of its child's prompt, so the parent's + # prompt is shorter. This makes the pass order-independent. + ordered = sorted(entries, key=lambda e: (len(e.prompt_token_ids), e.model_call_id)) + nodes: list[_Node] = [] + roots: list[_Node] = [] + quarantined: list[str] = [] + + nodes_by_call_id: dict[str, _Node] = {} + fallbacks: dict[str, int] = {} + + for entry in ordered: + prompt = list(entry.prompt_token_ids) + node = _Node(entry=entry, cumulative=prompt + list(entry.generation_token_ids)) + parent, ambiguous, note = _resolve_parent(node, nodes_by_call_id, nodes) + if note: + fallbacks[note] = fallbacks.get(note, 0) + 1 + if parent is not None: + node.parent = parent + if ambiguous: + # Two candidate parents with identical sequences: quarantine rather than guess. + node.quarantined = True + quarantined.append(entry.model_call_id) + parent.children.append(node) + else: + roots.append(node) + nodes.append(node) + nodes_by_call_id[entry.model_call_id] = node + + # Resolve retry siblings. A harness (Claude Code) retries on timeout / 5xx / dropped SSE, and the + # capture point records a call even if the client never received it, so a retry yields two nodes + # with identical prompt ids under the same parent and divergent generations. + # + # A recorded parent link settles this exactly: a later call names the sibling the harness actually + # kept, so the other is provably unused. Without one we fall back to "the sibling a later call + # extended wins". Neither can resolve a retry of the FINAL call -- there is no later call to name + # the survivor -- so that case is flagged as unresolved rather than tie-broken silently, and the + # caller masks the rollout instead of training on a generation the client may never have received. + unresolved_retries: list[str] = [] + siblings_by_parent: dict[int, list[_Node]] = {} + for node in nodes: + siblings_by_parent.setdefault(id(node.parent), []).append(node) + for group in siblings_by_parent.values(): + by_prompt: dict[tuple, list[_Node]] = {} + for node in group: + by_prompt.setdefault(tuple(node.entry.prompt_token_ids), []).append(node) + for retry_group in by_prompt.values(): + if len(retry_group) < 2: + continue + extended = [n for n in retry_group if n.children] + if extended: + keep = set(extended) + else: + keep = {min(retry_group, key=lambda n: n.entry.model_call_id)} + unresolved_retries.extend(n.entry.model_call_id for n in retry_group) + for node in retry_group: + if node not in keep and not node.quarantined: + node.quarantined = True + quarantined.append(node.entry.model_call_id) + + chains: list[Chain] = [] + + def walk(node: _Node, path: list[_Node]) -> None: + path = path + [node] + if not node.children: + if any(p.quarantined for p in path): + return + root = path[0] + chain = Chain(chain_id="", root_prompt=list(root.entry.prompt_token_ids)) + prev_cumulative = list(root.entry.prompt_token_ids) + for step, p in enumerate(path): + interstitial = [] if step == 0 else list(p.entry.prompt_token_ids[len(prev_cumulative) :]) + chain.links.append(ChainLink(entry=p.entry, interstitial=interstitial)) + prev_cumulative = list(p.entry.prompt_token_ids) + list(p.entry.generation_token_ids) + chains.append(chain) + return + for child in node.children: + walk(child, path) + + for root in roots: + walk(root, []) + + # Pick the main chain by how many tokens the policy generated in it, across all roots. + # + # Selecting from "the first root" is unsafe: entries are processed in increasing prompt length, + # so the first root is whichever root has the shortest prompt. A rollout's own first call starts + # with the harness system prompt and tool definitions and is large, so any auxiliary call with a + # short prompt would be selected instead, and the rollout itself relabelled a branch and dropped + # at delivery. That failure is silent: the delivered response is still contiguous and + # token-bearing. + # + # Generated-token mass does not depend on prompt length and needs no harness-specific knowledge. + # Ties break on root prompt length then first call id, so the choice is deterministic. + def generated_tokens(c: Chain) -> int: + return sum(len(link.entry.generation_token_ids) for link in c.links) + + def selection_key(c: Chain) -> tuple: + return (generated_tokens(c), len(c.root_prompt), c.links[0].entry.model_call_id if c.links else "") + + if chains: + main = max(chains, key=selection_key) + main.chain_id = "main" + branch = 0 + for c in chains: + if c is not main: + c.chain_id = f"branch-{branch}" + branch += 1 + + delivered = generated_tokens(main) if chains else 0 + captured = sum(len(e.generation_token_ids) for e in entries) + notes = { + "builder": "prefix_merging", + "roots": len(roots), + "chains": len(chains), + # What did not make it into the delivered chain. NeMo-RL takes one trajectory today, so + # sub-agent branches and everything after a context compaction are dropped; that is a + # deliberate limitation, but it must be visible rather than silent. + "generated_tokens_captured": captured, + "generated_tokens_delivered": delivered, + "delivered_fraction": round(delivered / captured, 4) if captured else 0.0, + "unresolved_retries": unresolved_retries, + } + if fallbacks: + notes["parent_link_fallbacks"] = fallbacks + return BuildOutput(chains=chains, quarantined=quarantined, notes=notes) + + +_BUILDERS: dict[str, Callable[[list[TokenEntry]], BuildOutput]] = { + "per_request": per_request, + "prefix_merging": prefix_merging, +} + + +def build_trajectories( + rollout_id: str, + entries: list[TokenEntry], + builder: str = "prefix_merging", + reward: float = 0.0, + reward_components: Optional[dict[str, float]] = None, +) -> list[Trajectory]: + """Build the trainable trajectories for one rollout from its token records.""" + if builder not in _BUILDERS: + raise ValueError(f"unknown builder {builder!r}; known: {sorted(_BUILDERS)}") + if not entries: + return [] + out = _BUILDERS[builder](entries) + total_calls = len(entries) + quarantined_fraction = (len(out.quarantined) / total_calls) if total_calls else 0.0 + trajectories: list[Trajectory] = [] + for chain in out.chains: + ids, mask, lps, spans = chain.flatten() + trained = sum(mask) + trajectories.append( + Trajectory( + rollout_id=rollout_id, + chain_id=chain.chain_id, + token_ids=ids, + loss_mask=mask, + log_probs=lps, + spans=spans, + reward=reward, + reward_components=reward_components, + provenance={ + "builder": out.notes.get("builder", builder), + "n_calls": len(chain.links), + # Metrics that make silent training loss visible: how much of the rollout was + # dropped, and how much of this chain is actually trained on. A reasoning model in + # a multi-call rollout can quietly collapse to turn-1-only; these surface it. + "quarantined_calls": len(out.quarantined), + "quarantined_fraction": round(quarantined_fraction, 4), + "trained_token_fraction": round(trained / len(ids), 4) if ids else 0.0, + "notes": out.notes, + }, + ) + ) + return trajectories + + +# --- Projection to the shape NeMo-RL consumes --- + + +def project_chain_to_output_items(chain: Chain) -> list[dict]: + """Project the chain into content-bearing Responses output items whose prompts are + contiguous. For each call, emit its captured output items (assistant text, tool + calls preserved) and set the contiguous prompt on the item that carries the + generation, so each generated item's prompt extends the previous one — the shape + NeMo-RL ingests, with the text its penalties read (section 7.2). Falls back to a + synthesized token-only item only when a call captured no content items.""" + items: list[dict] = [] + cumulative = list(chain.root_prompt) + for step, link in enumerate(chain.links): + cumulative = cumulative + (link.interstitial if step > 0 else []) + entry = link.entry + content_items = [dict(item) for item in (entry.output_items or [])] + generated = [item for item in content_items if item.get("generation_token_ids") is not None] + if not generated and content_items: + # No item carried token fields (unexpected); attach to the last so tokens are not lost. + generated = content_items[-1:] + if content_items: + for item in generated: + item["prompt_token_ids"] = list(cumulative) + item["generation_token_ids"] = list(entry.generation_token_ids) + item["generation_log_probs"] = list(entry.generation_log_probs) + if entry.routed_experts is not None: + item["routed_experts"] = entry.routed_experts + items.extend(content_items) + else: + item = { + "type": "message", + "prompt_token_ids": list(cumulative), + "generation_token_ids": list(entry.generation_token_ids), + "generation_log_probs": list(entry.generation_log_probs), + } + if entry.routed_experts is not None: + item["routed_experts"] = entry.routed_experts + items.append(item) + cumulative = cumulative + list(entry.generation_token_ids) + return items + + +def project_main_chain_response(rollout_id: str, out: BuildOutput, model: str = "") -> dict: + """Project the main chain into a Responses-shaped object with contiguous output items.""" + mains = [c for c in out.chains if c.chain_id == "main"] or out.chains[:1] + output = project_chain_to_output_items(mains[0]) if mains else [] + # Token fields ride only on generated items; a content-only leading item (e.g. assistant + # text emitted before a tool call) carries none. Read the usage counts from the items that + # actually have token fields so a leading content item does not KeyError or skew the totals. + generated = [item for item in output if item.get("generation_token_ids") is not None] + n_in = len(generated[0]["prompt_token_ids"]) if generated else 0 + n_out = sum(len(item["generation_token_ids"]) for item in generated) + return { + "id": f"proj-{rollout_id}", + "model": model, + "object": "response", + "output": output, + "usage": {"input_tokens": n_in, "output_tokens": n_out}, + } + + +def assert_nemo_rl_contiguity(response: dict) -> None: + """Enforce the invariant NeMo-RL's ingestion relies on: each output item's + prompt_token_ids must extend the tokens seen so far (prompt plus generation + of all prior items). Raises AssertionError otherwise.""" + seen: list[int] = [] + for item in response.get("output", []): + if not isinstance(item, dict) or item.get("generation_token_ids") is None: + continue + prompt = item.get("prompt_token_ids") or [] + if prompt[: len(seen)] != seen: + raise AssertionError( + "projection violates NeMo-RL prefix contiguity: an output item's prompt_token_ids " + "does not extend the tokens seen so far" + ) + seen = list(prompt) + list(item["generation_token_ids"]) diff --git a/nemo_gym/token_id_capture/consumer.py b/nemo_gym/token_id_capture/consumer.py new file mode 100644 index 0000000000..e5cc01db95 --- /dev/null +++ b/nemo_gym/token_id_capture/consumer.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The consumer that turns a rollout's captured tokens into trajectories. + +This is the single primitive both consumers call after a rollout finishes: Gym's +rollout collection (co-located, reading the token store's files) and a trainer's +finalizer (which passes a ``TokenSource``, e.g. HTTP or TransferQueue-backed). +The only difference between them is where the ``TokenEntry`` records come from; +the build and projection are identical. + +It is deliberately free of any rollout-record or model-server imports, so it +does not couple to those layers. The caller supplies the ``rollout_id`` (Gym's +rollout collection derives it from the record's task/rollout/attempt indices) +and the reward. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Optional + +from nemo_gym.token_id_capture.builder import ( + assert_nemo_rl_contiguity, + build_trajectories, + per_request, + prefix_merging, + project_main_chain_response, +) +from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.source import TokenSource +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +logger = logging.getLogger(__name__) + + +def token_id_capture_dirs_from_config(global_config_dict) -> list[Path]: + """Resolve the token store directory when training-token capture is enabled, else [].""" + from nemo_gym.token_id_capture.config import TokenIdCaptureConfig + + config = TokenIdCaptureConfig.model_validate(global_config_dict) + directory = config.resolved_dir() + return [directory] if (config.token_id_capture_enabled and directory is not None) else [] + + +def _assemble( + rollout_id: str, + entries: list[TokenEntry], + builder: str, + reward: float, + reward_components: Optional[dict[str, float]], + model: str, +) -> dict: + # A malformed capture must degrade this one rollout, not take down the caller. + # Both the contiguity assertion and the flattener raise, and the callers are a + # rollout-collection loop and NeMo-RL's training loop -- where an escaping + # exception kills a whole step's batch rather than dropping one sample. + try: + out = prefix_merging(entries) if builder == "prefix_merging" else per_request(entries) + response = project_main_chain_response(rollout_id, out, model=model) + assert_nemo_rl_contiguity(response) + trajectories = build_trajectories( + rollout_id, entries, builder=builder, reward=reward, reward_components=reward_components + ) + except (AssertionError, ValueError, KeyError, IndexError, TypeError) as error: + logger.warning( + "Could not build a trajectory for rollout %s from %d captured call(s): %s", + rollout_id, + len(entries), + error, + ) + return { + "rollout_id": rollout_id, + "builder": builder, + "trajectories": [], + "nemo_rl_response": None, + "mask_sample": True, + "error": f"{type(error).__name__}: {error}", + "metrics": {"n_calls": len(entries)}, + } + + notes = dict(out.notes) + # Surface what the build dropped. These were previously computed and thrown away, so a rollout + # that trained on one of five calls looked exactly like one that trained on all five. + metrics = { + "n_calls": len(entries), + "chains": notes.get("chains", len(out.chains)), + "quarantined_calls": len(out.quarantined), + "quarantined_fraction": round(len(out.quarantined) / len(entries), 4) if entries else 0.0, + "delivered_fraction": notes.get("delivered_fraction", 0.0), + "generated_tokens_captured": notes.get("generated_tokens_captured", 0), + "generated_tokens_delivered": notes.get("generated_tokens_delivered", 0), + "parent_link_fallbacks": notes.get("parent_link_fallbacks", {}), + } + unresolved = notes.get("unresolved_retries") or [] + return { + "rollout_id": rollout_id, + "builder": builder, + "trajectories": [t.model_dump() for t in trajectories], + "nemo_rl_response": response, + "metrics": metrics, + # A retry of the final call leaves two generations with no way to tell which one the client + # received. Training on the wrong one is silently off-policy, so the rollout is masked. + "mask_sample": bool(unresolved), + "unresolved_retries": list(unresolved), + } + + +def trajectories_for_rollout( + rollout_id: str, + token_capture_dirs: list[Path], + *, + builder: str = "prefix_merging", + reward: float = 0.0, + reward_components: Optional[dict[str, float]] = None, + model: str = "", +) -> Optional[dict]: + """Co-located path: read the rollout's tokens from the store files and build its trajectories. + + ``reward`` (scalar, for GRPO) and ``reward_components`` (named per-objective scores, for GDPO) + come from the verifier result and ride the trajectory; they are not read from the token store. + Returns ``None`` when no tokens were captured for the rollout (capture off, or a dialect the + engine returned no ids for). Mirrors how evaluation capture is merged into a rollout record. + """ + for directory in token_capture_dirs: + store = TokenCaptureStore(directory) + entries = store.read_entries(rollout_id) + if entries: + built = _assemble(rollout_id, entries, builder, reward, reward_components, model) + if store.is_incomplete(rollout_id): + # At least one call of this rollout failed to capture. The chain we built may look + # perfectly contiguous while being missing a turn, so mask rather than train on it. + built["mask_sample"] = True + built.setdefault("metrics", {})["capture_incomplete"] = True + return built + return None + + +async def trajectories_from_source( + rollout_id: str, + source: TokenSource, + *, + builder: str = "prefix_merging", + reward: float = 0.0, + reward_components: Optional[dict[str, float]] = None, + model: str = "", +) -> Optional[dict]: + """Non-co-located path: read the rollout's tokens through a ``TokenSource`` (HTTP, or a + trainer's own transport) and build its trajectories. Returns ``None`` when none were captured.""" + entries = await source.tokens_for(rollout_id) + if not entries: + return None + return _assemble(rollout_id, entries, builder, reward, reward_components, model) diff --git a/resources_servers/reasoning_gym/data/rg_ext_train_metrics.json b/resources_servers/reasoning_gym/data/rg_ext_train_metrics.json new file mode 100644 index 0000000000..7e08f912ec --- /dev/null +++ b/resources_servers/reasoning_gym/data/rg_ext_train_metrics.json @@ -0,0 +1,47 @@ +{ + "name": "train", + "type": "train", + "jsonl_fpath": "resources_servers/reasoning_gym/data/rg_ext_train.jsonl", + "num_repeats": 1, + "source": null, + "gitlab_identifier": null, + "huggingface_identifier": null, + "license": "Creative Commons Attribution 4.0 International", + "Number of examples": 20, + "Number of tools": { + "Total # non-null values": 0, + "Average": 0.0, + "Min": 0.0, + "Max": 0.0, + "Standard deviation": 0.0 + }, + "Json-dumped number of words (proxy for token count)": { + "Total # non-null values": 20, + "Average": 46.45, + "Min": 42.0, + "Max": 53.0, + "Standard deviation": 4.29 + }, + "Number of turns": { + "Total # non-null values": 20, + "Average": 1.0, + "Min": 1.0, + "Max": 1.0, + "Standard deviation": 0.0 + }, + "Temperature": { + "Total # non-null values": 0, + "Average": 0.0, + "Min": 0.0, + "Max": 0.0, + "Standard deviation": 0.0 + }, + "question": { + "unique_count": 20, + "total_count": 20 + }, + "answer": { + "unique_count": 20, + "total_count": 20 + } +} \ No newline at end of file diff --git a/responses_api_agents/claude_code_agent/predicate.sh b/responses_api_agents/claude_code_agent/predicate.sh new file mode 100755 index 0000000000..e69de29bb2 diff --git a/tests/unit_tests/test_trajectory_builder.py b/tests/unit_tests/test_trajectory_builder.py new file mode 100644 index 0000000000..4d4862e680 --- /dev/null +++ b/tests/unit_tests/test_trajectory_builder.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Trajectory builder: chaining, loss masks, and the NeMo-RL projection.""" + +import pytest + +from nemo_gym.token_id_capture import ( + Trajectory, + assert_nemo_rl_contiguity, + build_trajectories, + compute_digest, + prefix_merging, + project_main_chain_response, + stamp_lineage, + token_id_capture_dirs_from_config, + trajectories_for_rollout, +) +from nemo_gym.token_id_capture.records import TokenEntry +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +def _entry(mcid, prompt, gen, lp=None): + return TokenEntry( + rollout_id="t0-r0", + model_call_id=mcid, + model="m", + prompt_token_ids=prompt, + generation_token_ids=gen, + generation_log_probs=lp if lp is not None else [-0.1] * len(gen), + ) + + +# An append-only 3-call rollout: each call's prompt extends the prior prompt+generation +# plus interstitial tokens (tool output / new user turn). +CALL1 = _entry("c1", [1, 2, 3], [10, 11]) +CALL2 = _entry("c2", [1, 2, 3, 10, 11, 4, 5], [12]) +CALL3 = _entry("c3", [1, 2, 3, 10, 11, 4, 5, 12, 6], [13, 14]) +APPEND_ONLY = [CALL1, CALL2, CALL3] + + +def test_prefix_merging_builds_one_contiguous_main_chain(): + trajs = build_trajectories("t0-r0", APPEND_ONLY, builder="prefix_merging", reward=1.0) + assert len(trajs) == 1 + t = trajs[0] + assert t.chain_id == "main" + # The flat stream is the final cumulative sequence. + assert t.token_ids == [1, 2, 3, 10, 11, 4, 5, 12, 6, 13, 14] + # Generated tokens are masked 1, everything re-fed to a prompt is masked 0. + assert t.loss_mask == [0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1] + # Log probs are present exactly at the generated positions. + assert [lp is not None for lp in t.log_probs] == [bool(m) for m in t.loss_mask] + assert t.reward == 1.0 + assert t.provenance["n_calls"] == 3 + + +def test_order_independent(): + import random + + shuffled = list(APPEND_ONLY) + random.Random(0).shuffle(shuffled) + a = build_trajectories("t0-r0", APPEND_ONLY, builder="prefix_merging")[0] + b = build_trajectories("t0-r0", shuffled, builder="prefix_merging")[0] + assert a.token_ids == b.token_ids and a.loss_mask == b.loss_mask + + +def test_per_request_marks_the_same_generated_tokens(): + # Both builders must agree on which tokens were generated (mask 1). + def generated(trajs: list[Trajectory]): + out = [] + for t in trajs: + out += [tid for tid, m in zip(t.token_ids, t.loss_mask) if m == 1] + return sorted(out) + + merged = build_trajectories("t0-r0", APPEND_ONLY, builder="prefix_merging") + per_req = build_trajectories("t0-r0", APPEND_ONLY, builder="per_request") + assert len(per_req) == 3 + assert generated(merged) == generated(per_req) == sorted([10, 11, 12, 13, 14]) + + +def test_projection_is_nemo_rl_contiguous(): + out = prefix_merging(APPEND_ONLY) + response = project_main_chain_response("t0-r0", out, model="m") + assert [len(i["prompt_token_ids"]) for i in response["output"]] == [3, 7, 9] + assert response["usage"] == {"input_tokens": 3, "output_tokens": 5} + assert_nemo_rl_contiguity(response) # must not raise + + +def test_contiguity_assert_catches_a_gap(): + broken = { + "output": [ + {"type": "message", "prompt_token_ids": [1, 2, 3], "generation_token_ids": [10]}, + # prompt does not extend [1,2,3,10]: + {"type": "message", "prompt_token_ids": [1, 2, 3, 99], "generation_token_ids": [11]}, + ] + } + with pytest.raises(AssertionError): + assert_nemo_rl_contiguity(broken) + + +def _content_entry(mcid, prompt, gen, text): + lp = [-0.1] * len(gen) + return TokenEntry( + rollout_id="t0-r0", + model_call_id=mcid, + model="m", + prompt_token_ids=prompt, + generation_token_ids=gen, + generation_log_probs=lp, + output_items=[ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text}], + "prompt_token_ids": prompt, + "generation_token_ids": gen, + "generation_log_probs": lp, + } + ], + ) + + +def test_projection_carries_content_and_stays_contiguous(): + entries = [ + _content_entry("c1", [1, 2, 3], [10, 11], "first turn"), + _content_entry("c2", [1, 2, 3, 10, 11, 4, 5], [12], "second turn"), + ] + out = prefix_merging(entries) + resp = project_main_chain_response("t0-r0", out, model="m") + texts = [item["content"][0]["text"] for item in resp["output"]] + assert texts == ["first turn", "second turn"] # content preserved (not token-only) + assert [len(i["prompt_token_ids"]) for i in resp["output"]] == [3, 7] + assert_nemo_rl_contiguity(resp) # prompts still contiguous with content attached + + +def test_projection_handles_content_only_leading_item(): + # A single call whose output is an assistant text message (no token fields) followed by a + # tool call that carries the token fields -- the real shape when a model narrates before a + # tool call. Usage must be read from the token-bearing item, not output[0]. + entry = TokenEntry( + rollout_id="t0-r0", + model_call_id="c1", + model="m", + prompt_token_ids=[1, 2, 3], + generation_token_ids=[10, 11], + generation_log_probs=[-0.1, -0.1], + output_items=[ + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "let me check"}]}, + {"type": "function_call", "name": "grep", "arguments": "{}", "call_id": "x"}, + ], + ) + out = prefix_merging([entry]) + resp = project_main_chain_response("t0-r0", out, model="m") + assert resp["output"][0]["type"] == "message" # content-only leading item preserved + assert "prompt_token_ids" not in resp["output"][0] + assert resp["usage"] == {"input_tokens": 3, "output_tokens": 2} # counts from the token-bearing item + assert_nemo_rl_contiguity(resp) + + +def test_retry_sibling_is_dropped_and_main_chain_is_deterministic(): + # c2a and c2b are a retry pair (identical prompt, divergent generation). c3 extends c2a. + c1 = _entry("c1", [1, 2, 3], [10, 11]) + c2a = _entry("c2a", [1, 2, 3, 10, 11, 4], [12]) + c2b = _entry("c2b", [1, 2, 3, 10, 11, 4], [99]) + c3 = _entry("c3", [1, 2, 3, 10, 11, 4, 12, 5], [13]) + out = prefix_merging([c1, c2a, c2b, c3]) + assert "c2b" in out.quarantined # unextended retry sibling dropped + main = next(c for c in out.chains if c.chain_id == "main") + assert [link.entry.model_call_id for link in main.links] == ["c1", "c2a", "c3"] + assert_nemo_rl_contiguity(project_main_chain_response("t0-r0", out)) + + +def test_spans_mark_each_generation(): + trajs = build_trajectories("t0-r0", APPEND_ONLY, builder="prefix_merging") + t = trajs[0] + # One span per call, each covering exactly the mask-1 (generated) positions. + assert [call for _, _, call in t.spans] == ["c1", "c2", "c3"] + for start, end, _ in t.spans: + assert all(t.loss_mask[i] == 1 for i in range(start, end)) + assert t.provenance["trained_token_fraction"] > 0 + + +def test_consumer_reads_store_and_builds(tmp_path): + # The co-located consumer: write the rollout's tokens, then build from the store files. + store = TokenCaptureStore(tmp_path) + for e in APPEND_ONLY: + store.append(e.model_copy(update={"rollout_id": "t0-r0"})) + dirs = token_id_capture_dirs_from_config({"token_id_capture_enabled": True, "token_id_capture_dir": str(tmp_path)}) + assert dirs == [tmp_path] + merged = trajectories_for_rollout("t0-r0", dirs, builder="prefix_merging", reward=1.0) + assert merged is not None + assert merged["builder"] == "prefix_merging" + assert len(merged["trajectories"]) == 1 + assert merged["trajectories"][0]["token_ids"] == [1, 2, 3, 10, 11, 4, 5, 12, 6, 13, 14] + assert len(merged["nemo_rl_response"]["output"]) == 3 + + +def test_reward_components_ride_the_trajectory(tmp_path): + # Multi-objective (GDPO): the scalar reward and the named components both ride the trajectory, + # copied from the verifier result. Token records never carry them. + components = {"correctness": 1.0, "integer": 1.0, "format": 0.0} + trajs = build_trajectories("t0-r0", APPEND_ONLY, reward=2.0, reward_components=components) + assert trajs[0].reward == 2.0 + assert trajs[0].reward_components == components + # Single-objective (GRPO) leaves components None, so the trainer path is unchanged. + assert build_trajectories("t0-r0", APPEND_ONLY, reward=1.0)[0].reward_components is None + + store = TokenCaptureStore(tmp_path) + for e in APPEND_ONLY: + store.append(e) + dirs = token_id_capture_dirs_from_config({"token_id_capture_enabled": True, "token_id_capture_dir": str(tmp_path)}) + merged = trajectories_for_rollout("t0-r0", dirs, reward=2.0, reward_components=components) + assert merged["trajectories"][0]["reward_components"] == components + + +def test_consumer_noop_when_disabled_or_absent(tmp_path): + assert token_id_capture_dirs_from_config({}) == [] + assert trajectories_for_rollout("t0-r0", []) is None + # Enabled dir but no file for this rollout -> None (graceful no-op). + dirs = token_id_capture_dirs_from_config({"token_id_capture_enabled": True, "token_id_capture_dir": str(tmp_path)}) + assert trajectories_for_rollout("missing", dirs) is None + + +def test_ambiguous_parents_are_quarantined(): + # Two roots with identical prompt+generation, then a call extending that shared + # sequence: its parent is ambiguous, so the subtree is quarantined, not guessed. + a = _entry("a", [1, 2], [7, 8]) + b = _entry("b", [1, 2], [7, 8]) + child = _entry("child", [1, 2, 7, 8, 9], [20]) + out = prefix_merging([a, b, child]) + assert "child" in out.quarantined + # The quarantined child is excluded from every emitted chain. + for chain in out.chains: + assert all(link.entry.model_call_id != "child" for link in chain.links) + + +# --- side calls and chain selection ------------------------------------------- + + +def test_a_short_side_call_does_not_replace_the_rollout(): + """A conversation-title call must not become the delivered chain. + + Claude Code generates a title (and probes quota) on a tiny prompt, while the + rollout's first real call carries the full system prompt and tool + definitions. Entries are processed by increasing prompt length, so the title + call is the first root; selecting the main chain from the first root would + deliver the title and relabel the whole rollout a branch. Nothing would + error -- the trainer would receive a contiguous, token-bearing response + containing a generated title, with the rollout's reward attached. + """ + title = _entry("title", [9000, 9001], [7, 7, 7]) + real_1 = _entry("real1", list(range(100, 160)), [200, 201, 202, 203]) + real_2 = _entry("real2", list(range(100, 160)) + [200, 201, 202, 203, 500], [300, 301, 302]) + + out = prefix_merging([title, real_1, real_2]) + main = next(c for c in out.chains if c.chain_id == "main") + + assert [link.entry.model_call_id for link in main.links] == ["real1", "real2"] + assert out.notes["chains"] == 2 + # The dropped chain is reported rather than silently discarded. + assert out.notes["generated_tokens_captured"] == 10 + assert out.notes["generated_tokens_delivered"] == 7 + assert out.notes["delivered_fraction"] == 0.7 + + +def test_post_compaction_chain_is_reported_as_dropped(): + """A rewritten context starts a new root. Only one chain is delivered today, + so what is left behind has to show up in the metrics.""" + call_1 = _entry("c1", [1, 2, 3], [4, 5]) + call_2 = _entry("c2", [1, 2, 3, 4, 5, 6], [7]) + # Compaction: the prompt no longer extends anything captured. + call_3 = _entry("c3", [90, 91], [92, 93, 94, 95]) + + out = prefix_merging([call_1, call_2, call_3]) + assert out.notes["chains"] == 2 + assert out.notes["generated_tokens_captured"] == 7 + assert out.notes["delivered_fraction"] < 1.0 + + +# --- recorded parent links ---------------------------------------------------- + + +def _with_lineage(entry, parent_call_id=None): + stamp_lineage(entry, parent_call_id) + return entry + + +def test_recorded_parent_link_resolves_a_final_call_retry_exactly(): + """Two siblings share a prompt and differ only in their generation. + + Prefix inference cannot tell which one the harness kept, because both are + equally valid children. A recorded parent link on the next call names the + survivor, so the other is provably unused rather than tie-broken. + """ + root = _with_lineage(_entry("root", [1, 2], [3])) + kept = _with_lineage(_entry("kept", [1, 2, 3, 4], [5]), parent_call_id="root") + dropped = _with_lineage(_entry("dropped", [1, 2, 3, 4], [9]), parent_call_id="root") + # The next call continued `kept`, and says so. + nxt = _with_lineage(_entry("next", [1, 2, 3, 4, 5, 6], [7]), parent_call_id="kept") + + out = prefix_merging([root, kept, dropped, nxt]) + main = next(c for c in out.chains if c.chain_id == "main") + assert [link.entry.model_call_id for link in main.links] == ["root", "kept", "next"] + assert "dropped" in out.quarantined + # Resolved, so nothing is flagged for masking. + assert out.notes["unresolved_retries"] == [] + + +def test_unresolvable_final_retry_is_flagged_not_silently_tie_broken(): + """A retry of the LAST call has no successor to name the survivor. Neither + inference nor a parent link can resolve it, so it must be reported so the + caller can mask the rollout instead of training on a generation the client + may never have received.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + a = _with_lineage(_entry("a", [1, 2, 3, 4], [5]), parent_call_id="root") + b = _with_lineage(_entry("b", [1, 2, 3, 4], [9]), parent_call_id="root") + + out = prefix_merging([root, a, b]) + assert sorted(out.notes["unresolved_retries"]) == ["a", "b"] + + +def test_a_stale_parent_link_fails_verification_and_falls_back(): + """A rerun that appended onto a previous attempt's records must not merge two + attempts. The digest check catches the bad edge; the builder falls back to + inference and reports that it did.""" + root = _with_lineage(_entry("root", [1, 2], [3])) + child = _entry("child", [1, 2, 3, 4], [5]) + stamp_lineage(child, "root") + # Corrupt the recorded parent's digest, as a stale record would. + root.digest = compute_digest([42, 42, 42]) + + out = prefix_merging([root, child]) + assert out.notes["parent_link_fallbacks"] == {"parent_digest_mismatch": 1} + # Inference still finds the right parent, so the chain is intact. + main = next(c for c in out.chains if c.chain_id == "main") + assert [link.entry.model_call_id for link in main.links] == ["root", "child"] + + +def test_parent_link_and_inference_agree_on_a_clean_rollout(): + """Parity: with and without recorded links, the same rollout must stitch the + same way. This is what makes the lineage fields safe to add before anything + populates them.""" + plain = [ + _entry("c1", [1, 2, 3], [4, 5]), + _entry("c2", [1, 2, 3, 4, 5, 6], [7]), + _entry("c3", [1, 2, 3, 4, 5, 6, 7, 8], [9, 10]), + ] + linked = [ + _with_lineage(_entry("c1", [1, 2, 3], [4, 5])), + _with_lineage(_entry("c2", [1, 2, 3, 4, 5, 6], [7]), parent_call_id="c1"), + _with_lineage(_entry("c3", [1, 2, 3, 4, 5, 6, 7, 8], [9, 10]), parent_call_id="c2"), + ] + inferred = prefix_merging(plain) + recorded = prefix_merging(linked) + assert [c.flatten()[0] for c in inferred.chains] == [c.flatten()[0] for c in recorded.chains] + assert [c.flatten()[1] for c in inferred.chains] == [c.flatten()[1] for c in recorded.chains] + + +def test_malformed_capture_masks_the_rollout_instead_of_raising(tmp_path): + """The callers are a rollout-collection loop and NeMo-RL's training loop; an + escaping exception there kills a whole step's batch rather than dropping one + sample.""" + store = TokenCaptureStore(tmp_path) + bad = _entry("c1", [1, 2, 3], [4, 5]) + bad.generation_log_probs = [-0.1] # one log prob for two generated tokens + store.append(bad) + + built = trajectories_for_rollout("t0-r0", [tmp_path]) + assert built is not None + assert built["mask_sample"] is True + assert built["nemo_rl_response"] is None + assert "ValueError" in built["error"] + + +def test_incomplete_capture_masks_the_rollout(tmp_path): + """A rollout that lost a call can still stitch into a clean-looking chain -- + it is just missing a turn. The marker is what makes that visible.""" + store = TokenCaptureStore(tmp_path) + store.append(_entry("c1", [1, 2, 3], [4, 5])) + store.mark_incomplete("t0-r0", "c2") + + built = trajectories_for_rollout("t0-r0", [tmp_path]) + assert built["mask_sample"] is True + assert built["metrics"]["capture_incomplete"] is True + + +def test_clean_rollout_is_not_masked_and_reports_full_delivery(tmp_path): + store = TokenCaptureStore(tmp_path) + store.append(_entry("c1", [1, 2, 3], [4, 5])) + store.append(_entry("c2", [1, 2, 3, 4, 5, 6], [7])) + + built = trajectories_for_rollout("t0-r0", [tmp_path]) + assert built["mask_sample"] is False + assert built["metrics"]["delivered_fraction"] == 1.0 + assert built["metrics"]["quarantined_calls"] == 0 From 5c317c1a27f3412391574bf51267d699c30531ab Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:32:34 -0700 Subject: [PATCH 3/9] feat(token-id-capture): deliver rebuilt trajectories and retire consumed 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 --- nemo_gym/base_responses_api_agent.py | 40 +++++-- nemo_gym/rollout_collection.py | 111 +++++++++++++++++- nemo_gym/token_id_capture/__init__.py | 2 + nemo_gym/token_id_capture/consumer.py | 21 ++++ .../test_base_responses_api_agent.py | 33 ++++++ tests/unit_tests/test_rollout_collection.py | 56 +++++++++ 6 files changed, 253 insertions(+), 10 deletions(-) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index c6e76cd89f..f0b951578c 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -26,7 +26,11 @@ BaseVerifyResponse, ) from nemo_gym.config_types import ROLLOUT_PATH_PREFIX -from nemo_gym.global_config import OBSERVABILITY_ENABLED_KEY_NAME, get_first_server_config_dict +from nemo_gym.global_config import ( + OBSERVABILITY_ENABLED_KEY_NAME, + TOKEN_ID_CAPTURE_ENABLED_KEY_NAME, + get_first_server_config_dict, +) from nemo_gym.openai_utils import ( NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming, @@ -43,7 +47,13 @@ class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig): - pass + # Whether this agent's rollouts participate in training token capture. Native agents receive + # token ids inline on the model response and leave this off; opaque external harnesses (whose + # returned output carries no token ids) set it true so their model calls are correlated and + # captured into the token store, then rebuilt into a token-bearing response.output. The run-level + # token_id_capture_enabled switch still gates the capture infrastructure; this scopes which + # agents use it. + token_id_capture: bool = False class BaseResponsesAPIAgent(BaseServer): @@ -79,21 +89,33 @@ async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyRespo return app - def _model_call_capture_enabled(self) -> bool: - # Fail closed: an agent whose client carries no usable global config runs uncorrelated - # rather than erroring on every model call. + def _capture_correlation_enabled(self) -> bool: + """Whether the per-rollout ``/ng-rollout/`` correlation prefix should be applied. + + Two independent capture paths consume the same prefix: + - Eval model-call capture (``observability_enabled``), which applies to every agent. + - Training token capture (``token_id_capture_enabled``), which applies only to agents + that opt in with the per-agent ``token_id_capture`` flag -- native agents carry token + ids inline and do not need the store, so they do not emit the prefix for token capture. + + Fail closed: an agent whose client carries no usable global config runs uncorrelated + rather than erroring on every model call. + """ global_config = getattr(self.server_client, "global_config_dict", None) if not isinstance(global_config, Mapping): return False - return bool(global_config.get(OBSERVABILITY_ENABLED_KEY_NAME, False)) + token_capture = bool(global_config.get(TOKEN_ID_CAPTURE_ENABLED_KEY_NAME, False)) and bool( + getattr(self.config, "token_id_capture", False) + ) + return bool(global_config.get(OBSERVABILITY_ENABLED_KEY_NAME, False) or token_capture) def rollout_id_from_run(self, body: Any) -> Optional[str]: """Per-rollout capture id for a run-request (its task/rollout indices). - None when model-call capture (observability) is disabled or the body carries no indices, - so callers apply no correlation prefix in either case. + None when neither capture path is enabled or the body carries no indices, so callers apply + no correlation prefix in either case. """ - if not self._model_call_capture_enabled(): + if not self._capture_correlation_enabled(): return None return maybe_rollout_id_from_run_body(body) diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index d503687980..33817bd528 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -35,6 +35,7 @@ from nemo_gym.base_resources_server import AggregateMetrics, AggregateMetricsRequest from nemo_gym.base_responses_api_model import ( clear_model_call_captures_for_rollouts, + maybe_rollout_id_from_run_body, merge_model_call_capture_into_record, model_call_capture_dirs_from_config, ) @@ -46,6 +47,8 @@ ROLLOUT_INDEX_KEY_NAME, SKILLS_REF_KEY_NAME, TASK_INDEX_KEY_NAME, + TOKEN_ID_CAPTURE_KEY_NAME, + get_first_server_config_dict, get_global_config_dict, get_wandb_run, ) @@ -60,6 +63,12 @@ set_global_aiohttp_client, ) from nemo_gym.skills import SkillsConfig, load_skill_directory +from nemo_gym.token_id_capture import ( + TokenCaptureStore, + clear_token_captures_for_rollouts, + token_id_capture_dirs_from_config, + trajectories_for_rollout, +) # --------------------------------------------------------------------------- @@ -90,6 +99,10 @@ NG_FAILURE_CLASS_KEY = "_ng_failure_class" NG_NO_PERSIST_KEY = "_ng_no_persist" +# Per-rollout token-capture health, attached to the record so a run can see what the build dropped: +# chain count, quarantined fraction, delivered token fraction, and whether the rollout should be +# masked (an incomplete capture, or a final-call retry that cannot be resolved). +NG_TOKEN_CAPTURE_KEY = "_ng_token_capture" NG_TERMINAL_KEY = "_ng_failure_terminal" _DEFAULT_MAX_ROLLOUT_ATTEMPTS = 3 @@ -247,6 +260,20 @@ def materialized_jsonl_fpath(self) -> Path: return output_fpath.with_stem(output_fpath.stem + "_materialized_inputs").with_suffix(".jsonl") +def _agent_participates_in_token_capture(global_config: Any, agent_name: Optional[str]) -> bool: + """Whether the producing agent opted into training token capture (its ``token_id_capture`` + flag). Native agents leave it off -- they carry token ids inline and need no store rebuild -- so + this scopes the rebuild (and its token-less warnings) to external-harness rollouts. Defaults to + False when the agent or flag is absent.""" + if not agent_name: + return False + try: + agent_config = get_first_server_config_dict(global_config, agent_name) + except (KeyError, IndexError, TypeError): + return False + return bool((agent_config or {}).get(TOKEN_ID_CAPTURE_KEY_NAME, False)) + + def _rollout_request_debug_summary(row: Dict[str, Any]) -> Dict[str, Any]: agent_ref = row.get(AGENT_REF_KEY_NAME) or {} summary = { @@ -522,13 +549,24 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D # Resolve capture dirs once so each rollout's captured model calls can be folded # into its record below (uniform across agents; no-op when capture is off / dirs absent). - capture_dirs = model_call_capture_dirs_from_config(get_global_config_dict()) + global_config = get_global_config_dict() + capture_dirs = model_call_capture_dirs_from_config(global_config) + # Resolve the training-token store dir once, so each rollout's captured tokens can be built + # into a trajectory below. Independent of eval capture; empty (and a no-op) when token + # capture is off. + token_capture_dirs = token_id_capture_dirs_from_config(global_config) # Clear only rows about to be dispatched, after resume has assigned retry suffixes. This also # removes a kill-shaped attempt's partial capture when its rollout-attempt id is reused. if capture_dirs: print("Clearing existing model-call captures for rollouts being dispatched") clear_model_call_captures_for_rollouts(input_rows, capture_dirs) + if token_capture_dirs: + # Same reason, for the token store: rollout ids are deterministic and the store appends, + # so without this a fresh run would build a trajectory that merges a previous attempt's + # calls with this one's. + print("Clearing existing token captures for rollouts being dispatched") + clear_token_captures_for_rollouts(input_rows, token_capture_dirs) pcts_to_print = [20, 40, 60, 80, 90, 95, 98, 99, 100] counts_left = Counter(r[AGENT_REF_KEY_NAME]["name"] for r in input_rows) @@ -550,6 +588,77 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D if capture_dirs: merge_model_call_capture_into_record(result, capture_dirs) + # Build this rollout's captured tokens into a trajectory and attach it (no-op when token + # capture is off or no tokens were captured for the rollout). Never alters output/reward. + agent_name = (result.get(AGENT_REF_KEY_NAME) or {}).get("name") + if token_capture_dirs and _agent_participates_in_token_capture(global_config, agent_name): + rollout_id = maybe_rollout_id_from_run_body(result) + if rollout_id is None: + # Capture is on but the rollout carries no task/rollout indices, so its model + # calls could not be correlated. Fail loud: a silently token-less rollout would + # become an unusable (masked/empty) training sample. Usually means the agent did + # not apply the /ng-rollout correlation prefix to its model calls. + warnings.warn( + "training token capture is enabled but a rollout result carries no " + "task/rollout indices; its model calls were not correlated or captured and " + "it will be token-less. Check that the agent applies the /ng-rollout prefix.", + stacklevel=2, + ) + else: + response = result.get("response") if isinstance(result.get("response"), dict) else {} + built = trajectories_for_rollout( + rollout_id, + token_capture_dirs, + reward=float(result.get("reward") or 0.0), + reward_components=result.get("reward_components"), + model=str(response.get("model") or ""), + ) + if built is not None: + # Uniform delivery: replace the training-facing response's output with the + # merged, contiguous, content-bearing + token-bearing items, so NeMo-RL reads + # response.output the same way for native and blackbox rollouts (no sidecar, + # no per-agent branch). reward / reward_components stay on the record for GDPO. + projected = built["nemo_rl_response"] + if projected is not None: + if isinstance(result.get("response"), dict): + result["response"]["output"] = projected["output"] + else: + result["response"] = projected + # Carry what the build dropped onto the record. Without this the quarantine + # and delivered-token fractions are computed and discarded, and a rollout + # that trained on one of five calls looks like one that trained on all five. + record_metrics = dict(built.get("metrics") or {}) + if built.get("mask_sample"): + record_metrics["mask_sample"] = True + if built.get("error"): + record_metrics["error"] = built["error"] + result[NG_TOKEN_CAPTURE_KEY] = record_metrics + if built.get("mask_sample"): + warnings.warn( + f"rollout {rollout_id} was captured incompletely or ambiguously " + f"({record_metrics}); it is marked for masking rather than trained on.", + stacklevel=2, + ) + # Delete on consume. Rollouts write hundreds of KB each and the store + # appends, so keeping consumed files both grows the directory without bound + # and lets a later run with the same id append onto them. Failed builds keep + # their records: they are the only evidence of why the build failed. + if projected is not None and not os.environ.get("NG_KEEP_TOKCAP"): + for capture_dir in token_capture_dirs: + TokenCaptureStore(capture_dir).delete(rollout_id) + else: + # A derivable rollout id but an empty store means correlation broke somewhere + # between the agent and the capture middleware (e.g. an external harness or + # proxy that did not preserve the prefix). Surface it rather than emitting a + # token-less rollout into training. + warnings.warn( + f"training token capture is enabled but no tokens were captured for " + f"rollout {rollout_id}; its response.output was not rebuilt and it will " + "be token-less. The agent's model calls likely did not reach the capture " + "middleware correlated.", + stacklevel=2, + ) + no_persist = bool(result.get(NG_NO_PERSIST_KEY)) failure_class = result.get(NG_FAILURE_CLASS_KEY) diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index 2c3ed56793..d97fd0effa 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -41,6 +41,7 @@ ) from nemo_gym.token_id_capture.config import TokenIdCaptureConfig from nemo_gym.token_id_capture.consumer import ( + clear_token_captures_for_rollouts, token_id_capture_dirs_from_config, trajectories_for_rollout, trajectories_from_source, @@ -101,6 +102,7 @@ "Chain", "BuildOutput", "trajectories_for_rollout", + "clear_token_captures_for_rollouts", "trajectories_from_source", "token_id_capture_dirs_from_config", ] diff --git a/nemo_gym/token_id_capture/consumer.py b/nemo_gym/token_id_capture/consumer.py index e5cc01db95..01525eb939 100644 --- a/nemo_gym/token_id_capture/consumer.py +++ b/nemo_gym/token_id_capture/consumer.py @@ -57,6 +57,27 @@ def token_id_capture_dirs_from_config(global_config_dict) -> list[Path]: return [directory] if (config.token_id_capture_enabled and directory is not None) else [] +def clear_token_captures_for_rollouts(records: list, token_capture_dirs: list[Path]) -> None: + """Remove stale token records for rollouts about to be dispatched. + + Rollout ids are deterministic and ``TokenCaptureStore.append`` opens in "ab" + mode, so a rerun that reuses an id would append onto the previous attempt's + records and the builder would stitch two attempts into one trajectory. The + caller passes only the rows being dispatched, after any retry suffix has been + assigned. + """ + if not token_capture_dirs: + return + from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body + + for directory in token_capture_dirs: + store = TokenCaptureStore(directory) + for record in records: + rollout_id = maybe_rollout_id_from_run_body(record) + if rollout_id: + store.delete(rollout_id) + + def _assemble( rollout_id: str, entries: list[TokenEntry], diff --git a/tests/unit_tests/test_base_responses_api_agent.py b/tests/unit_tests/test_base_responses_api_agent.py index c6f0a1668f..a9dac86028 100644 --- a/tests/unit_tests/test_base_responses_api_agent.py +++ b/tests/unit_tests/test_base_responses_api_agent.py @@ -39,3 +39,36 @@ async def run(self, body=...): agent = TestSimpleResponsesAPIAgent(config=config, server_client=MagicMock(spec=ServerClient)) agent.setup_webserver() + + def _agent(self, global_config: dict, *, token_id_capture: bool = False) -> SimpleResponsesAPIAgent: + config = BaseResponsesAPIAgentConfig( + host="", port=0, entrypoint="", name="", token_id_capture=token_id_capture + ) + + class _Agent(SimpleResponsesAPIAgent): + async def responses(self, body=...): + raise NotImplementedError + + async def run(self, body=...): + raise NotImplementedError + + client = MagicMock(spec=ServerClient) + client.global_config_dict = global_config + return _Agent(config=config, server_client=client) + + def test_eval_capture_prefix_applies_to_every_agent(self) -> None: + # Eval capture (observability_enabled) correlates every agent, regardless of the per-agent + # token-capture opt-in. + body = {"_ng_task_index": 0, "_ng_rollout_index": 0} + assert self._agent({}).rollout_id_from_run(body) is None + assert self._agent({"observability_enabled": True}).rollout_id_from_run(body) == "0-0" + + def test_token_capture_prefix_is_scoped_to_participating_agents(self) -> None: + # Training token capture correlates a call only when the run-level switch is on AND the agent + # opted in -- native agents (opt-out) carry token ids inline and must not be correlated here. + body = {"_ng_task_index": 0, "_ng_rollout_index": 0} + gc = {"token_id_capture_enabled": True} + assert self._agent(gc, token_id_capture=False).rollout_id_from_run(body) is None + assert self._agent(gc, token_id_capture=True).rollout_id_from_run(body) == "0-0" + # The run-level switch is still required: opting in alone does nothing. + assert self._agent({}, token_id_capture=True).rollout_id_from_run(body) is None diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index d3145eb403..d71e823ad7 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -36,12 +36,25 @@ RolloutAggregationHelper, RolloutCollectionConfig, RolloutCollectionHelper, + _agent_participates_in_token_capture, _expand_input_glob, _failures_path_for, _get_max_rollout_attempts, _rollout_request_debug_summary, loads_jsonl_line, ) +from nemo_gym.token_id_capture import TokenCaptureStore, TokenEntry, clear_token_captures_for_rollouts + + +def test_agent_participates_in_token_capture_is_scoped_per_agent() -> None: + global_config = { + "ext_agent": {"responses_api_agents": {"claude_code_agent": {"token_id_capture": True}}}, + "native_agent": {"responses_api_agents": {"simple_agent": {}}}, + } + assert _agent_participates_in_token_capture(global_config, "ext_agent") is True + assert _agent_participates_in_token_capture(global_config, "native_agent") is False + assert _agent_participates_in_token_capture(global_config, "missing") is False + assert _agent_participates_in_token_capture(global_config, None) is False @pytest.fixture @@ -1357,3 +1370,46 @@ async def _noop(self, results, rows, output_fpath): # though output_jsonl_fpath is used to derive the metrics path. assert not output_fpath.exists() assert (tmp_path / "rollouts_aggregate_metrics.json").exists() + + +class TestTokenCaptureRetention: + """Delete-on-consume and stale-record clearing. + + Both directions matter because ``TokenCaptureStore.append`` opens in "ab" + mode and rollout ids are deterministic: without clearing, a rerun stitches + the previous attempt's calls together with this one's; without deleting, + the capture directory grows without bound across a run. + """ + + @staticmethod + def _entry(rollout_id: str, mcid: str) -> TokenEntry: + return TokenEntry( + rollout_id=rollout_id, + model_call_id=mcid, + prompt_token_ids=[1, 2, 3], + generation_token_ids=[4, 5], + generation_log_probs=[-0.1, -0.2], + ) + + def test_clear_removes_stale_records_before_dispatch(self, tmp_path: Path) -> None: + store = TokenCaptureStore(tmp_path) + store.append(self._entry("0-0", "old")) + store.mark_incomplete("0-0", "old") + rows = [{TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0}] + + clear_token_captures_for_rollouts(rows, [tmp_path]) + + assert store.read_entries("0-0") == [] + assert not store.is_incomplete("0-0") + + def test_clear_is_a_noop_without_capture_dirs(self, tmp_path: Path) -> None: + store = TokenCaptureStore(tmp_path) + store.append(self._entry("0-0", "keep")) + clear_token_captures_for_rollouts([{TASK_INDEX_KEY_NAME: 0, ROLLOUT_INDEX_KEY_NAME: 0}], []) + assert len(store.read_entries("0-0")) == 1 + + def test_clear_skips_rows_without_a_derivable_rollout_id(self, tmp_path: Path) -> None: + store = TokenCaptureStore(tmp_path) + store.append(self._entry("0-0", "keep")) + clear_token_captures_for_rollouts([{"unrelated": True}], [tmp_path]) + assert len(store.read_entries("0-0")) == 1 From 83e904d02e7e701725aa039aa4edab162933d916 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 28 Jul 2026 12:43:26 -0700 Subject: [PATCH 4/9] feat(token-id-capture): keep harness side calls out of the trajectory 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 Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 25 ++++++++- nemo_gym/token_id_capture/__init__.py | 2 + nemo_gym/token_id_capture/builder.py | 59 ++++++++++++++++++++ nemo_gym/token_id_capture/consumer.py | 32 ++++++++++- nemo_gym/token_id_capture/records.py | 9 +++ nemo_gym/token_id_capture/sink.py | 7 ++- tests/unit_tests/test_token_id_capture.py | 20 +++++++ tests/unit_tests/test_trajectory_builder.py | 61 +++++++++++++++++++++ 8 files changed, 209 insertions(+), 6 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 25a2191a12..a08fb77839 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -86,6 +86,27 @@ _ANTHROPIC_CONVERTER = AnthropicConverter() +def _request_facts(body: Any) -> dict[str, Any]: + """What the harness asked for, read off the parsed request body. + + Two things, both used to tell trajectory calls from side calls: the model the + harness requested (not the one the server served), and whether the request + declared tools. Reading it from the parsed request avoids touching the body a + second time in middleware. + """ + if body is None: + return {} + getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default) + model = getter("model", None) + tools = getter("tools", None) + facts: dict[str, Any] = {} + if isinstance(model, str) and model: + facts["requested_model"] = model + if tools is not None: + facts["has_tools"] = bool(tools) + return facts + + class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): pass @@ -210,7 +231,7 @@ async def _invoke_chat_completions( completion = await self.chat_completions(request=request, body=params) else: completion = await self.chat_completions(body=params) - await capture_tokens(completion) + await capture_tokens(completion, request_facts=_request_facts(params)) return completion async def messages(self, request: Request, body: dict = Body()): @@ -247,7 +268,7 @@ async def _invoke_responses( # Capture here rather than at the route: the streaming dispatch returns a StreamingResponse # and the Anthropic mapping drops the token fields, so this is the last point where the # assembled response still carries them, for every dialect. - await capture_tokens(response) + await capture_tokens(response, request_facts=_request_facts(params)) return response diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index d97fd0effa..cf357b98ff 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -35,6 +35,7 @@ Trajectory, assert_nemo_rl_contiguity, build_trajectories, + classify_side_calls, per_request, prefix_merging, project_main_chain_response, @@ -94,6 +95,7 @@ "capture_tokens", "CaptureTokenSource", "build_trajectories", + "classify_side_calls", "per_request", "prefix_merging", "project_main_chain_response", diff --git a/nemo_gym/token_id_capture/builder.py b/nemo_gym/token_id_capture/builder.py index e960fa6e78..295894b0a1 100644 --- a/nemo_gym/token_id_capture/builder.py +++ b/nemo_gym/token_id_capture/builder.py @@ -118,6 +118,65 @@ class Trajectory(BaseModel): provenance: dict = Field(default_factory=dict) +def classify_side_calls( + entries: list[TokenEntry], + side_call_model_patterns: tuple[str, ...] = (), +) -> tuple[list[TokenEntry], list[TokenEntry]]: + """Split a rollout's records into trajectory calls and side calls. + + A harness may make model calls that are not part of the trajectory being + trained -- a conversation title or a quota probe, for example. They are + genuine policy output (real token ids, real log probs), so nothing downstream + distinguishes them from trajectory calls. + + Not observed on the Claude Code path so far: runs report + side_calls_excluded = 0. This is a guard, not a fix for a measured problem. + + Two signals, neither requiring harness-specific code in the core: + + 1. An explicit pattern list (substring, case-insensitive) for deployments + that know their harness, e.g. ``("haiku",)``. + 2. Self-calibration on the requested model. A harness asks for a *different* + (usually smaller) model for these calls. Whichever model generated the + most tokens in this rollout is the policy model; calls asking for another + one are side calls. This needs no configuration and adapts to whatever the + harness happens to name its models. + + Records written before this field existed carry an empty ``requested_model``, + so they are all treated as trajectory calls and nothing changes for them. + """ + if not entries: + return [], [] + + patterns = tuple(p.lower() for p in side_call_model_patterns if p) + + def matches_pattern(entry: TokenEntry) -> bool: + name = (entry.requested_model or "").lower() + return bool(name) and any(p in name for p in patterns) + + kept = [e for e in entries if not matches_pattern(e)] + excluded = [e for e in entries if matches_pattern(e)] + + # Self-calibration: only meaningful when the records actually recorded what was requested and + # more than one model appears. A single-model rollout is left alone. + named = [e for e in kept if e.requested_model] + distinct = {e.requested_model for e in named} + if len(distinct) > 1: + by_model: dict[str, int] = {} + for entry in named: + by_model[entry.requested_model] = by_model.get(entry.requested_model, 0) + len(entry.generation_token_ids) + # Ties break on the name so the choice is deterministic. + policy_model = max(sorted(by_model), key=lambda m: (by_model[m], m)) + still_kept, extra = [], [] + for entry in kept: + (still_kept if (not entry.requested_model or entry.requested_model == policy_model) else extra).append( + entry + ) + kept, excluded = still_kept, excluded + extra + + return kept, excluded + + def per_request(entries: list[TokenEntry]) -> BuildOutput: ordered = sorted(entries, key=lambda e: (len(e.prompt_token_ids), e.model_call_id)) chains = [ diff --git a/nemo_gym/token_id_capture/consumer.py b/nemo_gym/token_id_capture/consumer.py index 01525eb939..010070da57 100644 --- a/nemo_gym/token_id_capture/consumer.py +++ b/nemo_gym/token_id_capture/consumer.py @@ -36,6 +36,7 @@ from nemo_gym.token_id_capture.builder import ( assert_nemo_rl_contiguity, build_trajectories, + classify_side_calls, per_request, prefix_merging, project_main_chain_response, @@ -85,7 +86,28 @@ def _assemble( reward: float, reward_components: Optional[dict[str, float]], model: str, + side_call_model_patterns: tuple[str, ...] = (), ) -> dict: + # Keep any harness side calls out of the trajectory. They are policy output and look + # trainable, so nothing downstream would notice them. + all_entries = entries + entries, side_calls = classify_side_calls(entries, side_call_model_patterns) + if not entries: + logger.warning( + "Rollout %s captured %d call(s), all classified as harness side calls; nothing to build.", + rollout_id, + len(all_entries), + ) + return { + "rollout_id": rollout_id, + "builder": builder, + "trajectories": [], + "nemo_rl_response": None, + "mask_sample": True, + "error": "all captured calls were side calls", + "metrics": {"n_calls": len(all_entries), "side_calls_excluded": len(side_calls)}, + } + # A malformed capture must degrade this one rollout, not take down the caller. # Both the contiguity assertion and the flattener raise, and the callers are a # rollout-collection loop and NeMo-RL's training loop -- where an escaping @@ -111,7 +133,7 @@ def _assemble( "nemo_rl_response": None, "mask_sample": True, "error": f"{type(error).__name__}: {error}", - "metrics": {"n_calls": len(entries)}, + "metrics": {"n_calls": len(entries), "side_calls_excluded": len(side_calls)}, } notes = dict(out.notes) @@ -126,6 +148,8 @@ def _assemble( "generated_tokens_captured": notes.get("generated_tokens_captured", 0), "generated_tokens_delivered": notes.get("generated_tokens_delivered", 0), "parent_link_fallbacks": notes.get("parent_link_fallbacks", {}), + # Excluded, not dropped silently. + "side_calls_excluded": len(side_calls), } unresolved = notes.get("unresolved_retries") or [] return { @@ -149,6 +173,7 @@ def trajectories_for_rollout( reward: float = 0.0, reward_components: Optional[dict[str, float]] = None, model: str = "", + side_call_model_patterns: tuple[str, ...] = (), ) -> Optional[dict]: """Co-located path: read the rollout's tokens from the store files and build its trajectories. @@ -161,7 +186,7 @@ def trajectories_for_rollout( store = TokenCaptureStore(directory) entries = store.read_entries(rollout_id) if entries: - built = _assemble(rollout_id, entries, builder, reward, reward_components, model) + built = _assemble(rollout_id, entries, builder, reward, reward_components, model, side_call_model_patterns) if store.is_incomplete(rollout_id): # At least one call of this rollout failed to capture. The chain we built may look # perfectly contiguous while being missing a turn, so mask rather than train on it. @@ -179,10 +204,11 @@ async def trajectories_from_source( reward: float = 0.0, reward_components: Optional[dict[str, float]] = None, model: str = "", + side_call_model_patterns: tuple[str, ...] = (), ) -> Optional[dict]: """Non-co-located path: read the rollout's tokens through a ``TokenSource`` (HTTP, or a trainer's own transport) and build its trajectories. Returns ``None`` when none were captured.""" entries = await source.tokens_for(rollout_id) if not entries: return None - return _assemble(rollout_id, entries, builder, reward, reward_components, model) + return _assemble(rollout_id, entries, builder, reward, reward_components, model, side_call_model_patterns) diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index abd1eb39f8..7f371ac32d 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -100,6 +100,15 @@ class TokenEntry(BaseModel): # Non-semantic; a cheap diagnostic for retry/sibling-branch cases. created_at: float = 0.0 + # --- What the harness asked for, used to tell trajectory calls from side calls. + # A harness may make model calls that are not part of the rollout -- a + # conversation title or a quota probe, for example. They carry real token ids + # and log probs, so nothing downstream distinguishes them. ``requested_model`` + # is what the harness asked for (not what the server served), since a harness + # commonly requests a smaller model for such calls. + requested_model: str = "" + has_tools: Optional[bool] = None + # --- Lineage. Optional: null when the model server could not identify the # parent, in which case the builder infers it from token prefixes instead. # diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index f4978145a2..3eb16549be 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -86,7 +86,11 @@ def reset_token_sink(token: Token) -> None: _TOKEN_SINK.reset(token) -async def capture_tokens(response: Any, parent_call_id: Optional[str] = None) -> None: +async def capture_tokens( + response: Any, + parent_call_id: Optional[str] = None, + request_facts: Optional[dict] = None, +) -> None: """Record a ``TokenEntry`` from a complete model response when a sink is set. ``response`` is a served response as a pydantic model or dict. No-op when no @@ -120,6 +124,7 @@ async def capture_tokens(response: Any, parent_call_id: Optional[str] = None) -> # reads is not token-only -- text-based penalties need it. output_items=response_to_output_items(payload), created_at=time.time(), + **(request_facts or {}), ) # cum_len/digest describe this call and are always computable; the parent # link is filled only when the model server resolved one. diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index d021d490c3..9eae219c56 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -470,3 +470,23 @@ def test_delete_removes_records_and_marker(tmp_path): assert not store.is_incomplete("gone-0") # Idempotent: consuming a rollout twice must not raise. store.delete("gone-0") + + +def test_capture_records_the_model_the_harness_asked_for(tmp_path): + """Claude Code asks for a small model when it generates a conversation title. + The record must keep what the *harness* requested, not what the server + served, or side calls are indistinguishable from trajectory calls.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post( + "/ng-rollout/facts0-roll0/v1/messages", + json={ + "model": "claude-3-5-haiku-20241022", + "max_tokens": 16, + "messages": [{"role": "user", "content": "title this"}], + }, + ) + assert resp.status_code == 200 + (entry,) = TokenCaptureStore(tmp_path).read_entries("facts0-roll0") + assert entry.requested_model == "claude-3-5-haiku-20241022" + # The served model is a different field and is unaffected. + assert entry.model != entry.requested_model diff --git a/tests/unit_tests/test_trajectory_builder.py b/tests/unit_tests/test_trajectory_builder.py index 4d4862e680..47b9765fa8 100644 --- a/tests/unit_tests/test_trajectory_builder.py +++ b/tests/unit_tests/test_trajectory_builder.py @@ -20,6 +20,7 @@ Trajectory, assert_nemo_rl_contiguity, build_trajectories, + classify_side_calls, compute_digest, prefix_merging, project_main_chain_response, @@ -404,3 +405,63 @@ def test_clean_rollout_is_not_masked_and_reports_full_delivery(tmp_path): assert built["mask_sample"] is False assert built["metrics"]["delivered_fraction"] == 1.0 assert built["metrics"]["quarantined_calls"] == 0 + + +# --- side calls --------------------------------------------------------------- + + +def _sc(mcid, prompt, gen, requested_model=""): + entry = _entry(mcid, prompt, gen) + entry.requested_model = requested_model + return entry + + +def test_side_calls_are_excluded_by_self_calibration(): + """No configuration: whichever model generated the most tokens is the policy + model, and calls asking for a different one are side calls.""" + real_1 = _sc("r1", [100, 101], [1, 2, 3, 4, 5], "big-policy-model") + real_2 = _sc("r2", [100, 101, 1, 2, 3, 4, 5, 6], [7, 8, 9], "big-policy-model") + title = _sc("title", [9000], [42], "tiny-title-model") + + kept, excluded = classify_side_calls([real_1, real_2, title]) + + assert [e.model_call_id for e in kept] == ["r1", "r2"] + assert [e.model_call_id for e in excluded] == ["title"] + + +def test_side_calls_are_excluded_by_explicit_pattern(): + real = _sc("r1", [100], [1, 2, 3], "claude-sonnet-4") + title = _sc("title", [9000], [42], "claude-3-5-haiku-20241022") + + kept, excluded = classify_side_calls([real, title], side_call_model_patterns=("haiku",)) + + assert [e.model_call_id for e in kept] == ["r1"] + assert [e.model_call_id for e in excluded] == ["title"] + + +def test_records_without_a_requested_model_are_all_kept(): + """Backward compatibility: records written before the field existed must not + be reclassified as side calls.""" + entries = [_entry("c1", [1, 2, 3], [4, 5]), _entry("c2", [1, 2, 3, 4, 5, 6], [7])] + kept, excluded = classify_side_calls(entries) + assert len(kept) == 2 and excluded == [] + + +def test_single_model_rollout_is_left_alone(): + entries = [_sc("c1", [1, 2, 3], [4, 5], "m"), _sc("c2", [1, 2, 3, 4, 5, 6], [7], "m")] + kept, excluded = classify_side_calls(entries) + assert len(kept) == 2 and excluded == [] + + +def test_consumer_excludes_side_calls_and_reports_the_count(tmp_path): + store = TokenCaptureStore(tmp_path) + store.append(_sc("r1", [100, 101], [1, 2, 3, 4, 5], "policy")) + store.append(_sc("r2", [100, 101, 1, 2, 3, 4, 5, 6], [7, 8, 9], "policy")) + store.append(_sc("title", [9000], [42], "titler")) + + built = trajectories_for_rollout("t0-r0", [tmp_path]) + + assert built["metrics"]["side_calls_excluded"] == 1 + assert built["metrics"]["n_calls"] == 2 + call_ids = [span[2] for span in built["trajectories"][0]["spans"]] + assert "title" not in call_ids From ecdccec20a808afd6bfb403e103ed52f08f25a67 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 28 Jul 2026 12:50:04 -0700 Subject: [PATCH 5/9] feat(token-id-capture): resolve each call's parent at request time 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 Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 32 ++- nemo_gym/token_id_capture/__init__.py | 12 + nemo_gym/token_id_capture/lineage.py | 262 ++++++++++++++++++++++ nemo_gym/token_id_capture/sink.py | 33 +++ tests/unit_tests/test_token_id_capture.py | 162 +++++++++++++ 5 files changed, 499 insertions(+), 2 deletions(-) create mode 100644 nemo_gym/token_id_capture/lineage.py diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index a08fb77839..1ca25aa3f1 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -107,6 +107,26 @@ def _request_facts(body: Any) -> dict[str, Any]: return facts +def _request_messages(body: Any) -> list[dict]: + """The conversation a request carries, across the three dialects. + + Used only to identify which recorded call this request continues (see + ``token_id_capture.lineage``): the assistant turns in it are the ones we + produced. Chat and Anthropic both use ``messages``; Responses carries + ``input``, which is a string for a first turn and a list of items after that. + """ + if body is None: + return [] + getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default) + messages = getter("messages", None) + if isinstance(messages, list): + return [m if isinstance(m, dict) else m.model_dump() for m in messages if m is not None] + items = getter("input", None) + if isinstance(items, list): + return [i if isinstance(i, dict) else i.model_dump() for i in items if i is not None] + return [] + + class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): pass @@ -231,7 +251,11 @@ async def _invoke_chat_completions( completion = await self.chat_completions(request=request, body=params) else: completion = await self.chat_completions(body=params) - await capture_tokens(completion, request_facts=_request_facts(params)) + await capture_tokens( + completion, + request_messages=_request_messages(params), + request_facts=_request_facts(params), + ) return completion async def messages(self, request: Request, body: dict = Body()): @@ -268,7 +292,11 @@ async def _invoke_responses( # Capture here rather than at the route: the streaming dispatch returns a StreamingResponse # and the Anthropic mapping drops the token fields, so this is the last point where the # assembled response still carries them, for every dialect. - await capture_tokens(response, request_facts=_request_facts(params)) + await capture_tokens( + response, + request_messages=_request_messages(params), + request_facts=_request_facts(params), + ) return response diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index cf357b98ff..d6b11c3f1c 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -41,6 +41,13 @@ project_main_chain_response, ) from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.lineage import ( + LineageIndex, + LineageNode, + RolloutLineage, + assistant_fingerprint, + canonicalize_tool_arguments, +) from nemo_gym.token_id_capture.consumer import ( clear_token_captures_for_rollouts, token_id_capture_dirs_from_config, @@ -94,6 +101,11 @@ "reset_token_sink", "capture_tokens", "CaptureTokenSource", + "LineageIndex", + "LineageNode", + "RolloutLineage", + "assistant_fingerprint", + "canonicalize_tool_arguments", "build_trajectories", "classify_side_calls", "per_request", diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py new file mode 100644 index 0000000000..42e0a54333 --- /dev/null +++ b/nemo_gym/token_id_capture/lineage.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Which recorded call does this request continue? + +The builder can infer that after the fact, by finding the earlier call whose +token sequence is the longest prefix of this call's prompt. That works, but it +cannot do two things: + +- **Resolve a retry.** Capture records a call once the response is assembled, + including one the client never received, so a harness retry leaves two records + with the same prompt and different generations. Both are equally valid + children, so inference has to guess between them. +- **Run before the call.** Inference happens after the rollout. Supplying the + engine the previous call's exact tokens (so the next prompt extends them by + construction instead of being re-rendered from text) has to happen *before* the + request is forwarded, when there is nothing yet to infer from. + +This module answers the question at request time instead, using only what the +harness already sends. A harness must echo the conversation to continue it, so +the assistant messages 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. + +The index is a map keyed by call, not a running cursor, so it is a tree and +forks cost nothing. Two sub-agents branching from one parent both resolve to it +and both get the same prefix -- which is exactly right, and is why the prefix to +supply is ``cum(parent)`` and not ``cum(previous call)``. Entries are added and +never mutated, so concurrent sub-agents cannot corrupt each other's lineage. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from typing import Any, Optional + + +_FINGERPRINT_DOMAIN = b"nemo-gym-lineage" + + +def canonicalize_tool_arguments(value: Any) -> str: + """Normalize a tool call's arguments for comparison only. + + Harnesses re-serialize tool-call arguments between turns -- compact one turn, + pretty-printed the next -- so the same call does not compare equal to itself. + Comparison is done on sorted-key, separator-normalized JSON; the model's + original string is what stays in the record and is never rewritten. + """ + if value is None: + return "" + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, ValueError): + return value.strip() + else: + parsed = value + try: + return json.dumps(parsed, sort_keys=True, separators=(",", ":")) + except (TypeError, ValueError): + return str(parsed) + + +def _text_of(content: Any) -> str: + """Flatten a message's *text* content, across the shapes the dialects use. + + Tool calls are deliberately NOT folded in here -- see ``_tools_of``. The two + dialects carry them in different places, so they have to be normalized + separately or the same turn hashes differently depending on which side is + looking at it. + """ + if content is None: + return "" + if isinstance(content, str): + return content + parts: list[str] = [] + if isinstance(content, list): + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict) and isinstance(block.get("text"), str): + parts.append(block["text"]) + return "\n".join(parts) + + +def _tools_of(message: dict) -> list[tuple[str, str]]: + """The tool calls in an assistant turn, as (name, canonical args). + + A turn we produced is recorded in OpenAI shape (``tool_calls`` with + ``function.name`` / ``function.arguments``), but Claude Code echoes it back + in Anthropic shape (``content`` blocks of ``type: tool_use`` with ``name`` / + ``input``). Both must reduce to the same list, or a tool-using turn -- the + only kind that produces multi-turn rollouts -- can never match itself. + """ + tools: list[tuple[str, str]] = [] + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + tools.append((str(block.get("name", "")), canonicalize_tool_arguments(block.get("input")))) + for call in message.get("tool_calls") or []: + function = (call or {}).get("function") or {} + tools.append((str(function.get("name", "")), canonicalize_tool_arguments(function.get("arguments")))) + return tools + + +def assistant_fingerprint(messages: list[dict]) -> str: + """Fingerprint the assistant turns of a request, in order. + + Only assistant turns are hashed: they are what *we* produced, so they + identify the call that produced the last one. User and tool content varies + with the environment and is irrelevant to lineage. + """ + hasher = hashlib.sha256(_FINGERPRINT_DOMAIN) + count = 0 + for message in messages or []: + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + count += 1 + hasher.update(b"\x00") + hasher.update(_text_of(message.get("content")).encode("utf-8")) + for name, arguments in _tools_of(message): + hasher.update(b"\x01") + hasher.update(name.encode("utf-8")) + hasher.update(arguments.encode("utf-8")) + if count == 0: + return "" + return hasher.hexdigest() + + +@dataclass +class LineageNode: + call_id: str + cum_tokens: list[int] + cum_len: int + digest: str + + +@dataclass +class RolloutLineage: + """Per-rollout call index. Append-only, so forks and concurrency are safe.""" + + by_fingerprint: dict[str, list[str]] = field(default_factory=dict) + by_call_id: dict[str, LineageNode] = field(default_factory=dict) + + def resolve(self, messages: list[dict]) -> Optional[LineageNode]: + """The call this request continues, or ``None``. + + ``None`` for a new root (nothing matches -- a fresh conversation, or one + the harness rewrote) and, deliberately, for an ambiguous match: if two + recorded calls produced byte-identical output we cannot tell which one + this continues, and guessing would attribute tokens to the wrong parent. + A unique match is required. + """ + fingerprint = assistant_fingerprint(messages) + if not fingerprint: + return None + call_ids = self.by_fingerprint.get(fingerprint) or [] + if len(call_ids) != 1: + return None + return self.by_call_id.get(call_ids[0]) + + def record(self, call_id: str, messages: list[dict], cum_tokens: list[int], digest: str) -> None: + """Index a completed call by the conversation a continuation of it would carry.""" + node = LineageNode(call_id=call_id, cum_tokens=list(cum_tokens), cum_len=len(cum_tokens), digest=digest) + self.by_call_id[call_id] = node + fingerprint = assistant_fingerprint(messages) + if fingerprint: + self.by_fingerprint.setdefault(fingerprint, []).append(call_id) + + +class LineageIndex: + """All live rollouts' lineage, bounded so an abandoned rollout cannot leak. + + A rollout that dies (harness crash, a batch discarded on a NaN-logprob + retry) is never read again, so eviction cannot rely on consumption. The cap + evicts whole rollouts in insertion order; losing one costs a fallback to + re-rendering, never a wrong answer. + """ + + def __init__(self, max_rollouts: int = 512) -> None: + self._max_rollouts = max_rollouts + self._rollouts: dict[str, RolloutLineage] = {} + + def for_rollout(self, rollout_id: str) -> RolloutLineage: + lineage = self._rollouts.get(rollout_id) + if lineage is None: + lineage = RolloutLineage() + self._rollouts[rollout_id] = lineage + while len(self._rollouts) > self._max_rollouts: + self._rollouts.pop(next(iter(self._rollouts))) + return lineage + + def drop(self, rollout_id: str) -> None: + self._rollouts.pop(rollout_id, None) + + def __len__(self) -> int: + return len(self._rollouts) + + +def assistant_turn_from_output_items(output_items: list[dict]) -> dict: + """Rebuild the assistant turn a harness will echo back, from a served response. + + Indexing a call means fingerprinting the conversation a *continuation* of it + would carry -- which is this request's messages plus the turn we just + produced. The next request only matches if that turn is reconstructed in the + shape the harness echoes: text content, and tool calls with their arguments. + + A mismatch is not a correctness problem, only a missed match: the lookup + falls back to a new root and the builder infers from token prefixes instead. + """ + text_parts: list[str] = [] + tool_calls: list[dict] = [] + for item in output_items or []: + if not isinstance(item, dict): + continue + kind = item.get("type") + if kind == "function_call": + tool_calls.append( + { + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", ""), + } + } + ) + continue + content = item.get("content") + # Responses items carry a list of blocks; a chat message carries a plain string. + if isinstance(content, str): + if content: + text_parts.append(content) + else: + for block in content or []: + if isinstance(block, dict) and isinstance(block.get("text"), str): + text_parts.append(block["text"]) + for block in item.get("summary") or []: + if isinstance(block, dict) and isinstance(block.get("text"), str): + text_parts.append(block["text"]) + for call in item.get("tool_calls") or []: + function = (call or {}).get("function") or {} + tool_calls.append( + {"function": {"name": function.get("name", ""), "arguments": function.get("arguments", "")}} + ) + turn: dict = {"role": "assistant", "content": "\n".join(text_parts)} + if tool_calls: + turn["tool_calls"] = tool_calls + return turn diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 3eb16549be..d5cf553845 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -37,9 +37,14 @@ from dataclasses import dataclass from typing import Any, Optional +from nemo_gym.token_id_capture.lineage import ( + LineageIndex, + assistant_turn_from_output_items, +) from nemo_gym.token_id_capture.protocols import TokenSink from nemo_gym.token_id_capture.records import ( TokenEntry, + cumulative_tokens, extract_token_fields, response_to_output_items, stamp_lineage, @@ -50,6 +55,16 @@ logger = logging.getLogger(__name__) +# Which recorded call each request continues, per rollout. Process-wide because the +# capture path is per request; bounded, so an abandoned rollout cannot leak. Losing an +# entry costs a fallback to prefix inference, never a wrong answer. +_LINEAGE = LineageIndex() + + +def lineage_index() -> LineageIndex: + return _LINEAGE + + @dataclass class CaptureContext: """What the capture middleware hands the model server for one call: which @@ -90,6 +105,7 @@ async def capture_tokens( response: Any, parent_call_id: Optional[str] = None, request_facts: Optional[dict] = None, + request_messages: Optional[list] = None, ) -> None: """Record a ``TokenEntry`` from a complete model response when a sink is set. @@ -111,6 +127,14 @@ async def capture_tokens( info = extract_token_fields(payload) if info is None: return + # Which call does this one continue? Resolved from the conversation the request already + # carries, so the harness needs to send nothing extra. A miss (new conversation, rewritten + # history, or two recorded calls with byte-identical output) leaves the link unset and the + # builder infers from token prefixes instead. + lineage = _LINEAGE.for_rollout(sink.rollout_id) + if parent_call_id is None and request_messages is not None: + parent = lineage.resolve(request_messages) + parent_call_id = parent.call_id if parent is not None else None try: entry = TokenEntry( rollout_id=sink.rollout_id, @@ -130,6 +154,15 @@ async def capture_tokens( # link is filled only when the model server resolved one. stamp_lineage(entry, parent_call_id) await sink.sink.put(entry) + # Index this call by the conversation a continuation of it would carry, so the next + # request resolves to it. + if request_messages is not None: + lineage.record( + sink.model_call_id, + list(request_messages) + [assistant_turn_from_output_items(entry.output_items)], + cumulative_tokens(entry), + entry.digest or "", + ) except Exception: # Capture is best-effort per call: a bad token payload must never fail the # model call and break the harness's run. But a rollout that lost a call diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 9eae219c56..830e2a89af 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -55,6 +55,7 @@ stamp_lineage, ) from nemo_gym.token_id_capture import reader as reader_module +from nemo_gym.token_id_capture.lineage import LineageIndex, RolloutLineage, assistant_fingerprint from nemo_gym.token_id_capture.reader import HttpTokenReader, LocalTokenReader from nemo_gym.token_id_capture.routes import make_token_store @@ -490,3 +491,164 @@ def test_capture_records_the_model_the_harness_asked_for(tmp_path): assert entry.requested_model == "claude-3-5-haiku-20241022" # The served model is a different field and is unaffected. assert entry.model != entry.requested_model + + +# --- lineage: which recorded call does a request continue? -------------------- + + +def test_fingerprint_ignores_non_assistant_turns(): + """Only assistant turns identify lineage: they are what we produced. User and + tool content varies with the environment and is irrelevant.""" + a = assistant_fingerprint([{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}]) + b = assistant_fingerprint([{"role": "user", "content": "DIFFERENT"}, {"role": "assistant", "content": "a"}]) + assert a == b != "" + # No assistant turn at all is a new conversation, not a match. + assert assistant_fingerprint([{"role": "user", "content": "q"}]) == "" + + +def test_fingerprint_survives_tool_argument_reserialization(): + """Harnesses re-serialize tool-call arguments between turns -- compact one + turn, pretty-printed the next. Without canonicalization the same call would + not compare equal to itself and every tool-using turn would miss.""" + compact = [ + {"role": "assistant", "content": "", "tool_calls": [{"function": {"name": "f", "arguments": '{"b":1,"a":2}'}}]} + ] + pretty = [ + { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "f", "arguments": '{\n "a": 2,\n "b": 1\n}'}}], + } + ] + assert assistant_fingerprint(compact) == assistant_fingerprint(pretty) + + +def test_lineage_resolves_the_parent_across_a_turn(): + lineage = RolloutLineage() + first_request = [{"role": "user", "content": "hello"}] + lineage.record("call-1", first_request + [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + + # The next request echoes the assistant turn, as any harness must to continue. + second_request = first_request + [{"role": "assistant", "content": "hi"}, {"role": "user", "content": "more"}] + parent = lineage.resolve(second_request) + assert parent is not None and parent.call_id == "call-1" + assert parent.cum_tokens == [1, 2, 3] and parent.cum_len == 3 + + +def test_lineage_misses_on_a_rewritten_history(): + """A compacted or rewritten context is a new root, not a wrong parent.""" + lineage = RolloutLineage() + lineage.record("call-1", [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + assert lineage.resolve([{"role": "assistant", "content": "a summary of the above"}]) is None + + +def test_lineage_refuses_an_ambiguous_parent(): + """Two recorded calls with byte-identical output cannot be told apart. Guessing + would attribute tokens to the wrong parent, so a unique match is required.""" + lineage = RolloutLineage() + messages = [{"role": "assistant", "content": "same"}] + lineage.record("call-a", messages, [1, 2], "da") + lineage.record("call-b", messages, [3, 4], "db") + assert lineage.resolve(messages) is None + + +def test_lineage_is_a_tree_so_forks_get_the_parent_not_the_previous_call(): + """Two sub-agents branching from one parent must BOTH resolve to that parent. + + A running cursor ("the last call") would hand the second branch a prefix + containing the first branch's generation, and the splice applies a supplied + prefix unconditionally -- so that would be silently wrong, not just wasteful. + """ + lineage = RolloutLineage() + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 4], + "da", + ) + + # The second branch continues the PARENT, not branch-a. + second = shared + [{"role": "user", "content": "b"}] + parent = lineage.resolve(second) + assert parent is not None and parent.call_id == "parent" + assert parent.cum_tokens == [1, 2, 3] + + +def test_lineage_index_is_bounded(): + """An abandoned rollout is never read again, so eviction cannot wait for + consumption. Losing an entry costs a fallback, never a wrong answer.""" + index = LineageIndex(max_rollouts=3) + for i in range(10): + index.for_rollout(f"r{i}") + assert len(index) == 3 + + +def test_served_calls_link_to_their_parent(tmp_path): + """End to end: a second call whose request echoes the first call's assistant + turn is recorded with parent_call_id pointing at it.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + first = [{"role": "user", "content": "hello"}] + client.post("/ng-rollout/lin0-roll0/v1/chat/completions", json={"messages": first}) + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 1 and entries[0].parent_call_id is None + + content = entries[0].output_items[0]["content"] + served_text = content if isinstance(content, str) else content[0]["text"] + second = first + [{"role": "assistant", "content": served_text}, {"role": "user", "content": "more"}] + client.post("/ng-rollout/lin0-roll0/v1/chat/completions", json={"messages": second}) + + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 2 + assert entries[1].parent_call_id == entries[0].model_call_id + + +def test_fingerprint_matches_across_openai_and_anthropic_tool_shapes(): + """A tool-using turn must match itself across dialects. + + We record the turn we produced in OpenAI shape (``tool_calls``), but Claude + Code echoes it back in Anthropic shape (``content`` blocks of type + ``tool_use``). If those hash differently, the parent is never resolved for + exactly the turns that create multi-turn rollouts -- which is what happened + in the first live multi-turn run: every record came back with + ``parent_call_id: None`` even though the calls chained perfectly. + """ + recorded = [ + { + "role": "assistant", + "content": "Let me compute that.", + "tool_calls": [{"function": {"name": "Bash", "arguments": '{"command":"echo 6"}'}}], + } + ] + echoed = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me compute that."}, + {"type": "tool_use", "name": "Bash", "input": {"command": "echo 6"}}, + ], + } + ] + assert assistant_fingerprint(recorded) == assistant_fingerprint(echoed) != "" + + +def test_lineage_resolves_a_tool_using_turn_echoed_in_anthropic_shape(): + lineage = RolloutLineage() + produced = { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "Bash", "arguments": '{"command":"factor 420"}'}}], + } + lineage.record("call-1", [{"role": "user", "content": "factor 420"}, produced], [1, 2, 3], "d1") + + # The harness continues the conversation, echoing the turn as Anthropic blocks + # and appending the tool result. + next_request = [ + {"role": "user", "content": "factor 420"}, + {"role": "assistant", "content": [{"type": "tool_use", "name": "Bash", "input": {"command": "factor 420"}}]}, + {"role": "user", "content": [{"type": "tool_result", "content": "420: 2 2 3 5 7"}]}, + ] + parent = lineage.resolve(next_request) + assert parent is not None and parent.call_id == "call-1" + assert parent.cum_tokens == [1, 2, 3] From bca52ad930847ea3abf01368b1bbd4595719acfa Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 28 Jul 2026 12:55:43 -0700 Subject: [PATCH 6/9] feat(vllm-model): supply the engine the previous call's exact tokens 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 --- nemo_gym/token_id_capture/__init__.py | 16 +- nemo_gym/token_id_capture/records.py | 6 + nemo_gym/token_id_capture/sink.py | 9 + responses_api_models/vllm_model/app.py | 71 ++++++- .../configs/vllm_model_supply_prefix.yaml | 11 + .../vllm_model/tests/test_app.py | 196 ++++++++++++++++++ 6 files changed, 301 insertions(+), 8 deletions(-) create mode 100644 responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index d6b11c3f1c..b435729a51 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -41,6 +41,12 @@ project_main_chain_response, ) from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.consumer import ( + clear_token_captures_for_rollouts, + token_id_capture_dirs_from_config, + trajectories_for_rollout, + trajectories_from_source, +) from nemo_gym.token_id_capture.lineage import ( LineageIndex, LineageNode, @@ -48,12 +54,6 @@ assistant_fingerprint, canonicalize_tool_arguments, ) -from nemo_gym.token_id_capture.consumer import ( - clear_token_captures_for_rollouts, - token_id_capture_dirs_from_config, - trajectories_for_rollout, - trajectories_from_source, -) from nemo_gym.token_id_capture.protocols import ( TokenSink, TokenSource, @@ -73,6 +73,8 @@ CaptureContext, TokenCaptureContext, capture_tokens, + current_capture_context, + lineage_index, reset_token_sink, set_token_sink, ) @@ -100,6 +102,8 @@ "set_token_sink", "reset_token_sink", "capture_tokens", + "current_capture_context", + "lineage_index", "CaptureTokenSource", "LineageIndex", "LineageNode", diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 7f371ac32d..b42248cac5 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -120,6 +120,12 @@ class TokenEntry(BaseModel): # O(N^2) prefix scan into an O(1) lookup, and let a claimed parent be # verified rather than trusted. parent_call_id: Optional[str] = None + # Whether the model server handed the engine this call's prefix verbatim + # (prefix supply) rather than letting the chat template re-render it. Recorded + # per call so a run can be audited after the fact: supply only fires on a + # unique, verified parent, so the supplied/total ratio is the honest measure + # of how often it applied versus fell back. + prefix_supplied: bool = False # len(prompt_token_ids) + len(generation_token_ids) for THIS call: the # length of the prefix a child of this call must start with. cum_len: Optional[int] = None diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index d5cf553845..92700f96a6 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -79,6 +79,9 @@ class CaptureContext: model_call_id: str store: TokenCaptureStore model: str = "" + # Set by the model server when it supplied this call's prefix; read back at + # capture time so the evidence lands on the record rather than only in a log. + prefix_supplied: bool = False @property def sink(self) -> TokenSink: @@ -97,6 +100,11 @@ def set_token_sink(sink: CaptureContext) -> Token: return _TOKEN_SINK.set(sink) +def current_capture_context() -> Optional[CaptureContext]: + """The capture context for the in-flight call, or None for untagged traffic.""" + return _TOKEN_SINK.get() + + def reset_token_sink(token: Token) -> None: _TOKEN_SINK.reset(token) @@ -148,6 +156,7 @@ async def capture_tokens( # reads is not token-only -- text-based penalties need it. output_items=response_to_output_items(payload), created_at=time.time(), + prefix_supplied=sink.prefix_supplied, **(request_facts or {}), ) # cum_len/digest describe this call and are always computable; the parent diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 792a2e2331..d462a40d8b 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -15,6 +15,7 @@ import asyncio import base64 import json +import logging import os from copy import deepcopy from time import time @@ -22,7 +23,10 @@ from aiohttp.client_exceptions import ClientResponseError from fastapi import Request -from pydantic import Field +from pydantic import Field, PrivateAttr + + +LOG = logging.getLogger(__name__) from nemo_gym.base_responses_api_model import ( BaseResponsesAPIModelConfig, @@ -44,6 +48,7 @@ split_responses_input_output_items, # noqa: F401 ) from nemo_gym.server_utils import SESSION_ID_KEY, is_nemo_gym_fastapi_entrypoint +from nemo_gym.token_id_capture import current_capture_context, lineage_index class VLLMModelConfig(BaseResponsesAPIModelConfig): @@ -59,6 +64,11 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # Whether or not the model can generate a reasoning output, and called again to produce additional reasoning output. sequential_reasoning_allowed: bool = True + # Supply the engine the exact tokens of the call this request continues, instead of letting the + # chat template re-render them from text. See _apply_prefix_supply below. Off by default: it + # requires a backend that honours required_prefix_token_ids (NeMo-RL's vLLM worker does). + supply_prefix_token_ids: bool = False + # As of Feb 2026, we default this to False since majority of open source models aren't responses native with the exception of GPT-OSS is_responses_native: bool = False @@ -448,6 +458,56 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict: # No user message found — create one with just the audio blocks. body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)}) + body_dict = self._apply_prefix_supply(body_dict) + + return body_dict + + # [supplied, eligible]. A plain list so it is mutable without a pydantic field. + _prefix_supply_counts: List[int] = PrivateAttr(default_factory=lambda: [0, 0]) + + def _apply_prefix_supply(self, body_dict: Dict[str, Any]) -> Dict[str, Any]: + """Hand the engine the previous call's exact tokens, so this prompt extends them. + + Without this the engine builds every prompt by re-rendering the whole + conversation through the chat template. Two things go wrong. Re-tokenizing + an assistant turn can produce a different (equally valid) split than the + one the model sampled, so the new prompt does not extend the old + prompt-plus-generation. And for a reasoning model the template drops + earlier thinking entirely, so the tokens are not merely re-split, they are + gone. Either way the trajectory cannot be chained and training silently + collapses to the first call. + + Supplying the parent's cumulative ids makes NeMo-RL's splice keep them + verbatim and append only the newly rendered tail, so contiguity holds by + construction and stripped reasoning is restored. + + The rule is conservative on purpose: the splice applies whatever it is + given without checking that it belongs to this conversation, so a wrong + prefix would silently generate from a conversation the harness never + asked for. Supply only on a unique, verified parent; otherwise send the + request untouched and let it start a new chain. A fallback costs a cold + KV cache and a shorter trained chain -- never a wrong answer. + """ + if not self.config.supply_prefix_token_ids: + return body_dict + self._prefix_supply_counts[1] += 1 + context = current_capture_context() + if context is None: + # Not a correlated rollout call, so there is no lineage to supply from. + return body_dict + parent = lineage_index().for_rollout(context.rollout_id).resolve(body_dict.get("messages") or []) + if parent is None or not parent.cum_tokens: + return body_dict + body_dict["required_prefix_token_ids"] = list(parent.cum_tokens) + # Record that this call's prefix was supplied, both on the durable record (so a run can be + # audited afterwards from the capture files) and as a running ratio in the log. Supply only + # fires on a unique, verified parent, so supplied/total is the honest measure of how often + # it applied rather than falling back to re-rendering. + context.prefix_supplied = True + self._prefix_supply_counts[0] += 1 + supplied, total = self._prefix_supply_counts[0], self._prefix_supply_counts[1] + if total and supplied % 10 == 0: + LOG.info("prefix supply: %d/%d calls supplied (%.0f%%)", supplied, total, 100.0 * supplied / total) return body_dict async def chat_completions( @@ -544,8 +604,15 @@ async def chat_completions( # when tokenizing, otherwise `prompt_token_ids` (and therefore logged # `prompt_str`) can be built with different chat template settings than # the actual generation request. + # + # ``required_prefix_token_ids`` is here for the same reason. Generation + # applies the supplied prefix, so without it here the recorded + # prompt_token_ids is a plain re-render of the conversation and does not + # extend the previous call -- the chain looks broken even though the + # engine generated from the right tokens. Measured: 0/9 links contiguous + # with supply on until this key was forwarded. tokenize_body_dict = dict() - for key in ("model", "messages", "tools", "chat_template_kwargs"): + for key in ("model", "messages", "tools", "chat_template_kwargs", "required_prefix_token_ids"): if key in body_dict: tokenize_body_dict[key] = body_dict[key] diff --git a/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml new file mode 100644 index 0000000000..910808f57f --- /dev/null +++ b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml @@ -0,0 +1,11 @@ +# A/B overlay: turn prefix supply ON (it is off by default in the shipped config). +# +# The model server hands the engine the exact tokens of the call this request continues +# (required_prefix_token_ids), so the new prompt EXTENDS the previous prompt+generation by +# construction instead of being re-rendered from text. Requires a unique, verified parent +# from the lineage index; otherwise it forwards the request untouched and starts a new +# chain -- wasteful, never wrong. +policy_model: + responses_api_models: + vllm_model: + supply_prefix_token_ids: true diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index 2ee728f090..6b911a119a 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/responses_api_models/vllm_model/tests/test_app.py @@ -53,6 +53,13 @@ NeMoGymSummary, ) from nemo_gym.server_utils import ServerClient +from nemo_gym.token_id_capture import ( + CaptureContext, + TokenCaptureStore, + lineage_index, + reset_token_sink, + set_token_sink, +) from responses_api_models.vllm_model.app import ( VLLMConverter, VLLMModel, @@ -4684,3 +4691,192 @@ async def mock_create_chat_completion(**kwargs): ) # The tokenize endpoint must not be reached once the contract check fails. mock_client.create_tokenize.assert_not_called() + + +class TestPrefixSupply: + """Supplying the engine the previous call's exact tokens. + + Without this the engine re-renders every prompt from text: a re-tokenized + assistant turn can split differently than the model sampled, and a reasoning + model's template drops earlier thinking entirely. Either way the new prompt + does not extend the old one and the trajectory cannot be chained. + """ + + @staticmethod + def _server(monkeypatch: MonkeyPatch, *, enabled: bool) -> VLLMModel: + config = VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=False, + uses_reasoning_parser=False, + supply_prefix_token_ids=enabled, + ) + get_global_config_dict_mock = MagicMock(return_value={}) + monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", get_global_config_dict_mock) + return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + + @staticmethod + def _armed(rollout_id: str, tmp_path) -> CaptureContext: + return CaptureContext(rollout_id=rollout_id, model_call_id="call-x", store=TokenCaptureStore(tmp_path)) + + def test_supplies_the_parents_cumulative_tokens(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + first_turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-0").record("parent", first_turn, [1, 2, 3, 4], "d") + + token = set_token_sink(self._armed("sup-0", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": first_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3, 4] + + def test_off_by_default(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=False) + turn = [{"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-1").record("parent", turn, [1, 2], "d") + + token = set_token_sink(self._armed("sup-1", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_uncorrelated_call_is_left_alone(self, monkeypatch: MonkeyPatch) -> None: + server = self._server(monkeypatch, enabled=True) + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "hello"}]}) + assert "required_prefix_token_ids" not in out + + def test_fingerprint_miss_falls_back_rather_than_supplying_something_wrong( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """A rewritten history must send the request untouched. The splice applies + whatever it is given without checking it belongs to this conversation, so + a wrong prefix would silently generate from a conversation the harness + never asked for.""" + server = self._server(monkeypatch, enabled=True) + lineage_index().for_rollout("sup-2").record("parent", [{"role": "assistant", "content": "hello"}], [1, 2], "d") + + token = set_token_sink(self._armed("sup-2", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "a summary"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_ambiguous_parent_is_not_supplied(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + turn = [{"role": "assistant", "content": "same"}] + lineage = lineage_index().for_rollout("sup-3") + lineage.record("a", turn, [1, 2], "da") + lineage.record("b", turn, [3, 4], "db") + + token = set_token_sink(self._armed("sup-3", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_a_fork_gets_the_parents_prefix_not_the_previous_calls(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + """Two sub-agents branching from one parent must both get cum(parent). + + A running cursor would hand the second branch a prefix containing the + first branch's generation -- which the splice would apply without + complaint, generating from a conversation that never happened. + """ + server = self._server(monkeypatch, enabled=True) + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage = lineage_index().for_rollout("sup-4") + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 9, 9], + "da", + ) + + token = set_token_sink(self._armed("sup-4", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": shared + [{"role": "user", "content": "b"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + + def test_reasoning_stripped_history_still_supplies_the_real_tokens( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """The case supplying exists for. + + A reasoning model's chat template drops earlier thinking when it renders + a later prompt, so the re-rendered prompt cannot extend the previous + prompt-plus-generation -- the tokens are gone, not merely re-split. The + supplied prefix comes from the recording, so those tokens are restored + and the chain survives. + """ + server = self._server(monkeypatch, enabled=True) + # The recorded turn included reasoning; what the harness echoes back does not. + recorded_turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "answer"}] + real_tokens_including_reasoning = [1, 2, 3, 4, 5, 6, 7] + lineage_index().for_rollout("sup-5").record("parent", recorded_turn, real_tokens_including_reasoning, "d") + + token = set_token_sink(self._armed("sup-5", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": recorded_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == real_tokens_including_reasoning + + +class TestPrefixSupplyAccounting: + """Supply must be auditable after the fact. + + Without this the only evidence supply fired is that chains happen to be + contiguous -- which they often are anyway, so a supply run and a non-supply + run look identical. That ambiguity made the first live supply experiment + inconclusive. + """ + + def test_supplied_call_is_marked_on_the_capture_context(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + lineage_index().for_rollout("acct-0").record("parent", turn, [1, 2, 3], "d") + + ctx = CaptureContext(rollout_id="acct-0", model_call_id="c", store=TokenCaptureStore(tmp_path)) + token = set_token_sink(ctx) + try: + out = server._apply_prefix_supply({"messages": turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + assert ctx.prefix_supplied is True + assert server._prefix_supply_counts == [1, 1] # 1 supplied of 1 eligible + + def test_fallback_counts_as_eligible_but_not_supplied(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + lineage_index().for_rollout("acct-1").record("parent", [{"role": "assistant", "content": "a"}], [1, 2], "d") + + ctx = CaptureContext(rollout_id="acct-1", model_call_id="c", store=TokenCaptureStore(tmp_path)) + token = set_token_sink(ctx) + try: + # A rewritten history: no unique parent, so supply must decline. + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "rewritten"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + assert ctx.prefix_supplied is False + assert server._prefix_supply_counts == [0, 1] # eligible, not supplied From 5a9e97a96ed645816d3c1e6d58723db950339b01 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 28 Jul 2026 12:58:22 -0700 Subject: [PATCH 7/9] feat(token-id-capture): let the token read route require a bearer token 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 Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 7 ++++- nemo_gym/token_id_capture/config.py | 7 +++++ nemo_gym/token_id_capture/routes.py | 33 ++++++++++++++++++++--- tests/unit_tests/test_token_id_capture.py | 25 +++++++++++++++++ 4 files changed, 68 insertions(+), 4 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 1ca25aa3f1..20a57f5e5e 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -69,6 +69,7 @@ ) from nemo_gym.token_id_capture import ( CaptureContext, + TokenIdCaptureConfig, capture_tokens, reset_token_sink, set_token_sink, @@ -1314,7 +1315,11 @@ def install_model_call_capture( token_store=token_store, ) if token_store is not None: - install_token_capture_routes(app, token_store) + install_token_capture_routes( + app, + token_store, + read_token=TokenIdCaptureConfig.model_validate(global_config_dict or {}).token_id_capture_read_token, + ) # --- Run-level capture helpers (rollout-collection side) --- diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 8c587822a7..08724cfba0 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -38,6 +38,13 @@ class TokenIdCaptureConfig(BaseModel): # Shared fallback directory (also used by evaluation capture). model_call_capture_dir: Optional[Path] = None + # Bearer token required by the token read route. The route serves a rollout's raw training + # tokens on the same app the harness calls to generate, which is fine inside a trusted cluster + # and not fine once the harness runs in a sandbox whose only egress is this server -- it could + # read its own training data, or another rollout's. Set this before the sandbox work; when unset + # the route stays open and logs a warning once. + token_id_capture_read_token: Optional[str] = None + @model_validator(mode="after") def _validate(self) -> "TokenIdCaptureConfig": if not self.token_id_capture_enabled: diff --git a/nemo_gym/token_id_capture/routes.py b/nemo_gym/token_id_capture/routes.py index dddcc400ce..c2844255be 100644 --- a/nemo_gym/token_id_capture/routes.py +++ b/nemo_gym/token_id_capture/routes.py @@ -24,14 +24,19 @@ from __future__ import annotations import asyncio +import logging +from hmac import compare_digest from typing import Any, Optional -from fastapi import APIRouter, HTTPException, Response +from fastapi import APIRouter, Header, HTTPException, Response from nemo_gym.token_id_capture.config import TokenIdCaptureConfig from nemo_gym.token_id_capture.store import TokenCaptureStore +logger = logging.getLogger(__name__) + + def make_token_store(global_config_dict: Any) -> Optional[TokenCaptureStore]: """Build the training-token store, or ``None`` when capture is disabled.""" config = TokenIdCaptureConfig.model_validate(global_config_dict) @@ -40,11 +45,33 @@ def make_token_store(global_config_dict: Any) -> Optional[TokenCaptureStore]: return TokenCaptureStore(config.resolved_dir()) -def install_token_capture_routes(app: Any, store: TokenCaptureStore) -> None: +def install_token_capture_routes(app: Any, store: TokenCaptureStore, read_token: Optional[str] = None) -> None: + """Register the token read route. + + The route serves a rollout's raw training tokens 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 it could read its own training data -- or another rollout's. + ``read_token`` requires a bearer token; when it is unset the route stays open + and warns once, so existing deployments keep working and the gap is visible. + """ + if not read_token: + logger.warning( + "The token-capture read route is unauthenticated. Anything that can reach this model " + "server can read captured training tokens for any rollout. Set token_id_capture_read_token " + "before running an agent harness in a sandbox whose only egress is this server." + ) + router = APIRouter() @router.get("/ng-capture/tokens/{rollout_id}") - async def get_tokens(rollout_id: str) -> Response: + async def get_tokens(rollout_id: str, authorization: Optional[str] = Header(default=None)) -> Response: + if read_token: + expected = f"Bearer {read_token}" + # Constant-time compare: the token is the only thing standing between an untrusted + # harness and every rollout's training data. + if authorization is None or not compare_digest(authorization, expected): + raise HTTPException(status_code=401, detail="token capture read route requires a bearer token") try: entries = await asyncio.to_thread(store.read_entries, rollout_id) except ValueError as exc: diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 830e2a89af..cc2fc21044 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -652,3 +652,28 @@ def test_lineage_resolves_a_tool_using_turn_echoed_in_anthropic_shape(): parent = lineage.resolve(next_request) assert parent is not None and parent.call_id == "call-1" assert parent.cum_tokens == [1, 2, 3] + + +def test_read_route_requires_a_token_when_one_is_configured(tmp_path): + """The route serves raw training tokens on the same app the harness calls to + generate. Inside a trusted cluster that is fine; once the harness runs in a + sandbox whose only egress is this server it could read its own training data, + or another rollout's.""" + config = dict(_both_enabled(tmp_path)) + config["token_id_capture_read_token"] = "s3cret" # pragma: allowlist secret + client = TestClient(_server(config).setup_webserver()) + client.post("/ng-rollout/auth0-roll0/v1/responses", json={"input": "hi"}) + + assert client.get("/ng-capture/tokens/auth0-roll0").status_code == 401 + assert client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer wrong"}).status_code == 401 + + ok = client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer s3cret"}) + assert ok.status_code == 200 + assert TokenEntry.model_validate_json(ok.text.splitlines()[0]).generation_token_ids == GTOKS + + +def test_read_route_stays_open_when_no_token_is_configured(tmp_path): + """Existing deployments keep working; the gap is logged rather than enforced.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/auth1-roll0/v1/responses", json={"input": "hi"}) + assert client.get("/ng-capture/tokens/auth1-roll0").status_code == 200 From f507b6f2d724516a8d99646dabc52182b90ea80c Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:33:05 -0700 Subject: [PATCH 8/9] feat(vllm-model): on-policy sampling pin via sampling_overrides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- responses_api_models/vllm_model/app.py | 13 +++++++++++++ .../vllm_model/configs/vllm_model_for_training.yaml | 7 +++++++ 2 files changed, 20 insertions(+) diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index d462a40d8b..62f1a31670 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -77,6 +77,13 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # Corresponds to the extra_body of OpenAI Client. extra_body: Optional[Dict[str, Any]] = None + # Sampling params to force on every chat request, overriding whatever the client sent. An + # external harness (e.g. a CLI agent) chooses its own temperature/top_p, but on-policy RL + # training requires generation to match the sampling distribution the policy is optimized under. + # The integrating training framework sets these to its generation sampling params; Gym only + # enforces them and holds no knowledge of any specific framework. Unset means no override. + sampling_overrides: Optional[Dict[str, Any]] = None + default_headers: Dict[str, str] = Field(default_factory=dict) # Optional prefix for resolving relative ``metadata.audio_path`` (or # entries in ``metadata.audio_paths``) against. Absolute paths are used @@ -458,6 +465,12 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict: # No user message found — create one with just the audio blocks. body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)}) + # Pin sampling params last so they win over anything the client sent. On-policy RL training + # requires the generation to match the training worker's sampling config; an external harness + # sets its own temperature/top_p, so force them here to keep captured rollouts on-policy. + if self.config.sampling_overrides: + body_dict.update(self.config.sampling_overrides) + body_dict = self._apply_prefix_supply(body_dict) return body_dict diff --git a/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml b/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml index 70727036c7..b0e30386c8 100644 --- a/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml +++ b/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml @@ -7,3 +7,10 @@ policy_model: model: ${policy_model_name} return_token_id_information: true uses_reasoning_parser: true + # On-policy training: force generation to the sampling params the policy is optimized under, + # overriding whatever an external harness requests, so captured rollouts stay on-policy. The + # integrating training framework supplies these via the generic keys below (Gym does not know + # any framework); the defaults are standard on-policy sampling when a framework sets nothing. + sampling_overrides: + temperature: ${oc.select:policy_generation_temperature,1.0} + top_p: ${oc.select:policy_generation_top_p,1.0} From 3860daf164589a61efcf1cb336caae4d76d7823d Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:33:40 -0700 Subject: [PATCH 9/9] feat(token-id-capture): Claude Code external-harness example 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 --- resources_servers/reasoning_gym/app.py | 9 +++++++++ .../configs/reasoning_gym_claude_code_agent.yaml | 3 +++ ...asoning_gym_claude_code_agent_model_server.yaml | 14 +++++++++++++- .../configs/claude_code_agent.yaml | 3 +++ 4 files changed, 28 insertions(+), 1 deletion(-) diff --git a/resources_servers/reasoning_gym/app.py b/resources_servers/reasoning_gym/app.py index 987d5407bc..1f490d0f42 100644 --- a/resources_servers/reasoning_gym/app.py +++ b/resources_servers/reasoning_gym/app.py @@ -108,6 +108,15 @@ def _extract_answer_from_response(self, response) -> str: if boxed_match: return boxed_match.group(1).strip() + # Conversational agents (e.g. Claude Code driving the policy) often do not use /\boxed + # and instead state the final answer as the last quoted string, e.g. + # So the final answer is: "Zoey is a fool, and Riley is a sage." + # Passing the whole chain-of-thought to the task scorer lets it parse an intermediate + # hypothesis and mis-score a correct rollout, so isolate that trailing quoted answer. + quoted = re.findall(r'["\u201c]([^"\u201d]{3,}?)["\u201d]', full_text) + if quoted: + return quoted[-1].strip() + # return full text if or \boxed{} fail return full_text.strip() if full_text.strip() else "" diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml index e224bcb71b..c62d4f74ab 100644 --- a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml @@ -14,6 +14,9 @@ reasoning_gym_claude_code_agent: resources_server: type: resources_servers name: reasoning_gym + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: ${anthropic_model_name} anthropic_api_key: ${anthropic_api_key} diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml index 10a60431aa..512efeeee5 100644 --- a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml @@ -32,6 +32,9 @@ reasoning_gym_claude_code_agent_model_server: model_server: type: responses_api_models name: policy_model + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: ${policy_model_name} anthropic_api_key: EMPTY # pragma: allowlist secret @@ -39,9 +42,18 @@ reasoning_gym_claude_code_agent_model_server: max_turns: 30 timeout: 300 thinking: disabled + # Restrict the CLI to Bash. Claude Code's system prompt is dominated by TOOL DEFINITIONS, and + # every tool it is allowed adds its schema to the prompt of every call. reasoning_gym only + # needs Bash (the system prompt below says so), so this cuts the fixed prompt cost + # substantially -- which is what keeps prompt + generation inside max_model_len. Measure the + # real cost after a run: the first captured call's prompt_token_ids IS the system prompt plus + # tools plus the first user message. + allowed_tools: "Bash" system_prompt: | You are a precise reasoning assistant. You have access to Bash to run Python for calculations. - For every problem: think step by step, use code to verify when helpful, and state your final answer clearly. + For every problem: think step by step, use code to verify when helpful, then end your reply + with ONLY your final answer inside ... tags, matching the exact format the + task asks for (reasoning_gym scores the text inside ). datasets: - name: example type: example diff --git a/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml index 5e3140bea7..41f50e34ab 100644 --- a/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml +++ b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml @@ -5,6 +5,9 @@ claude_code_agent: resources_server: type: resources_servers name: ??? + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: claude-sonnet-4-6 anthropic_api_key: ${anthropic_api_key}