Skip to content

Commit

Permalink
fixup! Add support for worker statuses
Browse files Browse the repository at this point in the history
  • Loading branch information
JamesWrigley committed Jan 16, 2025
1 parent 54b4cf6 commit 8ebd9bf
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 10 deletions.
12 changes: 6 additions & 6 deletions src/cluster.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1014,7 +1014,7 @@ not specified.
The callback will be called with the worker ID, the final
`Distributed.WorkerState` of the worker, and the last status of the worker as
set by [`setstatus`](@ref), e.g. `f(w::Int, state, status)`. `state` is an
set by [`setstatus!`](@ref), e.g. `f(w::Int, state, status)`. `state` is an
enum, a value of `WorkerState_terminated` means a graceful exit and a value of
`WorkerState_exterminated` means the worker died unexpectedly.
Expand Down Expand Up @@ -1210,7 +1210,7 @@ Identical to [`workers()`](@ref) except that the current worker is filtered out.
other_workers() = filter(!=(myid()), workers())

"""
setstatus(x, pid::Int=myid())
setstatus!(x, pid::Int=myid())
Set the status for worker `pid` to `x`. `x` may be any serializable object but
it's recommended to keep it small enough to cheaply send over a network. The
Expand All @@ -1223,22 +1223,22 @@ worker was last doing before it died.
# Examples
```julia-repl
julia> DistributedNext.setstatus("working on dataset 42")
julia> DistributedNext.setstatus!("working on dataset 42")
"working on dataset 42"
julia> DistributedNext.getstatus()
"working on dataset 42"
```
"""
function setstatus(x, pid::Int=myid())
function setstatus!(x, pid::Int=myid())
if pid procs()
throw(ArgumentError("Worker $(pid) does not exist, cannot set its status"))
end

if myid() == 1
@lock map_pid_statuses_lock map_pid_statuses[pid] = x
else
remotecall_fetch(setstatus, 1, x, myid())
remotecall_fetch(setstatus!, 1, x, myid())
end
end

Expand All @@ -1248,7 +1248,7 @@ _getstatus(pid) = @lock map_pid_statuses_lock get!(map_pid_statuses, pid, nothin
getstatus(pid::Int=myid())
Get the status for worker `pid`. If one was never explicitly set with
[`setstatus`](@ref) this will return `nothing`.
[`setstatus!`](@ref) this will return `nothing`.
"""
function getstatus(pid::Int=myid())
if pid procs()
Expand Down
8 changes: 4 additions & 4 deletions test/distributed_exec.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import Revise
using DistributedNext, Random, Serialization, Sockets
import DistributedNext
import DistributedNext: launch, manage, getstatus, setstatus
import DistributedNext: launch, manage, getstatus, setstatus!


@test cluster_cookie() isa String
Expand Down Expand Up @@ -1945,14 +1945,14 @@ end

# Test with the local worker
@test isnothing(getstatus())
setstatus("foo")
setstatus!("foo")
@test getstatus() == "foo"
@test_throws ArgumentError getstatus(2)

# Test with a remote worker
pid = only(addprocs(1))
@test isnothing(getstatus(pid))
remotecall_wait(setstatus, pid, "bar", pid)
remotecall_wait(setstatus!, pid, "bar", pid)
@test remotecall_fetch(getstatus, pid) == "bar"

rmprocs(pid)
Expand Down Expand Up @@ -2021,7 +2021,7 @@ end
last_status = nothing
exited_key = DistributedNext.add_worker_exited_callback((pid, state, status) -> (exit_state = state; last_status = status))
pid = only(addprocs(1))
setstatus("foo", pid)
setstatus!("foo", pid)

redirect_stderr(devnull) do
remote_do(exit, pid)
Expand Down

0 comments on commit 8ebd9bf

Please sign in to comment.