cluster-scripts

name

cluster-scripts

description

Generate Slurm job scripts and submission scripts for the Snellius HPC cluster from DVC pipeline stages. Use this skill whenever the user mentions "cluster", "snellius", "slurm", "sbatch", "submit job", "run on cluster", "HPC", or wants to run experiments on a remote GPU cluster. Also trigger when the user adds new DVC stages and might need corresponding cluster scripts, or asks about job submission, array jobs, or cluster resource allocation.

Generating Snellius Cluster Scripts from DVC Stages

This skill translates DVC pipeline stages (defined in dvc.yaml) into Slurm batch scripts for the Snellius HPC cluster. The project already has a working set of cluster scripts in cluster/ -- read them first to stay consistent with established conventions.

Before writing anything

1. Read dvc.yaml to understand which stages need cluster scripts
2. Read params.yaml to know which parameters each stage uses
3. Read the existing cluster scripts in cluster/ to match the project's conventions
4. Identify the DVC stage's cmd, params, deps, and outs -- these determine the job script's command, parameter reads, and output paths

Architecture: Two-tier orchestration

DVC defines the logical pipeline (what to run, with what params, producing what outputs). Cluster scripts define the physical execution (resource allocation, environment setup, Slurm orchestration). The cluster scripts replicate DVC's commands but replace uv run python with direct python calls (the venv has everything installed).

There are three kinds of scripts, each with a distinct role:

1. Job scripts (job_<name>.sh)

These are the actual Slurm batch scripts that run on compute nodes. Each one handles a single environment (like T-maze or epistemic maze) but dispatches to different stages via environment variables.

The dispatch pattern: Rather than writing one job script per DVC stage, group related stages into a single job script that uses a case statement on STAGE_TYPE (and optionally ANALYSIS, INFERENCE_MODE, STRATEGY) to select what to run. This keeps the cluster directory manageable and makes the relationship between stages explicit.

When to use a separate job script instead: If the stage needs fundamentally different Slurm resources (different memory, time limits, array jobs), it needs its own job script. For example, MiniGrid episodes need 64GB RAM and array jobs, while aggregation needs 1 CPU and 4GB -- these can't share a script because SBATCH directives are fixed at submission time (though they can be overridden with sbatch flags).

2. Submit scripts (submit_<name>.sh)

These run on the login node and orchestrate job submission. They:

• Check if outputs already exist (skip completed stages)
• Submit jobs with the right env vars via sbatch --export
• Wire up dependencies between jobs via --dependency=afterok:<jobid>
• Collect job IDs for dependency chaining

3. Top-level dispatcher (submit_all.sh)

Delegates to environment-specific submit scripts. Handles MiniGrid's array job logic inline because it's structurally different from the dispatch-based environments.

Template: Job script

#!/usr/bin/env bash
#SBATCH --job-name=aif-<environment>
#SBATCH --partition=gpu_a100
#SBATCH --gpus=1
#SBATCH --cpus-per-task=18
#SBATCH --mem=16G
#SBATCH --time=00:30:00
#SBATCH --output=logs/<environment>_%x_%j.out
#SBATCH --error=logs/<environment>_%x_%j.err

# <Brief description of what this script dispatches>
# Dispatches based on STAGE_TYPE env var:
#   <stage_type_1> -- requires <VAR1> (<values>)
#   <stage_type_2> -- requires <VAR2> (<values>)
#   ...

set -euo pipefail

PROJECT_DIR="${SLURM_SUBMIT_DIR:-.}"
cd "$PROJECT_DIR"
mkdir -p logs

export JAX_PLATFORMS="cuda"

source cluster/setup_env.sh

echo "Running <environment> ${STAGE_TYPE} on $(hostname) at $(date)"
python -c "import jax; print(f'JAX devices: {jax.devices()}')"

case "${STAGE_TYPE:?STAGE_TYPE not set}" in
    <stage_type_1>)
        python scripts/<environment>/<script>.py \
            --<flag> "${VAR:?VAR not set}" \
            --output-dir data/<environment>/<subdir>
        ;;
    *)
        echo "ERROR: unknown STAGE_TYPE='${STAGE_TYPE}'"
        exit 1
        ;;
esac

echo "<Environment> ${STAGE_TYPE} completed at $(date)"

Reading parameters from params.yaml

When a stage needs runtime parameters that DVC would normally interpolate via ${param.path}, read them with this helper:

read_param() {
    python -c "import yaml; p=yaml.safe_load(open('params.yaml')); print(p$1)"
}

# Usage:
N_EPISODES=$(read_param "['environment']['n_episodes']")

This approach is used instead of DVC parameter interpolation because compute nodes don't need DVC installed -- they just need params.yaml and pyyaml (which is a project dependency).

When to use read_param vs. hardcoded args: If the DVC stage's cmd uses ${param.path} interpolation, the cluster script needs read_param. If the DVC cmd has hardcoded flags (like --analysis curves), just hardcode them in the cluster script too.

Lightweight stages (T-maze/epistemic pattern)

DVC stages that are quick (seconds to minutes) and don't need parallelism. The DVC cmd maps directly to a python call in a case branch.

Pattern:

• One job script per environment with case dispatch
• STAGE_TYPE selects the category (convergence/experiment/figures)
• Additional env vars (ANALYSIS, INFERENCE_MODE, STRATEGY) select the specific stage
• All stages share the same Slurm resource allocation

Heavy/parallel stages (MiniGrid pattern)

DVC stages that are expensive per-episode and benefit from parallelism. These use Slurm array jobs.

#SBATCH --array=0-99
#SBATCH --output=logs/<name>_%A_%a.out
#SBATCH --error=logs/<name>_%A_%a.err

# Each array task runs one episode
python scripts/<env>/experiment.py \
    --episode-index "$SLURM_ARRAY_TASK_ID" \
    ...

Note the log format difference: %A_%a (array job ID + task ID) instead of %x_%j (job name + job ID).

Array jobs typically need a separate aggregation step (job_aggregate.sh) that runs after all episodes complete, submitted with --dependency=afterany:<array_job_id>.

Template: Submit script

#!/usr/bin/env bash
# Submit all <environment> DVC stages as individual SLURM jobs.
# <Brief description of dependency structure>
# Stages whose output files already exist are skipped.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR/.."
mkdir -p logs

JOB_SCRIPT="cluster/job_<environment>.sh"

# submit_stage <label> <env_vars> [dependency_flag]
submit_stage() {
    local label="$1"
    local env_vars="$2"
    local dep_flag="${3:-}"

    sbatch --parsable \
        --job-name="$label" \
        --export="ALL,${env_vars}" \
        ${dep_flag} \
        "$JOB_SCRIPT"
}

echo "=== <Environment> stages ==="

# --- Independent stages (parallel) ---
for variant in <variant1> <variant2> <variant3>; do
    OUTPUT="data/<environment>/${variant}/results.json"
    if [ -f "$OUTPUT" ]; then
        echo "  ${variant}: skipped (${OUTPUT} exists)"
        continue
    fi
    JID=$(submit_stage "<prefix>-${variant}" "STAGE_TYPE=<type>,VAR=${variant}")
    echo "  ${variant}: $JID"
done

# --- Dependent stages ---
# Collect job IDs, then submit downstream job with --dependency=afterok:<ids>

echo ""
echo "<Environment> jobs submitted. Monitor with: squeue -u \$USER"

Naming conventions

Thing	Pattern	Examples
Job script	job_<environment>.sh	job_tmaze.sh, job_epistemic.sh
Array job script	job_<mode>.sh	job_active.sh, job_marginal.sh
Submit script	submit_<environment>.sh	submit_tmaze.sh, submit_epistemic.sh
Job name (convergence)	<prefix>-conv-<analysis>	tm-conv-curves, ep-conv-lr_sweep
Job name (experiment)	<prefix>-exp-<mode>	tm-exp-active, ep-exp-planning
Job name (figures)	<prefix>-figures	tm-figures, ep-figures
SBATCH job-name	aif-<environment>	aif-tmaze, aif-epistemic
Log files (dispatch)	logs/<env>_%x_%j.{out,err}	logs/tmaze_tm-conv-curves_12345.out
Log files (array)	logs/<name>_%A_%a.{out,err}	logs/active_12345_5.out

Resource allocation guidelines

Workload type	GPU	CPUs	Memory	Time
Lightweight (T-maze, epistemic)	1 A100	18	16G	30min
Heavy episodes (MiniGrid)	1 A100	18	64G	30min-1hr
Convergence analysis	1 A100	18	32G	1hr
Aggregation / post-processing	1 A100	1	4G	10min
Smoke test	1 A100	1	4G	5min

Start with these defaults and adjust based on actual runtime. The partition is always gpu_a100.

Output caching in submit scripts

Before submitting a job, check if its output already exists:

OUTPUT="data/<environment>/<subdir>/results.json"
if [ -f "$OUTPUT" ]; then
    echo "  <stage>: skipped (${OUTPUT} exists)"
    continue
fi

For directory outputs (like figures), check the directory:

if [ -d "$OUTPUT_DIR" ] && [ -n "$(ls -A "$OUTPUT_DIR" 2>/dev/null)" ]; then
    echo "  <stage>: skipped (${OUTPUT_DIR}/ exists)"
fi

Dependency chaining

Collect job IDs from independent stages, then pass them to dependent stages:

# Collect IDs
DEP_JOBS=""
for ...; do
    JID=$(submit_stage ...)
    DEP_JOBS="${DEP_JOBS:+${DEP_JOBS}:}${JID}"
done

# Submit dependent stage
submit_stage "label" "STAGE_TYPE=figures" "--dependency=afterok:${DEP_JOBS}"

afterok = run only if all dependencies succeeded. Use afterany for jobs that should run regardless (like aggregation after array jobs, where some episodes may fail).

Resume-friendly array jobs (MiniGrid pattern)

For array jobs, check which episodes are missing before resubmitting:

missing=$(python scripts/<env>/find_missing_episodes.py "$output_dir" "$N_EPISODES" 2>/dev/null || echo "0-$((N_EPISODES - 1))")
if [ -z "$missing" ]; then
    echo "  all $N_EPISODES episodes complete, skipping"
    return
fi
sbatch --parsable --array="${missing}%${MAX_CONCURRENT}" "$job_script"

The %${MAX_CONCURRENT} throttle limits how many array tasks run simultaneously (courtesy to other cluster users).

Integrating with submit_all.sh

After creating a new environment's job + submit scripts, add a case to submit_all.sh:

case "$ENV" in
    <new_env>)
        bash "$SCRIPT_DIR/submit_<new_env>.sh"
        ;;
    all)
        # Add to the 'all' case:
        bash "$SCRIPT_DIR/submit_<new_env>.sh"
        echo ""
        ;;
esac

Checklist when adding a new environment

1. Read the DVC stages for the new environment in dvc.yaml
2. Identify which stages can share a job script (same resource needs) vs. need separate scripts
3. Write the job script with dispatch case for each stage type
4. Write the submit script with output caching and dependency chaining
5. Add the new environment to submit_all.sh
6. Update cluster/README.md with the new environment's stages, dependencies, and resource allocation
7. Update the smoke test if there are new modules to validate