Add support for worker state callbacks

Author: JamesWrigley

docs/src/_changelog.md

@@ -18,6 +18,7 @@ This documents notable changes in DistributedNext.jl. The format is based on
   incompatibilities from both libraries being used simultaneously ([#10]).
 - [`other_workers()`](@ref) and [`other_procs()`](@ref) were implemented and
   exported ([#18]).
+- Implemented callback support for workers being added/removed etc ([#17]).
 
 ## [v1.0.0] - 2024-12-02
 

docs/src/index.md

@@ -52,6 +52,17 @@ DistributedNext.cluster_cookie()
 DistributedNext.cluster_cookie(::Any)
 ```
 
+## Callbacks
+
+```@docs
+DistributedNext.add_worker_added_callback
+DistributedNext.remove_worker_added_callback
+DistributedNext.add_worker_exiting_callback
+DistributedNext.remove_worker_exiting_callback
+DistributedNext.add_worker_exited_callback
+DistributedNext.remove_worker_exited_callback
+```
+
 ## Cluster Manager Interface
 
 This interface provides a mechanism to launch and manage Julia workers on different cluster environments.

src/cluster.jl

@@ -461,12 +461,14 @@ function addprocs(manager::ClusterManager; kwargs...)
 
     cluster_mgmt_from_master_check()
 
-    lock(worker_lock)
-    try
-        addprocs_locked(manager::ClusterManager; kwargs...)
-    finally
-        unlock(worker_lock)
+    new_workers = @lock worker_lock addprocs_locked(manager::ClusterManager; kwargs...)
+    for worker in new_workers
+        for callback in values(worker_added_callbacks)
+            callback(worker)
+        end
     end
+
+    return new_workers
 end
 
 function addprocs_locked(manager::ClusterManager; kwargs...)
@@ -855,13 +857,96 @@ const HDR_COOKIE_LEN=16
 const map_pid_wrkr = Dict{Int, Union{Worker, LocalProcess}}()
 const map_sock_wrkr = IdDict()
 const map_del_wrkr = Set{Int}()
+const worker_added_callbacks = Dict{Any, Base.Callable}()
+const worker_exiting_callbacks = Dict{Any, Base.Callable}()
+const worker_exited_callbacks = Dict{Any, Base.Callable}()
 
 # whether process is a master or worker in a distributed setup
 myrole() = LPROCROLE[]
 function myrole!(proctype::Symbol)
     LPROCROLE[] = proctype
 end
 
+# Callbacks
+
+# We define the callback methods in a loop here and add docstrings for them afterwards
+for callback_type in (:added, :exiting, :exited)
+    let add_name = Symbol(:add_worker_, callback_type, :_callback),
+        remove_name = Symbol(:remove_worker_, callback_type, :_callback),
+        dict_name = Symbol(:worker_, callback_type, :_callbacks)
+
+        @eval begin
+            function $add_name(f::Base.Callable; key=nothing)
+                if !hasmethod(f, Tuple{Int})
+                    throw(ArgumentError("Callback function is invalid, it must be able to accept a single Int argument"))
+                end
+
+                if isnothing(key)
+                    key = Symbol(gensym(), nameof(f))
+                end
+
+                $dict_name[key] = f
+                return key
+            end
+
+            $remove_name(key) = delete!($dict_name, key)
+        end
+    end
+end
+
+"""
+    add_worker_added_callback(f::Base.Callable; key=nothing)
+
+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.
+"""
+function add_worker_added_callback end
+
+"""
+    remove_worker_added_callback(key)
+
+Remove the callback for `key`.
+"""
+function remove_worker_added_callback end
+
+"""
+    add_worker_exiting_callback(f::Base.Callable; key=nothing)
+
+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.
+
+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.
+"""
+function add_worker_exiting_callback end
+
+"""
+    remove_worker_exiting_callback(key)
+
+Remove the callback for `key`.
+"""
+function remove_worker_exiting_callback end
+
+"""
+    add_worker_exited_callback(f::Base.Callable; key=nothing)
+
+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.
+"""
+function add_worker_exited_callback end
+
+"""
+    remove_worker_exited_callback(key)
+
+Remove the callback for `key`.
+"""
+function remove_worker_exited_callback end
+
 # cluster management related API
 """
     myid()
@@ -1048,7 +1133,7 @@ function cluster_mgmt_from_master_check()
 end
 
 """
-    rmprocs(pids...; waitfor=typemax(Int))
+    rmprocs(pids...; waitfor=typemax(Int), callback_timeout=10)
 
 Remove the specified workers. Note that only process 1 can add or remove
 workers.
@@ -1062,6 +1147,10 @@ Argument `waitfor` specifies how long to wait for the workers to shut down:
     returned. The user should call [`wait`](@ref) on the task before invoking any other
     parallel calls.
 
+The `callback_timeout` specifies how long to wait for any callbacks to execute
+before continuing to remove the workers (see
+[`add_worker_exiting_callback()`](@ref)).
+
 # Examples
 ```julia-repl
 \$ julia -p 5
@@ -1078,24 +1167,36 @@ julia> workers()
  6
 ```
 """
-function rmprocs(pids...; waitfor=typemax(Int))
+function rmprocs(pids...; waitfor=typemax(Int), callback_timeout=10)
     cluster_mgmt_from_master_check()
 
     pids = vcat(pids...)
     if waitfor == 0
-        t = @async _rmprocs(pids, typemax(Int))
+        t = @async _rmprocs(pids, typemax(Int), callback_timeout)
         yield()
         return t
     else
-        _rmprocs(pids, waitfor)
+        _rmprocs(pids, waitfor, callback_timeout)
         # return a dummy task object that user code can wait on.
         return @async nothing
     end
 end
 
-function _rmprocs(pids, waitfor)
+function _rmprocs(pids, waitfor, callback_timeout)
     lock(worker_lock)
     try
+        # Run the callbacks
+        callback_tasks = Task[]
+        for pid in pids
+            for callback in values(worker_exiting_callbacks)
+                push!(callback_tasks, 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"
+        end
+
         rmprocset = Union{LocalProcess, Worker}[]
         for p in pids
             if p == 1
@@ -1241,6 +1342,14 @@ function deregister_worker(pg, pid)
             delete!(pg.refs, id)
         end
     end
+
+    # Call callbacks on the master
+    if myid() == 1
+        for callback in values(worker_exited_callbacks)
+            callback(pid)
+        end
+    end
+
     return
 end
 

test/distributed_exec.jl

@@ -1,6 +1,7 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
 using DistributedNext, Random, Serialization, Sockets
+import DistributedNext
 import DistributedNext: launch, manage
 
 
@@ -1934,6 +1935,47 @@ include("splitrange.jl")
     end
 end
 
+@testset "Worker state callbacks" begin
+    if nprocs() > 1
+        rmprocs(workers())
+    end
+
+    # 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))
+    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 added_workers == [pid]
+    rmprocs(workers())
+    @test exiting_workers == [pid]
+    @test exited_workers == [pid]
+
+    # 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
+    event = Base.Event()
+    callback_task = nothing
+    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)
+
+    notify(event)
+    wait(callback_task)
+
+    # Test that the previous callbacks were indeed removed
+    @test length(added_workers) == 1
+    @test length(exiting_workers) == 1
+    @test length(exited_workers) == 1
+end
+
 # Run topology tests last after removing all workers, since a given
 # cluster at any time only supports a single topology.
 if nprocs() > 1