dask_setup

HPC-tuned Dask helpers for NCI Gadi and other PBS/SLURM systems. Wraps dask.distributed.LocalCluster + Client with sensible defaults for single-node jobs, and extends to multi-node PBS/SLURM clusters via dask-jobqueue in v2.0.

pip install dask-setup              # single-node
pip install dask-setup[multinode]   # + PBS/SLURM via dask-jobqueue

conda install -c conda-forge dask-setup   # or via conda-forge

Quick Start

from dask_setup import setup_dask_client

# Pick a workload type and go
client, cluster, dask_tmp = setup_dask_client(mode="interactive", workload_type="cpu")   # heavy compute
client, cluster, dask_tmp = setup_dask_client(mode="interactive", workload_type="io")    # Zarr / object storage
client, cluster, dask_tmp = setup_dask_client(mode="interactive", workload_type="mixed") # both

dask_tmp is the path to the spill/temp directory (on $PBS_JOBFS if available). Pass it to Rechunker, Zarr, or anywhere else you want fast local I/O.

For more control, pass a DaskSetupConfig object or a named profile:

from dask_setup import setup_dask_client, DaskSetupConfig

config = DaskSetupConfig(
    workload_type="cpu",
    max_workers=8,
    reserve_mem_gb=32.0,
    spill_compression="lz4",
)
client, cluster, dask_tmp = setup_dask_client(config=config)

# Or use a built-in profile
client, cluster, dask_tmp = setup_dask_client(profile="climate_analysis")

Multi-node (v2.0):

from dask_setup import setup_dask_client, MultiNodeConfig

client, cluster, shared_tmp = setup_dask_client(
    mode="auto",   # detects PBS_JOBID / SLURM_JOB_ID, falls back to local
    multi_node_config=MultiNodeConfig(
        workers_per_node=4,
        cores_per_worker=12,
        mem_per_worker_gb=32.0,
        walltime="04:00:00",
    ),
)

Workload Types

Reading NetCDF? Use "cpu", not "io"

The split is not "I/O vs compute" — it is whether the library releases the GIL and is thread-safe. "io" runs one process with many threads, which only wins when threads can genuinely overlap.

HDF5 (underneath NetCDF4) is usually not built thread-safe, so xarray routes every NetCDF read through a single process-wide lock. Under "io" your threads do not become parallel readers — they queue on that one lock, one at a time, and zlib decompression is serialised inside it too. Separate processes each get their own lock, so "cpu" is what actually parallelises NetCDF reads.

Zarr is the opposite: it decompresses via numcodecs, which releases the GIL, and has no global lock, so threads scale well. Opening, concatenating and writing NetCDF is a "cpu" job even though it feels like I/O.

Rule of thumb: NetCDF in → "cpu". Zarr or object storage in → "io". Pass ds= and dask_setup will warn if you get this pair the wrong way round.

Also avoid open_mfdataset(..., parallel=True) together with "io": concurrent metadata reads can kill the worker outright, after which Dask silently recomputes the lost tasks and the job appears to hang rather than fail.

Key Parameters

Settings are layered, lowest to highest:

library defaults  <  config= or profile=  <  explicit keyword arguments

A parameter you leave unset inherits from the layer below. Passing a value always overrides, even when that value happens to equal the default — so reserve_mem_gb=50.0 overrides a profile that says 40.0.

config= and profile= share a layer rather than stacking: pass one or the other, and use explicit keyword arguments for the differences.

Common Patterns

Big xarray reductions (CPU-bound)

client, cluster, dask_tmp = setup_dask_client("cpu", reserve_mem_gb=60)
ds = ds.chunk({"time": 240, "y": 512, "x": 512})
out = ds.mean(("y", "x")).compute()

Opening many NetCDF/Zarr files (I/O-bound)

client, cluster, dask_tmp = setup_dask_client("io", reserve_mem_gb=40)
ds = xr.open_mfdataset(files, engine="netcdf4", chunks={}, parallel=True)

Rechunking with Rechunker (spill stays on fast local storage)

client, cluster, dask_tmp = setup_dask_client("cpu", reserve_mem_gb=60)
plan = rechunker.rechunk(
    ds.to_array().data,
    target_chunks={"time": 240, "y": 512, "x": 512},
    max_mem="6GB",
    target_store="out.zarr",
    temp_store=f"{dask_tmp}/tmp_rechunk.zarr",
)
plan.execute()

Writing to Zarr in time windows

client, cluster, dask_tmp = setup_dask_client("cpu", max_workers=1, reserve_mem_gb=50)
ds.isel(time=slice(0, 240)).to_zarr("out.zarr", mode="w", consolidated=True)
for start in range(240, ds.sizes["time"], 240):
    stop = min(start + 240, ds.sizes["time"])
    ds.isel(time=slice(start, stop)).to_zarr(
        "out.zarr", mode="a", region={"time": slice(start, stop)}
    )

Built-in Profiles

Use profile= to load a named configuration:

client, cluster, dask_tmp = setup_dask_client(profile="climate_analysis")

See Configuration for how to create and save your own profiles.

Dashboard

With dashboard=True (default), the cluster prints an SSH tunnel command:

Dask dashboard: http://127.0.0.1:<PORT>/status
Tunnel from your laptop:
  ssh -N -L 8787:<COMPUTE_HOST>:<PORT> gadi.nci.org.au
Then open: http://localhost:8787

The login host is inferred from the compute node's DNS domain, so it is correct at other sites too. Override it with $DASK_SETUP_LOGIN_HOST if the guess is wrong:

export DASK_SETUP_LOGIN_HOST=login.mycluster.edu

CLI

# Profile management
dask-setup list                                       # list all profiles
dask-setup show climate_analysis                      # show profile details
dask-setup create my_profile --from-profile zarr_io_heavy
dask-setup validate my_profile
dask-setup export climate_analysis -o profile.yaml
dask-setup import https://example.com/team.yaml
dask-setup delete my_profile

# Synthetic benchmark
dask-setup benchmark --profile development --size small --operation mean

# Generate a PBS/SLURM job script (v2.0)
dask-setup submit my_analysis.py --scheduler pbs \
    --workers-per-node 4 --cores-per-worker 12 --mem-per-worker 32 \
    --walltime 04:00:00 --queue normal --project ab01

Documentation

Full documentation lives in the GitHub wiki:

Installation

From PyPI:

pip install dask-setup

From conda-forge:

conda install -c conda-forge dask-setup

For multi-node PBS/SLURM support:

pip install dask-setup dask-jobqueue
# or
conda install -c conda-forge dask-setup dask-jobqueue

For GPU workloads (CuPy auto-detection):

pip install dask-setup cupy-cuda12x   # match your CUDA version

Contributing

Bug reports, feature requests, and pull requests are welcome. Please include tests and, for performance changes, benchmarks.

License

Apache-2.0 — see LICENSE for details.