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`	Run a :class:`FeatureProducer` in a background thread, draining a prompt source.
`PromptSource`	Zero-argument callable that supplies the next prompt batch.

Data

API

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

floatDefaults to 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

boolDefaults to 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

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

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

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

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

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

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.

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.

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

Protocol

Zero-argument callable that supplies the next prompt batch.

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

nemo_automodel.components.speculative.streaming.async_pipeline._PackedPromptBatch: TypeAlias = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torc...

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

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

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