This talk by Hendrick Mite from Coiled discusses the powerful combination of Dask and Xarray for processing massive geoscience datasets. Hendrick demonstrates using a 100-worker Dask cluster on AWS to analyze 18 terabytes of the National Water Model data, calculating the average depth to soil per US county. The talk explains how Xarray provides labeled multi-dimensional arrays, while Dask offers general-purpose parallelism and scales computations across a cluster, including managed services like Coiled. Key technical improvements in Dask's array handling are highlighted, specifically a new P2P rechunking system that drastically reduces task overhead and memory footprint for all-to-all communication patterns, and an improved algorithm for downscaling operations that prevents chunk explosion. Finally, the speaker emphasizes the need for representative, large-scale, end-to-end geospatial benchmarks to guide future development in Dask array and invites the audience to contribute real-world workloads.
Processing Geoscience Data at Scale with Dask and Xarray
0:00 Introduction: Speaker Hendrick Mite from Coiled introduces Dask and Xarray as a "power couple" for large-scale geoscience workloads.
0:41 Demo Setup: A Dask cluster of 100 workers on AWS is used for the demonstration.
1:12 Data Used: The demonstration uses the National Water Model data from Noah, focusing on the depth to soil variable. The original dataset is 250 TB, subset to 18 TB for 2018-2020 for the demo.
2:20 Combining Data: A smaller 1 GB dataset is used to map geospatial points to US counties.
2:51 Computation: The core computation involves grouping the depth to soil data by county and calculating the mean over time and points within each county. The final result is only 210 MB.
3:25 Observing the Computation: The Dask dashboard and the Coiled platform dashboard are used to visualize the computation's progress, task count (initially 110,000 tasks), worker utilization, network transfer, and memory usage.
6:00 Computation Completion: The 18 TB computation finishes, yielding the final 200 MB result, visualized as a map of the US with average depth to soil per county.
6:48 Cost: The 18 TB processing cost approximately $3 in cloud computing.
Dask (7:09): A library for general-purpose parallelism in Python. It provides parallel APIs for NumPy (dask.array) and Pandas (dask.dataframe) and can express complex, dynamic workloads.
Xarray (8:56): NumPy but with labels (dimensions and coordinates). This makes data handling more intuitive (e.g., filtering by time labels instead of axis numbers). Xarray can group multiple arrays with shared dimensions.
10:39 Dask and Xarray Together: Xarray translates computations into a Dask computation graph, which Dask then executes on various cluster types (Kubernetes, AWS, Slurm, managed services like Coiled), enabling computations close to cloud-stored data.
11:39 Game Changer: Running computations on a cloud cluster near the data avoids transferring large datasets locally and provides significantly more computational power than a desktop.
Rechunking (12:29): Discusses the difficulty of rechunking (changing how a dataset is sliced) large geospatial data, especially all-to-all patterns (e.g., space-chunked to time-chunked).
Task-based Rechunking Issues (14:05): Legacy task-based rechunking created huge numbers of tiny intermediate tasks (Cartesian product of inputs/outputs), overloading the scheduler and limiting scalability (N log N tasks).
P2P Rechunking (14:37): A new Peer-to-Peer system circumvents the task-based paradigm for data transfer, using an out-of-band system. This reduces tasks to O(N) and maintains a constant memory footprint.
P2P Demo (15:38): A comparison demo shows P2P completing much faster with stable memory compared to task-based rechunking which spilled data at 100 GB scale and couldn't even materialize the graph at 1 TB scale.
Downscaling (18:56): Discusses issues with chunk explosion when performing operations like rolling means, leading to many small chunks and high overhead.
Improved Downscaling Algorithm (20:52): A reimplemented algorithm keeps the number of chunks mostly stable, drastically improving performance.
Need for Benchmarks (22:10): Unlike tabular data (Pandas/Dask Dataframe), there's no standard industry benchmark for Dask/Xarray geospatial workloads.
Coiled's Benchmarking Suite (23:16): Coiled is building a suite focusing on end-to-end, large-scale, and representative real-world geospatial workloads.
Call for Contributions (24:10): The speaker invites geoscientists to propose their critical workloads to help guide Dask array improvements.
24:48 Conclusion: Dask and Xarray are powerful for terabyte-scale computations. Work is ongoing to unlock larger scales (tens/hundreds of TB). Geospatial benchmarks are being built to focus future development, and community input is requested.
Xarray for Multi-dimensional DataFrames (25:51): Xarray is built for grid data (arrays), not vector data frame computations like GeoPandas. There's a Dask integration for GeoPandas available.
Is P2P Default? (27:10): P2P is the default for all-to-all rechunking in recent versions (since approx. March/April 2024), but a heuristic decides when it makes sense versus task-based for adjacent communications.
Xarray Sparse Arrays (28:14): Speaker is unsure and refers to Xarray documentation.
Batching Tasks (28:40): The demo did not batch tasks; it showed the full large computation graph produced by Dask. Managing large graphs is a key future focus area.
29:17 Closing: Thanks the audience and encourages joining Discord for further discussion.