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.

The serving path awaits TokenSink.put before returning the model response. When put returns, the record must be durable and visible to every worker’s lineage resolver. The harness can then send a continuation to any worker without racing publication of the previous call. A transport that returns before cross-client visibility can produce intermittent unresolved samples under load.

A sink may additionally implement begin_call(rollout_id, model_call_id). It is an optional extension and deliberately not part of the TokenSink protocol. begin_call durably records a pre-dispatch intent. An intent with no matching entry at freeze must mask the rollout. That closes the window where the final call’s entry is lost without a trace. begin_call runs before generation, so the caller may fail the model call at zero compute cost.

nemo_gym.token_id_capture.conformance checks an external implementation against these contracts.

Module Contents

Classes

Name	Description
`CaptureLedger`	Store metadata for calls whose token data is staged externally.
`LineageMatch`	Describe a uniquely verified parent from a shared lineage store.
`LineageResolution`	Return one immutable request-time parent decision.
`LineageResolver`	Resolve request-time lineage from entries committed by a token sink.
`RolloutRemovalPayload`	The result of retiring or deleting rollouts, as plain wire data.
`RolloutRetiredError`	Reject reading the ledger of a rollout that was retired.
`TokenCaptureFrozenError`	Reject a write that arrived after its rollout’s capture was frozen.
`TokenCaptureRetiredError`	Reject a write that arrived after its rollout’s capture was retired.
`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_lineage_store`	Set (or clear) the process-wide request-time lineage store.
`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_lineage_store`	-
`installed_token_sink`	-
`installed_token_source`	-

Data

_INSTALLED_LINEAGE_STORE

_INSTALLED_SINK

_INSTALLED_SOURCE

API

class nemo_gym.token_id_capture.CaptureLedger()

Protocol

Bases: LineageResolver

Store metadata for calls whose token data is staged externally.

record publishes a successfully staged call for parent resolution. record_failure records a call that did not commit. manifest returns both kinds of rows without including token arrays. retire and delete remove the rows of rollouts the framework no longer needs. Records must be visible to every serving worker before record returns.

nemo_gym.token_id_capture.CaptureLedger.delete(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove the ledgers and fences of rollouts, as TokenCaptureStore.delete does.

Delete retired rollouts once nothing of those attempts can still write, for example after every Gym server has restarted or at the end of the run, which removes their fences. Delete a rollout ID before reusing it, so the new execution starts from an empty ledger. Deleting again is a no-op. Staged token data is untouched: the framework that staged it also removes it.

nemo_gym.token_id_capture.CaptureLedger.has_rows(
    rollout_id: str
) -> bool

async

Return whether any ledger row (committed or failed) exists.

A retired rollout raises RolloutRetiredError, so admission can refuse to capture its late calls.

nemo_gym.token_id_capture.CaptureLedger.manifest(
    rollout_id: str
) -> dict

async

Return the rollout’s token-free ledger as plain wire data.

The shape validates as staging.records.RolloutManifest: committed rows under records (each a CallRecord payload) and poison rows under failures. Cumulative token IDs never appear in the manifest. A retired rollout raises RolloutRetiredError instead of returning an empty manifest.

nemo_gym.token_id_capture.CaptureLedger.record(
    commit: nemo_gym.token_id_capture.staging.records.CaptureLedgerCommit
) -> None

async

Publish a completed call for later request-time resolution.

Repeating a model call ID with the same payload is a no-op. Reusing a model call ID with different data must fail. Return only after every serving worker can read the record.

nemo_gym.token_id_capture.CaptureLedger.record_failure(
    rollout_id: str,
    model_call_id: str,
    reason: str
) -> None

async

Record a call whose token capture did not complete.

Failure rows do not participate in parent resolution. manifest includes them under failures so finalization rejects the rollout.

nemo_gym.token_id_capture.CaptureLedger.retire(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove the ledgers of rollouts the framework is done with, and fence them.

Retire a finished rollout once its manifest has been read and the receipt built from it, and an abandoned rollout as soon as the framework gives up on it. Like TokenSource.drop on the complete-record store, retiring keeps a fence: later record and record_failure calls for the rollout are discarded instead of starting a new ledger. Even a finished rollout can still write: a client retry can run it twice, and the second execution writes after the first one returned. Retiring again is a no-op, so retrying a batch is safe. Staged token data is untouched: the framework that staged it also removes it.

class nemo_gym.token_id_capture.LineageMatch(
    model_call_id: str,
    cumulative_token_ids: tuple[int, ...],
    digest: str,
    staging_chain: tuple[str, ...] = (),
    prev_len: int = 0,
    chain_hash: str = ''
)

Dataclass

Describe a uniquely verified parent from a shared lineage store.

chain_hash

str = ''

cumulative_token_ids

tuple[int, ...]

digest

str

model_call_id

str

prev_len

int = 0

staging_chain

tuple[str, ...] = ()

class nemo_gym.token_id_capture.LineageResolution(
    status: nemo_gym.token_id_capture.records.ParentResolutionStatus,
    match: nemo_gym.token_id_capture.protocols.LineageMatch | None = None,
    reason: str = ''
)

Dataclass

Return one immutable request-time parent decision.

match

LineageMatch | None = None

reason

str = ''

status

ParentResolutionStatus

nemo_gym.token_id_capture.LineageResolution.__post_init__() -> None

class nemo_gym.token_id_capture.LineageResolver()

Protocol

Resolve request-time lineage from entries committed by a token sink.

This is a read-only view over sink-committed records. After TokenSink.put returns, a later resolve on any worker must see the entry. Visibility is required per rollout key only, so the store may shard by rollout. Implementations must never guess among candidates. resolve must return UNRESOLVED when it cannot prove one parent. The offline builder independently re-verifies every claimed link by digest. External implementations should embed Gym’s RolloutLineage matcher rather than reimplement the hashing. nemo_gym.token_id_capture.lineage is importable without the server stack. Callers of commit_entry outside capture_tokens must run stamp_continuation themselves. Unstamped entries are invisible to resolution.

nemo_gym.token_id_capture.LineageResolver.close() -> None

async

Release resources. Idempotent.

The store is read-only, so there is never pending work to flush.

nemo_gym.token_id_capture.LineageResolver.is_process_shared() -> bool

Return whether separate model-server workers share committed entries.

nemo_gym.token_id_capture.LineageResolver.resolve(
    rollout_id: str,
    request_items: list[dict]
) -> nemo_gym.token_id_capture.protocols.LineageResolution

async

Return whether the request is a root, resolved, or unresolved.

request_items are the unmodified harness items. The implementation must verify the recorded request context. A conflicting set of committed payloads for one call id must count as zero candidates. UNRESOLVED is always a safe answer; a wrong RESOLVED is caught later by digest verification.

class nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

Bases: typing.TypedDict

The result of retiring or deleting rollouts, as plain wire data.

It validates as staging.records.RolloutRemoval.

absent

list[str]

removed

list[str]

class nemo_gym.token_id_capture.protocols.RolloutRetiredError()

Bases: LookupError

Reject reading the ledger of a rollout that was retired.

A retired rollout’s ledger is gone, so an empty manifest would look like a rollout that made no calls. Reading it is a framework error, such as fetching a manifest after retiring the rollout. Deleting the rollout removes its fence, after which the ID reads as unused again.

class nemo_gym.token_id_capture.TokenCaptureFrozenError()

Bases: RuntimeError

Reject a write that arrived after its rollout’s capture was frozen.

Freezing seals a rollout’s capture verdict. A model call can still be in flight at that point. Its harness may have been killed at a timeout backstop. The sink fails the late write so no successful write follows a freeze. The caller drops the late record and must not mark the rollout incomplete. The freeze already judged completeness from the durable intent ledger. A post-freeze mark would mutate the consumed snapshot and break its retirement.

class nemo_gym.token_id_capture.TokenCaptureRetiredError()

Bases: TokenCaptureFrozenError

Reject a write that arrived after its rollout’s capture was retired.

A retired rollout stays fenced until it is deleted, and every later writer sees it as frozen, so callers handle this exactly like TokenCaptureFrozenError: drop the late record and do not mark the rollout incomplete. Typical sources are an abandoned attempt that keeps running and a duplicate execution of a finished rollout.

class nemo_gym.token_id_capture.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.TokenSink()

Protocol

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

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

async

Release resources. Idempotent.

There is never buffered unwritten data here. put guaranteed durability before it returned. A close that must flush records means put broke its contract.

nemo_gym.token_id_capture.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. It must succeed after freeze. It must change the observable version; that is what invalidates a stale retirement. Make it more available than put, for example through a local spill. put and mark_incomplete failing together is the silent-loss case.

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

async

Durably store one record before returning.

The entry carries its continuation lookup metadata. Durability and resolver visibility are the return condition, not an eventual goal. Any worker’s paired lineage resolver must see the entry once this method returns. Repeating the same call id with the same payload is a no-op. “Same payload” means the identical serialized entry, byte-for-byte, timestamps included. A retry must resend the same bytes; rebuilding the entry produces a conflict, not a retry. Reusing a call id with a different payload must fail. A transport without compare-and-swap may delegate that conflict to the reader. A resolver must then treat conflicting committed payloads for one call id as zero candidates. Writing after the rollout is frozen must fail with TokenCaptureFrozenError; see TokenSource.freeze.

This method may raise. The caller marks the rollout incomplete. A capture error never fails the model call.

class nemo_gym.token_id_capture.TokenSource()

Protocol

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

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

async

Release resources idempotently.

nemo_gym.token_id_capture.TokenSource.delete(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove rollouts’ records and fences.

Delete retired rollouts once nothing of those attempts can still write, or a rollout ID before reusing it. Deleting again is a no-op.

nemo_gym.token_id_capture.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. Transports without delete return True and own retention.

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

async

Freeze a rollout and return one atomic snapshot.

Freezing is idempotent. Entry order carries no meaning. Entries are unique per model_call_id. An at-least-once transport must dedupe identical copies before snapshotting. The fence is relaxed: “no successful write after freeze” is not required cluster-wide. A strict fence is unimplementable without compare-and-swap. A write racing freeze may therefore succeed durably. It must then bump the version. A conditional drop of the consumed snapshot then fails, and the evidence is retained. Attempt-scoped rollout ids are the sanctioned strategy for retirement without compare-and-swap.

nemo_gym.token_id_capture.TokenSource.retire(
    rollout_ids: collections.abc.Sequence[str]
) -> nemo_gym.token_id_capture.protocols.RolloutRemovalPayload

async

Remove rollouts’ records and keep a fence, whatever their state.

This is drop without the snapshot check, for rollouts the consumer is done with but won’t drop as a consumed snapshot, such as masked or failed captures and abandoned attempts. Later writes for a retired rollout must fail with TokenCaptureRetiredError. Retiring again is a no-op.

nemo_gym.token_id_capture.install_lineage_store(
    store: nemo_gym.token_id_capture.protocols.LineageResolver | None
) -> None

Set (or clear) the process-wide request-time lineage store.

nemo_gym.token_id_capture.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.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.installed_lineage_store() -> nemo_gym.token_id_capture.protocols.LineageResolver | None

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

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

nemo_gym.token_id_capture.protocols._INSTALLED_LINEAGE_STORE: LineageResolver | None = None

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

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