nemo_automodel.components.speculative.streaming

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.

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.

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.

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.

Release any backend resources; the producer itself holds 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:

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.

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

Parameters:

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

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.

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.

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.

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.

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.

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.

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.

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

Parameters:

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.

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

Parameters:

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

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

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.

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

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.

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

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.

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.

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

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.

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:

Raises:

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

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.