slurm-workflows: HPC workflow helpers for Slurm clusters

slurm-workflows lets you run Python functions on a Slurm cluster without sbatch scripts written by hand. It provides an interface inspired by concurrent.futures. The interface launches long-lived pilot workers. It then dispatches tasks to those workers. You pay Slurm's queue latency once per worker, not once per task.

Use it in three cases:

• You have many Python tasks to run on one cluster allocation.
• A sweep or a calibration has to spread across a pool of nodes.
• Per-worker state is expensive, and you want it to stay warm between tasks.

Features

• Pilot workers - pay Slurm's queue latency once per worker, then dispatch tasks to them.
• Dynamic scaling - grow or shrink a pool of workers at runtime.
• Stateful actors - keep expensive per-worker state (loaded models, database connections) warm across many tasks.
• Transparent serialization - cloudpickle serializes functions, arguments, and return values.
• Live monitoring tool - swtop, a terminal UI that shows the tasks, workers, nodes and jobs for a running workflow.
• Bayesian optimization - an optimizer built on botorch for optimization and calibration workflows.

Requirements

• Python >= 3.12
• Access to a Slurm cluster (sbatch, squeue, scancel on PATH)
• A running ds-service server, v5.1.0 or later

Installation

pip install -U slurm-workflows

To set up on UVA's Rivanna cluster, read How to install slurm-workflows on Rivanna.

Usage

from ds_service_client import DsServiceServer
from slurm_workflows import SlurmPilotExecutor


def square(x):
    return x * x


SETUP_SCRIPT = """
module load gcc/14.2.0
conda activate my-env
"""

with DsServiceServer(interface="ib0") as ds_service:
    ds_service.wait_until_ready()

    with SlurmPilotExecutor("my-run", ds_service.address) as executor:
        # 1. Describe a kind of worker. This submits nothing.
        executor.define_worker(
            name="cpu",
            sbatch_args=["-A my_alloc", "-p standard", "-t 01:00:00"],
            setup_script=SETUP_SCRIPT,
        )

        # 2. Launch 4 pilot jobs of that kind.
        executor.scale_workers("cpu", 4)

        # 3. Submit tasks to a named queue. Workers of that group pull from it.
        tasks = [executor.submit("cpu", square, i) for i in range(100)]

        # 4. Block until every result is in.
        executor.wait(tasks, desc="squaring")

# The executor canceled every pilot job at the end of the block.
print(sum(task.output for task in tasks))

The executor passes sbatch_args straight through to sbatch. As a result, any Slurm option works. You can submit tasks before the workers exist. The tasks wait on the queue until a pilot job starts and takes them.

Documentation

Tutorials

How-to guides

Reference

Explanation

For contributors

License

MIT - see LICENSE.