Skip to content

driver can hand one worker to two concurrent tasks #28

Description

@tmcgrath325

In the threaded driver(outfile, algorithms, img, mon), each frame's worker is selected by threadid() inside a @threads :dynamic loop. Under dynamic scheduling, two concurrently-running tasks can observe the same thread id, and are then handed the same AbstractWorker object and the same monitor Dict. Any worker holding mutable state is then mutated by two registrations at
once. This is easy to miss on a laptop, but happens more frequently as the number of threads is increased.

The code

https://github.com/HolyLab/RegisterDriver.jl/blob/master/src/RegisterDriver.jl#L125-L134

@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]])
        put!(results_ch, (movidx, deepcopy(tmp)))
    end
    yield()
end

Diagnosis from Claude:

aindices maps a thread id to a worker index, so algorithms[aindices[tid]] is shared by every task that observes tid. @threads :dynamic spawns tasks that may interleave on a thread whenever one of them reaches a yield point, and worker implementations reach several (I/O from println/@warn, FFTW planning, the blocking put! on the bounded channel). Two tasks that overlap this way run against one worker.

Reproducer

This replicates the dispatch only — no images, no registration — so it runs anywhere RegisterDriver is installed. It records which worker slot took each frame and when, then counts pairs of frames whose intervals overlap on the same slot.

using Base.Threads
using RegisterDriver: threadids

function probe(n; work_us = 2000)
    tpool = threadids()
    owner = fill(0, n); start = fill(0.0, n); stop = fill(0.0, n)
    ch = Channel{Tuple{Int,Int}}(32)
    writer = @async for _ in ch end

    @threads :dynamic for movidx in 1:n
        tid = threadid()
        if tid in tpool
            owner[movidx] = tid
            start[movidx] = time()
            t0 = time_ns()                       # stand-in for `worker`
            while (time_ns() - t0) < work_us * 1000
                yield()
            end
            stop[movidx] = time()
            put!(ch, (movidx, tid))
        end
        yield()
    end
    close(ch); wait(writer)

    overlaps = 0
    for i in 1:n, j in (i+1):n
        owner[i] == 0 && continue
        owner[i] == owner[j] || continue
        (start[i] < stop[j] && start[j] < stop[i]) && (overlaps += 1)
    end
    return (skipped = count(iszero, owner), overlaps)
end

println("nthreads=", nthreads())
for n in (16, 64, 200)
    println("  n=$n: ", probe(n))
end

Pairs of frames sharing a worker slot concurrently, Julia 1.12.6:

threads n=16 n=64 n=200
4 0 1 1
8 0 1 6
16 0 11 39

Observed effect

With RegisterWorkerApertures.Apertures, the shared state is algorithm.affinepenalty — a mutable struct whose λ is assigned during optimization — plus the shared monitor Dict. Registering a 3-D calcium-imaging recording (1031×1025×50, gridsize=(20,20,3)), roughly half of all 16-thread runs died with

nested task error: Initial value must be finite
  RegisterOptimize.jl:638  (fval0 = MOI.eval_objective(objective, uvec))

while 4-thread runs of the same configuration were consistently clean. The failure was intermittent: identical inputs failed twice and then passed four times.

I ruled out the data (finite, no blank planes), λ (fails at 1e-8 and 1e-2), mxshift (fails at 50, 20 and 10), worker reuse across frames (clean at 4 threads with 4 frames per worker), and RegisterMismatch's inner threading (fails with it both on and off). Replacing the dispatch with a worker pool in my own code made 8/8 runs at 16 threads clean and bit-identical, where results had previously varied run to run.

The crash is the visible half. The quiet half is that a collision which does not produce a NaN yields a deformation computed against another frame's penalty state, and is written to the output as if it were valid.

Suggested fix

Give each task exclusive ownership of a worker instead of deriving one from threadid():

if parallel
    next = Threads.Atomic{Int}(1)
    @sync for k in eachindex(algorithms, mon)
        Threads.@spawn while true
            movidx = Threads.atomic_add!(next, 1)
            movidx > n && break
            tmp = worker(algorithms[k], img, movidx, mon[k])
            put!(results_ch, (movidx, deepcopy(tmp)))
        end
    end
end

Every worker is used by exactly one task for the run's duration, so no synchronization is needed around the workers or their monitor dicts, and load balances the same way the current chunking does.

One consequence worth noting: threadids(), tpool and aindices become unnecessary for dispatch. If threadids() is exported mainly to let callers size algorithms to the set of ids @threads actually uses, that need goes away too — the pool works for any length(algorithms), and workertid no longer influences which images a worker receives.

One existing test needs adjusting: test/runtests.jl:88 asserts that every supplied worker handled at least one image. That is already unsatisfiable on master with more than 7 threads (the test image has 7 frames), and dynamic distribution makes it the wrong invariant. Asserting that every image was registered by one of the supplied workers holds in both cases.

Separate, smaller issue in the same function

driver calls init!(algorithms[1]) (line 80) and close!(algorithms[1]) (line 149) — only the first worker. Both are no-ops for CPU Apertures, but with dev >= 0 only worker 1 gets its cuda_objects, so every other worker hits a KeyError on :d_fixed in worker. That would make multi-worker GPU registration unusable, and whatever worker 1 acquires is the only thing released. foreach(init!, algorithms) / foreach(close!, algorithms) is the obvious fix.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions