loaderx 0.4.1__tar.gz → 0.5.4__tar.gz

loaderx-0.5.4/MANIFEST.in

@@ -0,0 +1,11 @@
+# The sdist ships sources, not binaries: the shared objects are platform
+# specific and `setup.py` rebuilds them with Zig when they are missing.
+include build.zig
+include build.zig.zon
+include LICENSE
+recursive-include src *.zig
+recursive-include tests *.py
+recursive-include scripts *.py
+
+exclude loaderx/lib/lib_here
+recursive-exclude loaderx/lib *.so *.dylib *.dll *.pdb *.lib

loaderx-0.5.4/PKG-INFO

@@ -0,0 +1,417 @@
+Metadata-Version: 2.4
+Name: loaderx
+Version: 0.5.4
+Summary: A compact, high-performance persistent record store with zero-copy batch gathering and transparent per-record compression, designed for single-machine AI training and serving pipelines
+Author-email: Ben0i0d <ben0i0d@foxmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 EOELAB AI Research
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
+Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
+Project-URL: Source, https://codeberg.org/eoelab/loaderx
+Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
+Keywords: flax,python,dataloader
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: Microsoft :: Windows
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: cffi
+Dynamic: license-file
+
+# Loaderx
+A compact, high-performance persistent record store with zero-copy batch gathering and transparent per-record compression, designed for single-machine AI training and serving pipelines
+
+```
+pip install loaderx
+```
+
+Wheels are published for Linux (glibc ≥ 2.17 and musl, x86-64 and arm64), macOS
+(≥ 11.0, Intel and Apple Silicon) and Windows (x64 and arm64). The bindings use
+cffi in ABI mode, so nothing links against the CPython ABI and one wheel per
+platform serves every supported Python.
+
+## Design Philosophy
+
+loaderx is built around several core principles:
+
+1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
+2. A strong focus on single-machine training workflows.
+3. We implement based on NumPy semantics, persisted through the Zrecord storage runtime.
+4. An **immortal (endless) step-based data loader**, rather than the traditional epoch-based design—better aligned with modern ML training practices.
+
+## Quick Start
+```python
+from loaderx import Dataset, DataLoader
+
+dataset = Dataset('train_data.zr')
+labelset = Dataset('train_label.zr')
+
+loader = DataLoader(dataset, labelset)
+
+for i, batch in enumerate(loader):
+    if i >= 256:
+        break
+
+print(batch['data'].shape)
+print(batch['label'].shape)
+```
+
+### Converting a NumPy tensor
+```python
+import numpy as np
+from loaderx import Dataset
+
+Dataset.write('train_data.zr', np.load('data.npy', mmap_mode='r'))
+Dataset.write('train_label.zr', np.load('label.npy', mmap_mode='r'))
+```
+One record per slice along axis 0; a 1-D array (the usual shape of a label set)
+becomes a dataset of scalars. The conversion streams in bounded chunks, so an
+mmapped array is never fully materialized.
+
+`Dataset` keeps a single `schema.json` inside the store recording the per-sample
+dtype and shape — that is the only loaderx-level metadata. Everything beneath it
+is plain Zrecord, reachable through `loaderx.Zrecord` when you want raw byte
+records (ragged samples, pre-encoded images) instead of tensors.
+
+### Removing samples
+```python
+ds = Dataset('train_data.zr')
+ds.delete([7, 12, 40])   # len shrinks; survivors stay addressable as 0..len
+ds.compact()             # reclaim the freed space on disk
+print(ds.stats())        # {'records': ..., 'live_bytes': ..., 'reclaimable': ...}
+```
+Deletion keeps the index space dense, which is what the sampler needs — but it
+does so by moving the tail into the hole, so surviving samples may change index.
+Do not hold indices across a delete. Both calls require exclusive access; do not
+run them while a `DataLoader` is iterating the dataset.
+
+### Integrating with JAX/Flax
+
+For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
+
+## Benchmarks
+
+`scripts/bench.py` measures loaderx against the alternatives in three layers —
+the index sampler, the record store, and the full data loader — each through the
+binding a client actually uses, so cffi, the GIL and the NumPy allocation are all
+inside the timings. The numbers below are one run on a 24-core machine, Python
+3.14, warm page cache. Reproduce with `python3 scripts/bench.py all`.
+
+**1. Sampler** — index generation on its own, IID (with replacement), 1M index
+space, against both NumPy random APIs.
+
+| sampler           | batch | per batch | indices   | vs default_rng |
+|-------------------|-------|-----------|-----------|----------------|
+| numpy randint     |   256 |   5.9 µs  |   44 M/s  |     0.64x      |
+| numpy default_rng |   256 |   3.7 µs  |   69 M/s  |     1.00x      |
+| **zsampler**      |   256 |   1.9 µs  |  138 M/s  |   **2.01x**    |
+| numpy randint     |  1024 |   7.2 µs  |  142 M/s  |     0.83x      |
+| numpy default_rng |  1024 |   6.0 µs  |  171 M/s  |     1.00x      |
+| **zsampler**      |  1024 |   3.0 µs  |  340 M/s  |   **1.99x**    |
+| numpy randint     |  8192 |  25.6 µs  |  320 M/s  |     0.65x      |
+| numpy default_rng |  8192 |  16.6 µs  |  493 M/s  |     1.00x      |
+| **zsampler**      |  8192 |   7.8 µs  | 1048 M/s  |   **2.13x**    |
+
+Zsampler draws indices about 2x faster than NumPy's faster API — more against the
+legacy `randint` most code still calls — and the lead is widest at the small
+batches a loader actually draws, where a batch is one cffi call either way and
+NumPy's vectorised fill has nothing to amortise. The one concession: zsampler's
+IID draw is modulo-biased where NumPy's is not, ~1e-13 over a u64 draw into a
+million-record space, which is below anything training will notice.
+
+**2. Store** — random batch gather, 20000 records × 12 KiB `uint8(3,64,64)`,
+batch 256 × 200, structured (mildly compressible) data.
+
+| store         | write      | gather      | per batch  | on disk | ratio |
+|---------------|------------|-------------|------------|---------|-------|
+| loaderx-flate |  362 MiB/s |  1671 MiB/s |   1.80 ms  | 209 MiB | 1.12x |
+| loaderx-raw   |  784 MiB/s | 21997 MiB/s |   0.14 ms  | 235 MiB | 1.00x |
+| npy-mmap      | 3064 MiB/s | 13692 MiB/s |   0.22 ms  | 234 MiB | 1.00x |
+| arrayrecord   |  374 MiB/s |   847 MiB/s |   3.54 ms  | 217 MiB | 1.08x |
+| hdf5          | 1327 MiB/s |   511 MiB/s |   5.87 ms  | 236 MiB | 0.99x |
+| hdf5-gzip     |   83 MiB/s |   206 MiB/s |  14.57 ms  | 211 MiB | 1.11x |
+| lmdb          |  677 MiB/s |  5337 MiB/s |   0.56 ms  | 313 MiB | 0.75x |
+| lmdb-zlib     |   71 MiB/s |   336 MiB/s |   8.93 ms  | 235 MiB | 1.00x |
+| blosc2        |   32 MiB/s |   226 MiB/s |  13.27 ms  | 216 MiB | 1.09x |
+| zarr          |   83 MiB/s |    28 MiB/s | 107.36 ms  | 217 MiB | 1.08x |
+| tensorstore   |  531 MiB/s |    98 MiB/s |  30.69 ms  | 214 MiB | 1.09x |
+
+The gather column is what the store is built for. loaderx-raw returns a scattered
+batch faster than everything here, a memory-mapped `.npy` included, because it
+copies straight into the destination across all cores with the GIL released;
+loaderx-flate keeps most of that speed *while* decompressing, and still gathers
+several times faster than any other compressed store (hdf5-gzip, lmdb-zlib,
+blosc2, zarr, tensorstore). The array stores lag on gather because a scattered
+single-record read fights a layout meant for contiguous slabs — which is exactly
+the access pattern a training loader has. Where loaderx does not lead is write: it
+pays compression and its own layout up front, a one-time cost it does not chase.
+Compression is modest (1.12x) only because this data is barely compressible; the
+row's point is that turning it on costs almost nothing on read.
+
+The reads are not all single-threaded: loaderx and blosc2 decompress across cores,
+the rest one record at a time. Zarr and TensorStore are sharded so `on disk`
+reflects compression rather than per-file block rounding, and TensorStore is given
+gzip (its zarr3 default is uncompressed). `on disk` is allocated blocks, which is
+why LMDB's preallocated map lands below 1.00x.
+
+**3. Loader** — the full input pipeline end to end (sample, fetch, collate, hand
+over a batch), same workload, 4 workers. Each loader reads from what it is built
+for: loaderx from Zrecord, Grain from ArrayRecord, torch from a memory-mapped
+`.npy`, jax-dataloader from an in-memory array.
+
+| loader         | batches/s | per batch | samples/s | throughput  |
+|----------------|-----------|-----------|-----------|-------------|
+| **loaderx**    |     724.2 |  1.38 ms  |   185,406 |  2173 MiB/s |
+| loaderx-raw    |    5271.0 |  0.19 ms  | 1,349,378 | 15813 MiB/s |
+| grain          |      67.3 | 14.85 ms  |    17,240 |   202 MiB/s |
+| torch          |     832.2 |  1.20 ms  |   213,038 |  2497 MiB/s |
+| jax-dataloader |    3284.6 |  0.30 ms  |   840,864 |  9854 MiB/s |
+
+End to end, loaderx-raw serves batches several times faster than torch or Grain
+and ahead of jax-dataloader — and it does so from disk-backed storage, where
+jax-dataloader's number is an in-memory ceiling with no storage path at all. The
+default loaderx trades some of that for on-disk compression and still lands in
+torch's range while reading compressed records. loaderx and jax use threads;
+torch and Grain use worker processes that escape the GIL but pay IPC per batch.
+loaderx never ends an epoch, so no step waits on a boundary; the others restart
+per epoch, a cost that is in their numbers as it would be in training.
+
+## Current Limitations
+* Single-host only; multi-host training is not supported.
+* A single sample must be at most 1 MiB, and a store holds at most 2^24
+  (16777216) records.
+* Metadata is mapped and used in place, so a store carries the host's byte
+  order and is not portable to a machine of the opposite endianness. Every
+  published platform is little-endian, so this only matters if you build for
+  one yourself.
+
+## Build
+```
+zig build                       # host shared objects, into loaderx/lib/
+zig build test                  # Zrecord suite, in both Debug and ReleaseFast
+python3 tests/test_loaderx.py   # Python layer, against whichever build is importable
+python3 scripts/bench.py        # throughput, against other stores
+```
+
+The Zig side is tested for behaviour only; throughput is measured from Python,
+through the binding a client actually uses. `scripts/bench.py` runs in three
+layers (`sampler`, `store`, `loader`, or `all`) and skips contenders that are not
+installed. See [Benchmarks](#benchmarks).
+
+### Publishing
+Zig cross-compiles every target from one machine, so releases need no CI matrix:
+
+```
+zig build dist                    # every platform, into zig-out/dist/<wheel tag>/
+python3 scripts/build_wheels.py   # one wheel per platform, plus the sdist
+```
+
+The dist directories are named after their Python wheel platform tag, so the tag
+mapping lives in exactly one place (`dist_targets` in `build.zig`). glibc and
+macOS minimums are pinned in the target triple, which is what makes
+`manylinux_2_17` and `macosx_11_0` honest rather than aspirational. Each wheel is
+checked after packing: it must carry this platform's libraries and no others.
+
+The sdist ships sources only. Installing from it runs `zig build` through
+`setup.py`, so it needs the Zig compiler; wheel users never hit that path.
+
+---
+
+# Zsampler
+Index Generator: a high-performance sampler implemented in Zig
+
+1. Sequential generation: indices are produced by traversing the index space in order.
+    * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
+2. Random generation: indices are sampled randomly from the index space.
+    * Global random: a set of samples is drawn randomly from the entire index space.
+
+---
+
+# Zrecord
+A record-based data runtime, focused on delivering extreme throughput and low latency.
+
+1. Zrecord is an unordered physical store made of N records. Records are
+   independent and carry no ordering, so every index and slice operation is
+   equivalent to a gather.
+2. Zrecord hands the client a dense index space: the live records are always
+   exactly `0..N`. Deletion preserves that by swapping the tail into the hole,
+   which means an index is stable only until something is deleted. A
+   multi-stream dataset is purely a client-side notion; the client keeps its
+   own stream index table pointing at records.
+3. Zrecord returns byte arrays. Type interpretation is the client's job.
+4. For performance, the IO model (`append | read | delete`) is batch-oriented;
+   a single-record operation is just the `batch_size == 1` case.
+5. Zrecord owns its memory internally — allocation and release are explicit.
+6. Each record is compressed and decompressed independently:
+
+```
+| tag   |  name   | algorithm |
+|-------|---------|-----------|
+|   0   |   raw   | none      |
+|   1   |  flate  | Deflate   |
+```
+
+7. Compression is transparent to the client:
+    * Compression runs concurrently across all cores. A record that fails to
+      shrink is stored raw automatically, so reading never costs more than the
+      data itself.
+    * Decompression writes straight into the caller's destination memory
+      (`gather`), with no intermediate buffer and no extra copy.
+
+## Persistence format
+Zrecord storage is metadata plus chunked data:
+```
+zrecord/
+  ├── meta.zr
+  ├── 0.zr
+  └── 1.zr
+```
+
+### Metadata (meta.zr)
+
+meta.zr is **used as the values it holds, not decoded into them**. Both structs
+below are `extern`, 16 bytes, and naturally aligned; an mmap is page aligned, so
+the header is a `*Header` and the record table is a `[]RecordLoc`. Looking up a
+record is `table[idx]`, and there is no serializer anywhere in the code.
+
+The fields are deliberately wider than the limits need — `chunk_id` is a `u16`
+for 4096 chunks, the lengths are `u32` for a 1 MiB cap. Bit-packing them would
+save 5 bytes per record and cost a hand-written codec on the hottest lookup
+path. The slack buys a machine-shaped array; it is worth it.
+
+**1. Header** — 16 bytes of global state.
+* `magic` (`ZREC`) and `version` are ordinary fields, so opening a directory
+  that is not a Zrecord store fails immediately instead of decoding garbage.
+* `length` is the total number of records | `tail_chunk`/`tail_offset` mark the
+  last write position.
+* There is no chunk count. Chunks are created in order and the frontier is
+  always in the last one, so the store holds exactly chunks `0..=tail_chunk` —
+  a count would be a second copy of that fact to keep in sync.
+
+```zig
+const Header = extern struct {
+    magic: [4]u8, version: u16,
+    tail_chunk: u16, tail_offset: u32,
+    length: u32,
+};
+```
+
+**2. Record table** — 16 bytes per entry, indexed directly. Mapping an index to
+a physical address is what makes random access efficient.
+* `chunk_id` is the containing chunk | `offset` is the position within it |
+  `phys_length`/`logic_length` are the stored and original sizes | `compress` is
+  the algorithm.
+
+```zig
+const Codec = enum(u8) { raw = 0, flate = 1, _ };
+const RecordLoc = extern struct {
+    offset: u32, phys_length: u32, logic_length: u32,
+    chunk_id: u16, compress: Codec, _reserved: u8,
+};
+```
+
+There is no liveness flag. Every entry below `length` is live, because deletion
+swaps the tail into the hole rather than tombstoning.
+
+**3. Maximum length.** Running out of space is effectively impossible, so the
+metadata is sized statically at the maximum.
+* Current design: up to 2^12 (4096) chunks, 2^32 (4 GiB) per chunk, 2^20 (1 MiB)
+  per record.
+* Derivation: N = chunks × 2^32 (chunk size) / 2^20 (max record size) ⇒ a
+  maximum length of 2^24 (16777216). Equality here is the point: the store can
+  never run out of chunks before it runs out of record slots, which is what
+  lets the table be sized statically. It is a compile-time assertion.
+* meta.zr is a fixed 16 + 2^24 × 16 bytes (256 MiB). It is sparse — a store
+  with one record allocates 4 KiB of it — and the mapping does not prefault.
+
+### Chunk data (x.zr)
+Densely packed record data.
+
+## Executor
+
+**1. Write.** Writes are append-only; everything else is offset redirection.
+* Append: compress concurrently → assign physical locations serially → flush
+  concurrently → commit metadata. Data is durable-ordered before the record
+  table, and the table before `length`, so a crash truncates rather than
+  corrupts. A record never straddles two chunks; one that would not fit rolls
+  over to a fresh chunk.
+* Delete: swap the last table entry into the deleted slot and drop the length by
+  one. A batch is applied in descending index order, so each swap pulls from a
+  slot no later target refers to. The index space stays dense — which is what
+  the sampler needs, since it draws uniformly from `0..N` and would otherwise
+  keep hitting holes. The deleted record's bytes become garbage.
+
+**2. Read.** Fill the destination memory concurrently, in place from the Python
+side (executed on async threads).
+* Committed records are immutable, so the read path is lock free; `length` is
+  published to readers through an atomic.
+
+**3. Concurrency model.** `Io.Group.async` shards work by CPU count, and shards
+beyond the limit run inline on the calling thread.
+* Shards receive contiguous blocks rather than a strided subset, keeping each
+  worker's reads and writes sequential.
+* Each shard reuses one compressor scratch (`flate.Compress` is ~230 KiB — far
+  too large to rebuild per record, let alone to place on a task stack).
+* Decompression runs in direct mode: matches are resolved against the bytes
+  already written into the destination, so there is no window buffer.
+
+**4. Garbage collection.** Deletion leaves the record's bytes stranded, so space
+is reclaimed by an offline `compact` that rewrites the live records in place.
+* Records are visited in *physical* order — which after swap-last deletes no
+  longer matches table order — and repacked densely in that same order. A record
+  therefore never moves to a higher address than it already had.
+* Two things follow. Writing a record can never land on the bytes of a record
+  not yet moved, so the rewrite is safe in place with no scratch copy of the
+  store. And each table entry can be updated the instant its record lands, so
+  the store is consistent at every point: an interrupted compaction leaves some
+  records moved and the rest where they were, and re-running finishes the job.
+* Records that are already in the right place are skipped, so a store with a
+  small amount of garbage near the end is cheap to compact.
+* Chunk files past the new frontier are closed and deleted.
+* `stats()` reports `live_bytes` against `chunk_bytes` so callers can decide when
+  it is worth running. Note the difference is an upper bound: a record never
+  straddles a chunk boundary, so up to one record's worth per chunk is slack
+  that compaction cannot remove.
+
+**5. File access.**
+* Metadata: a 192 MiB file created at init, accessed by mmap thereafter.
+* Chunk data: 4 GiB files created at init, accessed concurrently through
+  `readPositionalAll`/`writePositionalAll`.
+
+**Concurrency contract.** `gather` and `append` are safe to call concurrently
+from many threads. `delete` and `compact` mutate the table in ways a lock-free
+reader would observe half-applied, so they require exclusive access to the
+store.
+
+## Distributed extension
+**Not implemented yet; architecture noted in advance.**
+
+1. A cluster layer is added; a node becomes a shard holding several chunks.
+2. An indirection table maps a global path to a specific node, so the index path
+   becomes `idx → indirection[idx] = node, lidx → offset[lidx]`.