Dask_setup: HPC-tuned Dask for Gadi, Setonix and other PBS/SLURM systems

General
, ,
1

Hi all,

I’m excited to share dask_setup, an open-source Python package I’ve developed at the ARC Centre of Excellence for the Weather of the 21st Century. It makes running Dask on NCI Gadi, and on other PBS/SLURM systems such as Pawsey’s Setonix, as simple as possible.

Here are 10 things it can do for you:

  1. One-line cluster setup. setup_dask_client() detects your PBS/SLURM environment and launches a correctly configured Dask cluster, with no manual tuning needed.
  2. Automatic resource detection. It reads your job’s allocated CPUs and memory from the PBS/SLURM environment, so your cluster never tries to use more than your job was given.
  3. Smart worker topology. It picks the right balance of processes and threads for your workload:
    • "cpu" for compute and for reading NetCDF/HDF5
    • "io" for Zarr and object storage
    • "mixed" for pipelines that do both
    • "gpu" for CuPy/RAPIDS workloads
    • "auto" to let it decide based on your dataset
  4. Fast local spill. Dask spill files go to node-local scratch ($PBS_JOBFS on Gadi), which keeps spill I/O off the shared Lustre filesystem.
  5. Chunk size recommendations. recommend_chunks(ds) analyses your xarray dataset and the cluster’s memory to suggest chunk sizes. This reduces OOM errors and task-graph overhead.
  6. Multi-node support. Multi-node Dask clusters are launched via dask-jobqueue through the same simple API as single-node setup. There’s also a dask-setup submit CLI that generates the job script for you.
  7. Configuration profiles. It ships with built-in profiles (e.g. climate_analysis, zarr_io_heavy). You can also save your own as YAML and share them across scripts or with your team.
  8. Benchmarking tools. chunk_impact() and scaling_analysis() let you systematically test chunk sizes and worker counts to find the best setup for your data.
  9. Dynamic memory tuning. tune_memory_thresholds() watches how much spilling happens at runtime and adjusts Dask’s memory thresholds to suit.
  10. Dashboard access on HPC. It prints a ready-to-paste SSH tunnel command, so you can open the Dask dashboard in your local browser from inside a batch job.

Installation

bash

pip install dask-setup                    # single-node
pip install "dask-setup[multinode]"       # + PBS/SLURM multi-node via dask-jobqueue
# or
conda install -c conda-forge dask-setup

On NCI Gadi, it’s also available in the latest xp65 analysis environments:

bash

module use /g/data/xp65/public/modules/
module load conda/analysis3-26.10

Thanks @rbeucher for adding it.

Quick start

python

from dask_setup import setup_dask_client

client, cluster, dask_tmp = setup_dask_client(workload_type="cpu")
# ... your analysis here ...
client.close()
cluster.close()

Tip: if you’re reading NetCDF, use "cpu" rather than "io". HDF5 reads are serialised behind a single lock within a process, so parallelism has to come from multiple processes. The docs explain this in more detail.

Documentation: https://21centuryweather.github.io/dask_setup/ It covers configuration, multi-node, I/O optimisation, benchmarking and troubleshooting. There’s an intro notebook at Getting started with dask_setup — dask_setup 2.3.0 and other example notebooks for different situations.

Bug reports, feature requests and PRs are very welcome via GitHub Issues. I’m also happy to answer questions here. If there’s interest, I’m keen to run an online training session, so let me know below.

2

This sounds amazing @sam.green ! I would be interested in an online training session.