fixup! Add support for worker state callbacks

Author: JamesWrigley

src/cluster.jl

@@ -457,22 +457,34 @@ end
 ```
 """
 function addprocs(manager::ClusterManager; kwargs...)
+    params = merge(default_addprocs_params(manager), Dict{Symbol, Any}(kwargs))
+
     init_multi()
 
     cluster_mgmt_from_master_check()
 
-    new_workers = @lock worker_lock addprocs_locked(manager::ClusterManager; kwargs...)
+    new_workers = @lock worker_lock addprocs_locked(manager::ClusterManager, params)
+
+    callback_tasks = Dict{Any, Task}()
     for worker in new_workers
-        for callback in values(worker_added_callbacks)
-            callback(worker)
+        for (name, callback) in worker_added_callbacks
+            callback_tasks[name] = Threads.@spawn callback(worker)
         end
     end
 
+    running_callbacks = () -> ["'$(key)'" for (key, task) in callback_tasks if !istaskdone(task)]
+    while timedwait(() -> isempty(running_callbacks()), params[:callback_warning_interval]) === :timed_out
+        callbacks_str = join(running_callbacks(), ", ")
+        @warn "Waiting for these worker-added callbacks to finish: $(callbacks_str)"
+    end
+
+    # Wait on the tasks so that exceptions bubble up
+    wait.(values(callback_tasks))
+
     return new_workers
 end
 
-function addprocs_locked(manager::ClusterManager; kwargs...)
-    params = merge(default_addprocs_params(manager), Dict{Symbol,Any}(kwargs))
+function addprocs_locked(manager::ClusterManager, params)
     topology(Symbol(params[:topology]))
 
     if PGRP.topology !== :all_to_all
@@ -559,7 +571,8 @@ default_addprocs_params() = Dict{Symbol,Any}(
     :exeflags => ``,
     :env      => [],
     :enable_threaded_blas => false,
-    :lazy => true)
+    :lazy => true,
+    :callback_warning_interval => 10)
 
 
 function setup_launched_worker(manager, wconfig, launched_q)
@@ -872,6 +885,8 @@ end
 function _add_callback(f, key, dict)
     if !hasmethod(f, Tuple{Int})
         throw(ArgumentError("Callback function is invalid, it must be able to accept a single Int argument"))
+    elseif haskey(dict, key)
+        throw(ArgumentError("A callback function with key '$(key)' already exists"))
     end
 
     if isnothing(key)
@@ -889,14 +904,23 @@ _remove_callback(key, dict) = delete!(dict, key)
 
 Register a callback to be called on the master process whenever a worker is
 added. The callback will be called with the added worker ID,
-e.g. `f(w::Int)`. Returns a unique key for the callback.
+e.g. `f(w::Int)`. Chooses and returns a unique key for the callback if `key` is
+not specified.
+
+The worker-added callbacks will be executed concurrently. If one throws an
+exception it will not be caught and will bubble up through [`addprocs()`](@ref).
+
+Keep in mind that the callbacks will add to the time taken to launch workers; so
+try to either keep the callbacks fast to execute, or do the actual
+initialization asynchronously by spawning a task in the callback (beware of race
+conditions if you do this).
 """
 add_worker_added_callback(f::Base.Callable; key=nothing) = _add_callback(f, key, worker_added_callbacks)
 
 """
     remove_worker_added_callback(key)
 
-Remove the callback for `key`.
+Remove the callback for `key` that was added with [`add_worker_added_callback()`](@ref).
 """
 remove_worker_added_callback(key) = _remove_callback(key, worker_added_callbacks)
 
@@ -905,18 +929,19 @@ remove_worker_added_callback(key) = _remove_callback(key, worker_added_callbacks
 
 Register a callback to be called on the master process immediately before a
 worker is removed with [`rmprocs()`](@ref). The callback will be called with the
-worker ID, e.g. `f(w::Int)`. Returns a unique key for the callback.
+worker ID, e.g. `f(w::Int)`. Chooses and returns a unique key for the callback
+if `key` is not specified.
 
-All callbacks will be executed asynchronously and if they don't all finish
-before the `callback_timeout` passed to `rmprocs()` then the process will be
-removed anyway.
+All worker-exiting callbacks will be executed concurrently and if they don't
+all finish before the `callback_timeout` passed to `rmprocs()` then the process
+will be removed anyway.
 """
 add_worker_exiting_callback(f::Base.Callable; key=nothing) = _add_callback(f, key, worker_exiting_callbacks)
 
 """
     remove_worker_exiting_callback(key)
 
-Remove the callback for `key`.
+Remove the callback for `key` that was added with [`add_worker_exiting_callback()`](@ref).
 """
 remove_worker_exiting_callback(key) = _remove_callback(key, worker_exiting_callbacks)
 
@@ -926,14 +951,17 @@ remove_worker_exiting_callback(key) = _remove_callback(key, worker_exiting_callb
 Register a callback to be called on the master process when a worker has exited
 for any reason (i.e. not only because of [`rmprocs()`](@ref) but also the worker
 segfaulting etc). The callback will be called with the worker ID,
-e.g. `f(w::Int)`. Returns a unique key for the callback.
+e.g. `f(w::Int)`. Chooses and returns a unique key for the callback if `key` is
+not specified.
+
+If the callback throws an exception it will be caught and printed.
 """
 add_worker_exited_callback(f::Base.Callable; key=nothing) = _add_callback(f, key, worker_exited_callbacks)
 
 """
     remove_worker_exited_callback(key)
 
-Remove the callback for `key`.
+Remove the callback for `key` that was added with [`add_worker_exited_callback()`](@ref).
 """
 remove_worker_exited_callback(key) = _remove_callback(key, worker_exited_callbacks)
 
@@ -1176,15 +1204,17 @@ function _rmprocs(pids, waitfor, callback_timeout)
     lock(worker_lock)
     try
         # Run the callbacks
-        callback_tasks = Task[]
+        callback_tasks = Dict{Any, Task}()
         for pid in pids
-            for callback in values(worker_exiting_callbacks)
-                push!(callback_tasks, Threads.@spawn callback(pid))
+            for (name, callback) in worker_exiting_callbacks
+                callback_tasks[name] = Threads.@spawn callback(pid)
             end
         end
 
-        if timedwait(() -> all(istaskdone.(callback_tasks)), callback_timeout) === :timed_out
-            @warn "Some callbacks timed out, continuing to remove workers anyway"
+        if timedwait(() -> all(istaskdone.(values(callback_tasks))), callback_timeout) === :timed_out
+            timedout_callbacks = ["'$(key)'" for (key, task) in callback_tasks if !istaskdone(task)]
+            callbacks_str = join(timedout_callbacks, ", ")
+            @warn "Some worker-exiting callbacks have not yet finished, continuing to remove workers anyway. These are the callbacks still running: $(callbacks_str)"
         end
 
         rmprocset = Union{LocalProcess, Worker}[]
@@ -1335,8 +1365,12 @@ function deregister_worker(pg, pid)
 
     # Call callbacks on the master
     if myid() == 1
-        for callback in values(worker_exited_callbacks)
-            callback(pid)
+        for (name, callback) in worker_exited_callbacks
+            try
+                callback(pid)
+            catch ex
+                @error "Error when running worker-exited callback '$(name)'" exception=(ex, catch_backtrace())
+            end
         end
     end
 

test/distributed_exec.jl

@@ -1937,41 +1937,50 @@ include("splitrange.jl")
 end
 
 @testset "Worker state callbacks" begin
-    if nprocs() > 1
-        rmprocs(workers())
-    end
+    rmprocs(other_workers())
 
     # Smoke test to ensure that all the callbacks are executed
     added_workers = Int[]
     exiting_workers = Int[]
     exited_workers = Int[]
-    added_key = DistributedNext.add_worker_added_callback(pid -> push!(added_workers, pid))
+    added_key = DistributedNext.add_worker_added_callback(pid -> (push!(added_workers, pid); error("foo")))
     exiting_key = DistributedNext.add_worker_exiting_callback(pid -> push!(exiting_workers, pid))
     exited_key = DistributedNext.add_worker_exited_callback(pid -> push!(exited_workers, pid))
 
-    pid = only(addprocs(1))
+    # Test that the worker-added exception bubbles up
+    @test_throws TaskFailedException addprocs(1)
+
+    pid = only(workers())
     @test added_workers == [pid]
     rmprocs(workers())
     @test exiting_workers == [pid]
     @test exited_workers == [pid]
 
+    # Trying to reset an existing callback should fail
+    @test_throws ArgumentError DistributedNext.add_worker_added_callback(Returns(nothing); key=added_key)
+
     # Remove the callbacks
     DistributedNext.remove_worker_added_callback(added_key)
     DistributedNext.remove_worker_exiting_callback(exiting_key)
     DistributedNext.remove_worker_exited_callback(exited_key)
 
-    # Test that the `callback_timeout` option works
+    # Test that the worker-exiting `callback_timeout` option works and that we
+    # get warnings about slow worker-added callbacks.
     event = Base.Event()
     callback_task = nothing
+    added_key = DistributedNext.add_worker_added_callback(_ -> sleep(0.5))
     exiting_key = DistributedNext.add_worker_exiting_callback(_ -> (callback_task = current_task(); wait(event)))
-    addprocs(1)
 
-    @test_logs (:warn, r"Some callbacks timed out.+") rmprocs(workers(); callback_timeout=0.5)
+    @test_logs (:warn, r"Waiting for these worker-added callbacks.+") match_mode=:any addprocs(1; callback_warning_interval=0.05)
+    DistributedNext.remove_worker_added_callback(added_key)
+
+    @test_logs (:warn, r"Some worker-exiting callbacks have not yet finished.+") rmprocs(workers(); callback_timeout=0.5)
+    DistributedNext.remove_worker_exiting_callback(exiting_key)
 
     notify(event)
     wait(callback_task)
 
-    # Test that the previous callbacks were indeed removed
+    # Test that the initial callbacks were indeed removed
     @test length(added_workers) == 1
     @test length(exiting_workers) == 1
     @test length(exited_workers) == 1