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

## Module Contents

### Classes

| Name                                                                                | Description                                                                 |
| ----------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`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_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_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_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.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

Flush pending work and release resources idempotently.

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

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

```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.
Implementations that cannot delete return `True`.
Their owner remains responsible for 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.
No successful writes may occur after it returns.
Entry order carries no meaning.

```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_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_SINK: TokenSink | None = None
```

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