nemo_automodel.components.speculative.streaming.store

Pluggable data-plane transport for the speculative-training stream.

The :class:FeatureStore ABC abstracts where the supervision tensors actually live — an in-process dict for build/test, a POSIX shared mount for multi-node/colocated, or NCCL for GPU-to-GPU in later PRs. The :class:SampleRefQueue reads the store’s :meth:FeatureStore.health to decide whether to back off.

The contract is deliberately small (5 methods + 1 property) so PR 2 and later have an obvious surface to extend:

:meth:put — produce-side: stash tensors for a sample.
:meth:get — consume-side: materialize them; returns a :class:StoreHandle the consumer must hand to :meth:release once it’s done with them.
:meth:release — consume-once: free / drop the materials.
:meth:gc — sweep partially-released handles (a stale lease, a crashed consumer) so the store cannot leak.
:meth:health — ints only; the queue uses these for backpressure.
:meth:close — dispose store resources at shutdown.

Module Contents

Classes

Name	Description
`FeatureStore`	Abstract transport every feature-store backend implements.
`StoreHandle`	Opaque token the consumer must return to :meth:`FeatureStore.release`.
`StoreHealth`	Integers-only snapshot of store residency for backpressure decisions.

Functions

Name	Description
`_next_handle_id`	Mint a fresh :attr:`StoreHandle.handle_id` (module-level counter).

Data

__all__

_handle_id_counter

API

class nemo_automodel.components.speculative.streaming.store.FeatureStore()

Abstract

Abstract transport every feature-store backend implements.

Implementations must be safe to call from one thread per operation; cross-thread concurrency is the queue’s responsibility and outside this contract. The :class:~nemo_automodel.components.speculative.streaming.queue.SampleRefQueue uses :meth:health for backpressure; everything else is the produce/consume pair plus release / gc for lifecycle.

nemo_automodel.components.speculative.streaming.store.FeatureStore.close() -> None

abstract

Release every resource owned by the store (file handles, NCCL groups, …).

After close, all subsequent calls raise :class:RuntimeError. Idempotent with respect to an already-closed store.

nemo_automodel.components.speculative.streaming.store.FeatureStore.gc() -> int

abstract

Sweep stale entries (e.g. failed releases from a crashed consumer).

Returns the number of entries reclaimed. Called opportunistically by the :class:SampleRefQueue between leases and unconditionally by :meth:close.

nemo_automodel.components.speculative.streaming.store.FeatureStore.get(
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    device: torch.device | str | None = None
) -> tuple[dict[str, torch.Tensor], nemo_automodel.components.speculative.streaming.store.StoreHandle]

abstract

Materialize ref’s tensors on device and hand back a :class:StoreHandle.

The returned tensors are detached copies (clone for cuda tensors, plain detach for CPU views) so prefetch cannot observe aliasing through :meth:release. Returns one tensor per key in :attr:SampleRef.feature_keys, in the same insertion order, so the consumer can wire them straight into the per-algorithm batch.

nemo_automodel.components.speculative.streaming.store.FeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth

abstract

Return a ints-only :class: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 :meth:health from cached counters, not from the in-flight I/O thread.

nemo_automodel.components.speculative.streaming.store.FeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm,
    schema_version: int,
    target_model_version: str,
    draft_weight_version: str,
    num_tokens: int
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef

abstract

Stash tensors for sample_id and return a tensor-free :class:SampleRef.

The producer-side metadata (run_id, algorithm, schema_version, target_model_version, draft_weight_version, num_tokens) is the data the :class: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.

nemo_automodel.components.speculative.streaming.store.FeatureStore.release(
    handle: nemo_automodel.components.speculative.streaming.store.StoreHandle
) -> None

abstract

Free the resources backing handle; idempotent.

After this returns, the sample is no longer present in the store and a second :meth:get on the same :class:SampleRef raises :class: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.

class nemo_automodel.components.speculative.streaming.store.StoreHandle(
    store: 'FeatureStore',
    sample_id: str,
    ref: nemo_automodel.components.speculative.streaming.refs.SampleRef,
    handle_id: int = _next_handle_id()
)

Dataclass

Opaque token the consumer must return to :meth:FeatureStore.release.

Each :meth:FeatureStore.get mints a fresh handle with a unique :attr:handle_id. Two get calls on the same sample return two distinct handles; :meth:FeatureStore.release matches against handle_id so releasing one handle twice (or releasing a stale handle after a sibling has been acquired) cannot decrement a sibling’s outstanding count.

Holds the producing store, the sample id, the originating :class:SampleRef, and the per-get handle identity.

handle_id

int = field(default_factory=_next_handle_id)

ref

SampleRef

sample_id

str

store

'FeatureStore'

class nemo_automodel.components.speculative.streaming.store.StoreHealth(
    resident_bytes: int,
    capacity_bytes: int,
    sample_count: int,
    high_watermark_hit: bool,
    low_watermark_hit: bool
)

Dataclass

Integers-only snapshot of store residency for backpressure decisions.

capacity_bytes

int

high_watermark_hit

bool

low_watermark_hit

bool

resident_bytes

int

sample_count

int

nemo_automodel.components.speculative.streaming.store._next_handle_id() -> int

Mint a fresh :attr:StoreHandle.handle_id (module-level counter).

nemo_automodel.components.speculative.streaming.store.__all__ = ['FeatureStore', 'StoreHandle', 'StoreHealth']

nemo_automodel.components.speculative.streaming.store._handle_id_counter = itertools.count()