> This page is for version 0.3.0. > For other versions, use one of these documentation indexes: > - Main (default): https://docs.nvidia.com/nemo/gym/main/llms.txt > - 0.7.0: https://docs.nvidia.com/nemo/gym/v0.7.0/llms.txt > - 0.6.0: https://docs.nvidia.com/nemo/gym/v0.6.0/llms.txt > - 0.5.1: https://docs.nvidia.com/nemo/gym/v0.5.1/llms.txt > - 0.5.0: https://docs.nvidia.com/nemo/gym/v0.5.0/llms.txt > - 0.4.0: https://docs.nvidia.com/nemo/gym/v0.4.0/llms.txt > - 0.3.0: https://docs.nvidia.com/nemo/gym/v0.3.0/llms.txt > - 0.2.1: https://docs.nvidia.com/nemo/gym/v0.2.1/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/gym/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/gym/_mcp/server. # nemo_gym.models.token_capture.sink 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. ## Module Contents ### Classes | Name | Description | | ---------------------------------------------------------------------- | ---------------------------------------------- | | [`CaptureContext`](#nemo_gym-models-token_capture-sink-CaptureContext) | Describe one in-flight training-token capture. | ### Functions | Name | Description | | -------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | | [`_admit_candidate`](#nemo_gym-models-token_capture-sink-_admit_candidate) | Select a source-prefix candidate or a root for rewritten history. | | [`_capture_failed`](#nemo_gym-models-token_capture-sink-_capture_failed) | Report a capture failure without letting it reach the model call. | | [`_capture_missing`](#nemo_gym-models-token_capture-sink-_capture_missing) | Mark the rollout when a call this process should have recorded produced nothing. | | [`_count_resolution`](#nemo_gym-models-token_capture-sink-_count_resolution) | - | | [`_mark_incomplete`](#nemo_gym-models-token_capture-sink-_mark_incomplete) | Mark the rollout, or say loudly why it could not be marked. | | [`_note_missing_refuse_call`](#nemo_gym-models-token_capture-sink-_note_missing_refuse_call) | Warn once per worker that the sink cannot resolve a refused call's intent. | | [`_record_external_refusal`](#nemo_gym-models-token_capture-sink-_record_external_refusal) | Write the refusal as a definite ledger failure for an admitted worker-custody call. | | [`capture_engine_refusal`](#nemo_gym-models-token_capture-sink-capture_engine_refusal) | Record that the engine refused the in-flight call before it generated anything. | | [`capture_health_snapshot`](#nemo_gym-models-token_capture-sink-capture_health_snapshot) | Return worker-level capture health for metrics endpoints. | | [`capture_tokens`](#nemo_gym-models-token_capture-sink-capture_tokens) | Record a `TokenEntry` from a complete model response. | | [`commit_entry`](#nemo_gym-models-token_capture-sink-commit_entry) | Durably record a finished entry against the in-flight call. | | [`current_capture_context`](#nemo_gym-models-token_capture-sink-current_capture_context) | Return the capture context for the in-flight call. | | [`mark_external_staging_committed`](#nemo_gym-models-token_capture-sink-mark_external_staging_committed) | Mark the current call as durably recorded by a framework worker. | | [`register_call_intent`](#nemo_gym-models-token_capture-sink-register_call_intent) | Record durable call intent before dispatch starts generation. | | [`reset_token_sink`](#nemo_gym-models-token_capture-sink-reset_token_sink) | - | | [`resolve_parent`](#nemo_gym-models-token_capture-sink-resolve_parent) | Resolve which recorded call this request continues. | | [`set_token_sink`](#nemo_gym-models-token_capture-sink-set_token_sink) | - | ### Data [`NG_CAPTURE_FIELD`](#nemo_gym-models-token_capture-sink-NG_CAPTURE_FIELD) [`NG_COMMIT_COORDS_FIELD`](#nemo_gym-models-token_capture-sink-NG_COMMIT_COORDS_FIELD) [`_CAPTURE_CONTEXT`](#nemo_gym-models-token_capture-sink-_CAPTURE_CONTEXT) [`_CAPTURE_FAILURES`](#nemo_gym-models-token_capture-sink-_CAPTURE_FAILURES) [`_REFUSE_CALL_MISSING_NOTED`](#nemo_gym-models-token_capture-sink-_REFUSE_CALL_MISSING_NOTED) [`_RESOLUTION_COUNTS`](#nemo_gym-models-token_capture-sink-_RESOLUTION_COUNTS) [`_RESOLVER_UNAVAILABLE_NOTED`](#nemo_gym-models-token_capture-sink-_RESOLVER_UNAVAILABLE_NOTED) [`_STATS_LOCK`](#nemo_gym-models-token_capture-sink-_STATS_LOCK) [`logger`](#nemo_gym-models-token_capture-sink-logger) ### API ```python class nemo_gym.models.token_capture.sink.CaptureContext( capture_key: str, model_call_id: str, token_sink: nemo_gym.models.token_capture.protocols.TokenSink | None, lineage_store: nemo_gym.models.token_capture.protocols.LineageResolver | nemo_gym.models.token_capture.protocols.CaptureLedger | None = None, model: str = '', committed: bool = False, refused: bool = False, delta_records: bool = False, prefix_requested: bool = False, prefix_supplied: bool = False, parent_resolution: nemo_gym.models.token_capture.protocols.LineageResolution | None = None, external_staging: bool = False, parent_resolution_mode: nemo_gym.models.token_capture.protocols.ParentResolutionMode = 'fingerprint', admitted_at: float | None = None, capture_admission: nemo_gym.models.token_capture.staging.records.CaptureAdmission | None = None, parent_staging_chain: list[str] = list(), intent_registered: bool = False, parent_chain_hash: str = '', request_items: list[dict] | None = None, external_commit_coords: dict[str, typing.Any] | None = None, external_worker_response_seen: bool = False, admission_hook: collections.abc.Callable[[CaptureContext, CaptureAdmission], nemo_gym.models.token_capture.staging.records.CaptureAdmission] | None = None, publish_on_delivery: bool = False, deferred_commit: nemo_gym.models.token_capture.staging.records.CaptureLedgerCommit | None = None, engine_refusal_recorded: bool = False ) ``` Dataclass 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. Parent resolution runs once for each call. Prefix supply and token capture read the same immutable decision. **`admission_hook`** `Callable[[CaptureContext, CaptureAdmission], CaptureAdmission] | None = None` --- **`admitted_at`** `float | None = None` --- **`capture_admission`** `CaptureAdmission | None = None` --- **`capture_key`** `str` --- **`committed`** `bool = False` --- **`deferred_commit`** `CaptureLedgerCommit | None = None` --- **`delta_records`** `bool = False` --- **`engine_refusal_recorded`** `bool = False` --- **`external_commit_coords`** `dict[str, Any] | None = None` --- **`external_staging`** `bool = False` --- **`external_worker_response_seen`** `bool = False` --- **`intent_registered`** `bool = False` --- **`lineage_store`** `LineageResolver | CaptureLedger | None = None` --- **`model`** `str = ''` --- **`model_call_id`** `str` --- **`parent_call_id`** `str | None` --- **`parent_chain_hash`** `str = ''` --- **`parent_resolution`** `LineageResolution | None = None` --- **`parent_resolution_mode`** `ParentResolutionMode = 'fingerprint'` --- **`parent_staging_chain`** `list[str] = field(default_factory=list)` --- **`parent_tokens`** `list[int]` --- **`prefix_requested`** `bool = False` --- **`prefix_supplied`** `bool = False` --- **`publish_on_delivery`** `bool = False` --- **`refused`** `bool = False` --- **`request_items`** `list[dict] | None = None` --- **`token_sink`** `TokenSink | None` --- ```python nemo_gym.models.token_capture.sink._admit_candidate( context: nemo_gym.models.token_capture.sink.CaptureContext, request_items: list[dict] ) -> None ``` async Select a source-prefix candidate or a root for rewritten history. Match source-prefix evidence, including an older accepted call after a retry. RL still verifies the renderer and converted-message boundary. ```python nemo_gym.models.token_capture.sink._capture_failed( context: nemo_gym.models.token_capture.sink.CaptureContext, stage: str ) -> None ``` async 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. A call that arrives after its rollout was frozen or retired is a late call, not a capture failure. ```python nemo_gym.models.token_capture.sink._capture_missing( context: nemo_gym.models.token_capture.sink.CaptureContext, reason: str ) -> None ``` async 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. ```python nemo_gym.models.token_capture.sink._count_resolution( status_value: str ) -> None ``` ```python nemo_gym.models.token_capture.sink._mark_incomplete( context: nemo_gym.models.token_capture.sink.CaptureContext ) -> bool ``` async 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. Return False when the call is late (its rollout was frozen or retired), so callers can stay quiet. ```python nemo_gym.models.token_capture.sink._note_missing_refuse_call( context: nemo_gym.models.token_capture.sink.CaptureContext ) -> None ``` Warn once per worker that the sink cannot resolve a refused call's intent. ```python nemo_gym.models.token_capture.sink._record_external_refusal( context: nemo_gym.models.token_capture.sink.CaptureContext ) -> None ``` async Write the refusal as a definite ledger failure for an admitted worker-custody call. The ledger failure row carries the call id and `ENGINE_REFUSED_CALL_REASON` only: the staging wire models declare no column for the error code or for the request-time lineage, so a rollout rebuilt from staged records cannot tell a context-window refusal from an output-parameter one and reports no truncation verdict for it (NVIDIA-NeMo/Gym#4647 tracks carrying both). The row still names the call as one that served no completion, which the framework's receipt builder reads through `UNSERVED_CALL_FAILURE_REASONS`. ```python nemo_gym.models.token_capture.sink.capture_engine_refusal( code: str | None ) -> None ``` async Record that the engine refused the in-flight call before it generated anything. `code` is the error code the client received: `context_length_exceeded` for a prompt that does not fit the context window, `None` when the envelope carried no code. On the local capture path the sink's optional `record_engine_refusal` extension stores a `EngineRefusalRecord` and resolves the intent `register_call_intent` wrote, so the rollout freezes complete. The record carries the request-time lineage decision `resolve_parent` made before dispatch (the call the refused request continued), so the builder can tell a refusal that followed the delivered chain from one on another branch. A sink without the extension leaves the intent dangling, and the rollout is masked as a lost call would be; that is logged once per worker. Under worker custody the ledger gets a failure row with `ENGINE_REFUSED_CALL_REASON`. The call never reached the worker's staging, so the outcome is definite, unlike the uncommitted-call row the custody sweep would otherwise write; the sweep skips a call recorded here. This never raises. On the local capture path a refusal that cannot be recorded leaves the intent dangling, which masks the rollout; under worker custody the custody sweep writes the uncommitted-call row for the call instead. Either way the failure lands on the safe side. ```python nemo_gym.models.token_capture.sink.capture_health_snapshot() -> dict ``` Return worker-level capture health for metrics endpoints. ```python nemo_gym.models.token_capture.sink.capture_tokens( response: typing.Any, request_messages: list | None = None ) -> None ``` async 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. ```python nemo_gym.models.token_capture.sink.commit_entry( entry: nemo_gym.models.token_capture.records.TokenEntry, parent_resolution: nemo_gym.models.token_capture.protocols.LineageResolution | None = None ) -> None ``` async 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. ```python nemo_gym.models.token_capture.sink.current_capture_context() -> nemo_gym.models.token_capture.sink.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. ```python nemo_gym.models.token_capture.sink.mark_external_staging_committed( rollout_id: str, model_call_id: str ) -> None ``` Mark the current call as durably recorded by a framework worker. Call this only after the external staging sink has acknowledged the call. Identity validation prevents a delayed or cross-request acknowledgement from suppressing normal capture for a different request. ```python nemo_gym.models.token_capture.sink.register_call_intent() -> None ``` async Record durable call intent before dispatch starts generation. `begin_call` is an optional sink extension. A dangling intent identifies a lost entry. Failure happens before generation and propagates to the caller. The harness can retry without spending inference compute. Sinks without `begin_call` cannot report a missing final entry this way. ```python nemo_gym.models.token_capture.sink.reset_token_sink( token: contextvars.Token ) -> None ``` ```python nemo_gym.models.token_capture.sink.resolve_parent( request_messages: list | None ) -> None ``` async Resolve which recorded call this request continues. Use the request representation received from the harness. Resolve once before dialect conversion or dispatch. Prefix supply and capture then share one parent decision. Return without work for untagged traffic. Every attempted resolution records a root, resolved, or unresolved decision. An unresolved decision includes its reason. For external staging, parent resolution determines whether the worker may capture the call: * A unique parent creates a `token_in` admission. * A request with no prior assistant output creates a `text` root. * An unresolved request may create a `text` root only when the rollout has no ledger rows. * Every other result records a failure and leaves the call unadmitted. An unresolved continuation cannot become a new root. Doing so would train the earlier generated tokens as prompt tokens. ```python nemo_gym.models.token_capture.sink.set_token_sink( context: nemo_gym.models.token_capture.sink.CaptureContext ) -> contextvars.Token ``` ```python nemo_gym.models.token_capture.sink.NG_CAPTURE_FIELD = 'ng_capture' ``` ```python nemo_gym.models.token_capture.sink.NG_COMMIT_COORDS_FIELD = 'ng_commit_coords' ``` ```python nemo_gym.models.token_capture.sink._CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar('nemo_gym_capture_context', default=None) ``` ```python nemo_gym.models.token_capture.sink._CAPTURE_FAILURES = [0] ``` ```python nemo_gym.models.token_capture.sink._REFUSE_CALL_MISSING_NOTED = [False] ``` ```python nemo_gym.models.token_capture.sink._RESOLUTION_COUNTS = {'root': 0, 'resolved': 0, 'unresolved': 0} ``` ```python nemo_gym.models.token_capture.sink._RESOLVER_UNAVAILABLE_NOTED = [False] ``` ```python nemo_gym.models.token_capture.sink._STATS_LOCK = threading.Lock() ``` ```python nemo_gym.models.token_capture.sink.logger = logging.getLogger(__name__) ```