Sch: Add at-dagdebug helper macro

Author: jpsamaroo

src/sch/Sch.jl

@@ -467,6 +467,8 @@ function scheduler_init(ctx, state::ComputeState, d::Thunk, options, deps)
 end
 
 function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
+    @dagdebug nothing :global "Initializing scheduler" uid=state.uid
+
     safepoint(state)
 
     # Loop while we still have thunks to execute
@@ -480,12 +482,14 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
 
         isempty(state.running) && continue
         timespan_start(ctx, :take, 0, 0)
+        @dagdebug nothing :take "Waiting for results"
         chan_value = take!(state.chan) # get result of completed thunk
         timespan_finish(ctx, :take, 0, 0)
         if chan_value isa RescheduleSignal
             continue
         end
         pid, proc, thunk_id, (res, metadata) = chan_value
+        @dagdebug thunk_id :take "Got finished task"
         gproc = OSProc(pid)
         safepoint(state)
         lock(state.lock) do
@@ -559,6 +563,8 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
     return value, errored
 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)
@@ -697,6 +703,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
                     # FIXME: est_time_util = est_time_util isa MaxUtilization ? cap : est_time_util
                     push!(get!(()->Vector{Tuple{Thunk,<:Any,<:Any}}(), to_fire, (gproc, proc)), (task, est_time_util, est_alloc_util))
                     state.worker_time_pressure[gproc.pid][proc] += est_time_util
+                    @dagdebug task :schedule "Scheduling to $gproc -> $proc"
                     @goto pop_task
                 end
             end
@@ -1066,6 +1073,8 @@ function do_task(to_proc, comm)
     end
     timespan_finish(ctx, :storage_wait, thunk_id, (;f, to_proc, device=typeof(to_storage)))
 
+    @dagdebug thunk_id :execute "Moving data"
+
     # Initiate data transfers for function and arguments
     transfer_time = Threads.Atomic{UInt64}(0)
     transfer_size = Threads.Atomic{UInt64}(0)
@@ -1146,6 +1155,8 @@ function do_task(to_proc, comm)
     # FIXME
     #gcnum_start = Base.gc_num()
 
+    @dagdebug thunk_id :execute "Executing"
+
     result_meta = try
         # Set TLS variables
         Dagger.set_tls!((
@@ -1180,6 +1191,7 @@ function do_task(to_proc, comm)
         bt = catch_backtrace()
         RemoteException(myid(), CapturedException(ex, bt))
     end
+
     threadtime = cputhreadtime() - threadtime_start
     # FIXME: This is not a realistic measure of max. required memory
     #gc_allocd = min(max(UInt64(Base.gc_num().allocd) - UInt64(gcnum_start.allocd), UInt64(0)), UInt64(1024^4))
@@ -1189,6 +1201,9 @@ function do_task(to_proc, comm)
         pop!(TASKS_RUNNING, thunk_id)
         notify(TASK_SYNC)
     end
+
+    @dagdebug thunk_id :execute "Returning"
+
     # TODO: debug_storage("Releasing $to_storage_name")
     metadata = (
         time_pressure=real_time_util[],

src/sch/dynamic.jl

@@ -207,9 +207,11 @@ function _add_thunk!(ctx, state, task, tid, (f, args, kwargs, future, ref))
         thunk_id = ThunkID(thunk.id, thunk_ref)
         state.thunk_dict[thunk.id] = WeakThunk(thunk)
         reschedule_inputs!(state, thunk)
+        @dagdebug thunk :submit "Added to scheduler"
         if future !== nothing
             # Ensure we attach a future before the thunk is scheduled
             _register_future!(ctx, state, task, tid, (future, thunk_id, false))
+            @dagdebug thunk :submit "Registered future"
         end
         if ref !== nothing
             # Preserve the `EagerThunkFinalizer` through `thunk`

src/sch/util.jl

@@ -1,3 +1,37 @@
+const DAGDEBUG_CATEGORIES = Symbol[:global, :submit, :schedule, :scope,
+                                   :take, :execute]
+macro dagdebug(thunk, category, msg, args...)
+    cat_sym = category.value
+    @gensym id
+    debug_ex_id = :(@debug "[$($id)] ($($(repr(cat_sym)))) $($msg)" _module=Dagger _file=$(string(__source__.file)) _line=$(__source__.line))
+    append!(debug_ex_id.args, args)
+    debug_ex_noid = :(@debug "($($(repr(cat_sym)))) $($msg)" _module=Dagger _file=$(string(__source__.file)) _line=$(__source__.line))
+    append!(debug_ex_noid.args, args)
+    esc(quote
+        let $id = -1
+            if $thunk isa Integer
+                $id = Int($thunk)
+            elseif $thunk isa Thunk
+                $id = $thunk.id
+            elseif $thunk === nothing
+                $id = 0
+            else
+                @warn "Unsupported thunk argument to @dagdebug: $(typeof($thunk))"
+                $id = -1
+            end
+            if $id > 0
+                if $(QuoteNode(cat_sym)) in $DAGDEBUG_CATEGORIES
+                    $debug_ex_id
+                end
+            elseif $id == 0
+                if $(QuoteNode(cat_sym)) in $DAGDEBUG_CATEGORIES
+                    $debug_ex_noid
+                end
+            end
+        end
+    end)
+end
+
 """
     unwrap_nested_exception(err::Exception) -> Bool
 
@@ -267,13 +301,13 @@ function can_use_proc(task, gproc, proc, opts, scope)
         @warn "The `proclist` option is deprecated, please use scopes instead\nSee https://juliaparallel.org/Dagger.jl/stable/scopes/ for details" maxlog=1
         if opts.proclist isa Function
             if !Base.invokelatest(opts.proclist, proc)
-                @debug "[$(task.id)] Rejected $proc: proclist(proc) == false"
+                @dagdebug task :scope "Rejected $proc: proclist(proc) == false"
                 return false, scope
             end
             scope = constrain(scope, Dagger.ExactScope(proc))
         elseif opts.proclist isa Vector
             if !(typeof(proc) in opts.proclist)
-                @debug "[$(task.id)] Rejected $proc: !(typeof(proc) in proclist)"
+                @dagdebug task :scope "Rejected $proc: !(typeof(proc) in proclist)"
                 return false, scope
             end
             scope = constrain(scope,
@@ -282,7 +316,7 @@ function can_use_proc(task, gproc, proc, opts, scope)
             throw(SchedulingException("proclist must be a Function, Vector, or nothing"))
         end
         if scope isa Dagger.InvalidScope
-            @debug "[$(task.id)] Rejected $proc: Not contained in task scope ($scope)"
+            @dagdebug task :scope "Rejected $proc: Not contained in task scope ($scope)"
             return false, scope
         end
     end
@@ -291,26 +325,26 @@ function can_use_proc(task, gproc, proc, opts, scope)
     if opts.single !== nothing
         @warn "The `single` option is deprecated, please use scopes instead\nSee https://juliaparallel.org/Dagger.jl/stable/scopes/ for details" maxlog=1
         if gproc.pid != opts.single
-            @debug "[$(task.id)] Rejected $proc: gproc.pid ($(gproc.pid)) != single ($(opts.single))"
+            @dagdebug task :scope "Rejected $proc: gproc.pid ($(gproc.pid)) != single ($(opts.single))"
             return false, scope
         end
         scope = constrain(scope, Dagger.ProcessScope(opts.single))
         if scope isa Dagger.InvalidScope
-            @debug "[$(task.id)] Rejected $proc: Not contained in task scope ($scope)"
+            @dagdebug task :scope "Rejected $proc: Not contained in task scope ($scope)"
             return false, scope
         end
     end
 
     # Check against scope
     proc_scope = Dagger.ExactScope(proc)
     if constrain(scope, proc_scope) isa Dagger.InvalidScope
-        @debug "[$(task.id)] Rejected $proc: Not contained in task scope ($scope)"
+        @dagdebug task :scope "Rejected $proc: Not contained in task scope ($scope)"
         return false, scope
     end
 
     @label accept
 
-    @debug "[$(task.id)] Accepted $proc"
+    @dagdebug task :scope "Accepted $proc"
     return true, scope
 end