slurm-shim: How srun Launches Tasks on Open Cluster Scheduler

A lot of AI tooling only speaks SLURM. submitit calls sbatch. jax.distributed.initialize() finds its peers through SLURM_* variables. Most multi-node PyTorch examples start with srun torchrun. On a cluster running Open Cluster Scheduler (fka Sun Grid Engine), these tools need glue code, or they do not run at all.

slurm-shim removes that glue. It is a single Go binary that provides sbatch, srun, squeue, scancel, sacct, sinfo and scontrol on top of OCS. Job scripts stay as they are. OCS never cared much about what a job actually is, and a script written in SLURM's dialect is just one more thing for it to run.

This post gives a short overview of the scope. Then it goes into one part in detail: how srun starts the tasks of a step across hosts, and how every rank ends up with the right environment.

Scope

A cluster manager installs it once:

slurm-shim install            # prints the plan, changes nothing
slurm-shim install --apply    # creates the slurm-shim PE, wires starter_method
slurm-shim doctor             # checks the setup, keep the output for support

install --apply creates a dedicated slurm-shim parallel environment and sets the starter_method of the queues. It refuses to overwrite an existing start_proc_args or starter_method, because sites use those for their MPI integration.

After that, users submit ordinary SLURM scripts:

#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8

head=$(scontrol show hostnames | head -n1)
srun --ntasks-per-node=1 torchrun --nnodes=4 --nproc-per-node=8 \
  --rdzv-id=$SLURM_JOB_ID --rdzv-backend=c10d --rdzv-endpoint=$head:29500 train.py

What happens with it:

• sbatch turns the directives into a qsub. The partition becomes a queue, a PE and a slot count. --nodes and --ntasks-per-node pin the layout with qsub -par (see Placing Parallel Jobs with qsub -par). --time becomes h_rt, and --gpus-per-node becomes a request on the RSMAP complex. --array, --mem, --dependency and --signal are translated too, with a warning where OCS can only approximate SLURM.
• When the job starts, the PE hook fabricates the SLURM_* environment and the queue's starter_method sources it before the first line of the script.
• srun inside the job starts the tasks on the granted hosts through qrsh -inherit. That is the rest of this post.
• squeue, scancel, sacct, sinfo and scontrol are backed by qstat, qdel, qacct and qmod.

MPI is out of scope on purpose. OpenMPI, Intel MPI and MVAPICH already have a native tight integration in OCS, and that path is better than anything a shim can offer. srun --mpi=pmix fails by design. The same goes for tools that already talk to OCS directly, like Dask, Nextflow or Snakemake.

The shim is validated against live OCS 9.0.10 and 9.1.5 clusters. You need 9.1.5 or newer for full fidelity. GPU jobs are validated on multi-node NVIDIA L4 clusters: every rank opens exactly the devices OCS granted, and NCCL all-reduce runs across hosts. The framework recipes are validated as well (PyTorch DDP and FSDP with torchrun, Hugging Face Accelerate, DeepSpeed, Ray, vLLM, JAX, Flax, submitit, Hydra).

One srun, one stepper per host

srun runs on the master host of the PE job. It reads the allocation from layout.json, decides which ranks go where, and starts exactly one stepper process per host. The stepper on the master host is a plain fork/exec child. The steppers on all other hosts are started through qrsh -inherit, so sge_execd on that host owns the process. That is what makes accounting, qdel and the wallclock limit work for the remote ranks.

All steppers then connect back to a TCP listener that srun opened before starting anything. This return path is the whole control plane. The environment, the rank placement, the output, the signals and the exit codes all travel over it.

Two spawn paths, one control plane. The master host's stepper is an ordinary child of srun; remote steppers are handed to sge_execd through qrsh -inherit so Open Cluster Scheduler owns them. Both kinds connect back to the same token-authenticated listener.

srun --pty outside of a job takes a different path. It turns into an interactive qrsh session, and there is no stepper at all.

Before the first process starts

supervisor.launch() in internal/cli/srun/run.go runs these steps in this order:

1. Token. 32 random bytes, hex encoded. There is one token per step, and it is the only credential the control channel accepts.
2. Listener. A step that stays on the master host listens on 127.0.0.1 with an ephemeral port. A step that spans hosts listens on 0.0.0.0 inside a fixed port range, 61000 to 61439 by default, so a site can open it with one firewall rule. Leaving loopback is fine because the channel authenticates.
3. Preflight. For a multi-host step the PE must have control_slaves TRUE, checked with qconf -sp. A misconfigured PE fails here, not after half the hosts are running.
4. Spawn. One stepper per host, with a routing envelope in argv and the token in the environment. If any host fails, the steppers already started are killed and srun exits with 8.
5. Accept. srun waits until every host has authenticated, bounded by launch_timeout (60 seconds by default). A stepper that never connects fails the step instead of hanging it.

The listener exists before the first child, so no stepper can dial a port that is not bound yet.

A remote stepper cannot be told 0.0.0.0, it would dial its own loopback. srun keeps the bound port and hands out the routable address of the master host instead. On a filtered network, slurm-shim ports prints the rules the site needs, for firewalld, GCP VPC and nftables.

Three channels, deliberately split

The most consequential design decision in the launcher is what travels where. A process's argv is world-readable through /proc/<pid>/cmdline on a shared execution host, so argv carries routing and nothing else.

The environment never rides the command line, and never rides qrsh -V. buildQrshArgs deliberately omits -V; the job environment reaches each host only after that host has proved it holds the step token.

The resulting argv is short enough to read in full:

# internal/launch/qrsh.go, buildQrshArgs
qrsh -inherit -nostdin -noshell \
     -v SLURM_SHIM_TOKEN=<token> \
     <host> \
     $SGE_ROOT/slurm-shim/bin/slurm-shim stepper --envelope <base64>

-inherit is the tight integration flag. It attaches to the existing allocation instead of requesting a new one. -noshell execs the binary directly, without a login shell in between. QRSH_WRAPPER and SGE_RSH_COMMAND are removed from the environment of the qrsh child first, so neither can redirect the transport.

On the remote host this command also passes through the queue's starter_method. The starter has to recognize a stepper launch and exec it untouched. That match is a shell pattern Go cannot see, so a contract test feeds the real starter script the exact argv buildQrshArgs produces. If either side changes, the build breaks.

Why the token is not in argv. If it were, any user on the execution host could read it from /proc, dial the listener, claim a host name and receive that host's StepSpec. Under the default --export=ALL that is the submitter's entire environment. qrsh -v moves this exposure rather than removing it: OCS writes the token into the task's environment file in the execd spool directory. slurm-shim doctor checks whether other users can traverse the path to that directory.

The handshake

Once a stepper runs, everything between it and srun is length-prefixed frames on one socket. There is no RPC library and no request/response pairing, just an asynchronous stream in both directions. And it is the stepper that connects to srun, not the other way round. The launcher only has to get a process running, and the socket does the rest.

stepper -> srun     HELLO       token, host name
srun    -> stepper  SPEC        StepSpec as JSON
stepper -> srun     READY       output files open, ranks about to start
stepper -> srun     OUT         one line or 64 KiB chunk of rank output
srun    -> stepper  SIG         signal for every rank's process group
stepper -> srun     RANK_EXIT   exit code of one rank (RANK_FAIL if it never started)

HELLO has a 10 second deadline, and the token is compared in constant time. Every frame has a fixed ten byte header (type, flags, rank, payload length) and at most 1 MiB of payload. A frame is written in one go under a mutex, so output from different ranks never interleaves.

How a rank gets its environment

This is where the SLURM behavior comes from. JAX, torchrun and Accelerate read SLURM_PROCID, SLURM_LOCALID, SLURM_NTASKS or SLURM_JOB_NODELIST and trust what they find. Every rank has to see values that match the step, the host and its own position.

The rank environment is built in four layers. A later layer wins per key:

1. job environment   the job's environment, filtered by --export
                     (ALL by default, NONE, or a list of variables)
2. step shadows      SLURM_NTASKS, SLURM_TASKS_PER_NODE, SLURM_CPUS_PER_TASK,
                     SLURM_STEP_ID, SLURM_STEP_NODELIST, SLURM_STEP_NUM_NODES, ...
3. rank delta        SLURM_PROCID, SLURM_LOCALID, SLURM_NODEID,
                     SLURM_GTIDS, SLURM_CPUS_ON_NODE
4. host local        SLURMD_NODENAME, plus CUDA_VISIBLE_DEVICES or
                     ROCR_VISIBLE_DEVICES when the rank was granted GPUs

Layer 1 is what the starter_method fabricated when the job started: SLURM_JOB_ID, the compressed SLURM_JOB_NODELIST, SLURM_NNODES, SLURM_GPUS_ON_NODE and the rest. It describes the whole allocation. A step can be smaller than the allocation, so srun shadows the job level values with the geometry of this step in layer 2. Layers 1 and 2 go into StepSpec.Env and are the same for every rank on every host. Layer 3 is computed per rank and travels as EnvDelta in the rank list of each host. All three are decided by srun on the master host. The stepper only adds layer 4.

GPUs need a bit more care. Before layering, srun removes any device visibility variable the job inherited from a module file, a prolog or a container image. Otherwise an old value could stack on top of the mask the shim writes. It removes them rather than setting them empty, because an empty CUDA_VISIBLE_DEVICES means zero devices. Then exactly one variable is written per rank, with the devices from the job's RSMAP grant. With --gpus-per-task each rank gets its own slice of the host's grant.

Which variable that is gets resolved once on the master host from gpu.vendor and sent along as StepSpec.GPUEnvVar. The stepper never loads the config, so a config difference between nodes cannot change it, and it refuses any name other than the two it knows. Under gpu.isolation: cgroup the shim writes no mask at all and leaves device selection to the OCS cgroup setup.

srun --test-only -n 4 hostname shows the result without starting anything. The dry run report calls the same function the stepper uses for layers 3 and 4, so what it prints is what a rank really gets.

Finally, the stepper starts each rank through a small trampoline, slurm-shim rank-exec [--cpuset ...] [--chdir ...] -- <command>. It applies CPU affinity and the working directory, sets SLURM_TASK_PID to its own pid (which execve keeps) and then execs the user command. Each rank gets its own process group, so one kill reaches the whole tree. A status pipe on fd 3 tells a failed start apart from a real exit code of the user command.

When something fails

• OCS refuses the qrsh. The launcher watches each new qrsh for two seconds and classifies its stderr. If the remote execd does not know the job yet, it retries for 10 seconds. If slots are temporarily exhausted, it retries for up to 5 minutes and prints a SLURM style retry line. A missing JOB_ID fails the step at once.
• An output file cannot be opened. The stepper opens all -o/-e files before it sends READY. If one fails, the whole host fails before any rank starts.
• srun dies. The stepper treats a broken channel as the death of srun. It sends SIGTERM to all ranks, waits five seconds, then sends SIGKILL. No rank outlives its supervisor.
• A stepper dies. srun sees EOF and records a failure for every rank of that host that has not reported yet. The step does not hang.
• A rank fails under -K. The first non-zero exit code is latched and all ranks get SIGTERM. The step reports that first code, not the maximum.

In the other direction, srun forwards SIGINT, SIGTERM, SIGHUP, SIGUSR1, SIGUSR2 and SIGQUIT as SIG frames. A second SIGINT within one second becomes SIGKILL.

If you want to read the code, start with supervisor.launch(). Read it top to bottom and the other packages come up in the order you need them: internal/plan for placement, internal/launch for qrsh, internal/proto for the frames, internal/stepper for the ranks of one host, and internal/mux for putting the output back together.

Try it

You do not need a cluster. With Docker and a Go toolchain:

make cluster-up      # 3 node OCS 9.1.5 cluster in containers, shim installed
make demo            # srun fanning ranks across the nodes
make cluster-down

The code is at github.com/hpc-gridware/slurm-shim. Reports from real clusters, working or not, are very welcome as issues.