nemo_automodel.components.speculative.streaming.stores.shared_dir

File-backed :class:FeatureStore for cross-process / cross-node streaming.

:class:SharedDirFeatureStore writes one <sample_id>.safetensors file per produced sample into a shared directory. Reads return cloned tensors the consumer can mutate; release deletes the file once the last outstanding handle is dropped. Writes are atomic (tmp file + os.replace) so a concurrent reader never observes a partial write.

The store is the rendezvous for a multi-process run:

• Each process owns its own :class:SharedDirFeatureStore instance pointing at the same shared directory. Producers write files; consumers in other processes materialize them by sample_id without needing a process-local ownership record.
• The shared directory is the cross-process rendezvous; cross-rank routing of which rank reads which sample_id is left to distributed resharding layers above this store.
• Distributed parallelism (FSDP / CP / EP) lives in the trainer’s forward / backward; this store is rank-local state plus a shared filesystem.

Backpressure is identical to :class:LocalFeatureStore: sample and byte caps with the high / low-watermark hysteresis the queue reads. Residency (:meth:health and the max_samples / max_bytes caps) is scoped to the files this instance put, tracked in in-memory counters. That keeps :meth:health off the filesystem hot path (the queue polls it every poll_interval while a producer is paused) and gives each rank its full configured budget instead of 1/N of a directory-global measurement. Files other ranks wrote are not counted here; disk-global capacity is the deployment’s responsibility.

Module Contents

Classes

Functions

Data

__all__

logger

API

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
)

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

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

nemo_automodel.components.speculative.streaming.stores.shared_dir.__all__ = ['SharedDirFeatureStore']

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