nb2slurm — how it works (illustrative walkthrough)¶

This notebook is pseudo-code: it shows the intended end-to-end flow of nb2slurm using the real API. The cells are not meant to run yet — there is no cluster connected and the project notebooks don’t exist. Read it top to bottom as the story of how a user goes from one subject to many.

The idea: you have a notebook workflow that works for one subject (one catchment, one region, one number). nb2slurm turns it into the same workflow run for many subjects on a SLURM HPC — all driven from this notebook, no command line.

New to HPC/SLURM/conda? See docs/hpc-for-beginners.md first.

NOTE This is once you have a workflow in notebooks made already! If you are starting and you do not have your workflow (ready)? See docs/setup_notebooks.ipynb

0. Install (once)¶

On your laptop / cloud / login node:

# %pip install nb2slurm

Your settings — edit this one cell¶

This is the only cell most users need to touch. Everything below reuses these variables, so set them once and run the notebook top to bottom.

# =====================================================================
#  YOUR SETTINGS  —  edit these, then run the notebook top to bottom.
# =====================================================================

# --- Project ---------------------------------------------------------
project_name = "myproject"
notebooks = [
    "notebooks/0_settings.ipynb",  # first notebook: writes settings.json
    "notebooks/1_analysis.ipynb",
    "notebooks/2_more.ipynb",
]
varying = ["country", "region", "scenario"]  # the levels in jobs.json
jobs_json = "jobs.json"  # the single source of truth (see next section)

# --- HPC connection (SLURM over SSH) ---------------------------------
slurm_host = "spider.surfsara.nl"  # login-node hostname
username_on_slurm = "me"  # your cluster username
project_dir_on_slurm = f"/home/{username_on_slurm}/{project_name}"  # where the project lives on the cluster
ssh_key_file = "~/.ssh/id_ed25519"  # your private SSH key

# --- Resources per job ----------------------------------------------
time_for_job_to_run = "04:00:00"  # HH:MM:SS wall-clock limit
cpus_per_job = 2
nodes_per_job = 1
jobs_at_once = 3  # max running in parallel
output_directory = "output"  # could be "/scratch/me/runs"

# --- Conda environment (OPTIONAL; skip if the cluster provides one) --
conda_env_name = "myenv"
kernel_name = "myenv"  # Jupyter kernel papermill uses
conda_packages = ["xarray", "numpy"]  # EXAMPLE - list what your notebooks import
pip_packages = ["nb2slurm", "ewatercycle"]  # EXAMPLE - but always keep nb2slurm

# --- Data mounts (OPTIONAL rclone mounts) ---------------------------
data_mounts = [
    {"remote": "dcache:/climate-data/caravan", "mountpoint": "/scratch/caravan"},
]

1. The jobs file — one nested JSON drives everything¶

jobs.json is the single source of truth. It is a nested dict where each root-to-leaf path is one SLURM job, and the levels line up with varying (here: country -> region -> scenario). nb2slurm reads this one file to:

build the output directory tree (output/<country>/<region>/<scenario>/), and
submit the right jobs (one per leaf path).

You write this once (or generate it in Python — it is just a dict). You never build folders by hand in a notebook.

import json
import pathlib

example_jobs = {
    "Netherlands": {
        "north": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
        "south": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
    },
    "Germany": {
        "north": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
        "south": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
        "east": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
        "west": ["green_climate", "climate_as_we_are", "heavy_industrialization"],
    },
}
pathlib.Path(jobs_json).write_text(json.dumps(example_jobs, indent=2))
print(pathlib.Path(jobs_json).read_text())

2. The project layout nb2slurm assumes¶

Everything lives at the project root — notebooks/ and scripts/ are siblings (no project/ wrapper):

myproject/                 <- project root (= remote_dir on the cluster)
|- notebooks/
|  |- 0_settings.ipynb     <- parameterised by nb2slurm; writes settings.json
|  |- 1_analysis.ipynb     <- reads settings.json
|  |- 2_more.ipynb
|- jobs.json               <- the nested job/output hierarchy (you write this)
|- scripts/                <- GENERATED by wf.build()
|- output/                   <- output tree, mirrors jobs.json (created from it)
|- done/done.csv           <- which jobs finished (idempotency)
|- environment.yml         <- GENERATED by the Environment helper
|- control.ipynb           <- this notebook, see monte carlo example for a all-in-one control notebook

You only write notebooks/ and jobs.json. nb2slurm generates the rest.

3. The notebook contract¶

0_settings.ipynb does not build any folders — nb2slurm already made this job’s output dir from jobs.json, and passes it in as outdir. The first notebook just records the run’s details:

# --- cell tagged 'parameters' (papermill fills these per job) ---
country = "Netherlands"
region = "north"
scenario = "green_climate"
outdir = "."

import nb2slurm

nb2slurm.Settings.write(
    outdir,
    {
        "country": country,
        "region": region,
        "scenario": scenario,
        "outdir": outdir,
    },
)

Every later notebook has a parameters cell with just settings_path:

# --- cell tagged 'parameters' ---
settings_path = "settings.json"

import nb2slurm

settings = nb2slurm.Settings.load(settings_path)
# ... do the analysis for this one job, saving into settings["outdir"] ...

4. Describe the environment your code needs (optional)¶

Your libraries must be installed on the cluster as a conda environment, and papermill needs it registered as a Jupyter kernel. nb2slurm can build both for you.

This step is optional. Many clusters already provide Python through a module system or a shared environment. If so, skip Environment and instead point Workflow at an existing conda env (conda_env="hydro") or give it setup lines (setup=["module load 2023", "source /opt/envs/hydro/bin/activate"]). The only hard requirement is that kernel= names a Jupyter kernel that exists on the cluster.

import nb2slurm

env = nb2slurm.Environment(
    name=conda_env_name,
    kernel=kernel_name,  # must match Workflow(kernel=...)
    conda_packages=conda_packages,
    pip_packages=pip_packages,
)
print(env.to_yaml())  # preview the environment.yml that will be written

5. Describe the workflow¶

List your notebooks in run order, name the kernel, point at jobs.json, and declare the SLURM resources. Attaching environment=env keeps the names in sync.

wf = nb2slurm.Workflow(
    name=project_name,
    notebooks=notebooks,
    kernel=kernel_name,
    varying=varying,
    jobs_json=jobs_json,
    resources=dict(nodes=nodes_per_job, cpus=cpus_per_job, time=time_for_job_to_run),
    mounts=data_mounts,
    concurrency=jobs_at_once,
    output_dir=output_directory,
    environment=env,
)

6. Generate the SLURM / runner scripts¶

build() renders everything into scripts/ (and writes environment.yml).

for kind, path in wf.build().items():
    print(f"{kind:13s} -> {path}")

# scripts/run_workflow.py   papermill driver: skip-if-done -> nb0 -> rest -> mark done
# scripts/job.slurm         #SBATCH + conda activate + rclone mounts + run the driver
# scripts/submit_batch.sh   simple submitter (all at once) -- the easy one to read
# scripts/submit_jobs.sh    same, but throttles how many run at once
# scripts/cancel_jobs.sh    cancel jobs by name
# scripts/jobs.txt          flat job list generated from jobs.json (the bash scripts read this)
# scripts/structure.json    the resolved config
# environment.yml           the conda spec

7. Build the output directory tree (optional)¶

build_outputs() reads jobs.json and pre-creates the whole output/... tree. This is optional — each job also creates its own output dir at run time — but it lets you inspect the layout, or stage input files, before submitting.

for rel, path in wf.build_outputs().items():
    print(rel, "->", path)

8. Prove it works on ONE job, locally¶

The ‘duality of use’: before scaling out, check the chain still runs on your own machine. Just open the notebooks in Jupyter and run them top to bottom — these are the same notebooks SLURM will execute, only with a different subject and outdir injected per job.

To exercise the generated driver instead (skip-if-done, the settings hand-off, the notebook order), run it for a single job:

python scripts/run_workflow.py Netherlands north green_climate

9. Connect to the HPC¶

From here on, nb2slurm runs sbatch / squeue / scancel for you over SSH. The only one-time setup that is genuinely yours: an account + SSH key on the cluster.

cfg = nb2slurm.SSHConfig(
    host=slurm_host,
    user=username_on_slurm,
    remote_dir=project_dir_on_slurm,
    key_filename=ssh_key_file,
)

cfg.test_connection()  # quick OK/FAIL sanity check before pushing/submitting

# No key yet? nb2slurm.generate_key(key_type="ed25519") makes one and prints the
# public key to register with your HPC (its key-upload page, or authorized_keys).
# If the key has a passphrase, `ssh-add` it once so nb2slurm AND rsync connect
# without prompts (or pass passphrase= to SSHConfig for the paramiko calls).

10. Create the environment on the cluster (one-time, optional)¶

Skip this if you’re using a cluster-provided environment (conda_env= or setup= on the Workflow). Only needed when you defined an Environment above.

RUN ONCE successfully.

wf.create_environment(ssh=cfg)  # conda/mamba env create + ipykernel install
# streams progress live and runs non-interactively (no [Y/n] prompts). Idempotent:
# re-running updates the env in place.

"""If something goes wrong"""
# wf.remove_environment(ssh=cfg)                  # delete the env + its Jupyter kernel
# wf.create_environment(ssh=cfg, overwrite=True)  # force a clean rebuild from scratch

11. Preflight check¶

Confirm the cluster is ready, so a missing piece is a clear message now rather than a cryptic SLURM failure later.

wf.check(ssh=cfg)
# OK  project directory
# OK  notebooks present
# OK  scripts built
# OK  conda env 'myenv'
# OK  kernel 'myenv'

12. Submit the jobs¶

submit() reads jobs.json by default and launches one job per leaf path. Concurrency is throttled with job dependencies so you don’t flood the queue — or set concurrency=0 (on the Workflow, or per call wf.submit(..., concurrency=0)) to submit everything at once with no dependencies, ideal for a few quick jobs. Try dry_run=True first to see the exact sbatch commands.

wf.submit(ssh=cfg, dry_run=True)  # reads jobs.json, previews sbatch
job_ids = wf.submit(ssh=cfg)  # for real

# to run a custom subset instead of jobs.json:
#   wf.submit([("Netherlands", "north", "green_climate")], ssh=cfg)

13. Monitor¶

for row in wf.status(ssh=cfg):
    print(row)

14. Cancel (if needed)¶

wf.cancel(ssh=cfg)  # cancels the jobs submitted this session

15. Re-run safely¶

Each finished job is recorded in done/done.csv. Submitting again skips the ones already done — so you can grow jobs.json (add countries/regions/scenarios) and just re-run, or recover from a partial failure, without redoing work.

That’s the whole loop: edit jobs.json -> build -> check -> submit -> monitor, from one job to the full dataset, without leaving the notebook.