nemo_automodel.components.speculative.streaming.queue

Lease / ack / fail queue over :class:SampleRef for the streaming pipeline.

The :class:SampleRefQueue carries only references — no tensors — between a producer (target-side forward) and a consumer (draft-side trainer). Each “message” is a :class:SampleRef and is delivered exactly once: a consumer leases a ref, materializes its tensors via the :class:~nemo_automodel.components.speculative.streaming.store.FeatureStore, and then ACKs (release the lease and let the data-plane scrub the sample from the store) or FAILs (release the lease without scrubbing so a future consumer may retry).

A lease that is never ACK’d or FAIL’d within :attr:Lease.visibility_timeout is considered orphaned and is reclaimed by :meth:SampleRefQueue.reclaim_expired. That reclaim is what makes the queue safe to drive against a producer that may crash mid-flight (RFC §“Phased plan” PR 4’s “visibility-timeout redelivery”).

Backpressure is driven by the bound :class:FeatureStore’s :meth:FeatureStore.health (ints only — the queue never touches tensors in its hot path), with a high/low watermark hysteresis band so a fast producer cannot OOM the store and a slow producer cannot starve the trainer silently. The producer-side and consumer-side pause / resume transitions are tracked on the store via the same :attr:StoreHealth.high_watermark_hit / :attr:StoreHealth.low_watermark_hit flags, so a third party (e.g. an ops dashboard) can observe which side of the pipeline is the bottleneck without inspecting the queue internals.

Module Contents

Classes

Name	Description
`Lease`	Handle to a leased :class:`SampleRef`.
`SampleRefQueue`	Lease / ack / fail queue over :class:`SampleRef`.
`VisibilityTimeout`	How long an unacked :class:`Lease` is allowed to live before reclaim.

Functions

Name	Description
`_next_lease_id`	Mint a fresh :attr:`Lease.lease_id` (module-level counter).

Data

__all__

_lease_id_counter

logger

API

class nemo_automodel.components.speculative.streaming.queue.Lease(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    deadline: float,
    visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout,
    redelivery_count: int = 0,
    lease_id: int = _next_lease_id()
)

Dataclass

Handle to a leased :class:SampleRef.

Each :meth:SampleRefQueue.acquire call mints a fresh :class:Lease with a unique :attr:lease_id. The queue’s :meth:ack and :meth:fail verify the lease identity before mutating internal state, so a late ACK for a stale (reclaimed) lease cannot pop a newer active lease for the same sample_id.

deadline

float

lease_id

int = field(default_factory=_next_lease_id)

redelivery_count

int = 0

ref

SampleRef

visibility_timeout

VisibilityTimeout

class nemo_automodel.components.speculative.streaming.queue.SampleRefQueue(
    store: nemo_automodel.components.speculative.streaming.store.FeatureStore,
    visibility_timeout: nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout | None = None,
    high_watermark_bytes: int | None = None,
    low_watermark_bytes: int | None = None,
    on_pause: typing.Callable[[StoreHealth], None] | None = None,
    on_resume: typing.Callable[[StoreHealth], None] | None = None
)

Lease / ack / fail queue over :class:SampleRef.

Thread safety: a single :class:threading.Lock protects every list / counter, so a multi-producer / multi-consumer deployment works as long as only one thread at a time calls any one of the methods.

Parameters:

store

FeatureStore

The data-plane store the consumers will materialize against. The queue reads :meth:FeatureStore.health for backpressure.

visibility_timeout

VisibilityTimeout | NoneDefaults to None

How long a leased-but-not-acked ref can live before reclaim. Defaults to 30s; production deployments normally key this off the recipe’s per-step budget.

high_watermark_bytes

int | NoneDefaults to None

Optional resident-byte threshold for pausing. When None (default), the queue defers to :attr:StoreHealth.high_watermark_hit (i.e. the store’s own configured threshold). When set, the queue pauses whenever StoreHealth.resident_bytes >= high_watermark_bytes.

low_watermark_bytes

int | NoneDefaults to None

Optional resident-byte threshold for resuming. When None (default), the queue defers to :attr:StoreHealth.low_watermark_hit. When set, the queue resumes only after StoreHealth.resident_bytes <= low_watermark_bytes. Must be strictly less than high_watermark_bytes so the hysteresis band is non-empty.

on_pause / on_resume

Optional callbacks fired when the queue transitions high-watermark-paused -> resumed and back.

_active_by_sample

dict[str, int] = {}

_lock

= threading.Lock()

_outstanding

dict[int, Lease] = {}

_pending

list[SampleRef] = []

_pending_seen

set[str] = set()

_put_cv

= threading.Condition(self._lock)

_sample_counters

dict[str, int] = {}

_vt

= visibility_timeout or VisibilityTimeout()

is_closed

bool

Whether :meth:close has been called on this queue.

Consumers that pull :meth:acquire and receive None use this to disambiguate “drained, stop” (is_closed is True) from “transient empty poll, retry” (is_closed is False). Mirrors the Python queue.Queue separation between empty() and the lifecycle-shutdown signal.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_pause(
    health: nemo_automodel.components.speculative.streaming.store.StoreHealth
) -> bool

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 :attr:StoreHealth.high_watermark_hit (i.e. the store’s own configured threshold).

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue._should_resume(
    health: nemo_automodel.components.speculative.streaming.store.StoreHealth
) -> bool

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 :attr:StoreHealth.low_watermark_hit. Hysteresis is preserved either way: resume crosses the low threshold, pause crosses the high threshold.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.ack(
    lease: nemo_automodel.components.speculative.streaming.queue.Lease
) -> None

Mark a leased ref as successfully consumed and free its queue slot.

Verifies :attr: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 :meth:FeatureStore.get return value carries a :class:~nemo_automodel.components.speculative.streaming.store.StoreHandle that the consumer must hand to :meth:FeatureStore.release to drop the tensors. The queue’s responsibility ends at “lease no longer held”.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.acquire(
    poll_interval: float = 0.05
) -> nemo_automodel.components.speculative.streaming.queue.Lease | None

Lease the next ref; returns None when nothing is ready.

None is returned in two situations, which consumers disambiguate with :attr: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 :class:Lease is the only sanctioned way to access the ref’s tensors — :class:FeatureStore.get requires a :class:SampleRef, and that ref must come from a lease. The consumer MUST hand back the lease via :meth:ack (on success) or :meth:fail (on error) so the queue can reclaim the slot and the store can drop the sample.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.close() -> None

Mark the queue closed; :meth:acquire drains what remains, then returns None.

Closing does not discard already-enqueued refs: :meth: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. :attr:is_closed therefore reports the closed flag, not that the queue is already drained; :class:FeatureDataLoader polls it to know when a None from :meth:acquire is terminal.

Outstanding leases are left intact: their consumer still owns the tensors, and a leaked :meth:FeatureStore.release would push the store’s residency counter below zero. The store’s own :meth:close is the canonical place to drop residency.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.fail(
    lease: nemo_automodel.components.speculative.streaming.queue.Lease
) -> None

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

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.outstanding_count() -> int

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.pending_count() -> int

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef
) -> None

Enqueue ref for a future :meth:acquire.

Does not block on backpressure; producers that care should call :meth:put_blocks_until_below instead, which honors the high/low watermark hysteresis from :meth:FeatureStore.health.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.put_blocks_until_below(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    poll_interval: float = 0.05
) -> None

Enqueue ref, blocking the producer while the store is over its high watermark.

The producer is paused when :meth:_should_pause returns True (resident crossed the high threshold) and only resumed when :meth:_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:

ref

SampleRef

The reference to enqueue.

poll_interval

floatDefaults to 0.05

Seconds between backpressure checks when paused. Defaults to 50ms — well below typical step times, well above the cost of a Python-level :meth:FeatureStore.health call.

Raises:

RuntimeError: if the queue is closed while the producer is blocked, so a producer does not silently swallow a shutdown signal.

nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.reclaim_expired() -> int

Reclaim leases whose :attr:Lease.deadline has passed.

Each reclaimed lease is re-enqueued; :meth:acquire returns it on a future call with an incremented :attr:Lease.redelivery_count. Returns the number of leases reclaimed — a queue that is healthy returns 0 most of the time.

class nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout(
    seconds: float = 30.0
)

Dataclass

How long an unacked :class:Lease is allowed to live before reclaim.

Any positive value is accepted (sub-second values are useful in tests). Production deployments typically pick something an order of magnitude larger than the recipe’s per-step budget so a slow but healthy consumer does not see its leases reclaimed out from under it.

seconds

float = 30.0

nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout.__post_init__() -> None

nemo_automodel.components.speculative.streaming.queue._next_lease_id() -> int

Mint a fresh :attr:Lease.lease_id (module-level counter).

nemo_automodel.components.speculative.streaming.queue.__all__ = ['Lease', 'SampleRefQueue', 'VisibilityTimeout']

nemo_automodel.components.speculative.streaming.queue._lease_id_counter = itertools.count()

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