API Reference 

dask_setup.multinode.setup_interactive_cluster(workload_type='cpu', workers_per_node=None, threads_per_worker=None, wait_for_workers=True, worker_timeout=60.0, scheduler_port=8786, config=None)[source]

Set up a Dask cluster using resources already allocated in an interactive job.

This is the right function to call from a Jupyter notebook running inside a qsub -I (PBS) or salloc (SLURM) interactive session. Unlike setup_pbs_cluster() / setup_slurm_cluster(), it does not submit new batch jobs — it creates workers directly on the nodes that PBS or SLURM has already assigned to the current session.

Behaviour

Single allocated node — creates a LocalCluster using all cores and memory available in the allocation (reads PBS_NCPUS / PBS_MEM or SLURM equivalents via detect_resources()).
Multiple allocated nodes — creates a dask.distributed.SSHCluster across all nodes listed in PBS_NODEFILE / $SLURM_NODELIST. Passwordless SSH between nodes is required (standard on most HPC systems).

param workload_type:: Worker topology — "cpu", "io", "mixed", or "auto". Determines threads per worker and workers per node using the same logic as single-node setup_dask_client().
param workers_per_node:: Override the number of worker processes to start per node. When None (default), derived from workload_type and the core count.
param threads_per_worker:: Override threads per worker process. When None, derived from workload_type.
param wait_for_workers:: Block until at least one worker has connected. Defaults to True.
param worker_timeout:: Seconds to wait for workers. Defaults to 60 s (workers start almost immediately because no job submission is involved).
param scheduler_port:: Port for the Dask scheduler (multi-node SSH path only).
returns:: client — a connected dask.distributed.Client. cluster — LocalCluster or SSHCluster. tmp_path — path to the temporary directory, or "" when not set.
rtype:: (client, cluster, tmp_path)

Parameters:

workload_type (str)
workers_per_node (int | None)
threads_per_worker (int | None)
wait_for_workers (bool)
worker_timeout (float)
scheduler_port (int)
config (Any)

Return type:

tuple[Any, Any, str | None]

dask_setup.multinode.setup_pbs_cluster(config=None, *, workload_type=None, workers_per_node=None, cores_per_worker=None, mem_per_worker_gb=None, walltime=None, queue=None, project=None, job_extra_directives=None, n_nodes=None, shared_tmp_dir=None, env_extra=None, adaptive=None, min_jobs=None, max_jobs=None, n_workers=1, wait_for_workers=True, worker_timeout=300.0)[source]

Set up a multi-node Dask cluster backed by PBS/Torque.

Uses dask-jobqueue.PBSCluster under the hood, applying MultiNodeConfig defaults and the same error-handling conventions as single-node setup_dask_client().

Parameters:

config (MultiNodeConfig | None) – A MultiNodeConfig instance. If None, a default one is constructed from the keyword arguments.
workload_type (str | None) – Convenience overrides for any MultiNodeConfig field.
workers_per_node (int | None) – Convenience overrides for any MultiNodeConfig field.
… – Convenience overrides for any MultiNodeConfig field.
n_workers (int) – Number of worker jobs to submit immediately (before any adaptive scaling kicks in).
wait_for_workers (bool) – If True (default), block until at least one worker has connected to the scheduler before returning. This prevents the client from appearing to have 0 workers when PBS has not yet started the submitted jobs. Set to False to return immediately and poll client.wait_for_workers() yourself.
worker_timeout (float) – Seconds to wait for workers to come online when wait_for_workers is True. Defaults to 300 s (5 minutes). Raise this on queues with long wait times. A TimeoutError is raised (with a diagnostic message) if no workers connect within this window.
cores_per_worker (int | None)
mem_per_worker_gb (float | None)
walltime (str | None)
queue (str | None)
project (str | None)
job_extra_directives (list[str] | None)
n_nodes (int | None)
shared_tmp_dir (str | Path | None)
env_extra (list[str] | None)
adaptive (bool | None)
min_jobs (int | None)
max_jobs (int | None)

Returns:

client — a connected dask.distributed.Client. cluster — the live PBSCluster object. shared_tmp — a SharedTempDir or None.

Return type:

(client, cluster, shared_tmp)

dask_setup.multinode.setup_slurm_cluster(config=None, *, workload_type=None, workers_per_node=None, cores_per_worker=None, mem_per_worker_gb=None, walltime=None, queue=None, project=None, job_extra_directives=None, n_nodes=None, shared_tmp_dir=None, env_extra=None, adaptive=None, min_jobs=None, max_jobs=None, n_workers=1, wait_for_workers=True, worker_timeout=300.0)[source]

Set up a multi-node Dask cluster backed by SLURM.

Uses dask-jobqueue.SLURMCluster under the hood. Mirrors the API of setup_pbs_cluster().

Parameters:

config (MultiNodeConfig | None) – A MultiNodeConfig instance.
n_workers (int) – Number of worker jobs to submit immediately.
wait_for_workers (bool) – If True (default), block until at least one worker has connected to the scheduler before returning. Set to False to return immediately.
worker_timeout (float) – Seconds to wait for workers to come online. Defaults to 300 s.
workload_type (str | None)
workers_per_node (int | None)
cores_per_worker (int | None)
mem_per_worker_gb (float | None)
walltime (str | None)
queue (str | None)
project (str | None)
job_extra_directives (list[str] | None)
n_nodes (int | None)
shared_tmp_dir (str | Path | None)
env_extra (list[str] | None)
adaptive (bool | None)
min_jobs (int | None)
max_jobs (int | None)

Return type:

(client, cluster, shared_tmp)

dask_setup.multinode.generate_pbs_script(cfg, script_path, python_executable='python3', dask_setup_args='')[source]

Generate a PBS job script string that launches a dask_setup-based job.

Parameters:

cfg (MultiNodeConfig) – Multi-node configuration.
script_path (str) – Path to the user’s Python script to run.
python_executable (str) – Python interpreter to use (default python3).
dask_setup_args (str) – Additional arguments to pass to the script.

Returns:

Complete #PBS job script content.

Return type:

dask_setup.multinode.generate_slurm_script(cfg, script_path, python_executable='python3', dask_setup_args='')[source]

Generate a SLURM job script string.

Parameters:

cfg (MultiNodeConfig) – Multi-node configuration.
script_path (str) – Path to the user’s Python script to run.
python_executable (str) – Python interpreter to use (default python3).
dask_setup_args (str) – Additional arguments to pass to the script.

Returns:

Complete #SBATCH job script content.

Return type:

Chunking and xarray

dask_setup.xarray.recommend_chunks(ds, client=None, workload_type='auto', target_chunk_mb=(256, 512), safety_factor=0.6, verbose=False, storage_format=None, storage_path=None, chunk_domain=None, **kwargs)[source]

Recommend optimal chunk sizes for an xarray Dataset or DataArray.

Analyzes the dataset characteristics and available Dask worker resources to suggest chunk sizes that balance memory efficiency with good parallelization. Optionally considers storage format for I/O-optimized chunking.

Parameters:

ds (xr.Dataset | xr.DataArray) – xarray Dataset or DataArray to analyze
client (Client | None) – Active Dask client (to detect worker configuration), or None to use system defaults
workload_type (str) – Chunking strategy - “cpu” (compute-heavy), “io” (I/O-heavy), “mixed” (balanced), or “auto” (detect automatically)
target_chunk_mb (tuple[float, float]) – (min, max) target chunk size in MiB. Default: (256, 512)
safety_factor (float) – Fraction of worker memory to use for chunks. Default: 0.6 (60%)
verbose (bool) – If True, print detailed report and return ChunkRecommendation object. If False, return simple dict of chunk sizes.
storage_format (str | None) – Optional storage format hint (“zarr”, “netcdf”) for I/O optimization
storage_path (str | None) – Optional storage path/URL for format detection and location optimization
chunk_domain (str | None) –
Optional axis family to restrict chunking to.

"spatial" — only chunk spatial dimensions (lat, lon, x, y, …); all temporal dimensions (time, date, step, …) are loaded fully into a single chunk (equivalent to ds.chunk({"time": -1, ...})).

"temporal" — only chunk temporal dimensions; all spatial dimensions are loaded fully into a single chunk.

None (default) — no constraint, the existing behaviour applies and all dimensions are candidates for chunking.
**kwargs (Any) – Additional parameters (reserved for future use)

Returns:

A dict mapping dimension names to recommended chunk sizes (locked dimensions appear as -1) when verbose=False; a ChunkRecommendation object with full details when verbose=True.

Raises:

XarrayDependencyError – If xarray or numpy are not installed
ValueError – If chunk_domain is not one of the accepted values

Return type:

dict[str, int] | ChunkRecommendation

Examples

Basic usage: >>> chunks = recommend_chunks(ds, client) >>> ds_chunked = ds.chunk(chunks)

With detailed analysis: >>> recommendation = recommend_chunks(ds, client, verbose=True) >>> print(recommendation) # Shows detailed report >>> ds_chunked = ds.chunk(recommendation.chunks)

For I/O-heavy workloads: >>> chunks = recommend_chunks(ds, workload_type=”io”)

Chunk only spatial dims, load all time into memory: >>> chunks = recommend_chunks(ds, client, chunk_domain=”spatial”) >>> # chunks == {“lat”: 256, “lon”: 512, “time”: -1} >>> ds_chunked = ds.chunk(chunks)

Chunk only the time axis, load all spatial points into memory: >>> chunks = recommend_chunks(ds, client, chunk_domain=”temporal”) >>> # chunks == {“time”: 120, “lat”: -1, “lon”: -1} >>> ds_chunked = ds.chunk(chunks)

dask_setup.xarray.validate_chunks(ds, client=None, max_chunk_ratio=0.5, min_chunk_mb=10.0)[source]

Validate existing chunk sizes against cluster memory limits.

Compares each dimension’s current chunk size to the per-worker memory available in the cluster, and warns when chunks are dangerously large (OOM risk) or very small (high task-graph overhead).

This function is called automatically by setup_dask_client() when a dataset is passed via the ds= parameter.

Parameters:

ds (xr.Dataset or xr.DataArray) – Dataset with existing chunking to validate.
client (dask.distributed.Client or None) – Active Dask client used to read per-worker memory limits. When None, falls back to system psutil memory detection.
max_chunk_ratio (float) – Maximum acceptable fraction of per-worker memory per chunk. Default 0.5 — chunks larger than 50 % of worker memory trigger a warning.
min_chunk_mb (float) – Minimum acceptable chunk size in MiB. Default 10 MiB. Chunks smaller than this trigger a task-overhead warning.

Returns:

Warning messages (empty list if no issues found). The same messages are also emitted via warnings.warn() so they appear in logs.

Return type:

list[str]

Examples

issues = validate_chunks(ds, client)
if issues:
    chunks = recommend_chunks(ds, client)
    ds = ds.chunk(chunks)

class dask_setup.xarray.ChunkRecommendation(chunks, estimated_chunk_mb, total_chunks, warnings_list=None, dataset_info=None)[source]

Bases: object

Container for chunk recommendation results.

Parameters:

chunks (dict[str, int])
estimated_chunk_mb (float)
total_chunks (int)
warnings_list (list[str] | None)
dataset_info (dict[str, Any] | None)

dask_setup.rechunk.rechunk_dataset(ds, target_chunks, client, dask_tmp, output_path=None, max_mem='2GB', cleanup_temp=True)[source]

Rechunk an xarray Dataset using rechunker with HPC-friendly temp placement.

Wraps the rechunker library to safely rechunk large datasets to new chunk shapes. Both the intermediate Zarr temp store and the final output Zarr store are written to the Dask spill/temp directory (dask_tmp), which is typically $PBS_JOBFS fast local storage on NCI Gadi.

Parameters:

ds (xr.Dataset or xr.DataArray) – Dataset to rechunk. Must be backed by dask arrays (opened lazily).
target_chunks (dict[str, int]) – Target chunk sizes, e.g. {"time": 10, "lat": 200, "lon": 200}. Dimensions not listed keep their existing chunk sizes.
client (dask.distributed.Client) – Active Dask client used to execute the rechunking task graph.
dask_tmp (str or Path) – Dask spill/temp directory — the third element returned by setup_dask_client(). Intermediate and output stores are placed here unless output_path is specified.
output_path (str, Path, or None) – Destination Zarr store path for the rechunked dataset. If None, a timestamped subdirectory under dask_tmp is used. This path is not cleaned up automatically; it persists for the job’s lifetime.
max_mem (str) – Maximum memory rechunker may use per task, e.g. "1GB", "4GB". Defaults to "2GB". Should be at most 50 % of per-worker memory.
cleanup_temp (bool) – If True (default), the rechunker intermediate temp store is deleted after rechunking completes or fails. The output store is never deleted automatically.

Returns:

The rechunked dataset opened lazily from output_path.

Return type:

xr.Dataset or xr.DataArray

Raises:

ImportError – If rechunker or zarr are not installed.
RuntimeError – If rechunking fails for any reason. Partial output is cleaned up.

Notes

Install optional dependencies with:

pip install rechunker zarr

Examples

from dask_setup import setup_dask_client, recommend_chunks, rechunk_dataset

client, cluster, tmp = setup_dask_client(workload_type="io")
ds = xr.open_zarr("era5.zarr")
chunks = recommend_chunks(ds, client)
ds_rechunked = rechunk_dataset(ds, chunks, client, tmp)
ds_rechunked.to_zarr("era5_rechunked.zarr")

I/O patterns

dask_setup.io_patterns.recommend_io_chunks(ds, path_or_url=None, client=None, format_hint=None, access_pattern='auto', target_chunk_mb=(128, 512), storage_location='auto', verbose=False)[source]

Recommend I/O-optimized chunks for xarray datasets.

Parameters:

ds (xr.Dataset | xr.DataArray) – xarray Dataset or DataArray to optimize
path_or_url (str | None) – Optional path/URL to determine storage format and location
client (Client | None) – Optional Dask client for cluster info
format_hint (str | None) – Optional format override (“zarr”, “netcdf”, etc.)
access_pattern (str) – Access pattern hint (“sequential”, “random”, “streaming”, “compute”, “auto”)
target_chunk_mb (tuple[float, float]) – Target chunk size range in MiB
storage_location (str) – Storage location hint (“local”, “cloud”, “network”, “auto”)
verbose (bool) – If True, return full IORecommendation object

Returns:

IORecommendation object if verbose=True, else dict of chunk recommendations

Return type:

IORecommendation | dict[str, int]

dask_setup.io_patterns.detect_storage_format(path_or_url, ds=None)[source]

Detect storage format from path/URL and optional xarray object.

Parameters:

path_or_url (str) – File path or URL to analyze
ds (Dataset | DataArray | None) – Optional xarray Dataset/DataArray to check for format hints

Returns:

"zarr_v3", "zarr", "netcdf", "kerchunk", or "unknown". "zarr_v3" is returned for zarr-python ≥ 3.0 stores; "kerchunk" is returned for Kerchunk / VirtualiZarr reference stores (checked before zarr, since those stores present a zarr-like interface).

Return type:

Format name

class dask_setup.io_patterns.IORecommendation(format, chunks, compression, storage_options, access_pattern, estimated_throughput_mb_s, warnings, extra=<factory>)[source]

Bases: object

Container for I/O optimization recommendations.

Parameters:

format (str)
chunks (dict[str, int])
compression (dict[str, Any])
storage_options (dict[str, Any])
access_pattern (str)
estimated_throughput_mb_s (float)
warnings (list[str])
extra (dict[str, Any])

extra: dict[str, Any]: Format-specific extras (e.g. sharding config for Zarr v3).

class dask_setup.io_patterns.ZarrOptimizer(client=None)[source]

Bases: IOOptimizer

Optimizer for Zarr format storage.

Parameters:: client (Client | None)

detect_format(path_or_url, ds=None)[source]

Detect Zarr format from path or xarray backend.

Parameters:

path_or_url (str)
ds (Dataset | DataArray | None)

Return type:

optimize_chunks(ds, target_chunk_mb=(128, 512), access_pattern='auto')[source]

Optimize chunks for Zarr format.

Parameters:

ds (Dataset | DataArray)
target_chunk_mb (tuple[float, float])
access_pattern (str)

Return type:

optimize_compression(ds, storage_location='local')[source]

Recommend compression for Zarr.

Parameters:

ds (Dataset | DataArray)
storage_location (str)

Return type:

optimize_storage_options(path_or_url, access_pattern='auto')[source]

Optimize storage options for Zarr.

Parameters:

path_or_url (str)
access_pattern (str)

Return type:

class dask_setup.io_patterns.ZarrV3Optimizer(client=None)[source]

Bases: IOOptimizer

Optimizer for Zarr v3 format storage (zarr-python ≥ 3.0).

Handles the updated zarr.open() interface, sharding via zarr.codecs.ShardingCodec, and the new codec pipeline introduced in zarr>=3.0. When sharding is appropriate (large arrays), the recommended outer/inner chunk shapes are returned in IORecommendation.extra["sharding"].

Parameters:: client (Client | None)

detect_format(path_or_url, ds=None)[source]

Detect Zarr v3 from path pattern or dataset store metadata.

Parameters:

path_or_url (str)
ds (Dataset | DataArray | None)

Return type:

optimize_chunks(ds, target_chunk_mb=(128, 512), access_pattern='auto')[source]

Optimize outer chunk sizes for Zarr v3 (with sharding awareness).

Parameters:

ds (Dataset | DataArray)
target_chunk_mb (tuple[float, float])
access_pattern (str)

Return type:

optimize_compression(ds, storage_location='local')[source]

Recommend a zarr v3 codec-pipeline compression configuration.

Parameters:

ds (Dataset | DataArray)
storage_location (str)

Return type:

optimize_storage_options(path_or_url, access_pattern='auto')[source]

Optimize storage options for Zarr v3.

Parameters:

path_or_url (str)
access_pattern (str)

Return type:

class dask_setup.io_patterns.NetCDFOptimizer(client=None)[source]

Bases: IOOptimizer

Optimizer for NetCDF format storage.

Parameters:: client (Client | None)

detect_format(path_or_url, ds=None)[source]

Detect NetCDF format from path or xarray backend.

Parameters:

path_or_url (str)
ds (Dataset | DataArray | None)

Return type:

optimize_chunks(ds, target_chunk_mb=(64, 256), access_pattern='auto')[source]

Optimize chunks for NetCDF format.

Parameters:

ds (Dataset | DataArray)
target_chunk_mb (tuple[float, float])
access_pattern (str)

Return type:

optimize_compression(ds, storage_location='local')[source]

Recommend compression for NetCDF.

Parameters:

ds (Dataset | DataArray)
storage_location (str)

Return type:

optimize_storage_options(path_or_url, access_pattern='auto')[source]

Optimize storage options for NetCDF.

Parameters:

path_or_url (str)
access_pattern (str)

Return type:

class dask_setup.io_patterns.KerchunkOptimizer(client=None)[source]

Bases: IOOptimizer

Optimizer for datasets opened via Kerchunk or VirtualiZarr reference stores.

When a dataset is backed by a Kerchunk / VirtualiZarr reference filesystem, chunk boundaries are fixed by the byte ranges of the original files. Rechunking would require a full data copy, defeating the purpose of the virtual store. This optimizer detects the reference-filesystem pattern, returns the existing chunks unchanged, and adds an informational warning.

Parameters:: client (Client | None)

detect_format(path_or_url, ds=None)[source]

Detect Kerchunk/VirtualiZarr reference filesystem datasets.

Parameters:

path_or_url (str)
ds (Dataset | DataArray | None)

Return type:

optimize_chunks(ds, target_chunk_mb=(128, 512), access_pattern='auto')[source]

Return the existing chunks — rechunking a Kerchunk dataset is a no-op.

Parameters:

ds (Dataset | DataArray)
target_chunk_mb (tuple[float, float])
access_pattern (str)

Return type:

optimize_compression(ds, storage_location='local')[source]

Compression is fixed in the underlying source files.

Parameters:

ds (Dataset | DataArray)
storage_location (str)

Return type:

optimize_storage_options(path_or_url, access_pattern='auto')[source]

Return fsspec options appropriate for a Kerchunk reference store.

Parameters:

path_or_url (str)
access_pattern (str)

Return type:

dask_setup.parquet.recommend_parquet_chunks(df, client=None, path_or_url=None, target_partition_mb=(128.0, 512.0), compression='auto', storage_location='auto', verbose=False)[source]

Recommend partition sizes for Dask DataFrame Parquet workloads.

Parameters:

df (Any) – A Dask or pandas DataFrame (or any object with dtypes, columns, and optionally memory_usage()).
client (Client | None) – Optional Dask Client. When provided, per-worker memory limits are used to tighten the upper partition bound.
path_or_url (str | None) – Destination path or URL. Used to infer storage location and to recommend appropriate storage_options.
target_partition_mb (tuple[float, float]) – (min_mb, max_mb) target partition size range. Defaults to (128, 512).
compression (str) – Parquet compression codec. "auto" picks "snappy" for local storage and "zstd" for cloud / network.
storage_location (str) – One of "local", "cloud", "network", or "auto".
verbose (bool) – If True, return the full ParquetRecommendation and print a report. If False (default), return rows_per_partition only.

Returns:

ParquetRecommendation when verbose is True, otherwise the
recommended rows_per_partition as an int.

Return type:

ParquetRecommendation | int

class dask_setup.parquet.ParquetRecommendation(rows_per_partition, compression, storage_options, estimated_partition_mb, warnings, extra=<factory>)[source]

Bases: object

Container for Parquet / Arrow partition recommendations.

Parameters:

rows_per_partition (int)
compression (str)
storage_options (dict[str, Any])
estimated_partition_mb (float)
warnings (list[str])
extra (dict[str, Any])

rows_per_partition

Recommended number of rows per Dask partition.

Type:: int

compression

Recommended compression codec and options dict, suitable for passing to dask.dataframe.to_parquet(..., compression=...).

Type:: str

storage_options

Recommended storage_options dict for cloud or network targets.

Type:: dict[str, Any]

estimated_partition_mb

Estimated size of each partition in MiB after applying the recommended row count.

Type:: float

warnings

List of advisory messages (empty when all looks healthy).

Type:: list[str]

extra

Format-specific extras, e.g. row_group_size for Parquet v2.

Type:: dict[str, Any]

summary()[source]

Return a one-line human-readable summary.

Return type:: str

Workload inference and tuning

dask_setup.workload.infer_workload_type(ds=None)[source]

Infer the most appropriate workload type from a dataset’s structure.

Examines variable dtypes, dimension names, and array sizes to decide whether the workload is primarily:

CPU-bound (→ "cpu"): Many floating-point variables, scientific dimension names (time/lat/lon/lev), large arrays per variable.
I/O-bound (→ "io"): Mostly integer / byte / bool variables, many small variables typical of lookup tables or mask datasets.
Mixed (→ "mixed"): When the signal is ambiguous or no dataset is provided.

Parameters:: ds (xr.Dataset, xr.DataArray, or None) – Dataset to inspect. When None, returns "mixed" immediately.
Returns:: One of "cpu", "io", or "mixed".
Return type:: str

Notes

The heuristic is intentionally conservative — it returns "mixed" when the evidence is not strong (score difference ≤ 2). Users can always override with an explicit workload_type= argument.

The signal tables are tuned for CF-convention geoscience / climate datasets (the primary NCI Gadi use-case) but are general enough for most scientific array workloads.

Examples

import xarray as xr
from dask_setup import infer_workload_type

ds = xr.open_zarr("era5.zarr")
wt = infer_workload_type(ds)  # likely "cpu"
client, cluster, tmp = setup_dask_client(workload_type=wt)

dask_setup.tune.tune_memory_thresholds(client, strategy='auto', spill_threshold_gib=0.5, tighten_by=0.05, loosen_by=0.05)[source]

Adjust Dask worker memory thresholds based on observed spill behaviour.

Reads current spill metrics from the scheduler, then applies threshold adjustments on every worker:

No / low spilling (spill < spill_threshold_gib): tighten memory.target and memory.spill downward by tighten_by so workers retain more head-room before the GC kicks in.
Heavy spilling (spill ≥ spill_threshold_gib): loosen the same thresholds upward by loosen_by so workers buffer more data in RAM before reaching the spill threshold, reducing disk write amplification.

Adjustments are clamped to safe ranges: target ∈ [0.40, 0.90], spill ∈ [0.50, 0.95]. If the current value is already at the boundary no change is applied.

Parameters:

client (dask.distributed.Client) – A connected Dask client. The scheduler must still be reachable.
strategy ({"auto", "tighten", "loosen", "off"}) – "auto" (default) reads spill stats and chooses the direction. "tighten" / "loosen" force the direction regardless. "off" returns immediately without making any changes.
spill_threshold_gib (float) – Spill volume (GiB) above which "auto" mode switches to the loosen direction. Default: 0.5 GiB.
tighten_by (float) – Fraction to reduce thresholds by when tightening (default 0.05 = 5 pp).
loosen_by (float) – Fraction to increase thresholds by when loosening (default 0.05 = 5 pp).

Returns:

Snapshot of what changed and why. Call .summary() for a one-liner suitable for logging.

Return type:

MemoryTuneResult

Examples

from dask_setup import setup_dask_client, tune_memory_thresholds

client, cluster, tmp = setup_dask_client()

# Run a sample batch of tasks first
client.compute(ds.isel(time=slice(0, 10)).mean()).result()

# Now tune based on what was observed
result = tune_memory_thresholds(client)
print(result.summary())

class dask_setup.tune.MemoryTuneResult(strategy, old_target, new_target, old_spill, new_spill, spill_gib_observed, rationale, workers_updated)[source]

Bases: object

Result from a memory threshold tuning pass.

Parameters:

strategy (str)
old_target (float)
new_target (float)
old_spill (float)
new_spill (float)
spill_gib_observed (float)
rationale (str)
workers_updated (int)

strategy

Effective strategy used: "tighten", "loosen", or "off".

Type:: str

old_target

memory.target value before adjustment (fraction, 0–1).

Type:: float

new_target

memory.target value after adjustment.

Type:: float

old_spill

memory.spill value before adjustment.

Type:: float

new_spill

memory.spill value after adjustment.

Type:: float

spill_gib_observed

Total spill volume observed across all workers at tune time (GiB).

Type:: float

rationale

Human-readable explanation of the decision.

Type:: str

workers_updated

Number of workers whose thresholds were successfully updated. 0 if no change was needed or if the update failed.

Type:: int

summary()[source]

One-line human-readable summary of the tuning result.

Examples

"Memory thresholds adjusted on 4 worker(s): target 75%→70%, spill 85%→80% (low spill...)" "No changes made (strategy='off'; no changes made)"

Return type:: str

dask_setup.callbacks.register_worker_callbacks(client, on_worker_death=None, on_worker_added=None)[source]

Register callbacks for worker lifecycle events on a Dask cluster.

Installs a SchedulerPlugin that calls user-supplied functions when workers join or leave the cluster.

Parameters:

client (dask.distributed.Client) – A connected Dask client.
on_worker_death (callable(worker_addr: str) | None) – Called when a worker is removed from the scheduler — whether it crashed, was explicitly closed, or was scaled down. worker_addr is the worker’s TCP address string (e.g. "tcp://127.0.0.1:12345").
on_worker_added (callable(worker_addr: str) | None) – Called when a new worker connects to the scheduler.

Return type:

None

Notes

If neither callback is provided this function is a no-op.
The plugin is registered with name="dask_setup_worker_health". Calling this function a second time replaces the existing plugin (idempotent where supported by the installed Dask version).
Exceptions raised inside callbacks are caught and logged as warnings so that a buggy callback never crashes the scheduler.

Examples

import threading
from dask_setup import setup_dask_client, register_worker_callbacks

dead: list[str] = []
_lock = threading.Lock()

def on_death(addr: str) -> None:
    with _lock:
        dead.append(addr)

client, cluster, tmp = setup_dask_client()
register_worker_callbacks(client, on_worker_death=on_death)

# ... run tasks ...

if dead:
    print(f"Workers that died during the run: {dead}")

Reporting

dask_setup.reporting.cluster_report(client, start_time=None)[source]

Collect post-run statistics from a connected Dask client.

Queries the distributed scheduler for per-worker memory, spill volume, and task counts. All metric collection is done defensively — any field that cannot be retrieved (e.g. because this Dask version does not expose the metric) is left at its zero/empty default rather than raising.

Parameters:

client (dask.distributed.Client) – A connected Dask client. The scheduler must still be reachable; call this function before client.close().
start_time (float or None) – Wall-clock start time as returned by time.monotonic(). When provided, ClusterReport.wall_time_seconds is computed as time.monotonic() - start_time. Pass None (the default) to leave wall_time_seconds at 0.0.

Returns:

Snapshot of cluster metrics at the time of the call.

Return type:

ClusterReport

Examples

import time
from dask_setup import setup_dask_client, cluster_report

t0 = time.monotonic()
client, cluster, tmp = setup_dask_client()
# ... do work ...
report = cluster_report(client, start_time=t0)
print(report.summary_line())

class dask_setup.reporting.ClusterReport(wall_time_seconds=0.0, memory_per_worker_gib=<factory>, total_spill_gib=0.0, total_tasks=0)[source]

Bases: object

Summary statistics collected from a Dask cluster at close time.

Parameters:

wall_time_seconds (float)
memory_per_worker_gib (dict[str, float])
total_spill_gib (float)
total_tasks (int)

wall_time_seconds

Elapsed wall-clock time since the cluster was started (or since the start_time passed to cluster_report()). 0.0 if no start time was provided.

Type:: float

memory_per_worker_gib

Managed data held in memory (GiB) per worker address, sampled at report time. Data spilled to disk is excluded; it is counted in total_spill_gib instead. Keys are worker addresses (e.g. "tcp://127.0.0.1:12345").

Type:: dict[str, float]

total_spill_gib

Total bytes written to disk spill storage across all workers (GiB). 0.0 if spill metrics are unavailable in this Dask version.

Type:: float

total_tasks

Number of tasks present in the scheduler graph at report time. This is a snapshot — it is not an all-time cumulative count.

Type:: int

property wall_time_str: str: Human-readable elapsed wall time (e.g. "2h 5m 30s").

property peak_memory_gib: float: Maximum memory across all workers (GiB), or 0.0 if unknown.

property total_memory_gib: float: Sum of in-memory managed data across all workers (GiB).

summary_line()[source]

One-line summary suitable for logging or printing.

Example output:

wall=1m 23s | workers=4 | mem=3.21 GiB/max | spill=0.00 GiB | tasks=1842

Return type:: str

to_dict()[source]

Return a JSON-serialisable dictionary of all report fields.

Returns:: All fields including derived properties (wall_time_str, peak_memory_gib, total_memory_gib).
Return type:: dict

Benchmarking

dask_setup.benchmark.benchmark_config(configs, ds, operation='mean', *, repeats=1, warmup=False, fallback_on_detection_failure=True, verbose=False)[source]

A/B-test multiple configurations against the same xarray operation.

Each configuration is run in a fresh cluster so that memory state, worker topology, and spill behaviour are fully independent between runs.

Parameters:

configs (dict[str, DaskSetupConfig] | list[DaskSetupConfig] | DaskSetupConfig) –
Configurations to compare. May be:
- A dict[name, DaskSetupConfig] — explicit names.
- A list[DaskSetupConfig] — names are "config_0", etc.
- A single DaskSetupConfig — returns a one-element list.
ds (Any) – xarray Dataset or DataArray to benchmark against. Should not already be chunked — benchmark_config() will chunk it according to each config’s recommended chunks (or leave it as-is if no chunking is suggested).
operation (Callable[[Any], Any] | str) – Operation to time. Either a callable fn(ds) -> lazy_result or one of the string aliases: "mean", "sum", "std", "max", "min", "var", "rechunk".
repeats (int) – Number of timed repetitions per configuration. Mean and std-dev of wall time are stored in each BenchmarkResult.
warmup (bool) – If True, run one un-timed warmup pass before the timed repeats. Useful for eliminating cold-start overhead (e.g. JIT compilation).
fallback_on_detection_failure (bool) – Passed to setup_dask_client().
verbose (bool) – If True, print a summary line after each configuration completes.

Returns:

One entry per configuration, in the same order as configs.

Return type:

list[BenchmarkResult]

Examples

from dask_setup import DaskSetupConfig
from dask_setup.benchmark import benchmark_config

results = benchmark_config(
    {
        "io":  DaskSetupConfig(workload_type="io"),
        "cpu": DaskSetupConfig(workload_type="cpu"),
    },
    my_ds,
    operation="mean",
    repeats=3,
)
for r in sorted(results, key=lambda r: r.wall_time_seconds):
    print(r.summary_line())

dask_setup.benchmark.scaling_analysis(ds, operation='mean', worker_counts=(1, 2, 4, 8), *, base_config=None, repeats=1, warmup=False, fallback_on_detection_failure=True, mode='local', plot=False, verbose=False)[source]

Measure parallel scaling across different worker counts.

For each entry in worker_counts, a fresh local cluster is created with max_workers set to that count and the operation is timed. The resulting ScalingResult includes speedup and efficiency relative to the single-worker baseline.

Parameters:

ds (Any) – xarray Dataset or DataArray to benchmark.
operation (Callable[[Any], Any] | str) – Callable or string alias (see benchmark_config()).
worker_counts (Sequence[int]) – Sequence of worker counts to sweep. The first entry is used as the baseline for speedup calculations.
base_config (DaskSetupConfig | None) – Base DaskSetupConfig to use for each run. max_workers is overridden per run. Defaults to DaskSetupConfig() (library defaults).
repeats (int) – Timed repetitions per worker count.
warmup (bool) – Run one un-timed warmup pass before timing.
fallback_on_detection_failure (bool) – Passed to setup_dask_client().
mode (str) – Cluster mode passed to setup_dask_client(). Defaults to "local" so that the sweep always uses a LocalCluster on the current node — this prevents auto-detection from dispatching to PBS/SLURM and submitting child scheduler jobs, which would time out and produce NaN results. Pass "auto" only if you specifically want multi-node dispatch for each sweep step.
plot (bool) – If True, call ScalingResult.plot() and display the figure (requires matplotlib).
verbose (bool) – Print a summary line after each worker count.

Returns:

Scaling metrics for all worker counts.

Return type:

ScalingResult

dask_setup.benchmark.chunk_impact(ds, client, operation='mean', chunk_sizes=None, *, auto_chunks=True, repeats=1, warmup=False, plot=False, verbose=False)[source]

Sweep chunk sizes on a fixed cluster to find the throughput sweet spot.

Unlike benchmark_config() and scaling_analysis(), this function does not create or destroy clusters — it uses the client you provide. The dataset is rechunked before each measurement.

Parameters:

ds (Any) – xarray Dataset or DataArray. Must have known dimension sizes.
client (Client) – Connected Dask client. The cluster remains running throughout.
operation (Callable[[Any], Any] | str) – Callable or string alias (see benchmark_config()).
chunk_sizes (list[dict[str, int]] | None) – List of chunk-size dicts to try (e.g. [{"time": 10, "lat": 50}, {"time": 20, "lat": 100}]). If None and auto_chunks=True, a geometric series of five sizes is generated automatically.
auto_chunks (bool) – Generate chunk sizes automatically when chunk_sizes is None.
repeats (int) – Timed repetitions per chunk size.
warmup (bool) – Run one un-timed warmup pass per chunk size.
plot (bool) – Display the wall-time vs chunk-size plot (requires matplotlib).
verbose (bool) – Print a summary line after each chunk size is tested.

Returns:

Timing results for each chunk spec, plus the recommended spec.

Return type:

ChunkImpactResult

dask_setup.benchmark.run_synthetic_benchmark(profile_name='development', operation='mean', ds_size='small', repeats=1, fallback_on_detection_failure=True, verbose=False)[source]

Run a synthetic benchmark using dask.array (no xarray required).

Creates a random floating-point array of the given size, sets up a Dask cluster from the named profile, and times the requested operation.

This function is the backend for the dask-setup benchmark CLI subcommand.

Parameters:

profile_name (str) – Name of a builtin or user profile to use. "development" is the safe default for quick tests.
operation (str) – One of "mean", "sum", "std", "max", "min".
ds_size (str) – One of "tiny", "small", "medium", "large".
repeats (int) – Timed repetitions.
fallback_on_detection_failure (bool) – Passed to setup_dask_client().
verbose (bool) – Print progress messages.

Return type:

SyntheticBenchmarkResult

class dask_setup.benchmark.BenchmarkResult(name, wall_time_seconds, wall_time_std=0.0, peak_memory_gib=0.0, spill_gib=0.0, n_tasks=0, n_workers=0, tasks_per_second=0.0, errors=<factory>, extra=<factory>)[source]

Bases: object

Metrics for a single benchmarked configuration / parameter setting.

Parameters:

name (str)
wall_time_seconds (float)
wall_time_std (float)
peak_memory_gib (float)
spill_gib (float)
n_tasks (int)
n_workers (int)
tasks_per_second (float)
errors (list[str])
extra (dict[str, Any])

name

Label for this result (config name, worker count, chunk spec, etc.).

Type:: str

wall_time_seconds

Mean elapsed time for the operation across all repeats.

Type:: float

wall_time_std

Standard deviation of per-repeat wall times (0.0 for a single repeat).

Type:: float

peak_memory_gib

Highest total worker process memory (RSS summed across workers) observed during the timed runs, sampled every 0.2 s. If in-flight sampling was unavailable this falls back to a post-run reading and an explanatory note is appended to errors.

Type:: float

spill_gib

Total data written to disk spill storage during the run.

Type:: float

n_tasks

Number of dask-graph tasks the operation required.

Type:: int

n_workers

Number of workers active during this run.

Type:: int

tasks_per_second

Derived: n_tasks / wall_time_seconds (0.0 if either is unknown).

Type:: float

errors

Any non-fatal error messages collected during the run.

Type:: list[str]

extra

Additional metadata attached by the calling function (e.g. {"chunk_size_mb": 128} from chunk_impact()).

Type:: dict[str, Any]

summary_line()[source]

One-line human-readable summary.

Return type:: str

to_dict()[source]

Return a JSON-serialisable dictionary of all fields.

Return type:: dict[str, Any]

class dask_setup.benchmark.ScalingResult(results, worker_counts, speedups, efficiencies)[source]

Bases: object

Results of a worker-count scaling sweep.

Parameters:

results (list[BenchmarkResult])
worker_counts (list[int])
speedups (list[float])
efficiencies (list[float])

results

One entry per worker count, in the order they were benchmarked.

Type:: list[BenchmarkResult]

worker_counts

The worker counts that were swept.

Type:: list[int]

speedups

wall_times[0] / wall_times[i] relative to the single-worker run.

Type:: list[float]

efficiencies

speedups[i] / worker_counts[i] — ideal scaling gives 1.0.

Type:: list[float]

property wall_times: list[float]: Wall times in the same order as worker_counts.

best()[source]

Return the result with the shortest wall time.

Return type:: BenchmarkResult

summary()[source]

Multi-line summary table.

Return type:: str

to_dataframe()[source]

Return results as a pandas.DataFrame.

Raises ImportError if pandas is not installed.

Return type:: Any

plot(figsize=(10, 4), title='Scaling Analysis')[source]

Plot speedup and efficiency curves.

Returns the matplotlib Figure, or None if matplotlib is not installed.

Parameters:

figsize (tuple[float, float])
title (str)

Return type:

Any

class dask_setup.benchmark.ChunkImpactResult(results, chunk_specs, recommended_chunks=None)[source]

Bases: object

Results of a chunk-size sweep on a fixed cluster.

Parameters:

results (list[BenchmarkResult])
chunk_specs (list[dict[str, int]])
recommended_chunks (dict[str, int] | None)

results

One entry per chunk specification, in the order they were tried.

Type:: list[BenchmarkResult]

chunk_specs

The chunk specs that were swept. Alias for chunk_sizes.

Type:: list[dict[str, int]]

recommended_chunks

The chunk spec with the best (lowest) wall time.

Type:: dict[str, int]

property chunk_sizes: list[dict[str, int]]: Alias for chunk_specs for backward compatibility.

optimal()[source]

Return the result with the best (lowest) wall time.

Return type:: BenchmarkResult

to_dataframe()[source]

Return results as a pandas.DataFrame.

Raises ImportError if pandas is not installed.

Return type:: Any

plot(dim=None, figsize=(8, 4), title='Chunk Size Impact')[source]

Plot wall time vs chunk size for one dimension.

Parameters:

dim (str | None) – Dimension name to use as the x-axis. If None, uses the first dimension found in chunk_sizes.
Figure (Returns the matplotlib)
not (or None if matplotlib is)
found. (installed or no suitable dimension is)
figsize (tuple[float, float])
title (str)

Return type:

Any

Environment and logging

dask_setup.environment.is_jupyter()[source]

Return True if running inside a Jupyter kernel.

Covers JupyterLab, classic Jupyter Notebook, VSCode notebooks, Google Colab, and any other environment that uses a ZMQ kernel.

The result is cached after the first call so it is safe to call repeatedly inside hot paths.

Return type:: bool

Examples

from dask_setup.environment import is_jupyter

if is_jupyter():
    print("Running in Jupyter")

dask_setup.environment.get_environment_type()[source]

Return a string describing the detected runtime environment.

Returns:

str – One of "jupyter", "ipython", or "script".
The result is cached after the first call.

Return type:

dask_setup.logging.configure_logging(level='INFO', format_style='structured', use_color=None)[source]

Configure logging for dask_setup package.

Parameters:

level (str | int) – Logging level (string or logging constant)
format_style (str) – “structured” for human-readable, “json” for structured data
use_color (bool | None) – Whether to use color output (None = auto-detect)

Return type:

None

dask_setup.logging.get_logger(module_name)[source]

Convenience function to get a logger for a module.

Parameters:: module_name (str) – Name of the module (e.g., ‘client’, ‘resources’)
Returns:: DaskSetupLoggerAdapter instance
Return type:: DaskSetupLoggerAdapter

Errors

Custom exceptions for the dask_setup package.

exception dask_setup.exceptions.DaskSetupError[source]

Bases: Exception

Base exception for all dask_setup errors.

exception dask_setup.exceptions.InsufficientResourcesError(required_mem, available_mem, suggested_actions=None)[source]

Bases: DaskSetupError

Raised when system resources are insufficient for the requested configuration.

Parameters:

required_mem (float)
available_mem (float)
suggested_actions (list[str] | None)

Return type:

None

exception dask_setup.exceptions.InvalidConfigurationError[source]

Bases: DaskSetupError

Raised when configuration parameters are invalid.

exception dask_setup.exceptions.ResourceDetectionError[source]

Bases: DaskSetupError

Raised when resource detection fails.

class dask_setup.error_handling.ErrorContext(platform_info=<factory>, python_version=<factory>, total_memory_gb=<factory>, cpu_count=<factory>, hpc_environment=None, pbs_jobfs=<factory>, slurm_job_id=<factory>, available_dependencies=<factory>)[source]

Bases: object

Container for error context information.

Parameters:

platform_info (str)
python_version (str)
total_memory_gb (float)
cpu_count (int)
hpc_environment (str | None)
pbs_jobfs (str | None)
slurm_job_id (str | None)
available_dependencies (dict[str, str])

get_environment_summary()[source]

Get a summary of the current environment.

Return type:: str

class dask_setup.error_handling.EnhancedDaskSetupError(message, suggestions=None, context=None, documentation_url=None, error_code=None)[source]

Bases: DaskSetupError

Base class for enhanced dask_setup errors with rich context.

Parameters:

message (str)
suggestions (list[str] | None)
context (ErrorContext | None)
documentation_url (str | None)
error_code (str | None)

get_diagnostic_info()[source]

Get diagnostic information for error reporting.

Return type:: str

class dask_setup.error_handling.ConfigurationValidationError(field_errors, invalid_config=None, context=None)[source]

Enhanced configuration validation error with specific field guidance.

Parameters:

field_errors (dict[str, str])
invalid_config (dict[str, Any] | None)
context (ErrorContext | None)

class dask_setup.error_handling.ResourceConstraintError(resource_type, required, available, units='GB', context=None)[source]

Enhanced resource constraint error with HPC-aware suggestions.

Parameters:

resource_type (str)
required (float)
available (float)
units (str)
context (ErrorContext | None)

class dask_setup.error_handling.DependencyError(missing_package, feature, context=None, alternative_packages=None)[source]

Enhanced dependency error with installation guidance.

Parameters:

missing_package (str)
feature (str)
context (ErrorContext | None)
alternative_packages (list[str] | None)

class dask_setup.error_handling.StorageConfigurationError(storage_issue, storage_path=None, detected_format=None, context=None)[source]

Enhanced storage configuration error with format-specific guidance.

Parameters:

storage_issue (str)
storage_path (str | None)
detected_format (str | None)
context (ErrorContext | None)

class dask_setup.error_handling.ClusterSetupError(setup_issue, cluster_config=None, context=None)[source]