Pseudo Code nb2slurm¶

In these 4 notebooks: 0_config.ipynb, 1_build.ipynb, 2_submit.ipynb and 3_sync.ipynb we will explain the workings of nb2slurm in a ‘pseudo’ code way. In the montecarlo carlo example these four notebooks are condensed into one. But we do recommend to split them up if possible.

0. Configure your run¶

This is the only notebook you normally edit. Fill in the cells below, then run it top to bottom. It saves two files in the project root:

control_config.json — all your settings, read by 1_build, 2_submit and 3_sync (so they always agree).
jobs.json — the nested list of jobs to run.

Re-run this notebook whenever you change a setting.

Project¶

Name, the notebooks to run in order, and what varies per job.

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 (below)

The jobs to run¶

A nested dict: each root-to-leaf path is one SLURM job, and the levels line up with varying above. nb2slurm builds the matching output folders and submits one job per leaf.

Format: a dict nests one more level, a list at the bottom means several jobs sharing that parent, and None/[] ends the path. It’s plain JSON, so for small runs type it by hand (below) and for big runs generate it in Python (a comprehension, a CSV, an API query, …) — see the README for the rules. The only thing that matters is that the finished dict lands in jobs.json.

(this is a simplified 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"],
    },
}

HPC connection¶

Your cluster login details. The project dir is derived so it can’t drift.

No SSH key yet? nb2slurm.generate_key(key_type="ed25519") creates one and prints the public key to register with your HPC. If the key has a passphrase, ssh-add it (so rsync push/pull work too) or pass passphrase= to SSHConfig below.

username_on_hpc = "me"
hpc_host = "spider.surf.nl"
project_dir_on_hpc = f"/home/{username_on_hpc}/{project_name}"
ssh_key_file = "~/.ssh/id_ed25519"

Resources per job¶

What each SLURM job asks for, and where outputs go.

time_for_job_to_run = "04:00:00"  # HH:MM:SS wall-clock limit
cpus_per_job = 2
nodes_per_job = 1
output_directory = "output"  # could be "/scratch/me/output"

The time_for_job_to_run or wall-clock time lets the SLURM workload manager know how much time to allocate for your job. cpus_per_job are the number of cpu cores for the job. nodes_per_job are the number of nodes per job and is defaulted to 1. When it is one you request all the cores to be on the same compute node, which consists of 1 or 2 cpus. This is useful for high core count jobs like computational fluid dynamics for example. output_directory will define where the output of your project is stored, the output will follow the jobs structure.

We also provide the user with concurrency of the jobs in jobs_at_once. Meaning that the workload manager will have a number of jobs_at_once running per branch. This could take load of filesystems if needed. If set to 0 all jobs may run at once if there is available hardware.

jobs_at_once = 3  # max running in parallel, 0 means all jobs run at the same time with no concurrency

Conda environment (optional but recommended)¶

Skip if your cluster already provides Python with all the required packages (set conda_env/setup instead). Otherwise nb2slurm builds this environment + kernel for you in 1_build.

Need a different environment for one or two notebooks (e.g. a calibration step)? Add kernels={"notebooks/step_8.ipynb": "myenv2"} and extra_environments=[nb2slurm.Environment(name="myenv2", kernel="myenv2", ...)] to the Workflow in the Save cell below — see the README for details. The single-environment setup here covers the common case.

conda_env_name = "myenv"
kernel_name = "myenv"
conda_packages = ["xarray", "numpy"]  # EXAMPLE - list what your notebooks import
pip_packages = ["nb2slurm", "ewatercycle"]  # EXAMPLE - but always keep nb2slurm

Data mounts (optional)¶

rclone mounts should be set up before running. You can provide more if needed in this list as well.

data_mounts = [
    {"remote": "dcache:/climate-data/caravan", "mountpoint": "/scratch/caravan"},
]

Save¶

Builds the objects from your settings and writes control_config.json and jobs.json. The other notebooks load these — you don’t edit them by hand, but it is possible.

import json
import pathlib
import nb2slurm

env = nb2slurm.Environment(
    name=conda_env_name,
    kernel=kernel_name,
    conda_packages=conda_packages,
    pip_packages=pip_packages,
)
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,
)
cfg = nb2slurm.SSHConfig(
    host=hpc_host,
    user=username_on_hpc,
    remote_dir=project_dir_on_hpc,
    key_filename=ssh_key_file,
)

pathlib.Path("jobs.json").write_text(json.dumps(jobs, indent=2))
nb2slurm.save_config("control_config.json", workflow=wf, ssh=cfg)
print("saved control_config.json and jobs.json")