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 workers inside pilot jobs. It then dispatches tasks to those workers. You pay Slurm's scheduling latency once per pilot job, not once per task.

Use it in three cases:

• You have many Python tasks to run on one cluster allocation.
• An exploration or a search has to spread across a pool of workers.
• 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 job group. This submits nothing.
        executor.define_job_group(
            name="cpu",
            sbatch_args=["-A my_alloc", "-p standard", "-t 01:00:00"],
            setup_script=SETUP_SCRIPT,
        )

        # 2. Launch 4 pilot jobs of that job group.
        executor.scale_jobs("cpu", 4)

        # 3. Submit tasks to a named queue. That job group's workers claim them.
        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

You can submit tasks before the workers exist. They wait on the queue until a worker starts and claims them.

Documentation

For contributors

License

MIT. See LICENSE.