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

# nemo_automodel.components.speculative.streaming.store

Pluggable data-plane transport for the speculative-training stream.

The :class:`FeatureStore` ABC abstracts where the supervision tensors actually
live -- an in-process dict for build/test, a POSIX shared mount for
multi-node/colocated, or NCCL for GPU-to-GPU in later PRs. The
:class:`SampleRefQueue` reads the store's :meth:`FeatureStore.health` to
decide whether to back off.

The contract is deliberately small (5 methods + 1 property) so PR 2 and
later have an obvious surface to extend:

* :meth:`put` -- produce-side: stash tensors for a sample.
* :meth:`get`  -- consume-side: materialize them; returns a :class:`StoreHandle`
  the consumer must hand to :meth:`release` once it's done with them.
* :meth:`release` -- consume-once: free / drop the materials.
* :meth:`gc` -- sweep partially-released handles (a stale lease, a crashed
  consumer) so the store cannot leak.
* :meth:`health` -- ints only; the queue uses these for backpressure.
* :meth:`close` -- dispose store resources at shutdown.

## Module Contents

### Classes

| Name                                                                                  | Description                                                            |
| ------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`FeatureStore`](#nemo_automodel-components-speculative-streaming-store-FeatureStore) | Abstract transport every feature-store backend implements.             |
| [`StoreHandle`](#nemo_automodel-components-speculative-streaming-store-StoreHandle)   | Opaque token the consumer must return to :meth:`FeatureStore.release`. |
| [`StoreHealth`](#nemo_automodel-components-speculative-streaming-store-StoreHealth)   | Integers-only snapshot of store residency for backpressure decisions.  |

### Functions

| Name                                                                                        | Description                                                        |
| ------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`_next_handle_id`](#nemo_automodel-components-speculative-streaming-store-_next_handle_id) | Mint a fresh :attr:`StoreHandle.handle_id` (module-level counter). |

### Data

[`__all__`](#nemo_automodel-components-speculative-streaming-store-__all__)

[`_handle_id_counter`](#nemo_automodel-components-speculative-streaming-store-_handle_id_counter)

### API

```python
class nemo_automodel.components.speculative.streaming.store.FeatureStore()
```

Abstract

Abstract transport every feature-store backend implements.

Implementations must be safe to call from one thread per operation;
cross-thread concurrency is the queue's responsibility and outside this
contract. The :class:`~nemo_automodel.components.speculative.streaming.queue.SampleRefQueue`
uses :meth:`health` for backpressure; everything else is the
produce/consume pair plus `release` / `gc` for lifecycle.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.close() -> None
```

abstract

Release every resource owned by the store (file handles, NCCL groups, ...).

After `close`, all subsequent calls raise :class:`RuntimeError`.
Idempotent with respect to an already-closed store.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.gc() -> int
```

abstract

Sweep stale entries (e.g. failed releases from a crashed consumer).

Returns the number of entries reclaimed. Called opportunistically by
the :class:`SampleRefQueue` between leases and unconditionally by
:meth:`close`.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.get(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    device: torch.device | str | None = None
) -> tuple[dict[str, torch.Tensor], nemo_automodel.components.speculative.streaming.store.StoreHandle]
```

abstract

Materialize `ref`'s tensors on `device` and hand back a :class:`StoreHandle`.

The returned tensors are detached copies (`clone` for `cuda`
tensors, plain `detach` for CPU views) so prefetch cannot observe
aliasing through :meth:`release`. Returns one tensor per key in
:attr:`SampleRef.feature_keys`, in the same insertion order, so the
consumer can wire them straight into the per-algorithm batch.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth
```

abstract

Return a ints-only :class:`StoreHealth` snapshot for backpressure.

Must not block on tensor I/O (no `.cpu()`, no `.to()`); in-memory
counters are sufficient. Backends that do background I/O MUST serve
:meth:`health` from cached counters, not from the in-flight I/O
thread.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm,
    schema_version: int,
    target_model_version: str,
    draft_weight_version: str,
    num_tokens: int
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

abstract

Stash `tensors` for `sample_id` and return a tensor-free :class:`SampleRef`.

The producer-side metadata (`run_id`, `algorithm`, `schema_version`,
`target_model_version`, `draft_weight_version`, `num_tokens`) is
the data the :class:`SampleRef` carries on the control plane, so the
store must accept it here even though it does not inspect the values
beyond building the ref. Consumers see the ref unchanged.

Implementations must validate that the tensors match what they
declare (`dtype`, `shape` per feature, `numel * element_size`
summed across features == `ref.estimated_bytes`) and reject the
put with a specific exception when they do not, before any partial
write is observable.

```python
nemo_automodel.components.speculative.streaming.store.FeatureStore.release(
    handle: nemo_automodel.components.speculative.streaming.store.StoreHandle
) -> None
```

abstract

Free the resources backing `handle`; idempotent.

After this returns, the sample is no longer present in the store and
a second :meth:`get` on the same :class:`SampleRef` raises
:class:`KeyError`. `gc` retries releases that a previous call
rejected (e.g. transient I/O), so the queue can rely on
`gc + release` being idempotent + retriable.

```python
class nemo_automodel.components.speculative.streaming.store.StoreHandle(
    store: 'FeatureStore',
    sample_id: str,
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    handle_id: int = _next_handle_id()
)
```

Dataclass

Opaque token the consumer must return to :meth:`FeatureStore.release`.

Each :meth:`FeatureStore.get` mints a fresh handle with a unique
:attr:`handle_id`. Two `get` calls on the same sample return two
distinct handles; :meth:`FeatureStore.release` matches against
`handle_id` so releasing one handle twice (or releasing a stale
handle after a sibling has been acquired) cannot decrement a
sibling's outstanding count.

Holds the producing store, the sample id, the originating
:class:`SampleRef`, and the per-`get` handle identity.

**`handle_id`** `int = field(default_factory=_next_handle_id)`

---

**`ref`** `SampleRef`

---

**`sample_id`** `str`

---

**`store`** `'FeatureStore'`

---

```python
class nemo_automodel.components.speculative.streaming.store.StoreHealth(
    resident_bytes: int,
    capacity_bytes: int,
    sample_count: int,
    high_watermark_hit: bool,
    low_watermark_hit: bool
)
```

Dataclass

Integers-only snapshot of store residency for backpressure decisions.

**`capacity_bytes`** `int`

---

**`high_watermark_hit`** `bool`

---

**`low_watermark_hit`** `bool`

---

**`resident_bytes`** `int`

---

**`sample_count`** `int`

---

```python
nemo_automodel.components.speculative.streaming.store._next_handle_id() -> int
```

Mint a fresh :attr:`StoreHandle.handle_id` (module-level counter).

```python
nemo_automodel.components.speculative.streaming.store.__all__ = ['FeatureStore', 'StoreHandle', 'StoreHealth']
```

```python
nemo_automodel.components.speculative.streaming.store._handle_id_counter = itertools.count()
```