Metadata-Version: 2.4
Name: neprr-streaming
Version: 0.1.0
Classifier: Development Status :: 3 - Alpha
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Rust
Requires-Dist: numpy>=1.23
Requires-Dist: torch>=2.6
Requires-Dist: typing-extensions>=4.0
Requires-Dist: pandas>=2.2,<3 ; extra == 'spark'
Requires-Dist: pyarrow>=15 ; extra == 'spark'
Requires-Dist: pyspark>=3.5,<5 ; extra == 'spark'
Requires-Dist: pyarrow>=15 ; extra == 'writer'
Provides-Extra: spark
Provides-Extra: writer
License-File: LICENSE
License-File: NOTICE
License-File: LICENSES/NUMPY-BSD-3-CLAUSE.txt
License-File: LICENSES/PCG-MIT.txt
Summary: Rust-native, Python-compatible Mosaic Streaming loader and MDS writer
Author: Arcee AI
License-Expression: Apache-2.0
Requires-Python: >=3.10
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

# neprr-streaming

`neprr-streaming` is a Rust-native implementation of the Mosaic Streaming dataloader and MDS
writer.
It installs the drop-in Python import path `streaming` and targets Mosaic Streaming `0.11.0`
behavior for the explicitly supported MDS v2 subset.

The preferred mode uses `num_workers=0`. Each distributed training rank remains a process, while
one Tokio I/O worker and a bounded Rayon pool replace Python DataLoader workers, multiprocessing
shared memory, per-sample file opens, and Python cloud SDKs.

This is an alpha release. The Python distribution is named `neprr-streaming`.

Install into an environment without `mosaicml-streaming`: both distributions own the
`streaming` import path and cannot safely coexist. An environment containing the old
`arcee-streaming` distribution must also have that distribution removed before installing
`neprr-streaming`.

## Build and test

Rust `1.91.0` is pinned in `rust-toolchain.toml`. Python 3.12 is the contributor default, while
the ABI3 wheel supports Python 3.10+ and requires PyTorch 2.6 or newer.

```bash
uv sync --locked

cargo test -p streaming-core --all-features --locked
PYO3_NO_PYTHON=1 cargo clippy --workspace --all-targets --all-features --locked -- -D warnings
uv run --locked ruff check python/streaming python/tests scripts
uv run --locked pytest
```

Run `make develop` after changing Rust sources to rebuild the editable extension.

`uv.lock` is a standards-only PyPI contributor environment. The project metadata deliberately does
not declare a PyTorch package source: uv applies sources from Git dependencies while resolving
their consumers, so exporting a CPU index would silently force training projects off CUDA. Training
repositories select their own explicit CPU, CUDA, ROCm, or XPU index.

Build an ABI3 wheel locally:

```bash
# Linux x86_64: link against the declared manylinux_2_28 baseline.
uv run --locked maturin build --release --zig --out dist

# macOS arm64: use the native deployment target from pyproject.toml.
uv run --locked maturin build --release --out dist
```

The wheel exposes `streaming._streaming_rs` and supports CPython 3.10 or newer. Release targets are
Linux x86_64 (`manylinux_2_28`) and macOS arm64.

## Preferred usage

```python
from streaming import StreamingDataLoader, StreamingDataset, StreamingRuntimeConfig

dataset = StreamingDataset(
    remote="s3://my-bucket/pretraining",
    local="/mnt/nvme/mds-cache",
    split="train",
    batch_size=8,
    shuffle=True,
    shuffle_algo="py1e",
    validate_hash="sha1",
    runtime_config=StreamingRuntimeConfig(
        remote_mode="auto",
        download_concurrency=None,
        decode_threads=None,
        disk_threads=None,
        ready_queue_size=64,
    ),
)

with StreamingDataLoader(
    dataset,
    batch_size=8,
    num_workers=0,
    native_batching="auto",
) as loader:
    for batch in loader:
        train(batch)
```

## Rust MDS writing

The compatible `MDSWriter.write()` API buffers Python samples into Arrow batches. Production
Spark tasks should pass Arrow record batches directly so row encoding, shard rollover,
compression, hashing, and atomic local publication all run in Rust with the GIL released:

```python
from streaming import MDSWriter

with MDSWriter(
    out="/local_disk0/task-attempt",
    columns={"tokens": "ndarray:int32", "doc_starts": "ndarray:int32"},
    compression="zstd:16",
    hashes=["sha1", "sha256"],
) as writer:
    writer.write_record_batch(record_batch)
```

Install the optional Arrow dependency with `neprr-streaming[writer]`. Direct remote publication
is intentionally disabled at this layer. The Spark-aware `dataframe_to_mds` coordinator publishes
attempt-isolated shards to local/FUSE or S3, fully verifies them, deterministically assembles one
global v2 index plus additive v3 index, and writes the final dataset receipt last. See
[Rust MDS writing](docs/mds-writer.md) and the
[Spark DataFrame writer](docs/spark-dataframe-writer.md).

Predictive shard preparation defaults to `prefetch_mode="auto"`, which learns the batch period and
preparation tail while staying inside an evidence-calibrated five-second maximum runway and
normalized required-batch footprint. Use `prefetch_mode="fixed"` for the prior exact `predownload`
scheduling behavior, or `predownload=0` to disable prediction. A positive `predownload` is an
additional sample-horizon ceiling in automatic mode. Disk eviction becomes plan-aware only when
the caller supplies `cache_limit`; `None` remains unlimited.

With one stream and `shuffle=False`, the automatic default also keeps the next distinct physical
shard fully prepared after the first batch. That guaranteed boundary shard is not ordinary
five-second speculation: the runway caps work beyond it. Set an explicit positive `predownload` to
retain a strict sample horizon instead.

`native_batching="auto"` uses the native batch path only for `num_workers=0`, default collation,
a fixed batch size, homogeneous numeric samples, and a dataset that does not override Python sample
materialization hooks. For subclasses overriding `get_item`, `__getitem__`, or `__iter__`,
`"auto"` preserves the Python transform through the compatible path and `"required"` reports the
bypassed hooks as an error. Eligible batches are assembled into aligned Rust-owned storage and
transferred directly into CPU PyTorch tensors through DLPack, without an intermediate NumPy array
or tensor copy. Dynamic shapes and other ineligible samples retain compatible Python collation.
`pin_memory=True` uses PyTorch's standard pinned-memory copy after the zero-copy CPU handoff.
`"off"` always uses PyTorch collation. `num_workers>0` is retained as a PID-safe compatibility
fallback for `fork` and `spawn`.

For datasets whose written sample order defines the exact collective data-parallel sequence, set
`partition_algo="sequential_global"` with `shuffle=False`. This additive extension bypasses
canonical-node interleaving while preserving the default Mosaic-compatible `relaxed` behavior for
all existing callers. The epoch size and checkpoint position must be divisible by the complete
data-parallel microbatch. See [sequential-global partitioning](docs/sequential-global.md).

Automatic remote mode prioritizes cold first-batch work through a bounded bootstrap lane, then
uses adaptive resource-scoped concurrency, progress-aware response watchdogs, classified retries
with jitter and a retry budget, and node-local steady transfer slots. Legacy fixed retry/timeout behavior remains available with
`StreamingRuntimeConfig(remote_mode="legacy")`. See the [runtime/statistics guide](docs/runtime.md)
for policy semantics, container coordination, tuning overrides, and limitations.

## Portable checkpoint plans

The normal Mosaic checkpoint remains the same five-field dictionary and can still move unchanged
between `mosaicml-streaming`, Sapio's Python fork, and `neprr-streaming`. For planned epochs,
`StreamingDataLoader.export_resume_plan()` can additionally publish a versioned sibling sidecar;
after the ordinary state is restored, `load_resume_plan()` validates and atomically installs that
exact compact plan instead of regenerating shuffle, sampling, batching, and partition state.

The v1 fast path is exact for a global-batch-aligned checkpoint with unchanged topology and one
worker per rank. `mode="auto"` falls back to deterministic replanning on a missing, stale, corrupt,
oversized, or incompatible artifact. `mode="required"` fails before yielding a sample. The sidecar
is never embedded in `state_dict` and is never the only durable resume representation. Unmodified
Python Mosaic/Sapio ignores the sibling and replans; it can still load the checkpoint itself.

```python
state = loader.state_dict()
loader.export_resume_plan("checkpoint/streaming-plan-v1.sidecar", compression="auto")

resumed_loader.load_state_dict(state)
resumed_loader.load_resume_plan("checkpoint/streaming-plan-v1.sidecar", mode="auto")
```

## Storage URLs

- Local: `/data/mds` or `file:///data/mds`
- S3 and compatible stores: `s3://bucket/prefix`
- HF dataset repositories: `hf://datasets/owner/repository/prefix`
- HF Storage Buckets: `hf://buckets/owner/bucket/prefix`

S3 uses OpenDAL 0.58.1's AWS credential chain, including environment credentials, shared profiles,
web identity, `credential_process`, ECS, and EC2 IMDSv2. Cached AWS CLI v2 SSO tokens are supported
for profiles containing the SSO fields inline; profiles using a separate `[sso-session]` section
should export short-lived environment credentials until that configuration is validated. It
automatically probes the bucket region when it is not configured. `S3_ENDPOINT_URL`,
`AWS_ENDPOINT_URL_S3`, and `AWS_ENDPOINT_URL` are honored for compatible services. HF repositories
use `HF_TOKEN`, then `HUGGING_FACE_HUB_TOKEN`.

HF dataset repositories and the native HF Storage Bucket endpoint use `HF_TOKEN`, then
`HUGGING_FACE_HUB_TOKEN`, then the active token written by `hf auth login`. Private buckets require
a read-capable token; public buckets may be read
without one. Bucket access is read-only. The optional HF S3 gateway can be selected with
`ARCEE_STREAMING_HF_BUCKET_BACKEND=s3`; it uses the standard AWS credential chain and requires the
separate HFAK credentials generated by Hugging Face, not an ordinary Hub token.

## What is implemented

- MDS v2, uncompressed and Zstandard shards
- bytes, strings, JSON, numeric scalars, and static/dynamic ndarrays
- local, S3-compatible, HF dataset repository, and HF Storage Bucket reads
- stream `proportion`, `repeat`, and `choose`; balanced/fixed sampling and granularity
- relaxed partitioning, opt-in sequential-global partitioning, replication, all Mosaic batching
  and shuffle methods, sequential order, deterministic resume, and optional exact checkpoint-plan
  sidecars
- atomic whole-shard downloads, selected hash validation, mmap reads, cross-process shard locks,
  transactional LRU metadata, mmap-backed epoch plans, and bounded known-plan shard prefetch
- `StreamingDataset`, `StreamingDataLoader`, `Stream`, `LocalDataset`, runtime configuration,
  statistics, context managers, and fail-fast compatibility stubs
- local Mosaic-compatible `MDSWriter` with a Rust/Arrow batch path, atomic shard/index publication,
  Zstandard compression, and all supported hashes
- Spark `dataframe_to_mds` for text/token DataFrames with attempt-isolated retry semantics,
  deterministic partition-order assembly, dual v2/v3 publication, and receipt-gated integrity

See [installation and package ownership](docs/installation.md),
[the complete support matrix](docs/support-matrix.md), [architecture](docs/architecture.md),
[runtime/statistics guide](docs/runtime.md), [benchmark harness](docs/benchmarking.md),
[planning/resume performance](docs/planning-performance.md),
[predictive prefetch performance](docs/prefetch-performance.md),
[adaptive prefetch calibration](docs/prefetch-calibration.md),
[adaptive prefetch final evidence](evidence/adaptive-prefetch/README.md),
[Rust MDS writing](docs/mds-writer.md),
[Spark DataFrame writing](docs/spark-dataframe-writer.md),
[MDS repacking and layout benchmarks](docs/mds-layout-benchmark.md),
[sequential-global partitioning](docs/sequential-global.md),
[logical data-parallel topology](docs/data-parallel.md), and
[the post-v1 roadmap](docs/roadmap.md). Production validation is tracked in the
[compatibility matrix](docs/compatibility.md) and [hardening guide](docs/production-hardening.md).

## Compatibility oracles

`tests/oracles` is generated with the separately locked environment under
`tools/mosaic-reference`; it contains `mosaicml-streaming==0.11.0` but never co-installs it with
`neprr-streaming`. The fixtures cover binary decoding, all supported hashes, zstd, the production
`tokens`/`doc_starts` schema, heterogeneous dynamic-array fallback, exact stream weights, full
multi-stream epoch plans, elastic resume/checkpoint state, all shuffle algorithms, and relaxed partition vectors.
Mosaic is not a runtime dependency.

## License

Apache-2.0. Ported Mosaic algorithms retain attribution and SPDX headers. NumPy RNG behavior is
ported under its BSD-3-Clause license; see `NOTICE`.

