Add opt-in Rust OM4 data loader

[jder]

Summary

• Add an opt-in local data.loading.type: rust path for flat and compact OM4 Zarr stores behind the new CanonicalDataset / TrainingShard boundaries.
• Keep Python responsible for canonical semantics, sampling, DDP schedules, masking, normalization, and batch contracts while Rust owns persistent native reads, rollout-wide plane deduplication, bounded prefetch, pinned-buffer reuse, and explicit lifecycle cleanup.
• Present the change as five reviewable commits: canonical datasets, deterministic schedules, training-shard/batch preparation, explicit batch compatibility, and one atomic Rust backend integration.

Why

The current CPU loader can spend minutes in xarray/Python data-loading cold paths and shows large periodic data-wait stalls during realistic quarter-degree training. The Rust path keeps local Zarr readers open, loads each unique time/channel plane once, reuses pinned host buffers, and overlaps host reads plus CUDA batch preparation with model execution without changing sampler or dataset_id semantics.

Scope

This is opt-in and local-filesystem only. It supports training and validation for flat OM4 and OM4-compact. S3, LLC-specific canonicalization, inference, derived seasonal-climatology/anomaly channels, and changing heterogeneous batch semantics remain out of scope.

Quarter-degree two-GPU speed check

• Full v2-highres model, 84M params, qdeg OM4 v2, two RTX6000s on Torch gr101, batch size 1, gradient accumulation 4, rollout steps [4], boundary vars tauuo,tauvo,hfds.
• CPU epoch 1 came from job 14115987; the attempted two-epoch CPU/Rust job was stopped because CPU validation and epoch-2 cold loading would not leave time for the Rust half in the one-hour allocation.
• Rust epoch 1 came from corrected job 14117615, submitted immediately afterward on the same node with the same container and NCCL_P2P_DISABLE=1.

Loader	Train epoch total	Mean iter excluding step 0	Mean iter excluding step 0 and data waits >=1s	Validation total	Peak CPU / GPU memory
CPU	22.31 s/it	14.59 s	4.12 s	81.61 s/it	41.7 GB / 49.1 GB
Rust	4.13 s/it	3.24 s	3.21 s	0.59 s/it	16.3 GB / 53.9 GB

This is a 5.4x train-epoch improvement and 139x validation improvement for this one-epoch qdeg run. The non-stall training number is 1.3x faster, so the main win is removing CPU-loader cold loads and periodic stalls rather than changing steady GPU compute.

Validation

• Focused native/canonical/sampler suite: 85 passed.
• Standalone native-extension suite: 23 passed.
• Rust format, clippy, and unit tests: passed; 3 Rust unit tests.
• All pre-commit checks, including mypy and schema validation: passed.
• Broad non-manual/non-CUDA local run: 349 passed, 2 skipped, 10 xfailed; six setup errors were the known local FlashAttention libcudart.so.12 environment issue.

Notes for review

• The native OM4 canonical reader is composed: xarray remains the semantic/static/reference delegate, while the native capability serves optimized batch plane reads. No canonical manifest is introduced.
• Compatibility groups preserve first-seen dataset order, keeping DDP schedules rank-stable without changing process-local dataset_id semantics.
• Native physical rows are mapped through the store time coordinate, including when a canonical dataset is sliced before native decoration.
• Rust iterators use weak loader ownership for deterministic early-exit cleanup, and Trainer explicitly tears down training, validation, and inference workers.
• The generic loader factory and native loader constructor switch together because splitting that change would require a temporary TorchTrainDataset compatibility adapter solely to delete it in the next commit.