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.
`_count_resolution`	-
`_mark_incomplete`	Mark the rollout, or say loudly why it could not be marked.
`capture_health_snapshot`	Return worker-level capture health for metrics endpoints.
`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 durable call intent before dispatch starts generation.
`reset_token_sink`	-
`resolve_parent`	Resolve which recorded call this request continues.
`set_token_sink`	-

Data

_CAPTURE_CONTEXT

_CAPTURE_FAILURES

_RESOLUTION_COUNTS

_RESOLVER_UNAVAILABLE_NOTED

_STATS_LOCK

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,
    lineage_store: nemo_gym.token_id_capture.protocols.LineageStore | None = None,
    model: str = '',
    committed: bool = False,
    delta_records: bool = False,
    prefix_requested: bool = False,
    prefix_supplied: bool = False,
    parent_resolution: nemo_gym.token_id_capture.protocols.LineageResolution | None = None
)

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.

committed

bool = False

delta_records

bool = False

lineage_store

LineageStore | None = None

model

str = ''

model_call_id

str

parent_call_id

str | None

parent_resolution

LineageResolution | None = None

parent_tokens

list[int]

prefix_requested

bool = False

prefix_supplied

bool = False

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._count_resolution(
    status_value: str
) -> None

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_health_snapshot() -> dict

Return worker-level capture health for metrics endpoints.

nemo_gym.token_id_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.

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

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 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.

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

nemo_gym.token_id_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.

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._CAPTURE_FAILURES = [0]

nemo_gym.token_id_capture.sink._RESOLUTION_COUNTS = {'root': 0, 'resolved': 0, 'unresolved': 0}

nemo_gym.token_id_capture.sink._RESOLVER_UNAVAILABLE_NOTED = [False]

nemo_gym.token_id_capture.sink._STATS_LOCK = threading.Lock()

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