nemo_gym.token_id_capture.protocols

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.

Module Contents

Classes

Name	Description
`TokenCaptureSnapshot`	An immutable view of one rollout’s frozen capture records.
`TokenSink`	Receive captured records through Gym’s file store or a framework transport.
`TokenSource`	Where a trajectory builder freezes, reads, and retires records.

Functions

Name	Description
`install_token_sink`	Set (or clear, with `None`) the process-wide default sink.
`install_token_source`	Set (or clear) the caller-owned source in this process.
`installed_token_sink`	-
`installed_token_source`	-

Data

_INSTALLED_SINK

_INSTALLED_SOURCE

API

class nemo_gym.token_id_capture.protocols.TokenCaptureSnapshot(
    rollout_id: str,
    entries: tuple[nemo_gym.token_id_capture.records.TokenEntry, ...],
    incomplete: bool,
    snapshot_id: str,
    version: int
)

Dataclass

An immutable view of one rollout’s frozen capture records.

entries

tuple[TokenEntry, ...]

incomplete

bool

rollout_id

str

snapshot_id

str

version

int

class nemo_gym.token_id_capture.protocols.TokenSink()

Protocol

Receive captured records through Gym’s file store or a framework transport.

nemo_gym.token_id_capture.protocols.TokenSink.close() -> None

async

Flush pending work and release resources idempotently.

nemo_gym.token_id_capture.protocols.TokenSink.mark_incomplete(
    rollout_id: str,
    model_call_id: str = ''
) -> None

async

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.

nemo_gym.token_id_capture.protocols.TokenSink.put(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> None

async

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.

class nemo_gym.token_id_capture.protocols.TokenSource()

Protocol

Where a trajectory builder freezes, reads, and retires records.

nemo_gym.token_id_capture.protocols.TokenSource.close() -> None

async

Release resources idempotently.

nemo_gym.token_id_capture.protocols.TokenSource.drop(
    rollout_id: str,
    snapshot_id: str,
    version: int
) -> bool

async

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.

nemo_gym.token_id_capture.protocols.TokenSource.freeze(
    rollout_id: str
) -> nemo_gym.token_id_capture.protocols.TokenCaptureSnapshot

async

Freeze a rollout and return one atomic snapshot.

Freezing is idempotent. No successful writes may occur after it returns. Entry order carries no meaning.

nemo_gym.token_id_capture.protocols.install_token_sink(
    sink: nemo_gym.token_id_capture.protocols.TokenSink | None
) -> None

Set (or clear, with None) the process-wide default sink.

nemo_gym.token_id_capture.protocols.install_token_source(
    source: nemo_gym.token_id_capture.protocols.TokenSource | None
) -> None

Set (or clear) the caller-owned source in this process.

Gym does not close an installed source.

nemo_gym.token_id_capture.protocols.installed_token_sink() -> nemo_gym.token_id_capture.protocols.TokenSink | None

nemo_gym.token_id_capture.protocols.installed_token_source() -> nemo_gym.token_id_capture.protocols.TokenSource | None

nemo_gym.token_id_capture.protocols._INSTALLED_SINK: TokenSink | None = None

nemo_gym.token_id_capture.protocols._INSTALLED_SOURCE: TokenSource | None = None