From 3f82e45f7a6c1dcb15d7fa1891db7c652e49f1e4 Mon Sep 17 00:00:00 2001 From: tmcgrath325 Date: Mon, 3 Aug 2026 21:24:50 +0000 Subject: [PATCH 1/3] Give each worker its own task in the parallel driver Dispatching on `threadid()` inside `@threads :dynamic` made a worker's `workertid` decide which images reached it, with two consequences. Images whose task ran on a thread outside the set of worker ids were skipped by the `tid in tpool` guard: never registered, their slice of the output never written, and no error raised. Worker ids that match no thread -- `tid = i` over `1:nthreads()`, say, when the scheduler hands out 2:17 -- lost every image. Workers could also be shared. Tasks spawned per image interleave on a thread at any yield point, so two running concurrently may observe the same id and register against the same worker object and the same monitor dict, both of which they mutate in place. Collisions grow with thread count; registering a 200-image stack across 16 threads showed 39 pairs of images handled at once by one worker. `driver` now spawns one task per worker, each taking the next unclaimed image from an atomic counter, so a worker is only ever in use by one task and needs no locking. Images are distributed independently of `workertid`, and `algorithms` may have any length. Two testsets cover this. "images are distributed independently of workertid" gives the workers ids matching no thread and requires all 64 images to be registered; it reported 0 under the old dispatch. "workers are used exclusively" checks that `AlgExclusive` is never entered while already busy -- an invariant guard rather than a reproduction, since a dummy worker yields too briefly to collide reliably. The existing multi-worker test asserted that every worker handled at least one image. That holds only when images outnumber workers -- it already failed on more than seven threads -- and image distribution is now dynamic, so it instead checks that every image was registered by one of the workers supplied. Assisted-by: Claude Opus 5 --- src/RegisterDriver.jl | 37 +++++++++++++++++++-------------- test/WorkerDummy.jl | 27 +++++++++++++++++++++++- test/runtests.jl | 48 +++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 94 insertions(+), 18 deletions(-) diff --git a/src/RegisterDriver.jl b/src/RegisterDriver.jl index b48cf4b..a75c6c2 100644 --- a/src/RegisterDriver.jl +++ b/src/RegisterDriver.jl @@ -17,7 +17,7 @@ using ImageMetadata: ImageMetadata using JLD: JLD, jldopen using RegisterCore: RegisterCore, NumDenom using RegisterWorkerShell: RegisterWorkerShell, AbstractWorker, ArrayDecl, - close!, init!, load_mm_package, worker, workertid + close!, init!, load_mm_package, worker using SharedArrays: SharedArrays, SharedArray, sdata using StaticArrays: StaticArrays, StaticArray using Base.Threads: @threads, nthreads, threadid @@ -69,10 +69,6 @@ function driver(outfile::AbstractString, algorithms::AbstractVector, img, mon::A nalgs = length(algorithms) nummon = length(mon) nummon == nalgs || error("Number of monitors must equal number of workers") - numthreads = nthreads() - tpool = map(workertid, algorithms) - aindices = parallel ? Dict(map((alg, aidx) -> (workertid(alg) => aidx), algorithms, 1:length(algorithms))...) : - Dict(threadid() => 1) n = nimages(img) fs = FormatSpec("0$(ndigits(n))d") @@ -121,15 +117,24 @@ function driver(outfile::AbstractString, algorithms::AbstractVector, img, mon::A end if parallel - # writer_task shares the first thread, making static scheduling inefficient - @threads :dynamic for movidx in 1:n - tid = threadid() - if tid in tpool - println("thread $tid processing $movidx") - tmp = worker(algorithms[aindices[tid]], img, movidx, mon[aindices[tid]]) + # One task per worker, each taking the next unclaimed image. A worker + # must not be chosen by `threadid()`: tasks spawned per image can + # interleave on a thread at any yield point, so two running + # concurrently may observe the same id and would then share one + # worker — and one monitor dict — corrupting both registrations. + # Owning a worker for the task's whole life makes that impossible + # without any locking. + nextidx = Threads.Atomic{Int}(1) + @sync for k in eachindex(algorithms, mon) + Threads.@spawn while true + movidx = Threads.atomic_add!(nextidx, 1) + movidx > n && break + println("worker $k processing $movidx") + # `mon[k]` is reused for every image this worker handles, so + # the writer needs a snapshot rather than a live reference. + tmp = worker(algorithms[k], img, movidx, mon[k]) put!(results_ch, (movidx, deepcopy(tmp))) end - yield() end else for movidx in 1:n @@ -250,9 +255,11 @@ end Return the sorted list of thread IDs that Julia's scheduler actually assigns to tasks spawned with `@threads` and `Threads.@spawn`. -Julia's main thread (ID 1) typically does not execute worker tasks. The -returned IDs are useful for configuring `AbstractWorker` instances that pin -execution to a specific thread via the `workertid` field. +Julia's main thread (ID 1) typically does not execute worker tasks. + +[`driver`](@ref) does not need this: it gives each worker its own task and +distributes images between them, so `algorithms` may have any length and a +worker's `workertid` does not affect which images it receives. # Example diff --git a/test/WorkerDummy.jl b/test/WorkerDummy.jl index 96ffd0f..fae3a3e 100644 --- a/test/WorkerDummy.jl +++ b/test/WorkerDummy.jl @@ -4,7 +4,7 @@ module WorkerDummy using RegisterWorkerShell, Distributed import RegisterWorkerShell: worker -export Alg1, Alg2, Alg3, Alg4 +export Alg1, Alg2, Alg3, Alg4, AlgExclusive # Dispatch on the algorithm used to perform registration # Each algorithm has a container it uses for storage and communication @@ -77,4 +77,29 @@ function worker(algorithm::Alg4, moving, tindex, mon) return mon end +# AlgExclusive: detects a worker being used by two registrations at once. +# `driver` must never hand one worker to concurrently-running tasks, because a +# worker's fields and its monitor dict are mutated in place. +mutable struct AlgExclusive <: Alg + busy::Bool # set for the duration of a call, checked on entry + reentered::Bool # sticky: a second entry was seen while busy + ncalls::Int + workertid::Int +end +AlgExclusive(; tid = 1) = AlgExclusive(false, false, 0, tid) + +function worker(algorithm::AlgExclusive, moving, tindex, mon) + algorithm.busy && (algorithm.reentered = true) + algorithm.busy = true + algorithm.ncalls += 1 + # Yield points are what let two tasks interleave on one thread; a real + # worker reaches them through I/O, FFT planning and the like. + for _ in 1:20 + yield() + end + monitor!(mon, :tindex, tindex) + algorithm.busy = false + return mon +end + end # module diff --git a/test/runtests.jl b/test/runtests.jl index 4af382e..0424cef 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -83,9 +83,12 @@ end u0 = JLD.load(fn, "u0") @test tform[:, 4] == collect(range(1, stop = 12, length = 12) .+ 4) @test u0[:, :, 2] == fill(-2, (3, 3)) + # Images are distributed between the workers, so which worker takes which + # image is not fixed. What must hold is that every image was registered, + # each by one of the workers supplied. tid = JLD.load(fn, "workertid") - indx = unique(indexin(tid, tids)) - @test length(indx) == length(tids) && all(indx .> 0) + @test length(tid) == size(img, 3) + @test all(in(tids), tid) rm(fn) # Non-BitsType array (ComplexF32) alongside an unpackable string: exercises @@ -153,6 +156,47 @@ end rm(fn) end +@testset "workers are used exclusively" begin + # A worker's fields and monitor dict are mutated in place, so `driver` must + # never hand one to two concurrently-running tasks. Selecting the worker by + # `threadid()` did: tasks spawned per image interleave on a thread at any + # yield point, so two running at once could observe the same id. + workdir = tempname() + mkdir(workdir) + n = 64 + img = AxisArray(SharedArray{Float32}((16, 16, n)), :y, :x, :time) + + tids = threadids() + algs = [AlgExclusive(; tid = t) for t in tids] + mons = [Dict{Symbol,Any}(:tindex => 0) for _ in eachindex(tids)] + fn = joinpath(workdir, "exclusive.jld") + driver(fn, algs, img, mons; parallel = true) + + @test !any(a -> a.reentered, algs) + @test sum(a -> a.ncalls, algs) == n + @test JLD.load(fn, "tindex") == 1:n + rm(fn) +end + +@testset "images are distributed independently of workertid" begin + # A worker's `workertid` must not decide which images reach it. Dispatching + # on `threadid()` dropped, without any error, every image whose task ran on + # a thread outside the set of worker ids. + workdir = tempname() + mkdir(workdir) + n = 64 + img = AxisArray(SharedArray{Float32}((16, 16, n)), :y, :x, :time) + + algs = [AlgExclusive(; tid = 1000 + i) for i in 1:4] # ids matching no thread + mons = [Dict{Symbol,Any}(:tindex => 0) for _ in 1:4] + fn = joinpath(workdir, "anytid.jld") + driver(fn, algs, img, mons; parallel = true) + + @test sum(a -> a.ncalls, algs) == n + @test JLD.load(fn, "tindex") == 1:n + rm(fn) +end + @testset "nicehdf5 specializations" begin # Plain SharedArray → sdata sa = SharedArray{Float32}((3, 4)) From f096d72863cc4156714abcbffd7514fe1b53bfdf Mon Sep 17 00:00:00 2001 From: tmcgrath325 Date: Tue, 4 Aug 2026 00:58:56 +0000 Subject: [PATCH 2/3] Initialize and close every worker `driver` called `init!` and `close!` on `algorithms[1]` only, so workers 2:end registered images without their per-worker resources ever being set up, and whatever the first worker acquired was the only thing released. Algorithms with no-op lifecycle methods were unaffected; those holding device contexts or scratch buffers were not. `AlgLifecycle` records its own lifecycle calls and flags any image that reaches it before `init!`. Assisted-by: Claude Opus 5 --- src/RegisterDriver.jl | 11 +++++++---- test/WorkerDummy.jl | 26 ++++++++++++++++++++++++-- test/runtests.jl | 20 ++++++++++++++++++++ 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/RegisterDriver.jl b/src/RegisterDriver.jl index a75c6c2..b64ceb3 100644 --- a/src/RegisterDriver.jl +++ b/src/RegisterDriver.jl @@ -72,8 +72,11 @@ function driver(outfile::AbstractString, algorithms::AbstractVector, img, mon::A n = nimages(img) fs = FormatSpec("0$(ndigits(n))d") - println("Initializing algorithm") - init!(algorithms[1]) + println("Initializing algorithms") + # Every worker is initialized, not just the first: `init!` sets up per-worker + # resources (device contexts, scratch buffers), and every worker registers + # images of its own. + foreach(init!, algorithms) println("Working on algorithm and saving the result") jldopen(outfile, "w") do file @@ -150,8 +153,8 @@ function driver(outfile::AbstractString, algorithms::AbstractVector, img, mon::A wait(writer_task) end - println("Closing algorithm") - close!(algorithms[1]) + println("Closing algorithms") + foreach(close!, algorithms) return nothing end diff --git a/test/WorkerDummy.jl b/test/WorkerDummy.jl index fae3a3e..afd30d4 100644 --- a/test/WorkerDummy.jl +++ b/test/WorkerDummy.jl @@ -2,9 +2,9 @@ module WorkerDummy using RegisterWorkerShell, Distributed -import RegisterWorkerShell: worker +import RegisterWorkerShell: worker, init!, close! -export Alg1, Alg2, Alg3, Alg4, AlgExclusive +export Alg1, Alg2, Alg3, Alg4, AlgExclusive, AlgLifecycle # Dispatch on the algorithm used to perform registration # Each algorithm has a container it uses for storage and communication @@ -102,4 +102,26 @@ function worker(algorithm::AlgExclusive, moving, tindex, mon) return mon end +# AlgLifecycle: records its own `init!`/`close!` calls. A worker that registers +# images must have been initialized first, so `driver` owes every element of +# `algorithms` an `init!` and a matching `close!`. +mutable struct AlgLifecycle <: Alg + ninit::Int + nclose::Int + ncalls::Int + uninitialized::Bool # sticky: an image arrived before `init!` + workertid::Int +end +AlgLifecycle(; tid = 1) = AlgLifecycle(0, 0, 0, false, tid) + +init!(algorithm::AlgLifecycle) = (algorithm.ninit += 1; nothing) +close!(algorithm::AlgLifecycle) = (algorithm.nclose += 1; nothing) + +function worker(algorithm::AlgLifecycle, moving, tindex, mon) + algorithm.ninit == 0 && (algorithm.uninitialized = true) + algorithm.ncalls += 1 + monitor!(mon, :tindex, tindex) + return mon +end + end # module diff --git a/test/runtests.jl b/test/runtests.jl index 0424cef..6479c15 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -197,6 +197,26 @@ end rm(fn) end +@testset "every worker is initialized and closed" begin + # A worker registers images only after `init!` has set up its resources, so + # each element of `algorithms` needs its own `init!` and matching `close!`. + workdir = tempname() + mkdir(workdir) + n = 16 + img = AxisArray(SharedArray{Float32}((16, 16, n)), :y, :x, :time) + + algs = [AlgLifecycle(; tid = i) for i in 1:4] + mons = [Dict{Symbol,Any}(:tindex => 0) for _ in 1:4] + fn = joinpath(workdir, "lifecycle.jld") + driver(fn, algs, img, mons; parallel = true) + + @test all(a -> a.ninit == 1, algs) + @test all(a -> a.nclose == 1, algs) + @test !any(a -> a.uninitialized, algs) + @test sum(a -> a.ncalls, algs) == n + rm(fn) +end + @testset "nicehdf5 specializations" begin # Plain SharedArray → sdata sa = SharedArray{Float32}((3, 4)) From dd0a4f8be64b072e27ab7a1974d83b8089b245ec Mon Sep 17 00:00:00 2001 From: tmcgrath325 Date: Tue, 4 Aug 2026 01:55:09 +0000 Subject: [PATCH 3/3] Bump version to 1.0.3 Lets consumers require the parallel `driver` that gives each worker its own task and initializes every worker, rather than resolving to a version whose thread-id dispatch shares workers between concurrent tasks and silently skips frames. Assisted-by: Claude Opus 5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 1bc7154..55680d3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "RegisterDriver" uuid = "935ac36e-2656-11e9-1e3b-cbaa636797af" -version = "1.0.2" +version = "1.0.3" authors = ["Tim Holy "] [deps]