> 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.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`](#nemo_automodel-components-speculative-streaming-queue-Lease) | Handle to a leased :class:`SampleRef`. | | [`SampleRefQueue`](#nemo_automodel-components-speculative-streaming-queue-SampleRefQueue) | Lease / ack / fail queue over :class:`SampleRef`. | | [`VisibilityTimeout`](#nemo_automodel-components-speculative-streaming-queue-VisibilityTimeout) | How long an unacked :class:`Lease` is allowed to live before reclaim. | ### Functions | Name | Description | | ----------------------------------------------------------------------------------------- | ----------------------------------------------------------- | | [`_next_lease_id`](#nemo_automodel-components-speculative-streaming-queue-_next_lease_id) | Mint a fresh :attr:`Lease.lease_id` (module-level counter). | ### Data [`__all__`](#nemo_automodel-components-speculative-streaming-queue-__all__) [`_lease_id_counter`](#nemo_automodel-components-speculative-streaming-queue-_lease_id_counter) [`logger`](#nemo_automodel-components-speculative-streaming-queue-logger) ### API ```python 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` --- ```python 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 | None` — default: 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 | None` — default: 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 | None` — default: 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. --- ```python 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). ```python 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. ```python 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". ```python 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. ```python 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. ```python 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). ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.outstanding_count() -> int ``` ```python nemo_automodel.components.speculative.streaming.queue.SampleRefQueue.pending_count() -> int ``` ```python 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`. ```python 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`** `float` — default: 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. ```python 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. ```python 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` --- ```python nemo_automodel.components.speculative.streaming.queue.VisibilityTimeout.__post_init__() -> None ``` ```python nemo_automodel.components.speculative.streaming.queue._next_lease_id() -> int ``` Mint a fresh :attr:`Lease.lease_id` (module-level counter). ```python nemo_automodel.components.speculative.streaming.queue.__all__ = ['Lease', 'SampleRefQueue', 'VisibilityTimeout'] ``` ```python nemo_automodel.components.speculative.streaming.queue._lease_id_counter = itertools.count() ``` ```python nemo_automodel.components.speculative.streaming.queue.logger = logging.getLogger(__name__) ```