pytorch-reader: ZIP tensor reads serialize on a Mutex around the archive; 8 threads read no faster than 1

Summary

ZipSource holds Mutex<ZipArchive<BufReader<File>>> and takes the lock for the whole of every read_storage, so Tensor::read calls from different threads run one at a time. A loader that reads tensors in parallel gets exactly single-thread throughput from a ZIP checkpoint, which is every checkpoint torch.save has written since 1.6.

The legacy container does not have this: LegacySource locks only to look up the layout, then opens its own File for the read. The two containers also differ in a smaller way as a result: a ZIP reader keeps the file descriptor open from new until drop and reads fine after the file is unlinked, a legacy reader reopens by path and fails (test_os_errors_keep_their_kind pins that).

Measured

64 tensors of 1 M f32 (244 MB, stored entries, page-cached), best of 5, Apple M-series, 16 cores. Each row reads every tensor once.

No scaling at all through the reader; the positional-read baseline scales to about 4.7x. The 1-thread gap (22 vs 13 ms) is separate overhead in the ZIP path (by_name, the take adapter, the trailing-byte probe), noted but not the point here.

Fixture: torch.save({f"layer{i}.weight": torch.full((1_000_000,), float(i)) for i in range(64)}, "bench.pt").

use std::os::unix::fs::FileExt;
use std::time::Instant;
use pytorch_reader::PytorchReader;

fn main() {
    let path = std::env::args().nth(1).unwrap();
    let threads: usize = std::env::args().nth(2).map(|s| s.parse().unwrap()).unwrap_or(8);
    let reader = PytorchReader::new(&path).unwrap();
    let tensors: Vec<_> = reader.tensors().values().cloned().collect();
    for t in &tensors { t.read().unwrap(); } // warm the page cache

    let best = |f: &dyn Fn()| (0..5).map(|_| { let s = Instant::now(); f(); s.elapsed() }).min().unwrap();

    let parallel = best(&|| std::thread::scope(|s| {
        for chunk in tensors.chunks(tensors.len().div_ceil(threads)) {
            s.spawn(move || for t in chunk { std::hint::black_box(t.read().unwrap()); });
        }
    }));
    println!("Tensor::read, {threads} threads: {parallel:?}");

    let file = std::fs::File::open(&path).unwrap();
    let mut archive = zip::ZipArchive::new(std::fs::File::open(&path).unwrap()).unwrap();
    let ranges: Vec<(u64, usize)> = (0..tensors.len()).map(|i| {
        let e = archive.by_name(&format!("archive/data/{i}")).unwrap();
        (e.data_start().unwrap(), e.size() as usize)
    }).collect();
    let pread = best(&|| std::thread::scope(|s| {
        for chunk in ranges.chunks(ranges.len().div_ceil(threads)) {
            let file = &file;
            s.spawn(move || for &(off, len) in chunk {
                let mut buf = vec![0u8; len];
                file.read_exact_at(&mut buf, off).unwrap();
                std::hint::black_box(buf);
            });
        }
    }));
    println!("read_exact_at, {threads} threads: {pread:?}");
}

Options

Smallest change, matches the legacy container. Keep the ZipArchive for the central directory (names, data_start, sizes, compression method), lock only to look those up, then do the I/O on a handle the read owns: a fresh File::open per read as LegacySource does, or positional reads on one shared File (FileExt::read_at on Unix, seek_read on Windows), which needs no lock and no reopen. Stored entries are the whole of torch.save output, so this covers the real case; a deflated entry can keep the current locked stream path.

Keep the CRC behavior. Today a full-entry read reaches the entry's end and the zip crate checks the CRC (test_corrupted_storage_is_invalid_data pins "Invalid checksum" as InvalidData). A positional read bypasses that, so a whole-entry read would run crc32fast (already in the tree through zip) over the bytes against the directory's CRC. A windowed read skips it today too, so nothing changes there.

Either way the fd-lifetime difference between ZIP and legacy goes away, or gets documented if one of them is the intended behavior.

Why I think it matters before publishing

Lazy tensors are the crate's pitch, and parallel materialization is the obvious thing to do with them. Nothing in the API says the reads serialize, and a caller cannot work around it (the Mutex is inside Tensor's closure). Fixing it later is non-breaking, but the first release is where "does it scale" gets decided by users.