Add more Sch debug statements

And some fixes

Author: jpsamaroo

src/sch/Sch.jl

@@ -555,10 +555,12 @@ end
             @goto fire_tasks
         end
         task = popfirst!(state.ready)
+        @dagdebug task :schedule "Scheduling task"
         @maybelog ctx timespan_start(ctx, :schedule, (;uid=state.uid, thunk_id=task.id), (;thunk_id=task.id))
         if has_result(state, task)
             if haskey(state.errored, task)
                 # An error was eagerly propagated to this task
+                @dagdebug task :schedule "Task received upstream error, finishing"
                 finish_failed!(state, task)
             else
                 # This shouldn't have happened
@@ -663,9 +665,11 @@ end
                 end
             end
         end
+
         ex = SchedulingException("No processors available, try widening scope")
         store_result!(state, task, ex; error=true)
         set_failed!(state, task)
+        @dagdebug task :schedule "No processors available, skipping"
         sorted_procs_cleanup()
         costs_cleanup()
         @goto pop_task
@@ -732,6 +736,7 @@ function remove_dead_proc!(ctx, state, proc, options)
 end
 
 function finish_task!(ctx, state, node, thunk_failed)
+    @dagdebug node :finish "Finishing with $(thunk_failed ? "error" : "result")"
     pop!(state.running, node)
     delete!(state.running_on, node)
     if thunk_failed
@@ -1224,7 +1229,7 @@ function (dts::DoTaskSpec)()
         if unwrap_nested_exception(err) isa InvalidStateException || !isopen(return_queue)
             @dagdebug tid :execute "Return queue is closed, failing to put result" chan=return_queue exception=(err, catch_backtrace())
         else
-            rethrow(err)
+            rethrow()
         end
     end
     return

src/sch/util.jl

@@ -93,6 +93,7 @@ end
 function fill_registered_futures!(state, thunk, failed)
     if haskey(state.futures, thunk)
         # Notify any listening thunks
+        @dagdebug thunk :finish "Notifying $(length(state.futures[thunk])) futures"
         for future in state.futures[thunk]
             put!(future, load_result(state, thunk); error=failed)
         end
@@ -135,14 +136,20 @@ end
 
 "Schedules any dependents that may be ready to execute."
 function schedule_dependents!(state, thunk, failed)
+    @dagdebug thunk :finish "Checking dependents"
     if !haskey(state.waiting_data, thunk) || isempty(state.waiting_data[thunk])
         return
     end
+    ctr = 0
     for dep in state.waiting_data[thunk]
+        @dagdebug dep :schedule "Checking dependent"
         dep_isready = false
         if haskey(state.waiting, dep)
             set = state.waiting[dep]
             thunk in set && pop!(set, thunk)
+            if length(set) > 0
+                @dagdebug dep :schedule "Dependent has $(length(set)) upstreams"
+            end
             dep_isready = isempty(set)
             if dep_isready
                 delete!(state.waiting, dep)
@@ -151,11 +158,17 @@ function schedule_dependents!(state, thunk, failed)
             dep_isready = true
         end
         if dep_isready
+            ctr += 1
             if !failed
                 push!(state.ready, dep)
+                @dagdebug dep :schedule "Dependent is now ready"
+            else
+                set_failed!(state, thunk, dep)
+                @dagdebug dep :schedule "Dependent has transitively failed"
             end
         end
     end
+    @dagdebug thunk :finish "Marked $ctr dependents as $(failed ? "failed" : "ready")"
 end
 
 """
@@ -226,6 +239,8 @@ const RESCHEDULE_SYNCDEPS_SEEN_CACHE = TaskLocalValue{ReusableCache{Set{Thunk},N
 "Marks `thunk` and all dependent thunks as failed."
 function set_failed!(state, origin, thunk=origin)
     @assert islocked(state.lock)
+    has_result(state, thunk) && return
+    @dagdebug thunk :finish "Setting as failed"
     filter!(x -> x !== thunk, state.ready)
     # N.B. If origin === thunk, we assume that the caller has already set the error
     if origin !== thunk
@@ -554,6 +569,7 @@ end
         end
     end
 
+    # Estimate total cost for executing this task on each candidate processor
     for proc in procs
         gproc = get_parent(proc)
         chunks_filt = Iterators.filter(c->get_parent(processor(c)) != gproc, chunks)

src/submission.jl

@@ -185,7 +185,8 @@ const UID_TO_TID_CACHE = TaskLocalValue{ReusableCache{Dict{UInt64,Int},Nothing}}
                 #=FIXME:REALLOC=#
                 Sch.reschedule_syncdeps!(state, thunk)
                 old_fargs_cleanup() # reschedule_syncdeps! preserves all referenced tasks/chunks
-                @dagdebug thunk :submit "Added to scheduler"
+                n_upstreams = haskey(state.waiting, thunk) ? length(state.waiting[thunk]) : 0
+                @dagdebug thunk :submit "Added to scheduler with $n_upstreams upstreams"
                 if future !== nothing
                     # Ensure we attach a future before the thunk is scheduled
                     Sch._register_future!(ctx, state, task, tid, (future, thunk_id, false))

src/utils/dagdebug.jl

@@ -2,8 +2,8 @@ function istask end
 function task_id end
 
 const DAGDEBUG_CATEGORIES = Symbol[:global, :submit, :schedule, :scope,
-                                   :take, :execute, :move, :processor, :cancel,
-                                   :stream]
+                                   :take, :execute, :move, :processor, :finish,
+                                   :cancel, :stream]
 macro dagdebug(thunk, category, msg, args...)
     cat_sym = category.value
     @gensym id

src/utils/scopes.jl

@@ -26,7 +26,7 @@ compatible with the given scope.
 """
 compatible_processors(scope::AbstractScope=get_compute_scope(), ctx::Context=Sch.eager_context()) =
     compatible_processors(scope, procs(ctx))
-function compatible_processors(scope::AbstractScope, procs::Vector{Processor})
+function compatible_processors(scope::AbstractScope, procs::Vector{<:Processor})
     compat_procs = Set{Processor}()
     for gproc in procs
         # Fast-path in case entire process is incompatible

test/util.jl

@@ -14,7 +14,7 @@ end
 replace_obj!(ex::Symbol, obj) = Expr(:(.), obj, QuoteNode(ex))
 replace_obj!(ex, obj) = ex
 function _test_throws_unwrap(terr, ex; to_match=[])
-    @gensym oerr rerr
+    @gensym oerr rerr bt
     match_expr = Expr(:block)
     for m in to_match
         if m.head == :(=)
@@ -35,9 +35,11 @@ function _test_throws_unwrap(terr, ex; to_match=[])
         end
     end
     quote
+        $bt = nothing
         $oerr, $rerr = try
             nothing, $(esc(ex))
         catch err
+            $bt = catch_backtrace()
             (err, Dagger.Sch.unwrap_nested_exception(err))
         end
         if $terr isa Tuple
@@ -48,7 +50,7 @@ function _test_throws_unwrap(terr, ex; to_match=[])
             else
                 println("Full error:")
                 Base.showerror(stdout, $oerr)
-                Base.show_backtrace(stdout, backtrace())
+                Base.show_backtrace(stdout, $bt)
             end
         else
             @test $rerr isa $terr
@@ -57,7 +59,7 @@ function _test_throws_unwrap(terr, ex; to_match=[])
             else
                 println("Full error:")
                 Base.showerror(stdout, $oerr)
-                Base.show_backtrace(stdout, backtrace())
+                Base.show_backtrace(stdout, $bt)
             end
         end
     end