Repository navigation
MPI: tasks with acceleration=DistributedAcceleration() go missing when ranks spawn them independently #778
Description
We'd like to run some tasks only on the local rank inside an MPI program (a parallel loop over the rank's own threads), without every rank planning them. acceleration=Dagger.DistributedAcceleration() looked like the way to do that, and it works on one rank. With two ranks, some of the tasks never run:
# mpiexec -n 2 julia --project --threads=4 min.jl
using Dagger, MPI
Dagger.accelerate!(:mpi)
rank = MPI.Comm_rank(MPI.COMM_WORLD)
busy(i) = (sleep(0.1); i)
ts = [Dagger.@spawn(acceleration = Dagger.DistributedAcceleration(), busy(i)) for i in 1:12]
done = Threads.Atomic{Int}(0)
foreach(t -> Threads.@spawn((fetch(t); Threads.atomic_add!(done, 1))), ts)
t0 = time(); while done[] < 12 && time() - t0 < 20; sleep(0.1); end
println("rank $rank: $(done[])/12 tasks finished")With -n 1 all 12 finish. With -n 2 only 8–10 per rank do, and we see Hit probable hang on bcast_meta delivery and sometimes AssertionError: MPIRef rank mismatch. Dagger 0.22.7, MPI.jl 0.20.28, Julia 1.13; same on macOS (MPICH_jll) and on a Linux cluster (Open MPI).
We think there are two things going on:
- Some of these tasks get placed on
MPIProcessors, sometimes on the other rank, and then follow the SPMD execution path.accel_matches_proc(::DistributedAcceleration, proc)is alwaystrue, andcompatible_processors(::DistributedAcceleration, ...)skips the filter completely. The missing tasks are exactly the ones placed on MPI processors. - Local tasks take IDs from the same counter as SPMD tasks, and the MPI tags are derived from those IDs. If the ranks spawn different numbers of local tasks, later SPMD tasks get different IDs on each rank and hang (
Hit hang on bcast_meta deliverywith different tags on the two ranks).
I'll open a PR with a small fix for both: MPI processors are excluded for DistributedAcceleration, and local tasks under MPI get IDs from a separate range. With it, the reproducer runs fine (12/12 per rank, about 4× faster than spawning the tasks SPMD-style), local tasks with different counts per rank followed by normal SPMD work give correct results, and test/mpi/cpu.jl passes on 2 ranks × 2 threads.
We're not sure whether this is the intended way to use acceleration= under MPI, though. With the patch, a local task can use an SPMD result owned by its own rank; using one owned by another rank fails with an error, which seems reasonable. We don't know what the intended way is to pass a local task's result back into SPMD work.
Thanks a lot!
The analysis, the fix and this text were written with help from Claude (Anthropic); I checked the results.
Activity
jpsamaroo commentedon Oct 15, 2026
This is an edge case that we hadn't really considered when building the MPI and acceleration design - the (eventually) built design is that
Dagger.accelerate!is effectively a global switch (despite using a task-local variable to carry acceleration), and everything launched after it (and before anotherDagger.accelerate!call) must complete within that acceleration region, and must use that acceleration. So passingaccelerationas an option isn't really intended to be supported (we do frequently "smuggle" values through options like this, but we don't document that they're really internal, so it's not wrong for you to have expected it to work!).Maybe the design is bad and should be changed to support this - let's see what we'd have to be careful about when fixing it to make it work:
Everything in a comm must be convergent
Task IDs, input argument identities, types, control flow, etc. - all of this stuff needs to have the same value/identity on all participating ranks, when doing things on the same MPI comm. We rely on this in many, many places (see
check_uniformuses internally). However, if we use multiple comms, this theoretically isn't a problem - we can do different things on different comms, as long as things done within a comm stay convergent. But I'm not sure if this is a good path to fixing our problem, so moving on...Acceleration was originally intended to be task-local, but became process-global
This is a bigger problem -
Dagger.accelerate!changes the globally-available processors, and fiddles with MPI process state, so it can't be called concurrently on multiple tasks (which would have been the way I'd suggest to approach this). Specifically, this won't work correctly:Fixing this would thus require:
Contexts must become task-local
So that the processors available are limited by setting the current acceleration. They'll thus default to DistributedAcceleration processors, but can be switched to MPIAcceleration if desired.
MPI initialization/finalization must follow a global refcount
This is incremented with
Dagger.accelerate!(:mpi), and is decremented either by aDagger.accelerate!(:distributed)on the same task, or when the task is GC'd. Once it hits 0, we also need to wait for any tasks running under MPIAcceleration to complete, so that we don't finalize while they still have MPI calls to do.Using MPI task results as Distributed task arguments is dicey
Even when the MPI task is finished, its result may live on another rank, requiring MPIAcceleration-directed communication, which might or might not be doable. This needs testing. The other direction will probably work fine, but again, needs testing.
We still need to deal with task IDs
These are global and monotonically increase, and must be convergent, which becomes a big problem when suddenly the task ID counters become out of sync across ranks - I'm not sure how to tackle this one, so this will need some thinking.
Maybe other things too!
We'll need lots of tests to validate that this works under various concurrency and acceleration-switching scenarios. I'll throw this over to an agent and see what it thinks!