nemo_gym.models.token_capture.protocols
nemo_gym.models.token_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.
A sink may additionally implement
record_engine_refusal(rollout_id, *, model_call_id, code, parent_call_id=None, parent_resolution=None).
It is an optional extension beside begin_call and likewise not part of TokenSink.
record_engine_refusal durably records that the engine refused the call before it generated anything,
together with the request-time lineage decision (the recorded call the refused request continued
and that decision’s status), which the builder reads to place the refusal on one branch.
Such a call can never produce an entry, so the refusal resolves its pre-dispatch intent
instead of leaving it dangling, and the frozen snapshot reports it in engine_refusals.
After the rollout is frozen it must either fail with TokenCaptureFrozenError or leave the
frozen snapshot unchanged, its version and its engine_refusals included: a late refusal
cannot change the sample that was already delivered, so it must not stale the retirement of the
consumed snapshot. After the rollout is retired it must fail with TokenCaptureRetiredError,
like every other write.
nemo_gym.models.token_capture.conformance checks an external implementation against these contracts.
Module Contents
Classes
Functions
Data
API
Bases: ValueError
Refuse a request whose capture state cannot safely admit generation.
Bases: ValueError
Reject request history that capture cannot represent.
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.
Framework-owned capture additionally requires begin_call(rollout_id, model_call_id) and close_admissions(rollout_id). They share one atomic
admission boundary: close returns a manifest and rejects new calls, while
existing intents may still commit. These extensions do not freeze writes.
refuse_call(rollout_id, model_call_id) optionally settles an intent after
the worker acknowledges a refusal before generation or staging. Such calls
remain in the audit ledger but are excluded from the training manifest.
An error status or message alone is not this acknowledgement.
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.
Return whether any ledger row (committed or failed) exists.
A retired rollout raises RolloutRetiredError, so admission can refuse to capture its late calls.
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.
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.
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.
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.
Describe a uniquely verified parent from a shared lineage store.
Return one immutable request-time parent decision.
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.models.token_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.
Release resources. Idempotent.
The store is read-only, so there is never pending work to flush.
Return whether separate model-server workers share committed entries.
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.
Bases: typing.TypedDict
The result of retiring or deleting rollouts, as plain wire data.
It validates as staging.records.RolloutRemoval.
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.
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.
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.
An immutable view of one rollout’s frozen capture records.
engine_refusals holds the calls the engine refused, in the order they were recorded; a sink
without the record_engine_refusal extension leaves it empty.
Receive captured records through Gym’s file store or a framework transport.
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.
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.
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.
Where a trajectory builder freezes, reads, and retires records.
Release resources idempotently.
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.
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.
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.
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.
Set (or clear) the process-wide request-time lineage store.
Set (or clear, with None) the process-wide default sink.
Set (or clear) the caller-owned source in this process.
Gym does not close an installed source.