> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

Tensor-free control-plane contracts for the speculative-training stream.

The streaming pipeline (per the EAGLE-3 / DFlash / DSpark train-inference
disaggregation RFC, issue #3062 PR 1) splits every transferred sample into a
`refs` (control-plane, no tensors) and a `store` (data-plane, holds
the supervision tensors). The queue carries refs only; tensors live in the
store and are referenced by `nemo_automodel.components.speculative.streaming.refs.SampleRef.feature_keys`.

This package owns:

* `SampleRef` and
  `FeatureSpec`
  \-- the frozen, tensor-free reference carried on every control-plane hop.
* `assert_no_tensors`
  \-- the guard that enforces the no-tensor invariant on the control plane.
* `FeatureStore`
  \-- the pluggable data-plane transport (local dict, shared POSIX mount, NCCL,
  ...); `LocalFeatureStore`
  \-- the in-process implementation land-tested by PR 1.
* `SampleRefQueue`
  \-- the metadata-only lease/ack/fail queue between producers and consumers
  with visibility-timeout reclaim and watermark-based backpressure.
* `FeatureProducer`
  and `FeatureDataLoader`
  \-- the EAGLE-3 produce and consume sides that turn a target forward into
  refs and refs back into `Eagle3TargetBatch` instances.
* `AsyncFeaturePipeline`
  \-- runs the producer in a background thread so target forward and draft
  backward overlap; `SharedDirFeatureStore`
  adds a POSIX shared-mount data plane for cross-process producer/consumer.

## Subpackages

* **[`nemo_automodel.components.speculative.streaming.stores`](/nemo-automodel/nemo_automodel/components/speculative/streaming/stores)**

## Submodules

* **[`nemo_automodel.components.speculative.streaming.async_pipeline`](/nemo-automodel/nemo_automodel/components/speculative/streaming/async_pipeline)**
* **[`nemo_automodel.components.speculative.streaming.eagle3`](/nemo-automodel/nemo_automodel/components/speculative/streaming/eagle3)**
* **[`nemo_automodel.components.speculative.streaming.loader`](/nemo-automodel/nemo_automodel/components/speculative/streaming/loader)**
* **[`nemo_automodel.components.speculative.streaming.producer`](/nemo-automodel/nemo_automodel/components/speculative/streaming/producer)**
* **[`nemo_automodel.components.speculative.streaming.queue`](/nemo-automodel/nemo_automodel/components/speculative/streaming/queue)**
* **[`nemo_automodel.components.speculative.streaming.refs`](/nemo-automodel/nemo_automodel/components/speculative/streaming/refs)**
* **[`nemo_automodel.components.speculative.streaming.store`](/nemo-automodel/nemo_automodel/components/speculative/streaming/store)**

## Package Contents

### Classes

| Name                                                                                                                | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`AsyncFeaturePipeline`](#nemo_automodel-components-speculative-streaming-async_pipeline-AsyncFeaturePipeline)      | Run a `FeatureProducer` in a background thread, draining a prompt source. |
| [`FeatureAlgorithm`](#nemo_automodel-components-speculative-streaming-refs-FeatureAlgorithm)                        | Which speculative-decoding draft family produced this sample.             |
| [`FeatureDataLoader`](#nemo_automodel-components-speculative-streaming-loader-FeatureDataLoader)                    | Iterator over `Eagle3TargetBatch` materialized from a streaming queue.    |
| [`FeatureProducer`](#nemo_automodel-components-speculative-streaming-producer-FeatureProducer)                      | Run a target backend once and put the supervision into a `FeatureStore`.  |
| [`FeatureSpec`](#nemo_automodel-components-speculative-streaming-refs-FeatureSpec)                                  | Shape + dtype metadata for one named feature in a `SampleRef`.            |
| [`FeatureStore`](#nemo_automodel-components-speculative-streaming-store-FeatureStore)                               | Abstract transport every feature-store backend implements.                |
| [`Lease`](#nemo_automodel-components-speculative-streaming-queue-Lease)                                             | Handle to a leased `SampleRef`.                                           |
| [`LocalFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-local-LocalFeatureStore)              | In-process `FeatureStore` implementation.                                 |
| [`PromptSource`](#nemo_automodel-components-speculative-streaming-async_pipeline-PromptSource)                      | Zero-argument callable that supplies the next prompt batch.               |
| [`SampleRef`](#nemo_automodel-components-speculative-streaming-refs-SampleRef)                                      | Tensor-free reference to one produced sample in a feature store.          |
| [`SampleRefQueue`](#nemo_automodel-components-speculative-streaming-queue-SampleRefQueue)                           | Lease / ack / fail queue over `SampleRef`.                                |
| [`SharedDirFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-SharedDirFeatureStore) | `FeatureStore` backed by one `<sample_id>.safetensors` per sample.  |
| [`StoreHandle`](#nemo_automodel-components-speculative-streaming-store-StoreHandle)                                 | Opaque token the consumer must return to `FeatureStore.release`.          |
| [`StoreHealth`](#nemo_automodel-components-speculative-streaming-store-StoreHealth)                                 | Integers-only snapshot of store residency for backpressure decisions.     |
| [`VisibilityTimeout`](#nemo_automodel-components-speculative-streaming-queue-VisibilityTimeout)                     | How long an unacked `Lease` is allowed to live before reclaim.            |

### Functions

| Name                                                                                           | Description                                         |
| ---------------------------------------------------------------------------------------------- | --------------------------------------------------- |
| [`assert_no_tensors`](#nemo_automodel-components-speculative-streaming-refs-assert_no_tensors) | Recursively validate that `obj` carries no tensors. |

### API

```python
class nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline(
    producer: nemo_automodel.components.speculative.streaming.producer.FeatureProducer,
    queue: nemo_automodel.components.speculative.streaming.queue.SampleRefQueue,
    prompt_source: nemo_automodel.components.speculative.streaming.async_pipeline.PromptSource | typing.Iterator[nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | torch.Tensor],
    poll_interval: float = 0.1,
    stop_on_exhausted: bool = True
)
```

Run a `FeatureProducer` in a background thread, draining a prompt source.

**Lifecycle**

Construct, then `start` (or use the context manager).
The background thread is daemon; `stop` joins it within
`join_timeout` seconds. Outstanding leases are the trainer's
responsibility -- the pipeline does not drop them on close.

**Parameters:**

**`producer`** `FeatureProducer`

The `FeatureProducer` to invoke on each prompt.
It carries the wrapped target backend, the store, and the
per-call metadata.

---

**`queue`** `SampleRefQueue`

The `SampleRefQueue` to push the resulting
`SampleRef` onto. The queue's HWM/LWM hysteresis
paces the producer against the consumer.

---

**`prompt_source`** `PromptSource | Iterator[_PromptBatch | _PackedPromptBatch | torch.Tensor]`

A `PromptSource` -- either a zero-arg
callable or an `Iterator`. The callable is invoked
from the background thread; it must be thread-safe (e.g.
a `torch.utils.data.DataLoader` iterator is).

---

**`poll_interval`** `float` — default: 0.1

Seconds between `prompt_source` invocations
after exhaustion when `stop_on_exhausted` is `False`.
Defaults to 100ms -- cheap, lets the producer resume
quickly if more data lands.

---

**`stop_on_exhausted`** `bool` — default: True

When `True` (default), the background
thread exits as soon as `prompt_source` returns `None`
and `close` drains. When `False`, the thread keeps
polling so a streaming dataset can refill.

---

**`_error`** `BaseException | None = None`

---

**`_iterator`** `Iterator[_PromptBatch | _PackedPromptBatch | Tensor] | None = prompt_source`

---

**`_prompt_source`** `Callable[[], _PromptBatch | _PackedPromptBatch | None] = self._pull_from_iterator`

---

**`_stop_event`** `= threading.Event()`

---

**`_thread`** `Thread | None = None`

---

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.__enter__() -> 'AsyncFeaturePipeline'
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.__exit__(
    exc_type,
    exc,
    tb
) -> None
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline._pull_from_iterator() -> nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | None
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline._run() -> None
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.join(
    timeout: float | None = None
) -> None
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.start() -> None
```

Spawn the background producer thread.

Idempotent while the thread is alive: a second call is a no-op. The
thread is named `streaming-async-producer` so test failures and
runtime traces show which thread is hung.

The pipeline is single-use. When the producer thread exits -- prompt
source exhausted, `stop`, or an error -- the bound queue is
closed in `_run`'s `finally` so the consumer drains and stops. A
closed queue cannot reopen, so restarting is not supported; construct
a fresh pipeline (and queue) to stream again. Restart is rejected here
rather than failing later with an opaque "queue is closed" from a
blocking put.

```python
nemo_automodel.components.speculative.streaming.async_pipeline.AsyncFeaturePipeline.stop(
    join_timeout: float | None = 10.0
) -> None
```

Signal the background thread to exit and join it.

Outstanding leases are the trainer's; this method does not
ack them. Idempotent. If the background thread raises, the
exception is re-raised here so the trainer sees the failure.

```python
class nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm
```

**Bases:** `enum.Enum`

Which speculative-decoding draft family produced this sample.

The on-wire string value (`eagle3` / `dflash` / `dspark`) is the
stable schema key the consumer matches against. New algorithms land by
adding a member here plus a schema row in the RFC's "Feature schema per
algorithm" table; both moves are part of the same change.

**`DFLASH`** `= 'dflash'`

---

**`DSPARK`** `= 'dspark'`

---

**`EAGLE3`** `= 'eagle3'`

---

```python
class nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader(
    queue: nemo_automodel.components.speculative.streaming.queue.SampleRefQueue,
    store: nemo_automodel.components.speculative.streaming.store.FeatureStore,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm | None = None,
    acquire_poll_interval: float = 0.05
)
```

Iterator over `Eagle3TargetBatch` materialized from a streaming queue.

**Lifecycle**

Each iterator pull yields an `Eagle3TargetBatch` whose
tensors come from a fresh `store.get()`. The previous
batch's lease is ack'd and its store handle released on the
NEXT pull -- so the trainer can hold one batch across one
forward pass without it being freed mid-forward. A consumer
that wants eager reclamation (e.g. to free memory before
pulling the next batch) calls `consume_now` after
computing its loss. Iteration ends on `close` or when
the queue drains.

**Parameters:**

**`queue`** `SampleRefQueue`

The metadata-only queue the producer puts `SampleRef`
onto. `queue.close()` at any time cuts the iterator short.

---

**`store`** `FeatureStore`

The `FeatureStore` each lease will be materialized
through. Must match `ref.store_uri` for the leased refs.

---

**`algorithm`** `FeatureAlgorithm | None` — default: None

`FeatureAlgorithm` the loader runs the
per-algorithm schema check for; defaults to EAGLE-3.

---

**`_pending_handle`** `StoreHandle | None = None`

---

**`_pending_lease`** `Lease | None = None`

---

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__enter__()
```

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__exit__(
    exc_type,
    exc,
    tb
)
```

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__iter__()
```

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.__next__() -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch
```

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader._release_pending() -> None
```

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.close() -> None
```

Ack the most recent lease / release its store handle and shut the queue.

After `close` the iterator raises `StopIteration` on
the next pull. Idempotent so a trainer can use `with` safely.

```python
nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader.consume_now() -> None
```

Release the most recently yielded batch eagerly.

Useful for trainer hooks that want to free memory before
pulling the next batch (e.g. immediately after `backward()`).
Idempotent.

```python
class nemo_automodel.components.speculative.streaming.producer.FeatureProducer(
    target_backend: nemo_automodel.components.speculative.eagle.backend.Eagle3TargetBackend,
    store: nemo_automodel.components.speculative.streaming.store.FeatureStore,
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm | None = None,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    sample_id_factory: typing.Callable[[int], str] | None = None
)
```

Run a target backend once and put the supervision into a `FeatureStore`.

Thread safety: `produce` is not reentrant with itself on the same
producer instance; one thread should drive `produce` at a time.
Multiple producers over the same store are fine; the store's lock
serializes the puts.

**Parameters:**

**`target_backend`** `Eagle3TargetBackend`

an object exposing `generate_batch(input_ids,
attention_mask, loss_mask, ...) -> Eagle3TargetBatch` plus
the `get_input_embeddings` and `set_vocab_mapping`
accessors from
`Eagle3TargetBackend`.

---

**`store`** `FeatureStore`

The `FeatureStore` every `produce` call writes
into.

---

**`run_id`** `str`

Stable across the whole training run; mirrored onto every
`SampleRef`.

---

**`algorithm`** `FeatureAlgorithm | None` — default: None

Forced `FeatureAlgorithm`; defaults to picking
by the backend's runtime type.

---

**`target_model_version`** `str` — default: '0'

Monotonically increasing identifier of the
target-model weights; surfaced on each `SampleRef`.

---

**`draft_weight_version`** `str` — default: '0'

Same idea for the draft model's weights.

---

**`sample_id_factory`** `Callable[[int], str] | None` — default: None

Callable that produces the per-call sample id;
defaults to `"sample-{run_id}-{n}"`. Wrap with a hash of
the input batch if you want stable ids across runs.

---

**`_algorithm`**

---

**`_sample_id_factory`** `= sample_id_factory or self._default_sample_id`

---

**`_sample_seq`** `= 0`

---

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer._default_sample_id(
    n: int
) -> str
```

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer.close() -> None
```

Release any backend resources; the producer itself holds none.

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer.discard(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef
) -> None
```

Evict a produced sample whose ref will not be enqueued.

Used when a produced ref is dropped before it reaches the queue: a
spec-mismatch abort in `produce`, or the async pipeline shutting
down while a blocking put is in flight. `get` + `release` drops the
last handle so the store frees the sample; without it the sample's bytes
(or its shared-dir file) stay resident with no consumer able to release
them.

**Parameters:**

**`ref`** `SampleRef`

The reference whose backing sample should be evicted.

---

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer.get_input_embeddings() -> torch.nn.Module
```

Expose the wrapped backend's input-embedding module.

Mirrors `Eagle3TargetBackend.get_input_embeddings` so the
draft can seed its input embeddings from the target before the
first streamed batch arrives.

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer.produce(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

Run the wrapped backend once and stash its supervision into the store.

**Parameters:**

**`input_ids`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, `torch.long`.

---

**`attention_mask`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, `torch.long`.

---

**`loss_mask`** `torch.Tensor`

Tensor of shape `[batch, sequence]`, `torch.long`.

---

**`position_ids`** `torch.Tensor | None` — default: None

Per-document position ids when packing is on;
`None` for non-packed. Carried through unchanged to
the `Eagle3TargetBatch` so the trainer's loss
path stays bit-identical to the colocated path.

---

**`seq_lens`** `torch.Tensor | None` — default: None

Per-document sequence lengths when packing is on;
`None` for non-packed.

---

**`doc_remaining`** `torch.Tensor | None` — default: None

`[batch, sequence]` long tensor gating
cross-document TTT supervision when packing is on;
`None` otherwise.

---

**Returns:** `SampleRef`

class:`SampleRef` carrying a tensor-free description of

**Raises:**

* `MemoryError`: forwarded from the store when the put would
  exceed `max_samples` or `max_bytes`. Producers can
  wait for the store to drain through the queue's
  backpressure path (`put_blocks_until_below`).

```python
nemo_automodel.components.speculative.streaming.producer.FeatureProducer.set_vocab_mapping(
    selected_token_ids: torch.Tensor,
    selected_token_mask: torch.Tensor
) -> None
```

Thread the draft-vocab mapping through to the wrapped backend.

The colocated backend ignores the mapping and projects the
full-vocab `logits` trainer-side; a future draft-vocab
encoding producer supplies the mapping here so the target can
precompute `target_probs` + `position_mask` and the wire
never carries a full-vocab tensor.

```python
class nemo_automodel.components.speculative.streaming.refs.FeatureSpec(
    shape: tuple[int, ...],
    dtype: torch.dtype
)
```

Dataclass

Shape + dtype metadata for one named feature in a `SampleRef`.

Mirrors the `{dtype_code, shape}` pair the EAGLE-3 remote protocol
encodes in `nemo_automodel.components.speculative.eagle.remote.protocol.encode_nccl_metadata`
so the client can preallocate the NCCL receive buffer. Consumers here use
it to preallocate the receive tensor before calling
`FeatureStore.get`.

**`dtype`** `dtype`

The `torch.dtype` the producer will store under
`SampleRef.feature_keys[name]`. Validation in
`assert_no_tensors` rejects stray tensors in the ref but
a `torch.dtype` object on this field is allowed.

---

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

Logical shape with the same axis order the producer used.
`shape[0]` is conventionally batch and `shape[-1]` the
feature dimension; algorithms with packed layouts (THD, BHSD)
document their conventions in the algorithm row of the
RFC's "Feature schema per algorithm" table.

---

```python
nemo_automodel.components.speculative.streaming.refs.FeatureSpec.__post_init__() -> None
```

```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 `SampleRefQueue`
uses `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 `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 `SampleRefQueue` between leases and unconditionally by
`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 `StoreHandle`.

The returned tensors are detached copies (`clone` for `cuda`
tensors, plain `detach` for CPU views) so prefetch cannot observe
aliasing through `release`. Returns one tensor per key in
`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 `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
`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 `SampleRef`.

The producer-side metadata (`run_id`, `algorithm`, `schema_version`,
`target_model_version`, `draft_weight_version`, `num_tokens`) is
the data the `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 `get` on the same `SampleRef` raises
`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.queue.Lease(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    deadline: float,
    visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout,
    redelivery_count: int = 0,
    lease_id: int = _next_lease_id()
)
```

Dataclass

Handle to a leased `SampleRef`.

Each `SampleRefQueue.acquire` call mints a fresh `Lease`
with a unique `lease_id`. The queue's `ack` and
`fail` verify the lease identity before mutating internal
state, so a late ACK for a stale (reclaimed) lease cannot pop a
newer active lease for the same `sample_id`.

**`deadline`** `float`

Monotonic-clock timestamp at which this lease is
considered orphaned. Used by
`SampleRefQueue.reclaim_expired` to redeliver the ref.

---

**`lease_id`** `int = field(default_factory=_next_lease_id)`

Per-acquire unique identifier. The queue uses it as
the key in `_outstanding` and verifies it before any
ack/fail mutation.

---

**`redelivery_count`** `int = 0`

Number of times this ref has been leased and
re-leased (used for retry telemetry). Starts at 0.

---

**`ref`** `SampleRef`

The leased reference -- the only sanctioned way to materialize
its tensors is `FeatureStore.get`, which returns a
`StoreHandle`
the consumer must hand to `FeatureStore.release` once it is
done with them.

---

**`visibility_timeout`** `VisibilityTimeout`

The `VisibilityTimeout` that produced
this lease, kept here so the consumer can introspect it.

---

```python
class nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore(
    max_samples: int | None = 64,
    max_bytes: int | None = 256 * 1024 * 1024,
    high_watermark_bytes: int | None = 192 * 1024 * 1024,
    low_watermark_bytes: int | None = 64 * 1024 * 1024
)
```

**Bases:** [FeatureStore](/nemo-automodel/nemo_automodel/components/speculative/streaming/store#nemo_automodel-components-speculative-streaming-store-FeatureStore)

In-process `FeatureStore` implementation.

Thread safety: every public method holds a single `threading.Lock`,
so concurrent puts and gets from the same Python process are safe. Async
/ cross-process safety is the queue's responsibility and is out of scope
for PR 1.

**Parameters:**

**`max_samples`** `int | None` — default: 64

Hard cap on simultaneously-stored samples. `None` means
unbounded sample count (still bounded by `max_bytes`).

---

**`max_bytes`** `int | None` — default: 256 \* 1024 \* 1024

Hard cap on resident bytes. `None` means unbounded
(still bounded by `max_samples`). At least one of
`max_samples` / `max_bytes` must be set, otherwise a
misconfigured store silently behaves as unbounded.

---

**`high_watermark_bytes`** `int | None` — default: 192 \* 1024 \* 1024

Threshold above which `StoreHealth.high_watermark_hit`
is `True`. The producer pauses here.

---

**`low_watermark_bytes`** `int | None` — default: 64 \* 1024 \* 1024

Threshold below which `StoreHealth.low_watermark_hit`
is `True`. The producer resumes here. Must be
strictly less than `high_watermark_bytes`; a hysteresis band
of zero flaps the producer on every step.

---

**`_handle_refs`** `dict[int, str] = {}`

---

**`_high_watermark`**

---

**`_lock`** `= threading.Lock()`

---

**`_low_watermark`**

---

**`_resident_bytes`** `= 0`

---

**`_storage`** `dict[str, dict[str, Tensor]] = {}`

---

**`_uri_id`** `= uuid.uuid4().hex`

---

**`store_uri`** `str`

---

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

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore._tensor_bytes(
    tensor: torch.Tensor
) -> int
```

staticmethod

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.close() -> None
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.gc() -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.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]
```

Materialize `ref`'s features on `device` and hand back a `StoreHandle`.

**Parameters:**

**`ref`** `SampleRef`

The reference returned by `put` (typically via a
queue lease). `ref.store_uri` MUST equal this store's
`store_uri`; a mismatch raises `KeyError` so a
consumer cannot accidentally materialize a foreign ref.

---

**`device`** `torch.device | str | None` — default: None

Optional target device. `None` returns each feature
on the device it was put on; a non-`None` value
materializes every feature on that device via
`Tensor.to(device)` (a no-op when already in place).

---

**Returns:** `dict[str, torch.Tensor]`

A `(tensors, handle)` pair. `tensors` is a

**Raises:**

* `KeyError`: when `ref.store_uri` does not match this store,
  or when `ref.sample_id` is no longer present (released
  or never put).
* `RuntimeError`: when the stored tensor's shape or dtype
  differs from what the ref claims, or the store has been
  closed.

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth
```

```python
nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm = FeatureAlgorithm.EAGLE3,
    schema_version: int = 1,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    num_tokens: int = 0
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

Store `tensors` under `sample_id` and return a tensor-free `SampleRef`.

**Parameters:**

**`sample_id`** `str`

Stable identifier within `run_id`. Must be unique
in this store at put time; duplicates raise `ValueError`.

---

**`tensors`** `Mapping[str, torch.Tensor]`

Feature-name to tensor mapping. The store detaches,
clones, and makes each tensor contiguous before stashing it,
so the producer may keep mutating its source tensors after
the put returns without disturbing what a later
`get` hands out. The shape and dtype of each tensor
are captured into the returned `SampleRef`'s
`feature_specs`; the consumer uses those specs to
preallocate the receive buffer at `get` time, so
changing `tensors[name].shape` or `dtype` between put
and get without updating the ref will surface as a
`RuntimeError` on materialization.

---

**`run_id`** `str`

Same value on every ref of one run; surfaces on the
`SampleRef.run_id` so producers and consumers can
verify they are talking about the same run.

---

**`algorithm`** `FeatureAlgorithm` — default: FeatureAlgorithm.EAGLE3

Which draft family produced this sample; gates the
`SampleRef` required-features check.

---

**`schema_version`** `int` — default: 1

Bumped whenever the producer's feature set or
layout for `algorithm` changes incompatibly.

---

**`target_model_version`** `str` — default: '0'

Monotonically increasing identifier of
the target-model weights.

---

**`draft_weight_version`** `str` — default: '0'

Same idea for the draft model's weights.

---

**`num_tokens`** `int` — default: 0

Sum of attended tokens; used by the consumer for
empty / short loss-mask neutralization.

---

**Returns:** `SampleRef`

A tensor-free `SampleRef` carrying the per-feature

**Raises:**

* `MemoryError`: if the put would exceed `max_samples` or
  `max_bytes`. The producer is expected to retry after the
  store drains below the low watermark (see
  `health`).
* `RuntimeError`: if the store has been closed.
* `ValueError`: on bad input (empty sample id, empty tensors map,
  duplicate sample id).

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

```python
class nemo_automodel.components.speculative.streaming.async_pipeline.PromptSource()
```

Protocol

Zero-argument callable that supplies the next prompt batch.

```python
nemo_automodel.components.speculative.streaming.async_pipeline.PromptSource.__call__() -> nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch | nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch | None
```

Return the next prompt batch or `None` when the source is exhausted.

**Returns:** `_PromptBatch | _PackedPromptBatch | None`

A three-tuple `(input_ids, attention_mask, loss_mask)` where

```python
class nemo_automodel.components.speculative.streaming.refs.SampleRef(
    sample_id: str,
    run_id: str,
    store_uri: str,
    feature_keys: dict[str, str],
    feature_specs: dict[str, nemo_automodel.components.speculative.streaming.refs.FeatureSpec],
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm,
    schema_version: int,
    num_tokens: int,
    estimated_bytes: int,
    target_model_version: str,
    draft_weight_version: str
)
```

Dataclass

Tensor-free reference to one produced sample in a feature store.

**`algorithm`** `FeatureAlgorithm`

Which draft family produced this sample; the consumer
matches it against a registered `FeatureSpec` registry
to validate the ref before materializing.

---

**`draft_weight_version`** `str`

Same idea for the draft model's weights, so a
consumer can refuse to train against a ref produced before its
own weight snapshot. Starts at `"0"`; PR 4 wires the resync.

---

**`estimated_bytes`** `int`

Sum of every feature's `numel() * dtype_bytes`.
The `SampleRefQueue` reads this against
`StoreHealth.resident_bytes` for watermark hysteresis.

---

**`feature_keys`** `dict[str, str]`

Named features this sample contributes, mapped to
per-store keys (e.g. filenames for `SharedDirFeatureStore` in
PR 3, dict keys for `LocalFeatureStore`).

---

**`feature_specs`** `dict[str, FeatureSpec]`

Per-feature `FeatureSpec` so the consumer can
allocate the receive buffer before calling `FeatureStore.get`.

---

**`num_tokens`** `int`

Sum of attended tokens; used by the consumer for empty /
short loss-mask neutralization and by the queue for backpressure.

---

**`run_id`** `str`

Stable identifier across the whole training run. Same value
on every ref of one run so producers and consumers can verify
they are talking about the same run before materializing.

---

**`sample_id`** `str`

Stable identifier within `run_id`. The consumer uses it
to partition the stream across DP ranks (RFC §"Consumer-side DP
resharding") and to ACK / FAIL the corresponding lease.

---

**`schema_version`** `int`

Bumped whenever the producer's feature set or tensor
layout for `algorithm` changes incompatibly. The consumer uses
it as a hard gate on the ref.

---

**`store_uri`** `str`

`scheme://location` identifying the `FeatureStore`
back-end (e.g. `mem://local` for
`LocalFeatureStore`).
The store URI plus `feature_keys` form the lookup the consumer
makes; the store object itself is discovered through the URI at
materialization time, not stored on the ref.

---

**`target_model_version`** `str`

Monotonically increasing identifier of the
target-model weights that produced this sample. Used to reject
refs from a stale target (relevant once train-with-decode /
weight resync lands in PR 4).

---

```python
nemo_automodel.components.speculative.streaming.refs.SampleRef.__post_init__() -> None
```

```python
nemo_automodel.components.speculative.streaming.refs.SampleRef.feature_names() -> tuple[str, ...]
```

Return the feature names in insertion order.

The order is significant because some algorithms size the draft's
`fc` projection by the *order* of the aux hidden states, not just
their count; the producer chooses the order once, the ref freezes it,
and the consumer preserves it through `FeatureStore.get`.

```python
class nemo_automodel.components.speculative.streaming.queue.SampleRefQueue(
    store: nemo_automodel.components.speculative.streaming.store.FeatureStore,
    visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout | None = None,
    high_watermark_bytes: int | None = None,
    low_watermark_bytes: int | None = None,
    on_pause: typing.Callable[[StoreHealth], None] | None = None,
    on_resume: typing.Callable[[StoreHealth], None] | None = None
)
```

Lease / ack / fail queue over `SampleRef`.

Thread safety: a single `threading.Lock` protects every list /
counter, so a multi-producer / multi-consumer deployment works as long
as only one thread at a time calls any one of the methods.

**Parameters:**

**`store`** `FeatureStore`

The data-plane store the consumers will materialize against.
The queue reads `FeatureStore.health` for backpressure.

---

**`visibility_timeout`** `VisibilityTimeout | None` — default: None

How long a leased-but-not-acked ref can live
before reclaim. Defaults to 30s; production deployments
normally key this off the recipe's per-step budget.

---

**`high_watermark_bytes`** `int | None` — default: None

Optional resident-byte threshold for
pausing. When `None` (default), the queue defers to
`StoreHealth.high_watermark_hit` (i.e. the store's own
configured threshold). When set, the queue pauses whenever
`StoreHealth.resident_bytes >= high_watermark_bytes`.

---

**`low_watermark_bytes`** `int | None` — default: None

Optional resident-byte threshold for
resuming. When `None` (default), the queue defers to
`StoreHealth.low_watermark_hit`. When set, the queue
resumes only after `StoreHealth.resident_bytes <=
low_watermark_bytes`. Must be strictly less than
`high_watermark_bytes` so the hysteresis band is
non-empty.

---

**`on_pause / on_resume`**

Optional callbacks fired when the queue
transitions high-watermark-paused -> resumed and back.

---

**`_active_by_sample`** `dict[str, int] = {}`

---

**`_lock`** `= threading.Lock()`

---

**`_outstanding`** `dict[int, Lease] = {}`

---

**`_pending`** `list[SampleRef] = []`

---

**`_pending_seen`** `set[str] = set()`

---

**`_put_cv`** `= threading.Condition(self._lock)`

---

**`_sample_counters`** `dict[str, int] = {}`

---

**`_vt`** `= visibility_timeout or VisibilityTimeout()`

---

**`is_closed`** `bool`

Whether `close` has been called on this queue.

Consumers that pull `acquire` and receive `None` use
this to disambiguate "drained, stop" (`is_closed is True`)
from "transient empty poll, retry" (`is_closed is False`).
Mirrors the Python `queue.Queue` separation between
`empty()` and the lifecycle-shutdown signal.

---

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._gc_store() -> None
```

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_pause(
    health: nemo_automodel.components.speculative.streaming.store.StoreHealth
) -> bool
```

Whether the producer should pause against `health`.

When the queue ctor was given an explicit
`high_watermark_bytes`, that threshold wins; otherwise the
decision defers to `StoreHealth.high_watermark_hit`
(i.e. the store's own configured threshold).

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_resume(
    health: nemo_automodel.components.speculative.streaming.store.StoreHealth
) -> bool
```

Whether the producer should resume against `health`.

When the queue ctor was given an explicit
`low_watermark_bytes`, that threshold wins; otherwise the
decision defers to `StoreHealth.low_watermark_hit`.
Hysteresis is preserved either way: resume crosses the low
threshold, pause crosses the high threshold.

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.ack(
    lease: nemo_automodel.components.speculative.streaming.queue.Lease
) -> None
```

Mark a leased ref as successfully consumed and free its queue slot.

Verifies `Lease.lease_id` matches the live outstanding
entry for `lease.ref.sample_id`: a stale ACK for a lease
that has been reclaimed and re-leased is rejected (logged,
ignored) so the new consumer's live lease is not popped by
accident.

Does NOT touch the store -- the consumer's `FeatureStore.get`
return value carries a `StoreHandle`
that the consumer must hand to `FeatureStore.release` to drop
the tensors. The queue's responsibility ends at "lease no longer held".

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.acquire(
    poll_interval: float = 0.05
) -> nemo_automodel.components.speculative.streaming.queue.Lease | None
```

Lease the next ref; returns `None` when nothing is ready.

`None` is returned in two situations, which consumers
disambiguate with `is_closed`:

* `is_closed is True`: the queue has been shut down and is
  drained. The consumer should stop iterating.
* `is_closed is False`: a transient empty poll (the producer
  is briefly behind). The consumer should retry.

The returned `Lease` is the only sanctioned way to access
the ref's tensors -- `FeatureStore.get` requires a `SampleRef`,
and that ref must come from a lease. The consumer MUST hand back
the lease via `ack` (on success) or `fail` (on error)
so the queue can reclaim the slot and the store can drop the
sample.

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.close() -> None
```

Mark the queue closed; `acquire` drains what remains, then returns `None`.

Closing does not discard already-enqueued refs: `acquire` keeps
handing out pending refs until they are all leased, and only returns
`None` once the queue is closed *and* both pending and outstanding are
empty. `is_closed` therefore reports the closed flag, not that the
queue is already drained; `FeatureDataLoader` polls it to know
when a `None` from `acquire` is terminal.

Outstanding leases are left intact: their consumer still owns the
tensors, and a leaked `FeatureStore.release` would push the
store's residency counter below zero. The store's own `close`
is the canonical place to drop residency.

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.fail(
    lease: nemo_automodel.components.speculative.streaming.queue.Lease
) -> None
```

Return a leased ref to the pending queue, without dropping its tensors.

Verifies the lease identity before re-enqueuing: a stale
`fail` for a lease that has been reclaimed is a no-op. The
ref will be leased again (its `Lease.redelivery_count`
increments). Re-delivery is what makes the pipeline fault-tolerant
to a transient consumer error -- a permanently bad ref is the
consumer's problem (drop it after a bounded retry budget).

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.outstanding_count() -> int
```

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.pending_count() -> int
```

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef
) -> None
```

Enqueue `ref` for a future `acquire`.

Does not block on backpressure; producers that care should call
`put_blocks_until_below` instead, which honors the high/low
watermark hysteresis from `FeatureStore.health`.

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put_blocks_until_below(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    poll_interval: float = 0.05,
    abort_when: typing.Callable[[], bool] | None = None
) -> None
```

Enqueue `ref`, blocking the producer while the store is over its high watermark.

The producer is paused when `_should_pause` returns `True`
(resident crossed the high threshold) and only resumed when
`_should_resume` returns `True` (resident dropped back
below the low threshold). In the band between the two
thresholds the producer's existing paused / unpaused state is
preserved -- that hysteresis is what prevents flapping when the
producer is sitting near the high watermark.

**Parameters:**

**`ref`** `SampleRef`

The reference to enqueue.

---

**`poll_interval`** `float` — default: 0.05

Seconds between backpressure checks when paused.
Defaults to 50ms -- well below typical step times, well above
the cost of a Python-level `FeatureStore.health` call.

---

**`abort_when`** `Callable[[], bool] | None` — default: None

Optional callable checked on each loop iteration.
When it returns `True`, the put aborts with
`RuntimeError` so a shutdown signal can unblock a
producer waiting on backpressure without closing the queue
first.

---

**Raises:**

* `RuntimeError`: if the queue is closed while the producer is
  blocked, or if `abort_when` returns `True`.

```python
nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.reclaim_expired() -> int
```

Reclaim leases whose `Lease.deadline` has passed.

Each reclaimed lease is re-enqueued; `acquire` returns it on
a future call with an incremented `Lease.redelivery_count`.
Returns the number of leases reclaimed -- a queue that is healthy
returns 0 most of the time.

```python
class nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore(
    directory: str,
    max_samples: int | None = 64,
    max_bytes: int | None = 256 * 1024 * 1024,
    high_watermark_bytes: int | None = 192 * 1024 * 1024,
    low_watermark_bytes: int | None = 64 * 1024 * 1024
)
```

**Bases:** [FeatureStore](/nemo-automodel/nemo_automodel/components/speculative/streaming/store#nemo_automodel-components-speculative-streaming-store-FeatureStore)

`FeatureStore` backed by one `<sample_id>.safetensors` per sample.

Thread safety: every public method holds a single
`threading.Lock`. Cross-process / cross-node is supported
as long as distinct processes use distinct `sample_id` values;
the lock does not extend across processes.

**Parameters:**

**`directory`** `str`

Filesystem path used as the rendezvous. Created if it
does not exist. Concurrent producers and consumers in
separate processes / ranks coordinate via unique
`sample_id` values -- collision is the caller's problem.

---

**`max_samples, max_bytes, high_watermark_bytes, low_watermark_bytes`**

Same residency contract as `LocalFeatureStore`; the
queue's HWM/LWM hysteresis reads them off
`health`.

---

**`_handle_refs`** `dict[int, str] = {}`

---

**`_high_watermark`**

---

**`_lock`** `= threading.Lock()`

---

**`_low_watermark`**

---

**`_owned_files`** `dict[str, int] = {}`

---

**`store_uri`** `str`

---

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._atomic_write(
    path: str,
    tensors: typing.Mapping[str, torch.Tensor]
) -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._path_for(
    sample_id: str
) -> str
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.close() -> None
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.gc() -> int
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.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]
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth
```

```python
nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm = FeatureAlgorithm.EAGLE3,
    schema_version: int = 1,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    num_tokens: int = 0
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef
```

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

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

Each `FeatureStore.get` mints a fresh handle with a unique
`handle_id`. Two `get` calls on the same sample return two
distinct handles; `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
`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`

Configured hard cap from
`LocalFeatureStore`
(PR 3's shared-dir / PR 4's NCCL store report the analogous cap).

---

**`high_watermark_hit`** `bool`

True iff `resident_bytes >= high_watermark_bytes`
on the last `health` call. The queue pauses a producer that
sees this transition.

---

**`low_watermark_hit`** `bool`

True iff `resident_bytes <= low_watermark_bytes`
on the last `health` call. The queue resumes a producer
that has been paused and now sees this transition. The
hysteresis band between the two is what prevents flapping.

---

**`resident_bytes`** `int`

Bytes currently held in the store across un-acked
samples. Compared against `capacity_bytes` for the high/low
watermark hysteresis.

---

**`sample_count`** `int`

Number of un-acked samples currently held. Used for
the second cap (sample count) the RFC §"Open questions" Q2 keeps
alongside the byte backstop.

---

```python
class nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout(
    seconds: float = 30.0
)
```

Dataclass

How long an unacked `Lease` is allowed to live before reclaim.

Any positive value is accepted (sub-second values are useful in
tests). Production deployments typically pick something an order of
magnitude larger than the recipe's per-step budget so a slow but
healthy consumer does not see its leases reclaimed out from under
it.

**`seconds`** `float = 30.0`

---

```python
nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout.__post_init__() -> None
```

```python
nemo_automodel.components.speculative.streaming.refs.assert_no_tensors(
    obj: typing.Any,
    path: str = 'ref'
) -> None
```

Recursively validate that `obj` carries no tensors.

Walks dataclasses (any nested dataclass included), `dict` with `str`
keys, and `list` / `tuple`. Anything else must be a primitive
(`str`/`int`/`float`/`bool`/`None`/`bytes`) -- a tensor, numpy
array, or duck-typed tensor-like (`has data_ptr + is_cuda + numel`) at
any depth raises `ValueError`.

The check is structural, not nominal: a third-party tensor type that
quacks like one is rejected. PR 2's registry / schema validators layer on
top of this primitive guard, not instead of it.