HPC Integration - dask_setup Recipe

This notebook shows how dask_setup works with PBS (NCI Gadi) and SLURM schedulers: what it reads from the job environment, how to generate job scripts, how to reach the dashboard, and how single-node and multi-node runs differ.

Key Learning Points

  • Resource detection from scheduler environment variables

  • mode=: single-node LocalCluster vs multi-node dask-jobqueue clusters

  • Generating PBS/SLURM job scripts with generate_pbs_script() / generate_slurm_script()

  • Dashboard access via SSH tunnel, and DASK_SETUP_LOGIN_HOST

  • Job monitoring commands and HPC best practices

Everything here runs on a laptop: scheduler environments are simulated where needed.

Setup and Imports

import os
import tempfile
import time
from contextlib import contextmanager
from pathlib import Path

import dask.array as da

from dask_setup import (
    DaskClientContext,
    MultiNodeConfig,
    cluster_report,
    detect_cluster_mode,
    generate_pbs_script,
    generate_slurm_script,
)
from dask_setup.dashboard import LOGIN_HOST_ENV, get_login_host
from dask_setup.resources import detect_resources


@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


print("🚀 Setup complete!")

🚀 Setup complete!

HPC Environment Detection

dask_setup sizes the cluster from what the scheduler gave your job, not from the whole machine: a Gadi normalsr node has 104 cores, but a job with ncpus=48 must use 48. Detection order is SLURM → PBS → psutil (the whole machine).

On Gadi, $PBS_MEM is a raw byte count with no unit (e.g. 204010946560 for 190 GiB); dask_setup handles that and the 190gb form.

def show(label):
    r = detect_resources()
    print(f"{label:<28} via {r.detection_method:<7} {r.total_cores:>4} cores "
          f"{r.total_mem_bytes / 2**30:7.1f} GiB   mode={detect_cluster_mode()}")


show("This machine")

with simulated_env(PBS_JOBID="123.gadi-pbs", PBS_ENVIRONMENT="PBS_BATCH",
                   NCPUS="48", PBS_MEM="204010946560"):
    show("Gadi batch job (qsub)")

with simulated_env(PBS_JOBID="124.gadi-pbs", PBS_ENVIRONMENT="PBS_INTERACTIVE",
                   NCPUS="48", PBS_MEM="190gb"):
    show("Gadi interactive (qsub -I)")

with simulated_env(SLURM_JOB_ID="42", SLURM_BATCH_FLAG="1",
                   SLURM_CPUS_ON_NODE="32", SLURM_MEM_PER_NODE="128000"):
    show("SLURM batch job (sbatch)")

This machine                 via psutil     4 cores    15.6 GiB   mode=local
Gadi batch job (qsub)        via PBS       48 cores   190.0 GiB   mode=pbs
Gadi interactive (qsub -I)   via PBS       48 cores   177.0 GiB   mode=interactive
SLURM batch job (sbatch)     via SLURM     32 cores   125.0 GiB   mode=slurm

Single-Node vs Multi-Node: mode=

setup_dask_client(mode=...) decides what kind of cluster you get:

mode=

Cluster

"local"

A LocalCluster on the node you are on, sized from the job’s resources

"pbs" / "slurm"

Scheduler here; workers submitted as separate batch jobs via dask-jobqueue

"interactive"

Workers started on the nodes of an interactive allocation

"auto" (default)

Whatever detect_cluster_mode() returns, as in the mode= column above

⚠️ Inside a batch job, "auto" resolves to "pbs" / "slurm", the multi-node path. For the usual case of one job using the cores it was given, pass mode="local" explicitly:

client, cluster, dask_tmp = setup_dask_client(workload_type="cpu", mode="local")

Job Script Generation

MultiNodeConfig describes the resources for one job; generate_pbs_script() and generate_slurm_script() turn it into a job script that runs your Python file. The same config object drives multi-node clusters (setup_dask_client(mode="pbs", multi_node_config=...)).

env_extra lines are inserted verbatim before your script runs, so write full shell statements.

# One full Gadi normalsr node: 104 cores, ~500 GB
gadi = MultiNodeConfig(
    workload_type="cpu",
    workers_per_node=26,
    cores_per_worker=4,
    mem_per_worker_gb=18,
    walltime="04:00:00",
    queue="normalsr",
    project=os.environ.get("PROJECT", "ab12"),
    job_extra_directives=[
        "-l storage=gdata/xp65+scratch/ab12",
        "-l jobfs=400GB",  # spill space: dask_setup spills to $PBS_JOBFS
        "-l wd",
    ],
    env_extra=[
        "module use /g/data/xp65/public/modules",
        "module load conda/analysis3",
    ],
)
pbs_script = generate_pbs_script(gadi, script_path="analysis.py")
print(pbs_script)

# SLURM sites use their own partitions, accounts and module names
slurm = MultiNodeConfig(
    workload_type="cpu",
    workers_per_node=8,
    cores_per_worker=4,
    mem_per_worker_gb=15,
    walltime="04:00:00",
    queue="work",
    project="myaccount",
    env_extra=["module load python"],
)
slurm_script = generate_slurm_script(slurm, script_path="analysis.py")
print(slurm_script)

#!/bin/bash
#PBS -l ncpus=104,mem=468GB,walltime=04:00:00
#PBS -q normalsr
#PBS -P ab12
#PBS -l storage=gdata/xp65+scratch/ab12
#PBS -l jobfs=400GB
#PBS -l wd

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

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

#!/bin/bash
#SBATCH --ntasks=8
#SBATCH --cpus-per-task=4
#SBATCH --mem=120G
#SBATCH --time=04:00:00
#SBATCH --partition=work
#SBATCH --account=myaccount

# --- Environment ---
set -euo pipefail
module load python

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

Save Scripts and an Analysis Template

We write the scripts to a temporary directory; in practice, write them next to your analysis. The same scripts can be produced from the shell:

dask-setup submit analysis.py --scheduler pbs --queue normalsr --walltime 04:00:00 \
    --workers-per-node 26 --cores-per-worker 4 --mem-per-worker 18

output_dir = Path(tempfile.mkdtemp(prefix="dask_setup_hpc_"))

(output_dir / "dask_analysis.pbs").write_text(pbs_script)
(output_dir / "dask_analysis.slurm").write_text(slurm_script)

analysis_script = """#!/usr/bin/env python
\"\"\"Example dask_setup analysis for a single-node batch job.\"\"\"
import time
from pathlib import Path

import dask.array as da
import pandas as pd
import xarray as xr

from dask_setup import DaskClientContext, cluster_report


def main():
    start = time.monotonic()
    # mode="local": one LocalCluster using the cores and memory PBS/SLURM gave this job.
    # (mode="auto" would pick the multi-node backend inside a batch job.)
    with DaskClientContext(workload_type="cpu", mode="local") as (client, cluster, dask_tmp):
        n_time = 365 * 5
        shape, chunks = (n_time, 180, 360), (365, 90, 180)  # ~23 MB float32 chunks
        ds = xr.Dataset(
            {
                "temperature": (("time", "lat", "lon"), da.random.random(shape, chunks=chunks).astype("float32")),
                "precipitation": (("time", "lat", "lon"), da.random.random(shape, chunks=chunks).astype("float32")),
            },
            coords={"time": pd.date_range("2000-01-01", periods=n_time, freq="D")},
        )

        climatology = ds.groupby("time.month").mean().compute()

        out = Path("results")
        out.mkdir(exist_ok=True)
        climatology.to_netcdf(out / "monthly_climatology.nc")
        print(cluster_report(client, start_time=start).summary_line())


if __name__ == "__main__":
    main()
"""
(output_dir / "analysis.py").write_text(analysis_script)

for f in sorted(output_dir.iterdir()):
    print(f"📄 {f}")
print(f"\nSubmit with:  qsub {output_dir / 'dask_analysis.pbs'}")

📄 /tmp/dask_setup_hpc_fuvdtbh7/analysis.py
📄 /tmp/dask_setup_hpc_fuvdtbh7/dask_analysis.pbs
📄 /tmp/dask_setup_hpc_fuvdtbh7/dask_analysis.slurm

Submit with:  qsub /tmp/dask_setup_hpc_fuvdtbh7/dask_analysis.pbs

SSH Tunneling for Dashboard Access

With dashboard=True (the default), setup_dask_client() prints the dashboard address and the ssh -N -L ... command to run on your laptop. It fills in the login host from, in order:

  1. $DASK_SETUP_LOGIN_HOST, if set: put export DASK_SETUP_LOGIN_HOST=gadi.nci.org.au in your job script

  2. The node’s DNS domain (on gadi-cpu-clx-1234.gadi.nci.org.au that is gadi.nci.org.au)

  3. <login-node> as a placeholder when neither is available (e.g. on a laptop)

In Jupyter it shows a clickable link instead.

print(f"Login host here: {get_login_host()}")

with simulated_env(**{LOGIN_HOST_ENV: "gadi.nci.org.au"}):
    print(f"With {LOGIN_HOST_ENV} set: {get_login_host()}\n")
    with DaskClientContext(workload_type="cpu", max_workers=2, mode="local") as (client, cluster, tmp):
        pass  # the tunnel hint is printed when the cluster starts

Login host here: vfgwi3ehmjzeppmxekjceiydjg.gx.internal.cloudapp.net
With DASK_SETUP_LOGIN_HOST set: gadi.nci.org.au

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-2868
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

Test Cluster Creation

A quick end-to-end check you can run at the top of any job: start a cluster, run something, and print the run summary.

start = time.monotonic()
with DaskClientContext(workload_type="mixed", max_workers=4, dashboard=False, mode="local") as (
    client,
    cluster,
    temp_dir,
):
    workers = client.scheduler_info()["workers"]
    print(f"✅ {len(client.nthreads())} workers, spill directory {temp_dir}")
    for addr, w in list(workers.items())[:3]:
        print(f"   {addr} | {w['nthreads']} threads | {w['memory_limit'] / 2**30:.1f} GiB")

    x = da.random.random((4000, 4000), chunks=(1000, 1000))
    result = ((x + x.T) * 2).mean().compute()
    print(f"\nResult: {result:.6f}")
    print(cluster_report(client, start_time=start).summary_line())

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-2868
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
✅ 2 workers, spill directory /tmp/dask-2868
   tcp://127.0.0.1:33707 | 2 threads | 5.8 GiB
   tcp://127.0.0.1:46523 | 2 threads | 5.8 GiB

Result: 1.999528
wall=0.8s | workers=2 | mem=0.06 GiB/max

Job Monitoring Commands

Let’s learn about monitoring jobs and resources on different HPC schedulers.

def demonstrate_job_monitoring():
    """Show job monitoring commands for PBS and SLURM."""
    
    print("📊 JOB MONITORING COMMANDS")
    print("=" * 60)
    
    # Check current environment
    pbs_job = os.environ.get('PBS_JOBID')
    slurm_job = os.environ.get('SLURM_JOB_ID')
    
    if pbs_job:
        print(f"🐚 PBS JOB MONITORING (Current Job: {pbs_job})")
        print(f"─" * 50)
        
        monitoring_commands = {
            "Job Status": f"qstat -f {pbs_job}",
            "Resource Usage": f"qstat -f {pbs_job} | grep resources_used",
            "Queue Information": "qstat -Q",
            "All Your Jobs": "qstat -u $USER",
            "Node Information": "cat $PBS_NODEFILE",
            "JobFS Usage": "df -h $PBS_JOBFS",
            "Job History": f"tracejob {pbs_job}",
            "Detailed Queue Info": "pbsnodes -a | grep -A5 -B1 'state = free'"
        }
        
        for description, command in monitoring_commands.items():
            print(f"   {description:20}: {command}")
        
        # Try to get actual job information
        print(f"\n📈 CURRENT JOB INFORMATION:")
        try:
            result = subprocess.run(
                ['qstat', '-f', pbs_job], 
                capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                # Parse interesting fields
                lines = result.stdout.split('\n')
                relevant_lines = []
                for line in lines:
                    if any(keyword in line.lower() for keyword in 
                          ['job_name', 'queue', 'resources_used.walltime', 
                           'resources_used.mem', 'job_state', 'exec_host']):
                        relevant_lines.append(line.strip())
                
                for line in relevant_lines[:10]:  # Show first 10 relevant lines
                    print(f"   {line}")
            else:
                print(f"   qstat command failed (not available or job not found)")
        except Exception as e:
            print(f"   Could not run qstat: {e}")
            
    elif slurm_job:
        print(f"🔰 SLURM JOB MONITORING (Current Job: {slurm_job})")
        print(f"─" * 50)
        
        monitoring_commands = {
            "Job Status": f"squeue -j {slurm_job}",
            "Detailed Job Info": f"scontrol show job {slurm_job}",
            "Resource Usage": f"sstat -j {slurm_job} --format=AveCPU,AveRSS,MaxRSS",
            "Job Efficiency": f"seff {slurm_job}",
            "All Your Jobs": "squeue -u $USER",
            "Partition Info": "sinfo",
            "Node Information": f"squeue -j {slurm_job} -o %N",
            "Job History": f"sacct -j {slurm_job} --format=JobID,JobName,Partition,Account,AllocCPUS,State,ExitCode"
        }
        
        for description, command in monitoring_commands.items():
            print(f"   {description:20}: {command}")
        
        # Try to get actual job information
        print(f"\n📈 CURRENT JOB INFORMATION:")
        try:
            result = subprocess.run(
                ['squeue', '-j', slurm_job, '-o', '%i,%j,%t,%M,%l,%D,%C'], 
                capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                lines = result.stdout.strip().split('\n')
                if len(lines) >= 2:
                    headers = lines[0].split(',')
                    values = lines[1].split(',')
                    for header, value in zip(headers, values):
                        print(f"   {header}: {value}")
                else:
                    print(f"   Job info: {result.stdout}")
            else:
                print(f"   squeue command failed (not available or job not found)")
        except Exception as e:
            print(f"   Could not run squeue: {e}")
    
    else:
        print(f"💻 LOCAL SYSTEM MONITORING")
        print(f"─" * 30)
        
        monitoring_commands = {
            "Interactive Process Viewer": "htop",
            "Memory Usage": "free -h",
            "Disk Usage": "df -h",
            "System Load": "uptime",
            "Process Tree": "pstree -p",
            "Python Processes": "ps aux | grep python",
            "Dask Dashboard": "http://localhost:8787 (if enabled)"
        }
        
        for description, command in monitoring_commands.items():
            print(f"   {description:25}: {command}")
    
    # General monitoring tips
    print(f"\n💡 MONITORING BEST PRACTICES:")
    print(f"   • Check resource usage regularly during long jobs")
    print(f"   • Monitor memory usage to detect spilling")
    print(f"   • Watch for worker failures or restarts")
    print(f"   • Use dashboard task stream to identify bottlenecks")
    print(f"   • Set up job notifications for completion/failure")
    
    print(f"\n🚨 WARNING SIGNS TO WATCH FOR:")
    print(f"   • Memory usage near limits (>90%)")
    print(f"   • High I/O wait times")
    print(f"   • Frequent worker restarts")
    print(f"   • Tasks stuck in pending state")
    print(f"   • Unusual network traffic patterns")

# Demonstrate monitoring
demonstrate_job_monitoring()

📊 JOB MONITORING COMMANDS
============================================================
💻 LOCAL SYSTEM MONITORING
──────────────────────────────
   Interactive Process Viewer: htop
   Memory Usage             : free -h
   Disk Usage               : df -h
   System Load              : uptime
   Process Tree             : pstree -p
   Python Processes         : ps aux | grep python
   Dask Dashboard           : http://localhost:8787 (if enabled)

💡 MONITORING BEST PRACTICES:
   • Check resource usage regularly during long jobs
   • Monitor memory usage to detect spilling
   • Watch for worker failures or restarts
   • Use dashboard task stream to identify bottlenecks
   • Set up job notifications for completion/failure

🚨 WARNING SIGNS TO WATCH FOR:
   • Memory usage near limits (>90%)
   • High I/O wait times
   • Frequent worker restarts
   • Tasks stuck in pending state
   • Unusual network traffic patterns

HPC Best Practices

Requesting resources

  • Request jobfs for spill (#PBS -l jobfs=...); dask_setup spills to $PBS_JOBFS automatically

  • Size memory from your data: chunks × concurrent tasks, plus headroom for spill thresholds

  • Include gdata/xp65 in -l storage= if you load the conda/analysis3 modules

Choosing the topology

  • NetCDF/HDF5 input → workload_type="cpu"; Zarr / object storage → "io"

  • "io" always runs one worker process: max_workers does not apply to it

  • Tasks that need more memory than usable / n_workers → fewer workers (max_workers=) or smaller chunks. Raising reserve_mem_gb takes memory away from workers

Single node vs multi-node

  • One node is simpler and avoids network shuffles: pass mode="local" in batch jobs

  • Go multi-node (mode="pbs" + MultiNodeConfig) when one node’s memory or cores are not enough

Dashboard

  • Set DASK_SETUP_LOGIN_HOST in the job script so the printed tunnel command is copy-pasteable

  • Check the tunnel works on a short job before a long one

Robustness

  • Use DaskClientContext so workers are shut down even when the analysis fails

  • Print cluster_report(...).summary_line() at the end: wall time, memory and spill go in the job log

Conclusion

Key Takeaways:

  1. Detection: dask_setup sizes the cluster from $NCPUS/$PBS_MEM or $SLURM_*, not the whole machine

  2. Mode: mode="local" for one node; inside a batch job "auto" selects the multi-node backend

  3. Job scripts: generate_pbs_script() / generate_slurm_script() or dask-setup submit

  4. Dashboard: copy the printed tunnel command; set DASK_SETUP_LOGIN_HOST to complete it

  5. Spill: request jobfs and let dask_setup route spill to $PBS_JOBFS

Next Steps:

  • Recipe 08 runs a full analysis workflow that scales from laptop to Gadi

  • The Multi-Node wiki page covers mode="pbs" clusters in depth