Getting started with dask_setup

HPC-aware Dask clusters in one line, for NCI Gadi (PBS) and SLURM sites.

This notebook works through dask_setup from the one-line quick start down to profiles and multi-node job scripts. Every cell runs on a laptop; where behaviour differs on Gadi, the notebook simulates the PBS environment so you can see what would happen on a compute node.

#

Section

You will learn

1

Why dask_setup?

What goes wrong with a bare LocalCluster on HPC

2

Quick start

The one-line setup, and cleaning up

3

What it decided

Resource detection, worker topology, memory limits

4

Choosing a workload type

cpu vs io vs mixed, and the GIL

5

Context manager & reports

DaskClientContext, cluster_report()

6

Working with xarray

ds=, chunk recommendations, workload inference

7

Configuration & profiles

DaskSetupConfig, precedence, saved profiles, the CLI

8

Memory tuning

tune_memory_thresholds()

9

Beyond one node

MultiNodeConfig, PBS/SLURM job scripts

10

When things go wrong

Errors that tell you how to fix them

Prerequisites: pip install "dask-setup" xarray numpy (add dask-setup[multinode] for section 9 on a real cluster).

1. Why dask_setup?

A bare LocalCluster() on an HPC compute node usually makes three bad decisions:

  1. It sizes itself to the whole machine, not your job. A Gadi normalsr node has 104 cores, but if PBS gave you ncpus=48, psutil still reports 104. You oversubscribe the node and your job gets throttled or killed.

  2. It spills to the wrong disk. Spill files land in the current directory, which is often /scratch or /home on Lustre, the slowest and most contended filesystem on the machine. The node-local SSD at $PBS_JOBFS sits unused.

  3. It picks one topology for every workload. Reading NetCDF, crunching NumPy and writing Zarr want very different process/thread layouts.

setup_dask_client() reads the scheduler’s environment variables, routes spill to $PBS_JOBFS, and chooses the worker layout from the kind of work you tell it you’re doing.

import dask
import dask.array as da
import distributed

import dask_setup
from dask_setup import setup_dask_client

print(f"dask_setup  {dask_setup.__version__}")
print(f"dask        {dask.__version__}")
print(f"distributed {distributed.__version__}")

dask_setup  2.3.0
dask        2026.8.0
distributed 2026.8.0

2. Quick start

The whole API for most users is one call:

client, cluster, dask_tmp = setup_dask_client(workload_type="cpu", max_workers=4)

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 4 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~2.9 GiB
Compression: spill=auto | comm=False

It returns three things:

  • client: a normal dask.distributed.Client, so everything you already know about Dask works unchanged.

  • cluster: the underlying LocalCluster.

  • dask_tmp: the directory where workers spill to disk. On Gadi this is under $PBS_JOBFS.

With dashboard=True (the default) it also prints the dashboard address and, on a remote node, the exact ssh -N -L ... command to tunnel it back to your laptop.

max_workers=4 just keeps this demo small. On a real job, leave it out and dask_setup uses every core PBS/SLURM gave you.

⚠️ In a PBS/SLURM batch job, add mode="local". The default mode="auto" detects the batch job and switches to the multi-node backend (section 9), which submits more jobs.

client

Client

Client-ca797d9a-bd41-11f1-8d81-6045bd34f621

Connection method: Cluster object Cluster type: distributed.LocalCluster
Dashboard: http://127.0.0.1:39157/status

Cluster Info

 
print("Spill directory:", dask_tmp)

# Any Dask code now runs on this cluster
x = da.random.random((20_000, 20_000), chunks=(2_000, 2_000))
result = (x + x.T).mean().compute()
print(f"Result: {result:,.2f}")

Spill directory: /tmp/dask-3457

Result: 1.00

Always close the client and cluster when you’re done, otherwise worker processes stay behind in your job. Section 5 shows a context manager that does this for you.

client.close()
cluster.close()

3. What did it decide?

setup_dask_client() orchestrates a small pipeline of decision functions. You can call each one yourself to see (or debug) what it will do:

detect_resources()      SLURM env  →  PBS env  →  psutil
create_dask_temp_dir()  $PBS_JOBFS →  $TMPDIR  →  /tmp
decide_topology()       workload type + cores  →  n_workers × threads_per_worker
calculate_memory_spec() total RAM − reserve    →  memory limit per worker

3a. Resource detection

On a laptop this falls through to psutil:

from dask_setup.resources import detect_resources

res = detect_resources()
print(f"Detected via : {res.detection_method}")
print(f"Cores        : {res.total_cores}")
print(f"Memory       : {res.total_mem_bytes / 2**30:.1f} GiB")

Detected via : psutil
Cores        : 4
Memory       : 15.6 GiB

Now let’s pretend we’re inside a Gadi PBS job with ncpus=48, mem=190gb. PBS exports these as environment variables, and dask_setup reads them in preference to what the hardware reports.

Tip: use the same trick to test PBS behaviour from a login node or laptop: NCPUS=48 PBS_MEM=190gb PBS_JOBFS=/tmp python my_script.py

import os
from contextlib import contextmanager


@contextmanager
def simulated_env(**env):
    # Temporarily set environment variables, restoring the originals afterwards.
    saved = {k: os.environ.get(k) for k in env}
    os.environ.update(env)
    try:
        yield
    finally:
        for k, v in saved.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v


with simulated_env(NCPUS="48", PBS_MEM="190gb"):
    gadi = detect_resources()

print(f"Detected via : {gadi.detection_method}")
print(f"Cores        : {gadi.total_cores}")
print(f"Memory       : {gadi.total_mem_bytes / 2**30:.1f} GiB")

Detected via : PBS
Cores        : 48
Memory       : 177.0 GiB

3b. Worker topology

Given those 48 cores, here is the layout each workload type produces:

from dask_setup.topology import decide_topology

print(f"{'workload':<10}{'workers':>9}{'threads/worker':>16}{'processes':>11}")
print("-" * 46)
for wt in ["cpu", "mixed", "io"]:
    t = decide_topology(wt, total_cores=gadi.total_cores)
    print(f"{wt:<10}{t.n_workers:>9}{t.threads_per_worker:>16}{str(t.processes):>11}")

workload    workers  threads/worker  processes
----------------------------------------------
cpu              48               1       True
mixed            24               2       True
io                1              16      False

3c. Memory per worker

dask_setup holds back reserve_mem_gb for the OS, page cache and filesystem clients, then splits the rest evenly. Each worker gets a hard limit, and Dask’s spill thresholds (target 75%, spill 85%, pause 92%, terminate 98%) are measured against that limit.

from dask_setup.cluster import calculate_memory_spec

topo = decide_topology("cpu", total_cores=gadi.total_cores)
mem = calculate_memory_spec(gadi.total_mem_bytes, topo.n_workers, reserve_mem_gb=20)

print(f"Node memory      : {mem.total_mem_gib:.1f} GiB")
print(f"Reserved for OS  : {mem.reserved_mem_gb:.1f} GiB")
print(f"Usable by Dask   : {mem.usable_mem_gb:.1f} GiB")
print(f"Per worker (x{topo.n_workers}) : {mem.mem_per_worker_bytes / 2**30:.2f} GiB")

Node memory      : 177.0 GiB
Reserved for OS  : 20.0 GiB
Usable by Dask   : 157.0 GiB
Per worker (x48) : 3.27 GiB

4. Choosing a workload type

This is the most important choice you make, and the rule is about the GIL (Python’s Global Interpreter Lock), not about whether the work “feels like” I/O or compute.

workload_type

Layout

Use when the hot code…

Typical examples

"cpu"

many processes × 1 thread

holds the GIL

NetCDF / HDF5 reads, pure-Python loops, pandas apply

"io"

1 process × 8–16 threads

releases the GIL

Zarr reads/writes, NumPy/BLAS, cloud object storage

"mixed"

processes × 2 threads

does a bit of both

general analysis pipelines

"gpu"

1 process per GPU

runs on GPUs

CuPy / RAPIDS

"auto"

inferred

you pass ds= and want a guess

see section 6

⚠️ The NetCDF trap

It seems natural to call “reading lots of NetCDF files” an I/O workload. It isn’t, for Dask. The netCDF4/HDF5 C library serialises every read behind a single process-wide lock. With "io" you get one process and 16 threads, and only one of those threads can read at a time.

  • NetCDF / HDF5 input → "cpu" (each process has its own lock, so reads run in parallel)

  • Zarr input → "io" (Zarr’s codecs release the GIL, so threads really run concurrently)

dask_setup warns you if you pass a NetCDF-backed dataset with workload_type="io".

5. Context manager and run reports

DaskClientContext takes exactly the same arguments as setup_dask_client() and guarantees cleanup, even if your analysis raises halfway through. Prefer it in scripts and batch jobs.

DaskClientContext also logs a run summary on exit. cluster_report() gives you the same summary of the run (wall time, workers, memory, spill), which is handy at the end of a PBS job log.

import time

from dask_setup import DaskClientContext, cluster_report

start = time.monotonic()  # cluster_report measures against time.monotonic()
with DaskClientContext(workload_type="mixed", max_workers=2, dashboard=False) as (
    client,
    cluster,
    tmp,
):
    x = da.random.random((10_000, 10_000), chunks=(1_000, 1_000))
    std = x.std().compute()
    print(f"std = {std:.4f}")

    report = cluster_report(client, start_time=start)
    print(report.summary_line())

# Client and cluster are closed here, even if an exception was raised above

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 2 | threads/worker: 2 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~5.8 GiB
Compression: spill=auto | comm=False

std = 0.2887
wall=1.2s | workers=2 | mem=0.06 GiB/max

6. Working with xarray

Most dask_setup users are working with gridded climate/weather data in xarray. Pass your dataset as ds= and setup_dask_client() will:

  1. Validate its current chunks against per-worker memory (warning about OOM-sized or tiny chunks)

  2. Recommend chunk sizes for this cluster

  3. Return a 4-tuple: (client, cluster, dask_tmp, chunks)

Return arity: ds=None → 3 values, ds=<dataset> → 4 values. This is the one thing people trip over. Unpack accordingly.

Let’s build a synthetic “climate” dataset: 10 years of daily data on a 0.25° global grid, lazily. We chunk it one year per chunk, which is what you often get from xr.open_mfdataset over yearly files. Each chunk is ~1.4 GiB, far too big for comfortable parallel work.

import numpy as np
import pandas as pd
import xarray as xr

time_ = pd.date_range("2000-01-01", periods=3650, freq="D")
lat = np.linspace(-90, 90, 721)
lon = np.linspace(0, 359.75, 1440)

shape = (len(time_), len(lat), len(lon))
ds = xr.Dataset(
    {
        "tas": (("time", "lat", "lon"), da.zeros(shape, dtype="float32", chunks=(365, 721, 1440))),
        "pr": (("time", "lat", "lon"), da.zeros(shape, dtype="float32", chunks=(365, 721, 1440))),
    },
    coords={"time": time_, "lat": lat, "lon": lon},
)
print(f"Dataset size: {ds.nbytes / 2**30:.1f} GiB (lazy, nothing is allocated)")
chunk_shape = ds.tas.data.chunksize
print(f"Current chunk shape: {chunk_shape} = {np.prod(chunk_shape) * 4 / 2**30:.2f} GiB per chunk")

Dataset size: 28.2 GiB (lazy, nothing is allocated)
Current chunk shape: (365, 721, 1440) = 1.41 GiB per chunk

Let dask_setup guess the workload type

infer_workload_type() looks at dtypes, dimension names and sizes. It’s what workload_type="auto" uses under the hood.

from dask_setup import infer_workload_type

print("Inferred workload type:", infer_workload_type(ds))

Inferred workload type: cpu

Get a cluster and chunk recommendation together

client, cluster, dask_tmp, chunks = setup_dask_client(
    workload_type="cpu",
    max_workers=4,
    ds=ds,
    dashboard=False,
)
print("Recommended chunks:", chunks)

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 4 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~2.9 GiB
Compression: spill=auto | comm=False
Recommended chunks: {'time': 365, 'lat': 360, 'lon': 720}

ds_rechunked = ds.chunk(chunks)
new_shape = ds_rechunked.tas.data.chunksize
print(f"New chunk shape: {new_shape}")
print(f"Size per chunk:  {np.prod(new_shape) * 4 / 2**20:.0f} MiB")
print(
    f"Number of tasks per variable: {ds.tas.data.npartitions:,} → {ds_rechunked.tas.data.npartitions:,}"
)

New chunk shape: (365, 360, 720)
Size per chunk:  361 MiB
Number of tasks per variable: 10 → 60

The recommendation depends on the workload type. "cpu" cuts the data into square-ish blocks, which suits compute. "io" keeps each time step whole and slices only along time, which suits streaming reads and writes:

from dask_setup import recommend_chunks, validate_chunks

for wt in ["cpu", "io"]:
    print(f"{wt:>4}: {dict(recommend_chunks(ds, client=client, workload_type=wt))}")

 cpu: {'time': 365, 'lat': 360, 'lon': 720}
  io: {'time': 129}

You can also call the chunking helpers directly against a running client. verbose=True prints a report and returns a ChunkRecommendation object.

rec = recommend_chunks(ds, client=client, workload_type="cpu", verbose=True)
print(rec)

# validate_chunks() flags chunks that risk OOM (> half a worker's memory) or are tiny (< 10 MiB)
issues = validate_chunks(ds, client=client)
print("\nIssues with the ORIGINAL chunks:", issues or "none (big, but they still fit in a worker)")

📊 Xarray Chunking Recommendations
===================================
 Recommended chunks: {'time': 365, 'lat': 360, 'lon': 720}
 Estimated chunk size: 360.9 MiB
 Total chunks: 60

Current vs Recommended:
-------------------------
  lon: (1440,) → 720
  lat: (721,) → 360
  time: (365, 365, 365, 365, 365, 365, 365, 365, 365, 365) → 365

 Usage:
   ds_chunked = ds.chunk({'time': 365, 'lat': 360, 'lon': 720})
ChunkRecommendation(chunks={'time': 365, 'lat': 360, 'lon': 720}, estimated_chunk_mb=360.9, total_chunks=60)

Issues with the ORIGINAL chunks: none (big, but they still fit in a worker)

client.close()
cluster.close()

7. Configuration and profiles

Passing six keyword arguments in every notebook gets old fast. dask_setup has three layers of configuration, resolved lowest → highest priority:

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

An argument you pass explicitly always wins, even if it equals the default. Anything you leave out is inherited from the profile/config.

7a. A reusable config object

from dask_setup import DaskSetupConfig

my_config = DaskSetupConfig(
    workload_type="cpu",
    max_workers=4,
    reserve_mem_gb=4.0,
    dashboard=False,
    memory_spill=0.80,  # start spilling a little earlier than the 0.85 default
)

# Explicit keyword beats the config: we get 2 workers, not 4
with DaskClientContext(config=my_config, max_workers=2) as (client, cluster, tmp):
    print("Workers:", len(client.nthreads()))

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 2 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~5.8 GiB
Compression: spill=auto | comm=False
Workers: 2

7b. Built-in profiles

dask_setup ships five profiles for common situations:

from dask_setup import ConfigManager

mgr = ConfigManager()
for name, prof in mgr.list_profiles().items():
    c = prof.config
    tag = "builtin" if prof.builtin else "user/site"
    print(f"{name:<18} [{tag:<9}] workload={c.workload_type:<5}  {prof.description}")

climate_analysis   [builtin  ] workload=cpu    Optimized for climate data analysis with large arrays and heavy compute
zarr_io_heavy      [builtin  ] workload=io     Optimized for heavy Zarr I/O operations with many files
development        [builtin  ] workload=mixed  Lightweight configuration for development and testing
production         [builtin  ] workload=mixed  Robust production configuration with conservative memory usage
interactive        [builtin  ] workload=mixed  Optimized for interactive Jupyter notebook usage

Use one by name. profile="auto" picks one based on the detected resources (e.g. climate_analysis on a large Gadi node, development on a laptop).

with DaskClientContext(profile="development", dashboard=False) as (client, cluster, tmp):
    print(f"development profile → {len(client.nthreads())} worker(s)")

print("profile='auto' would pick:", mgr.auto_select_profile())

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 2 | threads/worker: 2 | processes: True
Memory: total ~15.6 GiB | usable ~7.6 GiB | per-worker ~3.8 GiB
Compression: spill=auto | comm=False
development profile → 2 worker(s)
profile='auto' would pick: interactive

7c. Your own profiles

Save a profile once, then use it from any notebook or script. Profiles are YAML files looked up in this order, so yours shadow the site’s and the site’s shadow the builtins:

  1. ~/.dask_setup/profiles/ (you)

  2. $DASK_SETUP_PROFILE_DIR or /etc/dask_setup/profiles/ (site-wide, set up by admins)

  3. Built-ins

To keep this tutorial from touching your real ~/.dask_setup, we point the manager at a throwaway directory. In practice you’d just use ConfigManager().

import tempfile
from pathlib import Path

from dask_setup.config import ConfigProfile

demo_dir = Path(tempfile.mkdtemp(prefix="dask_setup_demo_"))
demo_mgr = ConfigManager(config_dir=demo_dir)

demo_mgr.save_profile(
    ConfigProfile(
        name="era5_netcdf",
        config=DaskSetupConfig(
            workload_type="cpu",  # NetCDF → cpu (see section 4)
            reserve_mem_gb=20.0,
            memory_spill=0.80,
            description="ERA5 NetCDF processing on Gadi normalsr",
            tags=["era5", "netcdf", "gadi"],
        ),
    )
)

print((demo_dir / "profiles" / "era5_netcdf.yaml").read_text())

builtin: false
config:
  adaptive: false
  comm_compression: false
  dashboard: true
  dashboard_port: null
  description: ERA5 NetCDF processing on Gadi normalsr
  io_access_pattern: auto
  io_compression_level: null
  io_format: null
  io_storage_location: auto
  io_target_chunk_mb:
  - 128
  - 512
  max_mem_gb: null
  max_workers: null
  memory_pause: 0.92
  memory_spill: 0.8
  memory_target: 0.75
  memory_terminate: 0.98
  min_workers: null
  name: era5_netcdf
  reserve_mem_gb: 20.0
  silence_logs: false
  spill_compression: auto
  spill_threads: null
  suggest_chunks: false
  tags:
  - era5
  - netcdf
  - gadi
  temp_base_dir: null
  workload_type: cpu
created_at: '2026-10-01T02:42:55.117828'
modified_at: '2026-10-01T02:42:55.117828'
name: era5_netcdf
version: '1.7'

7d. The dask-setup CLI

The same profiles are manageable from the shell:

dask-setup list                    # all profiles, with where each comes from
dask-setup show climate_analysis   # full settings for one profile
dask-setup create my_profile       # interactive wizard
dask-setup validate                # check every profile parses
dask-setup schema > profile.json   # JSON Schema for editor autocomplete on YAML profiles

!dask-setup list

 Built-in Profiles:
  climate_analysis     - Optimized for climate data analysis with large arrays and heavy compute [climate, cpu-heavy, large-memory, analysis]
  development          - Lightweight configuration for development and testing [development, testing, lightweight, local]
  interactive          - Optimized for interactive Jupyter notebook usage [interactive, jupyter, notebook, balanced]
  production           - Robust production configuration with conservative memory usage [production, robust, conservative, adaptive]
  zarr_io_heavy        - Optimized for heavy Zarr I/O operations with many files [zarr, io-heavy, files, storage]

8. Memory tuning

If a job spills a lot, workers spend their time writing to disk. If it never spills, you’re leaving headroom unused. tune_memory_thresholds() reads the spill metrics from a running cluster and nudges memory.target / memory.spill on every worker:

  • little/no spill → tighten (spill earlier, keep more headroom against OOM)

  • heavy spill → loosen (let workers hold more in memory before spilling)

Call it between phases of a long pipeline, once there is real spill to react to. (setup_dask_client(adaptive_memory=True) runs it once at startup, before anything has spilled, so it only ever tightens.)

from dask_setup import tune_memory_thresholds

with DaskClientContext(workload_type="cpu", max_workers=2, dashboard=False) as (
    client,
    cluster,
    tmp,
):
    da.random.random((5_000, 5_000), chunks=1_000).sum().compute()
    result = tune_memory_thresholds(client)
    print(result.summary())

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 2 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~5.8 GiB
Compression: spill=auto | comm=False

Memory thresholds adjusted on 2 worker(s): target 75%→70%, spill 85%→80% (low spill (0.00 GiB < 0.50 GiB threshold); tightened thresholds for extra head-room)

9. Beyond one node

When one node isn’t enough, dask_setup drives dask-jobqueue to submit worker jobs to PBS or SLURM. The same setup_dask_client() call handles it through mode=:

mode=

What happens

"auto" (default)

Uses detect_cluster_mode(): "local" on a laptop or login node, "interactive" in qsub -I / salloc, "pbs" / "slurm" inside a batch job

"local"

Always a single-node LocalCluster (everything above)

"pbs" / "slurm"

Scheduler runs here; workers are submitted as separate batch jobs

"interactive"

Inside an interactive job (qsub -I, salloc)

So for the everyday case, one batch job using the cores it was given, pass mode="local".

Describe the worker jobs with a MultiNodeConfig:

from dask_setup import MultiNodeConfig, detect_cluster_mode

mn = MultiNodeConfig(
    workload_type="cpu",
    workers_per_node=12,
    cores_per_worker=4,
    mem_per_worker_gb=15,
    walltime="04:00:00",
    queue="normalsr",
    project="ab12",
    job_extra_directives=["-l storage=gdata/rt52+scratch/ab12", "-l jobfs=200GB"],
    env_extra=["module use /g/data/xp65/public/modules", "module load conda/analysis3"],
    adaptive=True,
    min_jobs=1,
    max_jobs=8,
)

print(f"Per job: {mn.total_cores_per_job} cores, {mn.total_mem_gb_per_job:.0f} GB")
print("Mode detected here:", detect_cluster_mode())

Per job: 48 cores, 180 GB
Mode detected here: local

On Gadi you’d then run:

client, cluster, dask_tmp = setup_dask_client(mode="pbs", multi_node_config=mn)

That needs a real PBS scheduler, so here we’ll generate the equivalent job script instead. This is also a quick way to submit a whole analysis script as a batch job:

from dask_setup import generate_pbs_script

print(generate_pbs_script(mn, script_path="analysis.py"))

#!/bin/bash
#PBS -l ncpus=48,mem=180GB,walltime=04:00:00
#PBS -q normalsr
#PBS -P ab12
#PBS -l storage=gdata/rt52+scratch/ab12
#PBS -l jobfs=200GB

# --- Environment ---
set -euo pipefail
module use /g/data/xp65/public/modules
module load conda/analysis3

# --- Run ---
python3 analysis.py

generate_slurm_script() does the same for SLURM sites, and dask-setup submit wraps both from the command line.

10. When things go wrong

dask_setup raises specific exception types whose messages name a concrete fix. All inherit from DaskSetupError, so you can catch that one type if you want to.

from dask_setup.exceptions import DaskSetupError

try:
    DaskSetupConfig(workload_type="gpu-please", max_workers=0)
except DaskSetupError as e:
    print(type(e).__name__)
    print(e)

ConfigurationValidationError
❌ [CONFIG_VALIDATION] Configuration validation failed:
   • workload_type: must be one of ['auto', 'cpu', 'gpu', 'io', 'mixed'] (got 'gpu-please') (current: gpu-please)
   • max_workers: must be positive (got 0) (current: 0)

💡 Suggestions:
   1. Use workload_type='io' for I/O intensive tasks, 'cpu' for compute, or 'mixed'
   2. Set max_workers to a value between 1 and 4

📖 Documentation: https://github.com/dask-contrib/dask_setup#configuration

Every problem is reported at once, not one per attempt, and each comes with a suggested fix.

Where it can, dask_setup adapts rather than fails. Here we ask for 4 workers but cap Dask at 2 GiB, which is not enough for 4 workers with at least 1 GiB each. The worker count drops to fit and the message tells you which setting to change:

with DaskClientContext(
    workload_type="cpu", max_workers=4, max_mem_gb=2, reserve_mem_gb=1, dashboard=False
) as (client, cluster, tmp):
    print("Workers actually started:", len(client.nthreads()))

WARNING  [client] Reducing worker count to fit available memory (requested_workers=4 | fitted_workers=1 | usable_mem_gib=1.0 | min_per_worker_gib=1.0 | hint=raise max_mem_gb or lower reserve_mem_gb for more workers)

[setup_dask_client] Reduced workers 4 -> 1 so each keeps at least 1.0 GiB (1.0 GiB usable).

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-3457
Workers: 1 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~1.0 GiB | per-worker ~1.0 GiB
Compression: spill=auto | comm=False
Workers actually started: 1

Debugging tip: turn on diagnostic logging to see every decision dask_setup makes:

export DASK_SETUP_LOG_LEVEL=DEBUG

or in Python, dask_setup.configure_logging(level="DEBUG").

Cheat sheet

from dask_setup import setup_dask_client, DaskClientContext

# Quick start (script or notebook)
client, cluster, dask_tmp = setup_dask_client(workload_type="cpu")

# Inside a PBS/SLURM batch job: stay on this node
client, cluster, dask_tmp = setup_dask_client(workload_type="cpu", mode="local")

# Safe cleanup
with DaskClientContext(workload_type="io") as (client, cluster, dask_tmp):
    ...

# With an xarray dataset → 4-tuple including recommended chunks
client, cluster, dask_tmp, chunks = setup_dask_client(workload_type="cpu", ds=ds)
ds = ds.chunk(chunks)

# Saved or built-in profile, overriding one setting
client, cluster, dask_tmp = setup_dask_client(profile="climate_analysis", max_workers=24)

# Multi-node on PBS
client, cluster, dask_tmp = setup_dask_client(mode="pbs", multi_node_config=mn)

Your data / code

workload_type

NetCDF / HDF5 files

"cpu"

Zarr stores, cloud object storage

"io"

NumPy-heavy maths on in-memory arrays

"cpu" or "io", so benchmark both

Mixed pipeline, not sure

"mixed"

Have a dataset, want a guess

"auto" with ds=

Exercises

  1. Re-run section 3 with simulated_env(NCPUS="104", PBS_MEM="500gb") (a full normalsr node). How many workers does "cpu" give you, and how much memory does each get?

  2. Change the synthetic dataset in section 6 to 1° resolution (181 × 360). How do the recommended chunks change, and why?

  3. Save a profile for your own project’s data and load it with setup_dask_client(profile=...).

  4. Generate a SLURM script for the same MultiNodeConfig with generate_slurm_script().

Further reading