Embarrassingly parallel and DAG workflows for science in Python

Python already ships with everything needed to run things in parallel on a single machine. Because of the GIL², you want processes, not threads, for anything CPU bound in pure Python (threads are sometimes fine if the heavy lifting happens in numpy/torch, which release the GIL).

concurrent.futures is the modern and nicer API, multiprocessing.Pool does essentially the same thing:

from concurrent.futures import ProcessPoolExecutor
from itertools import product

def run_one(subject, method):
    kspace, reference = load(subject)
    return evaluate(reconstruct(kspace, method), reference)

if __name__ == "__main__":
    tasks = list(product(subjects, methods))
    with ProcessPoolExecutor(max_workers=8) as executor:
        results = executor.map(run_one, *zip(*tasks))
    table = aggregate(dict(zip(tasks, results)))

It works, it has no dependencies, but:

• Everything sent to a worker must be pickled: no lambdas, no functions defined in the __main__ of a notebook, and large arrays are copied back and forth.
• The if __name__ = "main":= guard is mandatory, as workers will re-import your script (this is the default on macOS and Windows, and on Linux since Python 3.14).
• There is no notion of dependency: if you want a DAG, you have to call .result() on futures yourself, in the right order.

Joblib is the "just make my for loop parallel" library, and it does it very well. Also it is maintained by colleagues of mines, notably Thomas Moreau.

from joblib import Parallel, delayed

results = Parallel(n_jobs=-1)(
    delayed(run_one)(subject, method) for subject, method in tasks
)

Compared to the stdlib, you get for free:

• A more robust pickling (via cloudpickle, so it works in notebooks) and memory-mapping of large numpy arrays instead of copying them.
• Sane limits on the number of BLAS/OpenMP threads in each worker, so that 64 workers don't each spawn 64 threads.
• joblib.Memory, a disk cache that skips a function call if it has already been computed with the same arguments. For long experiments it's a life-saver.

It has no DAG support, and stays on one machine (unless you plug the Dask or Ray backend, see below).

Not a workflow engine per se, but worth mentioning because a lot of people use it without knowing: every n_jobs parameter in scikit-learn (GridSearchCV, cross_validate, random forests…) is joblib under the hood. If your "experiment" is a hyper-parameter search on an estimator, you may not need anything else. And with joblib.parallel_config(backend"dask")= the same code can run on a cluster.

Submitit is a small library to send Python functions as SLURM jobs, and get the results back as Python objects. No bash scripts, no sbatch, no parsing of output files.

import submitit

executor = submitit.AutoExecutor(folder="logs/%j")
executor.update_parameters(
    timeout_min=60, cpus_per_task=4, gpus_per_node=1,
    slurm_partition="gpu", slurm_array_parallelism=32,
)
jobs = executor.map_array(run_one, *zip(*tasks))  # a single SLURM job array
results = [job.result() for job in jobs]

What I like about it:

• map_array submits a SLURM job array, which is what your cluster admins would like you to do (instead of 1000 separate jobs).
• Logs, pickled inputs and outputs of each job end up in the folder, so it's easy to see what went wrong.
• AutoExecutor(cluster"debug")= runs everything in the current process, so you can put a breakpoint in your function before sending it to the cluster.

It has no DAG support either: each job is independent. You can chain stages manually (wait for the first array, then submit the second), which is often good enough.

Dask is the first one of this list that really thinks in graphs. With dask.delayed, you write your code as usual, but each call returns a lazy object instead of a result, and dask builds the DAG for you:

import dask

scores = {}
for subject in subjects:
    kspace, reference = dask.delayed(load, nout=2)(subject)  # loaded once per subject
    for method in methods:
        image = dask.delayed(reconstruct)(kspace, method)
        scores[subject, method] = dask.delayed(evaluate)(image, reference)
table = dask.delayed(aggregate)(scores)

table.visualize("dag.png")  # optional, requires graphviz
result = table.compute()

This is exactly the DAG drawn above, including the shared load. By default it runs on local threads/processes; add a dask.distributed.Client and you get a (very useful) live dashboard. With dask-jobqueue, the workers are SLURM jobs:

from dask.distributed import Client
from dask_jobqueue import SLURMCluster

cluster = SLURMCluster(cores=8, memory="32GB", walltime="02:00:00")
cluster.scale(jobs=10)
client = Client(cluster)  # everything computed from now on runs on the cluster

The downsides: debugging a failure inside a distributed worker is less pleasant than a plain traceback, and the scheduler has some overhead, so tasks should last at least a second or so. Dask is also a lot more than delayed (arrays, dataframes), which you may or may not need.

Ray comes from the ML world and is the big one of this list. The API is similar to dask's delayed: decorate a function with @ray.remote, call it with .remote(), and you get a reference to a future result, that you can pass to other remote functions to build the DAG:

import ray

ray.init()

load_r = ray.remote(num_returns=2)(load)
reconstruct_r = ray.remote(num_gpus=1)(reconstruct)
evaluate_r = ray.remote(evaluate)

refs = {}
for subject in subjects:
    kspace, reference = load_r.remote(subject)
    for method in methods:
        image = reconstruct_r.remote(kspace, method)
        refs[subject, method] = evaluate_r.remote(image, reference)
table = aggregate(dict(zip(refs, ray.get(list(refs.values())))))

Its strong points are the resource-aware scheduling (num_gpus=1 and Ray makes sure only one task per GPU runs at a time), a shared-memory object store that avoids copying large arrays, and a whole ecosystem on top (Ray Tune for hyper-parameter search, Ray Train…). The downside is that it is heavy, and running it on a SLURM cluster means starting a Ray head node and workers yourself inside an allocation, which is not really a one-liner.

If you work in neuroimaging, you have probably met Nipype, which wraps FSL, SPM, ANTs, FreeSurfer and co. into Python and chains them in workflows. Pydra is its successor, rewritten as a general purpose dataflow engine. I won't show code here, as its API changed quite a lot recently, but the key ideas are worth knowing:

• Each task output is cached on disk, indexed by a hash of its inputs: re-running a workflow only recomputes what changed.
• Splitters and combiners express the "for each subject, for each method" loops declaratively, and the combine step is our aggregate.
• The same workflow can run with a local, SLURM, SGE or Dask worker.

It's the most complete of the list for our use case, but also the steepest learning curve. If your pipeline is mostly calling existing neuroimaging tools, go for it. For pure Python number-crunching, dask is probably enough.