From bbd65e9adb4a811ef131b2187491aa5e4f1c5f43 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 4 Aug 2026 19:07:28 -0700 Subject: [PATCH 01/12] feat(token-id-capture): record the token ids a harness does not return An agent harness that drives its own model calls hands back a transcript with no token ids, because the wire formats it speaks have no field for them. RL trains on token ids, and re-tokenizing the returned text gives a sequence that differs from what the policy sampled by an unmeasured amount. The ids still exist inside the model server, for the moment before it converts the response to the harness's dialect and synthesizes a stream. Capture takes them there, keyed to the rollout that produced them, and writes one TokenEntry per model call. Calls are correlated to a rollout by the /ng-rollout/ path prefix already on main for evaluation capture, so this adds no second correlation scheme. The agent-side gate now serves both consumers, and a per-agent token_id_capture flag scopes which agents participate; native agents leave it off because they carry token ids on their own response items. The capture key is derived from a run request's task and rollout indices, which assumes each dispatch gets a distinct pair. A caller that restarts numbering per dispatch produces a repeated id, so two dispatches share one key and their calls stitch into one trajectory. An explicit _ng_rollout_id on the run body replaces the derivation, with the attempt suffix still applied on top. An id that would not survive the path segment is refused rather than rewritten, and the id pattern is defined once so the body check and the middleware cannot disagree. Settings live in one `token_id_capture` block rather than as flat keys, and it names where records go: token_id_capture: enabled: true dir: /tmp/ng_tokcap sink: my_pkg.sinks:MyDataPlaneSink `sink` is constructed once per server process at app startup. That matters at num_workers > 1: uvicorn is handed an app string and workers=N and spawns those workers, re-importing the app module rather than inheriting the launcher's memory, so a sink installed programmatically by a launcher does not exist in any worker. Measured, capture then falls back to the file store, or writes nothing at all when no directory is set, and logs no error either way. install_token_sink remains for programmatic use under the same constraint. The validator refuses combinations that would silently capture nothing: settings with `enabled: false`, a sink beside a directory, an unknown key, and a sink that cannot report a lost call. TokenSink and TokenSource are protocols in a module that imports no web framework, cluster runtime or tensor library, so an inference worker can write into its own data plane without pulling in the server stack. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_agent.py | 40 +- nemo_gym/base_responses_api_model.py | 145 ++- nemo_gym/global_config.py | 8 + nemo_gym/rollout_collection.py | 5 + nemo_gym/rollout_correlation.py | 48 +- nemo_gym/token_id_capture/__init__.py | 77 ++ nemo_gym/token_id_capture/config.py | 176 ++++ nemo_gym/token_id_capture/protocols.py | 122 +++ nemo_gym/token_id_capture/records.py | 170 ++++ nemo_gym/token_id_capture/sink.py | 190 ++++ nemo_gym/token_id_capture/store.py | 164 ++++ ...ng_gym_claude_code_agent_model_server.yaml | 4 + .../configs/claude_code_agent.yaml | 7 + .../test_base_responses_api_model.py | 42 + tests/unit_tests/test_rollout_collection.py | 51 ++ tests/unit_tests/test_token_id_capture.py | 843 ++++++++++++++++++ 16 files changed, 2054 insertions(+), 38 deletions(-) create mode 100644 nemo_gym/token_id_capture/__init__.py create mode 100644 nemo_gym/token_id_capture/config.py create mode 100644 nemo_gym/token_id_capture/protocols.py create mode 100644 nemo_gym/token_id_capture/records.py create mode 100644 nemo_gym/token_id_capture/sink.py create mode 100644 nemo_gym/token_id_capture/store.py create mode 100644 tests/unit_tests/test_token_id_capture.py diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 3633f562b0..579858115b 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -27,7 +27,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_BLOCK, + get_first_server_config_dict, +) from nemo_gym.openai_utils import ( NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming, @@ -46,6 +50,13 @@ class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig): skip_verification: bool = False skip_verification_reward: float = 0.0 + # 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): @@ -81,21 +92,34 @@ 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)) + block = global_config.get(TOKEN_ID_CAPTURE_BLOCK) or {} + token_capture = bool(isinstance(block, Mapping) and block.get("enabled", 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 846bb79f06..2dd4307b9b 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -67,6 +67,18 @@ BaseServer, SimpleServer, ) +from nemo_gym.token_id_capture import ( + CaptureContext, + capture_tokens, + installed_token_sink, + 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.config import token_id_capture_config +from nemo_gym.token_id_capture.store import make_token_store logger = logging.getLogger(__name__) @@ -76,6 +88,26 @@ _ANTHROPIC_CONVERTER = AnthropicConverter() +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 +122,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 +229,11 @@ async def _invoke_chat_completions( # only `body`. Dispatch on whichever this server declares so the shared dispatch works for # all of them. if "request" in inspect.signature(self.chat_completions).parameters: - return await self.chat_completions(request=request, body=params) - return await self.chat_completions(body=params) + completion = await self.chat_completions(request=request, body=params) + else: + completion = await self.chat_completions(body=params) + await capture_tokens(completion) + return completion async def messages(self, request: Request, body: dict = Body()): """Default Anthropic Messages <-> Responses mapping shared by every Gym model server. @@ -222,8 +262,14 @@ async def _invoke_responses( # `body`. Dispatch on whichever this server declares so the default messages() works for # all of them. if "request" in inspect.signature(self.responses).parameters: - return await self.responses(request=request, body=params) - return await self.responses(body=params) + response = await self.responses(request=request, body=params) + else: + response = await self.responses(body=params) + # Capture here rather than at the route: the streaming dispatch returns a StreamingResponse + # and the Anthropic mapping drops the token fields, so this is the last point where the + # assembled response still carries them, for every dialect. + await capture_tokens(response) + return response def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming: @@ -1026,10 +1072,22 @@ 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: CaptureStore | None, + model_server_name: str | None, + token_store: Any = None, + configured_sink: 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 + # Built from token_id_capture.sink, once, in this process. + self._configured_sink = configured_sink async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: if scope.get("type") != "http": @@ -1044,24 +1102,44 @@ 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. + # Destination order: a sink configured for this process, then one installed + # programmatically, then the file store. Both sink routes exist because a framework may + # send records to its own transport instead of disk; the configured one is preferred + # because it is built inside this process at app startup and so survives num_workers > 1, + # where a sink installed by a launcher does not reach the spawned workers at all. The + # installed sink is still resolved per request, so one installed after the app is built + # still takes effect. + token_sink = self._configured_sink or installed_token_sink() or self._token_store + if (self._store is None and token_sink 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 cannot: token ids are dropped on the SSE wire. + sink_token = None + if token_sink is not None: + sink_token = set_token_sink( + CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, sink=token_sink) + ) + + # 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]: @@ -1143,6 +1221,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 @@ -1210,21 +1292,38 @@ 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: str | None = 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 + # Built here, at app startup, so every uvicorn worker constructs its own. A sink installed by a + # launcher process is not inherited by spawned workers and would silently go unused. + configured_sink = ( + token_id_capture_config(global_config_dict).build_sink() 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, + configured_sink=configured_sink, ) diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 5a84fdf328..2447bf370e 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -95,6 +95,10 @@ QUERY_KEY_NAME = "query" OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" +# Run-wide training-token capture settings; see nemo_gym/token_id_capture/config.py. +TOKEN_ID_CAPTURE_BLOCK = "token_id_capture" +# 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" SKIP_VERIFICATION_KEY_NAME = "skip_verification" SKIP_VERIFICATION_REWARD_KEY_NAME = "skip_verification_reward" @@ -140,6 +144,10 @@ # Resume re-dispatch attempt counter (0 on the first attempt); distinguishes retries of the same # (task, rollout) so their captured model calls stay separable. ATTEMPT_INDEX_KEY_NAME = "_ng_attempt_index" +# Explicit capture id for a run request, used in place of the (task, rollout) derivation. Set it +# when the caller reuses task and rollout indices across dispatches, since the derived id would +# then repeat and two dispatches would share one capture key. +ROLLOUT_ID_KEY_NAME = "_ng_rollout_id" RESPONSES_CREATE_PARAMS_KEY_NAME = "responses_create_params" RESPONSE_KEY_NAME = "response" AGENT_REF_KEY_NAME = "agent_ref" diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index cbd8501788..d26f82dbb6 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -52,6 +52,7 @@ AGENT_REF_KEY_NAME, ATTEMPT_INDEX_KEY_NAME, RESPONSES_CREATE_PARAMS_KEY_NAME, + ROLLOUT_ID_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, SKILLS_REF_KEY_NAME, TASK_INDEX_KEY_NAME, @@ -816,6 +817,10 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D result[SKILLS_REF_KEY_NAME] = row[SKILLS_REF_KEY_NAME] if ATTEMPT_INDEX_KEY_NAME in row: result[ATTEMPT_INDEX_KEY_NAME] = row[ATTEMPT_INDEX_KEY_NAME] + if ROLLOUT_ID_KEY_NAME in row: + # Capture readback recomputes the id from the finished record, so an explicit id + # has to travel from the dispatched row onto the result the same way the indices do. + result[ROLLOUT_ID_KEY_NAME] = row[ROLLOUT_ID_KEY_NAME] # Fold this rollout's captured model calls into its record (uniform across agents; no-op # when capture is off). Never alters the harness output/reward already in `result`. diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py index c1d58a2694..8d7d85f0cc 100644 --- a/nemo_gym/rollout_correlation.py +++ b/nemo_gym/rollout_correlation.py @@ -23,6 +23,7 @@ from nemo_gym.config_types import ROLLOUT_PATH_PREFIX from nemo_gym.global_config import ( ATTEMPT_INDEX_KEY_NAME, + ROLLOUT_ID_KEY_NAME, ROLLOUT_INDEX_KEY_NAME, TASK_INDEX_KEY_NAME, ) @@ -30,21 +31,53 @@ _ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None) +# A capture id travels as a path segment in ``/ng-rollout/``, so it is limited to what a path +# segment carries unambiguously. Leading dots are excluded because the id is also a filename +# component in the capture stores. The middleware below matches on the same pattern, so an id this +# rejects is one that would not have survived the round trip anyway. +ROLLOUT_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: - """Build the capture key stamped by rollout collection.""" + """Build the capture key for a run request. + + An explicit ``_ng_rollout_id`` on the body wins. Otherwise the id is derived from the task and + rollout indices as ``"{task}-{rollout}"``. Both forms then take an ``-a{n}`` suffix for a + re-dispatch attempt past the first. + + The derivation is a contract, not an implementation detail: capture writers key records by the + id this returns and capture readers look records up by recomputing it from the finished rollout + record, so the two sides only agree while the rule is the same on both. Changing the format + invalidates any record already on disk. + + The derivation also assumes the caller gives each dispatch a distinct (task, rollout) pair. + A caller that restarts numbering, such as one running the same indices once per training step, + produces a repeated id and two dispatches then share a capture key, which stitches unrelated + calls into one trajectory. Set an explicit id to opt out of the derivation in that case. + """ if not isinstance(body, (BaseModel, Mapping)): return None def field(key: str) -> Any: return body.get(key) if isinstance(body, Mapping) else getattr(body, key, None) - task = field(TASK_INDEX_KEY_NAME) - rollout = field(ROLLOUT_INDEX_KEY_NAME) - if task is None or rollout is None: - return None + explicit = field(ROLLOUT_ID_KEY_NAME) + if explicit is not None: + # A malformed explicit id is refused rather than sanitized. Rewriting it would correlate + # calls under an id the caller never chose and cannot look up afterwards. + if not (isinstance(explicit, str) and ROLLOUT_ID_PATTERN.match(explicit)): + raise ValueError( + f"{ROLLOUT_ID_KEY_NAME} must be a string of letters, digits, dots, dashes or " + f"underscores starting with a letter or digit; got {explicit!r}" + ) + rollout_id = explicit + else: + task = field(TASK_INDEX_KEY_NAME) + rollout = field(ROLLOUT_INDEX_KEY_NAME) + if task is None or rollout is None: + return None + rollout_id = f"{task}-{rollout}" - rollout_id = f"{task}-{rollout}" attempt = field(ATTEMPT_INDEX_KEY_NAME) if attempt is not None and int(attempt) > 0: rollout_id = f"{rollout_id}-a{int(attempt)}" @@ -67,8 +100,9 @@ def rollout_context(rollout_id: Optional[str]) -> Iterator[None]: class RolloutContextMiddleware: """Strip a rollout prefix and expose it to downstream Gym calls for this request.""" + # Same id charset as ROLLOUT_ID_PATTERN, anchored between the prefix and the rest of the path. _PREFIX = re.compile( - rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P[A-Za-z0-9][A-Za-z0-9._-]*)(?P/.*)$" + rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P{ROLLOUT_ID_PATTERN.pattern.strip('^$')})(?P/.*)$" ) def __init__(self, app: Any) -> None: diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py new file mode 100644 index 0000000000..5b72148e67 --- /dev/null +++ b/nemo_gym/token_id_capture/__init__.py @@ -0,0 +1,77 @@ +# 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``). + +Records are read back through a ``TokenSource``. ``TokenCaptureStore`` is one, +and is what a reader sitting alongside the store uses. A framework staging +records through its own transport supplies its own source, which lives wherever +that transport does. What any source owes is an honest ``is_incomplete``: it is +how a consumer learns a rollout lost a call, and one that always answers False +trains on an incomplete rollout without knowing. +""" + +from nemo_gym.token_id_capture.config import TokenIdCaptureConfig +from nemo_gym.token_id_capture.protocols import ( + TokenSink, + TokenSource, + install_token_sink, + installed_token_sink, +) +from nemo_gym.token_id_capture.records import ( + TOKEN_ENTRY_RECORD_SCHEMA_VERSION, + TOKEN_FIELDS, + TokenEntry, + extract_token_fields, +) +from nemo_gym.token_id_capture.sink import ( + CaptureContext, + capture_tokens, + commit_entry, + reset_token_sink, + set_token_sink, +) +from nemo_gym.token_id_capture.store import TokenCaptureStore, make_token_store, validate_rollout_id + + +__all__ = [ + "TokenIdCaptureConfig", + "TokenEntry", + "TOKEN_ENTRY_RECORD_SCHEMA_VERSION", + "TOKEN_FIELDS", + "extract_token_fields", + "TokenCaptureStore", + "validate_rollout_id", + "make_token_store", + "TokenSink", + "TokenSource", + "install_token_sink", + "installed_token_sink", + "CaptureContext", + "set_token_sink", + "reset_token_sink", + "capture_tokens", + "commit_entry", +] diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py new file mode 100644 index 0000000000..975070bc92 --- /dev/null +++ b/nemo_gym/token_id_capture/config.py @@ -0,0 +1,176 @@ +# 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 settings for training-token capture, in one block. + +```yaml +env: + nemo_gym: + token_id_capture: + enabled: true + dir: /tmp/ng_tokcap # node-local; writer and reader share a node + sink: my_pkg.sinks:MyDataPlaneSink # optional; default is the file store at `dir` +``` + +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 ``dir`` is given, tokens are written alongside the eval +capture files in the top-level ``model_call_capture_dir``. + +The per-agent ``token_id_capture`` flag is a narrower, separate control: it scopes +which agents participate. Native agents leave it off because they already carry +token ids on their response items. + +Choosing where records go +------------------------- +``sink`` names a class implementing ``TokenSink``, as ``module.path:ClassName``. +It is constructed once per server process at app startup and replaces the file +store, so records go to a framework's own transport and never touch disk. + +Construction has to happen inside the serving process. A model server configured +with ``num_workers > 1`` is launched by uvicorn with an app string and +``workers=N``, and uvicorn spawns those workers with the ``spawn`` start method, +which re-imports the app module rather than inheriting the parent's memory. A +sink installed by a launcher script therefore does not exist in any worker, and +capture silently falls back to the file store, or writes nothing at all when no +``dir`` is set. Naming the sink here avoids that: each worker builds its own. + +``install_token_sink`` remains for programmatic use and is subject to the same +constraint, so call it at module import of the app, not from a parent process. +""" + +from __future__ import annotations + +import logging +from importlib import import_module +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, model_validator + +from nemo_gym.token_id_capture.protocols import TokenSink, installed_token_sink + + +logger = logging.getLogger(__name__) + +TOKEN_ID_CAPTURE_BLOCK = "token_id_capture" + + +class TokenIdCaptureSettings(BaseModel): + """The ``token_id_capture`` block.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + # Where the default file store writes. Falls back to ``model_call_capture_dir``. + dir: Path | None = None + # ``module.path:ClassName`` implementing TokenSink, constructed per server process. + sink: str | None = None + # Keyword arguments for that constructor: an endpoint, a client, credentials. A sink for a real + # transport needs wiring, and a zero-argument one could only get it from ambient state. Use + # ``${oc.env:VAR}`` for anything secret rather than writing it here. + sink_kwargs: dict[str, Any] = {} + + +class TokenIdCaptureConfig(BaseModel): + """The capture block plus the one top-level key it falls back to.""" + + model_config = ConfigDict(extra="ignore") + + token_id_capture: TokenIdCaptureSettings = TokenIdCaptureSettings() + # Shared with evaluation capture, which owns it, so it stays top-level. + model_call_capture_dir: Path | None = None + + @model_validator(mode="after") + def _validate(self) -> "TokenIdCaptureConfig": + block = self.token_id_capture + if not block.enabled: + # The rest of the block is left alone rather than rejected. Configs are templated, and + # setting a directory unconditionally while toggling `enabled` per run is ordinary. + return self + if block.sink is not None: + if block.dir is not None: + # Not an error: nothing is lost, the directory is simply never read. Worth saying + # once, because someone expecting files on disk will not find any. + logger.warning( + "token_id_capture.dir is set alongside token_id_capture.sink. The sink replaces " + "the file store, so %s will not be written to.", + block.dir, + ) + return self + directory = self.resolved_dir() + if directory is None: + # A process that installed a sink programmatically writes through that transport and + # never constructs the file store, so it has no directory to give. + if installed_token_sink() is not None: + return self + raise ValueError( + "token_id_capture.dir (or model_call_capture_dir) is required when " + "token_id_capture.enabled is true and no sink is configured or installed" + ) + if not directory.is_absolute(): + raise ValueError("training-token capture directory must be an absolute path") + return self + + @property + def enabled(self) -> bool: + return self.token_id_capture.enabled + + def resolved_dir(self) -> Path | None: + return self.token_id_capture.dir or self.model_call_capture_dir + + def build_sink(self) -> TokenSink | None: + """Construct the configured sink, or ``None`` when the file store is in use. + + Called once per server process at app startup, which is what makes this work under + ``num_workers > 1`` where a sink installed by a launcher does not reach the workers. + """ + target = self.token_id_capture.sink + if not self.token_id_capture.enabled or target is None: + return None + if ":" not in target: + raise ValueError(f"token_id_capture.sink must be 'module.path:ClassName' (got {target!r})") + module_path, _, class_name = target.partition(":") + try: + factory = getattr(import_module(module_path), class_name) + except (ImportError, AttributeError) as error: + raise ValueError(f"could not load token_id_capture.sink {target!r}: {error}") from error + try: + sink = factory(**self.token_id_capture.sink_kwargs) + except TypeError as error: + raise ValueError( + f"could not construct token_id_capture.sink {target!r} with " + f"sink_kwargs={sorted(self.token_id_capture.sink_kwargs)}: {error}" + ) from error + # Checked here rather than at first use: a sink that cannot record a failure makes an + # incomplete rollout look complete, and a startup error is better than that at step 400. + # + # isinstance against the protocol rather than a list of names written out here, so this + # keeps up when TokenSink gains a method. It only checks that the attributes exist, so the + # loop below adds the part it does not do. Neither checks signatures; nothing at runtime + # can, short of calling the methods. + missing = [name for name in sorted(TokenSink.__protocol_attrs__) if not callable(getattr(sink, name, None))] + if missing or not isinstance(sink, TokenSink): + raise ValueError( + f"token_id_capture.sink {target!r} does not satisfy TokenSink: " + f"{', '.join(missing) or 'attribute check failed'}" + ) + return sink + + +def token_id_capture_config(global_config_dict: Any) -> TokenIdCaptureConfig: + """Read the capture settings out of a global config dict.""" + return TokenIdCaptureConfig.model_validate(global_config_dict or {}) diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py new file mode 100644 index 0000000000..7e238bd8ed --- /dev/null +++ b/nemo_gym/token_id_capture/protocols.py @@ -0,0 +1,122 @@ +# 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. + +"""Interfaces for writing and reading captured training tokens. + +Gym owns the record shape, these protocols, and the code that builds a record. A +training framework supplies the implementations and runs them wherever its +tokens are produced. Neither side imports the other's transport. + +Placement of the write is therefore a deployment choice: + +- Gym owns serving (today): install the sink in the model server, which already + holds the assembled response, so there is no extra hop. +- A framework owns the inference worker: install the sink there, so bulk token + arrays go to the framework's data plane instead of riding back through Gym's + HTTP response. + +The capture code is the same in both cases. This module must stay free of +fastapi, ray, torch and aiohttp imports so a framework's worker can import it +without pulling in Gym's server stack. A unit test enforces that. +""" + +from __future__ import annotations + +from typing import 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. + + The record must be durable before this returns: 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. + """ + ... + + def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + """Record that a call of this rollout failed to capture. + + The rollout is now missing a turn, and a consumer must mask the sample rather + than train on a chain with a hole in it. The model call itself still succeeds, + so this is the only signal that anything went wrong: a sink that drops it makes + an incomplete rollout indistinguishable from a complete one. + + Synchronous, because the caller is a failure path that cannot await. A transport + that needs to send should queue here and flush elsewhere. + """ + ... + + +@runtime_checkable +class TokenSource(Protocol): + """Where a trajectory builder reads records from, and retires them afterwards.""" + + 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, using parent links or token-prefix relationships. + """ + ... + + async def drop(self, rollout_id: str) -> None: + """Retire a rollout's records once they have been consumed. + + A transport that cannot delete implements this as a no-op and leaves + retirement to whoever owns the storage. + """ + ... + + def is_incomplete(self, rollout_id: str) -> bool: + """Whether a call of this rollout failed to capture, as reported by + ``TokenSink.mark_incomplete``. + + The records that did arrive can stitch into a chain that looks perfectly + contiguous while missing a turn, so this is the only way to tell. A + transport that cannot tell returns False, the same way a transport that + cannot delete makes ``drop`` a no-op; the cost is that an incomplete + rollout of its is trained on rather than masked. + + Synchronous, to match ``mark_incomplete`` on the sink side. + """ + return False + + +# 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: TokenSink | None = None + + +def install_token_sink(sink: TokenSink | None) -> None: + """Set (or clear, with ``None``) the process-wide default sink.""" + global _INSTALLED_SINK + _INSTALLED_SINK = sink + + +def installed_token_sink() -> TokenSink | None: + return _INSTALLED_SINK diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py new file mode 100644 index 0000000000..cb97e51b31 --- /dev/null +++ b/nemo_gym/token_id_capture/records.py @@ -0,0 +1,170 @@ +# 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 + +from typing import Any + +from pydantic import BaseModel, ConfigDict, model_validator + + +# 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 whenever a field is added or its meaning changes. Writer and reader are different +# processes and may be different repositories, and records outlive a deploy, so a reader has to +# be able to refuse a record it was not built for. ``extra="allow"`` means an unknown shape +# otherwise decodes cleanly and corrupts training rows in silence. The field is present from the +# first version because a check added later cannot tell an old record from an unversioned one. +# +# 1 rollout and call identity, the token arrays, the output items and their carrier index +TOKEN_ENTRY_RECORD_SCHEMA_VERSION = 1 + + +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. Token + ids alone are not enough for a trainer that scores text, such as penalties for an + invalid tool call or a malformed thinking block. + + The token arrays are stored once, at the top level. The served response carries + them on the output item that produced the generation, and those per-item copies + are dropped on write: the builder overwrites them anyway, since an item's prompt + in a chained trajectory is the running cumulative sequence rather than the prompt + of the single call. ``token_item_index`` records which item they came off, so the + builder can put the chain-correct values back on the right one. + """ + + model_config = ConfigDict(extra="allow") + + schema_version: int = TOKEN_ENTRY_RECORD_SCHEMA_VERSION + 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: Any | None = None + # The served response's output items (Responses shape), content preserved, token + # arrays removed. + output_items: list[dict] = [] + # Index into ``output_items`` of the item the token arrays were taken off, or null + # when no item carried them. Records written before the arrays were de-duplicated + # leave this unset and still carry the arrays inline, which the builder handles. + token_item_index: int | None = None + # Non-semantic; a cheap diagnostic for retry/sibling-branch cases. + created_at: float = 0.0 + + @model_validator(mode="after") + def _refuse_a_newer_record(self) -> "TokenEntry": + """Decode a record older than this reader, refuse one newer. + + Older is safe: a field this reader does not have takes its default and the consumer + degrades, so a record written before parent links existed simply has none and the builder + matches token prefixes instead. + + Newer is not, and it is the direction ``extra="allow"`` hides. A field this reader cannot + see is kept and ignored, so a record whose tokens were written under rules this reader does + not know decodes clean and trains as though nothing were different. Refusing is loud: the + read fails, the caller marks that rollout unusable, and the run says which version it saw. + """ + if self.schema_version > TOKEN_ENTRY_RECORD_SCHEMA_VERSION: + raise ValueError( + f"token record is schema_version {self.schema_version}, but this reader understands " + f"up to {TOKEN_ENTRY_RECORD_SCHEMA_VERSION}. Upgrade the reader, or point it at " + "records written by a writer it matches." + ) + return self + + +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 strip_token_fields(items: list[dict]) -> tuple[list[dict], int | None]: + """Drop the token arrays from output items, keeping the content. + + Returns the stripped items and the index of the item the arrays came off, which + is the last one carrying them, matching what ``extract_token_fields`` reads. The + arrays are held once on the entry instead: storing them again per item roughly + doubles a record, and the per-item values are not the ones a trainer sees, since + the builder replaces an item's prompt with the chain's running sequence. + """ + index: int | None = None + stripped: list[dict] = [] + for position, item in enumerate(items): + if item.get("generation_token_ids") is not None: + index = position + stripped.append({key: value for key, value in item.items() if key not in TOKEN_FIELDS}) + return stripped, index + + +def extract_token_fields(response_json: dict) -> dict | None: + """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/sink.py b/nemo_gym/token_id_capture/sink.py new file mode 100644 index 0000000000..e4b2822346 --- /dev/null +++ b/nemo_gym/token_id_capture/sink.py @@ -0,0 +1,190 @@ +# 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, token ids included, for a moment 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 + +from nemo_gym.token_id_capture.protocols import TokenSink +from nemo_gym.token_id_capture.records import ( + TokenEntry, + extract_token_fields, + response_to_output_items, + strip_token_fields, +) + + +logger = logging.getLogger(__name__) + + +@dataclass +class CaptureContext: + """What the capture middleware hands the model server for one call: which + rollout and call this is, and where the record goes. + + ``sink`` is Gym's file store by default and anything satisfying ``TokenSink`` + otherwise, which is how a training framework redirects the write to its own data + plane without changing the capture code. It is typed as the protocol rather than + the file store so that redirection is a supported path and not an accident. + """ + + rollout_id: str + model_call_id: str + sink: TokenSink + model: str = "" + + +_TOKEN_SINK: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_token_sink", default=None) + + +def set_token_sink(sink: CaptureContext) -> Token: + return _TOKEN_SINK.set(sink) + + +def reset_token_sink(token: Token) -> None: + _TOKEN_SINK.reset(token) + + +async def capture_tokens(response: Any) -> 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 awaited, so the entry is durable before the model call returns and a + post-rollout reader always sees it, with no background writer to drain. + """ + sink = _TOKEN_SINK.get() + if sink is None: + return + # Everything that reads the response is guarded, not just the write. Decoding a payload and + # validating a record can fail on malformed token data exactly as writing it can, and the + # consequence is the same: the rollout is short a call. It is guarded here rather than left + # to the caller because the caller is the model server's own response path, so an exception + # escaping this function would fail the model call and break the harness's run. + try: + 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 + # Content only: the arrays live on the entry, not on the items as well. + content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) + + 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, since text-based penalties need it. + output_items=content_items, + token_item_index=token_item_index, + created_at=time.time(), + ) + except Exception: + _capture_failed(sink, "build") + return + await commit_entry(entry) + + +async def commit_entry(entry: TokenEntry) -> None: + """Durably record a finished entry against the in-flight call. + + Public and separate from ``capture_tokens`` because the two halves are useful apart. + ``capture_tokens`` reads the arrays off a served response; a framework that captures + engine-side already has them, and the response Gym sees may carry none at all, so it + needs this half without the extraction half. Forking it instead would duplicate the + ordering below, which is the part worth sharing. + + No-op when no sink is active. Never raises: capture is best effort per call, but a + rollout that lost a call is marked so a consumer masks it rather than training on a + chain with a hole. + """ + sink = _TOKEN_SINK.get() + if sink is None: + return + try: + await sink.sink.put(entry) + except Exception: + _capture_failed(sink, "write") + + +def _capture_failed(sink: CaptureContext, stage: str) -> None: + """Report a capture failure without letting it reach the model call. + + 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 it is marked, and delivery masks the sample rather than training on a + chain with a hole. Called only from an ``except`` block, so ``exc_info`` picks up the + active exception. + """ + logger.warning( + "Training-token capture failed to %s the record for model call %s of rollout %s.", + stage, + sink.model_call_id, + sink.rollout_id, + exc_info=True, + ) + _mark_incomplete(sink) + + +def _mark_incomplete(sink: CaptureContext) -> None: + """Mark the rollout, or say loudly why it could not be marked. + + A sink that does not implement ``mark_incomplete`` would otherwise raise inside the + failure path above and have the exception swallowed, leaving an incomplete rollout + that looks complete. That is the one outcome this whole path exists to prevent, so + it is logged at error rather than passed over. + """ + mark = getattr(sink.sink, "mark_incomplete", None) + if mark is None: + logger.error( + "Sink %s does not implement mark_incomplete. Rollout %s cannot be marked incomplete " + "and may be trained on with a missing call.", + type(sink.sink).__name__, + sink.rollout_id, + ) + return + try: + mark(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/store.py b/nemo_gym/token_id_capture/store.py new file mode 100644 index 0000000000..c999cc5351 --- /dev/null +++ b/nemo_gym/token_id_capture/store.py @@ -0,0 +1,164 @@ +# 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 or 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 +from typing import Any + +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, because 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, so 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; + # a framework swaps in its own without touching the capture path. + # + # Both offload to the default thread pool, which is shared process-wide and small + # (min(32, cpus + 4)). Serializing the entry dominates the cost rather than the write + # itself, so a long context is the case to watch if this ever shows up in a profile. + + 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 + + +def make_token_store(global_config_dict: Any) -> TokenCaptureStore | None: + """Build the training-token store, or ``None`` when this process is not writing one. + + ``None`` when capture is off, when no directory resolves, or when a sink is configured: the + records go to that transport instead and there is no file store to build. + """ + from nemo_gym.token_id_capture.config import TokenIdCaptureConfig + + config = TokenIdCaptureConfig.model_validate(global_config_dict) + if not config.enabled or config.token_id_capture.sink is not None: + return None + directory = config.resolved_dir() + return TokenCaptureStore(directory) if directory is not None else None 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..adcf80a523 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,10 @@ reasoning_gym_claude_code_agent_model_server: model_server: type: responses_api_models name: policy_model + # An external harness: its returned output has no token ids, so its model calls are + # captured and rebuilt when the run-level token_id_capture.enabled switch is on. Inert + # otherwise, so this is safe for the evaluation showcase above. + token_id_capture: true concurrency: 32 model: ${policy_model_name} anthropic_api_key: EMPTY # pragma: allowlist secret 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..c25b157115 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,13 @@ claude_code_agent: resources_server: type: resources_servers name: ??? + # This harness returns output with no token ids, so training on its rollouts requires + # capture: its model calls are correlated, captured, and rebuilt into a token-bearing + # response. On for this agent because it is an external harness, which is the case the + # flag exists to identify; a native agent leaves it off, since it carries token ids inline + # and rebuilding would replace them with a reconstruction. This costs nothing on its own: + # the run-level token_id_capture.enabled switch still has to be on for anything to happen. + token_id_capture: true concurrency: 32 model: claude-sonnet-4-6 anthropic_api_key: ${anthropic_api_key} diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 358c9bedfb..952950ce18 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -1087,6 +1087,48 @@ def test_maybe_rollout_id_from_run_body_attempt_suffix(): maybe_rollout_id_from_run_body({**base, "_ng_attempt_index": "invalid"}) +def test_maybe_rollout_id_from_run_body_prefers_an_explicit_id(): + from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body + + base = {"_ng_task_index": 3, "_ng_rollout_index": 2} + # A caller that restarts index numbering per dispatch derives the same id twice, which would + # give two dispatches one capture key. The explicit id is how it keeps them apart. + assert maybe_rollout_id_from_run_body({**base, "_ng_rollout_id": "s7-3-2"}) == "s7-3-2" + # A retry of an explicitly keyed rollout still keys separately from its first attempt, or the + # retry's calls would append onto the first attempt's records. + assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "s7-3-2", "_ng_attempt_index": 1}) == "s7-3-2-a1" + # The explicit id stands alone: no indices needed. + assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "abc"}) == "abc" + + +@pytest.mark.parametrize("bad", [".hidden", "has/slash", "has space", "", 7, None]) +def test_maybe_rollout_id_from_run_body_refuses_an_unusable_explicit_id(bad): + from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body + + body = {"_ng_task_index": 3, "_ng_rollout_index": 2, "_ng_rollout_id": bad} + if bad is None: + # Absent and null both mean "no explicit id", so the derivation still runs. + assert maybe_rollout_id_from_run_body(body) == "3-2" + return + # An id that cannot survive the round trip is refused rather than rewritten: correlating under + # a sanitized id would file records under a key the caller cannot look up afterwards. + with pytest.raises(ValueError): + maybe_rollout_id_from_run_body(body) + + +def test_explicit_rollout_ids_round_trip_through_the_path_prefix(): + from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body + from nemo_gym.rollout_correlation import RolloutContextMiddleware + + # The id becomes a path segment, so every id the body accepts has to be one the middleware + # gives back unchanged. A charset the two disagreed on would correlate calls to nothing. + for candidate in ["s7-3-2", "step7.task3", "a", "A_b-1.2"]: + rollout_id = maybe_rollout_id_from_run_body({"_ng_rollout_id": candidate}) + match = RolloutContextMiddleware._PREFIX.match(f"/ng-rollout/{rollout_id}/v1/responses") + assert match is not None and match.group("rollout_id") == candidate + assert match.group("rest") == "/v1/responses" + + def _capture_exchange(dialect, model_server, usage, response): return { "model_call_id": f"call-{model_server}", diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index efb314c342..91502bdc56 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -1000,6 +1000,57 @@ def run_examples(self, examples, *args, **kwargs): if redact_payloads: assert "data:image/png;base64,secret" not in orjson.dumps(results[0]).decode() + async def test_run_from_config_keys_capture_by_an_explicit_rollout_id( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from nemo_gym.base_responses_api_model import CaptureStore + from nemo_gym.global_config import ROLLOUT_ID_KEY_NAME + + capture_dir = tmp_path / "captures" + monkeypatch.setattr( + nemo_gym.rollout_collection, + "get_global_config_dict", + lambda: {"observability_enabled": True, "model_call_capture_dir": str(capture_dir)}, + ) + + # Indices that would derive "0-0" plus an explicit id. The explicit id has to win on both + # sides, or the writer and the reader key the same rollout differently and the readback + # finds nothing. + source_row = { + "responses_create_params": {"input": []}, + AGENT_REF_KEY_NAME: {"name": "agent"}, + ROLLOUT_ID_KEY_NAME: "step7.0-0", + } + input_fpath = tmp_path / "input.jsonl" + input_fpath.write_bytes(orjson.dumps(source_row) + b"\n") + config = RolloutCollectionConfig( + input_jsonl_fpath=str(input_fpath), + output_jsonl_fpath=str(tmp_path / "output.jsonl"), + resume_from_cache=False, + disable_aggregation=True, + ) + + store = CaptureStore(capture_dir) + + class Helper(RolloutCollectionHelper): + def run_examples(self, examples, *args, **kwargs): + [example] = examples + store.record( + "step7.0-0", + {"model_call_id": "call", "dialect": "responses", "request": {}, "response": {}}, + ) + future = Future() + future.set_result((example, {"response": {"usage": {}}})) + return [future] + + results = await Helper().run_from_config(config) + + assert results[0][ROLLOUT_ID_KEY_NAME] == "step7.0-0" + assert [call["model_call_id"] for call in results[0]["ng_model_call_capture"]["calls"]] == ["call"] + # Nothing was filed under the derived id, so the explicit id replaced it rather than + # sitting alongside it. + assert store.read("0-0") == [] + async def test_run_from_config_sorted(self, tmp_path: Path, empty_global_config: MagicMock) -> None: input_jsonl_fpath = tmp_path / "input.jsonl" samples = [ 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..481089990b --- /dev/null +++ b/tests/unit_tests/test_token_id_capture.py @@ -0,0 +1,843 @@ +# 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 json +import logging +import multiprocessing +import subprocess +import sys +from time import time +from unittest.mock import MagicMock, patch +from uuid import uuid4 + +import pytest +from fastapi import Body, Request +from fastapi.testclient import TestClient +from pydantic import ValidationError + +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 ( + TOKEN_ENTRY_RECORD_SCHEMA_VERSION, + TOKEN_FIELDS, + CaptureContext, + TokenCaptureStore, + TokenEntry, + TokenIdCaptureConfig, + commit_entry, + extract_token_fields, + install_token_sink, + reset_token_sink, + set_token_sink, +) +from nemo_gym.token_id_capture.protocols import TokenSource +from nemo_gym.token_id_capture.store 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") == [] + + +# --- config ------------------------------------------------------------------- + + +def _block(**kwargs) -> dict: + return {"token_id_capture": {"enabled": True, **kwargs}} + + +def test_config_disabled_needs_no_dir(): + cfg = TokenIdCaptureConfig.model_validate({}) + assert cfg.enabled is False + assert make_token_store({}) is None + + +def test_config_enabled_requires_absolute_dir(tmp_path): + with pytest.raises(ValueError): + TokenIdCaptureConfig.model_validate(_block()) + with pytest.raises(ValueError): + TokenIdCaptureConfig.model_validate(_block(dir="relative/dir")) + cfg = TokenIdCaptureConfig.model_validate(_block(dir=str(tmp_path))) + assert cfg.resolved_dir() == tmp_path + + +def test_config_falls_back_to_model_call_capture_dir(tmp_path): + cfg = TokenIdCaptureConfig.model_validate(_block() | {"model_call_capture_dir": str(tmp_path)}) + assert cfg.resolved_dir() == tmp_path + + +def test_config_keeps_settings_when_capture_is_off(tmp_path): + """Templated configs set a directory unconditionally and toggle `enabled` per run, so the + rest of the block is left alone rather than rejected.""" + cfg = TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": False, "dir": str(tmp_path)}}) + assert cfg.enabled is False + assert cfg.build_sink() is None + + +def test_config_warns_rather_than_fails_on_a_sink_beside_a_directory(caplog): + """Nothing is lost, the directory is just never read, but someone expecting files on disk + will not find any.""" + with caplog.at_level(logging.WARNING): + cfg = TokenIdCaptureConfig.model_validate(_block(sink=f"{__name__}:_ConfiguredSink", dir="/tmp/x")) + assert cfg.enabled is True + assert "will not be written to" in caplog.text + + +def test_config_rejects_an_unknown_key(): + """A typo in this block silently disables capture, so it is refused at startup.""" + with pytest.raises(ValueError): + TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": True, "dirr": "/tmp/x"}}) + + +# --- source / readers --------------------------------------------------------- + + +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, "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_token_arrays_are_stored_once(tmp_path): + """The served response carries the arrays on an output item; the record does not repeat them. + + Storing them again per item roughly doubles a record, and the per-item copy is not the + value a trainer reads: an item's prompt in a chained trajectory is the running sequence. + """ + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + client.post("/ng-rollout/task0-rollDedup/v1/responses", json={"input": "hi"}) + entry = TokenCaptureStore(tmp_path).read_entries("task0-rollDedup")[0] + assert entry.generation_token_ids == GTOKS + for item in entry.output_items: + assert not any(field in item for field in TOKEN_FIELDS) + # Content is kept; only the arrays move off. + assert entry.output_items[-1]["content"][0]["text"] == "hi from responses" + # Which item they came off, so a consumer can put the chain-correct values back. + assert entry.token_item_index == len(entry.output_items) - 1 + + +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, "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_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_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_concurrent_appends_to_one_rollout_stay_intact(tmp_path): + """Writes take an exclusive file lock, which is what keeps two writers from interleaving a + partial line. Under sharding the writers are separate processes, so the lock has to hold there + too; this covers the same code path with threads.""" + import concurrent.futures + + store = TokenCaptureStore(tmp_path) + entries = [ + TokenEntry( + rollout_id="r0", + model_call_id=f"call-{i}", + prompt_token_ids=list(range(200)), + generation_token_ids=[i] * 64, + generation_log_probs=[-0.1] * 64, + ) + for i in range(32) + ] + + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool: + list(pool.map(store.append, entries)) + + read_back = store.read_entries("r0") + assert len(read_back) == 32 + assert sorted(e.model_call_id for e in read_back) == sorted(e.model_call_id for e in entries) + # Every line parsed, so no write landed inside another. + assert all(len(e.generation_token_ids) == 64 for e in read_back) + + +# --- framework-owned sink: the documented extension point --------------------- + + +class _RecordingSink: + """A sink that is only a ``TokenSink``: no file store, no directory. + + Deliberately not a ``TokenCaptureStore`` subclass. A training framework whose sink is + its own transport has nothing on disk, and this is the shape the capture path has to + accept for ``install_token_sink`` to mean anything. + """ + + def __init__(self) -> None: + self.entries: list[TokenEntry] = [] + self.incomplete: list[tuple[str, str]] = [] + + async def put(self, entry: TokenEntry) -> None: + self.entries.append(entry) + + def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + self.incomplete.append((rollout_id, model_call_id)) + + +@pytest.fixture +def installed_sink(): + sink = _RecordingSink() + install_token_sink(sink) + try: + yield sink + finally: + install_token_sink(None) + + +def test_installed_sink_receives_entries_without_a_capture_dir(installed_sink): + """The framework path: capture on, no directory anywhere, records still arrive.""" + config = {"token_id_capture": {"enabled": True}} + client = TestClient(_server(config).setup_webserver()) + resp = client.post("/ng-rollout/task0-sink0/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + assert len(installed_sink.entries) == 1 + assert installed_sink.entries[0].generation_token_ids == GTOKS + assert installed_sink.entries[0].rollout_id == "task0-sink0" + + +def test_config_allows_no_directory_when_a_sink_is_installed(installed_sink): + """Requiring a directory would block the sink-only deployment the docstring describes.""" + assert TokenIdCaptureConfig.model_validate(_block()).resolved_dir() is None + + +def test_config_still_requires_a_directory_with_no_sink_installed(): + with pytest.raises(ValueError): + TokenIdCaptureConfig.model_validate(_block()) + + +def test_installed_sink_is_marked_incomplete_through_the_protocol(installed_sink, monkeypatch): + """A protocol-only sink must receive the incomplete signal. + + Reaching for a concrete store attribute here would raise inside the failure path and be + swallowed, leaving a rollout that lost a call looking complete. + """ + + async def boom(entry): + raise RuntimeError("transport down") + + monkeypatch.setattr(installed_sink, "put", boom) + client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) + resp = client.post("/ng-rollout/task0-sink1/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 # capture never fails the model call + assert installed_sink.incomplete == [("task0-sink1", installed_sink.incomplete[0][1])] + + +def test_a_sink_without_mark_incomplete_is_logged_not_swallowed(caplog): + """The signal cannot be lost quietly: that is the outcome the failure path exists to stop.""" + + class _PutOnlySink: + async def put(self, entry): + raise RuntimeError("transport down") + + install_token_sink(_PutOnlySink()) + try: + client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) + with caplog.at_level(logging.ERROR): + resp = client.post("/ng-rollout/task0-sink2/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + assert any("does not implement mark_incomplete" in r.message for r in caplog.records) + finally: + install_token_sink(None) + + +def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(installed_sink): + """Engine-side capture: the caller has the arrays, the served response does not. + + The commit half has to be reachable on its own, otherwise a framework in that position + forks the durability ordering rather than sharing it. + """ + entry = TokenEntry( + rollout_id="task0-sink3", + model_call_id="mc-1", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + token = set_token_sink(CaptureContext(rollout_id="task0-sink3", model_call_id="mc-1", sink=installed_sink)) + try: + asyncio.run(commit_entry(entry)) + finally: + reset_token_sink(token) + assert len(installed_sink.entries) == 1 + assert installed_sink.entries[0].generation_token_ids == GTOKS + + +def test_records_carry_a_schema_version(): + """Writer and reader are different processes and may be different repositories.""" + entry = TokenEntry( + rollout_id="r", + model_call_id="c", + prompt_token_ids=[1], + generation_token_ids=[2], + generation_log_probs=[-0.1], + ) + assert entry.schema_version == TOKEN_ENTRY_RECORD_SCHEMA_VERSION + assert "schema_version" in entry.model_dump_json() + + +def test_a_malformed_token_payload_does_not_fail_the_model_call(installed_sink): + """Building the record is guarded, not just writing it. + + ``capture_tokens`` is awaited directly on the model server's response path, so anything + it raises fails the model call. A payload whose token fields do not validate has to be + treated like any other capture failure: the call succeeds and the rollout is marked. + """ + entry_ctor = TokenEntry + + def _bad_entry(**kwargs): + # Stand in for a payload that fails validation, e.g. token ids that are not integers. + raise ValueError("prompt_token_ids: not a list of ints") + + with patch("nemo_gym.token_id_capture.sink.TokenEntry", _bad_entry): + client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) + resp = client.post("/ng-rollout/task0-bad0/v1/responses", json={"input": "hi"}) + + assert resp.status_code == 200, "a malformed token payload must not fail the model call" + assert installed_sink.entries == [], "nothing should have been written" + assert [r for r, _ in installed_sink.incomplete] == ["task0-bad0"], ( + "the rollout lost a call and must not look complete" + ) + assert entry_ctor is TokenEntry # patch scoped + + +@pytest.mark.parametrize("bad", ["", "a/b", "../escape", "a b"]) +def test_an_unsafe_rollout_id_is_rejected(tmp_path, bad): + """The id names the capture file, so it has to be a safe filename component: a separator + would let a rollout id write outside the store directory.""" + with pytest.raises(ValueError): + TokenCaptureStore(tmp_path).path_for(bad) + + +def test_a_record_is_readable_as_soon_as_put_returns(tmp_path): + """``put`` is awaited rather than backgrounded, so the record is on disk before the model + call returns. A reader in another process runs after the rollout and has no way to wait for + a writer, and delete-on-consume is only safe because nothing is still in flight.""" + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c1", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + ) + asyncio.run(store.put(entry)) + assert [e.model_call_id for e in asyncio.run(store.tokens_for("r0"))] == ["c1"] + + +def test_a_rollout_that_lost_a_call_is_distinguishable_from_a_complete_one(tmp_path): + """Capture failures do not fail the model call, so nothing downstream would otherwise know + a turn is missing. The chain built from what survived can look perfectly contiguous.""" + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c1", + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + ) + asyncio.run(store.put(entry)) + assert not store.is_incomplete("r0") + store.mark_incomplete("r0", "c2") + assert store.is_incomplete("r0") + + +# --- where records go, and surviving multiple server workers ------------------- + + +class _ConfiguredSink: + """Constructed by dotted path, so every server process builds its own.""" + + entries: list = [] + + async def put(self, entry) -> None: + type(self).entries.append(entry) + + def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + +class _NotASink: + async def put(self, entry) -> None: + pass + + +class _NotCallableSink: + put = "not a method" + + def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + +class _KwargSink: + def __init__(self, endpoint: str, shard: int = 0) -> None: + self.endpoint, self.shard = endpoint, shard + + async def put(self, entry) -> None: + pass + + def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + +def test_a_configured_sink_receives_entries(tmp_path): + _ConfiguredSink.entries = [] + config = {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + client = TestClient(_server(config).setup_webserver()) + + assert client.post("/ng-rollout/task0-cfg0/v1/responses", json={"input": "hi"}).status_code == 200 + + assert [e.rollout_id for e in _ConfiguredSink.entries] == ["task0-cfg0"] + assert _ConfiguredSink.entries[0].generation_token_ids == GTOKS + + +def test_a_configured_sink_wins_over_an_installed_one(installed_sink): + """Both routes exist; the configured one is preferred because it survives extra workers.""" + _ConfiguredSink.entries = [] + config = {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + client = TestClient(_server(config).setup_webserver()) + + assert client.post("/ng-rollout/task0-cfg1/v1/responses", json={"input": "hi"}).status_code == 200 + + assert len(_ConfiguredSink.entries) == 1 + assert installed_sink.entries == [] + + +def test_a_sink_receives_its_configured_kwargs(): + """A sink for a real transport needs an endpoint and a client; a zero-argument one could only + get them from ambient state.""" + config = TokenIdCaptureConfig.model_validate( + _block(sink=f"{__name__}:_KwargSink", sink_kwargs={"endpoint": "https://dp", "shard": 3}) + ) + sink = config.build_sink() + assert (sink.endpoint, sink.shard) == ("https://dp", 3) + + +def test_a_sink_given_kwargs_it_cannot_take_is_refused_at_startup(): + config = TokenIdCaptureConfig.model_validate(_block(sink=f"{__name__}:_KwargSink", sink_kwargs={"nope": 1})) + with pytest.raises(ValueError, match="sink_kwargs"): + config.build_sink() + + +def test_a_sink_that_cannot_report_failures_is_refused_at_startup(): + """Without mark_incomplete an incomplete rollout looks complete, so this fails at startup + rather than at whichever step first loses a call.""" + config = TokenIdCaptureConfig.model_validate( + {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_NotASink"}} + ) + with pytest.raises(ValueError, match="mark_incomplete"): + config.build_sink() + + +def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): + """isinstance against a Protocol only checks that the attributes exist, so callability is + checked too. Both are derived from the protocol rather than a list written out here, so the + check keeps up if TokenSink gains a method.""" + config = TokenIdCaptureConfig.model_validate( + {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_NotCallableSink"}} + ) + with pytest.raises(ValueError, match="put"): + config.build_sink() + + +@pytest.mark.parametrize( + "target, expected", + [("no_colon", "module.path:ClassName"), ("nemo_gym.token_id_capture:Nope", "could not load")], +) +def test_a_malformed_sink_path_is_refused_at_startup(target, expected): + config = TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": True, "sink": target}}) + with pytest.raises(ValueError, match=expected): + config.build_sink() + + +def test_a_programmatically_installed_sink_does_not_reach_a_spawned_worker(): + """A model server with num_workers > 1 is launched by uvicorn with an app string and + workers=N, and uvicorn spawns those workers (multiprocessing "spawn"), re-importing the app + module rather than inheriting the launcher's memory. ``install_token_sink`` sets a process + global, so it does not cross that boundary and capture silently falls back to the file store, + or writes nothing when no directory is set. + + This is why ``token_id_capture.sink`` is configuration rather than only a function call: it is + constructed inside each worker. The test pins the limitation so the reason for the config key + does not get lost. + """ + ctx = multiprocessing.get_context("spawn") # the context uvicorn uses + queue = ctx.Queue() + process = ctx.Process(target=_report_installed_sink, args=(queue,)) + process.start() + process.join(timeout=60) + + assert queue.get(timeout=10) == "None" + + +def _report_installed_sink(queue) -> None: + # Runs in the spawned process, which re-imports rather than inheriting. + from nemo_gym.token_id_capture import installed_token_sink + + queue.put(repr(installed_token_sink())) + + +def test_the_store_is_a_token_source(tmp_path): + """Records are read back through a TokenSource, and the file store is one. There is no + separate local reader: a wrapper over the store would only forward every call.""" + store = TokenCaptureStore(tmp_path) + assert isinstance(store, TokenSource) + + store.append( + TokenEntry( + rollout_id="r0", + model_call_id="c1", + prompt_token_ids=[1], + generation_token_ids=[2], + generation_log_probs=[-0.1], + ) + ) + assert [e.model_call_id for e in asyncio.run(store.tokens_for("r0"))] == ["c1"] + + # A colocated source can tell that a call failed to capture, which is what keeps an + # incomplete rollout from being trained on. + assert store.is_incomplete("r0") is False + store.mark_incomplete("r0", "c2") + assert store.is_incomplete("r0") is True + + +def _entry_fields(**overrides): + return dict( + rollout_id="r0", + model_call_id="c1", + prompt_token_ids=[1], + generation_token_ids=[2], + generation_log_probs=[-0.1], + **overrides, + ) + + +def test_a_record_older_than_this_reader_is_accepted(): + """A field this reader does not have takes its default and the consumer degrades: a record + written before parent links existed simply has none, and the builder matches prefixes.""" + entry = TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_RECORD_SCHEMA_VERSION - 1)) + assert entry.generation_token_ids == [2] + + +def test_a_record_newer_than_this_reader_is_refused(): + """The direction extra="allow" hides. A field this reader cannot see is kept and ignored, so + without this the record decodes clean and trains as though nothing were different.""" + with pytest.raises(ValidationError, match="this reader understands up to"): + TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_RECORD_SCHEMA_VERSION + 1)) + + +def test_a_newer_record_in_the_store_fails_the_read_rather_than_being_skipped(tmp_path): + """Read failure is the loud path: the caller marks that rollout unusable rather than training + on a partial set that looks complete.""" + store = TokenCaptureStore(tmp_path) + store.append(TokenEntry(**_entry_fields())) + path = next(tmp_path.glob("*.tokens.jsonl")) + record = json.loads(path.read_text().splitlines()[0]) + record["schema_version"] = TOKEN_ENTRY_RECORD_SCHEMA_VERSION + 1 + path.write_text(json.dumps(record) + "\n") + + with pytest.raises(ValidationError): + store.read_entries("r0") From 40056928dba6da06580bd8600fcd7b7bae9152f5 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 6 Aug 2026 17:47:12 -0700 Subject: [PATCH 02/12] fix(token-id-capture): mark a rollout incomplete when a call returns no token ids capture_tokens returned quietly when a response carried no token ids, so a rollout that lost a call looked identical to a complete one. The builder reads the gap between one call's tokens and the next call's prompt as tool output, which closes the chain over the hole: the missing call's generated tokens are delivered inside the next prompt at mask 0, and tokens the policy sampled train as if the environment had written them. Mark the rollout instead, on the same path an exception already takes. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/sink.py | 20 +++++++++++ tests/unit_tests/test_token_id_capture.py | 43 +++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index e4b2822346..5aa48eff48 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -99,9 +99,11 @@ async def capture_tokens(response: Any) -> None: elif isinstance(response, dict): payload = response else: + _capture_missing(sink, f"the response is a {type(response).__name__}") return info = extract_token_fields(payload) if info is None: + _capture_missing(sink, "the response carries no token ids") return # Content only: the arrays live on the entry, not on the items as well. content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) @@ -167,6 +169,24 @@ def _capture_failed(sink: CaptureContext, stage: str) -> None: _mark_incomplete(sink) +def _capture_missing(sink: CaptureContext, reason: str) -> None: + """Mark the rollout when a call produced no record and nothing raised. + + An active sink means this call belongs to a rollout being captured, so a response with no + token ids is a hole in the chain rather than traffic to skip. The builder reads the gap + between one call's tokens and the next call's prompt as tool output, so a skipped call's + generated tokens arrive inside the next prompt at mask 0, and tokens the policy sampled + train as if the environment had written them. + """ + logger.warning( + "Training-token capture has no token ids for model call %s of rollout %s: %s.", + sink.model_call_id, + sink.rollout_id, + reason, + ) + _mark_incomplete(sink) + + def _mark_incomplete(sink: CaptureContext) -> None: """Mark the rollout, or say loudly why it could not be marked. diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 481089990b..8165bd8941 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -408,6 +408,49 @@ async def boom(self, entry): assert store.is_incomplete("fail0-roll0") +class _SilentModel(_CapturingModel): + """A model server that answers normally but returns no token ids.""" + + async def responses( + self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body() + ) -> NeMoGymResponse: + response = _training_response("hi with no tokens") + for field in ("prompt_token_ids", "generation_token_ids", "generation_log_probs"): + setattr(response.output[0], field, None) + return response + + +def _silent_server(global_config_dict) -> SimpleResponsesAPIModel: + return _SilentModel( + config=BaseResponsesAPIModelConfig(host="0.0.0.0", port=8099, entrypoint="", name="srv"), + server_client=MagicMock(spec=ServerClient, global_config_dict=global_config_dict), + ) + + +def test_a_response_without_token_ids_marks_the_rollout_incomplete(tmp_path): + """A call that returns no token ids is a hole, not traffic to skip. + + Skipping it quietly is the worse failure: the rollout still looks complete, and the + call's generated tokens end up inside the next call's prompt, where they are trained + as if the environment had written them. + """ + client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post("/ng-rollout/silent0-roll0/v1/responses", json={"input": "hi"}) + # The model call itself still succeeds; capture never breaks the harness's run. + assert resp.status_code == 200 + + store = TokenCaptureStore(tmp_path) + assert store.read_entries("silent0-roll0") == [] + assert store.is_incomplete("silent0-roll0") + + +def test_untagged_traffic_without_token_ids_marks_nothing(tmp_path): + """No rollout prefix means no sink, so there is no rollout to call incomplete.""" + client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) + assert client.post("/v1/responses", json={"input": "hi"}).status_code == 200 + assert list(tmp_path.glob("**/*.incomplete")) == [] + + def test_delete_removes_records_and_marker(tmp_path): store = TokenCaptureStore(tmp_path) store.append( From 5f6eec8d0cff01cc9a3f922998278b687ae15d93 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Mon, 17 Aug 2026 17:09:08 -0700 Subject: [PATCH 03/12] fix(token-id-capture): make capture lifecycle fail closed Separate rollout correlation from capture intent and expose durable paired transport contracts so external frameworks can integrate without silent partial training data. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_agent.py | 36 ++- nemo_gym/base_responses_api_model.py | 54 ++-- nemo_gym/config_types.py | 1 + nemo_gym/server_utils.py | 12 +- nemo_gym/token_id_capture/__init__.py | 8 + nemo_gym/token_id_capture/config.py | 68 +++-- nemo_gym/token_id_capture/protocols.py | 80 ++--- nemo_gym/token_id_capture/records.py | 21 +- nemo_gym/token_id_capture/sink.py | 75 +++-- nemo_gym/token_id_capture/store.py | 191 +++++++++--- .../test_base_responses_api_model.py | 33 +- tests/unit_tests/test_token_id_capture.py | 289 +++++++++++++++--- 12 files changed, 670 insertions(+), 198 deletions(-) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 579858115b..5b8e263a09 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -26,7 +26,7 @@ BaseRunRequest, BaseVerifyResponse, ) -from nemo_gym.config_types import ROLLOUT_PATH_PREFIX +from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT from nemo_gym.global_config import ( OBSERVABILITY_ENABLED_KEY_NAME, TOKEN_ID_CAPTURE_BLOCK, @@ -76,6 +76,7 @@ def setup_webserver(self) -> FastAPI: # responses() recovers the rollout id from the path (see url_path_for_request) to correlate # its model calls. Same handler, so unprefixed calls are unaffected. app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses) + app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/{TOKEN_CAPTURE_PATH_SEGMENT}/v1/responses")(self.responses) run = self.run @@ -104,14 +105,24 @@ def _capture_correlation_enabled(self) -> bool: Fail closed: an agent whose client carries no usable global config runs uncorrelated rather than erroring on every model call. """ + return self._model_call_capture_enabled() or self._token_id_capture_enabled() + + def _model_call_capture_enabled(self) -> bool: + """Whether evaluation model-call observability is enabled.""" + 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)) + + def _token_id_capture_enabled(self) -> bool: + """Whether this agent explicitly opted into training-token capture.""" global_config = getattr(self.server_client, "global_config_dict", None) if not isinstance(global_config, Mapping): return False block = global_config.get(TOKEN_ID_CAPTURE_BLOCK) or {} - token_capture = bool(isinstance(block, Mapping) and block.get("enabled", False)) and bool( - getattr(self.config, "token_id_capture", False) + return bool(isinstance(block, Mapping) and block.get("enabled", False)) and bool( + getattr(getattr(self, "config", None), "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). @@ -131,7 +142,10 @@ def url_path_for_run(self, url_path: str, body: Any) -> str: handling ``/run`` — both direct model-server calls and self-calls to ``/v1/responses`` (the prefixed self-call route carries the id into ``responses()``). """ - return f"{rollout_path_prefix(self.rollout_id_from_run(body))}{url_path}" + return ( + f"{rollout_path_prefix(self.rollout_id_from_run(body), token_capture=self._token_id_capture_enabled())}" + f"{url_path}" + ) def base_url_for_run(self, base_url: str, body: Any) -> str: """A model-server base URL with the per-rollout capture-correlation prefix applied. @@ -140,7 +154,11 @@ def base_url_for_run(self, base_url: str, body: Any) -> str: harnesses that configure a client once instead of prefixing each call: same gating, applied to a server root URL (append the API-version suffix afterwards). """ - return apply_rollout_prefix(base_url, self.rollout_id_from_run(body)) + return apply_rollout_prefix( + base_url, + self.rollout_id_from_run(body), + token_capture=self._token_id_capture_enabled(), + ) def url_path_for_request(self, url_path: str, request: Optional[Request]) -> str: """Carry an inbound ``/ng-rollout/`` self-call prefix onto a downstream url_path. @@ -151,13 +169,15 @@ def url_path_for_request(self, url_path: str, request: Optional[Request]) -> str """ path_params = getattr(request, "path_params", None) rollout_id = path_params.get("rollout_id") if isinstance(path_params, Mapping) else None - return f"{rollout_path_prefix(rollout_id)}{url_path}" + request_path = getattr(getattr(request, "url", None), "path", "") + token_capture = f"/{TOKEN_CAPTURE_PATH_SEGMENT}/" in request_path + return f"{rollout_path_prefix(rollout_id, token_capture=token_capture)}{url_path}" def resolve_model_base_url(self, model_server_name: str, rollout_id: Optional[str] = None) -> str: """Resolve a model-server URL with an optional rollout prefix.""" server_config = get_first_server_config_dict(self.server_client.global_config_dict, model_server_name) base_url = self.server_client._build_server_base_url(server_config) - return f"{apply_rollout_prefix(base_url, rollout_id)}/v1" + return f"{apply_rollout_prefix(base_url, rollout_id, token_capture=SimpleResponsesAPIAgent._token_id_capture_enabled(self))}/v1" # TODO: right now there is no validation on the TypedDict NeMoGymResponseCreateParamsNonStreaming # We should explicitly add validation at this server level or we should explicitly not validate so that there is flexibility in this API. diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 2dd4307b9b..8f5afe8c1c 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -47,7 +47,7 @@ from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.chat_streaming import sanitize_streaming_chat_body, synthesize_chat_completion_sse -from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, ModelServerRef +from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT, ModelServerRef from nemo_gym.openai_utils import ( NeMoGymChatCompletion, NeMoGymChatCompletionCreateParamsNonStreaming, @@ -88,26 +88,6 @@ _ANTHROPIC_CONVERTER = AnthropicConverter() -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 @@ -779,7 +759,11 @@ def _consume_terminal_sse_event(buffer: bytearray, dialect: str) -> Optional[str # Consumer side of the URL-prefix protocol: strip /ng-rollout/ before routing, key capture by # . The constant + producer (apply_rollout_prefix) are in server_utils. -_ROLLOUT_PATH_RE = re.compile(rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P[^/]+)(?P/.*)$") +_ROLLOUT_PATH_RE = re.compile( + rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P[^/]+)" + rf"(?:/(?P{re.escape(TOKEN_CAPTURE_PATH_SEGMENT)}))?" + rf"(?P/.*)$" +) def make_capture_store(config: ModelCallCaptureConfig) -> Optional[CaptureStore]: @@ -1080,6 +1064,7 @@ def __init__( model_server_name: str | None, token_store: Any = None, configured_sink: Any = None, + token_capture_enabled: bool = False, ) -> None: self._app = app self._store = store @@ -1088,6 +1073,10 @@ def __init__( self._token_store = token_store # Built from token_id_capture.sink, once, in this process. self._configured_sink = configured_sink + # Capture can be on with no destination in this process, when a framework stages records + # from the inference worker instead. The identity and the parent resolution still have to + # happen here, so enablement rather than a destination decides whether this runs. + self._token_capture_enabled = token_capture_enabled async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: if scope.get("type") != "http": @@ -1096,9 +1085,11 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None path = scope.get("path", "") rollout_from_path: Optional[str] = None + token_capture_requested = False prefix_match = _ROLLOUT_PATH_RE.match(path) if prefix_match: rollout_from_path = prefix_match.group("rollout_id") + token_capture_requested = prefix_match.group("token_capture") is not None path = prefix_match.group("rest") scope = {**scope, "path": path, "raw_path": path.encode("utf-8")} @@ -1115,7 +1106,8 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None # installed sink is still resolved per request, so one installed after the app is built # still takes effect. token_sink = self._configured_sink or installed_token_sink() or self._token_store - if (self._store is None and token_sink is None) or rollout_from_path is None or dialect is None: + capture_wanted = token_capture_requested and (token_sink is not None or self._token_capture_enabled) + if (self._store is None and not capture_wanted) or rollout_from_path is None or dialect is None: await self._app(scope, receive, send) return @@ -1124,8 +1116,11 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None # Hand the model server a per-request token sink keyed to this call. It records token ids # from its complete response. The middleware cannot: token ids are dropped on the SSE wire. + # Set whenever capture is on, even with no destination here. The context carries the + # identity a staged record is keyed by and the parent a request continues, and both are + # resolved in this process whether or not it is the one that writes. sink_token = None - if token_sink is not None: + if capture_wanted: sink_token = set_token_sink( CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, sink=token_sink) ) @@ -1318,12 +1313,23 @@ def install_model_call_capture( configured_sink = ( token_id_capture_config(global_config_dict).build_sink() if global_config_dict is not None else None ) + owned_sinks = [sink for sink in (configured_sink, token_store) if sink is not None] + + async def _close_token_sinks() -> None: + for sink in owned_sinks: + await sink.close() + + if owned_sinks: + app.add_event_handler("shutdown", _close_token_sinks) app.add_middleware( _CaptureMiddleware, store=make_capture_store(config), model_server_name=model_server_name, token_store=token_store, configured_sink=configured_sink, + token_capture_enabled=( + token_id_capture_config(global_config_dict).enabled if global_config_dict is not None else False + ), ) diff --git a/nemo_gym/config_types.py b/nemo_gym/config_types.py index bbf04daf9b..5deac5d144 100644 --- a/nemo_gym/config_types.py +++ b/nemo_gym/config_types.py @@ -867,3 +867,4 @@ class AggregateMetrics(BaseModel): # Per-rollout model-call correlation. Callers place the rollout id in the model-server URL; # the capture middleware in base_responses_api_model.py strips this prefix before routing. ROLLOUT_PATH_PREFIX = "ng-rollout" +TOKEN_CAPTURE_PATH_SEGMENT = "token-capture" diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index c2ba7a674b..6701b8956d 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -55,6 +55,7 @@ from nemo_gym import WORKING_DIR from nemo_gym.config_types import ( ROLLOUT_PATH_PREFIX, + TOKEN_CAPTURE_PATH_SEGMENT, BaseRunServerInstanceConfig, BaseServerConfig, ) @@ -833,16 +834,19 @@ def get_server_url(server_name: str) -> str: return f"http://{model_server_config['host']}:{model_server_config['port']}" -def rollout_path_prefix(rollout_id: Optional[str]) -> str: +def rollout_path_prefix(rollout_id: Optional[str], *, token_capture: bool = False) -> str: """Return the leading model-server path prefix for a rollout, if available.""" - return f"/{ROLLOUT_PATH_PREFIX}/{rollout_id}" if rollout_id else "" + if not rollout_id: + return "" + capture_segment = f"/{TOKEN_CAPTURE_PATH_SEGMENT}" if token_capture else "" + return f"/{ROLLOUT_PATH_PREFIX}/{rollout_id}{capture_segment}" -def apply_rollout_prefix(base_url: str, rollout_id: Optional[str]) -> str: +def apply_rollout_prefix(base_url: str, rollout_id: Optional[str], *, token_capture: bool = False) -> str: """Append a rollout prefix to a model-server root URL.""" if not rollout_id: return base_url - return base_url.rstrip("/") + rollout_path_prefix(rollout_id) + return base_url.rstrip("/") + rollout_path_prefix(rollout_id, token_capture=token_capture) def setup_server_client(head_server_config: Optional[BaseServerConfig] = None) -> ServerClient: # pragma: no cover diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index 5b72148e67..558c01a15d 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -35,10 +35,13 @@ from nemo_gym.token_id_capture.config import TokenIdCaptureConfig from nemo_gym.token_id_capture.protocols import ( + TokenCaptureSnapshot, TokenSink, TokenSource, install_token_sink, + install_token_source, installed_token_sink, + installed_token_source, ) from nemo_gym.token_id_capture.records import ( TOKEN_ENTRY_RECORD_SCHEMA_VERSION, @@ -50,6 +53,7 @@ CaptureContext, capture_tokens, commit_entry, + current_capture_context, reset_token_sink, set_token_sink, ) @@ -67,11 +71,15 @@ "make_token_store", "TokenSink", "TokenSource", + "TokenCaptureSnapshot", "install_token_sink", + "install_token_source", "installed_token_sink", + "installed_token_source", "CaptureContext", "set_token_sink", "reset_token_sink", "capture_tokens", "commit_entry", + "current_capture_context", ] diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 975070bc92..3df8f7fe7a 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -59,9 +59,14 @@ from pathlib import Path from typing import Any -from pydantic import BaseModel, ConfigDict, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator -from nemo_gym.token_id_capture.protocols import TokenSink, installed_token_sink +from nemo_gym.token_id_capture.protocols import ( + TokenSink, + TokenSource, + installed_token_sink, + installed_token_source, +) logger = logging.getLogger(__name__) @@ -82,7 +87,12 @@ class TokenIdCaptureSettings(BaseModel): # Keyword arguments for that constructor: an endpoint, a client, credentials. A sink for a real # transport needs wiring, and a zero-argument one could only get it from ambient state. Use # ``${oc.env:VAR}`` for anything secret rather than writing it here. - sink_kwargs: dict[str, Any] = {} + sink_kwargs: dict[str, Any] = Field(default_factory=dict) + # Optional paired reader for framework-owned transports. + source: str | None = None + source_kwargs: dict[str, Any] = Field(default_factory=dict) + # Rebuild opaque-harness responses from captured records after the run. + rebuild_response: bool = True class TokenIdCaptureConfig(BaseModel): @@ -110,17 +120,24 @@ def _validate(self) -> "TokenIdCaptureConfig": "the file store, so %s will not be written to.", block.dir, ) + if block.rebuild_response and block.source is None and installed_token_source() is None: + raise ValueError( + "token_id_capture.source is required when a custom sink is used with rebuild_response=true" + ) return self directory = self.resolved_dir() if directory is None: # A process that installed a sink programmatically writes through that transport and # never constructs the file store, so it has no directory to give. - if installed_token_sink() is not None: + if installed_token_sink() is not None and ( + not block.rebuild_response or installed_token_source() is not None + ): return self - raise ValueError( - "token_id_capture.dir (or model_call_capture_dir) is required when " - "token_id_capture.enabled is true and no sink is configured or installed" - ) + if block.source is not None: + return self + if not block.rebuild_response: + return self + raise ValueError("token_id_capture requires a directory or paired source when rebuild_response=true") if not directory.is_absolute(): raise ValueError("training-token capture directory must be an absolute path") return self @@ -141,34 +158,39 @@ def build_sink(self) -> TokenSink | None: target = self.token_id_capture.sink if not self.token_id_capture.enabled or target is None: return None + return self._build_endpoint(target, self.token_id_capture.sink_kwargs, TokenSink, "sink") + + def build_source(self) -> TokenSource | None: + """Construct a configured framework-owned source.""" + target = self.token_id_capture.source + if not self.token_id_capture.enabled or target is None: + return None + return self._build_endpoint(target, self.token_id_capture.source_kwargs, TokenSource, "source") + + @staticmethod + def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: str): if ":" not in target: - raise ValueError(f"token_id_capture.sink must be 'module.path:ClassName' (got {target!r})") + raise ValueError(f"token_id_capture.{kind} must be 'module.path:ClassName' (got {target!r})") module_path, _, class_name = target.partition(":") try: factory = getattr(import_module(module_path), class_name) except (ImportError, AttributeError) as error: - raise ValueError(f"could not load token_id_capture.sink {target!r}: {error}") from error + raise ValueError(f"could not load token_id_capture.{kind} {target!r}: {error}") from error try: - sink = factory(**self.token_id_capture.sink_kwargs) + endpoint = factory(**kwargs) except TypeError as error: raise ValueError( - f"could not construct token_id_capture.sink {target!r} with " - f"sink_kwargs={sorted(self.token_id_capture.sink_kwargs)}: {error}" + f"could not construct token_id_capture.{kind} {target!r} with {kind}_kwargs={sorted(kwargs)}: {error}" ) from error - # Checked here rather than at first use: a sink that cannot record a failure makes an + # Checked here rather than at first use: an endpoint missing a lifecycle method makes an # incomplete rollout look complete, and a startup error is better than that at step 400. - # - # isinstance against the protocol rather than a list of names written out here, so this - # keeps up when TokenSink gains a method. It only checks that the attributes exist, so the - # loop below adds the part it does not do. Neither checks signatures; nothing at runtime - # can, short of calling the methods. - missing = [name for name in sorted(TokenSink.__protocol_attrs__) if not callable(getattr(sink, name, None))] - if missing or not isinstance(sink, TokenSink): + missing = [name for name in sorted(protocol.__protocol_attrs__) if not callable(getattr(endpoint, name, None))] + if missing or not isinstance(endpoint, protocol): raise ValueError( - f"token_id_capture.sink {target!r} does not satisfy TokenSink: " + f"token_id_capture.{kind} {target!r} does not satisfy {protocol.__name__}: " f"{', '.join(missing) or 'attribute check failed'}" ) - return sink + return endpoint def token_id_capture_config(global_config_dict: Any) -> TokenIdCaptureConfig: diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 7e238bd8ed..aed58d495d 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -34,82 +34,84 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Protocol, runtime_checkable from nemo_gym.token_id_capture.records import TokenEntry +@dataclass(frozen=True) +class TokenCaptureSnapshot: + """An immutable, sealed view of one rollout's capture records.""" + + rollout_id: str + entries: tuple[TokenEntry, ...] + incomplete: bool + seal_id: str + version: int + + @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. + """Durably store one record. - The record must be durable before this returns: 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. + Repeating the same call id with the same payload is a no-op. Reusing a + call id with a different payload or writing after seal must fail. - May raise. The caller counts the failure and marks the rollout, and - never fails the model call because of it. + May raise. The caller marks the rollout incomplete and never fails the + model call because of a capture error. """ ... - def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: - """Record that a call of this rollout failed to capture. + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + """Durably record that a call of this rollout failed to capture. The rollout is now missing a turn, and a consumer must mask the sample rather than train on a chain with a hole in it. The model call itself still succeeds, so this is the only signal that anything went wrong: a sink that drops it makes an incomplete rollout indistinguishable from a complete one. - - Synchronous, because the caller is a failure path that cannot await. A transport - that needs to send should queue here and flush elsewhere. """ ... + async def close(self) -> None: + """Flush pending work and release resources. Idempotent.""" + ... + @runtime_checkable class TokenSource(Protocol): - """Where a trajectory builder reads records from, and retires them afterwards.""" + """Where a trajectory builder seals, reads, and retires records.""" - async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: - """All records for a rollout, in any order. + async def seal(self, rollout_id: str) -> TokenCaptureSnapshot: + """Seal a rollout and return one atomic snapshot. - Order carries no meaning: calls run concurrently and may be served by - different workers. The builder recovers structure from the records - themselves, using parent links or token-prefix relationships. + Sealing is idempotent. No successful writes may occur after it returns. + Entry order carries no meaning. """ ... - async def drop(self, rollout_id: str) -> None: - """Retire a rollout's records once they have been consumed. + async def drop(self, rollout_id: str, *, seal_id: str, version: int) -> bool: + """Conditionally retire the exact sealed snapshot that was consumed. - A transport that cannot delete implements this as a no-op and leaves - retirement to whoever owns the storage. + Returns ``False`` if state changed after the snapshot. Implementations + that cannot delete return ``True`` and leave retention to their owner. """ ... - def is_incomplete(self, rollout_id: str) -> bool: - """Whether a call of this rollout failed to capture, as reported by - ``TokenSink.mark_incomplete``. - - The records that did arrive can stitch into a chain that looks perfectly - contiguous while missing a turn, so this is the only way to tell. A - transport that cannot tell returns False, the same way a transport that - cannot delete makes ``drop`` a no-op; the cost is that an incomplete - rollout of its is trained on rather than masked. - - Synchronous, to match ``mark_incomplete`` on the sink side. - """ - return False + async def close(self) -> None: + """Release resources. Idempotent.""" + ... # 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: TokenSink | None = None +_INSTALLED_SOURCE: TokenSource | None = None def install_token_sink(sink: TokenSink | None) -> None: @@ -120,3 +122,13 @@ def install_token_sink(sink: TokenSink | None) -> None: def installed_token_sink() -> TokenSink | None: return _INSTALLED_SINK + + +def install_token_source(source: TokenSource | None) -> None: + """Set (or clear, with ``None``) the process-wide default source.""" + global _INSTALLED_SOURCE + _INSTALLED_SOURCE = source + + +def installed_token_source() -> TokenSource | None: + return _INSTALLED_SOURCE diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index cb97e51b31..15826f3116 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -32,7 +32,7 @@ from typing import Any -from pydantic import BaseModel, ConfigDict, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator # The fields the model server attaches to a served response when token-id return @@ -78,7 +78,7 @@ class TokenEntry(BaseModel): routed_experts: Any | None = None # The served response's output items (Responses shape), content preserved, token # arrays removed. - output_items: list[dict] = [] + output_items: list[dict] = Field(default_factory=list) # Index into ``output_items`` of the item the token arrays were taken off, or null # when no item carried them. Records written before the arrays were de-duplicated # leave this unset and still carry the arrays inline, which the builder handles. @@ -105,6 +105,15 @@ def _refuse_a_newer_record(self) -> "TokenEntry": f"up to {TOKEN_ENTRY_RECORD_SCHEMA_VERSION}. Upgrade the reader, or point it at " "records written by a writer it matches." ) + if len(self.generation_token_ids) != len(self.generation_log_probs): + raise ValueError( + "generation_token_ids and generation_log_probs must have the same length " + f"(got {len(self.generation_token_ids)} and {len(self.generation_log_probs)})" + ) + if self.token_item_index is not None and not 0 <= self.token_item_index < len(self.output_items): + raise ValueError( + f"token_item_index {self.token_item_index} is outside output_items of length {len(self.output_items)}" + ) return self @@ -157,14 +166,18 @@ def extract_token_fields(response_json: dict) -> dict | None: carries token ids (e.g. token-id return is off, or an empty completion). """ candidates: list[dict] = [] + required = ("prompt_token_ids", "generation_token_ids", "generation_log_probs") for item in response_json.get("output") or []: - if isinstance(item, dict) and item.get("generation_token_ids") is not None: + if isinstance(item, dict) and any(item.get(field) is not None for field in required): 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: + if isinstance(message, dict) and any(message.get(field) is not None for field in required): candidates.append(message) if not candidates: return None source = candidates[-1] + missing = [field for field in required if source.get(field) is None] + if missing: + raise ValueError(f"partial token metadata is missing: {', '.join(missing)}") return {field: source.get(field) for field in TOKEN_FIELDS} diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 5aa48eff48..0903d7dddc 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -62,8 +62,14 @@ class CaptureContext: rollout_id: str model_call_id: str - sink: TokenSink + # None when capture is enabled but nothing in this process writes records. A framework that + # stages engine-side still needs the identity and the parent resolution this context carries, + # and there is no destination here to hand them to. + sink: TokenSink | None model: str = "" + # Set by ``commit_entry`` so the no-token-ids path can tell a call that was recorded by + # somebody else from one that was lost. + committed: bool = False _TOKEN_SINK: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_token_sink", default=None) @@ -73,6 +79,16 @@ def set_token_sink(sink: CaptureContext) -> Token: return _TOKEN_SINK.set(sink) +def current_capture_context() -> CaptureContext | None: + """The capture context for the in-flight call, or None for untagged traffic. + + The supported way to read the identity this call was minted with. + A framework that stages records from its inference worker keys them on that identity, + and this process is the only one that has it. + """ + return _TOKEN_SINK.get() + + def reset_token_sink(token: Token) -> None: _TOKEN_SINK.reset(token) @@ -99,11 +115,11 @@ async def capture_tokens(response: Any) -> None: elif isinstance(response, dict): payload = response else: - _capture_missing(sink, f"the response is a {type(response).__name__}") + await _capture_missing(sink, f"the response is a {type(response).__name__}") return info = extract_token_fields(payload) if info is None: - _capture_missing(sink, "the response carries no token ids") + await _capture_missing(sink, "the response carries no token ids") return # Content only: the arrays live on the entry, not on the items as well. content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) @@ -112,9 +128,9 @@ async def capture_tokens(response: Any) -> None: 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 [], + prompt_token_ids=info["prompt_token_ids"], + generation_token_ids=info["generation_token_ids"], + generation_log_probs=info["generation_log_probs"], routed_experts=info.get("routed_experts"), # Keep the content (assistant text, tool calls) so the trajectory the trainer # reads is not token-only, since text-based penalties need it. @@ -123,7 +139,7 @@ async def capture_tokens(response: Any) -> None: created_at=time.time(), ) except Exception: - _capture_failed(sink, "build") + await _capture_failed(sink, "build") return await commit_entry(entry) @@ -144,13 +160,25 @@ async def commit_entry(entry: TokenEntry) -> None: sink = _TOKEN_SINK.get() if sink is None: return + if entry.rollout_id != sink.rollout_id or entry.model_call_id != sink.model_call_id: + logger.warning( + "Training-token capture identity mismatch for model call %s of rollout %s.", + sink.model_call_id, + sink.rollout_id, + ) + await _mark_incomplete(sink) + return + if sink.sink is None: + sink.committed = True + return try: await sink.sink.put(entry) + sink.committed = True except Exception: - _capture_failed(sink, "write") + await _capture_failed(sink, "write") -def _capture_failed(sink: CaptureContext, stage: str) -> None: +async def _capture_failed(sink: CaptureContext, stage: str) -> None: """Report a capture failure without letting it reach the model call. Capture is best effort per call: a bad token payload must never fail the model call and @@ -166,28 +194,35 @@ def _capture_failed(sink: CaptureContext, stage: str) -> None: sink.rollout_id, exc_info=True, ) - _mark_incomplete(sink) + await _mark_incomplete(sink) -def _capture_missing(sink: CaptureContext, reason: str) -> None: - """Mark the rollout when a call produced no record and nothing raised. +async def _capture_missing(sink: CaptureContext, reason: str) -> None: + """Mark the rollout when a call this process should have recorded produced nothing. - An active sink means this call belongs to a rollout being captured, so a response with no - token ids is a hole in the chain rather than traffic to skip. The builder reads the gap - between one call's tokens and the next call's prompt as tool output, so a skipped call's - generated tokens arrive inside the next prompt at mask 0, and tokens the policy sampled - train as if the environment had written them. + A response with no token ids is a hole in the chain rather than traffic to skip. + The builder reads the gap between one call's tokens and the next call's prompt as tool output. + So a skipped call's generated tokens arrive inside the next prompt at mask 0, + and tokens the policy sampled train as if the environment had written them. + + Two cases are not holes and are left alone. + A call already committed through ``commit_entry`` was recorded by a caller that had the arrays + when this process did not. + A context with no sink means nothing here writes records at all, + so this process cannot tell a lost call from ordinary operation and the staging side owns that. """ + if sink.committed or sink.sink is None: + return logger.warning( "Training-token capture has no token ids for model call %s of rollout %s: %s.", sink.model_call_id, sink.rollout_id, reason, ) - _mark_incomplete(sink) + await _mark_incomplete(sink) -def _mark_incomplete(sink: CaptureContext) -> None: +async def _mark_incomplete(sink: CaptureContext) -> None: """Mark the rollout, or say loudly why it could not be marked. A sink that does not implement ``mark_incomplete`` would otherwise raise inside the @@ -205,6 +240,6 @@ def _mark_incomplete(sink: CaptureContext) -> None: ) return try: - mark(sink.rollout_id, sink.model_call_id) + await mark(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/store.py b/nemo_gym/token_id_capture/store.py index c999cc5351..2998793faf 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -34,11 +34,15 @@ import asyncio import fcntl import os +import tempfile +from contextlib import contextmanager from pathlib import Path from typing import Any +from uuid import uuid4 import orjson +from nemo_gym.token_id_capture.protocols import TokenCaptureSnapshot from nemo_gym.token_id_capture.records import TokenEntry @@ -67,38 +71,102 @@ 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. + def state_path_for(self, rollout_id: str) -> Path: + return self._root / f"{validate_rollout_id(rollout_id)}.tokens.state.json" - Capture is best effort per call, because 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. - """ + def lock_path_for(self, rollout_id: str) -> Path: + return self._root / f"{validate_rollout_id(rollout_id)}.tokens.lock" + + @contextmanager + def _locked(self, rollout_id: str, *, shared: bool = False): + with self.lock_path_for(rollout_id).open("a+b") as handle: + fcntl.flock(handle.fileno(), fcntl.LOCK_SH if shared else fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + def _read_state(self, rollout_id: str) -> dict[str, Any]: + path = self.state_path_for(rollout_id) + if not path.exists(): + return {"sealed": False, "incomplete": False, "seal_id": "", "version": 0} + state = orjson.loads(path.read_bytes()) + if not isinstance(state, dict): + raise ValueError(f"Invalid token-capture state for rollout {rollout_id}") + return state + + def _write_state(self, rollout_id: str, state: dict[str, Any]) -> None: + payload = orjson.dumps(state, option=orjson.OPT_SORT_KEYS | orjson.OPT_APPEND_NEWLINE) + with tempfile.NamedTemporaryFile(dir=self._root, prefix=".tokens-state-", delete=False) as handle: + temporary_path = Path(handle.name) + try: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + except BaseException: + temporary_path.unlink(missing_ok=True) + raise 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 + os.replace(temporary_path, self.state_path_for(rollout_id)) + self._fsync_root() + finally: + temporary_path.unlink(missing_ok=True) + + def _fsync_root(self) -> None: + descriptor = os.open(self._root, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + def _mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + with self._locked(rollout_id): + state = self._read_state(rollout_id) + state["incomplete"] = True + state["version"] = int(state.get("version", 0)) + 1 + self._write_state(rollout_id, state) + with self.incomplete_path_for(rollout_id).open("a", encoding="utf-8") as handle: + handle.write(f"{model_call_id}\n") + handle.flush() + os.fsync(handle.fileno()) + self._fsync_root() + + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + """Durably record that a call was lost.""" + await asyncio.to_thread(self._mark_incomplete, rollout_id, model_call_id) def is_incomplete(self, rollout_id: str) -> bool: - return self.incomplete_path_for(rollout_id).exists() + with self._locked(rollout_id, shared=True): + return bool(self._read_state(rollout_id).get("incomplete", False)) def append(self, entry: TokenEntry) -> None: - """Append one entry and fsync. Blocking file IO, so 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: + """Idempotently append one entry and fsync.""" + canonical = orjson.dumps(entry.model_dump(mode="json"), option=orjson.OPT_SORT_KEYS) + line = canonical + b"\n" + rollout_id = entry.rollout_id + with self._locked(rollout_id): + state = self._read_state(rollout_id) + if state.get("sealed", False): + raise RuntimeError(f"Token capture for rollout {rollout_id} is already sealed") + for existing in self._read_entries_unlocked(rollout_id): + if existing.model_call_id != entry.model_call_id: + continue + existing_bytes = orjson.dumps(existing.model_dump(mode="json"), option=orjson.OPT_SORT_KEYS) + if existing_bytes == canonical: + return + state["incomplete"] = True + state["version"] = int(state.get("version", 0)) + 1 + self._write_state(rollout_id, state) + raise ValueError( + f"Model call id {entry.model_call_id!r} was reused with a different payload " + f"for rollout {rollout_id!r}" + ) + with self.path_for(rollout_id).open("ab") as handle: handle.write(line) handle.flush() os.fsync(handle.fileno()) - finally: - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + state["version"] = int(state.get("version", 0)) + 1 + self._write_state(rollout_id, state) # --- TokenSink / TokenSource. The file store is Gym's default implementation of both; # a framework swaps in its own without touching the capture path. @@ -113,39 +181,78 @@ async def put(self, entry: TokenEntry) -> None: rollout never races a partial file.""" await asyncio.to_thread(self.append, entry) + async def seal(self, rollout_id: str) -> TokenCaptureSnapshot: + return await asyncio.to_thread(self._seal, rollout_id) + + def _seal(self, rollout_id: str) -> TokenCaptureSnapshot: + with self._locked(rollout_id): + state = self._read_state(rollout_id) + if not state.get("sealed", False): + state["sealed"] = True + state["seal_id"] = uuid4().hex + state["version"] = int(state.get("version", 0)) + 1 + self._write_state(rollout_id, state) + entries = tuple(self._read_entries_unlocked(rollout_id)) + return TokenCaptureSnapshot( + rollout_id=rollout_id, + entries=entries, + incomplete=bool(state.get("incomplete", False)), + seal_id=str(state["seal_id"]), + version=int(state["version"]), + ) + async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: - """``TokenSource``.""" + """Compatibility read for diagnostics. Consumers should use ``seal``.""" 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) + async def drop(self, rollout_id: str, *, seal_id: str, version: int) -> bool: + """Conditionally delete the sealed snapshot.""" + return await asyncio.to_thread(self._drop, rollout_id, seal_id, version) + + def _drop(self, rollout_id: str, seal_id: str, version: int) -> bool: + with self._locked(rollout_id): + state = self._read_state(rollout_id) + if ( + not state.get("sealed", False) + or state.get("seal_id") != seal_id + or int(state.get("version", 0)) != version + ): + return False + self.path_for(rollout_id).unlink(missing_ok=True) + self.incomplete_path_for(rollout_id).unlink(missing_ok=True) + self.state_path_for(rollout_id).unlink(missing_ok=True) + self._fsync_root() + return True + + async def close(self) -> None: + """The file store owns no persistent handles.""" def delete(self, rollout_id: str) -> None: - """Remove a rollout's records and its incomplete marker. + """Unconditionally remove a rollout's records. - 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. + This compatibility helper is for administrative cleanup. Normal + consumers use conditional ``drop``. """ - self.path_for(rollout_id).unlink(missing_ok=True) - self.incomplete_path_for(rollout_id).unlink(missing_ok=True) + with self._locked(rollout_id): + self.path_for(rollout_id).unlink(missing_ok=True) + self.incomplete_path_for(rollout_id).unlink(missing_ok=True) + self.state_path_for(rollout_id).unlink(missing_ok=True) + self._fsync_root() def read_entries(self, rollout_id: str) -> list[TokenEntry]: + with self._locked(rollout_id, shared=True): + return self._read_entries_unlocked(rollout_id) + + def _read_entries_unlocked(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) + for line in handle: + stripped = line.strip() + if stripped: + entries.append(TokenEntry.model_validate(orjson.loads(stripped))) return entries diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 952950ce18..8d47361004 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -808,7 +808,7 @@ def test_base_agent_resolve_model_base_url(monkeypatch): assert SimpleResponsesAPIAgent.resolve_model_base_url(agent, "model", None) == "http://h:1/v1" -def _make_base_agent(global_config): +def _make_base_agent(global_config, *, token_id_capture=False): from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig class _Agent(SimpleResponsesAPIAgent): @@ -820,7 +820,13 @@ async def run(self, body=...): server_client = MagicMock(spec=ServerClient) server_client.global_config_dict = global_config - config = BaseResponsesAPIAgentConfig(host="", port=0, entrypoint="", name="agent") + config = BaseResponsesAPIAgentConfig( + host="", + port=0, + entrypoint="", + name="agent", + token_id_capture=token_id_capture, + ) return _Agent(config=config, server_client=server_client) @@ -842,6 +848,20 @@ def test_base_agent_url_path_for_run_gates_on_observability_and_indices(): assert _make_base_agent(MagicMock()).url_path_for_run("/v1/responses", body) == "/v1/responses" +def test_base_agent_propagates_explicit_token_capture_intent(): + body = {TASK_INDEX_KEY_NAME: 3, ROLLOUT_INDEX_KEY_NAME: 1} + global_config = { + "observability_enabled": True, + "token_id_capture": {"enabled": True}, + } + opted_in = _make_base_agent(global_config, token_id_capture=True) + assert opted_in.url_path_for_run("/v1/responses", body) == "/ng-rollout/3-1/token-capture/v1/responses" + assert opted_in.base_url_for_run("http://h:1", body) == "http://h:1/ng-rollout/3-1/token-capture" + + opted_out = _make_base_agent(global_config, token_id_capture=False) + assert opted_out.url_path_for_run("/v1/responses", body) == "/ng-rollout/3-1/v1/responses" + + def test_base_agent_url_path_for_request_propagates_inbound_prefix(): agent = _make_base_agent({}) @@ -851,12 +871,21 @@ def test_base_agent_url_path_for_request_propagates_inbound_prefix(): assert agent.url_path_for_request("/v1/responses", SimpleNamespace()) == "/v1/responses" assert agent.url_path_for_request("/v1/responses", None) == "/v1/responses" + capture_prefixed = SimpleNamespace( + path_params={"rollout_id": "7-0"}, + url=SimpleNamespace(path="/ng-rollout/7-0/token-capture/v1/responses"), + ) + assert ( + agent.url_path_for_request("/v1/responses", capture_prefixed) == "/ng-rollout/7-0/token-capture/v1/responses" + ) + def test_base_agent_registers_prefixed_self_call_route(): from nemo_gym.server_utils import ROLLOUT_PATH_PREFIX routes = {route.path for route in _make_base_agent({}).setup_webserver().routes} assert f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses" in routes + assert f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/token-capture/v1/responses" in routes assert "/v1/responses" in routes diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 8165bd8941..3a49ef3a6a 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -55,7 +55,9 @@ TokenCaptureStore, TokenEntry, TokenIdCaptureConfig, + capture_tokens, commit_entry, + current_capture_context, extract_token_fields, install_token_sink, reset_token_sink, @@ -102,6 +104,32 @@ def test_extract_token_fields_absent_returns_none(): assert extract_token_fields({}) is None +def test_extract_token_fields_rejects_partial_metadata(): + with pytest.raises(ValueError, match="prompt_token_ids"): + extract_token_fields( + { + "output": [ + { + "type": "message", + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + } + ] + } + ) + + +def test_token_entry_rejects_mismatched_generation_arrays(): + with pytest.raises(ValidationError, match="same length"): + TokenEntry( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=[-0.1], + ) + + # --- store -------------------------------------------------------------------- @@ -123,11 +151,54 @@ def test_token_store_round_trip(tmp_path): assert store.read_entries("missing") == [] +def test_token_store_put_is_idempotent_and_conflicts_fail_closed(tmp_path): + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + asyncio.run(store.put(entry)) + asyncio.run(store.put(entry)) + assert store.read_entries("r0") == [entry] + + with pytest.raises(ValueError, match="reused with a different payload"): + asyncio.run(store.put(entry.model_copy(update={"generation_token_ids": [8, 9]}))) + assert store.is_incomplete("r0") + assert store.read_entries("r0") == [entry] + + +def test_token_store_seal_is_atomic_and_conditional_drop_is_race_safe(tmp_path): + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + asyncio.run(store.put(entry)) + snapshot = asyncio.run(store.seal("r0")) + assert snapshot.entries == (entry,) + assert asyncio.run(store.seal("r0")) == snapshot + + with pytest.raises(RuntimeError, match="already sealed"): + asyncio.run(store.put(entry.model_copy(update={"model_call_id": "late"}))) + asyncio.run(store.mark_incomplete("r0", "late")) + assert not asyncio.run(store.drop("r0", seal_id=snapshot.seal_id, version=snapshot.version)) + updated = asyncio.run(store.seal("r0")) + assert updated.incomplete + assert asyncio.run(store.drop("r0", seal_id=updated.seal_id, version=updated.version)) + assert store.read_entries("r0") == [] + + # --- config ------------------------------------------------------------------- def _block(**kwargs) -> dict: - return {"token_id_capture": {"enabled": True, **kwargs}} + return {"token_id_capture": {"enabled": True, "rebuild_response": False, **kwargs}} def test_config_disabled_needs_no_dir(): @@ -137,8 +208,8 @@ def test_config_disabled_needs_no_dir(): def test_config_enabled_requires_absolute_dir(tmp_path): - with pytest.raises(ValueError): - TokenIdCaptureConfig.model_validate(_block()) + """A directory that is set has to be absolute. A relative one silently resolves against + whatever the server's working directory happens to be.""" with pytest.raises(ValueError): TokenIdCaptureConfig.model_validate(_block(dir="relative/dir")) cfg = TokenIdCaptureConfig.model_validate(_block(dir=str(tmp_path))) @@ -173,6 +244,13 @@ def test_config_rejects_an_unknown_key(): TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": True, "dirr": "/tmp/x"}}) +def test_config_rejects_write_only_custom_capture(): + with pytest.raises(ValueError, match="source is required"): + TokenIdCaptureConfig.model_validate( + {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + ) + + # --- source / readers --------------------------------------------------------- @@ -256,7 +334,7 @@ def _both_enabled(tmp_path) -> dict: 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"}) + resp = client.post("/ng-rollout/task0-roll0/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll0") @@ -271,7 +349,7 @@ def test_responses_call_captures_tokens_joined_to_eval_record(tmp_path): 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"}) + client.post("/ng-rollout/task0-rollC/token-capture/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. @@ -287,7 +365,7 @@ def test_token_arrays_are_stored_once(tmp_path): value a trainer reads: an item's prompt in a chained trajectory is the running sequence. """ client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) - client.post("/ng-rollout/task0-rollDedup/v1/responses", json={"input": "hi"}) + client.post("/ng-rollout/task0-rollDedup/token-capture/v1/responses", json={"input": "hi"}) entry = TokenCaptureStore(tmp_path).read_entries("task0-rollDedup")[0] assert entry.generation_token_ids == GTOKS for item in entry.output_items: @@ -301,7 +379,7 @@ def test_token_arrays_are_stored_once(tmp_path): 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", + "/ng-rollout/task0-roll1/token-capture/v1/messages", json={"model": "claude-x", "max_tokens": 16, "messages": [{"role": "user", "content": "hello"}]}, ) assert resp.status_code == 200 @@ -314,7 +392,8 @@ def test_messages_call_captures_tokens(tmp_path): 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"}]} + "/ng-rollout/task0-roll2/token-capture/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hi"}]}, ) assert resp.status_code == 200 tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll2") @@ -324,13 +403,22 @@ def test_chat_completions_call_captures_tokens(tmp_path): def test_tokens_captured_even_when_eval_capture_disabled(tmp_path): config = {"token_id_capture": {"enabled": True, "dir": str(tmp_path)}} client = TestClient(_server(config).setup_webserver()) - resp = client.post("/ng-rollout/task1-roll0/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task1-roll0/token-capture/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_observability_prefix_does_not_enable_training_token_capture(tmp_path): + """Rollout correlation is neutral; token capture requires explicit path intent.""" + client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) + resp = client.post("/ng-rollout/observed-only/v1/responses", json={"input": "hi"}) + assert resp.status_code == 200 + assert TokenCaptureStore(tmp_path).read_entries("observed-only") == [] + assert len(read_model_call_records(CaptureStore(tmp_path), "observed-only")) == 1 + + 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"}) @@ -367,7 +455,7 @@ def test_streamed_messages_capture_tokens_absent_from_the_stream(tmp_path): client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) with client.stream( "POST", - "/ng-rollout/stream0-roll0/v1/messages", + "/ng-rollout/stream0-roll0/token-capture/v1/messages", json={ "model": "claude-x", "max_tokens": 16, @@ -402,7 +490,7 @@ async def boom(self, entry): 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"}) + resp = client.post("/ng-rollout/fail0-roll0/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert store.read_entries("fail0-roll0") == [] assert store.is_incomplete("fail0-roll0") @@ -435,7 +523,7 @@ def test_a_response_without_token_ids_marks_the_rollout_incomplete(tmp_path): as if the environment had written them. """ client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) - resp = client.post("/ng-rollout/silent0-roll0/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/silent0-roll0/token-capture/v1/responses", json={"input": "hi"}) # The model call itself still succeeds; capture never breaks the harness's run. assert resp.status_code == 200 @@ -444,6 +532,81 @@ def test_a_response_without_token_ids_marks_the_rollout_incomplete(tmp_path): assert store.is_incomplete("silent0-roll0") +def _external_mode(tmp_path) -> dict: + """Capture on, no destination in this process: records are staged elsewhere.""" + return { + "observability_enabled": False, + "token_id_capture": {"enabled": True, "rebuild_response": False}, + } + + +def test_external_mode_still_mints_identity_for_a_correlated_call(tmp_path): + """A framework staging from the inference worker keys its record on the identity minted here. + + Nothing in this process writes, so the machinery cannot be gated on having a destination. + """ + seen = {} + + class _Peek(_CapturingModel): + async def responses(self, request: Request, body=Body()) -> NeMoGymResponse: + context = current_capture_context() + seen["rollout_id"] = context.rollout_id if context else None + seen["model_call_id"] = context.model_call_id if context else None + seen["sink"] = context.sink if context else "no context" + return _training_response("hi") + + model = _Peek( + config=BaseResponsesAPIModelConfig(host="0.0.0.0", port=8099, entrypoint="", name="srv"), + server_client=MagicMock(spec=ServerClient, global_config_dict=_external_mode(tmp_path)), + ) + assert ( + TestClient(model.setup_webserver()) + .post("/ng-rollout/ext0-r0/token-capture/v1/responses", json={"input": "hi"}) + .status_code + == 200 + ) + + assert seen["rollout_id"] == "ext0-r0" + assert seen["model_call_id"], "a call id has to be minted for the staged record to key on" + assert seen["sink"] is None, "no destination in this process" + + +def test_external_mode_does_not_mark_a_token_less_response_incomplete(tmp_path): + """Under external staging a response with no token ids is ordinary, not a hole. + + This process writes nothing, so it cannot tell a lost call from normal operation, and + marking here would mask every rollout of the run. + """ + client = TestClient(_silent_server(_external_mode(tmp_path)).setup_webserver()) + assert client.post("/ng-rollout/ext1-r0/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 + assert list(tmp_path.glob("**/*.incomplete")) == [] + + +def test_a_committed_call_is_not_marked_even_without_token_ids(tmp_path): + """A caller that had the arrays when this process did not has already accounted for the call.""" + store = TokenCaptureStore(tmp_path) + context = CaptureContext(rollout_id="cm0-r0", model_call_id="c1", sink=store) + token = set_token_sink(context) + try: + asyncio.run( + commit_entry( + TokenEntry( + rollout_id="cm0-r0", + model_call_id="c1", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + ) + ) + assert context.committed is True + asyncio.run(capture_tokens({"output": [{"type": "message"}]})) + finally: + reset_token_sink(token) + + assert not store.is_incomplete("cm0-r0"), "the call was recorded, so it is not a hole" + + def test_untagged_traffic_without_token_ids_marks_nothing(tmp_path): """No rollout prefix means no sink, so there is no rollout to call incomplete.""" client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) @@ -462,7 +625,7 @@ def test_delete_removes_records_and_marker(tmp_path): generation_log_probs=LPS, ) ) - store.mark_incomplete("gone-0", "c") + asyncio.run(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() @@ -517,9 +680,12 @@ def __init__(self) -> None: async def put(self, entry: TokenEntry) -> None: self.entries.append(entry) - def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: self.incomplete.append((rollout_id, model_call_id)) + async def close(self) -> None: + pass + @pytest.fixture def installed_sink(): @@ -533,9 +699,9 @@ def installed_sink(): def test_installed_sink_receives_entries_without_a_capture_dir(installed_sink): """The framework path: capture on, no directory anywhere, records still arrive.""" - config = {"token_id_capture": {"enabled": True}} + config = {"token_id_capture": {"enabled": True, "rebuild_response": False}} client = TestClient(_server(config).setup_webserver()) - resp = client.post("/ng-rollout/task0-sink0/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-sink0/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert len(installed_sink.entries) == 1 assert installed_sink.entries[0].generation_token_ids == GTOKS @@ -547,9 +713,16 @@ def test_config_allows_no_directory_when_a_sink_is_installed(installed_sink): assert TokenIdCaptureConfig.model_validate(_block()).resolved_dir() is None -def test_config_still_requires_a_directory_with_no_sink_installed(): - with pytest.raises(ValueError): - TokenIdCaptureConfig.model_validate(_block()) +def test_config_allows_capture_with_no_destination_at_all(): + """A framework that stages records from the inference worker writes nothing here. + + It still needs capture on, because the identity a record is keyed by and the parent a + request continues are resolved in this process and nowhere else. + """ + cfg = TokenIdCaptureConfig.model_validate(_block()) + assert cfg.enabled is True + assert cfg.resolved_dir() is None + assert cfg.build_sink() is None def test_installed_sink_is_marked_incomplete_through_the_protocol(installed_sink, monkeypatch): @@ -563,8 +736,8 @@ async def boom(entry): raise RuntimeError("transport down") monkeypatch.setattr(installed_sink, "put", boom) - client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) - resp = client.post("/ng-rollout/task0-sink1/v1/responses", json={"input": "hi"}) + client = TestClient(_server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver()) + resp = client.post("/ng-rollout/task0-sink1/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 # capture never fails the model call assert installed_sink.incomplete == [("task0-sink1", installed_sink.incomplete[0][1])] @@ -578,9 +751,11 @@ async def put(self, entry): install_token_sink(_PutOnlySink()) try: - client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) + client = TestClient( + _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() + ) with caplog.at_level(logging.ERROR): - resp = client.post("/ng-rollout/task0-sink2/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-sink2/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert any("does not implement mark_incomplete" in r.message for r in caplog.records) finally: @@ -636,8 +811,10 @@ def _bad_entry(**kwargs): raise ValueError("prompt_token_ids: not a list of ints") with patch("nemo_gym.token_id_capture.sink.TokenEntry", _bad_entry): - client = TestClient(_server({"token_id_capture": {"enabled": True}}).setup_webserver()) - resp = client.post("/ng-rollout/task0-bad0/v1/responses", json={"input": "hi"}) + client = TestClient( + _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() + ) + resp = client.post("/ng-rollout/task0-bad0/token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200, "a malformed token payload must not fail the model call" assert installed_sink.entries == [], "nothing should have been written" @@ -684,7 +861,7 @@ def test_a_rollout_that_lost_a_call_is_distinguishable_from_a_complete_one(tmp_p ) asyncio.run(store.put(entry)) assert not store.is_incomplete("r0") - store.mark_incomplete("r0", "c2") + asyncio.run(store.mark_incomplete("r0", "c2")) assert store.is_incomplete("r0") @@ -699,7 +876,10 @@ class _ConfiguredSink: async def put(self, entry) -> None: type(self).entries.append(entry) - def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + async def close(self) -> None: pass @@ -707,11 +887,17 @@ class _NotASink: async def put(self, entry) -> None: pass + async def close(self) -> None: + pass + class _NotCallableSink: put = "not a method" - def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + async def close(self) -> None: pass @@ -722,16 +908,25 @@ def __init__(self, endpoint: str, shard: int = 0) -> None: async def put(self, entry) -> None: pass - def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + async def close(self) -> None: pass def test_a_configured_sink_receives_entries(tmp_path): _ConfiguredSink.entries = [] - config = {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + config = { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "sink": f"{__name__}:_ConfiguredSink", + } + } client = TestClient(_server(config).setup_webserver()) - assert client.post("/ng-rollout/task0-cfg0/v1/responses", json={"input": "hi"}).status_code == 200 + assert client.post("/ng-rollout/task0-cfg0/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 assert [e.rollout_id for e in _ConfiguredSink.entries] == ["task0-cfg0"] assert _ConfiguredSink.entries[0].generation_token_ids == GTOKS @@ -740,10 +935,16 @@ def test_a_configured_sink_receives_entries(tmp_path): def test_a_configured_sink_wins_over_an_installed_one(installed_sink): """Both routes exist; the configured one is preferred because it survives extra workers.""" _ConfiguredSink.entries = [] - config = {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + config = { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "sink": f"{__name__}:_ConfiguredSink", + } + } client = TestClient(_server(config).setup_webserver()) - assert client.post("/ng-rollout/task0-cfg1/v1/responses", json={"input": "hi"}).status_code == 200 + assert client.post("/ng-rollout/task0-cfg1/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 assert len(_ConfiguredSink.entries) == 1 assert installed_sink.entries == [] @@ -769,7 +970,13 @@ def test_a_sink_that_cannot_report_failures_is_refused_at_startup(): """Without mark_incomplete an incomplete rollout looks complete, so this fails at startup rather than at whichever step first loses a call.""" config = TokenIdCaptureConfig.model_validate( - {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_NotASink"}} + { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "sink": f"{__name__}:_NotASink", + } + } ) with pytest.raises(ValueError, match="mark_incomplete"): config.build_sink() @@ -780,7 +987,13 @@ def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): checked too. Both are derived from the protocol rather than a list written out here, so the check keeps up if TokenSink gains a method.""" config = TokenIdCaptureConfig.model_validate( - {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_NotCallableSink"}} + { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "sink": f"{__name__}:_NotCallableSink", + } + } ) with pytest.raises(ValueError, match="put"): config.build_sink() @@ -791,7 +1004,9 @@ def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): [("no_colon", "module.path:ClassName"), ("nemo_gym.token_id_capture:Nope", "could not load")], ) def test_a_malformed_sink_path_is_refused_at_startup(target, expected): - config = TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": True, "sink": target}}) + config = TokenIdCaptureConfig.model_validate( + {"token_id_capture": {"enabled": True, "rebuild_response": False, "sink": target}} + ) with pytest.raises(ValueError, match=expected): config.build_sink() @@ -843,7 +1058,7 @@ def test_the_store_is_a_token_source(tmp_path): # A colocated source can tell that a call failed to capture, which is what keeps an # incomplete rollout from being trained on. assert store.is_incomplete("r0") is False - store.mark_incomplete("r0", "c2") + asyncio.run(store.mark_incomplete("r0", "c2")) assert store.is_incomplete("r0") is True From cc82a329ee91b0cb67f196c10a56b05e8088ed4f Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Mon, 17 Aug 2026 17:12:41 -0700 Subject: [PATCH 04/12] fix(token-id-capture): close sinks through app lifespan Use the current FastAPI lifespan API so configured capture transports are closed without breaking server startup. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 8f5afe8c1c..d44d222821 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -35,6 +35,7 @@ import re import time from abc import abstractmethod +from contextlib import asynccontextmanager from pathlib import Path from typing import Any, Mapping, Optional from uuid import uuid4 @@ -1320,7 +1321,17 @@ async def _close_token_sinks() -> None: await sink.close() if owned_sinks: - app.add_event_handler("shutdown", _close_token_sinks) + original_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def _capture_lifespan(application): + try: + async with original_lifespan(application) as state: + yield state + finally: + await _close_token_sinks() + + app.router.lifespan_context = _capture_lifespan app.add_middleware( _CaptureMiddleware, store=make_capture_store(config), From b2736863614ff2eb42651d42a9bcd0af8a4f3326 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Mon, 17 Aug 2026 23:11:19 -0700 Subject: [PATCH 05/12] fix(token-id-capture): harden training capture intent Make agent selection explicit without changing evaluation defaults, and avoid rescanning growing token payloads on every durable write. Rename snapshot and URL contracts so their lifecycle and training purpose are clear. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_agent.py | 21 +-- nemo_gym/config_types.py | 2 +- nemo_gym/global_config.py | 5 +- nemo_gym/token_id_capture/__init__.py | 11 +- nemo_gym/token_id_capture/config.py | 9 +- nemo_gym/token_id_capture/protocols.py | 46 ++++--- nemo_gym/token_id_capture/records.py | 28 ++-- nemo_gym/token_id_capture/store.py | 121 ++++++++++++++---- ...ng_gym_claude_code_agent_model_server.yaml | 4 - responses_api_agents/claude_code_agent/app.py | 13 +- .../claude_code_agent/tests/test_app.py | 21 ++- .../test_base_responses_api_model.py | 24 +++- tests/unit_tests/test_token_id_capture.py | 106 +++++++++++---- 13 files changed, 286 insertions(+), 125 deletions(-) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 5b8e263a09..69b9995de3 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -50,12 +50,11 @@ class BaseResponsesAPIAgentConfig(BaseRunServerInstanceConfig): skip_verification: bool = False skip_verification_reward: float = 0.0 - # 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. + # Whether this agent's rollouts participate in training token capture. + # Native agents already receive token ids inline and normally leave this disabled. + # Opaque external harnesses enable it because their returned output has no token ids. + # The run-level ``token_id_capture.enabled`` setting gates the capture infrastructure. + # The run-level ``token_id_capture.all_agents`` setting overrides this agent-level choice. token_id_capture: bool = False @@ -72,9 +71,9 @@ def setup_webserver(self) -> FastAPI: self.setup_session_middleware(app) app.post("/v1/responses")(self.responses) - # Prefixed twin of /v1/responses: a self-call made with url_path_for_run() lands here, and - # responses() recovers the rollout id from the path (see url_path_for_request) to correlate - # its model calls. Same handler, so unprefixed calls are unaffected. + # A self-call made with ``url_path_for_run`` lands on a prefixed twin. + # ``responses`` recovers the rollout id from the path. + # The same handler serves prefixed and unprefixed calls. app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses")(self.responses) app.post(f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/{TOKEN_CAPTURE_PATH_SEGMENT}/v1/responses")(self.responses) @@ -120,7 +119,9 @@ def _token_id_capture_enabled(self) -> bool: if not isinstance(global_config, Mapping): return False block = global_config.get(TOKEN_ID_CAPTURE_BLOCK) or {} - return bool(isinstance(block, Mapping) and block.get("enabled", False)) and bool( + if not isinstance(block, Mapping) or not block.get("enabled", False): + return False + return bool(block.get("all_agents", False)) or bool( getattr(getattr(self, "config", None), "token_id_capture", False) ) diff --git a/nemo_gym/config_types.py b/nemo_gym/config_types.py index 5deac5d144..9e98a060d8 100644 --- a/nemo_gym/config_types.py +++ b/nemo_gym/config_types.py @@ -867,4 +867,4 @@ class AggregateMetrics(BaseModel): # Per-rollout model-call correlation. Callers place the rollout id in the model-server URL; # the capture middleware in base_responses_api_model.py strips this prefix before routing. ROLLOUT_PATH_PREFIX = "ng-rollout" -TOKEN_CAPTURE_PATH_SEGMENT = "token-capture" +TOKEN_CAPTURE_PATH_SEGMENT = "training-token-capture" diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 2447bf370e..8ddf5f8cd4 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -95,10 +95,9 @@ QUERY_KEY_NAME = "query" OBSERVABILITY_ENABLED_KEY_NAME = "observability_enabled" MODEL_CALL_CAPTURE_DIR_KEY_NAME = "model_call_capture_dir" -# Run-wide training-token capture settings; see nemo_gym/token_id_capture/config.py. +# Run-wide training-token capture settings. +# See ``nemo_gym/token_id_capture/config.py``. TOKEN_ID_CAPTURE_BLOCK = "token_id_capture" -# 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" SKIP_VERIFICATION_KEY_NAME = "skip_verification" SKIP_VERIFICATION_REWARD_KEY_NAME = "skip_verification_reward" diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index 558c01a15d..51ebd73dcd 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -25,12 +25,11 @@ imports the record, the protocols, and the capture core to write into its own data plane (see ``protocols.py``). -Records are read back through a ``TokenSource``. ``TokenCaptureStore`` is one, -and is what a reader sitting alongside the store uses. A framework staging -records through its own transport supplies its own source, which lives wherever -that transport does. What any source owes is an honest ``is_incomplete``: it is -how a consumer learns a rollout lost a call, and one that always answers False -trains on an incomplete rollout without knowing. +Records are read back through a ``TokenSource``. +``TokenCaptureStore`` is Gym's local implementation. +A framework staging records through its own transport supplies its own source. +The source freezes an atomic snapshot containing both entries and incomplete state. +This prevents a consumer from training on a rollout that lost a model call. """ from nemo_gym.token_id_capture.config import TokenIdCaptureConfig diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 3df8f7fe7a..a2468e32c0 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -80,13 +80,16 @@ class TokenIdCaptureSettings(BaseModel): model_config = ConfigDict(extra="forbid") enabled: bool = False + # Capture model calls from every agent. + # The default keeps capture scoped by each agent's ``token_id_capture`` setting. + all_agents: bool = False # Where the default file store writes. Falls back to ``model_call_capture_dir``. dir: Path | None = None # ``module.path:ClassName`` implementing TokenSink, constructed per server process. sink: str | None = None - # Keyword arguments for that constructor: an endpoint, a client, credentials. A sink for a real - # transport needs wiring, and a zero-argument one could only get it from ambient state. Use - # ``${oc.env:VAR}`` for anything secret rather than writing it here. + # Keyword arguments for that constructor. + # A real transport needs explicit endpoint, client, or credential wiring. + # Use ``${oc.env:VAR}`` for secrets instead of writing them here. sink_kwargs: dict[str, Any] = Field(default_factory=dict) # Optional paired reader for framework-owned transports. source: str | None = None diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index aed58d495d..bee1e708e1 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -42,38 +42,42 @@ @dataclass(frozen=True) class TokenCaptureSnapshot: - """An immutable, sealed view of one rollout's capture records.""" + """An immutable view of one rollout's frozen capture records.""" rollout_id: str entries: tuple[TokenEntry, ...] incomplete: bool - seal_id: str + snapshot_id: str version: int @runtime_checkable class TokenSink(Protocol): - """Where captured records go. Implemented by Gym's file store, or by a - framework over its own transport.""" + """Where captured records go. + + Gym's file store and framework-owned transports implement this protocol. + """ async def put(self, entry: TokenEntry) -> None: """Durably store one record. - Repeating the same call id with the same payload is a no-op. Reusing a - call id with a different payload or writing after seal must fail. + Repeating the same call id with the same payload is a no-op. + Reusing a call id with a different payload must fail. + Writing after the rollout is frozen must fail. - May raise. The caller marks the rollout incomplete and never fails the - model call because of a capture error. + This method may raise. + The caller marks the rollout incomplete. + A capture error never fails the model call. """ ... async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: """Durably record that a call of this rollout failed to capture. - The rollout is now missing a turn, and a consumer must mask the sample rather - than train on a chain with a hole in it. The model call itself still succeeds, - so this is the only signal that anything went wrong: a sink that drops it makes - an incomplete rollout indistinguishable from a complete one. + The rollout is now missing a turn. + A consumer must mask the sample instead of training on a chain with a hole. + The model call itself still succeeds. + This marker is therefore the durable signal that capture failed. """ ... @@ -84,21 +88,23 @@ async def close(self) -> None: @runtime_checkable class TokenSource(Protocol): - """Where a trajectory builder seals, reads, and retires records.""" + """Where a trajectory builder freezes, reads, and retires records.""" - async def seal(self, rollout_id: str) -> TokenCaptureSnapshot: - """Seal a rollout and return one atomic snapshot. + async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: + """Freeze a rollout and return one atomic snapshot. - Sealing is idempotent. No successful writes may occur after it returns. + Freezing is idempotent. + No successful writes may occur after it returns. Entry order carries no meaning. """ ... - async def drop(self, rollout_id: str, *, seal_id: str, version: int) -> bool: - """Conditionally retire the exact sealed snapshot that was consumed. + async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: + """Conditionally retire the exact frozen snapshot that was consumed. - Returns ``False`` if state changed after the snapshot. Implementations - that cannot delete return ``True`` and leave retention to their owner. + Return ``False`` if state changed after the snapshot. + Implementations that cannot delete return ``True``. + Their owner remains responsible for retention. """ ... diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 15826f3116..5d8281ae1f 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -142,28 +142,28 @@ def response_to_output_items(payload: dict) -> list[dict]: def strip_token_fields(items: list[dict]) -> tuple[list[dict], int | None]: """Drop the token arrays from output items, keeping the content. - Returns the stripped items and the index of the item the arrays came off, which - is the last one carrying them, matching what ``extract_token_fields`` reads. The - arrays are held once on the entry instead: storing them again per item roughly - doubles a record, and the per-item values are not the ones a trainer sees, since - the builder replaces an item's prompt with the chain's running sequence. + Return the stripped items and the index of their token-bearing item. + Capture requires exactly one token-bearing item. + The arrays are held once on the entry. + Storing them again per item would roughly double the record size. """ - index: int | None = None + indices: list[int] = [] stripped: list[dict] = [] for position, item in enumerate(items): if item.get("generation_token_ids") is not None: - index = position + indices.append(position) stripped.append({key: value for key, value in item.items() if key not in TOKEN_FIELDS}) - return stripped, index + if len(indices) > 1: + raise ValueError("multiple output items carry token metadata") + return stripped, indices[0] if indices else None def extract_token_fields(response_json: dict) -> dict | None: """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). + Handle Responses output items and Chat Completions messages. + Exactly one item may carry token metadata. + Return ``None`` when no item carries token ids. """ candidates: list[dict] = [] required = ("prompt_token_ids", "generation_token_ids", "generation_log_probs") @@ -176,7 +176,9 @@ def extract_token_fields(response_json: dict) -> dict | None: candidates.append(message) if not candidates: return None - source = candidates[-1] + if len(candidates) > 1: + raise ValueError("multiple response items carry token metadata") + source = candidates[0] missing = [field for field in required if source.get(field) is None] if missing: raise ValueError(f"partial token metadata is missing: {', '.join(missing)}") diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index 2998793faf..b637b2e156 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -33,6 +33,7 @@ import asyncio import fcntl +import hashlib import os import tempfile from contextlib import contextmanager @@ -89,12 +90,71 @@ def _locked(self, rollout_id: str, *, shared: bool = False): def _read_state(self, rollout_id: str) -> dict[str, Any]: path = self.state_path_for(rollout_id) if not path.exists(): - return {"sealed": False, "incomplete": False, "seal_id": "", "version": 0} + return { + "frozen": False, + "incomplete": False, + "snapshot_id": "", + "version": 0, + "entry_digests": {}, + "indexed_size": 0, + } state = orjson.loads(path.read_bytes()) if not isinstance(state, dict): raise ValueError(f"Invalid token-capture state for rollout {rollout_id}") return state + @staticmethod + def _entry_digest(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + def _sync_entry_index(self, rollout_id: str, state: dict[str, Any]) -> bool: + """Reconcile an entry index with any durable JSONL tail. + + The JSONL write is durable before its state update. + A process can therefore stop with one unindexed entry. + Normal writes use the state index without parsing prior token arrays. + Recovery parses only the unindexed tail. + """ + path = self.path_for(rollout_id) + file_size = path.stat().st_size if path.exists() else 0 + stored_index = state.get("entry_digests") + stored_size = state.get("indexed_size") + legacy_state = not isinstance(stored_index, dict) or not isinstance(stored_size, int) + entry_digests = dict(stored_index) if isinstance(stored_index, dict) else {} + indexed_size = stored_size if isinstance(stored_size, int) else 0 + if indexed_size < 0 or indexed_size > file_size: + raise ValueError(f"Invalid token-capture index offset for rollout {rollout_id}") + if indexed_size == file_size and not legacy_state: + return False + + recovered = 0 + if path.exists(): + with path.open("rb") as handle: + handle.seek(indexed_size) + for line in handle: + payload = line.strip() + if not payload: + continue + entry = TokenEntry.model_validate(orjson.loads(payload)) + digest = self._entry_digest(payload) + existing = entry_digests.get(entry.model_call_id) + if existing is not None and existing != digest: + state["incomplete"] = True + state["version"] = int(state.get("version", 0)) + 1 + self._write_state(rollout_id, state) + raise ValueError( + f"Model call id {entry.model_call_id!r} has conflicting durable payloads " + f"for rollout {rollout_id!r}" + ) + entry_digests[entry.model_call_id] = digest + recovered += 1 + + state["entry_digests"] = entry_digests + state["indexed_size"] = file_size + if recovered and not legacy_state: + state["version"] = int(state.get("version", 0)) + recovered + return True + def _write_state(self, rollout_id: str, state: dict[str, Any]) -> None: payload = orjson.dumps(state, option=orjson.OPT_SORT_KEYS | orjson.OPT_APPEND_NEWLINE) with tempfile.NamedTemporaryFile(dir=self._root, prefix=".tokens-state-", delete=False) as handle: @@ -143,16 +203,19 @@ def append(self, entry: TokenEntry) -> None: """Idempotently append one entry and fsync.""" canonical = orjson.dumps(entry.model_dump(mode="json"), option=orjson.OPT_SORT_KEYS) line = canonical + b"\n" + digest = self._entry_digest(canonical) rollout_id = entry.rollout_id with self._locked(rollout_id): state = self._read_state(rollout_id) - if state.get("sealed", False): - raise RuntimeError(f"Token capture for rollout {rollout_id} is already sealed") - for existing in self._read_entries_unlocked(rollout_id): - if existing.model_call_id != entry.model_call_id: - continue - existing_bytes = orjson.dumps(existing.model_dump(mode="json"), option=orjson.OPT_SORT_KEYS) - if existing_bytes == canonical: + if state.get("frozen", False): + raise RuntimeError(f"Token capture for rollout {rollout_id} is already frozen") + index_changed = self._sync_entry_index(rollout_id, state) + entry_digests = state["entry_digests"] + existing_digest = entry_digests.get(entry.model_call_id) + if existing_digest is not None: + if existing_digest == digest: + if index_changed: + self._write_state(rollout_id, state) return state["incomplete"] = True state["version"] = int(state.get("version", 0)) + 1 @@ -165,15 +228,15 @@ def append(self, entry: TokenEntry) -> None: handle.write(line) handle.flush() os.fsync(handle.fileno()) + state["indexed_size"] = handle.tell() + entry_digests[entry.model_call_id] = digest state["version"] = int(state.get("version", 0)) + 1 self._write_state(rollout_id, state) - # --- TokenSink / TokenSource. The file store is Gym's default implementation of both; - # a framework swaps in its own without touching the capture path. + # The file store is Gym's default TokenSink and TokenSource. + # A framework can replace it without changing the capture path. # - # Both offload to the default thread pool, which is shared process-wide and small - # (min(32, cpus + 4)). Serializing the entry dominates the cost rather than the write - # itself, so a long context is the case to watch if this ever shows up in a profile. + # Both interfaces offload blocking work to the process-wide default thread pool. async def put(self, entry: TokenEntry) -> None: """``TokenSink``: durable on return. The blocking append is offloaded so @@ -181,40 +244,44 @@ async def put(self, entry: TokenEntry) -> None: rollout never races a partial file.""" await asyncio.to_thread(self.append, entry) - async def seal(self, rollout_id: str) -> TokenCaptureSnapshot: - return await asyncio.to_thread(self._seal, rollout_id) + async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: + return await asyncio.to_thread(self.freeze_now, rollout_id) - def _seal(self, rollout_id: str) -> TokenCaptureSnapshot: + def freeze_now(self, rollout_id: str) -> TokenCaptureSnapshot: + """Synchronously freeze one rollout and return its stable snapshot.""" with self._locked(rollout_id): state = self._read_state(rollout_id) - if not state.get("sealed", False): - state["sealed"] = True - state["seal_id"] = uuid4().hex + index_changed = self._sync_entry_index(rollout_id, state) + if not state.get("frozen", False): + state["frozen"] = True + state["snapshot_id"] = uuid4().hex state["version"] = int(state.get("version", 0)) + 1 self._write_state(rollout_id, state) + elif index_changed: + self._write_state(rollout_id, state) entries = tuple(self._read_entries_unlocked(rollout_id)) return TokenCaptureSnapshot( rollout_id=rollout_id, entries=entries, incomplete=bool(state.get("incomplete", False)), - seal_id=str(state["seal_id"]), + snapshot_id=str(state["snapshot_id"]), version=int(state["version"]), ) async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: - """Compatibility read for diagnostics. Consumers should use ``seal``.""" + """Compatibility read for diagnostics. Consumers should use ``freeze``.""" return await asyncio.to_thread(self.read_entries, rollout_id) - async def drop(self, rollout_id: str, *, seal_id: str, version: int) -> bool: - """Conditionally delete the sealed snapshot.""" - return await asyncio.to_thread(self._drop, rollout_id, seal_id, version) + async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: + """Conditionally delete the frozen snapshot.""" + return await asyncio.to_thread(self._drop, rollout_id, snapshot_id, version) - def _drop(self, rollout_id: str, seal_id: str, version: int) -> bool: + def _drop(self, rollout_id: str, snapshot_id: str, version: int) -> bool: with self._locked(rollout_id): state = self._read_state(rollout_id) if ( - not state.get("sealed", False) - or state.get("seal_id") != seal_id + not state.get("frozen", False) + or state.get("snapshot_id") != snapshot_id or int(state.get("version", 0)) != version ): return False 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 adcf80a523..10a60431aa 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,10 +32,6 @@ reasoning_gym_claude_code_agent_model_server: model_server: type: responses_api_models name: policy_model - # An external harness: its returned output has no token ids, so its model calls are - # captured and rebuilt when the run-level token_id_capture.enabled switch is on. Inert - # otherwise, so this is safe for the evaluation showcase above. - token_id_capture: true concurrency: 32 model: ${policy_model_name} anthropic_api_key: EMPTY # pragma: allowlist secret diff --git a/responses_api_agents/claude_code_agent/app.py b/responses_api_agents/claude_code_agent/app.py index dd85306962..915a9dbad7 100644 --- a/responses_api_agents/claude_code_agent/app.py +++ b/responses_api_agents/claude_code_agent/app.py @@ -318,13 +318,18 @@ def _resolve_base_url(self) -> str: return self.config.anthropic_base_url or "" def _resolve_call_base_url(self, rollout_id: Optional[str]) -> str: - """Base URL for the CLI's model calls, with the per-rollout capture prefix applied only when a - Gym model server is configured. A real Anthropic endpoint (``model_server`` unset) has no - prefix-stripping middleware, so prefixing it would 404 every call. + """Return the CLI model-call URL with its rollout prefix. + + Apply the prefix only for a configured Gym model server. + A real Anthropic endpoint has no prefix-stripping middleware. """ base_url = self._resolve_base_url() if base_url and self.config.model_server: - base_url = apply_rollout_prefix(base_url, rollout_id) + base_url = apply_rollout_prefix( + base_url, + rollout_id, + token_capture=self._token_id_capture_enabled(), + ) return base_url def _build_settings(self) -> dict[str, Any]: diff --git a/responses_api_agents/claude_code_agent/tests/test_app.py b/responses_api_agents/claude_code_agent/tests/test_app.py index 1625d753ee..472163cb79 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -67,8 +67,8 @@ def _config(**kwargs) -> ClaudeCodeAgentConfig: def _make_agent(**kwargs) -> ClaudeCodeAgent: - # Patch only the external side effect (claude-code install/version check) so the real - # model_post_init still runs — it initializes the model's private attrs and the semaphore. + # Patch only the Claude Code installation check. + # The real model initialization still configures private attributes and the semaphore. with patch("responses_api_agents.claude_code_agent.app.ensure_claude_code"): return ClaudeCodeAgent(config=_config(**kwargs), server_client=MagicMock(spec=ServerClient)) @@ -837,17 +837,28 @@ async def fake_exec(*cmd, **kwargs): def test_base_url_correlation(self, tmp_path: Path) -> None: agent = _make_agent(model_server=ModelServerRef(type="responses_api_models", name="policy_model")) base_url = self._run_and_capture_base_url(agent, tmp_path, rollout_id="task3-roll1") - # CLI appends /v1/messages -> server strips /ng-rollout/ and keys capture by it. + # The CLI appends ``/v1/messages``. + # The server strips the rollout prefix and uses its id for correlation. assert base_url == "http://model-server:9000/ng-rollout/task3-roll1" with patch.object(agent, "_resolve_base_url", return_value="http://model-server:9000"): assert agent._resolve_call_base_url(None) == "http://model-server:9000" - # Real Anthropic endpoint (no model server): never prefixed -- it has no stripping middleware, - # so a prefix would 404 every /v1/messages call. + # A real Anthropic endpoint has no prefix-stripping middleware. anthropic = _make_agent(anthropic_base_url="https://api.anthropic.com") assert anthropic._resolve_call_base_url("t3-r1") == "https://api.anthropic.com" + def test_training_capture_intent_reaches_the_cli_base_url(self, tmp_path: Path) -> None: + agent = _make_agent( + model_server=ModelServerRef(type="responses_api_models", name="policy_model"), + token_id_capture=True, + ) + agent.server_client.global_config_dict = {"token_id_capture": {"enabled": True}} + + base_url = self._run_and_capture_base_url(agent, tmp_path, rollout_id="task3-roll1") + + assert base_url == "http://model-server:9000/ng-rollout/task3-roll1/training-token-capture" + class TestExtractInstruction: def test_user_only(self) -> None: diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 8d47361004..053db7f1e1 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -855,13 +855,26 @@ def test_base_agent_propagates_explicit_token_capture_intent(): "token_id_capture": {"enabled": True}, } opted_in = _make_base_agent(global_config, token_id_capture=True) - assert opted_in.url_path_for_run("/v1/responses", body) == "/ng-rollout/3-1/token-capture/v1/responses" - assert opted_in.base_url_for_run("http://h:1", body) == "http://h:1/ng-rollout/3-1/token-capture" + assert opted_in.url_path_for_run("/v1/responses", body) == "/ng-rollout/3-1/training-token-capture/v1/responses" + assert opted_in.base_url_for_run("http://h:1", body) == "http://h:1/ng-rollout/3-1/training-token-capture" opted_out = _make_base_agent(global_config, token_id_capture=False) assert opted_out.url_path_for_run("/v1/responses", body) == "/ng-rollout/3-1/v1/responses" +def test_base_agent_all_agents_overrides_the_agent_opt_in(): + body = {TASK_INDEX_KEY_NAME: 3, ROLLOUT_INDEX_KEY_NAME: 1} + global_config = { + "token_id_capture": { + "enabled": True, + "all_agents": True, + } + } + agent = _make_base_agent(global_config, token_id_capture=False) + + assert agent.url_path_for_run("/v1/responses", body) == ("/ng-rollout/3-1/training-token-capture/v1/responses") + + def test_base_agent_url_path_for_request_propagates_inbound_prefix(): agent = _make_base_agent({}) @@ -873,10 +886,11 @@ def test_base_agent_url_path_for_request_propagates_inbound_prefix(): capture_prefixed = SimpleNamespace( path_params={"rollout_id": "7-0"}, - url=SimpleNamespace(path="/ng-rollout/7-0/token-capture/v1/responses"), + url=SimpleNamespace(path="/ng-rollout/7-0/training-token-capture/v1/responses"), ) assert ( - agent.url_path_for_request("/v1/responses", capture_prefixed) == "/ng-rollout/7-0/token-capture/v1/responses" + agent.url_path_for_request("/v1/responses", capture_prefixed) + == "/ng-rollout/7-0/training-token-capture/v1/responses" ) @@ -885,7 +899,7 @@ def test_base_agent_registers_prefixed_self_call_route(): routes = {route.path for route in _make_base_agent({}).setup_webserver().routes} assert f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/v1/responses" in routes - assert f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/token-capture/v1/responses" in routes + assert f"/{ROLLOUT_PATH_PREFIX}/{{rollout_id}}/training-token-capture/v1/responses" in routes assert "/v1/responses" in routes diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 3a49ef3a6a..1c3f246e4a 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -30,6 +30,7 @@ from unittest.mock import MagicMock, patch from uuid import uuid4 +import orjson import pytest from fastapi import Body, Request from fastapi.testclient import TestClient @@ -119,6 +120,16 @@ def test_extract_token_fields_rejects_partial_metadata(): ) +def test_extract_token_fields_rejects_multiple_carriers(): + carrier = { + "prompt_token_ids": PTOKS, + "generation_token_ids": GTOKS, + "generation_log_probs": LPS, + } + with pytest.raises(ValueError, match="multiple response items"): + extract_token_fields({"output": [carrier, carrier]}) + + def test_token_entry_rejects_mismatched_generation_arrays(): with pytest.raises(ValidationError, match="same length"): TokenEntry( @@ -170,7 +181,46 @@ def test_token_store_put_is_idempotent_and_conflicts_fail_closed(tmp_path): assert store.read_entries("r0") == [entry] -def test_token_store_seal_is_atomic_and_conditional_drop_is_race_safe(tmp_path): +def test_token_store_append_uses_the_compact_entry_index(tmp_path, monkeypatch): + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + store.append(entry) + + def fail_if_rescanned(_rollout_id): + raise AssertionError("append rescanned prior token records") + + monkeypatch.setattr(store, "_read_entries_unlocked", fail_if_rescanned) + store.append(entry.model_copy(update={"model_call_id": "c1"})) + store.append(entry) + + +def test_token_store_recovers_an_unindexed_durable_tail(tmp_path): + store = TokenCaptureStore(tmp_path) + entry = TokenEntry( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + payload = orjson.dumps(entry.model_dump(mode="json"), option=orjson.OPT_SORT_KEYS) + b"\n" + store.path_for("r0").write_bytes(payload) + + store.append(entry) + + assert store.read_entries("r0") == [entry] + state = orjson.loads(store.state_path_for("r0").read_bytes()) + assert state["indexed_size"] == len(payload) + assert set(state["entry_digests"]) == {"c0"} + + +def test_token_store_freeze_is_atomic_and_conditional_drop_is_race_safe(tmp_path): store = TokenCaptureStore(tmp_path) entry = TokenEntry( rollout_id="r0", @@ -180,17 +230,17 @@ def test_token_store_seal_is_atomic_and_conditional_drop_is_race_safe(tmp_path): generation_log_probs=LPS, ) asyncio.run(store.put(entry)) - snapshot = asyncio.run(store.seal("r0")) + snapshot = asyncio.run(store.freeze("r0")) assert snapshot.entries == (entry,) - assert asyncio.run(store.seal("r0")) == snapshot + assert asyncio.run(store.freeze("r0")) == snapshot - with pytest.raises(RuntimeError, match="already sealed"): + with pytest.raises(RuntimeError, match="already frozen"): asyncio.run(store.put(entry.model_copy(update={"model_call_id": "late"}))) asyncio.run(store.mark_incomplete("r0", "late")) - assert not asyncio.run(store.drop("r0", seal_id=snapshot.seal_id, version=snapshot.version)) - updated = asyncio.run(store.seal("r0")) + assert not asyncio.run(store.drop("r0", snapshot_id=snapshot.snapshot_id, version=snapshot.version)) + updated = asyncio.run(store.freeze("r0")) assert updated.incomplete - assert asyncio.run(store.drop("r0", seal_id=updated.seal_id, version=updated.version)) + assert asyncio.run(store.drop("r0", snapshot_id=updated.snapshot_id, version=updated.version)) assert store.read_entries("r0") == [] @@ -334,7 +384,7 @@ def _both_enabled(tmp_path) -> dict: 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/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-roll0/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 tokens = TokenCaptureStore(tmp_path).read_entries("task0-roll0") @@ -349,7 +399,7 @@ def test_responses_call_captures_tokens_joined_to_eval_record(tmp_path): def test_captured_entry_carries_content(tmp_path): client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) - client.post("/ng-rollout/task0-rollC/token-capture/v1/responses", json={"input": "hi"}) + client.post("/ng-rollout/task0-rollC/training-token-capture/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. @@ -365,7 +415,7 @@ def test_token_arrays_are_stored_once(tmp_path): value a trainer reads: an item's prompt in a chained trajectory is the running sequence. """ client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) - client.post("/ng-rollout/task0-rollDedup/token-capture/v1/responses", json={"input": "hi"}) + client.post("/ng-rollout/task0-rollDedup/training-token-capture/v1/responses", json={"input": "hi"}) entry = TokenCaptureStore(tmp_path).read_entries("task0-rollDedup")[0] assert entry.generation_token_ids == GTOKS for item in entry.output_items: @@ -379,7 +429,7 @@ def test_token_arrays_are_stored_once(tmp_path): def test_messages_call_captures_tokens(tmp_path): client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) resp = client.post( - "/ng-rollout/task0-roll1/token-capture/v1/messages", + "/ng-rollout/task0-roll1/training-token-capture/v1/messages", json={"model": "claude-x", "max_tokens": 16, "messages": [{"role": "user", "content": "hello"}]}, ) assert resp.status_code == 200 @@ -392,7 +442,7 @@ def test_messages_call_captures_tokens(tmp_path): 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/token-capture/v1/chat/completions", + "/ng-rollout/task0-roll2/training-token-capture/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]}, ) assert resp.status_code == 200 @@ -403,7 +453,7 @@ def test_chat_completions_call_captures_tokens(tmp_path): def test_tokens_captured_even_when_eval_capture_disabled(tmp_path): config = {"token_id_capture": {"enabled": True, "dir": str(tmp_path)}} client = TestClient(_server(config).setup_webserver()) - resp = client.post("/ng-rollout/task1-roll0/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task1-roll0/training-token-capture/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. @@ -455,7 +505,7 @@ def test_streamed_messages_capture_tokens_absent_from_the_stream(tmp_path): client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) with client.stream( "POST", - "/ng-rollout/stream0-roll0/token-capture/v1/messages", + "/ng-rollout/stream0-roll0/training-token-capture/v1/messages", json={ "model": "claude-x", "max_tokens": 16, @@ -490,7 +540,7 @@ async def boom(self, entry): 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/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/fail0-roll0/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert store.read_entries("fail0-roll0") == [] assert store.is_incomplete("fail0-roll0") @@ -523,7 +573,7 @@ def test_a_response_without_token_ids_marks_the_rollout_incomplete(tmp_path): as if the environment had written them. """ client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) - resp = client.post("/ng-rollout/silent0-roll0/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/silent0-roll0/training-token-capture/v1/responses", json={"input": "hi"}) # The model call itself still succeeds; capture never breaks the harness's run. assert resp.status_code == 200 @@ -561,7 +611,7 @@ async def responses(self, request: Request, body=Body()) -> NeMoGymResponse: ) assert ( TestClient(model.setup_webserver()) - .post("/ng-rollout/ext0-r0/token-capture/v1/responses", json={"input": "hi"}) + .post("/ng-rollout/ext0-r0/training-token-capture/v1/responses", json={"input": "hi"}) .status_code == 200 ) @@ -578,7 +628,9 @@ def test_external_mode_does_not_mark_a_token_less_response_incomplete(tmp_path): marking here would mask every rollout of the run. """ client = TestClient(_silent_server(_external_mode(tmp_path)).setup_webserver()) - assert client.post("/ng-rollout/ext1-r0/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 + assert ( + client.post("/ng-rollout/ext1-r0/training-token-capture/v1/responses", json={"input": "hi"}).status_code == 200 + ) assert list(tmp_path.glob("**/*.incomplete")) == [] @@ -701,7 +753,7 @@ def test_installed_sink_receives_entries_without_a_capture_dir(installed_sink): """The framework path: capture on, no directory anywhere, records still arrive.""" config = {"token_id_capture": {"enabled": True, "rebuild_response": False}} client = TestClient(_server(config).setup_webserver()) - resp = client.post("/ng-rollout/task0-sink0/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-sink0/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert len(installed_sink.entries) == 1 assert installed_sink.entries[0].generation_token_ids == GTOKS @@ -737,7 +789,7 @@ async def boom(entry): monkeypatch.setattr(installed_sink, "put", boom) client = TestClient(_server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver()) - resp = client.post("/ng-rollout/task0-sink1/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-sink1/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 # capture never fails the model call assert installed_sink.incomplete == [("task0-sink1", installed_sink.incomplete[0][1])] @@ -755,7 +807,7 @@ async def put(self, entry): _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() ) with caplog.at_level(logging.ERROR): - resp = client.post("/ng-rollout/task0-sink2/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-sink2/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200 assert any("does not implement mark_incomplete" in r.message for r in caplog.records) finally: @@ -814,7 +866,7 @@ def _bad_entry(**kwargs): client = TestClient( _server({"token_id_capture": {"enabled": True, "rebuild_response": False}}).setup_webserver() ) - resp = client.post("/ng-rollout/task0-bad0/token-capture/v1/responses", json={"input": "hi"}) + resp = client.post("/ng-rollout/task0-bad0/training-token-capture/v1/responses", json={"input": "hi"}) assert resp.status_code == 200, "a malformed token payload must not fail the model call" assert installed_sink.entries == [], "nothing should have been written" @@ -926,7 +978,10 @@ def test_a_configured_sink_receives_entries(tmp_path): } client = TestClient(_server(config).setup_webserver()) - assert client.post("/ng-rollout/task0-cfg0/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 + assert ( + client.post("/ng-rollout/task0-cfg0/training-token-capture/v1/responses", json={"input": "hi"}).status_code + == 200 + ) assert [e.rollout_id for e in _ConfiguredSink.entries] == ["task0-cfg0"] assert _ConfiguredSink.entries[0].generation_token_ids == GTOKS @@ -944,7 +999,10 @@ def test_a_configured_sink_wins_over_an_installed_one(installed_sink): } client = TestClient(_server(config).setup_webserver()) - assert client.post("/ng-rollout/task0-cfg1/token-capture/v1/responses", json={"input": "hi"}).status_code == 200 + assert ( + client.post("/ng-rollout/task0-cfg1/training-token-capture/v1/responses", json={"input": "hi"}).status_code + == 200 + ) assert len(_ConfiguredSink.entries) == 1 assert installed_sink.entries == [] From d91191f77672f07d08218ee03621f3e2272071cb Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Mon, 17 Aug 2026 23:47:23 -0700 Subject: [PATCH 06/12] docs(token-id-capture): clarify capture contracts in code Keep comments and docstrings aligned with static agent selection, the training-specific route marker, and frozen source snapshots. Use short standalone sentences throughout the capture path. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_agent.py | 44 ++--- nemo_gym/base_responses_api_model.py | 74 ++++--- nemo_gym/global_config.py | 6 +- nemo_gym/rollout_collection.py | 4 +- nemo_gym/rollout_correlation.py | 34 ++-- nemo_gym/token_id_capture/__init__.py | 27 ++- nemo_gym/token_id_capture/config.py | 70 ++++--- nemo_gym/token_id_capture/protocols.py | 34 ++-- nemo_gym/token_id_capture/records.py | 88 ++++----- nemo_gym/token_id_capture/sink.py | 108 +++++------ nemo_gym/token_id_capture/store.py | 39 ++-- .../configs/claude_code_agent.yaml | 12 +- .../test_base_responses_api_model.py | 16 +- tests/unit_tests/test_rollout_collection.py | 10 +- tests/unit_tests/test_token_id_capture.py | 181 +++++++++--------- 15 files changed, 339 insertions(+), 408 deletions(-) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 69b9995de3..9bebad799f 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -93,16 +93,13 @@ async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyRespo return app def _capture_correlation_enabled(self) -> bool: - """Whether the per-rollout ``/ng-rollout/`` correlation prefix should be applied. + """Return whether this agent needs rollout correlation. - 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. + Evaluation uses ``/ng-rollout//...`` for every agent. + Training capture uses ``/ng-rollout//training-token-capture/...``. + Training capture requires ``token_id_capture.enabled``. + It also requires the static agent flag or run-level ``all_agents``. + Missing global configuration disables correlation. """ return self._model_call_capture_enabled() or self._token_id_capture_enabled() @@ -126,22 +123,21 @@ def _token_id_capture_enabled(self) -> bool: ) def rollout_id_from_run(self, body: Any) -> Optional[str]: - """Per-rollout capture id for a run-request (its task/rollout indices). + """Return the capture id for a run request. - None when neither capture path is enabled or the body carries no indices, so callers apply - no correlation prefix in either case. + Return ``None`` when capture is disabled. + Return ``None`` when the body has no usable identity. """ if not self._capture_correlation_enabled(): return None return maybe_rollout_id_from_run_body(body) def url_path_for_run(self, url_path: str, body: Any) -> str: - """A downstream url_path with the per-rollout capture-correlation prefix applied. + """Apply this run's capture path to a downstream URL path. - Returns ``/ng-rollout/`` when observability is enabled and the run body - carries task/rollout indices; otherwise ``url_path`` unchanged. Use for calls made while - handling ``/run`` — both direct model-server calls and self-calls to ``/v1/responses`` - (the prefixed self-call route carries the id into ``responses()``). + Evaluation uses ``/ng-rollout//...``. + Training capture uses ``/ng-rollout//training-token-capture/...``. + Calls without a rollout id remain unchanged. """ return ( f"{rollout_path_prefix(self.rollout_id_from_run(body), token_capture=self._token_id_capture_enabled())}" @@ -149,11 +145,9 @@ def url_path_for_run(self, url_path: str, body: Any) -> str: ) def base_url_for_run(self, base_url: str, body: Any) -> str: - """A model-server base URL with the per-rollout capture-correlation prefix applied. + """Apply this run's capture path to a model-server root URL. - ``base_url_for_run`` is the base-URL counterpart of ``url_path_for_run`` for SDK-style - harnesses that configure a client once instead of prefixing each call: same gating, applied - to a server root URL (append the API-version suffix afterwards). + Append the API-version suffix after this method returns. """ return apply_rollout_prefix( base_url, @@ -162,11 +156,11 @@ def base_url_for_run(self, base_url: str, body: Any) -> str: ) def url_path_for_request(self, url_path: str, request: Optional[Request]) -> str: - """Carry an inbound ``/ng-rollout/`` self-call prefix onto a downstream url_path. + """Carry an inbound capture path onto a downstream URL path. - Agents whose model calls happen inside ``responses()`` receive the correlation id as the - ``rollout_id`` path parameter of the prefixed self-call route; this re-applies it to the - outgoing model call. Unprefixed requests pass through unchanged. + Prefixed self-calls expose the rollout id as a path parameter. + Training-capture requests preserve their dedicated path segment. + Unprefixed requests remain unchanged. """ path_params = getattr(request, "path_params", None) rollout_id = path_params.get("rollout_id") if isinstance(path_params, Mapping) else None diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index d44d222821..7ca2e994c6 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -76,8 +76,8 @@ 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). +# The store factory needs Gym's server stack. +# The leaf package does not re-export it. from nemo_gym.token_id_capture.config import token_id_capture_config from nemo_gym.token_id_capture.store import make_token_store @@ -246,9 +246,9 @@ async def _invoke_responses( 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. + # Capture before streaming dispatch wraps the response. + # Anthropic mapping drops the token fields. + # The assembled response still carries them here for every dialect. await capture_tokens(response) return response @@ -1070,13 +1070,13 @@ def __init__( 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. + # This store records training tokens for correlated training-capture calls. self._token_store = token_store # Built from token_id_capture.sink, once, in this process. self._configured_sink = configured_sink - # Capture can be on with no destination in this process, when a framework stages records - # from the inference worker instead. The identity and the parent resolution still have to - # happen here, so enablement rather than a destination decides whether this runs. + # Capture may have no destination in this process. + # A framework may stage records from its inference worker. + # This process still resolves the capture identity. self._token_capture_enabled = token_capture_enabled async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: @@ -1096,16 +1096,15 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None dialect = _OBSERVED_PATHS.get(path) - # 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. + # Forward when no active store needs this correlated endpoint. + # The prefix is already stripped. # An unprefixed call is forwarded rather than mixed with unrelated calls under a shared key. - # Destination order: a sink configured for this process, then one installed - # programmatically, then the file store. Both sink routes exist because a framework may - # send records to its own transport instead of disk; the configured one is preferred - # because it is built inside this process at app startup and so survives num_workers > 1, - # where a sink installed by a launcher does not reach the spawned workers at all. The - # installed sink is still resolved per request, so one installed after the app is built - # still takes effect. + # Prefer the configured sink. + # Then use the installed sink. + # Finally use the file store. + # Configured sinks are built in each server process. + # Launcher-installed sinks do not reach spawned workers. + # Installed sinks are resolved for each request. token_sink = self._configured_sink or installed_token_sink() or self._token_store capture_wanted = token_capture_requested and (token_sink is not None or self._token_capture_enabled) if (self._store is None and not capture_wanted) or rollout_from_path is None or dialect is None: @@ -1115,19 +1114,19 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None 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 cannot: token ids are dropped on the SSE wire. - # Set whenever capture is on, even with no destination here. The context carries the - # identity a staged record is keyed by and the parent a request continues, and both are - # resolved in this process whether or not it is the one that writes. + # Give the model server a token sink keyed to this call. + # The sink records token ids from the complete response. + # Middleware cannot recover token ids from SSE. + # The context exists even without a local destination. + # External staging uses the identity resolved here. sink_token = None if capture_wanted: sink_token = set_token_sink( CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, sink=token_sink) ) - # Training-token capture only: no evaluation record, so skip the response buffering entirely - # and just forward with the sink live. + # Training-only capture has no evaluation record. + # Forward without buffering while the sink is active. if self._store is None: try: await self._app(scope, receive, send) @@ -1296,21 +1295,20 @@ def install_model_call_capture( ) -> 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 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. + Always strip ``/ng-rollout//...`` before routing. + Evaluation capture records requests and responses for that path. + Non-terminal SSE chunks continue immediately. + The terminal event follows the durable evaluation write. + Training capture uses ``/ng-rollout//training-token-capture/...``. + That path provides a request-scoped token sink. + The model server records token ids from its complete response. + Consumers access records through ``TokenSource.freeze``. + There is no HTTP token reader. """ token_store = make_token_store(global_config_dict) if global_config_dict is not None else None - # Built here, at app startup, so every uvicorn worker constructs its own. A sink installed by a - # launcher process is not inherited by spawned workers and would silently go unused. + # Build this sink at app startup. + # Each uvicorn worker constructs its own sink. + # Spawned workers do not inherit a launcher-installed sink. configured_sink = ( token_id_capture_config(global_config_dict).build_sink() if global_config_dict is not None else None ) diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 8ddf5f8cd4..066579dbb4 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -143,9 +143,9 @@ # Resume re-dispatch attempt counter (0 on the first attempt); distinguishes retries of the same # (task, rollout) so their captured model calls stay separable. ATTEMPT_INDEX_KEY_NAME = "_ng_attempt_index" -# Explicit capture id for a run request, used in place of the (task, rollout) derivation. Set it -# when the caller reuses task and rollout indices across dispatches, since the derived id would -# then repeat and two dispatches would share one capture key. +# An explicit capture id replaces the task and rollout derivation. +# Set it when dispatches reuse task and rollout indices. +# Otherwise two dispatches would share one capture key. ROLLOUT_ID_KEY_NAME = "_ng_rollout_id" RESPONSES_CREATE_PARAMS_KEY_NAME = "responses_create_params" RESPONSE_KEY_NAME = "response" diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index d26f82dbb6..3399a6ca7a 100644 --- a/nemo_gym/rollout_collection.py +++ b/nemo_gym/rollout_collection.py @@ -818,8 +818,8 @@ async def run_from_config(self, config: RolloutCollectionConfig) -> Tuple[List[D if ATTEMPT_INDEX_KEY_NAME in row: result[ATTEMPT_INDEX_KEY_NAME] = row[ATTEMPT_INDEX_KEY_NAME] if ROLLOUT_ID_KEY_NAME in row: - # Capture readback recomputes the id from the finished record, so an explicit id - # has to travel from the dispatched row onto the result the same way the indices do. + # Capture readback recomputes the id from the finished record. + # Preserve an explicit id on the result just like the indices. result[ROLLOUT_ID_KEY_NAME] = row[ROLLOUT_ID_KEY_NAME] # Fold this rollout's captured model calls into its record (uniform across agents; no-op diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py index 8d7d85f0cc..91874c5708 100644 --- a/nemo_gym/rollout_correlation.py +++ b/nemo_gym/rollout_correlation.py @@ -31,29 +31,22 @@ _ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None) -# A capture id travels as a path segment in ``/ng-rollout/``, so it is limited to what a path -# segment carries unambiguously. Leading dots are excluded because the id is also a filename -# component in the capture stores. The middleware below matches on the same pattern, so an id this -# rejects is one that would not have survived the round trip anyway. +# A capture id is a path segment in ``/ng-rollout//...``. +# Restrict it to characters that survive a path round trip. +# Exclude leading dots because stores also use the id as a filename component. +# Middleware uses the same pattern. ROLLOUT_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") def maybe_rollout_id_from_run_body(body: BaseModel | Mapping[str, Any] | None) -> Optional[str]: """Build the capture key for a run request. - An explicit ``_ng_rollout_id`` on the body wins. Otherwise the id is derived from the task and - rollout indices as ``"{task}-{rollout}"``. Both forms then take an ``-a{n}`` suffix for a - re-dispatch attempt past the first. - - The derivation is a contract, not an implementation detail: capture writers key records by the - id this returns and capture readers look records up by recomputing it from the finished rollout - record, so the two sides only agree while the rule is the same on both. Changing the format - invalidates any record already on disk. - - The derivation also assumes the caller gives each dispatch a distinct (task, rollout) pair. - A caller that restarts numbering, such as one running the same indices once per training step, - produces a repeated id and two dispatches then share a capture key, which stitches unrelated - calls into one trajectory. Set an explicit id to opt out of the derivation in that case. + An explicit ``_ng_rollout_id`` takes precedence. + Otherwise derive ``"{task}-{rollout}"`` from the task and rollout indices. + Re-dispatch attempts append ``-a{n}``. + Writers and consumers must use this same identity. + Reused task and rollout indices produce a repeated capture key. + Use an explicit id when numbering restarts across dispatches. """ if not isinstance(body, (BaseModel, Mapping)): return None @@ -63,8 +56,8 @@ def field(key: str) -> Any: explicit = field(ROLLOUT_ID_KEY_NAME) if explicit is not None: - # A malformed explicit id is refused rather than sanitized. Rewriting it would correlate - # calls under an id the caller never chose and cannot look up afterwards. + # Reject malformed explicit ids instead of sanitizing them. + # Rewriting would create a key the caller cannot look up. if not (isinstance(explicit, str) and ROLLOUT_ID_PATTERN.match(explicit)): raise ValueError( f"{ROLLOUT_ID_KEY_NAME} must be a string of letters, digits, dots, dashes or " @@ -100,7 +93,8 @@ def rollout_context(rollout_id: Optional[str]) -> Iterator[None]: class RolloutContextMiddleware: """Strip a rollout prefix and expose it to downstream Gym calls for this request.""" - # Same id charset as ROLLOUT_ID_PATTERN, anchored between the prefix and the rest of the path. + # Match the same id characters as ``ROLLOUT_ID_PATTERN``. + # Anchor the id between the prefix and the remaining path. _PREFIX = re.compile( rf"^/{re.escape(ROLLOUT_PATH_PREFIX)}/(?P{ROLLOUT_ID_PATTERN.pattern.strip('^$')})(?P/.*)$" ) diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index 51ebd73dcd..b522ae240e 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -13,23 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Training-token capture: produce, store, read, and source ``TokenEntry`` records. +"""Provide the core training-token capture interfaces. -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``). - -Records are read back through a ``TokenSource``. -``TokenCaptureStore`` is Gym's local implementation. -A framework staging records through its own transport supplies its own source. -The source freezes an atomic snapshot containing both entries and incomplete state. -This prevents a consumer from training on a rollout that lost a model call. +Training capture is separate from evaluation capture. +Middleware sets a request-scoped token sink. +The model server records a ``TokenEntry`` from its complete response. +Consumers call ``TokenSource.freeze`` for an atomic snapshot. +The snapshot includes entries and incomplete state. +Its ``snapshot_id`` identifies the exact frozen state. +``TokenCaptureStore`` is Gym's local sink and source implementation. +Framework transports may provide their own sink and source. +There is no HTTP token reader. +This leaf package avoids imports from Gym's server stack. """ from nemo_gym.token_id_capture.config import TokenIdCaptureConfig diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index a2468e32c0..0e971114ae 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -13,43 +13,40 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Run-wide settings for training-token capture, in one block. +"""Define run-wide training-token capture settings. ```yaml env: nemo_gym: token_id_capture: enabled: true - dir: /tmp/ng_tokcap # node-local; writer and reader share a node - sink: my_pkg.sinks:MyDataPlaneSink # optional; default is the file store at `dir` + dir: /tmp/ng_tokcap # The writer and consumer share this node-local directory. + sink: my_pkg.sinks:MyDataPlaneSink # This optional sink replaces the file store. ``` -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 ``dir`` is given, tokens are written alongside the eval -capture files in the top-level ``model_call_capture_dir``. - -The per-agent ``token_id_capture`` flag is a narrower, separate control: it scopes -which agents participate. Native agents leave it off because they already carry -token ids on their response items. +Evaluation capture uses ``/ng-rollout//...``. +Training capture uses ``/ng-rollout//training-token-capture/...``. +Training capture records token ids and log probabilities. +Evaluation capture records request and response summaries. +A run can enable either path independently. +Training capture applies through the static agent flag or run-level ``all_agents``. +Native agents normally leave the static flag disabled. +Their responses already carry token ids. +The top-level ``model_call_capture_dir`` is the fallback file-store directory. Choosing where records go ------------------------- ``sink`` names a class implementing ``TokenSink``, as ``module.path:ClassName``. -It is constructed once per server process at app startup and replaces the file -store, so records go to a framework's own transport and never touch disk. - -Construction has to happen inside the serving process. A model server configured -with ``num_workers > 1`` is launched by uvicorn with an app string and -``workers=N``, and uvicorn spawns those workers with the ``spawn`` start method, -which re-imports the app module rather than inheriting the parent's memory. A -sink installed by a launcher script therefore does not exist in any worker, and -capture silently falls back to the file store, or writes nothing at all when no -``dir`` is set. Naming the sink here avoids that: each worker builds its own. - -``install_token_sink`` remains for programmatic use and is subject to the same -constraint, so call it at module import of the app, not from a parent process. +Each server process constructs its sink at app startup. +A configured sink replaces the file store. +The paired ``source`` implements ``TokenSource`` for consumers. +Consumers call ``TokenSource.freeze`` to obtain an atomic snapshot. +Consumers retire that exact snapshot with its ``snapshot_id`` and version. +There is no HTTP token reader. +Uvicorn workers use spawned processes. +They do not inherit a sink installed by a launcher. +Configure the sink here so each worker builds its own. +Programmatic installation must occur inside the serving process. """ from __future__ import annotations @@ -111,13 +108,13 @@ class TokenIdCaptureConfig(BaseModel): def _validate(self) -> "TokenIdCaptureConfig": block = self.token_id_capture if not block.enabled: - # The rest of the block is left alone rather than rejected. Configs are templated, and - # setting a directory unconditionally while toggling `enabled` per run is ordinary. + # Keep inactive settings for templated configurations. + # A run may toggle only ``enabled``. return self if block.sink is not None: if block.dir is not None: - # Not an error: nothing is lost, the directory is simply never read. Worth saying - # once, because someone expecting files on disk will not find any. + # The custom sink replaces the configured directory. + # Warn because no files will appear there. logger.warning( "token_id_capture.dir is set alongside token_id_capture.sink. The sink replaces " "the file store, so %s will not be written to.", @@ -130,8 +127,8 @@ def _validate(self) -> "TokenIdCaptureConfig": return self directory = self.resolved_dir() if directory is None: - # A process that installed a sink programmatically writes through that transport and - # never constructs the file store, so it has no directory to give. + # A programmatic sink replaces the file store. + # That process does not need a directory. if installed_token_sink() is not None and ( not block.rebuild_response or installed_token_source() is not None ): @@ -153,10 +150,11 @@ def resolved_dir(self) -> Path | None: return self.token_id_capture.dir or self.model_call_capture_dir def build_sink(self) -> TokenSink | None: - """Construct the configured sink, or ``None`` when the file store is in use. + """Construct the configured sink. - Called once per server process at app startup, which is what makes this work under - ``num_workers > 1`` where a sink installed by a launcher does not reach the workers. + Return ``None`` when the file store is in use. + Call this once in each server process. + Launcher-installed sinks do not reach spawned workers. """ target = self.token_id_capture.sink if not self.token_id_capture.enabled or target is None: @@ -185,8 +183,8 @@ def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: s raise ValueError( f"could not construct token_id_capture.{kind} {target!r} with {kind}_kwargs={sorted(kwargs)}: {error}" ) from error - # Checked here rather than at first use: an endpoint missing a lifecycle method makes an - # incomplete rollout look complete, and a startup error is better than that at step 400. + # Validate the endpoint at startup. + # A missing lifecycle method can make incomplete capture look complete. missing = [name for name in sorted(protocol.__protocol_attrs__) if not callable(getattr(endpoint, name, None))] if missing or not isinstance(endpoint, protocol): raise ValueError( diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index bee1e708e1..5cadec8d97 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -13,23 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Interfaces for writing and reading captured training tokens. - -Gym owns the record shape, these protocols, and the code that builds a record. A -training framework supplies the implementations and runs them wherever its -tokens are produced. Neither side imports the other's transport. - -Placement of the write is therefore a deployment choice: - -- Gym owns serving (today): install the sink in the model server, which already - holds the assembled response, so there is no extra hop. -- A framework owns the inference worker: install the sink there, so bulk token - arrays go to the framework's data plane instead of riding back through Gym's - HTTP response. - -The capture code is the same in both cases. This module must stay free of -fastapi, ray, torch and aiohttp imports so a framework's worker can import it -without pulling in Gym's server stack. A unit test enforces that. +"""Define interfaces for captured training tokens. + +Gym owns the record shape and capture protocols. +A training framework may implement the transport. +The sink may run in a Gym model server. +It may instead run in a framework inference worker. +Engine-side placement keeps token arrays off Gym's HTTP response. +Consumers read through ``TokenSource.freeze``. +They identify the frozen state with ``snapshot_id``. +There is no HTTP token reader. +This module avoids FastAPI, Ray, Torch, and aiohttp imports. """ from __future__ import annotations @@ -113,9 +107,9 @@ async def close(self) -> None: ... -# 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. +# Install these defaults once in the process that owns them. +# The owner may be a Gym model server or a framework inference worker. +# Request-scoped sinks take precedence. _INSTALLED_SINK: TokenSink | None = None _INSTALLED_SOURCE: TokenSource | None = None diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 5d8281ae1f..4861b3e672 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -13,19 +13,14 @@ # 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. +"""Define training-token records extracted from served responses. + +A ``TokenEntry`` contains one model call's training data. +It stores the exact prompt token ids. +It stores generated token ids and their log probabilities. +Evaluation uses a separate ``ModelCallRecord``. +Evaluation records do not carry token arrays. +Both records share a ``model_call_id``. """ from __future__ import annotations @@ -35,35 +30,30 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator -# 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. +# These fields carry token metadata on a served response. +# ``routed_experts`` is optional for MoE backends. TOKEN_FIELDS = ("prompt_token_ids", "generation_token_ids", "generation_log_probs", "routed_experts") -# Bumped whenever a field is added or its meaning changes. Writer and reader are different -# processes and may be different repositories, and records outlive a deploy, so a reader has to -# be able to refuse a record it was not built for. ``extra="allow"`` means an unknown shape -# otherwise decodes cleanly and corrupts training rows in silence. The field is present from the -# first version because a check added later cannot tell an old record from an unversioned one. +# Increment this version when a field or its meaning changes. +# Writers and readers may run in different processes or repositories. +# Records may outlive a deployment. +# Readers must reject unsupported newer records. +# ``extra="allow"`` otherwise hides unknown fields. # # 1 rollout and call identity, the token arrays, the output items and their carrier index TOKEN_ENTRY_RECORD_SCHEMA_VERSION = 1 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. Token - ids alone are not enough for a trainer that scores text, such as penalties for an - invalid tool call or a malformed thinking block. - - The token arrays are stored once, at the top level. The served response carries - them on the output item that produced the generation, and those per-item copies - are dropped on write: the builder overwrites them anyway, since an item's prompt - in a chained trajectory is the running cumulative sequence rather than the prompt - of the single call. ``token_item_index`` records which item they came off, so the - builder can put the chain-correct values back on the right one. + """Store one model call's content and token metadata. + + The rollout id identifies the training sample. + The model call id joins evaluation context. + ``output_items`` preserves assistant text and tool calls. + Text-based penalties require that content. + Token arrays are stored once at the top level. + ``token_item_index`` identifies their original output item. + A trajectory builder can restore chain-correct token fields there. """ model_config = ConfigDict(extra="allow") @@ -76,28 +66,22 @@ class TokenEntry(BaseModel): generation_token_ids: list[int] generation_log_probs: list[float] routed_experts: Any | None = None - # The served response's output items (Responses shape), content preserved, token - # arrays removed. + # Preserve response output items without token arrays. output_items: list[dict] = Field(default_factory=list) - # Index into ``output_items`` of the item the token arrays were taken off, or null - # when no item carried them. Records written before the arrays were de-duplicated - # leave this unset and still carry the arrays inline, which the builder handles. + # This index identifies the item that carried token arrays. + # ``None`` means no item carried them. + # Older records may keep arrays inline and leave this unset. token_item_index: int | None = None - # Non-semantic; a cheap diagnostic for retry/sibling-branch cases. + # This non-semantic timestamp helps diagnose retries and sibling branches. created_at: float = 0.0 @model_validator(mode="after") def _refuse_a_newer_record(self) -> "TokenEntry": - """Decode a record older than this reader, refuse one newer. + """Accept older records and reject newer records. - Older is safe: a field this reader does not have takes its default and the consumer - degrades, so a record written before parent links existed simply has none and the builder - matches token prefixes instead. - - Newer is not, and it is the direction ``extra="allow"`` hides. A field this reader cannot - see is kept and ignored, so a record whose tokens were written under rules this reader does - not know decodes clean and trains as though nothing were different. Refusing is loud: the - read fails, the caller marks that rollout unusable, and the run says which version it saw. + Missing older fields use their defaults. + Unknown newer fields may change token semantics. + Rejecting them prevents silent training corruption. """ if self.schema_version > TOKEN_ENTRY_RECORD_SCHEMA_VERSION: raise ValueError( @@ -120,9 +104,9 @@ def _refuse_a_newer_record(self) -> "TokenEntry": 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. + Responses payloads already carry ``output``. + Chat payloads carry ``choices[*].message``. + Wrap each assistant message as a Responses ``message`` item. """ output = payload.get("output") if isinstance(output, list) and output: diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 0903d7dddc..10cc598ab2 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -13,20 +13,15 @@ # 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, token ids included, for a moment 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. +"""Capture training tokens from one complete model response. + +Streaming responses omit token ids from the wire. +The model server still holds the complete response before streaming. +Middleware provides a request-scoped token sink. +The model server passes its complete response to ``capture_tokens``. +The sink writes a ``TokenEntry``. +Its ``model_call_id`` joins the corresponding evaluation record. +Untagged traffic has no capture context. """ from __future__ import annotations @@ -51,24 +46,20 @@ @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. + """Describe one in-flight training-token capture. - ``sink`` is Gym's file store by default and anything satisfying ``TokenSink`` - otherwise, which is how a training framework redirects the write to its own data - plane without changing the capture code. It is typed as the protocol rather than - the file store so that redirection is a supported path and not an accident. + The context identifies the rollout and model call. + ``sink`` receives the resulting record. + A framework may provide any ``TokenSink`` implementation. """ rollout_id: str model_call_id: str - # None when capture is enabled but nothing in this process writes records. A framework that - # stages engine-side still needs the identity and the parent resolution this context carries, - # and there is no destination here to hand them to. + # ``None`` means another process owns record staging. + # The context still carries the capture identity. sink: TokenSink | None model: str = "" - # Set by ``commit_entry`` so the no-token-ids path can tell a call that was recorded by - # somebody else from one that was lost. + # ``commit_entry`` sets this after another capture path records the call. committed: bool = False @@ -80,11 +71,10 @@ def set_token_sink(sink: CaptureContext) -> Token: def current_capture_context() -> CaptureContext | None: - """The capture context for the in-flight call, or None for untagged traffic. + """Return the capture context for the in-flight call. - The supported way to read the identity this call was minted with. - A framework that stages records from its inference worker keys them on that identity, - and this process is the only one that has it. + Return ``None`` for untagged traffic. + Framework inference workers use this identity for staged records. """ return _TOKEN_SINK.get() @@ -94,21 +84,19 @@ def reset_token_sink(token: Token) -> None: async def capture_tokens(response: Any) -> None: - """Record a ``TokenEntry`` from a complete model response when a sink is set. + """Record a ``TokenEntry`` from a complete model response. - ``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 awaited, so the entry is durable before the model call returns and a - post-rollout reader always sees it, with no background writer to drain. + Accept a Pydantic model or dictionary. + Return without work when no capture context exists. + Mark local capture incomplete when required token ids are absent. + Await the write before the model call returns. """ sink = _TOKEN_SINK.get() if sink is None: return - # Everything that reads the response is guarded, not just the write. Decoding a payload and - # validating a record can fail on malformed token data exactly as writing it can, and the - # consequence is the same: the rollout is short a call. It is guarded here rather than left - # to the caller because the caller is the model server's own response path, so an exception - # escaping this function would fail the model call and break the harness's run. + # Guard response decoding and record validation. + # Either failure leaves the rollout short one call. + # Capture errors must not fail the model call. try: if hasattr(response, "model_dump"): payload = response.model_dump() @@ -132,8 +120,7 @@ async def capture_tokens(response: Any) -> None: generation_token_ids=info["generation_token_ids"], generation_log_probs=info["generation_log_probs"], routed_experts=info.get("routed_experts"), - # Keep the content (assistant text, tool calls) so the trajectory the trainer - # reads is not token-only, since text-based penalties need it. + # Preserve content for text-based training penalties. output_items=content_items, token_item_index=token_item_index, created_at=time.time(), @@ -147,15 +134,12 @@ async def capture_tokens(response: Any) -> None: async def commit_entry(entry: TokenEntry) -> None: """Durably record a finished entry against the in-flight call. - Public and separate from ``capture_tokens`` because the two halves are useful apart. - ``capture_tokens`` reads the arrays off a served response; a framework that captures - engine-side already has them, and the response Gym sees may carry none at all, so it - needs this half without the extraction half. Forking it instead would duplicate the - ordering below, which is the part worth sharing. - - No-op when no sink is active. Never raises: capture is best effort per call, but a - rollout that lost a call is marked so a consumer masks it rather than training on a - chain with a hole. + ``capture_tokens`` extracts arrays from a served response. + Engine-side capture may already have those arrays. + Engine-side callers can use this method directly. + Return without work when no capture context exists. + Capture failures mark the rollout incomplete. + This method never fails the model call. """ sink = _TOKEN_SINK.get() if sink is None: @@ -181,11 +165,9 @@ async def commit_entry(entry: TokenEntry) -> None: async def _capture_failed(sink: CaptureContext, stage: str) -> None: """Report a capture failure without letting it reach the model call. - 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 it is marked, and delivery masks the sample rather than training on a - chain with a hole. Called only from an ``except`` block, so ``exc_info`` picks up the - active exception. + Bad token payloads must not fail the model call. + Mark the rollout so consumers can mask the sample. + Call this only from an ``except`` block. """ logger.warning( "Training-token capture failed to %s the record for model call %s of rollout %s.", @@ -202,14 +184,12 @@ async def _capture_missing(sink: CaptureContext, reason: str) -> None: A response with no token ids is a hole in the chain rather than traffic to skip. The builder reads the gap between one call's tokens and the next call's prompt as tool output. - So a skipped call's generated tokens arrive inside the next prompt at mask 0, - and tokens the policy sampled train as if the environment had written them. + A skipped call's generated tokens then enter the next prompt with mask zero. + Policy tokens would train as if the environment produced them. Two cases are not holes and are left alone. - A call already committed through ``commit_entry`` was recorded by a caller that had the arrays - when this process did not. - A context with no sink means nothing here writes records at all, - so this process cannot tell a lost call from ordinary operation and the staging side owns that. + A committed call was recorded by another capture path. + A context without a sink delegates completeness to external staging. """ if sink.committed or sink.sink is None: return @@ -225,10 +205,8 @@ async def _capture_missing(sink: CaptureContext, reason: str) -> None: async def _mark_incomplete(sink: CaptureContext) -> None: """Mark the rollout, or say loudly why it could not be marked. - A sink that does not implement ``mark_incomplete`` would otherwise raise inside the - failure path above and have the exception swallowed, leaving an incomplete rollout - that looks complete. That is the one outcome this whole path exists to prevent, so - it is logged at error rather than passed over. + A missing ``mark_incomplete`` method can hide incomplete capture. + Log that condition as an error. """ mark = getattr(sink.sink, "mark_incomplete", None) if mark is None: diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index b637b2e156..e712d23ac3 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -13,20 +13,13 @@ # 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 or 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. +"""Store training ``TokenEntry`` records by rollout. + +Each rollout uses one ``.tokens.jsonl`` file. +Evaluation records use a separate file. +Each write uses ``fsync``. +A per-rollout file lock serializes writers to the same rollout. +Different rollouts can write concurrently. """ from __future__ import annotations @@ -239,9 +232,10 @@ def append(self, entry: TokenEntry) -> None: # Both interfaces offload blocking work to the process-wide default thread pool. 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.""" + """Store an entry durably without blocking the event loop. + + Await the append so later consumers cannot race a partial file. + """ await asyncio.to_thread(self.append, entry) async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: @@ -297,8 +291,8 @@ async def close(self) -> None: def delete(self, rollout_id: str) -> None: """Unconditionally remove a rollout's records. - This compatibility helper is for administrative cleanup. Normal - consumers use conditional ``drop``. + This compatibility helper supports administrative cleanup. + Normal consumers use conditional ``drop``. """ with self._locked(rollout_id): self.path_for(rollout_id).unlink(missing_ok=True) @@ -324,10 +318,11 @@ def _read_entries_unlocked(self, rollout_id: str) -> list[TokenEntry]: def make_token_store(global_config_dict: Any) -> TokenCaptureStore | None: - """Build the training-token store, or ``None`` when this process is not writing one. + """Build the training-token file store. - ``None`` when capture is off, when no directory resolves, or when a sink is configured: the - records go to that transport instead and there is no file store to build. + Return ``None`` when capture is disabled. + Return ``None`` when no directory resolves. + Return ``None`` when a custom sink owns the records. """ from nemo_gym.token_id_capture.config import TokenIdCaptureConfig 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 c25b157115..7e32214f7e 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,12 +5,12 @@ claude_code_agent: resources_server: type: resources_servers name: ??? - # This harness returns output with no token ids, so training on its rollouts requires - # capture: its model calls are correlated, captured, and rebuilt into a token-bearing - # response. On for this agent because it is an external harness, which is the case the - # flag exists to identify; a native agent leaves it off, since it carries token ids inline - # and rebuilding would replace them with a reconstruction. This costs nothing on its own: - # the run-level token_id_capture.enabled switch still has to be on for anything to happen. + # This harness returns output without token ids. + # Its static flag opts the agent into training capture. + # The run-level ``token_id_capture.enabled`` setting must also be enabled. + # Run-level ``all_agents`` can opt in every agent instead. + # Native agents normally leave this flag disabled. + # Their responses already carry token ids. token_id_capture: true concurrency: 32 model: claude-sonnet-4-6 diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 053db7f1e1..f3da67c4fa 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -1134,11 +1134,11 @@ def test_maybe_rollout_id_from_run_body_prefers_an_explicit_id(): from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body base = {"_ng_task_index": 3, "_ng_rollout_index": 2} - # A caller that restarts index numbering per dispatch derives the same id twice, which would - # give two dispatches one capture key. The explicit id is how it keeps them apart. + # Restarted index numbering derives the same id twice. + # An explicit id keeps the dispatches separate. assert maybe_rollout_id_from_run_body({**base, "_ng_rollout_id": "s7-3-2"}) == "s7-3-2" - # A retry of an explicitly keyed rollout still keys separately from its first attempt, or the - # retry's calls would append onto the first attempt's records. + # A retry of an explicitly keyed rollout still gets a distinct key. + # Otherwise retry calls would append to the first attempt. assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "s7-3-2", "_ng_attempt_index": 1}) == "s7-3-2-a1" # The explicit id stands alone: no indices needed. assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "abc"}) == "abc" @@ -1153,8 +1153,8 @@ def test_maybe_rollout_id_from_run_body_refuses_an_unusable_explicit_id(bad): # Absent and null both mean "no explicit id", so the derivation still runs. assert maybe_rollout_id_from_run_body(body) == "3-2" return - # An id that cannot survive the round trip is refused rather than rewritten: correlating under - # a sanitized id would file records under a key the caller cannot look up afterwards. + # Reject ids that cannot survive the path round trip. + # Sanitizing would create a key the caller cannot look up. with pytest.raises(ValueError): maybe_rollout_id_from_run_body(body) @@ -1163,8 +1163,8 @@ def test_explicit_rollout_ids_round_trip_through_the_path_prefix(): from nemo_gym.base_responses_api_model import maybe_rollout_id_from_run_body from nemo_gym.rollout_correlation import RolloutContextMiddleware - # The id becomes a path segment, so every id the body accepts has to be one the middleware - # gives back unchanged. A charset the two disagreed on would correlate calls to nothing. + # The id becomes a path segment. + # Middleware must return every accepted id unchanged. for candidate in ["s7-3-2", "step7.task3", "a", "A_b-1.2"]: rollout_id = maybe_rollout_id_from_run_body({"_ng_rollout_id": candidate}) match = RolloutContextMiddleware._PREFIX.match(f"/ng-rollout/{rollout_id}/v1/responses") diff --git a/tests/unit_tests/test_rollout_collection.py b/tests/unit_tests/test_rollout_collection.py index 91502bdc56..8ce96be4f7 100644 --- a/tests/unit_tests/test_rollout_collection.py +++ b/tests/unit_tests/test_rollout_collection.py @@ -1013,9 +1013,9 @@ async def test_run_from_config_keys_capture_by_an_explicit_rollout_id( lambda: {"observability_enabled": True, "model_call_capture_dir": str(capture_dir)}, ) - # Indices that would derive "0-0" plus an explicit id. The explicit id has to win on both - # sides, or the writer and the reader key the same rollout differently and the readback - # finds nothing. + # These indices would derive ``0-0``. + # The explicit id must win for both writer and consumer. + # Otherwise readback finds no matching capture. source_row = { "responses_create_params": {"input": []}, AGENT_REF_KEY_NAME: {"name": "agent"}, @@ -1047,8 +1047,8 @@ def run_examples(self, examples, *args, **kwargs): assert results[0][ROLLOUT_ID_KEY_NAME] == "step7.0-0" assert [call["model_call_id"] for call in results[0]["ng_model_call_capture"]["calls"]] == ["call"] - # Nothing was filed under the derived id, so the explicit id replaced it rather than - # sitting alongside it. + # No capture uses the derived id. + # The explicit id replaces it. assert store.read("0-0") == [] async def test_run_from_config_sorted(self, tmp_path: Path, empty_global_config: MagicMock) -> None: diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 1c3f246e4a..db12259a7a 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -12,12 +12,14 @@ # 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``. +"""Test training-token capture records, stores, and sources. + +Served-path tests build a real ``SimpleResponsesAPIModel``. +Middleware mints a ``model_call_id``. +Middleware sets a request-scoped token sink. +The model server records a ``TokenEntry``. +Consumers read records through ``TokenSource.freeze``. +There is no HTTP token reader. """ import asyncio @@ -258,8 +260,10 @@ def test_config_disabled_needs_no_dir(): def test_config_enabled_requires_absolute_dir(tmp_path): - """A directory that is set has to be absolute. A relative one silently resolves against - whatever the server's working directory happens to be.""" + """Reject a relative capture directory. + + Relative paths depend on the server working directory. + """ with pytest.raises(ValueError): TokenIdCaptureConfig.model_validate(_block(dir="relative/dir")) cfg = TokenIdCaptureConfig.model_validate(_block(dir=str(tmp_path))) @@ -272,16 +276,17 @@ def test_config_falls_back_to_model_call_capture_dir(tmp_path): def test_config_keeps_settings_when_capture_is_off(tmp_path): - """Templated configs set a directory unconditionally and toggle `enabled` per run, so the - rest of the block is left alone rather than rejected.""" + """Allow inactive settings in templated configurations. + + A run may toggle only ``enabled``. + """ cfg = TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": False, "dir": str(tmp_path)}}) assert cfg.enabled is False assert cfg.build_sink() is None def test_config_warns_rather_than_fails_on_a_sink_beside_a_directory(caplog): - """Nothing is lost, the directory is just never read, but someone expecting files on disk - will not find any.""" + """Warn when a custom sink replaces the configured directory.""" with caplog.at_level(logging.WARNING): cfg = TokenIdCaptureConfig.model_validate(_block(sink=f"{__name__}:_ConfiguredSink", dir="/tmp/x")) assert cfg.enabled is True @@ -409,10 +414,11 @@ def test_captured_entry_carries_content(tmp_path): def test_token_arrays_are_stored_once(tmp_path): - """The served response carries the arrays on an output item; the record does not repeat them. + """Store token arrays once and preserve response content. - Storing them again per item roughly doubles a record, and the per-item copy is not the - value a trainer reads: an item's prompt in a chained trajectory is the running sequence. + Served responses carry arrays on an output item. + Captured records move them to the entry. + Chained trajectories rebuild each item's running prompt. """ client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) client.post("/ng-rollout/task0-rollDedup/training-token-capture/v1/responses", json={"input": "hi"}) @@ -478,13 +484,11 @@ def test_uncorrelated_call_captures_nothing(tmp_path): def test_package_is_dependency_free_leaf(): - """``nemo_gym.token_id_capture`` must import without Gym's server stack. + """Keep token capture independent of 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. + Framework inference workers import the record and protocols. + They must not import Ray, FastAPI, or uvicorn through this package. + A subprocess isolates this check from earlier test imports. """ heavy = ("ray", "fastapi", "uvicorn", "aiohttp", "requests", "torch") program = ( @@ -496,11 +500,11 @@ def test_package_is_dependency_free_leaf(): def test_streamed_messages_capture_tokens_absent_from_the_stream(tmp_path): - """The Claude Code shape: streamed /v1/messages. + """Capture tokens before streaming Anthropic 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. + Token ids exist only on the assembled response. + Anthropic conversion omits them from SSE. + This test covers the complete served path. """ client = TestClient(_server(_both_enabled(tmp_path)).setup_webserver()) with client.stream( @@ -527,10 +531,10 @@ def test_streamed_messages_capture_tokens_absent_from_the_stream(tmp_path): def test_capture_failure_marks_the_rollout_incomplete(tmp_path, monkeypatch): - """A lost call must not leave the rollout looking complete. + """Mark a rollout incomplete when capture loses a call. - 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". + A bad payload must not break the model call. + Consumers must still detect the missing record. """ store = TokenCaptureStore(tmp_path) @@ -566,11 +570,10 @@ def _silent_server(global_config_dict) -> SimpleResponsesAPIModel: def test_a_response_without_token_ids_marks_the_rollout_incomplete(tmp_path): - """A call that returns no token ids is a hole, not traffic to skip. + """Treat missing token ids as an incomplete capture. - Skipping it quietly is the worse failure: the rollout still looks complete, and the - call's generated tokens end up inside the next call's prompt, where they are trained - as if the environment had written them. + Silent omission makes the rollout look complete. + Generated tokens may then enter the next prompt with mask zero. """ client = TestClient(_silent_server(_both_enabled(tmp_path)).setup_webserver()) resp = client.post("/ng-rollout/silent0-roll0/training-token-capture/v1/responses", json={"input": "hi"}) @@ -591,9 +594,10 @@ def _external_mode(tmp_path) -> dict: def test_external_mode_still_mints_identity_for_a_correlated_call(tmp_path): - """A framework staging from the inference worker keys its record on the identity minted here. + """Create capture identity without a local destination. - Nothing in this process writes, so the machinery cannot be gated on having a destination. + Framework inference workers use the identity minted here. + Local destination availability must not gate that identity. """ seen = {} @@ -622,10 +626,10 @@ async def responses(self, request: Request, body=Body()) -> NeMoGymResponse: def test_external_mode_does_not_mark_a_token_less_response_incomplete(tmp_path): - """Under external staging a response with no token ids is ordinary, not a hole. + """Leave completeness to external staging without a local destination. - This process writes nothing, so it cannot tell a lost call from normal operation, and - marking here would mask every rollout of the run. + This process cannot distinguish a lost call from normal external capture. + Marking locally would mask every rollout. """ client = TestClient(_silent_server(_external_mode(tmp_path)).setup_webserver()) assert ( @@ -687,9 +691,10 @@ def test_delete_removes_records_and_marker(tmp_path): def test_concurrent_appends_to_one_rollout_stay_intact(tmp_path): - """Writes take an exclusive file lock, which is what keeps two writers from interleaving a - partial line. Under sharding the writers are separate processes, so the lock has to hold there - too; this covers the same code path with threads.""" + """Keep concurrent writers from interleaving partial records. + + The exclusive file lock covers threads and processes. + """ import concurrent.futures store = TokenCaptureStore(tmp_path) @@ -718,11 +723,10 @@ def test_concurrent_appends_to_one_rollout_stay_intact(tmp_path): class _RecordingSink: - """A sink that is only a ``TokenSink``: no file store, no directory. + """Implement ``TokenSink`` without a file store. - Deliberately not a ``TokenCaptureStore`` subclass. A training framework whose sink is - its own transport has nothing on disk, and this is the shape the capture path has to - accept for ``install_token_sink`` to mean anything. + Framework transports may keep no local files. + The capture path must accept this protocol-only implementation. """ def __init__(self) -> None: @@ -766,10 +770,9 @@ def test_config_allows_no_directory_when_a_sink_is_installed(installed_sink): def test_config_allows_capture_with_no_destination_at_all(): - """A framework that stages records from the inference worker writes nothing here. + """Allow external staging without a local store. - It still needs capture on, because the identity a record is keyed by and the parent a - request continues are resolved in this process and nowhere else. + This process still resolves the capture identity. """ cfg = TokenIdCaptureConfig.model_validate(_block()) assert cfg.enabled is True @@ -778,10 +781,9 @@ def test_config_allows_capture_with_no_destination_at_all(): def test_installed_sink_is_marked_incomplete_through_the_protocol(installed_sink, monkeypatch): - """A protocol-only sink must receive the incomplete signal. + """Send incomplete state through the ``TokenSink`` protocol. - Reaching for a concrete store attribute here would raise inside the failure path and be - swallowed, leaving a rollout that lost a call looking complete. + Capture code must not require concrete store attributes. """ async def boom(entry): @@ -815,10 +817,10 @@ async def put(self, entry): def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(installed_sink): - """Engine-side capture: the caller has the arrays, the served response does not. + """Allow engine-side capture to commit an existing entry. - The commit half has to be reachable on its own, otherwise a framework in that position - forks the durability ordering rather than sharing it. + Engine-side callers already have the token arrays. + They should share the standard durability path. """ entry = TokenEntry( rollout_id="task0-sink3", @@ -850,11 +852,11 @@ def test_records_carry_a_schema_version(): def test_a_malformed_token_payload_does_not_fail_the_model_call(installed_sink): - """Building the record is guarded, not just writing it. + """Guard record construction failures. - ``capture_tokens`` is awaited directly on the model server's response path, so anything - it raises fails the model call. A payload whose token fields do not validate has to be - treated like any other capture failure: the call succeeds and the rollout is marked. + ``capture_tokens`` runs on the model response path. + Invalid token fields must not fail the model call. + The rollout must still be marked incomplete. """ entry_ctor = TokenEntry @@ -878,16 +880,17 @@ def _bad_entry(**kwargs): @pytest.mark.parametrize("bad", ["", "a/b", "../escape", "a b"]) def test_an_unsafe_rollout_id_is_rejected(tmp_path, bad): - """The id names the capture file, so it has to be a safe filename component: a separator - would let a rollout id write outside the store directory.""" + """Reject rollout ids that could escape the store directory.""" with pytest.raises(ValueError): TokenCaptureStore(tmp_path).path_for(bad) def test_a_record_is_readable_as_soon_as_put_returns(tmp_path): - """``put`` is awaited rather than backgrounded, so the record is on disk before the model - call returns. A reader in another process runs after the rollout and has no way to wait for - a writer, and delete-on-consume is only safe because nothing is still in flight.""" + """Make ``put`` durable before it returns. + + Consumers may run in another process after rollout completion. + Conditional deletion requires all writes to be finished. + """ store = TokenCaptureStore(tmp_path) entry = TokenEntry( rollout_id="r0", @@ -901,8 +904,11 @@ def test_a_record_is_readable_as_soon_as_put_returns(tmp_path): def test_a_rollout_that_lost_a_call_is_distinguishable_from_a_complete_one(tmp_path): - """Capture failures do not fail the model call, so nothing downstream would otherwise know - a turn is missing. The chain built from what survived can look perfectly contiguous.""" + """Expose incomplete capture to consumers. + + Capture failures do not fail model calls. + Surviving records may otherwise look contiguous. + """ store = TokenCaptureStore(tmp_path) entry = TokenEntry( rollout_id="r0", @@ -1009,8 +1015,7 @@ def test_a_configured_sink_wins_over_an_installed_one(installed_sink): def test_a_sink_receives_its_configured_kwargs(): - """A sink for a real transport needs an endpoint and a client; a zero-argument one could only - get them from ambient state.""" + """Require explicit constructor wiring for a transport sink.""" config = TokenIdCaptureConfig.model_validate( _block(sink=f"{__name__}:_KwargSink", sink_kwargs={"endpoint": "https://dp", "shard": 3}) ) @@ -1025,8 +1030,7 @@ def test_a_sink_given_kwargs_it_cannot_take_is_refused_at_startup(): def test_a_sink_that_cannot_report_failures_is_refused_at_startup(): - """Without mark_incomplete an incomplete rollout looks complete, so this fails at startup - rather than at whichever step first loses a call.""" + """Reject sinks that cannot mark incomplete capture.""" config = TokenIdCaptureConfig.model_validate( { "token_id_capture": { @@ -1041,9 +1045,11 @@ def test_a_sink_that_cannot_report_failures_is_refused_at_startup(): def test_a_sink_whose_protocol_member_is_not_callable_is_refused(): - """isinstance against a Protocol only checks that the attributes exist, so callability is - checked too. Both are derived from the protocol rather than a list written out here, so the - check keeps up if TokenSink gains a method.""" + """Require callable methods for the ``TokenSink`` protocol. + + Attribute presence alone is insufficient. + Derive the checks from the protocol. + """ config = TokenIdCaptureConfig.model_validate( { "token_id_capture": { @@ -1070,15 +1076,11 @@ def test_a_malformed_sink_path_is_refused_at_startup(target, expected): def test_a_programmatically_installed_sink_does_not_reach_a_spawned_worker(): - """A model server with num_workers > 1 is launched by uvicorn with an app string and - workers=N, and uvicorn spawns those workers (multiprocessing "spawn"), re-importing the app - module rather than inheriting the launcher's memory. ``install_token_sink`` sets a process - global, so it does not cross that boundary and capture silently falls back to the file store, - or writes nothing when no directory is set. - - This is why ``token_id_capture.sink`` is configuration rather than only a function call: it is - constructed inside each worker. The test pins the limitation so the reason for the config key - does not get lost. + """Build configured sinks inside spawned workers. + + Uvicorn workers re-import the app module. + They do not inherit launcher process globals. + Each worker must construct its configured sink. """ ctx = multiprocessing.get_context("spawn") # the context uvicorn uses queue = ctx.Queue() @@ -1097,8 +1099,10 @@ def _report_installed_sink(queue) -> None: def test_the_store_is_a_token_source(tmp_path): - """Records are read back through a TokenSource, and the file store is one. There is no - separate local reader: a wrapper over the store would only forward every call.""" + """Use the file store as the local ``TokenSource``. + + A separate local reader would only forward each call. + """ store = TokenCaptureStore(tmp_path) assert isinstance(store, TokenSource) @@ -1113,8 +1117,8 @@ def test_the_store_is_a_token_source(tmp_path): ) assert [e.model_call_id for e in asyncio.run(store.tokens_for("r0"))] == ["c1"] - # A colocated source can tell that a call failed to capture, which is what keeps an - # incomplete rollout from being trained on. + # A colocated source can detect a capture failure. + # This prevents training on an incomplete rollout. assert store.is_incomplete("r0") is False asyncio.run(store.mark_incomplete("r0", "c2")) assert store.is_incomplete("r0") is True @@ -1132,22 +1136,19 @@ def _entry_fields(**overrides): def test_a_record_older_than_this_reader_is_accepted(): - """A field this reader does not have takes its default and the consumer degrades: a record - written before parent links existed simply has none, and the builder matches prefixes.""" + """Use defaults for fields absent from older records.""" entry = TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_RECORD_SCHEMA_VERSION - 1)) assert entry.generation_token_ids == [2] def test_a_record_newer_than_this_reader_is_refused(): - """The direction extra="allow" hides. A field this reader cannot see is kept and ignored, so - without this the record decodes clean and trains as though nothing were different.""" + """Reject newer records hidden by ``extra="allow"``.""" with pytest.raises(ValidationError, match="this reader understands up to"): TokenEntry(**_entry_fields(schema_version=TOKEN_ENTRY_RECORD_SCHEMA_VERSION + 1)) def test_a_newer_record_in_the_store_fails_the_read_rather_than_being_skipped(tmp_path): - """Read failure is the loud path: the caller marks that rollout unusable rather than training - on a partial set that looks complete.""" + """Fail loudly instead of training on a partial newer record.""" store = TokenCaptureStore(tmp_path) store.append(TokenEntry(**_entry_fields())) path = next(tmp_path.glob("*.tokens.jsonl")) From e85bdee549dbbf4fbdf475d066140effb15afe8d Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 18 Aug 2026 00:23:52 -0700 Subject: [PATCH 07/12] docs(token-id-capture): finish capture comment cleanup Keep remaining protocol and storage comments to one complete thought per line. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/config.py | 3 ++- nemo_gym/token_id_capture/protocols.py | 10 ++++++++-- nemo_gym/token_id_capture/sink.py | 3 ++- nemo_gym/token_id_capture/store.py | 5 ++++- tests/unit_tests/test_base_responses_api_model.py | 2 +- 5 files changed, 17 insertions(+), 6 deletions(-) diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 0e971114ae..3f4486a2ef 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -80,7 +80,8 @@ class TokenIdCaptureSettings(BaseModel): # Capture model calls from every agent. # The default keeps capture scoped by each agent's ``token_id_capture`` setting. all_agents: bool = False - # Where the default file store writes. Falls back to ``model_call_capture_dir``. + # Where the default file store writes. + # Falls back to ``model_call_capture_dir``. dir: Path | None = None # ``module.path:ClassName`` implementing TokenSink, constructed per server process. sink: str | None = None diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 5cadec8d97..216bd7f818 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -76,7 +76,10 @@ async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> Non ... async def close(self) -> None: - """Flush pending work and release resources. Idempotent.""" + """Flush pending work and release resources. + + This operation is idempotent. + """ ... @@ -103,7 +106,10 @@ async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool ... async def close(self) -> None: - """Release resources. Idempotent.""" + """Release resources. + + This operation is idempotent. + """ ... diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 10cc598ab2..2c62e36e6a 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -109,7 +109,8 @@ async def capture_tokens(response: Any) -> None: if info is None: await _capture_missing(sink, "the response carries no token ids") return - # Content only: the arrays live on the entry, not on the items as well. + # Keep content on the output items. + # Store token arrays only on the entry. content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) entry = TokenEntry( diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index e712d23ac3..360ffee02c 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -263,7 +263,10 @@ def freeze_now(self, rollout_id: str) -> TokenCaptureSnapshot: ) async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: - """Compatibility read for diagnostics. Consumers should use ``freeze``.""" + """Read records for compatibility diagnostics. + + Consumers should use ``freeze``. + """ return await asyncio.to_thread(self.read_entries, rollout_id) async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index f3da67c4fa..073c151e29 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -1140,7 +1140,7 @@ def test_maybe_rollout_id_from_run_body_prefers_an_explicit_id(): # A retry of an explicitly keyed rollout still gets a distinct key. # Otherwise retry calls would append to the first attempt. assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "s7-3-2", "_ng_attempt_index": 1}) == "s7-3-2-a1" - # The explicit id stands alone: no indices needed. + # The explicit id requires no indices. assert maybe_rollout_id_from_run_body({"_ng_rollout_id": "abc"}) == "abc" From 93c54a905513db79e05b5b6f7c43882cccd43057 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 18 Aug 2026 00:38:30 -0700 Subject: [PATCH 08/12] fix(token-id-capture): reject writes after retirement Keep a frozen state tombstone after conditional drop so a late writer from the retired attempt cannot recreate its records. Explicit pre-dispatch cleanup starts the next attempt. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/store.py | 5 ++++- tests/unit_tests/test_token_id_capture.py | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index 360ffee02c..02a145411b 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -284,7 +284,10 @@ def _drop(self, rollout_id: str, snapshot_id: str, version: int) -> bool: return False self.path_for(rollout_id).unlink(missing_ok=True) self.incomplete_path_for(rollout_id).unlink(missing_ok=True) - self.state_path_for(rollout_id).unlink(missing_ok=True) + # Keep a frozen tombstone until explicit pre-dispatch cleanup. + # A late writer from this attempt must still observe the freeze. + state["retired"] = True + self._write_state(rollout_id, state) self._fsync_root() return True diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index db12259a7a..b8967d4054 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -244,6 +244,14 @@ def test_token_store_freeze_is_atomic_and_conditional_drop_is_race_safe(tmp_path assert updated.incomplete assert asyncio.run(store.drop("r0", snapshot_id=updated.snapshot_id, version=updated.version)) assert store.read_entries("r0") == [] + assert orjson.loads(store.state_path_for("r0").read_bytes())["retired"] is True + with pytest.raises(RuntimeError, match="already frozen"): + asyncio.run(store.put(entry.model_copy(update={"model_call_id": "late-after-drop"}))) + + store.delete("r0") + replacement = entry.model_copy(update={"model_call_id": "replacement"}) + asyncio.run(store.put(replacement)) + assert store.read_entries("r0") == [replacement] # --- config ------------------------------------------------------------------- From 6a902fd865112dd241b8e4daab5f6bc66038d874 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 18 Aug 2026 11:08:57 -0700 Subject: [PATCH 09/12] fix(token-id-capture): keep source ownership with consumers Remove framework source construction from Gym configuration so consumers create and inject sources in their own process. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/config.py | 27 ++++------------------- nemo_gym/token_id_capture/protocols.py | 5 ++++- tests/unit_tests/test_token_id_capture.py | 20 ++++++++++++++--- 3 files changed, 25 insertions(+), 27 deletions(-) diff --git a/nemo_gym/token_id_capture/config.py b/nemo_gym/token_id_capture/config.py index 3f4486a2ef..119d06857d 100644 --- a/nemo_gym/token_id_capture/config.py +++ b/nemo_gym/token_id_capture/config.py @@ -38,8 +38,9 @@ ------------------------- ``sink`` names a class implementing ``TokenSink``, as ``module.path:ClassName``. Each server process constructs its sink at app startup. +A framework must make that class importable in the server process. A configured sink replaces the file store. -The paired ``source`` implements ``TokenSource`` for consumers. +Consumers construct and inject their ``TokenSource`` in their own process. Consumers call ``TokenSource.freeze`` to obtain an atomic snapshot. Consumers retire that exact snapshot with its ``snapshot_id`` and version. There is no HTTP token reader. @@ -60,9 +61,7 @@ from nemo_gym.token_id_capture.protocols import ( TokenSink, - TokenSource, installed_token_sink, - installed_token_source, ) @@ -89,9 +88,6 @@ class TokenIdCaptureSettings(BaseModel): # A real transport needs explicit endpoint, client, or credential wiring. # Use ``${oc.env:VAR}`` for secrets instead of writing them here. sink_kwargs: dict[str, Any] = Field(default_factory=dict) - # Optional paired reader for framework-owned transports. - source: str | None = None - source_kwargs: dict[str, Any] = Field(default_factory=dict) # Rebuild opaque-harness responses from captured records after the run. rebuild_response: bool = True @@ -121,24 +117,16 @@ def _validate(self) -> "TokenIdCaptureConfig": "the file store, so %s will not be written to.", block.dir, ) - if block.rebuild_response and block.source is None and installed_token_source() is None: - raise ValueError( - "token_id_capture.source is required when a custom sink is used with rebuild_response=true" - ) return self directory = self.resolved_dir() if directory is None: # A programmatic sink replaces the file store. # That process does not need a directory. - if installed_token_sink() is not None and ( - not block.rebuild_response or installed_token_source() is not None - ): - return self - if block.source is not None: + if installed_token_sink() is not None: return self if not block.rebuild_response: return self - raise ValueError("token_id_capture requires a directory or paired source when rebuild_response=true") + raise ValueError("token_id_capture requires a directory or sink") if not directory.is_absolute(): raise ValueError("training-token capture directory must be an absolute path") return self @@ -162,13 +150,6 @@ def build_sink(self) -> TokenSink | None: return None return self._build_endpoint(target, self.token_id_capture.sink_kwargs, TokenSink, "sink") - def build_source(self) -> TokenSource | None: - """Construct a configured framework-owned source.""" - target = self.token_id_capture.source - if not self.token_id_capture.enabled or target is None: - return None - return self._build_endpoint(target, self.token_id_capture.source_kwargs, TokenSource, "source") - @staticmethod def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: str): if ":" not in target: diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 216bd7f818..25181825f3 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -131,7 +131,10 @@ def installed_token_sink() -> TokenSink | None: def install_token_source(source: TokenSource | None) -> None: - """Set (or clear, with ``None``) the process-wide default source.""" + """Set (or clear) the caller-owned source in this process. + + Gym does not close an installed source. + """ global _INSTALLED_SOURCE _INSTALLED_SOURCE = source diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index b8967d4054..1266bfcbe8 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -307,10 +307,24 @@ def test_config_rejects_an_unknown_key(): TokenIdCaptureConfig.model_validate({"token_id_capture": {"enabled": True, "dirr": "/tmp/x"}}) -def test_config_rejects_write_only_custom_capture(): - with pytest.raises(ValueError, match="source is required"): +def test_config_accepts_a_sink_without_constructing_the_consumer_source(): + config = TokenIdCaptureConfig.model_validate( + {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + ) + assert config.token_id_capture.sink == f"{__name__}:_ConfiguredSink" + + +@pytest.mark.parametrize( + ("key", "value"), + [ + ("source", "framework.capture:Source"), + ("source_kwargs", {"endpoint": "transport://tokens"}), + ], +) +def test_config_rejects_framework_source_construction(key, value): + with pytest.raises(ValueError, match="Extra inputs are not permitted"): TokenIdCaptureConfig.model_validate( - {"token_id_capture": {"enabled": True, "sink": f"{__name__}:_ConfiguredSink"}} + {"token_id_capture": {"enabled": True, "rebuild_response": False, key: value}} ) From bb40ee257537df8936baf5d4eff52c7d66cb7e6e Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 18 Aug 2026 11:31:31 -0700 Subject: [PATCH 10/12] fix(token-id-capture): reserve the run config block Keep token_id_capture out of server discovery so env prefetch does not treat run-wide capture settings as a server. Signed-off-by: Ananth Subramaniam --- nemo_gym/global_config.py | 1 + tests/unit_tests/test_base_responses_api_model.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 066579dbb4..3c62d52dd0 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -132,6 +132,7 @@ QUERY_KEY_NAME, OBSERVABILITY_ENABLED_KEY_NAME, MODEL_CALL_CAPTURE_DIR_KEY_NAME, + TOKEN_ID_CAPTURE_BLOCK, COMPONENT_NAME_KEY_NAME, SKIP_VERIFICATION_KEY_NAME, SKIP_VERIFICATION_REWARD_KEY_NAME, diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 073c151e29..2b1d485c73 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -587,7 +587,9 @@ def text(self): # --- capture-store config + init failure --- def test_model_call_capture_keys_are_reserved_global_config(): - assert {"observability_enabled", "model_call_capture_dir"} <= set(NEMO_GYM_RESERVED_TOP_LEVEL_KEYS) + assert {"observability_enabled", "model_call_capture_dir", "token_id_capture"} <= set( + NEMO_GYM_RESERVED_TOP_LEVEL_KEYS + ) def test_model_call_capture_config_requires_absolute_dir_when_enabled(tmp_path, monkeypatch): From f1fef591758782d31b25a282c3051008022cb4dc Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 18 Aug 2026 13:41:04 -0700 Subject: [PATCH 11/12] fix(token-id-capture): clarify capture ownership and defaults Name the request context separately from its sink, remove speculative compatibility prose, and keep Claude Code capture opt-in rather than enabled by default. Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_model.py | 2 +- nemo_gym/token_id_capture/protocols.py | 16 +--- nemo_gym/token_id_capture/records.py | 1 - nemo_gym/token_id_capture/sink.py | 82 +++++++++---------- nemo_gym/token_id_capture/store.py | 3 +- .../configs/claude_code_agent.yaml | 7 -- .../claude_code_agent/tests/test_app.py | 1 + tests/unit_tests/test_token_id_capture.py | 6 +- 8 files changed, 51 insertions(+), 67 deletions(-) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 7ca2e994c6..062181e59c 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -1122,7 +1122,7 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None sink_token = None if capture_wanted: sink_token = set_token_sink( - CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, sink=token_sink) + CaptureContext(rollout_id=rollout_id, model_call_id=model_call_id, token_sink=token_sink) ) # Training-only capture has no evaluation record. diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 25181825f3..20c7637905 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -22,7 +22,6 @@ Engine-side placement keeps token arrays off Gym's HTTP response. Consumers read through ``TokenSource.freeze``. They identify the frozen state with ``snapshot_id``. -There is no HTTP token reader. This module avoids FastAPI, Ray, Torch, and aiohttp imports. """ @@ -47,10 +46,7 @@ class TokenCaptureSnapshot: @runtime_checkable class TokenSink(Protocol): - """Where captured records go. - - Gym's file store and framework-owned transports implement this protocol. - """ + """Receive captured records through Gym's file store or a framework transport.""" async def put(self, entry: TokenEntry) -> None: """Durably store one record. @@ -76,10 +72,7 @@ async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> Non ... async def close(self) -> None: - """Flush pending work and release resources. - - This operation is idempotent. - """ + """Flush pending work and release resources idempotently.""" ... @@ -106,10 +99,7 @@ async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool ... async def close(self) -> None: - """Release resources. - - This operation is idempotent. - """ + """Release resources idempotently.""" ... diff --git a/nemo_gym/token_id_capture/records.py b/nemo_gym/token_id_capture/records.py index 4861b3e672..7f77b3e725 100644 --- a/nemo_gym/token_id_capture/records.py +++ b/nemo_gym/token_id_capture/records.py @@ -70,7 +70,6 @@ class TokenEntry(BaseModel): output_items: list[dict] = Field(default_factory=list) # This index identifies the item that carried token arrays. # ``None`` means no item carried them. - # Older records may keep arrays inline and leave this unset. token_item_index: int | None = None # This non-semantic timestamp helps diagnose retries and sibling branches. created_at: float = 0.0 diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 2c62e36e6a..53a595cdcb 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -49,7 +49,7 @@ class CaptureContext: """Describe one in-flight training-token capture. The context identifies the rollout and model call. - ``sink`` receives the resulting record. + ``token_sink`` receives the resulting record. A framework may provide any ``TokenSink`` implementation. """ @@ -57,17 +57,17 @@ class CaptureContext: model_call_id: str # ``None`` means another process owns record staging. # The context still carries the capture identity. - sink: TokenSink | None + token_sink: TokenSink | None model: str = "" # ``commit_entry`` sets this after another capture path records the call. committed: bool = False -_TOKEN_SINK: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_token_sink", default=None) +_CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_capture_context", default=None) -def set_token_sink(sink: CaptureContext) -> Token: - return _TOKEN_SINK.set(sink) +def set_token_sink(context: CaptureContext) -> Token: + return _CAPTURE_CONTEXT.set(context) def current_capture_context() -> CaptureContext | None: @@ -76,11 +76,11 @@ def current_capture_context() -> CaptureContext | None: Return ``None`` for untagged traffic. Framework inference workers use this identity for staged records. """ - return _TOKEN_SINK.get() + return _CAPTURE_CONTEXT.get() def reset_token_sink(token: Token) -> None: - _TOKEN_SINK.reset(token) + _CAPTURE_CONTEXT.reset(token) async def capture_tokens(response: Any) -> None: @@ -91,8 +91,8 @@ async def capture_tokens(response: Any) -> None: Mark local capture incomplete when required token ids are absent. Await the write before the model call returns. """ - sink = _TOKEN_SINK.get() - if sink is None: + context = _CAPTURE_CONTEXT.get() + if context is None: return # Guard response decoding and record validation. # Either failure leaves the rollout short one call. @@ -103,20 +103,20 @@ async def capture_tokens(response: Any) -> None: elif isinstance(response, dict): payload = response else: - await _capture_missing(sink, f"the response is a {type(response).__name__}") + await _capture_missing(context, f"the response is a {type(response).__name__}") return info = extract_token_fields(payload) if info is None: - await _capture_missing(sink, "the response carries no token ids") + await _capture_missing(context, "the response carries no token ids") return # Keep content on the output items. # Store token arrays only on the entry. content_items, token_item_index = strip_token_fields(response_to_output_items(payload)) entry = TokenEntry( - rollout_id=sink.rollout_id, - model_call_id=sink.model_call_id, - model=sink.model or str(payload.get("model") or ""), + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + model=context.model or str(payload.get("model") or ""), prompt_token_ids=info["prompt_token_ids"], generation_token_ids=info["generation_token_ids"], generation_log_probs=info["generation_log_probs"], @@ -127,7 +127,7 @@ async def capture_tokens(response: Any) -> None: created_at=time.time(), ) except Exception: - await _capture_failed(sink, "build") + await _capture_failed(context, "build") return await commit_entry(entry) @@ -142,28 +142,28 @@ async def commit_entry(entry: TokenEntry) -> None: Capture failures mark the rollout incomplete. This method never fails the model call. """ - sink = _TOKEN_SINK.get() - if sink is None: + context = _CAPTURE_CONTEXT.get() + if context is None: return - if entry.rollout_id != sink.rollout_id or entry.model_call_id != sink.model_call_id: + if entry.rollout_id != context.rollout_id or entry.model_call_id != context.model_call_id: logger.warning( "Training-token capture identity mismatch for model call %s of rollout %s.", - sink.model_call_id, - sink.rollout_id, + context.model_call_id, + context.rollout_id, ) - await _mark_incomplete(sink) + await _mark_incomplete(context) return - if sink.sink is None: - sink.committed = True + if context.token_sink is None: + context.committed = True return try: - await sink.sink.put(entry) - sink.committed = True + await context.token_sink.put(entry) + context.committed = True except Exception: - await _capture_failed(sink, "write") + await _capture_failed(context, "write") -async def _capture_failed(sink: CaptureContext, stage: str) -> None: +async def _capture_failed(context: CaptureContext, stage: str) -> None: """Report a capture failure without letting it reach the model call. Bad token payloads must not fail the model call. @@ -173,14 +173,14 @@ async def _capture_failed(sink: CaptureContext, stage: str) -> None: logger.warning( "Training-token capture failed to %s the record for model call %s of rollout %s.", stage, - sink.model_call_id, - sink.rollout_id, + context.model_call_id, + context.rollout_id, exc_info=True, ) - await _mark_incomplete(sink) + await _mark_incomplete(context) -async def _capture_missing(sink: CaptureContext, reason: str) -> None: +async def _capture_missing(context: CaptureContext, reason: str) -> None: """Mark the rollout when a call this process should have recorded produced nothing. A response with no token ids is a hole in the chain rather than traffic to skip. @@ -192,33 +192,33 @@ async def _capture_missing(sink: CaptureContext, reason: str) -> None: A committed call was recorded by another capture path. A context without a sink delegates completeness to external staging. """ - if sink.committed or sink.sink is None: + if context.committed or context.token_sink is None: return logger.warning( "Training-token capture has no token ids for model call %s of rollout %s: %s.", - sink.model_call_id, - sink.rollout_id, + context.model_call_id, + context.rollout_id, reason, ) - await _mark_incomplete(sink) + await _mark_incomplete(context) -async def _mark_incomplete(sink: CaptureContext) -> None: +async def _mark_incomplete(context: CaptureContext) -> None: """Mark the rollout, or say loudly why it could not be marked. A missing ``mark_incomplete`` method can hide incomplete capture. Log that condition as an error. """ - mark = getattr(sink.sink, "mark_incomplete", None) + mark = getattr(context.token_sink, "mark_incomplete", None) if mark is None: logger.error( "Sink %s does not implement mark_incomplete. Rollout %s cannot be marked incomplete " "and may be trained on with a missing call.", - type(sink.sink).__name__, - sink.rollout_id, + type(context.token_sink).__name__, + context.rollout_id, ) return try: - await mark(sink.rollout_id, sink.model_call_id) + await mark(context.rollout_id, context.model_call_id) except Exception: - logger.warning("Could not mark rollout %s incomplete.", sink.rollout_id, exc_info=True) + logger.warning("Could not mark rollout %s incomplete.", context.rollout_id, exc_info=True) diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index 02a145411b..8735e6d041 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -270,7 +270,7 @@ async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: return await asyncio.to_thread(self.read_entries, rollout_id) async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: - """Conditionally delete the frozen snapshot.""" + """Delete snapshot payloads while retaining its tombstone and lock.""" return await asyncio.to_thread(self._drop, rollout_id, snapshot_id, version) def _drop(self, rollout_id: str, snapshot_id: str, version: int) -> bool: @@ -299,6 +299,7 @@ def delete(self, rollout_id: str) -> None: This compatibility helper supports administrative cleanup. Normal consumers use conditional ``drop``. + The lock file remains so concurrent callers keep using one inode. """ with self._locked(rollout_id): self.path_for(rollout_id).unlink(missing_ok=True) 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 7e32214f7e..5e3140bea7 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,13 +5,6 @@ claude_code_agent: resources_server: type: resources_servers name: ??? - # This harness returns output without token ids. - # Its static flag opts the agent into training capture. - # The run-level ``token_id_capture.enabled`` setting must also be enabled. - # Run-level ``all_agents`` can opt in every agent instead. - # Native agents normally leave this flag disabled. - # Their responses already carry token ids. - 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/tests/test_app.py b/responses_api_agents/claude_code_agent/tests/test_app.py index 472163cb79..5587dae04d 100644 --- a/responses_api_agents/claude_code_agent/tests/test_app.py +++ b/responses_api_agents/claude_code_agent/tests/test_app.py @@ -1041,5 +1041,6 @@ def test_config_yaml_parses(self) -> None: assert "claude_code_agent" in data inner = data["claude_code_agent"]["responses_api_agents"]["claude_code_agent"] assert inner["entrypoint"] == "app.py" + assert inner.get("token_id_capture", False) is False assert inner["concurrency"] == 32 assert inner["max_turns"] == 30 diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index 1266bfcbe8..c107c03361 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -628,7 +628,7 @@ async def responses(self, request: Request, body=Body()) -> NeMoGymResponse: context = current_capture_context() seen["rollout_id"] = context.rollout_id if context else None seen["model_call_id"] = context.model_call_id if context else None - seen["sink"] = context.sink if context else "no context" + seen["sink"] = context.token_sink if context else "no context" return _training_response("hi") model = _Peek( @@ -663,7 +663,7 @@ def test_external_mode_does_not_mark_a_token_less_response_incomplete(tmp_path): def test_a_committed_call_is_not_marked_even_without_token_ids(tmp_path): """A caller that had the arrays when this process did not has already accounted for the call.""" store = TokenCaptureStore(tmp_path) - context = CaptureContext(rollout_id="cm0-r0", model_call_id="c1", sink=store) + context = CaptureContext(rollout_id="cm0-r0", model_call_id="c1", token_sink=store) token = set_token_sink(context) try: asyncio.run( @@ -851,7 +851,7 @@ def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(instal generation_token_ids=GTOKS, generation_log_probs=LPS, ) - token = set_token_sink(CaptureContext(rollout_id="task0-sink3", model_call_id="mc-1", sink=installed_sink)) + token = set_token_sink(CaptureContext(rollout_id="task0-sink3", model_call_id="mc-1", token_sink=installed_sink)) try: asyncio.run(commit_entry(entry)) finally: From 90cde043cb1e081d66bdff669e793d0eab542e16 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 20 Aug 2026 14:30:05 -0700 Subject: [PATCH 12/12] fix(token-id-capture): harden retirement and agent overrides Signed-off-by: Ananth Subramaniam --- nemo_gym/base_responses_api_agent.py | 2 +- nemo_gym/token_id_capture/store.py | 2 ++ tests/unit_tests/test_base_responses_api_model.py | 8 +++++++- tests/unit_tests/test_token_id_capture.py | 9 ++++++++- 4 files changed, 18 insertions(+), 3 deletions(-) diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 9bebad799f..bf8c91c854 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -172,7 +172,7 @@ def resolve_model_base_url(self, model_server_name: str, rollout_id: Optional[st """Resolve a model-server URL with an optional rollout prefix.""" server_config = get_first_server_config_dict(self.server_client.global_config_dict, model_server_name) base_url = self.server_client._build_server_base_url(server_config) - return f"{apply_rollout_prefix(base_url, rollout_id, token_capture=SimpleResponsesAPIAgent._token_id_capture_enabled(self))}/v1" + return f"{apply_rollout_prefix(base_url, rollout_id, token_capture=self._token_id_capture_enabled())}/v1" # TODO: right now there is no validation on the TypedDict NeMoGymResponseCreateParamsNonStreaming # We should explicitly add validation at this server level or we should explicitly not validate so that there is flexibility in this API. diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index 8735e6d041..d6d81bcc55 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -286,6 +286,8 @@ def _drop(self, rollout_id: str, snapshot_id: str, version: int) -> bool: self.incomplete_path_for(rollout_id).unlink(missing_ok=True) # Keep a frozen tombstone until explicit pre-dispatch cleanup. # A late writer from this attempt must still observe the freeze. + state["indexed_size"] = 0 + state["entry_digests"] = {} state["retired"] = True self._write_state(rollout_id, state) self._fsync_root() diff --git a/tests/unit_tests/test_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 2b1d485c73..0c9c5fb1f0 100644 --- a/tests/unit_tests/test_base_responses_api_model.py +++ b/tests/unit_tests/test_base_responses_api_model.py @@ -803,11 +803,17 @@ def test_base_agent_resolve_model_base_url(monkeypatch): server_client=SimpleNamespace( global_config_dict={}, _build_server_base_url=lambda _config: "http://h:1", - ) + ), + _token_id_capture_enabled=lambda: False, ) assert SimpleResponsesAPIAgent.resolve_model_base_url(agent, "model", "rid") == "http://h:1/ng-rollout/rid/v1" assert SimpleResponsesAPIAgent.resolve_model_base_url(agent, "model", None) == "http://h:1/v1" + agent._token_id_capture_enabled = lambda: True + assert ( + SimpleResponsesAPIAgent.resolve_model_base_url(agent, "model", "rid") + == "http://h:1/ng-rollout/rid/training-token-capture/v1" + ) def _make_base_agent(global_config, *, token_id_capture=False): diff --git a/tests/unit_tests/test_token_id_capture.py b/tests/unit_tests/test_token_id_capture.py index c107c03361..b13f074325 100644 --- a/tests/unit_tests/test_token_id_capture.py +++ b/tests/unit_tests/test_token_id_capture.py @@ -244,7 +244,14 @@ def test_token_store_freeze_is_atomic_and_conditional_drop_is_race_safe(tmp_path assert updated.incomplete assert asyncio.run(store.drop("r0", snapshot_id=updated.snapshot_id, version=updated.version)) assert store.read_entries("r0") == [] - assert orjson.loads(store.state_path_for("r0").read_bytes())["retired"] is True + state = orjson.loads(store.state_path_for("r0").read_bytes()) + assert state["retired"] is True + assert state["indexed_size"] == 0 + assert state["entry_digests"] == {} + retired = asyncio.run(store.freeze("r0")) + assert retired.entries == () + assert retired.snapshot_id == updated.snapshot_id + assert retired.version == updated.version with pytest.raises(RuntimeError, match="already frozen"): asyncio.run(store.put(entry.model_copy(update={"model_call_id": "late-after-drop"})))