From 767c48834bc96f92c71752169000495487659d63 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 23 Jul 2026 08:31:30 -0700 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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