Task names and GraphVizExt improvements

at-spawn: Add `name` option for custom task names
logging: Allow `DTask` in `logs_annotate!`
logging: Allow `DArray` in `logs_annotate!` (labels array chunks)
logging: Add task result, DTask-to-Thunk, UID-to-TID logging events
logging: Add `all_task_deps` option to `enable_logging!` for convenience
GraphVizExt: Properly label tasks and objects
GraphVizExt: Show box for task and oval for object
GraphVizExt: Improve connectivity checks
GraphVizExt/PlotsExt: Change default colors to Tab20
PlotsExt: Add processor and scheduler visualization modes
PlotsExt: Remove unused plotting implementation
PlotsExt: Ensure X axis is globally aligned

Author: jpsamaroo

ext/GraphVizExt.jl

@@ -15,11 +15,20 @@ import Dagger.TimespanLogging: Timespan
 _name_to_color(name::AbstractString, colors) =
     colors[mod1(hash(name), length(colors))]
 _name_to_color(name::AbstractString, ::Nothing) = "black"
-_default_colors = ["red", "orange", "green", "blue", "purple", "pink", "silver"]
+tab20_colors = [
+    "#1f77b4", "#aec7e8", "#ff7f0e", "#ffbb78",
+    "#2ca02c", "#98df8a", "#d62728", "#ff9896",
+    "#9467bd", "#c5b0d5", "#8c564b", "#c49c94",
+    "#e377c2", "#f7b6d2", "#7f7f7f", "#c7c7c7",
+    "#bcbd22", "#dbdb8d", "#17becf", "#9edae5"
+]
+_default_colors = tab20_colors
 
-function logs_to_df(logs::Dict; colors=_default_colors, name_to_color=_name_to_color, color_by=:fn)
+function logs_to_df(logs::Dict, ::Val{:execution}; colors=_default_colors, name_to_color=_name_to_color, color_by=:fn)
     # Generate function names
     fn_names = Dict{Int, String}()
+    dtask_names = Dict{UInt, String}()
+    uid_to_tid = Dict{UInt, Int}()
     for w in keys(logs)
         for idx in 1:length(logs[w][:core])
             category = logs[w][:core][idx].category::Symbol
@@ -32,10 +41,33 @@ function logs_to_df(logs::Dict; colors=_default_colors, name_to_color=_name_to_c
                     @warn "Task names missing from logs"
                     fn_names[tid] = "unknown"
                 end
+                if haskey(logs[w], :taskuidtotid)
+                    uid_tid = logs[w][:taskuidtotid][idx]
+                    if uid_tid !== nothing
+                        uid, tid = uid_tid::Pair{UInt,Int}
+                        uid_to_tid[uid] = tid
+                    end
+                end
+            elseif category == :data_annotation && kind == :start
+                id = logs[w][:id][idx]::NamedTuple
+                name = String(id.name)
+                obj = id.objectid::Dagger.LoggedMutableObject
+                objid = obj.objid
+                if obj.kind == :task
+                    dtask_names[objid] = name
+                end
             end
         end
     end
 
+    # Process DTasks-to-Thunk mappings
+    for (uid, tid) in uid_to_tid
+        # Patch in custom name
+        if haskey(dtask_names, uid)
+            fn_names[tid] = dtask_names[uid]
+        end
+    end
+
     # FIXME: Color eltype
     df = DataFrame(proc=Processor[], proc_name=String[], fn_name=String[], tid=Int[], t_start=UInt64[], t_end=UInt64[], color=Any[])
     Dagger.logs_event_pairs(logs) do w, start_idx, finish_idx
@@ -59,43 +91,42 @@ function logs_to_df(logs::Dict; colors=_default_colors, name_to_color=_name_to_c
     end
     return df
 end
-
-# Implementation by Przemyslaw Szufel
-function Dagger.render_logs(logs::Dict, ::Val{:plots_gantt_ps})
-    df = logs_to_df(logs)
-
-    proc_names = sort!(unique(df.proc_name))
-    proc_idx = Dict{String,Int}()
-    for name in proc_names
-        proc_idx[name] = findfirst(==(name), proc_names)
-    end
-    proc_idxs = map(name->proc_idx[name], proc_names)
-    tvals = zeros(UInt64, length(proc_names))
-    plt = bar(orientation=:h, yticks=(1:length(proc_names), proc_names), linewidth=0,yflip=true,color=:green,legend=nothing)
-    xlabel!(plt, "Time in seconds")
-    dfc = deepcopy(df)
-    while nrow(dfc) > 0
-        rowslast = DataFrame([g[findmax(g.t_end)[2],:] for g in groupby(dfc, :proc_name)])
-        tvals .= .0
-        for i in 1:nrow(rowslast)
-            tvals[proc_idx[rowslast.proc_name[i]]] = rowslast.t_end[i]
+function logs_to_df(logs::Dict, ::Val{:processor}; colors=_default_colors, name_to_color=_name_to_color, kwargs...)
+    # Collect processor events
+    # FIXME: Color eltype
+    df = DataFrame(proc=Processor[], proc_name=String[], category=String[], t_start=UInt64[], t_end=UInt64[], color=Any[])
+    Dagger.logs_event_pairs(logs) do w, start_idx, finish_idx
+        category = logs[w][:core][start_idx].category
+        if category in (:compute, :storage_wait, :storage_safe_scan, :proc_run_fetch, :proc_steal_local)
+            proc = logs[w][:id][start_idx].processor::Processor
+            proc_name = Dagger.short_name(proc)
+            t_start = logs[w][:core][start_idx].timestamp::UInt64
+            t_end = logs[w][:core][finish_idx].timestamp::UInt64
+            category_str = string(category)
+            color = name_to_color(category_str, colors)
+            push!(df, (;proc, proc_name, category=category_str, t_start, t_end, color))
         end
-        #setindex!.(Ref(tvals), rowslast.t_end, getindex.(proc_idx, rowslast.proc_name))
-        bar!(plt, tvals[proc_idxs], orientation=:h, linewidth=0.5,yflip=true,color=:green)
-        tvals .= .0
-        for i in 1:nrow(rowslast)
-            tvals[proc_idx[rowslast.proc_name[i]]] = rowslast.t_start[i]
+    end
+    return df
+end
+function logs_to_df(logs::Dict, ::Val{:scheduler}; colors=_default_colors, name_to_color=_name_to_color, kwargs...)
+    # Collect scheduler events
+    # FIXME: Color eltype
+    df = DataFrame(category=String[], t_start=UInt64[], t_end=UInt64[], color=Any[])
+    Dagger.logs_event_pairs(logs) do w, start_idx, finish_idx
+        category = logs[w][:core][start_idx].category
+        if category in (:scheduler_init, :scheduler_exit, :init_proc, :remove_proc, :add_thunk, :handle_fault, :schedule, :fire, :enqueue, :take, :finish)
+            t_start = logs[w][:core][start_idx].timestamp::UInt64
+            t_end = logs[w][:core][finish_idx].timestamp::UInt64
+            category_str = string(category)
+            color = name_to_color(category_str, colors)
+            push!(df, (;category=category_str, t_start, t_end, color))
         end
-        #setindex!.(Ref(tvals), rowslast.t_start, proc_idx[rowslast.proc_name])
-        bar!(plt, tvals[proc_idxs], orientation=:h, linewidth=0.5,linecolor=:white,yflip=true,color=:white)
-        annotate!.(Ref(plt),(rowslast.t_start .+ rowslast.t_end) ./ 2,  findfirst.( .==(rowslast.proc_name), Ref(proc_names)),  text.(string.(rowslast.tid),9,rotation=90 ))
-        dfc = dfc[ .! (dfc.tid .∈ Ref(rowslast.tid) ), : ]
     end
-    # FIXME: theoretic_optimal = simulate_polling(df)[1] + minimum(df.t_start)
-    theoretic_optimal = minimum(df.t_start)
-    plot!(plt, [theoretic_optimal,theoretic_optimal], [0, length(proc_names)+1],width=3,color=:black,style=:dot)
-    return plt
+    return df
 end
+logs_to_df(logs::Dict, ::Val{target}; kwargs...) where target =
+        throw(ArgumentError("Invalid target: $(repr(target))"))
 
 # Implementation adapted from:
 # https://discourse.julialang.org/t/how-to-make-a-gantt-plot-with-plots-jl/95165/7
@@ -105,30 +136,54 @@ end
 Render a Gantt chart of task execution in `logs` using Plots. `kwargs` are passed to `plot` directly.
 """
 function Dagger.render_logs(logs::Dict, ::Val{:plots_gantt};
+                            target=:execution,
                             colors=_default_colors, name_to_color=_name_to_color,
                             color_by=:fn, kwargs...)
-    df = logs_to_df(logs; colors, name_to_color, color_by)
+    df = logs_to_df(logs, Val{target}(); colors, name_to_color, color_by)
+    y_elem = if target == :execution || target == :processor
+        :proc_name
+    elseif target == :scheduler
+        :category
+    end
+    ylabel = if target == :execution || target == :processor
+        "Processor"
+    elseif target == :scheduler
+        "Category"
+    end
+    sort!(df, y_elem)
+
+    global_t_start, global_t_end = extrema(logs[1][:core][idx].timestamp for idx in 1:length(logs[1][:core]))
 
     rect(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
 
-    t_init = minimum(df.t_start)
-    t_start = (df.t_start .- t_init) ./ 1e9
-    t_end = (df.t_end .- t_init) ./ 1e9
+    t_start = (df.t_start .- global_t_start) ./ 1e9
+    t_end = (df.t_end .- global_t_start) ./ 1e9
     duration = t_end .- t_start
-    u = unique(df.proc_name)
+    u = unique(getproperty(df, y_elem))
     dy = Dict(u .=> 1:length(u))
-    r = [rect(t1, 1, t2, dy[t3]) for (t1,t2,t3) in zip(duration, t_start, df.proc_name)]
-    labels = permutedims(df.fn_name)
-    # Deduplicate labels
-    for idx in 1:length(labels)
-        if findfirst(other_idx->labels[other_idx]==labels[idx], 1:length(labels)) < idx
-            labels[idx] = ""
+    r = [rect(t1, 1, t2, dy[t3]) for (t1,t2,t3) in zip(duration, t_start, getproperty(df, y_elem))]
+    if target == :execution
+        labels = permutedims(df.fn_name)
+    elseif target == :processor
+        labels = permutedims(df.category)
+    elseif target == :scheduler
+        labels = nothing
+    end
+
+    if labels !== nothing
+        # Deduplicate labels
+        for idx in 1:length(labels)
+            if findfirst(other_idx->labels[other_idx]==labels[idx], 1:length(labels)) < idx
+                labels[idx] = ""
+            end
         end
     end
 
     return plot(r; color=permutedims(df.color), labels,
                 yticks=(1.5:(nrow(df) + 0.5), u),
-                xlabel="Time (seconds)", ylabel="Processor",
+                xlabel="Time (seconds)", ylabel,
+                xlim=(0.0, (global_t_end - global_t_start) / 1e9),
+                legendalpha=0,
                 kwargs...)
 end
 

ext/PlotsExt.jl

@@ -524,6 +524,13 @@ function Base.:(==)(x::AbstractArray{T,N}, y::ArrayOp{S,N}) where {T,S,N}
     return collect(x) == y
 end
 
+function logs_annotate!(ctx::Context, A::DArray, name::Union{String,Symbol})
+    for (idx, chunk) in enumerate(A.chunks)
+        sd = A.subdomains[idx]
+        Dagger.logs_annotate!(ctx, chunk, name*'['*join(sd.indexes, ',')*']')
+    end
+end
+
 # TODO: Allow `f` to return proc
 mapchunk(f, chunk) = tochunk(f(poolget(chunk.handle)))
 function mapchunks(f, d::DArray{T,N,F}) where {T,N,F}

src/array/darray.jl

@@ -604,9 +604,9 @@ function scheduler_run(ctx, state::ComputeState, d::Thunk, options)
                 end
             end
 
-            timespan_start(ctx, :finish, (;thunk_id), (;thunk_id))
+            timespan_start(ctx, :finish, (;thunk_id), (;thunk_id, result=res))
             finish_task!(ctx, state, node, thunk_failed)
-            timespan_finish(ctx, :finish, (;thunk_id), (;thunk_id))
+            timespan_finish(ctx, :finish, (;thunk_id), (;thunk_id, result=res))
 
             delete_unused_tasks!(state)
         end
@@ -1684,7 +1684,7 @@ function do_task(to_proc, task_desc)
     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))
-    timespan_finish(ctx, :compute, (;thunk_id, processor=to_proc), (;f))
+    timespan_finish(ctx, :compute, (;thunk_id, processor=to_proc), (;f, result=result_meta))
     lock(TASK_SYNC) do
         real_time_util[] -= est_time_util
         pop!(TASKS_RUNNING, thunk_id)

src/sch/Sch.jl

@@ -26,7 +26,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
 
     id = next_id()
 
-    timespan_start(ctx, :add_thunk, (;thunk_id=id), (;f, args, options))
+    timespan_start(ctx, :add_thunk, (;thunk_id=id), (;f, args, options, uid))
 
     # Lookup DTask/ThunkID -> Thunk
     old_args = copy(args)
@@ -129,7 +129,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
             put!(state.chan, Sch.RescheduleSignal())
         end
 
-        timespan_finish(ctx, :add_thunk, (;thunk_id=id), (;f, args, options))
+        timespan_finish(ctx, :add_thunk, (;thunk_id=id), (;f, args, options, uid))
 
         return thunk_id
     end
@@ -224,10 +224,13 @@ function eager_spawn(spec::DTaskSpec)
     future = ThunkFuture()
     finalizer_ref = poolset(DTaskFinalizer(uid); device=MemPool.CPURAMDevice())
 
-    # Return unlaunched DTask
+    # Create unlaunched DTask
     return DTask(uid, future, finalizer_ref)
 end
 function eager_launch!((spec, task)::Pair{DTaskSpec,DTask})
+    # Assign a name, if specified
+    eager_assign_name!(spec, task)
+
     # Lookup DTask -> ThunkID
     local args, options
     lock(Sch.EAGER_ID_MAP) do id_map
@@ -244,6 +247,11 @@ end
 function eager_launch!(specs::Vector{Pair{DTaskSpec,DTask}})
     ntasks = length(specs)
 
+    # Assign a name, if specified
+    for (spec, task) in specs
+        eager_assign_name!(spec, task)
+    end
+
     uids = [task.uid for (_, task) in specs]
     futures = [task.future for (_, task) in specs]
     finalizer_refs = [task.finalizer_ref for (_, task) in specs]
@@ -263,3 +271,11 @@ function eager_launch!(specs::Vector{Pair{DTaskSpec,DTask}})
         task.thunk_ref = thunk_ids[i].ref
     end
 end
+
+function eager_assign_name!(spec::DTaskSpec, task::DTask)
+    # Assign a name, if specified
+    if haskey(spec.options, :name)
+        Dagger.logs_annotate!(task, spec.options.name)
+        spec.options = (;filter(x -> x[1] != :name, Base.pairs(spec.options))...)
+    end
+end