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.

Installation

A run needs:

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

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))

328350

Tasks can be submitted before the workers exist. They wait on the queue until a pilot job starts and takes them.

Documentation

For contributors

License

MIT. See LICENSE.