slurm-workflows: HPC workflow helpers for Slurm clusters

slurm-workflows runs Python functions on a Slurm cluster from a Python script. 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. Every later task runs on a worker that started earlier.

Use it in three cases:

• You have one Python function to run over many inputs on one cluster allocation.
• You have many Python tasks to run, some of which wait on others.
• Per-worker state is expensive to build, and you want to keep it between tasks.

For a search over a parameter space, use slurm-workflows-optimize, which builds on this package.

Installation

A run needs Slurm's sbatch, squeue and scancel, and the ds-service binary, on PATH.

pip install -U slurm-workflows

To install 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, the Python script that uses the executor, 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 miniforge/26.3.2
conda activate slurm-workflows
"""

# The network interface the workers can reach.
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. Map `square` over the inputs on the "cpu" job group's workers.
        #    The call blocks until every value is back, in input order.
        squares = executor.map("cpu", square, range(100))

# The executor canceled every pilot job at the end of the block.
print(sum(squares))

work directory: '/home/<user>/.cache/slurm-workflows/my-run/<timestamp>'
328350

map blocks, and hands its inputs to the workers as they start. For one value instead of a list, map_reduce folds the values on the workers. For a handle per task, submit enqueues one task and returns at once, even before the workers exist. To choose between them, read The pilot-job model.

Documentation

For contributors

License

MIT. See LICENSE.