nemo_gym.token_id_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`	Describe one in-flight training-token capture.

Functions

Name	Description
`_capture_failed`	Report a capture failure without letting it reach the model call.
`_capture_missing`	Mark the rollout when a call this process should have recorded produced nothing.
`_mark_incomplete`	Mark the rollout, or say loudly why it could not be marked.
`capture_tokens`	Record a `TokenEntry` from a complete model response.
`commit_entry`	Durably record a finished entry against the in-flight call.
`current_capture_context`	Return the capture context for the in-flight call.
`register_call_intent`	Record that the captured call is about to be dispatched.
`reset_token_sink`	-
`set_token_sink`	-

Data

_CAPTURE_CONTEXT

logger

API

class nemo_gym.token_id_capture.sink.CaptureContext(
    rollout_id: str,
    model_call_id: str,
    token_sink: nemo_gym.token_id_capture.protocols.TokenSink | None,
    model: str = '',
    committed: 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.

committed

bool = False

model

str = ''

model_call_id

str

rollout_id

str

token_sink

TokenSink | None

nemo_gym.token_id_capture.sink._capture_failed(
    context: nemo_gym.token_id_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.

nemo_gym.token_id_capture.sink._capture_missing(
    context: nemo_gym.token_id_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.

nemo_gym.token_id_capture.sink._mark_incomplete(
    context: nemo_gym.token_id_capture.sink.CaptureContext
) -> None

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.

nemo_gym.token_id_capture.sink.capture_tokens(
    response: typing.Any
) -> 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.

nemo_gym.token_id_capture.sink.commit_entry(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> 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.

nemo_gym.token_id_capture.sink.current_capture_context() -> nemo_gym.token_id_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.

nemo_gym.token_id_capture.sink.register_call_intent() -> None

async

Record that the captured call is about to be dispatched.

begin_call is an optional sink extension. It lets a source detect a call whose entry was lost. A failure happens before generation and must fail the model call.

nemo_gym.token_id_capture.sink.reset_token_sink(
    token: contextvars.Token
) -> None

nemo_gym.token_id_capture.sink.set_token_sink(
    context: nemo_gym.token_id_capture.sink.CaptureContext
) -> contextvars.Token

nemo_gym.token_id_capture.sink._CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar('nemo_gym_capture_context', default=None)

nemo_gym.token_id_capture.sink.logger = logging.getLogger(__name__)