MPI: spread plain @spawn tasks across ranks and threads

AGENTS.md

@@ -659,3 +659,30 @@ lesson.
     Compare unrelated PR logs and check upstream releases before changing the
     test bootstrap: installing Python only for its suite does not cause either
     failure. Both fixes together pass Dagger's Python suite on Julia 1.11.
+
+64. **Uniform eager placement must vary by task, not by ready-arrival order.**
+    MPI ignores rank-local load and capacity, so sorting every task's candidates
+    identically sends all plain `@spawn` work to rank 0's first thread. Rotate
+    the canonical candidates using the already-uniform thunk ID, rather than
+    a scheduler cursor: ready tasks can arrive in different orders across ranks.
+    With input-owner affinity, rotate only that rank's candidate prefix so
+    locality survives; if a scope excludes the owner, rotate the permitted
+    candidates instead. Calls without a task must retain canonical ordering
+    for Datadeps and processor grids. The issue #760 reproduction changed from
+    16 serial tasks on one processor to four tasks on each of four processors.
+
+65. **A dynamic candidate scan can box even an unused backend argument.**
+    Calling `processor_owner_rank(accel, proc)` from a `Vector{Processor}`
+    loop boxed the isbits `MPIAcceleration` once per candidate: four candidates
+    added four allocations / 64 bytes per eager placement. An out-of-line scan
+    with `@nospecialize(accel)` preserves runtime hook dispatch while passing
+    the already-boxed backend through the loop. Only that scan boundary now
+    boxes (one allocation / 16 bytes), and the no-affinity path stays unchanged.
+    This is a pass-through planning argument; do not erase types needed by
+    execution or arithmetic to achieve the same effect.
+
+66. **An MPI test timeout can leave ranks in separate process groups.**
+    Timing out `test/mpi/run.jl` killed its launcher but left both Julia ranks
+    running; MPICH had placed each rank in its own process group. After a
+    timeout, inspect rank PIDs/PGIDs and stop those groups too. Killing only
+    the wrapper's group is not sufficient for this launcher.

src/acceleration.jl

@@ -179,6 +179,18 @@ bind_moved_argument(::Acceleration, original, moved) = moved
 "Logical execution rank of `proc` used for uniform task placement preference."
 processor_owner_rank(accel::Acceleration, proc::Processor) = first(fire_order_key(proc))
 
+# This scan dispatches on abstract Processor elements. Pass the acceleration
+# through unspecialized so an isbits backend is not boxed for every candidate;
+# keep the barrier out of line so inlining does not reintroduce those boxes.
+@noinline function preferred_processor_count(procs, @nospecialize(accel), rank)
+    n = 0
+    for proc in procs
+        processor_owner_rank(accel, proc) == rank || break
+        n += 1
+    end
+    return n
+end
+
 "Deterministic ordering key for `proc` under uniform execution."
 processor_order_key(accel::Acceleration, proc::Processor) =
     (fire_order_key(proc)..., short_name(proc))
@@ -204,6 +216,10 @@ task_processor_preference(accel::Acceleration, task, state) = nothing
 Re-sort `procs` with a deterministic, acceleration-dispatched key when
 `uniform_execution(accel)` so every rank gets the same order (and with a
 `task` the acceleration may prefer the rank owning task argument data).
+Rotate the candidates within that preferred rank (or all candidates when
+there is no eligible preferred rank) using the task ID, so independent tasks
+spread across processors without consulting rank-local load or ready order.
+Calls without a task retain the canonical order used by Datadeps and proc grids.
 When not under uniform execution, leaves `procs` unchanged.
 """
 function select_processors_uniform!(procs, accel::Acceleration,
@@ -214,6 +230,19 @@ function select_processors_uniform!(procs, accel::Acceleration,
         pref !== nothing && processor_owner_rank(accel, p) != pref,
         processor_order_key(accel, p)...,
     ))
+    task === nothing && return procs
+    n = length(procs)
+    if pref !== nothing
+        n = preferred_processor_count(procs, accel, pref)
+        n == 0 && (n = length(procs))
+    end
+    if n > 1
+        offset = mod(task.id - 1, n)
+        # Rotate only the preferred prefix, in place and without scratch.
+        reverse!(procs, 1, offset)
+        reverse!(procs, offset + 1, n)
+        reverse!(procs, 1, n)
+    end
     return procs
 end
 

test/mpi/cpu.jl

@@ -145,6 +145,8 @@ end
     MPI.Barrier(comm)
 end
 
+include("spawn.jl")
+
 @testset "Datadeps" begin
 
 include("threads.jl")

test/mpi/spawn.jl

@@ -0,0 +1,86 @@
+# Plain eager tasks need a uniform placement decision even when ranks make
+# progress at different speeds. Check actual execution, scopes, and locality.
+mpi_spawn_location(x=nothing) = (MPI.Comm_rank(MPI.COMM_WORLD), Threads.threadid(), x)
+mpi_spawn_value() = (sleep(0.05); 42)
+
+function check_spawn_placements(tasks, allowed, value=nothing)
+    placements = Dagger.Processor[]
+    for task in tasks
+        proc = fetch(task; raw=true).processor
+        push!(placements, proc)
+        @test proc in allowed
+        @test Dagger.check_uniform(proc)
+        @test fetch(task) == (proc.rank, proc.innerProc.tid, value)
+    end
+    @test [count(==(p), placements) for p in allowed] == fill(4, length(allowed))
+end
+
+@testset "Plain spawn placement" begin
+    cpus = filter(p -> p.innerProc isa Dagger.ThreadProc, mpi_procs())
+    target = nranks - 1
+    rank_cpus = filter(p -> p.rank == target, cpus)
+    last_cpus = [last(filter(p -> p.rank == r, cpus)) for r in 0:nranks-1]
+    thread_scope = Dagger.UnionScope(map(Dagger.ExactScope, last_cpus))
+    cases = (
+        ("default", (;), cpus),
+        ("all ranks", (; scope=Dagger.scope(mpi_ranks=:)), cpus),
+        ("rank scope", (; scope=Dagger.scope(mpi_rank=target)), rank_cpus),
+        ("compute scope", (; compute_scope=thread_scope), last_cpus),
+        ("exact scope", (; scope=Dagger.ExactScope(last(cpus))), [last(cpus)]),
+    )
+    @testset "Canonical processor listing" begin
+        procs = Dagger.Processor[reverse(cpus)...]
+        @test @inferred(Dagger.select_processors_uniform!(procs, accel)) === procs
+        @test procs == cpus
+        task = Dagger.Thunk(mpi_spawn_location; id=10)
+        @test @inferred(Dagger.select_processors_uniform!(procs, accel, task, nothing)) === procs
+        empty!(procs)
+        @test Dagger.select_processors_uniform!(procs, accel, task, nothing) === procs
+    end
+    # Start the scheduler before checking complete rounds of consecutive IDs.
+    fetch(Dagger.@spawn mpi_spawn_location())
+    @testset "uniformity checks=$checks" for checks in (true, false)
+        try
+            @testset "$label" for (label, options, allowed) in cases
+                Dagger.check_uniformity!(checks)
+                tasks = Dagger.with_options(; options...) do
+                    [Dagger.@spawn(mpi_spawn_location()) for _ in 1:4length(allowed)]
+                end
+                foreach(wait, tasks)
+                Dagger.check_uniformity!(true)
+                check_spawn_placements(tasks, allowed)
+            end
+
+            @testset "Chunk owner preference" begin
+                chunk = Dagger.tochunk(42, proc_for_rank(target), space_for_rank(target))
+                Dagger.check_uniformity!(checks)
+                tasks = [Dagger.@spawn(mpi_spawn_location(chunk)) for _ in 1:4length(rank_cpus)]
+                foreach(wait, tasks)
+                Dagger.check_uniformity!(true)
+                check_spawn_placements(tasks, rank_cpus, 42)
+
+                # A scope can exclude the input's owner. Fall back to the
+                # permitted rank, spreading across all its threads.
+                allowed = filter(p -> p.rank == 0, cpus)
+                Dagger.check_uniformity!(checks)
+                tasks = Dagger.with_options(; scope=Dagger.scope(mpi_rank=0)) do
+                    [Dagger.@spawn(mpi_spawn_location(chunk)) for _ in 1:4length(allowed)]
+                end
+                foreach(wait, tasks)
+                Dagger.check_uniformity!(true)
+                check_spawn_placements(tasks, allowed, 42)
+            end
+
+            @testset "DTask owner preference" begin
+                Dagger.check_uniformity!(checks)
+                producer = Dagger.@spawn scope=Dagger.ExactScope(last(rank_cpus)) mpi_spawn_value()
+                tasks = [Dagger.@spawn(mpi_spawn_location(producer)) for _ in 1:4length(rank_cpus)]
+                foreach(wait, tasks)
+                Dagger.check_uniformity!(true)
+                check_spawn_placements(tasks, rank_cpus, 42)
+            end
+        finally
+            Dagger.check_uniformity!(true)
+        end
+    end
+end