Monte Carlo π with nb2slurm

A complete, runnable nb2slurm example. The workflow estimates π by Monte Carlo: throw random darts at the unit square and count how many land inside the quarter circle. We run it for several sample sizes and random seeds — each (n_samples, seed) pair is one SLURM job.

This control notebook builds the scripts, shows how to test the chain locally, and then syncs everything to an HPC and runs it there. The workflow notebooks live in notebooks/, the helper in scripts/montecarlo.py, and the jobs in jobs.json.

user_on_hpc = "me"  # your username on the HPC — reused to build the ssh config below
hpc_service_link = "spider.surfsara.nl"  # the link of your HPC provider

1. Describe the workflow

We will be using varying=["n_samples", "seed"] as a root-to-leaf, this matches the two levels in jobs.json.

We also describe the conda environment the jobs run in on the HPC. nb2slurm writes an environment.yml, builds the env on the cluster, and registers a Jupyter kernel of the same name so papermill can execute the notebooks. This workflow only needs matplotlib (the plot step) and nb2slurm itself. The kernel name must match Workflow(kernel=...).

import nb2slurm

# the conda env + Jupyter kernel the jobs run in on the HPC.
# the job needs papermill (to run the notebooks) and nb2slurm (the runner and the
# notebooks import it); installing nb2slurm pulls papermill in as a dependency.
env = nb2slurm.Environment(
    name="montecarlo",
    kernel="montecarlo",  # must match Workflow(kernel=...)
    python="3.11",
    conda_packages=["matplotlib"],
    pip_packages=["nb2slurm", "papermill"],
)

wf = nb2slurm.Workflow(
    name="monte_carlo_pi",
    notebooks=[
        "notebooks/0_settings.ipynb",
        "notebooks/1_simulate.ipynb",
        "notebooks/2_plot.ipynb",
    ],
    kernel="montecarlo",  # papermill runs the notebooks in this kernel
    varying=["n_samples", "seed"],
    jobs_json="jobs.json",
    resources=dict(  # Since it is an easy job there is no need for much cores/time, guestimate your wall-time by running the entire workflow locally
        nodes=1, cpus=1, time="00:05:00"
    ),
    environment=env,  # conda_env is taken from env.name
    concurrency=0,  # 6 quick, independent jobs for this example -> submit all at once, no dependencies
)

2. The jobs

jobs.json is the grid you define — there’s no magic, it’s just a file you create. It’s a nested JSON whose depth matches varying=["n_samples", "seed"]: each top-level key is an n_samples value mapping to the list of seeds to run for it. Every root-to-leaf path is one SLURM job (and one output/<n_samples>/<seed>/ folder), e.g.

{"1000": ["1", "2", "3"], "1000000": ["1", "2", "3"]}

You can hand-write that, but it’s easier to build it from Python so it’s parametrizable — set your levels once and generate every combination:

import json

# choose your levels here — this is the only bit you edit to change the grid
n_samples = [1000, 1_000_000]
seeds = [1, 2, 3]

# full grid: run every seed for every sample size (values are strings, as on the
# SLURM command line; the notebooks cast them back to int)
jobs = {str(n): [str(s) for s in seeds] for n in n_samples}

with open("jobs.json", "w") as f:
    json.dump(jobs, f, indent=2)

print(json.dumps(jobs, indent=2))

Workflow.jobs_from_json() reads that file back and expands every root-to-leaf path into one job (one tuple per (n_samples, seed)):

for job in wf.jobs_from_json():
    print(job)
# -> ('1000', '1') ('1000', '2') ('1000', '3') ('1000000', '1') ...

3. Generate the scripts

Renders scripts/run_workflow.py, job.slurm, the submitters, jobs.txt and — because we gave the workflow an environment — the environment.yml that builds the conda env on the cluster.

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

4. Test it locally

No special runner needed: just open the workflow notebooks and run them. In Jupyter, open notebooks/0_settings.ipynb, then 1_simulate.ipynb, then 2_plot.ipynb, and run each top to bottom (or Run All). Their parameters cells hold sensible defaults (n_samples=1000, seed=1, a local outdir), so they write settings.json, result.json and scatter.png you can inspect straight away.

These are the same notebooks SLURM will execute — nb2slurm just injects a different (n_samples, seed) per job and points outdir at output/<n_samples>/<seed>/ so jobs don’t collide. If the chain works in Jupyter it will work on the cluster. That’s the duality: develop locally, scale out unchanged.

5. Send it to the HPC and run there

The full cluster cycle — nb2slurm hides the ssh/rsync/conda/sbatch:

  1. wf.push(ssh=cfg) — sync your source up to the cluster (notebooks, scripts, jobs.json, environment.yml). Outputs are never uploaded.

  2. wf.create_environment(ssh=cfg) — build the conda env from environment.yml and register its Jupyter kernel. Run once (re-running just updates it).

  3. wf.check(ssh=cfg) — preflight: remote dir, notebooks, scripts, env + kernel.

  4. wf.submit(ssh=cfg) — one SLURM job per (n_samples, seed); each job conda activates the env before running.

  5. wf.status(ssh=cfg) — monitor the queue.

  6. wf.pull(ssh=cfg) — bring output/ (and the done/ ledger) back down.

First describe your login node. Building the SSHConfig does no I/O, so this cell is safe to run as-is; edit the values to match your account. (conda/mamba must be available on the login node — module load it first if your cluster needs that.)

cfg = nb2slurm.SSHConfig(
    host=hpc_service_link,
    user=user_on_hpc,
    remote_dir=f"/home/{user_on_hpc}/monte_carlo_pi",
    key_filename="~/.ssh/id_rsa",
)

First-time SSH setup (run once)

No SSH key yet? nb2slurm.generate_key() makes an RSA keypair locally (it skips if one already exists) and prints the public key. nb2slurm can’t install it for you — most HPCs disable password login, so there’s no way in until your key is registered. You add the public key on your cluster’s side:

  • SURF / Spider / Snellius and many academic clusters: paste it into the key-upload page of their portal (or send it to the service desk) — see your HPC’s docs for “add SSH public key”.

  • If your cluster still allows password login, append it to ~/.ssh/authorized_keys on a login node yourself.

nb2slurm.public_key() reprints the line any time you need to copy it again.

Unlocking a passphrase-protected key. On Snellius (and most clusters) the “password” you’re prompted for is your key’s passphrase, decrypted locally — not a server login. Two ways to handle it:

  • ssh-agent (recommended): run ssh-add ~/.ssh/id_rsa once in a terminal. Then both nb2slurm and rsync (push/pull) authenticate through the agent — no secret in the notebook, no hangs.

  • passphrase= on the config: e.g. SSHConfig(..., passphrase="…"). This unlocks the paramiko calls (test_connection/submit/status/check), but push/pull go through the rsync CLI and still need the agent — so prefer ssh-add for the full workflow.

Leave the key-generation cell commented once your key is registered.

# make a keypair locally (skips if the key already exists) and print the public key.
# ed25519 is the modern, recommended type; the function default is key_type="rsa".
# nb2slurm.generate_key(key_type="ed25519")            # -> ~/.ssh/id_ed25519
# nb2slurm.generate_key(path="~/.ssh/id_rsa", key_type="rsa")

# reprint the public key any time, to copy into your HPC's key-upload page
# print(nb2slurm.public_key("~/.ssh/id_ed25519"))

Check the connection

cfg.test_connection() runs hostname && whoami on the cluster and prints a clear OK/FAIL — a quick sanity check before pushing anything. If your key is encrypted and not unlocked, it tells you to add passphrase= or run ssh-add. It fails fast (bounded timeout) rather than hanging on an unreachable host.

# cfg.test_connection()

Preview the sync and the submit (dry runs)

dry_run=True prints the exact rsync / sbatch commands without touching the network, so you can see precisely what will be sent and run — no cluster needed.

push syncs the project up, excluding output/, done/, .git and caches, so re-pushing notebook edits can never clobber results already on the cluster.

wf.push(ssh=cfg, dry_run=True)

submit issues one sbatch per subset. We set concurrency=0 on the workflow, so all six are queued at once with no dependencies. For a big run, set concurrency=N (on the workflow, or per call wf.submit(..., concurrency=N)) to cap how many run simultaneously — job N then waits for job N‑concurrency via a SLURM afterany dependency.

wf.submit(ssh=cfg, dry_run=True)

The live run

The cluster cycle, one step per cell — run them top to bottom the first time, then re-run individual cells as needed (e.g. status until the queue drains). Left commented so the notebook stays runnable without a cluster; uncomment as you go.

1 · Push — upload your source (notebooks, scripts, environment.yml) to the cluster. output/ is never uploaded, so this can’t overwrite results already there.

# wf.push(ssh=cfg)

2 · Build the environment (one-time) — create the conda env from the uploaded environment.yml and register its kernel. Streams progress live; re-run to update it in place (see Managing the environment below).

# wf.create_environment(ssh=cfg)

3 · Check, then 4 · Submit — preflight the cluster (dir, notebooks, scripts, env, kernel), then launch one job per (n_samples, seed). With concurrency=0 all six are queued at once.

# wf.check(ssh=cfg)       # 3. preflight
# wf.submit(ssh=cfg)      # 4. one SLURM job per (n_samples, seed)

5 · Monitor — the live queue as a list of dicts. Re-run this cell until it comes back empty.

# wf.status(ssh=cfg)      # 5. monitor the queue (re-run until it empties)

6 · Pull results — bring output/ (and the done/ ledger) back down. Safe to re-run as more jobs finish; it only fetches results, never your local notebooks. Then head to analyse_subsets.ipynb.

# wf.pull(ssh=cfg)        # 6. bring output/ (+ done/) back down

Managing the environment

create_environment is idempotent: the first call builds the env, and running it again updates it in place (no Overwrite? prompt). If a build gets interrupted and leaves a half-finished env, or you want a guaranteed clean slate, delete it and rebuild:

  • wf.remove_environment(ssh=cfg) — delete the conda env and its Jupyter kernel. Safe to call when nothing is there yet.

  • wf.create_environment(ssh=cfg, overwrite=True) — remove-then-create in one go.

  • wf.environment.exists(ssh=cfg) — True/False, without changing anything.

# wf.environment.exists(ssh=cfg)                  # check without changing anything -> True/False
# wf.remove_environment(ssh=cfg)                  # delete the env + its Jupyter kernel
# wf.create_environment(ssh=cfg, overwrite=True)  # force a clean rebuild from scratch

Rerunning: clearing the done ledger

done.csv is what makes re-submitting skip subsets that already finished — handy normally, but if you change what 1_simulate.ipynb/2_plot.ipynb actually do (not just their inputs), old subsets won’t rerun on their own since nb2slurm has no way to know the logic changed. Clear the ledger to force everything to rerun:

  • wf.reset_done(ssh=cfg) — deletes done_csv on the cluster (or locally if you omit ssh). The next wf.submit(ssh=cfg) reruns every (n_samples, seed).

To rerun only a subset, edit done.csv by hand (one key per row, e.g. 1000_1) or pass an explicit items= list to submit() instead.

# wf.reset_done(ssh=cfg)   # clears done.csv on the cluster so submit() reruns everything

6. Analyse the results

After wf.pull(...) brings every subset’s output/<n_samples>/<seed>/ back, open analyse_subsets.ipynb to compare them all — π estimate per seed and the $1/\sqrt{n}$ convergence across sample sizes.

What made this nb2slurm-compatible?

  • 0_settings.ipynb has a parameters cell (n_samples, seed, outdir) and writes settings.json.

  • 1_simulate / 2_plot have a settings_path parameters cell and load it.

  • All outputs are written under settings["outdir"], so parallel jobs don’t collide.

  • The simulate step is idempotent (skips if result.json exists).

  • The helper is imported from scripts/ with a sys.path fallback.

  • The plot uses a headless backend + display().

See docs/setup_notebooks.ipynb for the general checklist.