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

Background-thread prefetch pipeline for speculative-decoding draft training.

:class:`AsyncFeaturePipeline` wraps a :class:`FeatureProducer` and a
:class:`SampleRefQueue` and runs the target forward in a background
thread, pushing every produced :class:`SampleRef` onto the queue
through :meth:`SampleRefQueue.put_blocks_until_below` so the queue's
HWM/LWM hysteresis governs the producer's pacing.

The trainer-side :class:`~nemo_automodel.components.speculative.streaming.loader.FeatureDataLoader`
iterates the queue; the trainer consumes `Eagle3TargetBatch` instances.
The queue is filled by a background thread rather than by the trainer's
main thread, so target-side forward and draft-side backward overlap.

Distributed-training note: the pipeline is per-rank. Each rank owns its own
:class:`FeatureProducer`, :class:`SampleRefQueue`, and :class:`FeatureStore`;
FSDP / CP / EP happen inside the trainer's forward / backward. Cross-rank
sample routing is handled above this layer.

## Module Contents

### Classes

| Name                                                                                                           | Description                                                                      |
| -------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`AsyncFeaturePipeline`](#nemo_automodel-components-speculative-streaming-async_pipeline-AsyncFeaturePipeline) | Run a :class:`FeatureProducer` in a background thread, draining a prompt source. |
| [`PromptSource`](#nemo_automodel-components-speculative-streaming-async_pipeline-PromptSource)                 | Zero-argument callable that supplies the next prompt batch.                      |

### Data

[`_PackedPromptBatch`](#nemo_automodel-components-speculative-streaming-async_pipeline-_PackedPromptBatch)

[`_PromptBatch`](#nemo_automodel-components-speculative-streaming-async_pipeline-_PromptBatch)

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

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

### 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 :class:`FeatureProducer` in a background thread, draining a prompt source.

**Parameters:**

**`producer`** `FeatureProducer`

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

---

**`queue`** `SampleRefQueue`

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

---

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

A :class:`PromptSource` -- either a zero-arg
callable or an :class:`Iterator`. The callable is invoked
from the background thread; it must be thread-safe (e.g.
a :class:`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 :meth:`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, :meth:`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.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
nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch: TypeAlias = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torc...
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline._PromptBatch: TypeAlias = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

```python
nemo_automodel.components.speculative.streaming.async_pipeline.__all__ = ['AsyncFeaturePipeline', 'PromptSource']
```

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