Move support for addprocs!()/rmprocs!() exclusively to the eager API

This ends up being simpler than supporting both.

Author: JamesWrigley

Diff for: src/context.jl

@@ -1,3 +1,5 @@
+export Context, addprocs!, rmprocs!
+
 """
     Context(xs::Vector{OSProc}) -> Context
     Context(xs::Vector{Int}) -> Context
@@ -15,21 +17,18 @@ Special fields include:
 mutable struct Context
     procs::Vector{Processor}
     proc_lock::ReentrantLock
-    proc_notify::Threads.Condition
     log_sink::Any
     profile::Bool
     options
 end
 
 function Context(procs::Vector{P}=Processor[OSProc(w) for w in procs()];
-        proc_lock=ReentrantLock(), proc_notify=Threads.Condition(),
-        log_sink=TimespanLogging.NoOpLog(), log_file=nothing, profile=false,
-        options=nothing) where {P<:Processor}
+                 proc_lock=ReentrantLock(), log_sink=TimespanLogging.NoOpLog(),
+                 log_file=nothing, profile=false, options=nothing) where {P<:Processor}
     if log_file !== nothing
         @warn "`log_file` is no longer supported\nPlease instead load `GraphViz.jl` and use `render_logs(logs, :graphviz)`."
     end
-    Context(procs, proc_lock, proc_notify, log_sink,
-            profile, options)
+    Context(procs, proc_lock, log_sink, profile, options)
 end
 Context(xs::Vector{Int}; kwargs...) = Context(map(OSProc, xs); kwargs...)
 Context(ctx::Context, xs::Vector=copy(procs(ctx))) = # make a copy
@@ -62,41 +61,90 @@ procs(ctx::Context) = lock(ctx) do
 end
 
 """
-    addprocs!(ctx::Context, xs)
+    addprocs!(xs) -> Processor[]
 
-Add new workers `xs` to `ctx`.
+Add new workers `xs` to the eager scheduler and returns the ones that were
+actually added.
 
 Workers will typically be assigned new tasks in the next scheduling iteration
 if scheduling is ongoing.
 
 Workers can be either `Processor`s or the underlying process IDs as `Integer`s.
 """
-addprocs!(ctx::Context, xs::AbstractVector{<:Integer}) = addprocs!(ctx, map(OSProc, xs))
-function addprocs!(ctx::Context, xs::AbstractVector{<:OSProc})
-    lock(ctx) do
-        append!(ctx.procs, xs)
+addprocs!(xs::AbstractVector{<:Integer}) = addprocs!(map(OSProc, xs))
+function addprocs!(xs::AbstractVector{<:Processor})
+    Sch.init_eager()
+
+    ctx = Sch.eager_context()
+    state = Sch.EAGER_STATE[]
+
+    to_add = setdiff(xs, procs(ctx))
+
+    timespan_start(ctx, :addprocs!, nothing, nothing)
+
+    # Initialize new procs
+    for p in to_add
+        Sch.init_proc(state, p, ctx.log_sink)
+
+        # Empty the processor cache list and force reschedule
+        lock(state.lock) do
+            state.procs_cache_list[] = nothing
+        end
+        put!(state.chan, Sch.RescheduleSignal())
     end
-    lock(ctx.proc_notify) do
-        notify(ctx.proc_notify)
+
+    lock(ctx) do
+        append!(ctx.procs, to_add)
     end
+    timespan_finish(ctx, :addprocs!, nothing, nothing)
+
+    return to_add
 end
 
 """
-    rmprocs!(ctx::Context, xs)
+    rmprocs!(xs) -> Processor[]
 
-Remove the specified workers `xs` from `ctx`.
+Remove the specified workers `xs` from the eager scheduler and returns the ones
+that were actuall added.
 
 Workers will typically finish all their assigned tasks if scheduling is ongoing
 but will not be assigned new tasks after removal.
 
 Workers can be either `Processor`s or the underlying process IDs as `Integer`s.
 """
-rmprocs!(ctx::Context, xs::AbstractVector{<:Integer}) = rmprocs!(ctx, map(OSProc, xs))
-function rmprocs!(ctx::Context, xs::AbstractVector{<:OSProc})
-    lock(ctx) do
-        filter!(p -> (p ∉ xs), ctx.procs)
+rmprocs!(xs::AbstractVector{<:Integer}) = rmprocs!(map(OSProc, xs))
+function rmprocs!(xs::AbstractVector{<:Processor})
+    Sch.init_eager()
+
+    ctx = Sch.eager_context()
+    state = Sch.EAGER_STATE[]
+
+    _rmprocs!(ctx, state, xs; allow_empty=false)
+end
+
+function _rmprocs!(ctx::Context, state, xs; allow_empty=true)
+
+    if !allow_empty && Set(procs(ctx)) == Set(xs)
+        throw(ArgumentError("Refusing to remove all processors from the Dagger.Context, at least one must be present."))
+    end
+
+    to_remove = intersect(xs, procs(ctx))
+
+    timespan_start(ctx, :rmprocs!, nothing, nothing)
+
+    for p in to_remove
+        Sch.cleanup_proc(state, p, ctx.log_sink)
+
+        # Empty the processor cache list
+        lock(state.lock) do
+            state.procs_cache_list[] = nothing
+        end
     end
-    lock(ctx.proc_notify) do
-        notify(ctx.proc_notify)
+
+    lock(ctx) do
+        filter!(p -> (p ∉ to_remove), ctx.procs)
     end
+    timespan_finish(ctx, :rmprocs!, nothing, nothing)
+
+    return to_remove
 end

Diff for: src/processor.jl

@@ -1,4 +1,4 @@
-export OSProc, Context, addprocs!, rmprocs!
+export OSProc
 
 import Base: @invokelatest
 

Diff for: src/sch/Sch.jl

@@ -415,7 +415,7 @@ function cleanup_proc(state, p, log_sink)
             # We allow ProcessExitedException's, which means that the worker
             # shutdown halfway through cleanup.
             if !(ex isa ProcessExitedException)
-                throw(ex)
+                rethrow()
             end
         end
     end
@@ -515,15 +515,6 @@ function scheduler_init(ctx, state::ComputeState, d::Thunk, options, deps)
             end
         end
     end
-
-    # Listen for new workers
-    @async begin
-        try
-            monitor_procs_changed!(ctx, state)
-        catch err
-            @error "Error assigning workers" exception=(err,catch_backtrace())
-        end
-    end
 end
 
 function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
@@ -643,11 +634,6 @@ function scheduler_exit(ctx, state::ComputeState, options)
         end
         empty!(state.futures)
     end
-
-    # Let the context procs handler clean itself up
-    lock(ctx.proc_notify) do
-        notify(ctx.proc_notify)
-    end
 end
 
 function procs_to_use(ctx, options=ctx.options)
@@ -865,56 +851,9 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
     end
 end
 
-"""
-Monitors for workers being added/removed to/from `ctx`, sets up or tears down
-per-worker state, and notifies the scheduler so that work can be reassigned.
-"""
-function monitor_procs_changed!(ctx, state)
-    # Load current set of procs
-    old_ps = procs_to_use(ctx)
-
-    while !state.halt.set
-        # Wait for the notification that procs have changed
-        lock(ctx.proc_notify) do
-            wait(ctx.proc_notify)
-        end
-
-        timespan_start(ctx, :assign_procs, nothing, nothing)
-
-        # Load new set of procs
-        new_ps = procs_to_use(ctx)
-
-        # Initialize new procs
-        diffps = setdiff(new_ps, old_ps)
-        for p in diffps
-            init_proc(state, p, ctx.log_sink)
-
-            # Empty the processor cache list and force reschedule
-            lock(state.lock) do
-                state.procs_cache_list[] = nothing
-            end
-            put!(state.chan, RescheduleSignal())
-        end
-
-        # Cleanup removed procs
-        diffps = setdiff(old_ps, new_ps)
-        for p in diffps
-            cleanup_proc(state, p, ctx.log_sink)
-
-            # Empty the processor cache list
-            lock(state.lock) do
-                state.procs_cache_list[] = nothing
-            end
-        end
-
-        timespan_finish(ctx, :assign_procs, nothing, nothing)
-        old_ps = new_ps
-    end
-end
-
 function remove_dead_proc!(ctx, state, proc, options=ctx.options)
     @assert options.single !== proc.pid "Single worker failed, cannot continue."
-    rmprocs!(ctx, [proc])
+    Dagger._rmprocs!(ctx, state, [proc])
     delete!(state.worker_time_pressure, proc.pid)
     delete!(state.worker_storage_pressure, proc.pid)
     delete!(state.worker_storage_capacity, proc.pid)

Diff for: test/processors.jl

@@ -61,21 +61,6 @@ end
         addcb()
     end
 
-    @testset "Modify workers in Context" begin
-        ps = addprocs(4, exeflags="--project")
-        @everywhere ps using Dagger
-
-        ctx = Context(ps[1:2])
-
-        Dagger.addprocs!(ctx, ps[3:end])
-        @test map(p -> p.pid, procs(ctx)) == ps
-
-        Dagger.rmprocs!(ctx, ps[3:end])
-        @test map(p -> p.pid, procs(ctx)) == ps[1:2]
-
-        wait(rmprocs(ps))
-    end
-
     @testset "Callable as Thunk function" begin
         @everywhere begin
             struct ABC end