diff --git a/nemo_gym/base_responses_api_agent.py b/nemo_gym/base_responses_api_agent.py index 3633f562b0..bf8c91c854 100644 --- a/nemo_gym/base_responses_api_agent.py +++ b/nemo_gym/base_responses_api_agent.py @@ -26,8 +26,12 @@ BaseRunRequest, 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.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT +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,12 @@ 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 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 class BaseResponsesAPIAgent(BaseServer): @@ -61,10 +71,11 @@ 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) run = self.run @@ -81,59 +92,87 @@ async def run_with_rollout_context(*args: Any, **kwargs: Any) -> BaseVerifyRespo return app + def _capture_correlation_enabled(self) -> bool: + """Return whether this agent needs rollout correlation. + + 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() + 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. + """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 {} + 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) + ) + 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 model-call capture (observability) is disabled 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._model_call_capture_enabled(): + 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))}{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. + """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, 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. + """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 - 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=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/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 846bb79f06..062181e59c 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 @@ -47,7 +48,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, @@ -67,6 +68,18 @@ BaseServer, SimpleServer, ) +from nemo_gym.token_id_capture import ( + CaptureContext, + capture_tokens, + installed_token_sink, + reset_token_sink, + set_token_sink, +) + +# 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 logger = logging.getLogger(__name__) @@ -90,7 +103,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 +210,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 +243,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 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 def _validate_responses_params(body: dict) -> NeMoGymResponseCreateParamsNonStreaming: @@ -733,7 +760,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]: @@ -1026,10 +1057,27 @@ 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, + token_capture_enabled: bool = False, + ) -> None: self._app = app self._store = store self._model_server_name = model_server_name + # 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 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: if scope.get("type") != "http": @@ -1038,30 +1086,55 @@ 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")} - # 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: + # 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. + # 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: 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 + + # 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, token_sink=token_sink) + ) + + # 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) + finally: + if sink_token is not None: + reset_token_sink(sink_token) + return + request_body = bytearray() async def _receive() -> dict[str, Any]: @@ -1143,6 +1216,10 @@ async def _flush_deferred_response() -> None: finally: await _flush_deferred_response() raise + finally: + # The sink is only needed while the model server produces the response. + if sink_token is not None: + reset_token_sink(sink_token) completed_at = time.time() latency_ms = (time.perf_counter() - start) * 1000.0 @@ -1210,21 +1287,58 @@ 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). + 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 + # 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 + ) + 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: + 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), 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..9e98a060d8 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 = "training-token-capture" diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 5a84fdf328..3c62d52dd0 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -95,6 +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``. +TOKEN_ID_CAPTURE_BLOCK = "token_id_capture" COMPONENT_NAME_KEY_NAME = "component_name" SKIP_VERIFICATION_KEY_NAME = "skip_verification" SKIP_VERIFICATION_REWARD_KEY_NAME = "skip_verification_reward" @@ -129,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, @@ -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" +# 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" AGENT_REF_KEY_NAME = "agent_ref" diff --git a/nemo_gym/rollout_collection.py b/nemo_gym/rollout_collection.py index cbd8501788..3399a6ca7a 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. + # 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 # 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..91874c5708 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,46 @@ _ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_rollout_id", default=None) +# 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 stamped by rollout collection.""" + """Build the capture key for a run request. + + 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 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: + # 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 " + 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 +93,10 @@ 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.""" + # 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[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/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 new file mode 100644 index 0000000000..b522ae240e --- /dev/null +++ b/nemo_gym/token_id_capture/__init__.py @@ -0,0 +1,79 @@ +# 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. + +"""Provide the core training-token capture interfaces. + +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 +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, + TOKEN_FIELDS, + TokenEntry, + extract_token_fields, +) +from nemo_gym.token_id_capture.sink import ( + CaptureContext, + capture_tokens, + commit_entry, + current_capture_context, + 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", + "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 new file mode 100644 index 0000000000..119d06857d --- /dev/null +++ b/nemo_gym/token_id_capture/config.py @@ -0,0 +1,181 @@ +# 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. + +"""Define run-wide training-token capture settings. + +```yaml +env: + nemo_gym: + token_id_capture: + enabled: true + 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. +``` + +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``. +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. +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. +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 + +import logging +from importlib import import_module +from pathlib import Path +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, 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 + # 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. + # 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) + # Rebuild opaque-harness responses from captured records after the run. + rebuild_response: bool = True + + +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: + # 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: + # 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.", + block.dir, + ) + 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: + return self + if not block.rebuild_response: + return self + 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 + + @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. + + 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: + return None + return self._build_endpoint(target, self.token_id_capture.sink_kwargs, TokenSink, "sink") + + @staticmethod + def _build_endpoint(target: str, kwargs: dict[str, Any], protocol: type, kind: str): + if ":" not in target: + 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.{kind} {target!r}: {error}") from error + try: + endpoint = factory(**kwargs) + except TypeError as error: + raise ValueError( + f"could not construct token_id_capture.{kind} {target!r} with {kind}_kwargs={sorted(kwargs)}: {error}" + ) from error + # 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( + f"token_id_capture.{kind} {target!r} does not satisfy {protocol.__name__}: " + f"{', '.join(missing) or 'attribute check failed'}" + ) + return endpoint + + +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..20c7637905 --- /dev/null +++ b/nemo_gym/token_id_capture/protocols.py @@ -0,0 +1,133 @@ +# 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. + +"""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``. +This module avoids FastAPI, Ray, Torch, and aiohttp imports. +""" + +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 view of one rollout's frozen capture records.""" + + rollout_id: str + entries: tuple[TokenEntry, ...] + incomplete: bool + snapshot_id: str + version: int + + +@runtime_checkable +class TokenSink(Protocol): + """Receive captured records through Gym's file store or a framework transport.""" + + 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 must fail. + Writing after the rollout is frozen must fail. + + 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. + 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. + """ + ... + + async def close(self) -> None: + """Flush pending work and release resources idempotently.""" + ... + + +@runtime_checkable +class TokenSource(Protocol): + """Where a trajectory builder freezes, reads, and retires records.""" + + async def freeze(self, rollout_id: str) -> TokenCaptureSnapshot: + """Freeze a rollout and return one atomic snapshot. + + Freezing is idempotent. + No successful writes may occur after it returns. + Entry order carries no meaning. + """ + ... + + async def drop(self, rollout_id: str, *, snapshot_id: str, version: int) -> bool: + """Conditionally retire the exact frozen snapshot that was consumed. + + Return ``False`` if state changed after the snapshot. + Implementations that cannot delete return ``True``. + Their owner remains responsible for retention. + """ + ... + + async def close(self) -> None: + """Release resources idempotently.""" + ... + + +# 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 + + +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 + + +def install_token_source(source: TokenSource | None) -> None: + """Set (or clear) the caller-owned source in this process. + + Gym does not close an installed 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 new file mode 100644 index 0000000000..7f77b3e725 --- /dev/null +++ b/nemo_gym/token_id_capture/records.py @@ -0,0 +1,168 @@ +# 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. + +"""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 + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +# 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") + +# 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): + """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") + + 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 + # Preserve response output items without token arrays. + output_items: list[dict] = Field(default_factory=list) + # This index identifies the item that carried token arrays. + # ``None`` means no item carried them. + token_item_index: int | None = None + # 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": + """Accept older records and reject newer records. + + 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( + 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." + ) + 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 + + +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``. + Wrap each assistant message as a Responses ``message`` item. + """ + 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. + + 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. + """ + indices: list[int] = [] + stripped: list[dict] = [] + for position, item in enumerate(items): + if item.get("generation_token_ids") is not None: + indices.append(position) + stripped.append({key: value for key, value in item.items() if key not in TOKEN_FIELDS}) + 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. + + 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") + for item in response_json.get("output") or []: + 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 any(message.get(field) is not None for field in required): + candidates.append(message) + if not candidates: + return None + 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)}") + 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..53a595cdcb --- /dev/null +++ b/nemo_gym/token_id_capture/sink.py @@ -0,0 +1,224 @@ +# 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. + +"""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 + +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: + """Describe one in-flight training-token capture. + + The context identifies the rollout and model call. + ``token_sink`` receives the resulting record. + A framework may provide any ``TokenSink`` implementation. + """ + + rollout_id: str + model_call_id: str + # ``None`` means another process owns record staging. + # The context still carries the capture identity. + token_sink: TokenSink | None + model: str = "" + # ``commit_entry`` sets this after another capture path records the call. + committed: bool = False + + +_CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar("nemo_gym_capture_context", default=None) + + +def set_token_sink(context: CaptureContext) -> Token: + return _CAPTURE_CONTEXT.set(context) + + +def current_capture_context() -> CaptureContext | None: + """Return the capture context for the in-flight call. + + Return ``None`` for untagged traffic. + Framework inference workers use this identity for staged records. + """ + return _CAPTURE_CONTEXT.get() + + +def reset_token_sink(token: Token) -> None: + _CAPTURE_CONTEXT.reset(token) + + +async def capture_tokens(response: Any) -> None: + """Record a ``TokenEntry`` from a complete model response. + + 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. + """ + context = _CAPTURE_CONTEXT.get() + if context is None: + return + # 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() + elif isinstance(response, dict): + payload = response + else: + 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(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=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"], + routed_experts=info.get("routed_experts"), + # Preserve content for text-based training penalties. + output_items=content_items, + token_item_index=token_item_index, + created_at=time.time(), + ) + except Exception: + await _capture_failed(context, "build") + return + await commit_entry(entry) + + +async def commit_entry(entry: TokenEntry) -> None: + """Durably record a finished entry against the in-flight call. + + ``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. + """ + context = _CAPTURE_CONTEXT.get() + if context is None: + return + 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.", + context.model_call_id, + context.rollout_id, + ) + await _mark_incomplete(context) + return + if context.token_sink is None: + context.committed = True + return + try: + await context.token_sink.put(entry) + context.committed = True + except Exception: + await _capture_failed(context, "write") + + +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. + 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.", + stage, + context.model_call_id, + context.rollout_id, + exc_info=True, + ) + await _mark_incomplete(context) + + +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. + The builder reads the gap between one call's tokens and the next call's prompt as tool output. + 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 committed call was recorded by another capture path. + A context without a sink delegates completeness to external staging. + """ + 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.", + context.model_call_id, + context.rollout_id, + reason, + ) + await _mark_incomplete(context) + + +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(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(context.token_sink).__name__, + context.rollout_id, + ) + return + try: + await mark(context.rollout_id, context.model_call_id) + except Exception: + 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 new file mode 100644 index 0000000000..d6d81bcc55 --- /dev/null +++ b/nemo_gym/token_id_capture/store.py @@ -0,0 +1,342 @@ +# 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. + +"""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 + +import asyncio +import fcntl +import hashlib +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 + + +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 state_path_for(self, rollout_id: str) -> Path: + return self._root / f"{validate_rollout_id(rollout_id)}.tokens.state.json" + + 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 { + "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: + 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: + 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: + with self._locked(rollout_id, shared=True): + return bool(self._read_state(rollout_id).get("incomplete", False)) + + 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("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 + 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()) + 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) + + # The file store is Gym's default TokenSink and TokenSource. + # A framework can replace it without changing the capture path. + # + # Both interfaces offload blocking work to the process-wide default thread pool. + + async def put(self, entry: TokenEntry) -> None: + """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: + return await asyncio.to_thread(self.freeze_now, rollout_id) + + 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) + 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)), + snapshot_id=str(state["snapshot_id"]), + version=int(state["version"]), + ) + + async def tokens_for(self, rollout_id: str) -> list[TokenEntry]: + """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: + """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: + with self._locked(rollout_id): + state = self._read_state(rollout_id) + if ( + not state.get("frozen", False) + or state.get("snapshot_id") != snapshot_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) + # 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() + return True + + async def close(self) -> None: + """The file store owns no persistent handles.""" + + def delete(self, rollout_id: str) -> None: + """Unconditionally remove a rollout's records. + + 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) + 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: + for line in handle: + stripped = line.strip() + if stripped: + entries.append(TokenEntry.model_validate(orjson.loads(stripped))) + return entries + + +def make_token_store(global_config_dict: Any) -> TokenCaptureStore | None: + """Build the training-token file store. + + 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 + + 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/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..5587dae04d 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: @@ -1030,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_base_responses_api_model.py b/tests/unit_tests/test_base_responses_api_model.py index 358c9bedfb..0c9c5fb1f0 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): @@ -801,14 +803,20 @@ 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): +def _make_base_agent(global_config, *, token_id_capture=False): from nemo_gym.base_responses_api_agent import BaseResponsesAPIAgentConfig class _Agent(SimpleResponsesAPIAgent): @@ -820,7 +828,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 +856,33 @@ 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/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({}) @@ -851,12 +892,22 @@ 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/training-token-capture/v1/responses"), + ) + assert ( + agent.url_path_for_request("/v1/responses", capture_prefixed) + == "/ng-rollout/7-0/training-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}}/training-token-capture/v1/responses" in routes assert "/v1/responses" in routes @@ -1087,6 +1138,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} + # 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 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 requires no indices. + 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 + # 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) + + +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. + # 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") + 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..8ce96be4f7 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)}, + ) + + # 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"}, + 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"] + # 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: 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..b13f074325 --- /dev/null +++ b/tests/unit_tests/test_token_id_capture.py @@ -0,0 +1,1189 @@ +# 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. +"""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 +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 orjson +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, + capture_tokens, + commit_entry, + current_capture_context, + 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 + + +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_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( + rollout_id="r0", + model_call_id="c0", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=[-0.1], + ) + + +# --- 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") == [] + + +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_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", + 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.freeze("r0")) + assert snapshot.entries == (entry,) + assert asyncio.run(store.freeze("r0")) == snapshot + + 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", snapshot_id=snapshot.snapshot_id, version=snapshot.version)) + updated = asyncio.run(store.freeze("r0")) + assert updated.incomplete + assert asyncio.run(store.drop("r0", snapshot_id=updated.snapshot_id, version=updated.version)) + assert store.read_entries("r0") == [] + 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"}))) + + store.delete("r0") + replacement = entry.model_copy(update={"model_call_id": "replacement"}) + asyncio.run(store.put(replacement)) + assert store.read_entries("r0") == [replacement] + + +# --- config ------------------------------------------------------------------- + + +def _block(**kwargs) -> dict: + return {"token_id_capture": {"enabled": True, "rebuild_response": False, **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): + """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))) + 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): + """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): + """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 + 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"}}) + + +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, "rebuild_response": False, key: value}} + ) + + +# --- 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/training-token-capture/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/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. + 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): + """Store token arrays once and preserve response content. + + 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"}) + 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/training-token-capture/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/training-token-capture/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/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. + 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"}) + 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(): + """Keep token capture independent of Gym's server stack. + + 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 = ( + 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): + """Capture tokens before streaming Anthropic messages. + + 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( + "POST", + "/ng-rollout/stream0-roll0/training-token-capture/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): + """Mark a rollout incomplete when capture loses a call. + + A bad payload must not break the model call. + Consumers must still detect the missing record. + """ + 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/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") + + +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): + """Treat missing token ids as an incomplete capture. + + 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"}) + # 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 _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): + """Create capture identity without a local destination. + + Framework inference workers use the identity minted here. + Local destination availability must not gate that identity. + """ + 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.token_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/training-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): + """Leave completeness to external staging without a local destination. + + 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 ( + client.post("/ng-rollout/ext1-r0/training-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", token_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()) + 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( + TokenEntry( + rollout_id="gone-0", + model_call_id="c", + prompt_token_ids=PTOKS, + generation_token_ids=GTOKS, + generation_log_probs=LPS, + ) + ) + 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() + 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): + """Keep concurrent writers from interleaving partial records. + + The exclusive file lock covers threads and processes. + """ + 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: + """Implement ``TokenSink`` without a file store. + + Framework transports may keep no local files. + The capture path must accept this protocol-only implementation. + """ + + 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) + + 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(): + 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, "rebuild_response": False}} + client = TestClient(_server(config).setup_webserver()) + 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 + 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_allows_capture_with_no_destination_at_all(): + """Allow external staging without a local store. + + This process still resolves the capture identity. + """ + 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): + """Send incomplete state through the ``TokenSink`` protocol. + + Capture code must not require concrete store attributes. + """ + + async def boom(entry): + raise RuntimeError("transport down") + + 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/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])] + + +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, "rebuild_response": False}}).setup_webserver() + ) + with caplog.at_level(logging.ERROR): + 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: + install_token_sink(None) + + +def test_commit_entry_records_a_call_with_no_token_fields_on_the_response(installed_sink): + """Allow engine-side capture to commit an existing entry. + + Engine-side callers already have the token arrays. + They should share the standard durability path. + """ + 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", token_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): + """Guard record construction failures. + + ``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 + + 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, "rebuild_response": False}}).setup_webserver() + ) + 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" + 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): + """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): + """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", + 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): + """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", + 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") + asyncio.run(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) + + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + async def close(self) -> None: + pass + + +class _NotASink: + async def put(self, entry) -> None: + pass + + async def close(self) -> None: + pass + + +class _NotCallableSink: + put = "not a method" + + async def mark_incomplete(self, rollout_id: str, model_call_id: str = "") -> None: + pass + + async def close(self) -> 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 + + 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, + "rebuild_response": False, + "sink": f"{__name__}:_ConfiguredSink", + } + } + client = TestClient(_server(config).setup_webserver()) + + 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 + + +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, + "rebuild_response": False, + "sink": f"{__name__}:_ConfiguredSink", + } + } + client = TestClient(_server(config).setup_webserver()) + + 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 == [] + + +def test_a_sink_receives_its_configured_kwargs(): + """Require explicit constructor wiring for a transport sink.""" + 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(): + """Reject sinks that cannot mark incomplete capture.""" + config = TokenIdCaptureConfig.model_validate( + { + "token_id_capture": { + "enabled": True, + "rebuild_response": False, + "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(): + """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": { + "enabled": True, + "rebuild_response": False, + "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, "rebuild_response": False, "sink": target}} + ) + with pytest.raises(ValueError, match=expected): + config.build_sink() + + +def test_a_programmatically_installed_sink_does_not_reach_a_spawned_worker(): + """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() + 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): + """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) + + 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 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 + + +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(): + """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(): + """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): + """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")) + 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")