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

Capture training tokens from one complete model response.

Streaming responses omit token ids from the wire.
The model server still holds the complete response before streaming.
Middleware provides a request-scoped token sink.
The model server passes its complete response to `capture_tokens`.
The sink writes a `TokenEntry`.
Its `model_call_id` joins the corresponding evaluation record.
Untagged traffic has no capture context.

## Module Contents

### Classes

| Name                                                               | Description                                    |
| ------------------------------------------------------------------ | ---------------------------------------------- |
| [`CaptureContext`](#nemo_gym-token_id_capture-sink-CaptureContext) | Describe one in-flight training-token capture. |

### Functions

| Name                                                                                 | Description                                                                      |
| ------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------- |
| [`_capture_failed`](#nemo_gym-token_id_capture-sink-_capture_failed)                 | Report a capture failure without letting it reach the model call.                |
| [`_capture_missing`](#nemo_gym-token_id_capture-sink-_capture_missing)               | Mark the rollout when a call this process should have recorded produced nothing. |
| [`_mark_incomplete`](#nemo_gym-token_id_capture-sink-_mark_incomplete)               | Mark the rollout, or say loudly why it could not be marked.                      |
| [`capture_tokens`](#nemo_gym-token_id_capture-sink-capture_tokens)                   | Record a `TokenEntry` from a complete model response.                            |
| [`commit_entry`](#nemo_gym-token_id_capture-sink-commit_entry)                       | Durably record a finished entry against the in-flight call.                      |
| [`current_capture_context`](#nemo_gym-token_id_capture-sink-current_capture_context) | Return the capture context for the in-flight call.                               |
| [`register_call_intent`](#nemo_gym-token_id_capture-sink-register_call_intent)       | Record that the captured call is about to be dispatched.                         |
| [`reset_token_sink`](#nemo_gym-token_id_capture-sink-reset_token_sink)               | -                                                                                |
| [`set_token_sink`](#nemo_gym-token_id_capture-sink-set_token_sink)                   | -                                                                                |

### Data

[`_CAPTURE_CONTEXT`](#nemo_gym-token_id_capture-sink-_CAPTURE_CONTEXT)

[`logger`](#nemo_gym-token_id_capture-sink-logger)

### API

```python
class nemo_gym.token_id_capture.sink.CaptureContext(
    rollout_id: str,
    model_call_id: str,
    token_sink: nemo_gym.token_id_capture.protocols.TokenSink | None,
    model: str = '',
    committed: bool = False
)
```

Dataclass

Describe one in-flight training-token capture.

The context identifies the rollout and model call.
`token_sink` receives the resulting record.
A framework may provide any `TokenSink` implementation.

**`committed`** `bool = False`

---

**`model`** `str = ''`

---

**`model_call_id`** `str`

---

**`rollout_id`** `str`

---

**`token_sink`** `TokenSink | None`

---

```python
nemo_gym.token_id_capture.sink._capture_failed(
    context: nemo_gym.token_id_capture.sink.CaptureContext,
    stage: str
) -> None
```

async

Report a capture failure without letting it reach the model call.

Bad token payloads must not fail the model call.
Mark the rollout so consumers can mask the sample.
Call this only from an `except` block.

```python
nemo_gym.token_id_capture.sink._capture_missing(
    context: nemo_gym.token_id_capture.sink.CaptureContext,
    reason: str
) -> None
```

async

Mark the rollout when a call this process should have recorded produced nothing.

A response with no token ids is a hole in the chain rather than traffic to skip.
The builder reads the gap between one call's tokens and the next call's prompt as tool output.
A skipped call's generated tokens then enter the next prompt with mask zero.
Policy tokens would train as if the environment produced them.

Two cases are not holes and are left alone.
A committed call was recorded by another capture path.
A context without a sink delegates completeness to external staging.

```python
nemo_gym.token_id_capture.sink._mark_incomplete(
    context: nemo_gym.token_id_capture.sink.CaptureContext
) -> None
```

async

Mark the rollout, or say loudly why it could not be marked.

A missing `mark_incomplete` method can hide incomplete capture.
Log that condition as an error.

```python
nemo_gym.token_id_capture.sink.capture_tokens(
    response: typing.Any
) -> None
```

async

Record a `TokenEntry` from a complete model response.

Accept a Pydantic model or dictionary.
Return without work when no capture context exists.
Mark local capture incomplete when required token ids are absent.
Await the write before the model call returns.

```python
nemo_gym.token_id_capture.sink.commit_entry(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> None
```

async

Durably record a finished entry against the in-flight call.

`capture_tokens` extracts arrays from a served response.
Engine-side capture may already have those arrays.
Engine-side callers can use this method directly.
Return without work when no capture context exists.
Capture failures mark the rollout incomplete.
This method never fails the model call.

```python
nemo_gym.token_id_capture.sink.current_capture_context() -> nemo_gym.token_id_capture.sink.CaptureContext | None
```

Return the capture context for the in-flight call.

Return `None` for untagged traffic.
Framework inference workers use this identity for staged records.

```python
nemo_gym.token_id_capture.sink.register_call_intent() -> None
```

async

Record that the captured call is about to be dispatched.

`begin_call` is an optional sink extension.
It lets a source detect a call whose entry was lost.
A failure happens before generation and must fail the model call.

```python
nemo_gym.token_id_capture.sink.reset_token_sink(
    token: contextvars.Token
) -> None
```

```python
nemo_gym.token_id_capture.sink.set_token_sink(
    context: nemo_gym.token_id_capture.sink.CaptureContext
) -> contextvars.Token
```

```python
nemo_gym.token_id_capture.sink._CAPTURE_CONTEXT: ContextVar[CaptureContext | None] = ContextVar('nemo_gym_capture_context', default=None)
```

```python
nemo_gym.token_id_capture.sink.logger = logging.getLogger(__name__)
```