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] diff --git a/src/RegisterDriver.jl b/src/RegisterDriver.jl index b48cf4b..b64ceb3 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,15 +69,14 @@ 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") - 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 @@ -121,15 +120,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 @@ -145,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 @@ -250,9 +258,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..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 +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 @@ -77,4 +77,51 @@ 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 + +# 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 4af382e..6479c15 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,67 @@ 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 "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))