nemo_automodel.components.speculative.streaming.stores

Pluggable feature-store backends for the streaming data plane.

LocalFeatureStore is the in-process reference backend. SharedDirFeatureStore adds cross-process rendezvous on a shared POSIX mount. A future NcclFeatureStore can plug in behind the same FeatureStore contract.

Submodules

Package Contents

Classes

Name	Description
`LocalFeatureStore`	In-process `FeatureStore` implementation.
`SharedDirFeatureStore`	`FeatureStore` backed by one `<sample_id>.safetensors` per sample.

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 FeatureStore implementation.

Thread safety: every public method holds a single 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 StoreHealth.high_watermark_hit is True. The producer pauses here.

low_watermark_bytes

int | NoneDefaults to 64 * 1024 * 1024

Threshold below which 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 StoreHandle.

Parameters:

ref

SampleRef

The reference returned by put (typically via a queue lease). ref.store_uri MUST equal this store’s 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 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 get hands out. The shape and dtype of each tensor are captured into the returned SampleRef’s feature_specs; the consumer uses those specs to preallocate the receive buffer at 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 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 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 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).

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

class nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore(
    directory: str,
    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

FeatureStore backed by one <sample_id>.safetensors per sample.

Thread safety: every public method holds a single threading.Lock. Cross-process / cross-node is supported as long as distinct processes use distinct sample_id values; the lock does not extend across processes.

Parameters:

directory

str

Filesystem path used as the rendezvous. Created if it does not exist. Concurrent producers and consumers in separate processes / ranks coordinate via unique sample_id values — collision is the caller’s problem.

max_samples, max_bytes, high_watermark_bytes, low_watermark_bytes

Same residency contract as LocalFeatureStore; the queue’s HWM/LWM hysteresis reads them off health.

_handle_refs

dict[int, str] = {}

_high_watermark

_lock

= threading.Lock()

_low_watermark

_owned_files

dict[str, int] = {}

store_uri

str

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._atomic_write(
    path: str,
    tensors: typing.Mapping[str, torch.Tensor]
) -> int

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore._path_for(
    sample_id: str
) -> str

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.close() -> None

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.gc() -> int

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.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]

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.health() -> nemo_automodel.components.speculative.streaming.store.StoreHealth

nemo_automodel.components.speculative.streaming.stores.shared_dir.SharedDirFeatureStore.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

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