This is a very simple Python script for launching Slurm array jobs. It's currently tailored to my use case and makes quite a few assumptions, for example:

• The slurm script consists of some header lines followed by a single main command
• The main command takes arguments using the argparse syntax, i.e. --arg1 val1 --arg2 val2
• Currently no support for positional arguments to the main command

That said, it should be easy to modify the script to suit your needs.

Easiest is probably

pip install git+https://github.com/ejnnr/slurm_launcher.git

This will install the script as slurm_launcher in your path. (I personally alias this to slurm but didn't make that the default to make clashes less likely.)

Simple example:

slurm_launcher my_project.train --verbose --lr 0.01 --batch_size 32

This will generate the following Slurm script and submit it to the queue after asking you for confirmation:

#!/bin/bash
#SBATCH --output=slurm/%A_%a.out

srun python -m my_project.train --verbose --lr 0.01 --batch_size 32

The main selling point is that you can automatically do grid searches by specifying multiple comma-separated values for an argument:

slurm_launcher my_project.train --verbose --lr 0.01,0.001 --batch_size 32,64

will generate

#!/bin/bash
#SBATCH --output=slurm/%A_%a.out
#SBATCH --array=0-3

PARAMS_ARRAY=(
    "--verbose --lr 0.01 --batch_size 32"
    "--verbose --lr 0.01 --batch_size 64"
    "--verbose --lr 0.001 --batch_size 32"
    "--verbose --lr 0.001 --batch_size 64"
)
srun python -m my_project.train ${PARAMS_ARRAY[$SLURM_ARRAY_TASK_ID]}

You can set arbitrary sbatch arguments on the fly using --slurm.<arg> <val>:

slurm_launcher my_project.train --epochs 10 --slurm.mem 8gb

Finally, you can also define arguments by putting a file called slurm.yaml in the current directory. See src/slurm_launcher/slurm.yaml for the default configuration and possible options.

In particular, you may want to set a preamble:

launcher:
  preamble: |
    cd /path/to/my/project
    source venv/bin/activate

This will be run before the main srun command.

You can also override how the main command should be launched via launcher.cmd_prefix. The default is "srun python -m ", but you may want e.g. "srun python3 " instead, or "srun python -m my_project.", or something other than Python.