Parallel Loops
TanayLabUtilities.ParallelLoops
—
Module
Parallel loops.
TanayLabUtilities.ParallelLoops.parallel_loop_wo_rng
—
Function
parallel_loop_wo_rng(
body::Function,
indices::AbstractVector{<:Integer};
name::AbstractString = ".loop",
policy::Symbol = :greedy_sticky,
order::Maybe{AbstractVector{<:Integer}} = nothing,
weights::Maybe{AbstractVector{<:Integer}} = nothing,
nested::Bool = false,
progress::Maybe{AbstractProgress} = nothing,
)::Nothing
Run the
body
in parallel, passing it the iteration
index
. The
policy
is one of:
-
:greedy_sticky-nthreads()sticky worker tasks spawned via@threads :static for _ in 1:nthreads(), each pulling the next index position from anAtomic{Int}counter. Combines the dynamic load balancing of:greedywith the threadid stability of:static: tasks never migrate, sothreadid()-indexed scratch is safe inside the body. Pair with a heaviest-firstorder(orweights) to attack the tail-latency problem - heavy items dispatch first, light items fill the tail. -
:greedy,:dynamic,:static- passed straight to@threads. Note that:greedyand:dynamiccreate non-sticky tasks that may migrate across threads at yield points, sothreadid()-indexed scratch is unsafe under those. -
:serial- run the loop on the calling thread (useful for debugging).
If
order
is specified, it must be a permutation of the
positions
of
indices
- that is, a length-
length(indices)
vector whose values are a permutation of
1:length(indices)
. The body still receives the corresponding
indices[order[k]]
value at each step k; only the visit order is altered. Typical usage is heaviest-work-first (so the longest items dispatch as early as possible): pass
sortperm(weight_per_index; rev = true)
. Under
:greedy
,
:dynamic
,
:greedy_sticky
, or
:serial
, the iteration visits the positions in
order
as given. Under
:static
,
order
is reshuffled in place in a round-robin fashion each contiguous chunk gets one item from each weight quantile, balancing the static partition. Mutates the caller's
order
array.
If
weights
is specified, it must be a length-
length(indices)
vector of non-negative integers giving the estimated work for each position in
indices
. When
order
is not given,
weights
computes it to be
sortperm(weights; rev = true)
so the heaviest items dispatch first. When
progress
is also given, each iteration advances the bar by the visited position's
weights
entry instead of by 1; the caller is responsible for sizing the progress total to
sum(weights)
so the percentage reflects work done rather than items done.
weights
and
order
are independent and may be combined: pass both to keep an explicit visit order while reporting work-weighted progress. Mutually exclusive with
progress_chunk
(the unweighted-throttling alternative).
If this is invoked from inside another parallel loop and
nested
is
false
(the default),
policy
is ignored and the loop is executed serially. This makes functions safe to compose: parallel at the top level, serial when called from inside another parallel loop. Pass
nested = true
to opt out of the safety demotion - the loop then uses
policy
as if at the top level. The caller is responsible for the scratch contract: under
:static
/
:greedy_sticky
each spawned sub-task has a stable
threadid()
, but multiple outer iterations may be inside their nested loops at the same time, so per-thread scratch indexed by inner
threadid()
must be local to the outer iteration (not a global per-thread array shared across all outer iterations).
The
name
is used for
flame_timed
for the whole loop. If the loop is parallel, the whole loop is a line in the serial flame file, and each iteration is a line in the parallel flame file. The serial file therefore gives a view of the elapsed time of the top-level loops, to identify what actually matters, and the parallel file shows the internal breakdown of the computations in each loop. It is sadly impossible to include threaded tasks in the same flamegraph in a meaningful way (at least not without extending the flamegraph format and visualization).
If
progress
is specified, it is updated for each iteration. If
progress_chunk
is specified, then
indices
must be a range
1:n
,
order
must be
nothing
, and we throttle calls to the progress bar update to one every
progress_chunk
.
The code inside the loop body should not use any random number generation. Use
parallel_loop_with_rng
if random number generation is needed.
using Random
size = 10
for policy in (:serial, :greedy, :dynamic, :static, :greedy_sticky)
results = zeros(Int, size)
parallel_loop_wo_rng(1:size; policy) do index
results[index] = index
return nothing
end
@assert results == collect(1:size)
end
# output
TanayLabUtilities.ParallelLoops.parallel_loop_with_rng
—
Function
parallel_loop_with_rng(
body::Function,
indices::AbstractVector{<:Integer};
name::AbstractString = ".loop",
policy::Symbol = :greedy,
order::Maybe{AbstractVector{<:Integer}} = nothing,
weights::Maybe{AbstractVector{<:Integer}} = nothing,
progress::Maybe{AbstractProgress} = nothing,
seed::Maybe{Integer} = nothing,
rng::Maybe{AbstractRNG} = nothing
)::Nothing
Run the
body
in parallel, passing it the iteration
index
and a separate
rng
that is seeded to a reproducible state regardless of the allocation of tasks to threads. A copy of this
rng
is given to each iteration, after being reset to
seed + index
for reproducibility. If no
seed
is specified, it is just sampled
rng
before the loop starts. If the
rng
isn't given, then this uses (and sets for each iteration) the
default_rng()
. In this case passing it to the body is redundant but is still done for consistency.
The
policy
is passed to
@threads
if it is one of (the default
:greedy
,
:dynamic
, or
:static
). If it is
:serial
, then the loop is not run in parallel (useful for debugging).
If
order
and/or
weights
are specified, they are forwarded to
parallel_loop_wo_rng
- see there for the position-permutation semantics, the in-place round-robin shuffle under
:static
, and the work-weighted progress accounting. The body still receives the same
indices[order[k]]
value at each step k, and the seed is
seed + indices[order[k]]
so reproducibility is independent of the visit order.
If this is invoked from inside another parallel loop and
nested
is
false
(the default),
policy
is ignored and the loop is executed serially. This makes functions safe to compose: parallel at the top level, serial when called from inside another parallel loop. Pass
nested = true
to opt out of the safety demotion - the loop then uses
policy
as if at the top level. The caller is responsible for the scratch contract: under
:static
/
:greedy_sticky
each spawned sub-task has a stable
threadid()
, but multiple outer iterations may be inside their nested loops at the same time, so per-thread scratch indexed by inner
threadid()
must be local to the outer iteration (not a global per-thread array shared across all outer iterations).
The
name
is used for
flame_timed
for the whole loop. If the loop is parallel, the whole loop is a line in the serial flame file, and each iteration is a line in the parallel flame file. The serial file therefore gives a view of the elapsed time of the top-level loops, to identify what actually matters, and the parallel file shows the internal breakdown of the computations in each loop. It is sadly impossible to include threaded tasks in the same flamegraph in a meaningful way (at least not without extending the flamegraph format and visualization).
If
progress
is specified, it is updated for each iteration.
Yes, the
TaskLocalRNG
is supposed to do this, but, it actually depends on the way tasks are allocated to threads. The implementation here will give the same results regardless of the thread scheduling policy. Sigh.
using Random
size = 10
function collect_rng(rng::AbstractRNG)::Vector{Float64}
results = zeros(Float64, size)
parallel_loop_with_rng(1:size; rng) do index, rng
results[index] = rand(rng)
end
@assert results[1] != results[2]
return results
end
@assert collect_rng(MersenneTwister(1)) == collect_rng(MersenneTwister(1))
function collect_default_rng()::Vector{Float64}
results = zeros(Float64, size)
parallel_loop_with_rng(1:size; seed = 123456, policy = :dynamic) do index, _
results[index] = rand()
end
@assert results[1] != results[2]
return results
end
@assert collect_default_rng() == collect_default_rng()
# output
TanayLabUtilities.ParallelLoops.is_in_parallel_loop
—
Function
is_in_parallel_loop()::Bool
Return whether we are inside a
parallel_loop_wo_rng
or
parallel_loop_with_rng
.
TanayLabUtilities.ParallelLoops.DebugProgress
—
Function
DebugProgress(n::Integer; kwargs...)::Maybe{Progress}
Same as
Progress
in
ProgressMeter
, but returns
nothing
if debug is not enabled for the modules calling
DebugProgress
.
TanayLabUtilities.ParallelLoops.DebugProgressUnknown
—
Function
DebugProgressUnknown(; kwargs...)::Maybe{ProgressUnknown}
Same as
ProgressUnknown
in
ProgressMeter
, but returns
nothing
if debug is not enabled for the modules calling
DebugProgressUnknown
. Use this for progress bars where the total amount of work is not known in advance (for example, a planning phase that walks a data set to discover the set of properties to act on).