Add optional TensorStore backend for Xarray data loading

[alxmrs]

Summary

• add an opt-in xarray-tensorstore backend to the existing Xarray/CanonicalSource data path
• start selected TensorStore reads asynchronously before Xarray materializes a batch
• preserve each variable's native Zarr chunks when the Dask inference path is requested
• support local Zarr v2, AWS S3, and anonymous HTTP/S3-compatible endpoints (including the public OSN stores)
• keep zarr-python as the default and ship TensorStore as the optional samudra[tensorstore] extra

This gives Samudra a native-code Zarr option without changing the established Xarray data-loading boundary.

Enable it with:

data:
  xarray_backend: tensorstore

Local benchmark

Cached local Zarr v2 reads through _XarrayCanonicalReader, 25 measured iterations after an initial read:

Request	Zarr-Python median	TensorStore median	Speedup
3 channels × 2 times	1.795 ms	0.447 ms	4.0×
16 channels × 4 times	18.884 ms	3.723 ms	5.1×

The backends returned exactly equal arrays. These numbers measure hot local reads, not network throughput or full training-step performance.

I also exercised the complete DataConfig -> CanonicalSource -> read() path against the public two-degree OM4 OSN store. It returned a (2, 4, 90, 180) float32 batch successfully via TensorStore.

Limitations

The public xarray-tensorstore API cannot pass an explicit authenticated custom S3 endpoint through both its Zarr-Python metadata reader and TensorStore array reader. Those locations fail loudly with guidance to use an AWS s3:// location, anonymous HTTP endpoint, or the existing zarr-python backend.

Validation

• uv run pytest -m "not manual and not cuda" — 472 passed, 2 skipped, 10 xfailed
• uvx pre-commit run --all-files
• public OSN OM4 end-to-end smoke test

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9670f80b57

ℹ️ 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".

[alxmrs]

@codex may I have your review for this PR?

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Bravo.

Reviewed commit: a70d9c5f56

ℹ️ 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".

[alxmrs]

@fomo-bot how does this data loading implementation compare to #800? In your best estimate, what would you recommend? What are the tradeoffs in choosing one approach vs the other?

[alxmrs]

@fomo-bot in making your comparison, feel free to check out both or either branches to make spikes/prototypes to more fairly compare the approaches.

[fomo-bot]

I’d adopt #871 as the incremental backend option, while carrying forward #800’s batch planning and prefetch design. Before committing to maintaining the Rust reader, I’d compare TensorStore inside that same pipeline. I made that prototype, and the result is promising.

I compared the current heads (a70d9c5f and 5d519f54) using the same local OM4 data and GPU. This measured loading through preparation of a GPU-ready batch, including normalization and masking.

Implementation	Time per GPU-ready batch	Peak GPU allocation, loader only
Existing Zarr-Python loader	160 ms	725 MiB
#871: TensorStore through Xarray	109 ms	725 MiB
#800: Rust with its optimized pipeline	24 ms	1,031 MiB
Prototype: TensorStore inside #800’s pipeline	23 ms	1,031 MiB

These are medians of three 16-batch runs after warming the same reads: one GB10, local uncompressed one-degree Zarr, batch size 1, hist=1, four rollout steps, and 77 prognostic plus three boundary channels. The existing loaders used four persistent workers; #800 and the prototype used eight native read threads and two-batch host prefetch. All 16 checked output batches matched across implementations.

The prototype replaces #800’s native array reader with TensorStore reads into its existing pinned buffers. It retains #800’s planning, preprocessing, and CUDA prefetch. It is experimental code, not functionality currently provided by either PR.

These measurements do not establish full-training speedups or performance on compressed quarter-degree data, remote storage, or multi-GPU systems.

The tradeoffs are:

• Add optional TensorStore backend for Xarray data loading #871 is a small change with broad applicability. It accelerates reads while retaining the established Xarray and CanonicalSource path, including the inference route. It supports local stores, AWS S3, and public OSN endpoints; authenticated custom S3 endpoints remain unsupported. It preserves the existing worker, collation, and per-rollout-step loading costs. Implementation.

• Add opt-in Rust OM4 data loader #800 optimizes more of the pipeline. It reads overlapping input/label planes once across the rollout, reuses pinned buffers, transfers and preprocesses unique planes before gathering repeated positions, and prepares batches on a dedicated CUDA stream. The existing loader already performs GPU preprocessing; Add opt-in Rust OM4 data loader #800 reduces repeated work and changes its scheduling. The extra prefetched device batch costs GPU memory, as the table shows. Implementation and ownership details.

• Add opt-in Rust OM4 data loader #800 currently has narrower coverage. Its accelerated path supports local flat/compact OM4 training and validation. Remote stores, LLC, inference, and derived anomaly channels are outside its current scope. In particular, configurations using hfds_anomalies need attention. Its reported quarter-degree training experiment provides stronger evidence for that workload than this PR’s hot-read microbenchmark does. Scope and training results.

• Maintenance favors an upstream reader, with qualifications. TensorStore already supplies a native C++ implementation; choosing it avoids owning the Rust extension’s build, packaging, and reader integration. We would still own the batching and CUDA-buffer lifecycle code. Also, xarray-tensorstore depends on internal Xarray APIs and has limitations around CF-encoded data variables, so compatibility coverage remains necessary. Upstream limitations.

My interpretation of the prototype is that #800’s pipeline design contributes substantial value independently of Rust. On this workload, changing its reader to TensorStore preserved performance. That supports making its batching/prefetch machinery usable with either reader.

For immediate local OM4 throughput, #800 is the stronger implementation today. The next deciding experiment should compare Rust and the TensorStore prototype on the same compressed quarter-degree dataset and actual multi-GPU training run, measuring epoch time, data-wait stalls, and host/device memory. If TensorStore remains comparable there, I would favor it as the native reader and retain #800’s pipeline optimizations.

Debug info

• Repository: m2lines/Samudra
• Trigger: Add optional TensorStore backend for Xarray data loading #871 (comment)
• Request ID: 8b9a73d4-b2e3-44f7-af7b-c0c1e26995d4
• Accepted: true

[alxmrs]

@fomo-bot wow, this is a promising result and a nuanced new direction. Will you draft a new PR that showcases your prototype using tensorstore inside #800's pipeline?

[fomo-bot]

Opened draft PR #874, stacked on #800.



It includes the selectable TensorStore reader, shared batching/prefetch pipeline, parity tests, and a reproducible benchmark. The TensorStore path runs without the Rust extension.



Validation: 553 CPU tests passed, 7 CUDA tests passed, and all pre-commit checks passed. GitHub CI is running; I’ll follow up on its results.

Debug info

• Repository: m2lines/Samudra
• Trigger: Add optional TensorStore backend for Xarray data loading #871 (comment)
• Request ID: 7ecbade4-10c6-4629-8285-278e3137da54
• Accepted: true

[fomo-bot]

Draft PR #874 has passed CI at 4e466dc6: CPU, GPU, data tests, pre-commit, Rust extension tests on both architectures, and the x86_64 container build.



I updated its validation record, including the workflow-configured skips.

Debug info

• Repository: m2lines/Samudra
• Trigger: Add optional TensorStore backend for Xarray data loading #871
• Request ID: 85a0ee43-5b43-42b6-b162-d98f62224a33
• Accepted: true