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.

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:

# 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

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:

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 and the Spark DataFrame writer.

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.

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

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, the complete support matrix, architecture, runtime/statistics guide, benchmark harness, planning/resume performance, predictive prefetch performance, adaptive prefetch calibration, adaptive prefetch final evidence, Rust MDS writing, Spark DataFrame writing, MDS repacking and layout benchmarks, sequential-global partitioning, logical data-parallel topology, and the post-v1 roadmap. Production validation is tracked in the compatibility matrix and hardening guide.

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.