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)
• The ds-service binary on PATH

pip install -U slurm-workflows

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

Usage

Replace the Slurm account (-A), the partition (-p) and the setup script with the ones for your cluster. The driver must run on a node with the ib0 interface. For another interface, read How to run the ds-service server.

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

work directory: '/home/<user>/.cache/slurm-workflows/my-run/<timestamp>'
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.