Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 105 additions & 43 deletions src/cluster.jl
Original file line number Diff line number Diff line change
Expand Up @@ -103,14 +103,15 @@ end
@enum WorkerState W_CREATED W_CONNECTED W_TERMINATING W_TERMINATED W_UNKNOWN_STATE
mutable struct Worker
id::Int
msg_lock::Threads.ReentrantLock # Lock for del_msgs, add_msgs, and gcflag
msg_lock::ReentrantLock # guards `del_msgs`, `add_msgs`, and `gcflag`
del_msgs::Array{Any,1} # XXX: Could del_msgs and add_msgs be Channels?
add_msgs::Array{Any,1}
@atomic gcflag::Bool
# `state` changes are published via the module-global `worker_state_cond`;
# wait on that (re-checking this atomic) rather than a per-worker condition
@atomic state::WorkerState
c_state::Threads.Condition # wait for state changes, lock for state
ct_time::Float64 # creation time
conn_func::Any # used to setup connections lazily
ct_time::Float64 # creation time
conn_func::Any # used to setup connections lazily

r_stream::IO
w_stream::IO
Expand Down Expand Up @@ -142,7 +143,7 @@ mutable struct Worker
if haskey(map_pid_wrkr, id)
return map_pid_wrkr[id]
end
w=new(id, Threads.ReentrantLock(), [], [], false, W_CREATED, Threads.Condition(), time(), conn_func)
w=new(id, ReentrantLock(), [], [], false, W_CREATED, time(), conn_func)
w.initialized = Event()
register_worker(w)
w
Expand All @@ -152,10 +153,10 @@ mutable struct Worker
end

function set_worker_state(w, state)
lock(w.c_state) do
@atomic w.state = state
notify(w.c_state; all=true)
end
@atomic w.state = state
# Wake every task waiting on worker state/membership: `wait_for_conn`,
# topology setup in `create_worker`, `_rmprocs`, `check_master_connect`, etc.
@lock worker_state_cond notify(worker_state_cond; all=true)
end

function check_worker_state(w::Worker)
Expand Down Expand Up @@ -202,10 +203,14 @@ function wait_for_conn(w)
timeout = worker_timeout() - (time() - w.ct_time)
timeout <= 0 && error("peer $(w.id) has not connected to $(myid())")

if timedwait(() -> (@atomic w.state) === W_CONNECTED, timeout) === :timed_out
# Notify any waiters on the state and throw
@lock w.c_state notify(w.c_state)
error("peer $(w.id) didn't connect to $(myid()) within $timeout seconds")
deadline = time_ns() + timeout * 1e9
@lock worker_state_cond begin
while (@atomic w.state) !== W_CONNECTED
wait_or_deadline(worker_state_cond, deadline) || break
end
if (@atomic w.state) !== W_CONNECTED
error("peer $(w.id) didn't connect to $(myid()) within $(worker_timeout()) seconds")
end
end
end
nothing
Expand All @@ -223,6 +228,32 @@ end

worker_timeout() = parse(Float64, get(ENV, "JULIA_WORKER_TIMEOUT", "60.0"))

# Wait for `notify` on the `Threads.Condition`. Triggers its own notification
# when the absolute `deadline` (nanoseconds on the `time_ns()` clock; `Inf`
# waits forever) arrives.
#
# Return `false` once the deadline has passed (so callers can
# `wait_or_deadline(...) || break`) and `true` otherwise. The deadline causes
# spurious wakeups for other waiters, so callers must re-check their own
# condition in a loop.
function wait_or_deadline(c::Threads.Condition, deadline::Real)
remaining = (deadline - time_ns()) / 1e9
if !isfinite(remaining)
wait(c)
return true
end
remaining <= 0 && return false
timer = Timer(remaining) do _
@lock c notify(c; all=true)
end
try
wait(c)
finally
close(timer)
end
return true
end


## worker creation and setup ##
"""
Expand Down Expand Up @@ -323,7 +354,7 @@ function read_worker_host_port(io::IO)

# Wait at most for JULIA_WORKER_TIMEOUT seconds to read host:port
# info from the worker
timeout = worker_timeout() * 1e9
deadline = t0 + worker_timeout() * 1e9
# We expect the first line to contain the host:port string. However, as
# the worker may be launched via ssh or a cluster manager like SLURM,
# ignore any informational / warning lines printed by the launch command.
Expand All @@ -332,12 +363,20 @@ function read_worker_host_port(io::IO)

ntries = 1000
leader = String[]
# notified when a spawned `readline` task finishes, so we can wait for a line
# to arrive (or the overall timeout to elapse) without polling
readcond = Threads.Condition()
try
while ntries > 0
readtask = @async readline(io)
yield()
while !istaskdone(readtask) && ((time_ns() - t0) < timeout)
sleep(0.05)
readtask = @async try
readline(io)
finally
@lock readcond notify(readcond; all=true)
end
@lock readcond begin
while !istaskdone(readtask)
wait_or_deadline(readcond, deadline) || break
end
end
!istaskdone(readtask) && break

Expand Down Expand Up @@ -665,9 +704,9 @@ function create_worker(manager::ClusterManager, wconfig::WorkerConfig)
for jw in PGRP.workers
if (jw.id != 1) && (jw.id < w.id)
# wait for wl to join
if (@atomic jw.state) === W_CREATED

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Access to a threaded value outside of a lock

lock(jw.c_state) do
wait(jw.c_state)
@lock worker_state_cond begin
while (@atomic jw.state) === W_CREATED
wait(worker_state_cond)
end
end
push!(join_list, jw)
Expand All @@ -679,23 +718,27 @@ function create_worker(manager::ClusterManager, wconfig::WorkerConfig)
filterfunc(x) = (x.id != 1) && isdefined(x, :config) &&
(notnothing(x.config.ident) in something(wconfig.connect_idents, []))

wlist = filter(filterfunc, PGRP.workers)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Access to a threaded value outside of a lock

waittime = 0
while wconfig.connect_idents !== nothing &&
length(wlist) < length(wconfig.connect_idents)
if waittime >= timeout
error("peer workers did not connect within $timeout seconds")
deadline = time_ns() + timeout * 1e9
# Recompute `wlist` under the lock before each wait so we never miss a
# `register_worker` notification that races with the filter (a lost
# wakeup would stall this until the deadline).
local wlist
@lock worker_state_cond begin
while true
wlist = filter(filterfunc, PGRP.workers)
(wconfig.connect_idents === nothing ||
length(wlist) >= length(wconfig.connect_idents)) && break
if !wait_or_deadline(worker_state_cond, deadline)
error("peer workers did not connect within $timeout seconds")
end
end
sleep(1.0)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unnecessary delay for more than 1 ms (for :custom topology)

waittime += 1
wlist = filter(filterfunc, PGRP.workers)
end

for wl in wlist
lock(wl.c_state) do
if (@atomic wl.state) === W_CREATED
@lock worker_state_cond begin
while (@atomic wl.state) === W_CREATED
# wait for wl to join
wait(wl.c_state)
wait(worker_state_cond)
end
end
push!(join_list, wl)
Expand Down Expand Up @@ -761,12 +804,18 @@ function check_master_connect()
if ccall(:jl_running_on_valgrind,Cint,()) != 0
return
end

errormonitor(
@async begin
timeout = worker_timeout()
if timedwait(() -> haskey(map_pid_wrkr, 1), timeout) === :timed_out

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Access to threaded value outside of a lock

print(stderr, "Master process (id 1) could not connect within $(timeout) seconds.\nexiting.\n")
timeout = worker_timeout() * 1e9
deadline = time_ns() + timeout
@lock worker_state_cond begin
while !haskey(map_pid_wrkr, 1)
wait_or_deadline(worker_state_cond, deadline) || break
end
end

if !haskey(map_pid_wrkr, 1)
print(stderr, "Master process (id 1) could not connect within $(timeout/1e9) seconds.\nexiting.\n")
exit(1)
end
end
Expand Down Expand Up @@ -848,6 +897,11 @@ const LPROCROLE = Ref{Symbol}(:master)
const HDR_VERSION_LEN=16
const HDR_COOKIE_LEN=16
const map_pid_wrkr = Dict{Int, Union{Worker, LocalProcess}}()
# Notified whenever a worker is registered (added to `map_pid_wrkr`/`PGRP.workers`)
# or changes state (via `set_worker_state`), so tasks can wait for peers or the
# master to join the cluster. The lock guards reads/writes of the membership
# collections done under it.
const worker_state_cond = Threads.Condition()
const map_sock_wrkr = IdDict()
const map_del_wrkr = Set{Int}()

Expand Down Expand Up @@ -1082,10 +1136,15 @@ function _rmprocs(pids, waitfor)
end
end

start = time_ns()
while (time_ns() - start) < waitfor*1e9
all(w -> (@atomic w.state) === W_TERMINATED, rmprocset) && break

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

O(n^2) loop when O(n) condition was sufficient

sleep(min(0.1, waitfor - (time_ns() - start)/1e9))
# `waitfor == typemax(Int)` means wait indefinitely; otherwise bound the
# total wait across all workers by a single shared deadline
deadline = waitfor == typemax(Int) ? Inf : time_ns() + waitfor * 1e9
for w in rmprocset
@lock worker_state_cond begin
while (@atomic w.state) !== W_TERMINATED
wait_or_deadline(worker_state_cond, deadline) || break
end
end
end

unremoved = [wrkr.id for wrkr in filter(w -> (@atomic w.state) !== W_TERMINATED, rmprocset)]
Expand Down Expand Up @@ -1155,8 +1214,11 @@ end

register_worker(w) = register_worker(PGRP, w)
function register_worker(pg, w)
push!(pg.workers, w)
map_pid_wrkr[w.id] = w
@lock worker_state_cond begin

@vtjnash vtjnash Jul 15, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Failure to lock before changing state

push!(pg.workers, w)
map_pid_wrkr[w.id] = w
notify(worker_state_cond; all=true)
end
end

function register_worker_streams(w)
Expand Down
4 changes: 2 additions & 2 deletions src/remotecall.jl
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ function send_del_client_no_lock(rr)
end

function publish_del_msg!(w::Worker, msg)
lock(w.msg_lock) do
@lock w.msg_lock begin
push!(w.del_msgs, msg)
@atomic w.gcflag = true
end
Expand Down Expand Up @@ -354,7 +354,7 @@ function send_add_client(rr::AbstractRemoteRef, i)
# to the processor that owns the remote ref. it will add_client
# itself inside deserialize().
w = worker_from_id(rr.where)
lock(w.msg_lock) do
@lock w.msg_lock begin
push!(w.add_msgs, (remoteref_id(rr), i))
@atomic w.gcflag = true
end
Expand Down
Loading