nemo_automodel.components.speculative.streaming.stores.local

In-process feature store for the speculative-training stream.

The local store keeps tensors in a Python dict under the :class:~nemo_automodel.components.speculative.streaming.refs.SampleRef.sample_id, with a resident-byte counter so the queue can drive backpressure from the same store it puts into. It is the build-and-test surface for the entire streaming pipeline: enough to wire up a colocated producer and consumer without a network, without a shared mount, and without GPUDirect.

Residency policy (RFC §“Open questions” Q2 answer: bytes as the hard backstop, sample count as a soft cap). When the next put would exceed either cap, :meth:put raises :class:MemoryError — the producer is then expected to retry after the consumer drains below the low watermark.

Module Contents

Classes

Name	Description
`LocalFeatureStore`	In-process :class:`FeatureStore` implementation.

Data

__all__

logger

API

class nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore(
    max_samples: int | None = 64,
    max_bytes: int | None = 256 * 1024 * 1024,
    high_watermark_bytes: int | None = 192 * 1024 * 1024,
    low_watermark_bytes: int | None = 64 * 1024 * 1024
)

Bases: FeatureStore

In-process :class:FeatureStore implementation.

Thread safety: every public method holds a single :class:threading.Lock, so concurrent puts and gets from the same Python process are safe. Async / cross-process safety is the queue’s responsibility and is out of scope for PR 1.

Parameters:

max_samples

int | NoneDefaults to 64

Hard cap on simultaneously-stored samples. None means unbounded sample count (still bounded by max_bytes).

max_bytes

int | NoneDefaults to 256 * 1024 * 1024

Hard cap on resident bytes. None means unbounded (still bounded by max_samples). At least one of max_samples / max_bytes must be set, otherwise a misconfigured store silently behaves as unbounded.

high_watermark_bytes

int | NoneDefaults to 192 * 1024 * 1024

Threshold above which :attr:StoreHealth.high_watermark_hit is True. The producer pauses here.

low_watermark_bytes

int | NoneDefaults to 64 * 1024 * 1024

Threshold below which :attr:StoreHealth.low_watermark_hit is True. The producer resumes here. Must be strictly less than high_watermark_bytes; a hysteresis band of zero flaps the producer on every step.

_handle_refs

dict[int, str] = {}

_high_watermark

_lock

= threading.Lock()

_low_watermark

_resident_bytes

= 0

_storage

dict[str, dict[str, Tensor]] = {}

_uri_id

= uuid.uuid4().hex

store_uri

str

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

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore._tensor_bytes(
    tensor: torch.Tensor
) -> int

staticmethod

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.close() -> None

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.gc() -> int

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.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]

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

Parameters:

ref

SampleRef

The reference returned by :meth:put (typically via a queue lease). ref.store_uri MUST equal this store’s :attr:store_uri; a mismatch raises KeyError so a consumer cannot accidentally materialize a foreign ref.

device

torch.device | str | NoneDefaults to None

Optional target device. None returns each feature on the device it was put on; a non-None value materializes every feature on that device via Tensor.to(device) (a no-op when already in place).

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.

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.put(
    sample_id: str,
    tensors: typing.Mapping[str, torch.Tensor],
    run_id: str,
    algorithm: nemo_automodel.components.speculative.streaming.refs.FeatureAlgorithm = FeatureAlgorithm.EAGLE3,
    schema_version: int = 1,
    target_model_version: str = '0',
    draft_weight_version: str = '0',
    num_tokens: int = 0
) -> nemo_automodel.components.speculative.streaming.refs.SampleRef

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

Parameters:

sample_id

str

Stable identifier within run_id. Must be unique in this store at put time; duplicates raise ValueError.

tensors

Mapping[str, torch.Tensor]

Feature-name to tensor mapping. The store detaches, clones, and makes each tensor contiguous before stashing it, so the producer may keep mutating its source tensors after the put returns without disturbing what a later :meth:get hands out. The shape and dtype of each tensor are captured into the returned :class:SampleRef’s feature_specs; the consumer uses those specs to preallocate the receive buffer at :meth:get time, so changing tensors[name].shape or dtype between put and get without updating the ref will surface as a RuntimeError on materialization.

run_id

str

Same value on every ref of one run; surfaces on the :class:SampleRef.run_id so producers and consumers can verify they are talking about the same run.

algorithm

FeatureAlgorithmDefaults to FeatureAlgorithm.EAGLE3

Which draft family produced this sample; gates the :class:SampleRef required-features check.

schema_version

intDefaults to 1

Bumped whenever the producer’s feature set or layout for algorithm changes incompatibly.

target_model_version

strDefaults to '0'

Monotonically increasing identifier of the target-model weights.

draft_weight_version

strDefaults to '0'

Same idea for the draft model’s weights.

num_tokens

intDefaults to 0

Sum of attended tokens; used by the consumer for empty / short loss-mask neutralization.

Returns: SampleRef

A tensor-free :class: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 :meth:health).
RuntimeError: if the store has been closed.
ValueError: on bad input (empty sample id, empty tensors map, duplicate sample id).

nemo_automodel.components.speculative.streaming.stores.local.LocalFeatureStore.release(
    handle: nemo_automodel.components.speculative.streaming.store.StoreHandle
) -> None

nemo_automodel.components.speculative.streaming.stores.local.__all__ = ['LocalFeatureStore']

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