Sch: Improve behavior on internal error

Author: jpsamaroo

src/sch/Sch.jl

@@ -575,12 +575,23 @@ end
 function scheduler_exit(ctx, state::ComputeState, options)
     @dagdebug nothing :global "Tearing down scheduler" uid=state.uid
 
-    close(state.chan)
-    notify(state.halt)
     @sync for p in procs_to_use(ctx)
         @async cleanup_proc(state, p, ctx.log_sink)
     end
 
+    lock(state.lock) do
+        close(state.chan)
+        notify(state.halt)
+
+        # Notify any waiting tasks
+        for (_, futures) in state.futures
+            for future in futures
+                put!(future, SchedulingException("Scheduler exited"); error=true)
+            end
+        end
+        empty!(state.futures)
+    end
+
     # Let the context procs handler clean itself up
     lock(ctx.proc_notify) do
         notify(ctx.proc_notify)

src/sch/eager.jl

@@ -1,5 +1,5 @@
 const EAGER_INIT = Ref{Bool}(false)
-const EAGER_THUNK_CHAN = Channel(typemax(Int))
+const EAGER_THUNK_CHAN = Ref{Channel{Any}}()
 const EAGER_FORCE_KILL = Ref{Bool}(false)
 const EAGER_ID_MAP = Dict{UInt64,Int}()
 const EAGER_CONTEXT = Ref{Union{Context,Nothing}}(nothing)
@@ -15,15 +15,11 @@ end
 function init_eager()
     EAGER_INIT[] && return
     EAGER_INIT[] = true
+    EAGER_THUNK_CHAN[] = Channel(typemax(Int))
     ctx = eager_context()
     @async try
         sopts = SchedulerOptions(;allow_errors=true)
-        scope = Dagger.ExactScope(Dagger.ThreadProc(1, 1))
-        atexit() do
-            EAGER_FORCE_KILL[] = true
-            close(EAGER_THUNK_CHAN)
-        end
-        opts = Dagger.Options((;scope,
+        opts = Dagger.Options((;scope=Dagger.ExactScope(Dagger.ThreadProc(1, 1)),
                                 occupancy=Dict(Dagger.ThreadProc=>0)))
         Dagger.compute(ctx, Dagger.delayed(eager_thunk, opts)();
                        options=sopts)
@@ -37,6 +33,8 @@ function init_eager()
         write(stderr, iob)
     finally
         EAGER_INIT[] = false
+        EAGER_FORCE_KILL[] = true
+        close(EAGER_THUNK_CHAN[])
     end
 end
 
@@ -92,9 +90,10 @@ function eager_thunk()
         nothing
     end
     tls = Dagger.get_tls()
-    while isopen(EAGER_THUNK_CHAN)
+    chan = EAGER_THUNK_CHAN[]
+    while isopen(chan)
         try
-            added_future, future, uid, ref, f, args, opts = take!(EAGER_THUNK_CHAN)
+            added_future, future, uid, ref, f, args, opts = take!(chan)
             # preserve inputs until they enter the scheduler
             tid = GC.@preserve args begin
                 _args = map(args) do pos_x

src/thunk.jl

@@ -313,7 +313,7 @@ function _spawn(f, options::Options, args...)
     finalizer_ref = poolset(EagerThunkFinalizer(uid); device=MemPool.CPURAMDevice())
     added_future = Future()
     propagates = keys(options.options)
-    put!(Dagger.Sch.EAGER_THUNK_CHAN, (added_future, future, uid, finalizer_ref, f, (args...,), (;propagates, options.options...,)))
+    put!(Dagger.Sch.EAGER_THUNK_CHAN[], (added_future, future, uid, finalizer_ref, f, (args...,), (;propagates, options.options...,)))
     thunk_ref = fetch(added_future)
     return (uid, future, finalizer_ref, thunk_ref)
 end