> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/gym/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/gym/_mcp/server.

# 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                                                                 |
| ----------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`LineageMatch`](#nemo_gym-token_id_capture-protocols-LineageMatch)                 | Describe a uniquely verified parent from a shared lineage store.            |
| [`LineageResolution`](#nemo_gym-token_id_capture-protocols-LineageResolution)       | Return one immutable request-time parent decision.                          |
| [`LineageStore`](#nemo_gym-token_id_capture-protocols-LineageStore)                 | Resolve request-time lineage from entries committed by a token sink.        |
| [`TokenCaptureSnapshot`](#nemo_gym-token_id_capture-protocols-TokenCaptureSnapshot) | An immutable view of one rollout's frozen capture records.                  |
| [`TokenSink`](#nemo_gym-token_id_capture-protocols-TokenSink)                       | Receive captured records through Gym's file store or a framework transport. |
| [`TokenSource`](#nemo_gym-token_id_capture-protocols-TokenSource)                   | Where a trajectory builder freezes, reads, and retires records.             |

### Functions

| Name                                                                                      | Description                                                 |
| ----------------------------------------------------------------------------------------- | ----------------------------------------------------------- |
| [`install_lineage_store`](#nemo_gym-token_id_capture-protocols-install_lineage_store)     | Set (or clear) the process-wide request-time lineage store. |
| [`install_token_sink`](#nemo_gym-token_id_capture-protocols-install_token_sink)           | Set (or clear, with `None`) the process-wide default sink.  |
| [`install_token_source`](#nemo_gym-token_id_capture-protocols-install_token_source)       | Set (or clear) the caller-owned source in this process.     |
| [`installed_lineage_store`](#nemo_gym-token_id_capture-protocols-installed_lineage_store) | -                                                           |
| [`installed_token_sink`](#nemo_gym-token_id_capture-protocols-installed_token_sink)       | -                                                           |
| [`installed_token_source`](#nemo_gym-token_id_capture-protocols-installed_token_source)   | -                                                           |

### Data

[`_INSTALLED_LINEAGE_STORE`](#nemo_gym-token_id_capture-protocols-_INSTALLED_LINEAGE_STORE)

[`_INSTALLED_SINK`](#nemo_gym-token_id_capture-protocols-_INSTALLED_SINK)

[`_INSTALLED_SOURCE`](#nemo_gym-token_id_capture-protocols-_INSTALLED_SOURCE)

### API

```python
class nemo_gym.token_id_capture.protocols.LineageMatch(
    model_call_id: str,
    cumulative_token_ids: tuple[int, ...],
    digest: str
)
```

Dataclass

Describe a uniquely verified parent from a shared lineage store.

**`cumulative_token_ids`** `tuple[int, ...]`

---

**`digest`** `str`

---

**`model_call_id`** `str`

---

```python
class nemo_gym.token_id_capture.protocols.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`

---

```python
nemo_gym.token_id_capture.protocols.LineageResolution.__post_init__() -> None
```

```python
class nemo_gym.token_id_capture.protocols.LineageStore()
```

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.

```python
nemo_gym.token_id_capture.protocols.LineageStore.close() -> None
```

async

Release resources. Idempotent.

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

```python
nemo_gym.token_id_capture.protocols.LineageStore.is_process_shared() -> bool
```

Return whether separate model-server workers share committed entries.

```python
nemo_gym.token_id_capture.protocols.LineageStore.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.

```python
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`

---

```python
class nemo_gym.token_id_capture.protocols.TokenSink()
```

Protocol

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

```python
nemo_gym.token_id_capture.protocols.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.

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

```python
nemo_gym.token_id_capture.protocols.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 freeze must fail or bump the frozen version; see `TokenSource.freeze`.

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

```python
class nemo_gym.token_id_capture.protocols.TokenSource()
```

Protocol

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

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

async

Release resources idempotently.

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

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

```python
nemo_gym.token_id_capture.protocols.install_lineage_store(
    store: nemo_gym.token_id_capture.protocols.LineageStore | None
) -> None
```

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

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

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

```python
nemo_gym.token_id_capture.protocols.installed_lineage_store() -> nemo_gym.token_id_capture.protocols.LineageStore | None
```

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

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

```python
nemo_gym.token_id_capture.protocols._INSTALLED_LINEAGE_STORE: LineageStore | None = None
```

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

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