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

Type	Topology	Best for
"cpu"	Many processes, 1 thread each	NumPy/Numba math, xarray reductions, and reading NetCDF/HDF5
"io"	1 process, 8–16 threads	Zarr, object storage (S3), HTTP — libraries that release the GIL
"mixed"	Processes with 2 threads each	Pipelines that both read and compute
"gpu"	1 process per GPU, up to 8 threads	CuPy/RAPIDS CUDA workloads
"auto"	Inferred from dataset	Let dask_setup decide based on your data

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

Parameter	Default	Description
workload_type	"io"	Worker topology: "cpu", "io", "mixed", "gpu", "auto"
max_workers	all cores	Hard cap on worker count
reserve_mem_gb	auto	Memory held back for OS/cache (GiB): 20% of RAM, clamped to [4, 50]
max_mem_gb	total RAM	Upper bound on Dask's total memory use
dashboard	True	Start dashboard and print SSH tunnel hint
profile	None	Named config profile
config	None	Pre-built DaskSetupConfig object — same layer as profile; if both are given, profile wins
mode	"interactive"	"interactive" uses the current allocation and never submits jobs (a single node or no job at all is a plain LocalCluster); "auto" picks "pbs"/"slurm" in batch jobs, which submit worker jobs; also "local", "pbs", "slurm"
multi_node_config	None	MultiNodeConfig for PBS/SLURM multi-node jobs (v2.0)

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")

Profile	Workload	Reserve	Notes
climate_analysis	cpu	60 GB	Heavy compute with large arrays
zarr_io_heavy	io	40 GB	Many Zarr files
development	mixed	8 GB	Local testing, 2 workers max
production	mixed	80 GB	Adaptive, dashboard off
interactive	mixed	20 GB	Jupyter notebooks, 4 workers

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, including an API reference generated from the docstrings, is at 21centuryweather.github.io/dask_setup:

Page	What's covered
Configuration	DaskSetupConfig, profiles, site-wide profiles, profile inheritance, CLI, JSON Schema
Multi-Node	MultiNodeConfig, PBS/SLURM cluster setup, GPU topology, shared temp dirs, dask-setup submit
IO-Optimization	recommend_chunks, recommend_io_chunks, Zarr v3, Kerchunk, Parquet/Arrow, storage-aware chunking
Benchmarking	benchmark_config, scaling_analysis, chunk_impact, dask-setup benchmark
Internals	Resource detection, topology decisions, temp/spill routing, module layout
Troubleshooting	Common errors, OOM, multi-node issues, migration guide
User Feedback	Questions from users, the answers, and what changed as a result
API Reference	Every public function and class

---

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.