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/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 99426db471..20a57f5e5e 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -67,6 +67,17 @@ BaseServer, SimpleServer, ) +from nemo_gym.token_id_capture import ( + CaptureContext, + TokenIdCaptureConfig, + 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__) @@ -76,6 +87,47 @@ _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 + + +def _request_messages(body: Any) -> list[dict]: + """The conversation a request carries, across the three dialects. + + Used only to identify which recorded call this request continues (see + ``token_id_capture.lineage``): the assistant turns in it are the ones we + produced. Chat and Anthropic both use ``messages``; Responses carries + ``input``, which is a string for a first turn and a list of items after that. + """ + if body is None: + return [] + getter = body.get if isinstance(body, dict) else lambda key, default=None: getattr(body, key, default) + messages = getter("messages", None) + if isinstance(messages, list): + return [m if isinstance(m, dict) else m.model_dump() for m in messages if m is not None] + items = getter("input", None) + if isinstance(items, list): + return [i if isinstance(i, dict) else i.model_dump() for i in items if i is not None] + return [] + + class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): pass @@ -90,7 +142,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 +249,15 @@ 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, + request_messages=_request_messages(params), + request_facts=_request_facts(params), + ) + return completion async def messages(self, request: Request, body: dict = Body()): """Default Anthropic Messages <-> Responses mapping shared by every Gym model server. @@ -206,6 +270,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 +287,18 @@ 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, + request_messages=_request_messages(params), + request_facts=_request_facts(params), + ) + return response def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming: @@ -1003,10 +1078,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 +1105,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 +1216,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 +1287,39 @@ 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, + read_token=TokenIdCaptureConfig.model_validate(global_config_dict or {}).token_id_capture_read_token, + ) # --- Run-level capture helpers (rollout-collection side) --- diff --git a/nemo_gym/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/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 new file mode 100644 index 0000000000..b435729a51 --- /dev/null +++ b/nemo_gym/token_id_capture/__init__.py @@ -0,0 +1,126 @@ +# 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 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 +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. +""" + +from nemo_gym.token_id_capture.builder import ( + BuildOutput, + Chain, + Trajectory, + assert_nemo_rl_contiguity, + build_trajectories, + classify_side_calls, + 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 ( + clear_token_captures_for_rollouts, + token_id_capture_dirs_from_config, + trajectories_for_rollout, + trajectories_from_source, +) +from nemo_gym.token_id_capture.lineage import ( + LineageIndex, + LineageNode, + RolloutLineage, + assistant_fingerprint, + canonicalize_tool_arguments, +) +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, + current_capture_context, + lineage_index, + 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", + "current_capture_context", + "lineage_index", + "CaptureTokenSource", + "LineageIndex", + "LineageNode", + "RolloutLineage", + "assistant_fingerprint", + "canonicalize_tool_arguments", + "build_trajectories", + "classify_side_calls", + "per_request", + "prefix_merging", + "project_main_chain_response", + "assert_nemo_rl_contiguity", + "Trajectory", + "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/builder.py b/nemo_gym/token_id_capture/builder.py new file mode 100644 index 0000000000..295894b0a1 --- /dev/null +++ b/nemo_gym/token_id_capture/builder.py @@ -0,0 +1,505 @@ +# 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 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 = [ + 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/config.py b/nemo_gym/token_id_capture/config.py new file mode 100644 index 0000000000..08724cfba0 --- /dev/null +++ b/nemo_gym/token_id_capture/config.py @@ -0,0 +1,62 @@ +# 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 + + # Bearer token required by the token read route. The route serves a rollout's raw training + # tokens on the same app the harness calls to generate, which is fine inside a trusted cluster + # and not fine once the harness runs in a sandbox whose only egress is this server -- it could + # read its own training data, or another rollout's. Set this before the sandbox work; when unset + # the route stays open and logs a warning once. + token_id_capture_read_token: Optional[str] = None + + @model_validator(mode="after") + def _validate(self) -> "TokenIdCaptureConfig": + if not self.token_id_capture_enabled: + 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/consumer.py b/nemo_gym/token_id_capture/consumer.py new file mode 100644 index 0000000000..010070da57 --- /dev/null +++ b/nemo_gym/token_id_capture/consumer.py @@ -0,0 +1,214 @@ +# 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, + classify_side_calls, + 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 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], + builder: str, + 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 + # 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), "side_calls_excluded": len(side_calls)}, + } + + 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", {}), + # Excluded, not dropped silently. + "side_calls_excluded": len(side_calls), + } + 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 = "", + side_call_model_patterns: tuple[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, 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. + 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 = "", + 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, side_call_model_patterns) diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py new file mode 100644 index 0000000000..42e0a54333 --- /dev/null +++ b/nemo_gym/token_id_capture/lineage.py @@ -0,0 +1,262 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Which recorded call does this request continue? + +The builder can infer that after the fact, by finding the earlier call whose +token sequence is the longest prefix of this call's prompt. That works, but it +cannot do two things: + +- **Resolve a retry.** Capture records a call once the response is assembled, + including one the client never received, so a harness retry leaves two records + with the same prompt and different generations. Both are equally valid + children, so inference has to guess between them. +- **Run before the call.** Inference happens after the rollout. Supplying the + engine the previous call's exact tokens (so the next prompt extends them by + construction instead of being re-rendered from text) has to happen *before* the + request is forwarded, when there is nothing yet to infer from. + +This module answers the question at request time instead, using only what the +harness already sends. A harness must echo the conversation to continue it, so +the assistant messages in a request are the ones we produced: hashing them in +order identifies the call that produced the last one. Nothing is added to the +wire and nothing depends on the harness preserving a field we invented. + +The index is a map keyed by call, not a running cursor, so it is a tree and +forks cost nothing. Two sub-agents branching from one parent both resolve to it +and both get the same prefix -- which is exactly right, and is why the prefix to +supply is ``cum(parent)`` and not ``cum(previous call)``. Entries are added and +never mutated, so concurrent sub-agents cannot corrupt each other's lineage. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from typing import Any, Optional + + +_FINGERPRINT_DOMAIN = b"nemo-gym-lineage" + + +def canonicalize_tool_arguments(value: Any) -> str: + """Normalize a tool call's arguments for comparison only. + + Harnesses re-serialize tool-call arguments between turns -- compact one turn, + pretty-printed the next -- so the same call does not compare equal to itself. + Comparison is done on sorted-key, separator-normalized JSON; the model's + original string is what stays in the record and is never rewritten. + """ + if value is None: + return "" + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, ValueError): + return value.strip() + else: + parsed = value + try: + return json.dumps(parsed, sort_keys=True, separators=(",", ":")) + except (TypeError, ValueError): + return str(parsed) + + +def _text_of(content: Any) -> str: + """Flatten a message's *text* content, across the shapes the dialects use. + + Tool calls are deliberately NOT folded in here -- see ``_tools_of``. The two + dialects carry them in different places, so they have to be normalized + separately or the same turn hashes differently depending on which side is + looking at it. + """ + if content is None: + return "" + if isinstance(content, str): + return content + parts: list[str] = [] + if isinstance(content, list): + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict) and isinstance(block.get("text"), str): + parts.append(block["text"]) + return "\n".join(parts) + + +def _tools_of(message: dict) -> list[tuple[str, str]]: + """The tool calls in an assistant turn, as (name, canonical args). + + A turn we produced is recorded in OpenAI shape (``tool_calls`` with + ``function.name`` / ``function.arguments``), but Claude Code echoes it back + in Anthropic shape (``content`` blocks of ``type: tool_use`` with ``name`` / + ``input``). Both must reduce to the same list, or a tool-using turn -- the + only kind that produces multi-turn rollouts -- can never match itself. + """ + tools: list[tuple[str, str]] = [] + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_use": + tools.append((str(block.get("name", "")), canonicalize_tool_arguments(block.get("input")))) + for call in message.get("tool_calls") or []: + function = (call or {}).get("function") or {} + tools.append((str(function.get("name", "")), canonicalize_tool_arguments(function.get("arguments")))) + return tools + + +def assistant_fingerprint(messages: list[dict]) -> str: + """Fingerprint the assistant turns of a request, in order. + + Only assistant turns are hashed: they are what *we* produced, so they + identify the call that produced the last one. User and tool content varies + with the environment and is irrelevant to lineage. + """ + hasher = hashlib.sha256(_FINGERPRINT_DOMAIN) + count = 0 + for message in messages or []: + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + count += 1 + hasher.update(b"\x00") + hasher.update(_text_of(message.get("content")).encode("utf-8")) + for name, arguments in _tools_of(message): + hasher.update(b"\x01") + hasher.update(name.encode("utf-8")) + hasher.update(arguments.encode("utf-8")) + if count == 0: + return "" + return hasher.hexdigest() + + +@dataclass +class LineageNode: + call_id: str + cum_tokens: list[int] + cum_len: int + digest: str + + +@dataclass +class RolloutLineage: + """Per-rollout call index. Append-only, so forks and concurrency are safe.""" + + by_fingerprint: dict[str, list[str]] = field(default_factory=dict) + by_call_id: dict[str, LineageNode] = field(default_factory=dict) + + def resolve(self, messages: list[dict]) -> Optional[LineageNode]: + """The call this request continues, or ``None``. + + ``None`` for a new root (nothing matches -- a fresh conversation, or one + the harness rewrote) and, deliberately, for an ambiguous match: if two + recorded calls produced byte-identical output we cannot tell which one + this continues, and guessing would attribute tokens to the wrong parent. + A unique match is required. + """ + fingerprint = assistant_fingerprint(messages) + if not fingerprint: + return None + call_ids = self.by_fingerprint.get(fingerprint) or [] + if len(call_ids) != 1: + return None + return self.by_call_id.get(call_ids[0]) + + def record(self, call_id: str, messages: list[dict], cum_tokens: list[int], digest: str) -> None: + """Index a completed call by the conversation a continuation of it would carry.""" + node = LineageNode(call_id=call_id, cum_tokens=list(cum_tokens), cum_len=len(cum_tokens), digest=digest) + self.by_call_id[call_id] = node + fingerprint = assistant_fingerprint(messages) + if fingerprint: + self.by_fingerprint.setdefault(fingerprint, []).append(call_id) + + +class LineageIndex: + """All live rollouts' lineage, bounded so an abandoned rollout cannot leak. + + A rollout that dies (harness crash, a batch discarded on a NaN-logprob + retry) is never read again, so eviction cannot rely on consumption. The cap + evicts whole rollouts in insertion order; losing one costs a fallback to + re-rendering, never a wrong answer. + """ + + def __init__(self, max_rollouts: int = 512) -> None: + self._max_rollouts = max_rollouts + self._rollouts: dict[str, RolloutLineage] = {} + + def for_rollout(self, rollout_id: str) -> RolloutLineage: + lineage = self._rollouts.get(rollout_id) + if lineage is None: + lineage = RolloutLineage() + self._rollouts[rollout_id] = lineage + while len(self._rollouts) > self._max_rollouts: + self._rollouts.pop(next(iter(self._rollouts))) + return lineage + + def drop(self, rollout_id: str) -> None: + self._rollouts.pop(rollout_id, None) + + def __len__(self) -> int: + return len(self._rollouts) + + +def assistant_turn_from_output_items(output_items: list[dict]) -> dict: + """Rebuild the assistant turn a harness will echo back, from a served response. + + Indexing a call means fingerprinting the conversation a *continuation* of it + would carry -- which is this request's messages plus the turn we just + produced. The next request only matches if that turn is reconstructed in the + shape the harness echoes: text content, and tool calls with their arguments. + + A mismatch is not a correctness problem, only a missed match: the lookup + falls back to a new root and the builder infers from token prefixes instead. + """ + text_parts: list[str] = [] + tool_calls: list[dict] = [] + for item in output_items or []: + if not isinstance(item, dict): + continue + kind = item.get("type") + if kind == "function_call": + tool_calls.append( + { + "function": { + "name": item.get("name", ""), + "arguments": item.get("arguments", ""), + } + } + ) + continue + content = item.get("content") + # Responses items carry a list of blocks; a chat message carries a plain string. + if isinstance(content, str): + if content: + text_parts.append(content) + else: + for block in content or []: + if isinstance(block, dict) and isinstance(block.get("text"), str): + text_parts.append(block["text"]) + for block in item.get("summary") or []: + if isinstance(block, dict) and isinstance(block.get("text"), str): + text_parts.append(block["text"]) + for call in item.get("tool_calls") or []: + function = (call or {}).get("function") or {} + tool_calls.append( + {"function": {"name": function.get("name", ""), "arguments": function.get("arguments", "")}} + ) + turn: dict = {"role": "assistant", "content": "\n".join(text_parts)} + if tool_calls: + turn["tool_calls"] = tool_calls + return turn diff --git a/nemo_gym/token_id_capture/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..b42248cac5 --- /dev/null +++ b/nemo_gym/token_id_capture/records.py @@ -0,0 +1,195 @@ +# 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 + + # --- 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. + # + # 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 + # Whether the model server handed the engine this call's prefix verbatim + # (prefix supply) rather than letting the chat template re-render it. Recorded + # per call so a run can be audited after the fact: supply only fires on a + # unique, verified parent, so the supplied/total ratio is the honest measure + # of how often it applied versus fell back. + prefix_supplied: bool = False + # len(prompt_token_ids) + len(generation_token_ids) for THIS call: the + # length of the prefix a child of this call must start with. + cum_len: Optional[int] = None + # 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..c2844255be --- /dev/null +++ b/nemo_gym/token_id_capture/routes.py @@ -0,0 +1,82 @@ +# 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 +import logging +from hmac import compare_digest +from typing import Any, Optional + +from fastapi import APIRouter, Header, HTTPException, Response + +from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +logger = logging.getLogger(__name__) + + +def make_token_store(global_config_dict: Any) -> Optional[TokenCaptureStore]: + """Build the training-token store, or ``None`` when capture is disabled.""" + config = TokenIdCaptureConfig.model_validate(global_config_dict) + if not config.token_id_capture_enabled: + return None + return TokenCaptureStore(config.resolved_dir()) + + +def install_token_capture_routes(app: Any, store: TokenCaptureStore, read_token: Optional[str] = None) -> None: + """Register the token read route. + + The route serves a rollout's raw training tokens on the same app the harness + calls to generate. That is acceptable inside a trusted cluster; it is not + acceptable once the harness runs in a sandbox whose only egress is this + server, because it could read its own training data -- or another rollout's. + ``read_token`` requires a bearer token; when it is unset the route stays open + and warns once, so existing deployments keep working and the gap is visible. + """ + if not read_token: + logger.warning( + "The token-capture read route is unauthenticated. Anything that can reach this model " + "server can read captured training tokens for any rollout. Set token_id_capture_read_token " + "before running an agent harness in a sandbox whose only egress is this server." + ) + + router = APIRouter() + + @router.get("/ng-capture/tokens/{rollout_id}") + async def get_tokens(rollout_id: str, authorization: Optional[str] = Header(default=None)) -> Response: + if read_token: + expected = f"Bearer {read_token}" + # Constant-time compare: the token is the only thing standing between an untrusted + # harness and every rollout's training data. + if authorization is None or not compare_digest(authorization, expected): + raise HTTPException(status_code=401, detail="token capture read route requires a bearer token") + try: + entries = await asyncio.to_thread(store.read_entries, rollout_id) + except ValueError as exc: + 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..92700f96a6 --- /dev/null +++ b/nemo_gym/token_id_capture/sink.py @@ -0,0 +1,189 @@ +# 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.lineage import ( + LineageIndex, + assistant_turn_from_output_items, +) +from nemo_gym.token_id_capture.protocols import TokenSink +from nemo_gym.token_id_capture.records import ( + TokenEntry, + cumulative_tokens, + extract_token_fields, + response_to_output_items, + stamp_lineage, +) +from nemo_gym.token_id_capture.store import TokenCaptureStore + + +logger = logging.getLogger(__name__) + + +# Which recorded call each request continues, per rollout. Process-wide because the +# capture path is per request; bounded, so an abandoned rollout cannot leak. Losing an +# entry costs a fallback to prefix inference, never a wrong answer. +_LINEAGE = LineageIndex() + + +def lineage_index() -> LineageIndex: + return _LINEAGE + + +@dataclass +class CaptureContext: + """What the capture middleware hands the model server for one call: which + 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 = "" + # Set by the model server when it supplied this call's prefix; read back at + # capture time so the evidence lands on the record rather than only in a log. + prefix_supplied: bool = False + + @property + def sink(self) -> TokenSink: + 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 current_capture_context() -> Optional[CaptureContext]: + """The capture context for the in-flight call, or None for untagged traffic.""" + return _TOKEN_SINK.get() + + +def reset_token_sink(token: Token) -> None: + _TOKEN_SINK.reset(token) + + +async def capture_tokens( + response: Any, + parent_call_id: Optional[str] = None, + request_facts: Optional[dict] = None, + request_messages: Optional[list] = None, +) -> None: + """Record a ``TokenEntry`` from a complete model response when a sink is set. + + ``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 + # Which call does this one continue? Resolved from the conversation the request already + # carries, so the harness needs to send nothing extra. A miss (new conversation, rewritten + # history, or two recorded calls with byte-identical output) leaves the link unset and the + # builder infers from token prefixes instead. + lineage = _LINEAGE.for_rollout(sink.rollout_id) + if parent_call_id is None and request_messages is not None: + parent = lineage.resolve(request_messages) + parent_call_id = parent.call_id if parent is not None else None + try: + entry = TokenEntry( + rollout_id=sink.rollout_id, + 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(), + prefix_supplied=sink.prefix_supplied, + **(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. + stamp_lineage(entry, parent_call_id) + await sink.sink.put(entry) + # Index this call by the conversation a continuation of it would carry, so the next + # request resolves to it. + if request_messages is not None: + lineage.record( + sink.model_call_id, + list(request_messages) + [assistant_turn_from_output_items(entry.output_items)], + cumulative_tokens(entry), + entry.digest or "", + ) + except Exception: + # Capture is best-effort per call: a bad token payload must never fail the + # model call and break the harness's run. But a rollout that lost a call + # 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/resources_servers/reasoning_gym/app.py b/resources_servers/reasoning_gym/app.py index 987d5407bc..1f490d0f42 100644 --- a/resources_servers/reasoning_gym/app.py +++ b/resources_servers/reasoning_gym/app.py @@ -108,6 +108,15 @@ def _extract_answer_from_response(self, response) -> str: if boxed_match: return boxed_match.group(1).strip() + # Conversational agents (e.g. Claude Code driving the policy) often do not use /\boxed + # and instead state the final answer as the last quoted string, e.g. + # So the final answer is: "Zoey is a fool, and Riley is a sage." + # Passing the whole chain-of-thought to the task scorer lets it parse an intermediate + # hypothesis and mis-score a correct rollout, so isolate that trailing quoted answer. + quoted = re.findall(r'["\u201c]([^"\u201d]{3,}?)["\u201d]', full_text) + if quoted: + return quoted[-1].strip() + # return full text if or \boxed{} fail return full_text.strip() if full_text.strip() else "" diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml index e224bcb71b..c62d4f74ab 100644 --- a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent.yaml @@ -14,6 +14,9 @@ reasoning_gym_claude_code_agent: resources_server: type: resources_servers name: reasoning_gym + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: ${anthropic_model_name} anthropic_api_key: ${anthropic_api_key} diff --git a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml index 10a60431aa..512efeeee5 100644 --- a/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml +++ b/resources_servers/reasoning_gym/configs/reasoning_gym_claude_code_agent_model_server.yaml @@ -32,6 +32,9 @@ reasoning_gym_claude_code_agent_model_server: model_server: type: responses_api_models name: policy_model + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: ${policy_model_name} anthropic_api_key: EMPTY # pragma: allowlist secret @@ -39,9 +42,18 @@ reasoning_gym_claude_code_agent_model_server: max_turns: 30 timeout: 300 thinking: disabled + # Restrict the CLI to Bash. Claude Code's system prompt is dominated by TOOL DEFINITIONS, and + # every tool it is allowed adds its schema to the prompt of every call. reasoning_gym only + # needs Bash (the system prompt below says so), so this cuts the fixed prompt cost + # substantially -- which is what keeps prompt + generation inside max_model_len. Measure the + # real cost after a run: the first captured call's prompt_token_ids IS the system prompt plus + # tools plus the first user message. + allowed_tools: "Bash" system_prompt: | You are a precise reasoning assistant. You have access to Bash to run Python for calculations. - For every problem: think step by step, use code to verify when helpful, and state your final answer clearly. + For every problem: think step by step, use code to verify when helpful, then end your reply + with ONLY your final answer inside ... tags, matching the exact format the + task asks for (reasoning_gym scores the text inside ). datasets: - name: example type: example diff --git a/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/configs/claude_code_agent.yaml b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml index 5e3140bea7..41f50e34ab 100644 --- a/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml +++ b/responses_api_agents/claude_code_agent/configs/claude_code_agent.yaml @@ -5,6 +5,9 @@ claude_code_agent: resources_server: type: resources_servers name: ??? + # External harness: returned output carries no token ids, so opt into token capture so this + # agent's model calls are correlated, captured, and rebuilt into a token-bearing response. + token_id_capture: true concurrency: 32 model: claude-sonnet-4-6 anthropic_api_key: ${anthropic_api_key} 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/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index 792a2e2331..62f1a31670 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -15,6 +15,7 @@ import asyncio import base64 import json +import logging import os from copy import deepcopy from time import time @@ -22,7 +23,10 @@ from aiohttp.client_exceptions import ClientResponseError from fastapi import Request -from pydantic import Field +from pydantic import Field, PrivateAttr + + +LOG = logging.getLogger(__name__) from nemo_gym.base_responses_api_model import ( BaseResponsesAPIModelConfig, @@ -44,6 +48,7 @@ split_responses_input_output_items, # noqa: F401 ) from nemo_gym.server_utils import SESSION_ID_KEY, is_nemo_gym_fastapi_entrypoint +from nemo_gym.token_id_capture import current_capture_context, lineage_index class VLLMModelConfig(BaseResponsesAPIModelConfig): @@ -59,6 +64,11 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # Whether or not the model can generate a reasoning output, and called again to produce additional reasoning output. sequential_reasoning_allowed: bool = True + # Supply the engine the exact tokens of the call this request continues, instead of letting the + # chat template re-render them from text. See _apply_prefix_supply below. Off by default: it + # requires a backend that honours required_prefix_token_ids (NeMo-RL's vLLM worker does). + supply_prefix_token_ids: bool = False + # As of Feb 2026, we default this to False since majority of open source models aren't responses native with the exception of GPT-OSS is_responses_native: bool = False @@ -67,6 +77,13 @@ class VLLMModelConfig(BaseResponsesAPIModelConfig): # Corresponds to the extra_body of OpenAI Client. extra_body: Optional[Dict[str, Any]] = None + # Sampling params to force on every chat request, overriding whatever the client sent. An + # external harness (e.g. a CLI agent) chooses its own temperature/top_p, but on-policy RL + # training requires generation to match the sampling distribution the policy is optimized under. + # The integrating training framework sets these to its generation sampling params; Gym only + # enforces them and holds no knowledge of any specific framework. Unset means no override. + sampling_overrides: Optional[Dict[str, Any]] = None + default_headers: Dict[str, str] = Field(default_factory=dict) # Optional prefix for resolving relative ``metadata.audio_path`` (or # entries in ``metadata.audio_paths``) against. Absolute paths are used @@ -448,6 +465,62 @@ def _preprocess_chat_completion_create_params(self, request: Request, body_dict: # No user message found — create one with just the audio blocks. body_dict.setdefault("messages", []).append({"role": "user", "content": list(audio_blocks)}) + # Pin sampling params last so they win over anything the client sent. On-policy RL training + # requires the generation to match the training worker's sampling config; an external harness + # sets its own temperature/top_p, so force them here to keep captured rollouts on-policy. + if self.config.sampling_overrides: + body_dict.update(self.config.sampling_overrides) + + body_dict = self._apply_prefix_supply(body_dict) + + return body_dict + + # [supplied, eligible]. A plain list so it is mutable without a pydantic field. + _prefix_supply_counts: List[int] = PrivateAttr(default_factory=lambda: [0, 0]) + + def _apply_prefix_supply(self, body_dict: Dict[str, Any]) -> Dict[str, Any]: + """Hand the engine the previous call's exact tokens, so this prompt extends them. + + Without this the engine builds every prompt by re-rendering the whole + conversation through the chat template. Two things go wrong. Re-tokenizing + an assistant turn can produce a different (equally valid) split than the + one the model sampled, so the new prompt does not extend the old + prompt-plus-generation. And for a reasoning model the template drops + earlier thinking entirely, so the tokens are not merely re-split, they are + gone. Either way the trajectory cannot be chained and training silently + collapses to the first call. + + Supplying the parent's cumulative ids makes NeMo-RL's splice keep them + verbatim and append only the newly rendered tail, so contiguity holds by + construction and stripped reasoning is restored. + + The rule is conservative on purpose: the splice applies whatever it is + given without checking that it belongs to this conversation, so a wrong + prefix would silently generate from a conversation the harness never + asked for. Supply only on a unique, verified parent; otherwise send the + request untouched and let it start a new chain. A fallback costs a cold + KV cache and a shorter trained chain -- never a wrong answer. + """ + if not self.config.supply_prefix_token_ids: + return body_dict + self._prefix_supply_counts[1] += 1 + context = current_capture_context() + if context is None: + # Not a correlated rollout call, so there is no lineage to supply from. + return body_dict + parent = lineage_index().for_rollout(context.rollout_id).resolve(body_dict.get("messages") or []) + if parent is None or not parent.cum_tokens: + return body_dict + body_dict["required_prefix_token_ids"] = list(parent.cum_tokens) + # Record that this call's prefix was supplied, both on the durable record (so a run can be + # audited afterwards from the capture files) and as a running ratio in the log. Supply only + # fires on a unique, verified parent, so supplied/total is the honest measure of how often + # it applied rather than falling back to re-rendering. + context.prefix_supplied = True + self._prefix_supply_counts[0] += 1 + supplied, total = self._prefix_supply_counts[0], self._prefix_supply_counts[1] + if total and supplied % 10 == 0: + LOG.info("prefix supply: %d/%d calls supplied (%.0f%%)", supplied, total, 100.0 * supplied / total) return body_dict async def chat_completions( @@ -544,8 +617,15 @@ async def chat_completions( # when tokenizing, otherwise `prompt_token_ids` (and therefore logged # `prompt_str`) can be built with different chat template settings than # the actual generation request. + # + # ``required_prefix_token_ids`` is here for the same reason. Generation + # applies the supplied prefix, so without it here the recorded + # prompt_token_ids is a plain re-render of the conversation and does not + # extend the previous call -- the chain looks broken even though the + # engine generated from the right tokens. Measured: 0/9 links contiguous + # with supply on until this key was forwarded. tokenize_body_dict = dict() - for key in ("model", "messages", "tools", "chat_template_kwargs"): + for key in ("model", "messages", "tools", "chat_template_kwargs", "required_prefix_token_ids"): if key in body_dict: tokenize_body_dict[key] = body_dict[key] diff --git a/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml b/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml index 70727036c7..b0e30386c8 100644 --- a/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml +++ b/responses_api_models/vllm_model/configs/vllm_model_for_training.yaml @@ -7,3 +7,10 @@ policy_model: model: ${policy_model_name} return_token_id_information: true uses_reasoning_parser: true + # On-policy training: force generation to the sampling params the policy is optimized under, + # overriding whatever an external harness requests, so captured rollouts stay on-policy. The + # integrating training framework supplies these via the generic keys below (Gym does not know + # any framework); the defaults are standard on-policy sampling when a framework sets nothing. + sampling_overrides: + temperature: ${oc.select:policy_generation_temperature,1.0} + top_p: ${oc.select:policy_generation_top_p,1.0} diff --git a/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml new file mode 100644 index 0000000000..910808f57f --- /dev/null +++ b/responses_api_models/vllm_model/configs/vllm_model_supply_prefix.yaml @@ -0,0 +1,11 @@ +# A/B overlay: turn prefix supply ON (it is off by default in the shipped config). +# +# The model server hands the engine the exact tokens of the call this request continues +# (required_prefix_token_ids), so the new prompt EXTENDS the previous prompt+generation by +# construction instead of being re-rendered from text. Requires a unique, verified parent +# from the lineage index; otherwise it forwards the request untouched and starts a new +# chain -- wasteful, never wrong. +policy_model: + responses_api_models: + vllm_model: + supply_prefix_token_ids: true diff --git a/responses_api_models/vllm_model/tests/test_app.py b/responses_api_models/vllm_model/tests/test_app.py index 2ee728f090..6b911a119a 100644 --- a/responses_api_models/vllm_model/tests/test_app.py +++ b/responses_api_models/vllm_model/tests/test_app.py @@ -53,6 +53,13 @@ NeMoGymSummary, ) from nemo_gym.server_utils import ServerClient +from nemo_gym.token_id_capture import ( + CaptureContext, + TokenCaptureStore, + lineage_index, + reset_token_sink, + set_token_sink, +) from responses_api_models.vllm_model.app import ( VLLMConverter, VLLMModel, @@ -4684,3 +4691,192 @@ async def mock_create_chat_completion(**kwargs): ) # The tokenize endpoint must not be reached once the contract check fails. mock_client.create_tokenize.assert_not_called() + + +class TestPrefixSupply: + """Supplying the engine the previous call's exact tokens. + + Without this the engine re-renders every prompt from text: a re-tokenized + assistant turn can split differently than the model sampled, and a reasoning + model's template drops earlier thinking entirely. Either way the new prompt + does not extend the old one and the trajectory cannot be chained. + """ + + @staticmethod + def _server(monkeypatch: MonkeyPatch, *, enabled: bool) -> VLLMModel: + config = VLLMModelConfig( + host="0.0.0.0", + port=8081, + base_url="http://api.openai.com/v1", + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="", + return_token_id_information=False, + uses_reasoning_parser=False, + supply_prefix_token_ids=enabled, + ) + get_global_config_dict_mock = MagicMock(return_value={}) + monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", get_global_config_dict_mock) + return VLLMModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + + @staticmethod + def _armed(rollout_id: str, tmp_path) -> CaptureContext: + return CaptureContext(rollout_id=rollout_id, model_call_id="call-x", store=TokenCaptureStore(tmp_path)) + + def test_supplies_the_parents_cumulative_tokens(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + first_turn = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-0").record("parent", first_turn, [1, 2, 3, 4], "d") + + token = set_token_sink(self._armed("sup-0", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": first_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3, 4] + + def test_off_by_default(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=False) + turn = [{"role": "assistant", "content": "hello"}] + lineage_index().for_rollout("sup-1").record("parent", turn, [1, 2], "d") + + token = set_token_sink(self._armed("sup-1", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_uncorrelated_call_is_left_alone(self, monkeypatch: MonkeyPatch) -> None: + server = self._server(monkeypatch, enabled=True) + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "hello"}]}) + assert "required_prefix_token_ids" not in out + + def test_fingerprint_miss_falls_back_rather_than_supplying_something_wrong( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """A rewritten history must send the request untouched. The splice applies + whatever it is given without checking it belongs to this conversation, so + a wrong prefix would silently generate from a conversation the harness + never asked for.""" + server = self._server(monkeypatch, enabled=True) + lineage_index().for_rollout("sup-2").record("parent", [{"role": "assistant", "content": "hello"}], [1, 2], "d") + + token = set_token_sink(self._armed("sup-2", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "a summary"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_ambiguous_parent_is_not_supplied(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = self._server(monkeypatch, enabled=True) + turn = [{"role": "assistant", "content": "same"}] + lineage = lineage_index().for_rollout("sup-3") + lineage.record("a", turn, [1, 2], "da") + lineage.record("b", turn, [3, 4], "db") + + token = set_token_sink(self._armed("sup-3", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": turn}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + + def test_a_fork_gets_the_parents_prefix_not_the_previous_calls(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + """Two sub-agents branching from one parent must both get cum(parent). + + A running cursor would hand the second branch a prefix containing the + first branch's generation -- which the splice would apply without + complaint, generating from a conversation that never happened. + """ + server = self._server(monkeypatch, enabled=True) + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage = lineage_index().for_rollout("sup-4") + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 9, 9], + "da", + ) + + token = set_token_sink(self._armed("sup-4", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": shared + [{"role": "user", "content": "b"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + + def test_reasoning_stripped_history_still_supplies_the_real_tokens( + self, monkeypatch: MonkeyPatch, tmp_path + ) -> None: + """The case supplying exists for. + + A reasoning model's chat template drops earlier thinking when it renders + a later prompt, so the re-rendered prompt cannot extend the previous + prompt-plus-generation -- the tokens are gone, not merely re-split. The + supplied prefix comes from the recording, so those tokens are restored + and the chain survives. + """ + server = self._server(monkeypatch, enabled=True) + # The recorded turn included reasoning; what the harness echoes back does not. + recorded_turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "answer"}] + real_tokens_including_reasoning = [1, 2, 3, 4, 5, 6, 7] + lineage_index().for_rollout("sup-5").record("parent", recorded_turn, real_tokens_including_reasoning, "d") + + token = set_token_sink(self._armed("sup-5", tmp_path)) + try: + out = server._apply_prefix_supply({"messages": recorded_turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == real_tokens_including_reasoning + + +class TestPrefixSupplyAccounting: + """Supply must be auditable after the fact. + + Without this the only evidence supply fired is that chains happen to be + contiguous -- which they often are anyway, so a supply run and a non-supply + run look identical. That ambiguity made the first live supply experiment + inconclusive. + """ + + def test_supplied_call_is_marked_on_the_capture_context(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + turn = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}] + lineage_index().for_rollout("acct-0").record("parent", turn, [1, 2, 3], "d") + + ctx = CaptureContext(rollout_id="acct-0", model_call_id="c", store=TokenCaptureStore(tmp_path)) + token = set_token_sink(ctx) + try: + out = server._apply_prefix_supply({"messages": turn + [{"role": "user", "content": "next"}]}) + finally: + reset_token_sink(token) + + assert out["required_prefix_token_ids"] == [1, 2, 3] + assert ctx.prefix_supplied is True + assert server._prefix_supply_counts == [1, 1] # 1 supplied of 1 eligible + + def test_fallback_counts_as_eligible_but_not_supplied(self, monkeypatch: MonkeyPatch, tmp_path) -> None: + server = TestPrefixSupply._server(monkeypatch, enabled=True) + lineage_index().for_rollout("acct-1").record("parent", [{"role": "assistant", "content": "a"}], [1, 2], "d") + + ctx = CaptureContext(rollout_id="acct-1", model_call_id="c", store=TokenCaptureStore(tmp_path)) + token = set_token_sink(ctx) + try: + # A rewritten history: no unique parent, so supply must decline. + out = server._apply_prefix_supply({"messages": [{"role": "assistant", "content": "rewritten"}]}) + finally: + reset_token_sink(token) + + assert "required_prefix_token_ids" not in out + assert ctx.prefix_supplied is False + assert server._prefix_supply_counts == [0, 1] # eligible, not supplied 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 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..cc2fc21044 --- /dev/null +++ b/tests/unit_tests/test_token_id_capture.py @@ -0,0 +1,679 @@ +# 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.lineage import LineageIndex, RolloutLineage, assistant_fingerprint +from nemo_gym.token_id_capture.reader import HttpTokenReader, LocalTokenReader +from nemo_gym.token_id_capture.routes import make_token_store + + +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") + + +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 + + +# --- lineage: which recorded call does a request continue? -------------------- + + +def test_fingerprint_ignores_non_assistant_turns(): + """Only assistant turns identify lineage: they are what we produced. User and + tool content varies with the environment and is irrelevant.""" + a = assistant_fingerprint([{"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}]) + b = assistant_fingerprint([{"role": "user", "content": "DIFFERENT"}, {"role": "assistant", "content": "a"}]) + assert a == b != "" + # No assistant turn at all is a new conversation, not a match. + assert assistant_fingerprint([{"role": "user", "content": "q"}]) == "" + + +def test_fingerprint_survives_tool_argument_reserialization(): + """Harnesses re-serialize tool-call arguments between turns -- compact one + turn, pretty-printed the next. Without canonicalization the same call would + not compare equal to itself and every tool-using turn would miss.""" + compact = [ + {"role": "assistant", "content": "", "tool_calls": [{"function": {"name": "f", "arguments": '{"b":1,"a":2}'}}]} + ] + pretty = [ + { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "f", "arguments": '{\n "a": 2,\n "b": 1\n}'}}], + } + ] + assert assistant_fingerprint(compact) == assistant_fingerprint(pretty) + + +def test_lineage_resolves_the_parent_across_a_turn(): + lineage = RolloutLineage() + first_request = [{"role": "user", "content": "hello"}] + lineage.record("call-1", first_request + [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + + # The next request echoes the assistant turn, as any harness must to continue. + second_request = first_request + [{"role": "assistant", "content": "hi"}, {"role": "user", "content": "more"}] + parent = lineage.resolve(second_request) + assert parent is not None and parent.call_id == "call-1" + assert parent.cum_tokens == [1, 2, 3] and parent.cum_len == 3 + + +def test_lineage_misses_on_a_rewritten_history(): + """A compacted or rewritten context is a new root, not a wrong parent.""" + lineage = RolloutLineage() + lineage.record("call-1", [{"role": "assistant", "content": "hi"}], [1, 2, 3], "d1") + assert lineage.resolve([{"role": "assistant", "content": "a summary of the above"}]) is None + + +def test_lineage_refuses_an_ambiguous_parent(): + """Two recorded calls with byte-identical output cannot be told apart. Guessing + would attribute tokens to the wrong parent, so a unique match is required.""" + lineage = RolloutLineage() + messages = [{"role": "assistant", "content": "same"}] + lineage.record("call-a", messages, [1, 2], "da") + lineage.record("call-b", messages, [3, 4], "db") + assert lineage.resolve(messages) is None + + +def test_lineage_is_a_tree_so_forks_get_the_parent_not_the_previous_call(): + """Two sub-agents branching from one parent must BOTH resolve to that parent. + + A running cursor ("the last call") would hand the second branch a prefix + containing the first branch's generation, and the splice applies a supplied + prefix unconditionally -- so that would be silently wrong, not just wasteful. + """ + lineage = RolloutLineage() + shared = [{"role": "user", "content": "q"}, {"role": "assistant", "content": "plan"}] + lineage.record("parent", shared, [1, 2, 3], "dp") + lineage.record( + "branch-a", + shared + [{"role": "user", "content": "a"}, {"role": "assistant", "content": "A"}], + [1, 2, 3, 4], + "da", + ) + + # The second branch continues the PARENT, not branch-a. + second = shared + [{"role": "user", "content": "b"}] + parent = lineage.resolve(second) + assert parent is not None and parent.call_id == "parent" + assert parent.cum_tokens == [1, 2, 3] + + +def test_lineage_index_is_bounded(): + """An abandoned rollout is never read again, so eviction cannot wait for + consumption. Losing an entry costs a fallback, never a wrong answer.""" + index = LineageIndex(max_rollouts=3) + for i in range(10): + index.for_rollout(f"r{i}") + assert len(index) == 3 + + +def test_served_calls_link_to_their_parent(tmp_path): + """End to end: a second call whose request echoes the first call's assistant + turn is recorded with parent_call_id pointing at it.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + first = [{"role": "user", "content": "hello"}] + client.post("/ng-rollout/lin0-roll0/v1/chat/completions", json={"messages": first}) + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 1 and entries[0].parent_call_id is None + + content = entries[0].output_items[0]["content"] + served_text = content if isinstance(content, str) else content[0]["text"] + second = first + [{"role": "assistant", "content": served_text}, {"role": "user", "content": "more"}] + client.post("/ng-rollout/lin0-roll0/v1/chat/completions", json={"messages": second}) + + entries = TokenCaptureStore(tmp_path).read_entries("lin0-roll0") + assert len(entries) == 2 + assert entries[1].parent_call_id == entries[0].model_call_id + + +def test_fingerprint_matches_across_openai_and_anthropic_tool_shapes(): + """A tool-using turn must match itself across dialects. + + We record the turn we produced in OpenAI shape (``tool_calls``), but Claude + Code echoes it back in Anthropic shape (``content`` blocks of type + ``tool_use``). If those hash differently, the parent is never resolved for + exactly the turns that create multi-turn rollouts -- which is what happened + in the first live multi-turn run: every record came back with + ``parent_call_id: None`` even though the calls chained perfectly. + """ + recorded = [ + { + "role": "assistant", + "content": "Let me compute that.", + "tool_calls": [{"function": {"name": "Bash", "arguments": '{"command":"echo 6"}'}}], + } + ] + echoed = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me compute that."}, + {"type": "tool_use", "name": "Bash", "input": {"command": "echo 6"}}, + ], + } + ] + assert assistant_fingerprint(recorded) == assistant_fingerprint(echoed) != "" + + +def test_lineage_resolves_a_tool_using_turn_echoed_in_anthropic_shape(): + lineage = RolloutLineage() + produced = { + "role": "assistant", + "content": "", + "tool_calls": [{"function": {"name": "Bash", "arguments": '{"command":"factor 420"}'}}], + } + lineage.record("call-1", [{"role": "user", "content": "factor 420"}, produced], [1, 2, 3], "d1") + + # The harness continues the conversation, echoing the turn as Anthropic blocks + # and appending the tool result. + next_request = [ + {"role": "user", "content": "factor 420"}, + {"role": "assistant", "content": [{"type": "tool_use", "name": "Bash", "input": {"command": "factor 420"}}]}, + {"role": "user", "content": [{"type": "tool_result", "content": "420: 2 2 3 5 7"}]}, + ] + parent = lineage.resolve(next_request) + assert parent is not None and parent.call_id == "call-1" + assert parent.cum_tokens == [1, 2, 3] + + +def test_read_route_requires_a_token_when_one_is_configured(tmp_path): + """The route serves raw training tokens on the same app the harness calls to + generate. Inside a trusted cluster that is fine; once the harness runs in a + sandbox whose only egress is this server it could read its own training data, + or another rollout's.""" + config = dict(_both_enabled(tmp_path)) + config["token_id_capture_read_token"] = "s3cret" # pragma: allowlist secret + client = TestClient(_server(config).setup_webserver()) + client.post("/ng-rollout/auth0-roll0/v1/responses", json={"input": "hi"}) + + assert client.get("/ng-capture/tokens/auth0-roll0").status_code == 401 + assert client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer wrong"}).status_code == 401 + + ok = client.get("/ng-capture/tokens/auth0-roll0", headers={"Authorization": "Bearer s3cret"}) + assert ok.status_code == 200 + assert TokenEntry.model_validate_json(ok.text.splitlines()[0]).generation_token_ids == GTOKS + + +def test_read_route_stays_open_when_no_token_is_configured(tmp_path): + """Existing deployments keep working; the gap is logged rather than enforced.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/auth1-roll0/v1/responses", json={"input": "hi"}) + assert client.get("/ng-capture/tokens/auth1-roll0").status_code == 200 diff --git a/tests/unit_tests/test_trajectory_builder.py b/tests/unit_tests/test_trajectory_builder.py new file mode 100644 index 0000000000..47b9765fa8 --- /dev/null +++ b/tests/unit_tests/test_trajectory_builder.py @@ -0,0 +1,467 @@ +# 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, + classify_side_calls, + 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 + + +# --- 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