Migrating to sbatch from srun #

So far we have shown how to run jobs with srun, which lets you see results directly in your console. However, once you are comfortable using Pegasus through srun, we recommend moving to sbatch. It forgoes interactivity to gain several key advantages:

Jobs submitted through sbatch run independently of your console. No need for screen or tmux sessions to keep them alive. They even survive a login node crash.
Console output is automatically written to a file, by default slurm-$SLURM_JOB_ID.out in the current directory.
Define job arrays to run a series of experiments. Ideal for hyperparameter sweeps and the like. No need to manually submit and manage many jobs, or writing custom scripts to do so. Plus, arrays are much nicer to the scheduler than many individual jobs.
Define default values for command line arguments directly in your sbatch script.

Hello sbatch #

With sbatch you submit jobs as shell scripts that are executed on compute nodes. Below is a simple hello world example. Note that Slurm will not accept files without the shebang line (#!/bin/bash).

#!/bin/bash
echo "hello world!"

Save this to hello_sbatch and submit your first batch job with sbatch --time 1 --mem 16M hello_sbatch. All console output is saved to slurm-$SLURM_JOB_ID.out in the current directory.

More complex scripts #

sbatch scripts are fully fledged shell scripts that Slurm executes for you. The shebang at the top determine which shell is used.

When --container-image is defined, the script is executed inside the container, using a shell provided by the container. You can add alternative shells like zsh or fish to the image and use them with sbatch.

The example below starts a long-running command in the background, runs some other commands with loops and branching, and finally waits for all background commands to finish.

#!/bin/bash
# let's set the following defaults (can be overriden from the commandline):
#SBATCH --job-name hello_sbatch_complex
#SBATCH --cpus-per-task 2
#SBATCH --mem 256M
#SBATCH --time 1
#SBATCH --container-image=/enroot/debian+13.5.sqsh

# a long-running command that we send to the background
SLEEPY_TIME=3
(sleep $SLEEPY_TIME && echo "long-running command finished!") &

echo "let's count down!"
for i in {5..1..-1}; do
   echo $i
done

# randomly hello world in spanish or english
if (($RANDOM > 16383)); then
   echo "¡Hola mundo!"
else
   echo "Hello world!"
fi

# wait for the long-running command to finish
wait

Note that we define defaults for command line parameters using #SBATCH directives at the top of the file. You can override them from the command line.

All #SBATCH directives need to be in the preamble at the top of the file. Slurm stops parsing your script once it encounters the first command, so any directives after that point will be ignored.

You might find the available envrionment variables useful to integrate into your script.

Limitations of #SBATCH directives #

While convenient, #SBATCH directives are merely parsed by Slurm instead of interpreted by the shell. For example #SBATCH --container-workdir="`pwd`" sets the workdir to the literal string `pwd` instead of the current directory as you would expect.

The simplest solution is to pass these arguments to sbatch on the command line like you would with srun, i.e., sbatch --container-workdir="`pwd`" ....

--container-workdir="`pwd`" in particular can be replaced by adding cd "$SLURM_SUBMIT_DIR" to your script to change to the directory.

For the input, output, and error option, sbatch provides special filename patterns that let you specify the names of these files based on, e.g., id, name, hostname, and array index.

Job arrays #

As mentioned earlier, one very handy feature that sbatch offers over srun is job arrays. These are especially useful if you want to run a series of similar jobs. The example below runs 5 jobs in total (index 0-4 inclusive), where at most 3 run in parallel (%3).

#!/bin/bash
#SBATCH --array 0-4%3
#SBATCH --job-name sbatch_array_test

echo "hello world from job array index $SLURM_ARRAY_TASK_ID"

As you can see, the job can access its array index via the $SLURM_ARRAY_TASK_ID environment variable. All output is by default saved to slurm-$SLURM_ARRAY_JOB_ID-$SLURM_ARRAY_TASK_ID.out in the current directory. See the array option, as well as available envrionment variables for more details.

Different parameters per array job #

In the example above, we used $SLURM_ARRAY_TASK_ID to make each job in the array print something different. Similarly, it is your responsibility to assign different parameters to each job in the array. There are many ways to achieve this. You can take inspiration from the following patterns:

Create one file per job, either including task index in the name

#!/bin/bash
#SBATCH --array=0-42
python train.py "parameters-$SLURM_ARRAY_TASK_ID.json"

or globbing into a bash array

#!/bin/bash
#SBATCH --array=0-42
FILES=(
    inputs/*.json
)
python train.py FILES[$SLURM_ARRAY_TASK_ID]

Create a single JSON file with a list of parameters, indexed by task index.
```
#!/bin/bash
#SBATCH --array=0-42
python train.py --parameters parameters.json --index $SLURM_ARRAY_TASK_ID"
```
Then in Python, json.load the parameters file and retrieve the index.

Define sets of parameters in your script.

#!/bin/bash
#SBATCH --array=0-2

# Define sets of parameter for each task
params_0=(--learning-rate 0.1 --epochs 30000 --mode "train really good")
params_1=(--learning-rate 1.5 --verbose)
params_2=(--model vit --force --discombobulation-factor 7.3)

# Retrieve parameters by name
param_names=(params_0 params_1 params_2)
declare -n params=${param_names[$SLURM_ARRAY_TASK_ID]}

python train.py "${params[@]}"

Generate parameters in code, e.g., with a ParameterGrid from scikit-learn. Pass the index to Python as before.

#!/bin/bash
#SBATCH --array=0-42
python train.py --index $SLURM_ARRAY_TASK_ID"

In Python, define the grid and obtain the parameters for the given index.

from sklearn.model_selection import ParameterGrid
...
grid = ParameterGrid({
    'learning_rate': [0.1, 0.2, 10],
    'epochs': [90, 120, 180],
    'optimizer': ['Adam', 'SGD']
})
train(grid[index])

srun in sbatch #

For some time Pyxis supports --container-image and other arguments related to Enroot with sbatch. As mentioned earlier, defining a container image executes the entire script inside a container. If you need to run commands on the compute node directly, or need to use different container images for different steps, you can also use srun inside sbatch.

#!/bin/bash
#SBATCH --job-name=srun-in-sbatch

# run commands directly on the compute node
enroot import -o my-image.sqsh -- docker://my-namespace/my-image

# use multiple srun commands to run different container images
srun \
  --container-image=my-image.sqsh \
  python generate.py
srun \
  --container-image=/enroot/nvcr.io_nvidia_pytorch_26.04-py3.sqsh \
  --container-workdir="`pwd`" \
  --container-mounts=/netscratch:/netscratch,/ds:/ds:ro,"`pwd`":"`pwd`" \
  python train.py

sbatch without a script file #

If you really detest writing separate script files for your experiments, you can also use the --wrap option to create a command line only version. One advantage is that the shell interprets your string before handing it off to sbatch, so you can use variables and subshells again.

sbatch \
  --container-image=/enroot/nvcr.io_nvidia_pytorch_26.04-py3.sqsh \
  --container-workdir="`pwd`" \
  --container-mounts=/netscratch:/netscratch,/ds:/ds:ro,"`pwd`":"`pwd`" \
  --wrap \
    "echo \"hello world! array index: \$SLURM_ARRAY_TASK_ID\""

Be careful around variable expansion. Note how $SLURM_ARRAY_TASK_ID is escaped to prevents its expansion at submit time, where it is undefined.

sbatch further accepts scripts on standard input. Like with --wrap, the shell interprets the given string before passing it on to sbatch, so we regain access to environment variables and subshells:

sbatch <<EOF
#!/bin/sh
#SBATCH --container-image=/enroot/nvcr.io_nvidia_pytorch_26.04-py3.sqsh
#SBATCH --container-workdir="`pwd`"
#SBATCH --container-mounts=/netscratch:/netscratch,/ds:/ds:ro,"`pwd`":"`pwd`"

python train.py
EOF