Worker

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 <tim.holy@gmail.com>"]
 
 [deps]

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
 

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

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))