Add support for worker statuses

Author: JamesWrigley

Project.toml

@@ -4,6 +4,7 @@ version = "1.1.1"
 
 [deps]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 
@@ -20,6 +21,7 @@ LibSSH = "0.7"
 LinearAlgebra = "1"
 Random = "1"
 Revise = "3.7.0"
+ScopedValues = "1.6.0"
 Serialization = "1"
 Sockets = "1"
 Test = "1"

docs/src/_changelog.md

@@ -12,6 +12,8 @@ This documents notable changes in DistributedNext.jl. The format is based on
 ### Added
 - Implemented callback support for workers being added/removed etc ([#17]).
 - Added a package extension to support Revise.jl ([#17]).
+- Added support for setting worker statuses with [`setstatus`](@ref) and
+  [`getstatus`](@ref) ([#17]).
 
 ### Fixed
 - Modified the default implementations of methods like `take!` and `wait` on

docs/src/index.md

@@ -14,6 +14,10 @@ DistributedNext.rmprocs
 DistributedNext.interrupt
 DistributedNext.myid
 DistributedNext.pmap
+DistributedNext.getstatus
+DistributedNext.@getstatus
+DistributedNext.setstatus!
+DistributedNext.@setstatus!
 DistributedNext.RemoteException
 DistributedNext.ProcessExitedException
 DistributedNext.Future

src/DistributedNext.jl

@@ -23,6 +23,12 @@ using Serialization, Sockets
 import Serialization: serialize, deserialize
 import Sockets: connect, wait_connected
 
+@static if VERSION < v"1.11"
+    using ScopedValues: ScopedValue, @with
+else
+    using Base.ScopedValues: ScopedValue, @with
+end
+
 # NOTE: clusterserialize.jl imports additional symbols from Serialization for use
 
 export

src/cluster.jl

@@ -891,6 +891,8 @@ const HDR_COOKIE_LEN = 16
 const map_pid_wrkr = Lockable(Dict{Int, Union{Worker, LocalProcess}}())
 const map_sock_wrkr = Lockable(IdDict())
 const map_del_wrkr = Lockable(Set{Int}())
+const _exited_callback_pid = ScopedValue{Int}(-1)
+const map_pid_statuses = Lockable(Dict{Int, Any}())
 const worker_starting_callbacks = Dict{Any, Base.Callable}()
 const worker_started_callbacks = Dict{Any, Base.Callable}()
 const worker_exiting_callbacks = Dict{Any, Base.Callable}()
@@ -1042,9 +1044,9 @@ segfaulting etc). Chooses and returns a unique key for the callback if `key` is
 not specified.
 
 The callback will be called with the worker ID and the final
-`Distributed.WorkerState` of the worker, e.g. `f(w::Int, state)`. `state` is an
-enum, a value of `WorkerState_terminated` means a graceful exit and a value of
-`WorkerState_exterminated` means the worker died unexpectedly.
+`Distributed.WorkerState` of the worker, e.g. `f(w::Int, state)`. `state` is
+an enum, a value of `WorkerState_terminated` means a graceful exit and a value
+of `WorkerState_exterminated` means the worker died unexpectedly.
 
 All worker-exited callbacks will be executed concurrently. If a callback throws
 an exception it will be caught and printed.
@@ -1238,6 +1240,112 @@ Identical to [`workers()`](@ref) except that the current worker is filtered out.
 """
 other_workers() = filter(!=(myid()), workers())
 
+"""
+    @setstatus! x
+    @setstatus! x pid
+
+Set the status for the calling module on worker `pid` (defaults to the current
+worker) to `x`. `x` may be any serializable object but it's recommended to keep
+it small enough to cheaply send over a network. Statuses can be retrieved inside
+worker-exited callbacks (see [`add_worker_exited_callback`](@ref)) before the
+worker is fully deregistered.
+
+Statuses are keyed by the calling `Module`, so multiple libraries can
+independently track their own status on the same worker without conflicting.
+
+This can be handy if you want a way to know what a worker is doing at any given
+time, or (in combination with a worker-exited callback) for knowing what a
+worker was last doing before it died.
+
+# Examples
+```julia-repl
+julia> DistributedNext.@setstatus! "working on dataset 42"
+"working on dataset 42"
+
+julia> DistributedNext.@getstatus
+"working on dataset 42"
+```
+
+See also [`setstatus!`](@ref) for the function form that accepts an explicit module key.
+"""
+macro setstatus!(x)
+    mod = __module__
+    :(setstatus!($(esc(x)), $mod))
+end
+
+macro setstatus!(x, pid)
+    mod = __module__
+    :(setstatus!($(esc(x)), $mod, $(esc(pid))))
+end
+
+"""
+    setstatus!(x, mod::Module, pid::Int=myid())
+
+Function form of [`@setstatus!`](@ref). Sets the status for module `mod` on
+worker `pid` to `x`.
+"""
+function setstatus!(x, mod::Module, pid::Int=myid())
+    if !id_in_procs(pid)
+        throw(ArgumentError("Worker $(pid) does not exist, cannot set its status"))
+    end
+
+    if myid() == 1
+        @lock map_pid_statuses begin
+            statuses = get!(map_pid_statuses[], pid, Dict{Module, Any}())
+            statuses[mod] = x
+        end
+    else
+        remotecall_fetch(setstatus!, 1, x, mod, myid())
+    end
+end
+
+function _getstatus(pid, mod)
+    @lock map_pid_statuses begin
+        statuses = get(map_pid_statuses[], pid, nothing)
+        isnothing(statuses) ? nothing : get(statuses, mod, nothing)
+    end
+end
+
+"""
+    @getstatus
+    @getstatus pid
+
+Get the status set by the calling module for worker `pid` (defaults to the
+current worker). If one was never explicitly set with [`@setstatus!`](@ref)
+this will return `nothing`.
+
+See also [`getstatus`](@ref) for the function form.
+"""
+macro getstatus()
+    mod = __module__
+    :(getstatus($mod))
+end
+macro getstatus(pid)
+    mod = __module__
+    :(getstatus($mod, $(esc(pid))))
+end
+
+"""
+    getstatus(mod::Module, pid::Int=myid())
+
+Function form of [`@getstatus`](@ref). Gets the status for module `mod` on
+worker `pid`. Returns `nothing` if no status was set.
+"""
+function getstatus(mod::Module, pid::Int=myid())
+    # During the worker-exited callbacks this function may be called, at which
+    # point it will not exist in procs(). Thus we check whether the function is
+    # being called for an exited worker and allow it if so.
+    if !id_in_procs(pid) && _exited_callback_pid[] != pid
+        throw(ArgumentError("Worker $(pid) does not exist, cannot get its status"))
+    end
+
+    if myid() == 1
+        _getstatus(pid, mod)
+    else
+        remotecall_fetch(getstatus, 1, mod, pid)
+    end
+end
+
 function cluster_mgmt_from_master_check()
     if myid() != 1
         throw(ErrorException("Only process 1 can add and remove workers"))
@@ -1463,15 +1571,22 @@ function deregister_worker(pg, pid)
         end
     end
 
-    # Call callbacks on the master
     if myid() == 1
-        for (name, callback) in worker_exited_callbacks
-            try
-                callback(pid, w.state)
-            catch ex
-                @error "Error when running worker-exited callback '$(name)'" exception=(ex, catch_backtrace())
-            end
+        params = default_addprocs_params(w.manager)
+        warning_interval = params[:callback_warning_interval]
+
+        # Call callbacks on the master, with the scoped value set so that
+        # getstatus() can be called for the exiting worker without failing the
+        # pid check. We go to some effort to make sure this works after
+        # deregistering the worker because if it's called beforehand the worker
+        # will incorrectly be shown in e.g. procs().
+        @with _exited_callback_pid => pid begin
+            _run_callbacks_concurrently("worker-exited", worker_exited_callbacks,
+                                        warning_interval, [(pid, w.state)]; catch_exceptions=true)
         end
+
+        # Delete its statuses
+        @lock map_pid_statuses delete!(map_pid_statuses[], pid)
     end
 
     return

test/distributed_exec.jl

@@ -3,7 +3,7 @@
 import Revise
 using DistributedNext, Random, Serialization, Sockets
 import DistributedNext
-import DistributedNext: launch, manage
+import DistributedNext: launch, manage, getstatus, setstatus!, @getstatus, @setstatus!
 
 
 @test cluster_cookie() isa String
@@ -1826,7 +1826,9 @@ end
     let julia = `$(Base.julia_cmd()) --startup-file=no`; mktempdir() do tmp
         pkg_project = joinpath(Base.pkgdir(DistributedNext), "Project.toml")
         project = mkdir(joinpath(tmp, "project"))
-        depots = [mkdir(joinpath(tmp, "depot1")), mkdir(joinpath(tmp, "depot2"))]
+        # Keep the writable depot in the depots list so that external
+        # dependencies (i.e. ScopedValues.jl) can be loaded.
+        depots = [mkdir(joinpath(tmp, "depot1")), mkdir(joinpath(tmp, "depot2")), Base.DEPOT_PATH[1]]
         load_path = [mkdir(joinpath(tmp, "load_path")), "@stdlib", "@", pkg_project]
         pathsep = Sys.iswindows() ? ";" : ":"
         env = Dict(
@@ -1935,15 +1937,15 @@ end
         project = mktempdir()
         env = Dict(
             "JULIA_LOAD_PATH" => string(LOAD_PATH[1], $(repr(pathsep)), "@stdlib", $(repr(pathsep)), "$(escaped_pkg_project)"),
-            "JULIA_DEPOT_PATH" => DEPOT_PATH[1],
+            "JULIA_DEPOT_PATH" => DEPOT_PATH[end],
             "TMPDIR" => ENV["TMPDIR"],
         )
         addprocs(1; env = env, exeflags = `--project=\$(project)`)
         env["JULIA_PROJECT"] = project
         addprocs(1; env = env)
         """ * funcscode * """
         for w in workers()
-            @test remotecall_fetch(depot_path, w)          == [DEPOT_PATH[1]]
+            @test remotecall_fetch(depot_path, w)          == [DEPOT_PATH[end]]
             @test remotecall_fetch(load_path, w)           == [LOAD_PATH[1], "@stdlib", "$(escaped_pkg_project)"]
             @test remotecall_fetch(active_project, w)      == project
             @test remotecall_fetch(Base.active_project, w) == joinpath(project, "Project.toml")
@@ -1983,7 +1985,40 @@ include("splitrange.jl")
     end
 end
 
+@testset "Worker statuses" begin
+    rmprocs(other_workers())
+
+    # Test with the local worker using macros
+    @test isnothing(@getstatus())
+    @setstatus!("foo")
+    @test @getstatus() == "foo"
+    @test_throws ArgumentError getstatus(Main, 2)
+
+    # Test with a remote worker using the function form
+    pid = only(addprocs(1))
+    @test isnothing(getstatus(Main, pid))
+    remotecall_wait(setstatus!, pid, "bar", Main, pid)
+    @test remotecall_fetch(getstatus, pid, Main) == "bar"
+
+    # Test that different modules have independent statuses
+    setstatus!("from_main", Main, pid)
+    setstatus!("from_distributed", DistributedNext, pid)
+    @test getstatus(Main, pid) == "from_main"
+    @test getstatus(DistributedNext, pid) == "from_distributed"
+
+    rmprocs(pid)
+end
+
 @testset "Worker state callbacks" begin
+    # Helper function to wait for a worker to have been completely deregistered
+    # (including worker-exited callbacks finished) by waiting for the workers
+    # status to have been deleted. Only works if the worker has a status of
+    # course.
+    function wait_for_deregistration(pid)
+        statuses = DistributedNext.map_pid_statuses
+        @test timedwait(() -> @lock(statuses, !haskey(statuses[], pid)), 10) == :ok
+    end
+
     rmprocs(other_workers())
 
     # Adding a callback with an invalid signature should fail
@@ -2040,16 +2075,34 @@ end
     @test length(exiting_workers) == 1
     @test length(exited_workers) == 1
 
-    # Test that workers that were killed forcefully are detected as such
+    # Test that workers that were killed forcefully are detected as such, and
+    # that statuses can be retrieved in the callback.
     exit_state = nothing
-    exited_key = DistributedNext.add_worker_exited_callback((pid, state) -> exit_state = state)
+    last_status = nothing
+    exited_key = DistributedNext.add_worker_exited_callback((pid, state) -> (exit_state = state; last_status = @getstatus(pid)))
     pid = only(addprocs(1))
+    @setstatus!("foo", pid)
 
+    # Kill the process with stderr redirected so the error messages don't
+    # unnecessarily show up in the logs.
     redirect_stderr(devnull) do
         remote_do(exit, pid)
-        timedwait(() -> !isnothing(exit_state), 10)
+        wait_for_deregistration(pid)
     end
     @test exit_state == DistributedNext.WorkerState_exterminated
+    @test last_status == "foo"
+    DistributedNext.remove_worker_exited_callback(exited_key)
+
+    # Test that exceptions in worker-exited callbacks are caught
+    exited_key = DistributedNext.add_worker_exited_callback((pid, state) -> error("foo"))
+    @test_logs (:error, r"Error when running worker-exited callback.+") match_mode=:any begin
+        pid = only(addprocs(1))
+        # Set a dummy status so that wait_for_deregistration() works
+        @setstatus!("foo", pid)
+        rmprocs(pid)
+
+        wait_for_deregistration(pid)
+    end
     DistributedNext.remove_worker_exited_callback(exited_key)
 end
 

test/runtests.jl

@@ -35,5 +35,5 @@ include("managers.jl")
 include("distributed_stdlib_detection.jl")
 
 @testset "Aqua" begin
-    Aqua.test_all(DistributedNext)
+    Aqua.test_all(DistributedNext; stale_deps=(; ignore=[:ScopedValues]))
 end