HPC patterns

How to take a Study/Experiment from your laptop to a cluster: declare its dependencies, emit a self-contained kit, provision the environment, submit a Slurm array job, and collect results. This page collects the HPC-relevant flags that span run and workflow and shows the canonical cluster patterns.

End to end (Slurm)

Take a large parameter sweep: a 4 × 39 × 10 grid over intrinsic frequency, coupling and conduction speed (1 560 cells × 10 trials). It runs as a Slurm array job where each task simulates a chunk of the grid and the backend vmap/pmap-s its share:

# 1. Emit the kit locally (tvbo finds the recipe's code/ itself — no PYTHONPATH prefix).
#    Bake env activation into the kit via `setup` so each array task self-activates.
#    Hardcode the conda base (from `conda info --base`): a batch shell has no `conda`
#    command, so `$(conda info --base)` would fail on the node — use the real path.
tvbo workflow slurm mystudy.yaml \
    --experiment "Full IF x K x v grid" \
    --backend tvboptim \
    --set chunk=40 \
    --set slurm.partition=medium \
    --set slurm.time=08:00:00 \
    --set slurm.mem=16G \
    --set slurm.mail_type=END,FAIL \
    --set slurm.mail_user=you@example.org \
    --set 'slurm.setup=["source /opt/conda/etc/profile.d/conda.sh","conda activate mystudy_env"]' \
    -o ./cluster_kit

# 2. Ship it to the cluster
rsync -a ./cluster_kit/ hpc:~/runs/mysweep/

# 3. On the cluster: provision the env once, then launch with one tvbo command.
#    `submit` runs the same sbatch array + finalize chain — no sbatch by hand.
ssh hpc
cd ~/runs/mysweep
conda env create -f environment.yml      # or: pip install -r requirements.txt
tvbo workflow submit .                     # 40 array tasks + gather; email on END/FAIL

Because setup is baked into run.sbatch, each array task activates mystudy_env itself — the run does not depend on whatever shell you submitted from.

The generated run.sbatch runs, per array task:

tvbo run spec/experiment.yaml --backend=tvboptim \
    --slurm-chunk=$SLURM_ARRAY_TASK_ID/40 \
    -o out/.../$SLURM_ARRAY_JOB_ID/$SLURM_ARRAY_TASK_ID

It points at the kit’s own frozen spec/experiment.yaml, so the job is self-contained. Any custom callable/builder modules the recipe references (callable: {module: my_analysis}, builder: {module: my_networks}) are copied into the kit’s code/ and put on PYTHONPATH, so no external recipe or code is needed at run time. Installed dependencies still come from the emitted requirements.txt / environment.yml (they are not bundled).

Chunking a sweep into array tasks (chunk / --slurm-chunk)

Every cell of the sweep’s cartesian product gets a deterministic index j; --slurm-chunk i/N runs cell j iff j % N == i. There are two layers:

  • --set chunk=N (planning): how many Slurm array tasks to split the sweep into. The emitted run.sbatch gets #SBATCH --array=0-(N-1) and passes --slurm-chunk=$SLURM_ARRAY_TASK_ID/N.
  • --slurm-chunk i/N (execution): the 1/N share of the cells one tvbo run invocation actually runs.

This works even when every axis is backend-vectorized. For tvboptim, a 4 × 39 × 10 grid is vectorized inside the backend, so with no chunking it is a single (memory-heavy) vmap job; --set chunk=40 shards it into 40 array tasks that each vmap ~39 cells:

tvbo workflow plan mystudy.yaml --experiment "Full IF x K x v grid" \
    --backend tvboptim --set chunk=40
# …
# chunk                : 40  →  40 array task(s)

Pick N so each task’s batch fits node memory. More tasks = smaller batches = lower peak memory but more scheduling overhead. The Koller 4 × 39 × 10 grid at 10 trials is ~1.6 GB for two cells, so the whole sweep in one batch would need ~1.25 TB; that is what chunking avoids.

How a task runs its share

--slurm-chunk i/N slices the backend’s own vectorized batch, not a Python cell loop. On tvboptim the sweep is a Space, and task i runs Space[i::N]: the strided subset of grid points, still vmapped, with peak memory scaled by 1/N. Each task saves its slice as a flat point-dimension array (one entry per cell it ran), so the shards partition the grid and reassemble by value (see Two-stage runs).

Only a backend that vectorizes every swept axis can shard this way, and tvbo run checks that against the backend’s ontology capabilities (the same vectorize_axes the planner uses). A backend that fans an axis instead (e.g. tvb, which vectorizes nothing) rejects --slurm-chunk and points you to the Snakemake or Nextflow emitter, where each cell is its own task. See Choosing an engine for that split.

Parallelism inside a task (parallel_mode)

Within one array task the backend parallelises its cells/trials via the exploration’s parallel_mode (set on Exploration.parallel_mode in the recipe, or leave auto):

Mode Behaviour Use when
vmap all cells batched at once (fast, n × working-set memory) small batch fits memory
lax_map cells run sequentially (bounded memory) large batch / CPU node
pmap sharded across devices (jax.pmap) multi-GPU / multi-device node
auto vmap when the estimated batch fits, else lax_map default

Combine --set chunk=N (across nodes) with parallel_mode (within a node) for the fastest memory-safe layout: e.g. GPU nodes with pmap, or many CPU tasks with auto.

Environment provisioning (environment.requirements)

Declare the run’s dependencies on the experiment’s schema-native environment.requirements, and the kit emits a matching requirements.txt (pip) and environment.yml (conda):

# in the experiment
environment:
  name: mystudy_env
  requirements:
    - {name: tvboptim, package: tvboptim}
    - {name: jax, package: jax}
    - {name: libigl, package: libigl, version_spec: ">=2.5"}
    - {name: potpourri3d, package: potpourri3d}
    - {name: scikit-sparse, package: scikit-sparse}

→ emitted environment.yml:

name: mystudy-experiment
channels: [conda-forge, defaults]
dependencies:
  - python>=3.10
  - pip
  - pip:
      - tvboptim
      - jax
      - libigl>=2.5
      - potpourri3d
      - scikit-sparse

Each requirement’s source_url/url (if set) is emitted verbatim (for git or wheel URLs); otherwise package + version_spec. Provision on the cluster with conda env create -f environment.yml or pip install -r requirements.txt.

Email notification

tvbo workflow slurm … \
    --set slurm.mail_type=END,FAIL \
    --set slurm.mail_user=you@example.org

emits #SBATCH --mail-type=END,FAIL and #SBATCH --mail-user=….

Resources (--set slurm.*)

Any key under slurm.* becomes an #SBATCH directive:

tvbo workflow slurm … \
    --set slurm.account=brain \
    --set slurm.partition=medium \
    --set slurm.time=08:00:00 \
    --set slurm.mem=16G \
    --set slurm.cpus_per_task=8 \
    --set slurm.gres=gpu:1

Supported keys: account, partition, time, mem, cpus_per_task, gres, mail_type, mail_user, plus modules (a list of module load …) and venv (activated before the run).

Engine dispatch (one command, end to end)

tvbo workflow run slurm mystudy.yaml --experiment 2,3,20,30 --backend tvboptim

tvbo workflow run <engine> emits the kit and submits it. A comma-separated --experiment (2,3,20,30) emits and submits one kit/job per experiment — on Slurm they run in parallel; the ids match an experiment’s key, name, label, or numeric id. The lower-level equivalents are tvbo run mystudy.yaml --engine slurm (emit + sbatch) and tvbo workflow slurm … -o kit then sbatch kit/run.sbatch. For snakemake / nextflow, tvbo workflow run invokes the engine and blocks until it finishes; the two-step forms write the artefact but leave you to run snakemake --cores all / nextflow run main.nf (Slurm’s sbatch is fire-and-forget).

Containers (--set container=IMAGE)

tvbo workflow slurm … --set container=ghcr.io/the-virtual-brain/tvbo:0.7.0

The generated run.sbatch wraps the run in singularity exec <image> …. A container is an alternative to environment.yml when the cluster prefers images.

Two-stage runs (simulate → analyse)

A large exploration usually declares no observations:. The array job produces the raw simulation results, and the figure-level metrics are computed in a second pass by analysis scripts over those results (mirroring the common simulate-then-analyse HPC pattern). Submit the simulation array first, then a dependent analysis job:

sim=$(sbatch --parsable run.sbatch)
sbatch --dependency=afterok:$sim analysis.sbatch

Each array task writes its own slice. Because --slurm-chunk i/N slices the backend’s vectorized batch (tvboptim: Space[i::N]), task i runs only its cells and saves them as a flat point-dim array (netCDF), each cell carrying its swept parameter values as coordinates. The analysis pass stitches the slices back into the full grid by value with tvbo.data.reassemble_shards, so the result is correct no matter how the sweep was sharded:

from tvbo.data import reassemble_shards

grid = reassemble_shards("out/.../$SLURM_ARRAY_JOB_ID", to_grid=True)
# grid: xr.DataArray with one dim per swept parameter, addressed by value
# grid.sel(**{"Kuramoto.omega_mean_hz": 10.0, "conduction_speed": 3.0})

Provenance in the kit’s README

Every kit’s README.md ends with a reproducibility block: backend, container, cell/chunk/array-task counts, vectorized vs. workflow-fanned axes, and every --set KEY=VALUE with its source. The kit documents itself.

See also