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

Functions

Name	Description
`_materialize_batch`	Build an :class:`Eagle3TargetBatch` from the store’s tensors.

Data

__all__

logger

API

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

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

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

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

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

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

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.

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.

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.

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

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