From 70067db0f727756d6963ebc69960e792d30f596f Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Mon, 12 Oct 2026 11:14:56 -0700 Subject: [PATCH] Spread MPI eager tasks without rank-local scheduling state Uniform placement ignored local load and chose the same first processor for every plain spawn. Rotate the canonical candidates by the uniform thunk ID so placement is independent of each rank's readiness order. Keep input-owner affinity within its preferred rank, honor restricted scopes, and retain the canonical listing for Datadeps/proc-grid calls without a task. Reproduced #760 on Julia 1.12.4 / MPI 0.20.27: 16 CPU-bound tasks on 2 ranks x 2 threads went from 16 serial tasks on rank 0/thread 1 in 4.77s to four tasks on each processor in 1.28s. One rank x four threads likewise went from 4.74s to 1.29s, using all four threads. GC deltas after 10 warmups + GC.gc(), minima over five runs of 1,000 selections with identical reversed input lists (four processors): - No affinity: 323,000 allocations / 13,840,000 bytes before and after. - Owner affinity: 323,000 / 13,968,000 before; 324,000 / 13,984,000 after. The affinity scan adds one 16-byte backend box per task; an out-of-line pass-through barrier avoids boxing it separately for each candidate. Validation: new placement tests fail on the original code (14 failures per rank); final 2x2 placement tests pass (476 assertions per rank), as do Datadeps thread placement, 1x4 placement, return-type inference checks, and test/allocations.jl without changing bounds. The broader MPI CPU suite exceeded its 10-minute limit compiling high-dimensional stencils; no test failure was reported before the timeout. Record placement, boxing, and MPI timeout cleanup lessons in AGENTS.md. --- AGENTS.md | 27 ++++++++++++++ src/acceleration.jl | 29 +++++++++++++++ test/mpi/cpu.jl | 2 ++ test/mpi/spawn.jl | 86 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 144 insertions(+) create mode 100644 test/mpi/spawn.jl diff --git a/AGENTS.md b/AGENTS.md index 2ee908bf8..0385d5384 100644 --- a/AGENTS.md +++ b/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. diff --git a/src/acceleration.jl b/src/acceleration.jl index 1b4092c68..409bd2c5a 100644 --- a/src/acceleration.jl +++ b/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 diff --git a/test/mpi/cpu.jl b/test/mpi/cpu.jl index fd07b7658..65c2697ea 100644 --- a/test/mpi/cpu.jl +++ b/test/mpi/cpu.jl @@ -145,6 +145,8 @@ end MPI.Barrier(comm) end +include("spawn.jl") + @testset "Datadeps" begin include("threads.jl") diff --git a/test/mpi/spawn.jl b/test/mpi/spawn.jl new file mode 100644 index 000000000..fc74028bc --- /dev/null +++ b/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