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

| Name                                                                                                                | Description                                                                     |
| ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`SharedDirFeatureStore`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-SharedDirFeatureStore) | :class:`FeatureStore` backed by one `<sample_id>.safetensors` per sample. |

### Functions

| Name                                                                                                | Description |
| --------------------------------------------------------------------------------------------------- | ----------- |
| [`_tensor_bytes`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-_tensor_bytes) | -           |

### Data

[`__all__`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-__all__)

[`logger`](#nemo_automodel-components-speculative-streaming-stores-shared_dir-logger)

### API

```python
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](/nemo-automodel/nemo_automodel/components/speculative/streaming/store#nemo_automodel-components-speculative-streaming-store-FeatureStore)

:class:`FeatureStore` backed by one `<sample_id>.safetensors` per sample.

Thread safety: every public method holds a single
:class:`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 :class:`LocalFeatureStore`; the
queue's HWM/LWM hysteresis reads them off
:meth:`health`.

---

**`_handle_refs`** `dict[int, str] = {}`

---

**`_high_watermark`**

---

**`_lock`** `= threading.Lock()`

---

**`_low_watermark`**

---

**`_owned_files`** `dict[str, int] = {}`

---

**`store_uri`** `str`

---

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

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

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

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

```python
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]
```

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

```python
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
```

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

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

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

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