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.

[jder]

Oh, one other thing about this: it's probably worth refreshing the torch docs in repo after this lands to reflect the (new, much smaller) resource requirements and resulting torch flexibility in general (e.g. preemptible compute is much more available with smaller requests)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9cd36b1bdc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Reject derived channels before canonicalization

When loading.type='rust' is used with a derived boundary configuration—including the shipped configs/data/om4.yaml—this check runs in prepare() only after BaseDataSourceConfig._build_source() calls canonicalize_datasets(). If the store does not already contain hfds_anomalies, canonicalization computes the anomaly over the full hfds time series before this block discards the result and raises; for quarter-degree data, a simple unsupported-configuration error therefore triggers a multi-gigabyte read and calculation. Validate the configured layout before opening or canonicalizing the datasets.

AGENTS.md reference: AGENTS.md:L303-L307

Useful? React with 👍 / 👎.