heshamfs/materials-simulation-skills/slurm-job-script-generator

SLURM Job Script Generator

Goal

Generate a correct, copy-pasteable SLURM job script (.sbatch) for running a simulation, and surface common configuration mistakes (bad walltime format, conflicting memory flags, oversubscription hints).

Requirements

• Python 3.10+
• No external dependencies (Python standard library only)
• Works on Linux, macOS, and Windows (script generation only)

Inputs to Gather

Decision Guidance

MPI vs MPI+OpenMP layout

Does the code use OpenMP / threading?
├── NO  → Use MPI-only: cpus-per-task=1
└── YES → Use hybrid: set cpus-per-task = threads per MPI rank
          and export OMP_NUM_THREADS = cpus-per-task

Rule of thumb: if you see diminishing strong-scaling efficiency at high MPI ranks, try fewer ranks with more threads per rank (and measure).

Memory flag selection

• Use either --mem (per node) or --mem-per-cpu (per CPU), not both.
• Follow your cluster’s documentation; some sites enforce one style.
• SLURM --mem units are integer MB by default, or an integer with suffix K/M/G/T (and --mem=0 commonly means “all memory on node”).

Launcher selection (avoid double-wrapping)

• By default (--launcher srun) the generator prepends srun --ntasks=N --cpus-per-task=T to your run command so it inherits SLURM task placement.
• If your run command already starts with `srun`, `mpirun`, or `mpiexec`, pass `--launcher none` so the generator does not double-wrap it. Wrapping srun around mpirun launches N independent copies of mpirun (each spawning its own MPI world); wrapping srun around srun is malformed.
• As a safety net, the generator auto-detects a leading launcher (srun, mpirun, mpiexec, mpiexec.hydra, orterun, aprun, jsrun) and falls back to no-wrap, emitting a warning that recommends --launcher none.

GPU layout

• Ranks-per-GPU = total_ranks / (nodes * gpus_per_node). When --gpus-per-node is set the generator reports results.derived.total_gpus and results.derived.ranks_per_gpu.
• Aim to map ranks evenly to devices. If ntasks is not divisible by the total number of GPUs, the generator emits a "task-to-GPU ratio is not an integer" warning; either adjust ntasks/GPUs or document intentional sharing (e.g. NVIDIA MPS).
• Consider binding options such as --gpu-bind=closest or --ntasks-per-gpu (passed via your run command / --srun-extra). GPU and QoS policies are site-specific — confirm with your cluster docs.

Script Outputs (JSON Fields)

results.derived reports ntasks, ntasks_per_node, cpus_total_requested, and (when applicable) cores_per_node, cpus_per_node_requested, total_gpus, and ranks_per_gpu. results.warnings may include CPU oversubscription, task-to-GPU ratio, and double-launcher warnings.

Workflow

1. Gather cluster constraints (partition/account, GPU policy, memory policy).
2. Choose a process layout (MPI-only vs hybrid MPI+OpenMP).
3. Generate the script with slurm_script_generator.py.
4. Inspect warnings (conflicts, suspicious layouts).
5. Save the generated script as job.sbatch.
6. Submit with sbatch job.sbatch and monitor with squeue.

CLI Examples

# Preview a job script (prints to stdout)
python3 skills/hpc-deployment/slurm-job-script-generator/scripts/slurm_script_generator.py \
  --job-name phasefield \
  --time 00:10:00 \
  --partition compute \
  --nodes 1 \
  --ntasks-per-node 8 \
  --cpus-per-task 2 \
  --mem 16G \
  --module gcc/12 \
  --module openmpi/4.1 \
  -- \
  ./simulate --config config.json
 
# Write to a file and also emit structured JSON
python3 skills/hpc-deployment/slurm-job-script-generator/scripts/slurm_script_generator.py \
  --job-name phasefield \
  --time 00:10:00 \
  --nodes 1 \
  --ntasks 16 \
  --cpus-per-task 1 \
  --out job.sbatch \
  --json \
  -- \
  /bin/echo hello

Conversational Workflow Example

User: I need an sbatch script for my MPI simulation. I want 2 nodes, 64 ranks per node, 2 OpenMP threads per rank, and 2 hours.

Agent workflow:

1. Confirm partition/account and whether GPUs are needed.
2. Generate a hybrid job script:

``bash python3 scripts/slurm_script_generator.py --job-name run --time 02:00:00 --nodes 2 --ntasks-per-node 64 --cpus-per-task 2 -- ./simulate ``

1. Explain the mapping:

• Total ranks = 128
• Threads per rank = 2 (OMP_NUM_THREADS=2)

1. If the user provides node core counts, sanity-check oversubscription using --cores-per-node.

Error Handling

Verification checklist

• [ ] Confirmed the generated results.script places every #SBATCH directive immediately after the shebang and before set -euo pipefail (open the script and check the first real command line) — a directive after the first command is silently ignored by SLURM.
• [ ] Inspected results.warnings and confirmed it is empty, or recorded each warning (CPU oversubscription, non-integer task-to-GPU ratio, double-launcher) with a deliberate justification for ignoring it.
• [ ] Recorded results.derived.cpus_total_requested (= ntasks * cpus-per-task) and, when --cores-per-node was supplied, confirmed cpus_per_node_requested <= cores_per_node so the node is not oversubscribed.
• [ ] For hybrid runs, verified the script's export OMP_NUM_THREADS value equals results.derived cpus_per_task (the generator sets them equal — confirm that matches the intended threads-per-rank).
• [ ] For GPU jobs, recorded results.derived.total_gpus and ranks_per_gpu and confirmed ranks_per_gpu is the intended integer (or documented intentional sharing such as MPS).
• [ ] Verified results.run_line is not double-wrapped: if the run command already starts with srun/mpirun/mpiexec/orterun/aprun/jsrun, confirmed --launcher none was used (or the auto-detect warning fired) so SLURM does not launch N independent copies.
• [ ] Cross-checked the partition, account, QoS, memory style (--mem vs --mem-per-cpu), and GPU directive against the actual cluster's documented policy — the generator only validates internal consistency, never site policy.

Common pitfalls & rationalizations

Security

Input Validation

• --time is validated against strict HH:MM:SS or D-HH:MM:SS format via regex (minutes/seconds in [00,59])
• --nodes, --ntasks, --ntasks-per-node, --cpus-per-task, --gpus-per-node, --cores-per-node are validated as positive integers with generous upper bounds (e.g. nodes ≤ 100000, ntasks ≤ 10000000, cpus-per-task ≤ 4096, gpus-per-node ≤ 64); the derived total ntasks = nodes * ntasks-per-node is also bounds-checked
• --mem and --mem-per-cpu are validated against SLURM's accepted format (^[0-9]+([KMGT])?$); providing both simultaneously is rejected
• --job-name is validated against ^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$ (no spaces or shell metacharacters)
• --partition, --account, --qos, --constraint, --reservation, and --gpu-type are validated against a safe-character allowlist ^[A-Za-z0-9][A-Za-z0-9._:,+-]{0,127}$ before being emitted into #SBATCH directives
• --module values are validated against ^[A-Za-z0-9][A-Za-z0-9._/+-]{0,127}$ to prevent shell injection (no ;, |, &, backticks, $, or whitespace)
• --env keys must be valid shell identifiers (^[A-Za-z_][A-Za-z0-9_]*$); values are shell-quoted in the generated export lines
• --srun-extra is tokenized with shlex.split and each token is re-quoted with shlex.quote, so it cannot inject shell syntax (;, |, &, $, ...) into the run line
• Invalid input causes the CLI to print a message to stderr and exit with code 2

File Access

• The script reads no external files; all inputs are provided via CLI arguments
• --out writes the generated sbatch script to a single specified file path
• The generated script is a plain-text shell script with #SBATCH directives; it contains no dynamically generated code