Skip to content

MPI: keep tasks with acceleration=DistributedAcceleration() rank-local - #779

Open
AnnaLena77 wants to merge 1 commit into
JuliaParallel:masterfrom
AnnaLena77:fix/mpi-rank-local-tasks
Open

AnnaLena77 wants to merge 1 commit into
JuliaParallel:masterfrom
AnnaLena77:fix/mpi-rank-local-tasks

Conversation

@AnnaLena77

Copy link
Copy Markdown

Fixes #778.

Under accelerate!(:mpi), tasks spawned with acceleration=DistributedAcceleration() (or inside with_options(; acceleration=DistributedAcceleration())) are meant to run only on the rank that spawns them. Two things broke that when ranks spawn such tasks independently:

  • Placement: accel_matches_proc(::DistributedAcceleration, proc) is true for every processor, and compatible_processors(::DistributedAcceleration, ...) skipped the filter completely. Some of these tasks were placed on MPIProcessors, sometimes on another rank, and then never ran. This PR excludes MPI processors for DistributedAcceleration in MPIExt and removes the shortcut, so the generic filter applies.
  • Task IDs: these tasks took IDs from the same counter as uniform tasks, and MPI tags are derived from those IDs. With different numbers of local tasks per rank, later uniform tasks got different IDs on each rank and hung. Under uniform execution, tasks with a non-uniform acceleration now take IDs from a separate range (LOCAL_ID_COUNTER, starting at typemax(Int) ÷ 2).

Tests: test/mpi/rank_local.jl (included in test/mpi/cpu.jl) spawns a different number of local tasks on every rank, both ways, checks that they run on ThreadProcs of their own rank, then runs uniform work with cross-rank data, and checks that a local task can use a uniform result owned by its own rank. Without the fix it fails (tasks on MPI processors, Non-uniform value, hangs); with it, it passes. The full test/mpi/cpu.jl passes on 2 ranks × 2 threads (MPICH_jll, macOS).

Not covered: a local task using a uniform result owned by another rank fails with an error, and passing a local task's result back into uniform work isn't addressed here.

Written with help from Claude (Anthropic); I checked the results.

@jpsamaroo jpsamaroo left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this! I'm not approving or rejecting this PR - I need some time to figure out the design, and then we can come back to this.

Comment thread src/dtask.jl
return UInt64(Threads.atomic_add!(LOCAL_ID_COUNTER, 1))
end
return eager_next_id()
end

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes me quite uncomfortable - I want to think more on the design before committing to something like this.

Comment thread src/acceleration.jl
# accel_matches_proc is identically true for DistributedAcceleration; skip the
# second Set construction
compatible_processors(accel::DistributedAcceleration, scope::AbstractScope, procs::Vector{<:Processor}) =
compatible_processors(scope, procs)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this will cause an increase in allocations, as it was a fast-path avoiding Set allocation.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

MPI: tasks with acceleration=DistributedAcceleration() go missing when ranks spawn them independently

2 participants