Xarray GPU optimization #771

negin513 · 2025-05-01T05:01:37Z

Contributors: @negin513, @weiji14 , @TomAugspurger , @maxrjoes, @akshaysubr, @kafitzgerald

vercel · 2025-05-01T05:01:41Z

@negin513 is attempting to deploy a commit to the xarray Team on Vercel.

A member of the Team first needs to authorize it.

for more information, see https://pre-commit.ci

TomAugspurger

Thanks for writing this up!

src/posts/gpu-pipeline/index.md

Co-authored-by: Tom Augspurger <[email protected]>

src/posts/gpu-pipeline/index.md

dcherian · 2025-05-08T17:10:35Z

src/posts/gpu-pipeline/index.md

+  - name: Katelyn Fitzgerald
+    github: kafitzgerald
+
+summary: 'How to accelerate AI/ML workflows in Earth Sciences with GPU-native Xarray and Zarr.'


Can we make this more direct? "X% speedup" or "XMBps throughput"?

src/posts/gpu-pipeline/index.md

dcherian · 2025-05-08T17:16:36Z

src/posts/gpu-pipeline/index.md

+(TODO ongoing work) Eventually with this [cupy-xarray Pull Request merged](https://github.com/xarray-contrib/cupy-xarray/pull/70) (based on earlier work at https://xarray.dev/blog/xarray-kvikio), this can be simplified to:
+
+```python
+import cupy_xarray
+
+ds = xr.open_dataset(filename_or_obj="/tmp/air-temp.zarr", engine="kvikio")
+assert isinstance(ds.air.data, cp.ndarray)
+```


This could go in a future work section at the end

Yeah, I'm not sure if this API is feasible or even desirable (have tried to implement this in xarray-contrib/cupy-xarray#70, but no luck yet patching the buffer protocol). So ok to move this towards the end.

src/posts/gpu-pipeline/index.md

dcherian · 2025-05-08T17:18:00Z

src/posts/gpu-pipeline/index.md

+- Consider using GPU Direct Storage (GDS) for optimal performance, but be aware of the setup and configuration required.
+- GPU Direct Storage (GDS) can be an improvement for data-intensive workflows, but requires some setup and configuration.
+- NVIDIA DALI is a powerful tool for optimizing data loading, but requires some effort to integrate into existing workflows.
+- GPU-based decompression is a promising area for future work, but requires further development and testing.


dcherian · 2025-05-08T17:19:58Z

src/posts/gpu-pipeline/index.md

@@ -0,0 +1,223 @@
+---
+title: 'Accelerating AI/ML Workflows in Earth Sciences with GPU-Native Xarray and Zarr (and more!)'


Suggested change

title: 'Accelerating AI/ML Workflows in Earth Sciences with GPU-Native Xarray and Zarr (and more!)'

title: 'GPU-Native Earth Science AI/ML Workflows Xarray, Zarr, DALI, and nvcomp'

better SEO this way?

src/posts/gpu-pipeline/index.md

weiji14 · 2025-05-08T23:20:52Z

src/posts/gpu-pipeline/index.md

+(TODO ongoing work) Eventually with this [cupy-xarray Pull Request merged](https://github.com/xarray-contrib/cupy-xarray/pull/70) (based on earlier work at https://xarray.dev/blog/xarray-kvikio), this can be simplified to:
+
+```python
+import cupy_xarray
+
+ds = xr.open_dataset(filename_or_obj="/tmp/air-temp.zarr", engine="kvikio")
+assert isinstance(ds.air.data, cp.ndarray)
+```


Yeah, I'm not sure if this API is feasible or even desirable (have tried to implement this in xarray-contrib/cupy-xarray#70, but no luck yet patching the buffer protocol). So ok to move this towards the end.

src/posts/gpu-pipeline/index.md

for more information, see https://pre-commit.ci

weiji14

Awesome work, this is coming along really nicely already! Just some minor nitpicks, but hope that we can publish this next month!

src/posts/gpu-pipeline/index.md

kafitzgerald

Thanks so much for putting this together!

Mostly just a few minor suggestions from my end beyond the existing comments / questions.

src/posts/gpu-pipeline/index.md

kafitzgerald · 2025-06-02T15:53:27Z

src/posts/gpu-pipeline/index.md

+
+## TL;DR
+
+Earth science AI/ML workflows are often bottlenecked by slow data loading, leaving GPUs underutilized while CPUs struggle to feed large climate datasets like ERA5. In this blog post, we discuss how to build a GPU-native pipeline using Zarr v3, CuPy, KvikIO, and NVIDIA DALI to accelerate data throughput. We walk through profiling results, chunking strategies, direct-to-GPU data reads, and GPU-accelerated preprocessing, all aimed at maximizing GPU usage and minimizing I/O overhead.


Suggested change

Earth science AI/ML workflows are often bottlenecked by slow data loading, leaving GPUs underutilized while CPUs struggle to feed large climate datasets like ERA5. In this blog post, we discuss how to build a GPU-native pipeline using Zarr v3, CuPy, KvikIO, and NVIDIA DALI to accelerate data throughput. We walk through profiling results, chunking strategies, direct-to-GPU data reads, and GPU-accelerated preprocessing, all aimed at maximizing GPU usage and minimizing I/O overhead.

Earth science AI/ML workflows are often limited by slow data loading, leaving GPUs underutilized while CPUs struggle to feed large climate datasets like ERA5. In this blog post, we discuss how to build a GPU-native pipeline using Zarr v3, CuPy, KvikIO, and NVIDIA DALI to accelerate data throughput. We walk through profiling results, chunking strategies, direct-to-GPU data reads, and GPU-accelerated preprocessing, all aimed at maximizing GPU usage and minimizing I/O overhead.

Not committed to this - just trying to vary the language a bit.

src/posts/gpu-pipeline/index.md

netlify · 2025-06-12T07:39:41Z

✅ Deploy Preview for xarraydev ready!

Name	Link
🔨 Latest commit	`17352a3`
🔍 Latest deploy log	https://app.netlify.com/projects/xarraydev/deploys/684aa60b8599730008229284
😎 Deploy Preview	https://deploy-preview-771--xarraydev.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

To edit notification comments on pull requests, go to your Netlify project configuration.

for more information, see https://pre-commit.ci

Co-authored-by: Max Jones <[email protected]>

for more information, see https://pre-commit.ci

Co-authored-by: Wei Ji <[email protected]>

Co-authored-by: Katelyn FitzGerald <[email protected]>

for more information, see https://pre-commit.ci

Co-authored-by: Max Jones <[email protected]>

negin513 added 5 commits April 30, 2025 18:32

first draft

f69796c

adding headers

32a8e32

adding baseline image

d2f7e0d

update blog post

d23c74f

update chunking

64b45e1

[pre-commit.ci] auto fixes from pre-commit.com hooks

95e5d65

for more information, see https://pre-commit.ci

maxrjones mentioned this pull request May 4, 2025

Publish Xarray blog post on NCAR hackathon NASA-IMPACT/veda-odd#166

Open

TomAugspurger reviewed May 5, 2025

View reviewed changes

Apply suggestions from code review

b52e1e7

Co-authored-by: Tom Augspurger <[email protected]>

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

src/posts/gpu-pipeline/index.md Outdated Show resolved Hide resolved

dcherian reviewed May 8, 2025

View reviewed changes

weiji14 reviewed May 8, 2025

View reviewed changes

negin513 added 7 commits May 12, 2025 13:46

moving profiling_screenshot1.png over

a2416b3

adding profiling screenshot

67f9c29

update

5d168be

screenshot 1 added

8c292d3

moving baseline png

e9195c8

adding pngs for the plots

304acb0

updates

eccf86c

[pre-commit.ci] auto fixes from pre-commit.com hooks

d0b4856

for more information, see https://pre-commit.ci

weiji14 reviewed May 24, 2025

View reviewed changes

kafitzgerald reviewed Jun 2, 2025

View reviewed changes

negin513 added 3 commits June 12, 2025 01:37

updates & clean ups of the the blog post

4d8124d

improved performance chart

7ab3039

merge conflict

aac8647

pre-commit-ci bot and others added 23 commits June 12, 2025 07:40

[pre-commit.ci] auto fixes from pre-commit.com hooks

7e31db2

for more information, see https://pre-commit.ci

comment addressed

9135c67

Update src/posts/gpu-pipeline/index.md

a7ddc66

Co-authored-by: Max Jones <[email protected]>

update dali

3728b1a

max's comments

75fc193

update benchmark

257a4b5

[pre-commit.ci] auto fixes from pre-commit.com hooks

ca5ff74

for more information, see https://pre-commit.ci

Update src/posts/gpu-pipeline/index.md

50773a9

Co-authored-by: Wei Ji <[email protected]>

update to blogpost

6f666a9

Update src/posts/gpu-pipeline/index.md

8be9c11

Co-authored-by: Wei Ji <[email protected]>

Update src/posts/gpu-pipeline/index.md

c017e15

Co-authored-by: Wei Ji <[email protected]>

Update src/posts/gpu-pipeline/index.md

508af65

Co-authored-by: Wei Ji <[email protected]>

Update src/posts/gpu-pipeline/index.md

43e066a

Co-authored-by: Katelyn FitzGerald <[email protected]>

Update src/posts/gpu-pipeline/index.md

c842b44

Co-authored-by: Katelyn FitzGerald <[email protected]>

Update src/posts/gpu-pipeline/index.md

a51f276

Co-authored-by: Katelyn FitzGerald <[email protected]>

Update src/posts/gpu-pipeline/index.md

bc3cbed

Co-authored-by: Katelyn FitzGerald <[email protected]>

Update src/posts/gpu-pipeline/index.md

993f772

Co-authored-by: Katelyn FitzGerald <[email protected]>

Update src/posts/gpu-pipeline/index.md

a5c387e

Co-authored-by: Katelyn FitzGerald <[email protected]>

updates

6359e46

address comments

4355691

update thank you messages

b13036a

[pre-commit.ci] auto fixes from pre-commit.com hooks

b35db32

for more information, see https://pre-commit.ci

Update src/posts/gpu-pipeline/index.md

17352a3

Co-authored-by: Max Jones <[email protected]>

		@@ -0,0 +1,223 @@
		---
		title: 'Accelerating AI/ML Workflows in Earth Sciences with GPU-Native Xarray and Zarr (and more!)'

	title: 'Accelerating AI/ML Workflows in Earth Sciences with GPU-Native Xarray and Zarr (and more!)'
	title: 'GPU-Native Earth Science AI/ML Workflows Xarray, Zarr, DALI, and nvcomp'


		## TL;DR

		Earth science AI/ML workflows are often bottlenecked by slow data loading, leaving GPUs underutilized while CPUs struggle to feed large climate datasets like ERA5. In this blog post, we discuss how to build a GPU-native pipeline using Zarr v3, CuPy, KvikIO, and NVIDIA DALI to accelerate data throughput. We walk through profiling results, chunking strategies, direct-to-GPU data reads, and GPU-accelerated preprocessing, all aimed at maximizing GPU usage and minimizing I/O overhead.

Xarray GPU optimization #771

Are you sure you want to change the base?

Xarray GPU optimization #771

Uh oh!

Conversation

negin513 commented May 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vercel bot commented May 1, 2025

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcherian May 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dcherian May 8, 2025

Choose a reason for hiding this comment

Uh oh!

weiji14 May 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dcherian May 8, 2025

Choose a reason for hiding this comment

Uh oh!

dcherian May 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

weiji14 May 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

weiji14 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kafitzgerald left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kafitzgerald Jun 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

netlify bot commented Jun 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for xarraydev ready!

Uh oh!

Uh oh!

negin513 commented May 1, 2025 •

edited

Loading

dcherian May 8, 2025 •

edited

Loading

netlify bot commented Jun 12, 2025 •

edited

Loading