Prototype TensorStore inside the native OM4 batch pipeline

[fomo-bot]

This draft lets data.loading.type: tensorstore use the rollout-wide plane deduplication, pinned-buffer reuse, host prefetch, and CUDA batch preparation from #800. Xarray remains the semantic delegate; TensorStore reads directly into the shared pipeline's output buffers without requiring the Rust extension.

Stacked on #800 (u/jder/rust-loader) so the review shows only the prototype integration. Related to the comparison in #871; this uses direct TensorStore reads rather than #871's Xarray backend option.

What is included

• A small Om4IoRuntime / Om4PlaneReader interface used by the existing Rust implementation and the new TensorStore reader.
• Local float32 flat and compact OM4 Zarr v2 support, shared native concurrency settings, and the existing sampler and buffer-lifetime contracts.
• Explicit rejection of unsupported CF encodings; outstanding reads drain before an error releases a caller-owned buffer.
• CPU/CUDA parity coverage, including compressed fixtures, sliced dates, channel/level order, overlapping rollouts, masks, and NaNs. Tests prohibit Rust extension loading on the TensorStore path.
• A reproducible scripts/benchmark_native_loading.py harness and docs/tensorstore-native-loader.md walkthrough.
• Correction of two stale fixture paths in Add opt-in Rust OM4 data loader #800's existing trainer integration tests.

Local experiment

GB10, uncompressed one-degree OM4 (180 × 360), batch size 1, history 1, four rollout steps, 77 prognostic plus three boundary channels. Median of three warmed 16-batch passes through the configuration and loader factories:

Native reader	GPU-ready batch	Peak loader GPU allocation
Rust	23.02 ms	1,031 MiB
TensorStore	26.88 ms	1,031 MiB

All 16 prepared batches matched the CPU reference for each reader. These are loader-only measurements, including transfer and preprocessing, without model execution. They do not establish full-training speedups or equivalence on compressed quarter-degree / multi-GPU workloads.

Validation

• Focused native/config/sampler CPU suite: 141 passed.

• Native-loader CUDA suite: 7 passed.

• uvx pre-commit run --all-files: passed.

• Clean --no-dev --extra tensorstore installation: Rust extension absent; local OM4 native read passed.

• Broad non-manual/non-CUDA suite: 553 passed, 2 skipped, 10 xfailed; 75 deselected.

• GitHub CI at 4e466dc6: CPU, GPU, data tests, pre-commit, Rust extension tests on x86_64/ARM64, and the x86_64 container build passed. All workflows and report jobs reached terminal success.

• ARM64 container and separate container CPU/GPU test jobs were skipped by workflow conditions.

Draft scope

The pipeline retains its existing rust_data.py module and RustTrainDataLoader name to keep this comparison focused. Remote storage, LLC, native inference, derived anomaly channels, broad renaming, and a decision to replace Rust are outside this draft. The next experiment is matched compressed quarter-degree training with epoch time, data-wait, and host/device memory measurements.