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

Streaming consumer for speculative-decoding draft training.

:class:`FeatureDataLoader` is a Python iterator over
:class:`Eagle3TargetBatch`. It pulls one :class:`SampleRef` lease at a
time from a :class:`SampleRefQueue`, materializes the tensors through a
:class:`FeatureStore`, hands the trainer a fresh
:class:`Eagle3TargetBatch`, and releases the *previous* lease on every
`__next__` call -- so the trainer can hold one batch across one full
forward without it being freed mid-forward.

The loader is per-rank and the leases stay on-rank; the queue + store
live in the same Python process as the consumer. FSDP / CP / EP
parallelism lives inside the trainer's forward / backward and is
unaffected by the loader's lifecycle.

## Module Contents

### Classes

| Name                                                                                             | Description                                                                   |
| ------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------- |
| [`FeatureDataLoader`](#nemo_automodel-components-speculative-streaming-loader-FeatureDataLoader) | Iterator over :class:`Eagle3TargetBatch` materialized from a streaming queue. |

### Functions

| Name                                                                                               | Description                                                   |
| -------------------------------------------------------------------------------------------------- | ------------------------------------------------------------- |
| [`_materialize_batch`](#nemo_automodel-components-speculative-streaming-loader-_materialize_batch) | Build an :class:`Eagle3TargetBatch` from the store's tensors. |

### Data

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

[`logger`](#nemo_automodel-components-speculative-streaming-loader-logger)

### API

```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 :class:`Eagle3TargetBatch` materialized from a streaming queue.

**Parameters:**

**`queue`** `SampleRefQueue`

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

---

**`store`** `FeatureStore`

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

---

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

:class:`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 :class:`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
nemo_automodel.components.speculative.streaming.loader._materialize_batch(
    algorithm,
    tensors: dict[str, torch.Tensor]
) -> nemo_automodel.components.speculative.eagle.target.Eagle3TargetBatch
```

Build an :class:`Eagle3TargetBatch` from the store's tensors.

Pulled into a module-level helper so the per-algorithm batch-building
logic lives in one place and adding DFlash / DSpark only requires a
new branch here.

```python
nemo_automodel.components.speculative.streaming.loader.__all__ = ['FeatureDataLoader']
```

```python
nemo_automodel.components.speculative.streaming.loader.logger = logging.getLogger(__name__)
```