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

Author: jpsamaroo

Diff for: ext/GraphVizExt.jl

@@ -7,7 +7,7 @@ else
 end
 
 import Dagger
-import Dagger: DTask, Chunk, Processor
+import Dagger: DTask, Chunk, Processor, LoggedMutableObject
 import Dagger.TimespanLogging: Timespan
 import Graphs: SimpleDiGraph, add_edge!, add_vertex!, inneighbors, outneighbors, vertices, is_directed, edges, nv, src, dst
 
@@ -24,6 +24,8 @@ function pretty_time(t; digits=3)
     end
 end
 
+sanitize_label(label::String) = replace(label, "\"" => "\\\"")
+
 _name_to_color(name::AbstractString, colors) =
     colors[mod1(hash(name), length(colors))]
 _name_to_color(name::AbstractString, ::Nothing) = "black"
@@ -51,12 +53,22 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                             colors=_default_colors, name_to_color=_name_to_color)
     # Lookup all relevant task/argument dependencies and values in logs
     g = SimpleDiGraph()
+
     tid_to_vertex = Dict{Int,Int}()
-    task_names = String[]
+    tid_to_auto_name = Dict{Int,String}()
+    tid_to_name = Dict{Int,String}()
     tid_to_proc = Dict{Int,Processor}()
-    arg_names = Dict{UInt,String}()
+
+    objid_to_vertex = Dict{UInt,Int}()
+    objid_to_name = Dict{UInt,String}()
+
     task_args = Dict{Int,Vector{Pair{Union{Int,Symbol},UInt}}}()
-    arg_moves = Dict{Int,Vector{Pair{Union{Int,Symbol},Tuple{UInt,UInt}}}}()
+    task_arg_moves = Dict{Int,Vector{Pair{Union{Int,Symbol},Tuple{UInt,UInt}}}}()
+    task_result = Dict{Int,UInt}()
+
+    uid_to_tid = Dict{UInt,Int}()
+    dtasks_to_patch = Set{UInt}()
+
     for w in keys(logs)
         for idx in 1:length(logs[w][:core])
             category = logs[w][:core][idx].category
@@ -67,40 +79,143 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                 taskdeps = logs[w][:taskdeps][idx]::Pair{Int,Vector{Int}}
                 taskname = logs[w][:tasknames][idx]::String
                 tid, deps = taskdeps
-                add_vertex!(g)
-                tid_to_vertex[tid] = nv(g)
-                push!(task_names, taskname)
+                v = get!(tid_to_vertex, tid) do
+                    add_vertex!(g)
+                    tid_to_vertex[tid] = nv(g)
+                    nv(g)
+                end
+                tid_to_auto_name[tid] = taskname
                 for dep in deps
-                    haskey(tid_to_vertex, dep) || continue
-                    add_edge!(g, tid_to_vertex[dep], nv(g))
+                    dep_v = get!(tid_to_vertex, dep) do
+                        add_vertex!(g)
+                        tid_to_vertex[dep] = nv(g)
+                        nv(g)
+                    end
+                    add_edge!(g, dep_v, v)
+                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 == :compute && kind == :start
                 id::NamedTuple
                 tid = id.thunk_id
                 proc = id.processor
                 tid_to_proc[tid] = proc
+            elseif category == :compute && kind == :finish
+                if haskey(logs[w], :taskresult)
+                    result_info = logs[w][:taskresult][idx]
+                    result_info === nothing && continue
+                    tid, obj = result_info::Pair{Int,LoggedMutableObject}
+                    objid = obj.objid
+                    task_result[tid] = objid
+                    v = get!(objid_to_vertex, objid) do
+                        add_vertex!(g)
+                        objid_to_vertex[objid] = nv(g)
+                        nv(g)
+                    end
+                    add_edge!(g, tid_to_vertex[tid], v)
+                end
             elseif category == :move && kind == :finish
                 if haskey(logs[w], :taskargs)
-                    id, args = logs[w][:taskargs][idx]::Pair{Int,<:Vector}
-                    append!(get!(Vector{Pair{Union{Int,Symbol},UInt}}, task_args, id), args)
+                    tid, args = logs[w][:taskargs][idx]::Pair{Int,<:Vector}
+                    args = map(arg->arg[1]=>arg[2].objid, args)
+                    append!(get!(Vector{Pair{Union{Int,Symbol},UInt}}, task_args, tid), args)
+                    for arg in args
+                        objid = arg[2]
+                        arg_id = get!(objid_to_vertex, objid) do
+                            add_vertex!(g)
+                            objid_to_vertex[objid] = nv(g)
+                            nv(g)
+                        end
+                        add_edge!(g, arg_id, tid_to_vertex[tid])
+                    end
                 end
                 if haskey(logs[w], :taskargmoves)
                     move_info = logs[w][:taskargmoves][idx]
                     move_info === nothing && continue
-                    tid, pos, pre_objid, post_objid = move_info
-                    v = get!(Vector{Pair{Union{Int,Symbol},Tuple{UInt,UInt}}}, arg_moves, tid)
+                    tid, pos, pre_obj, post_obj = move_info
+                    v = get!(Vector{Pair{Union{Int,Symbol},Tuple{UInt,UInt}}}, task_arg_moves, tid)
+                    pre_objid = pre_obj.objid
+                    post_objid = post_obj.objid
                     push!(v, pos => (pre_objid, post_objid))
+                    pre_arg_id = get!(objid_to_vertex, pre_objid) do
+                        add_vertex!(g)
+                        objid_to_vertex[pre_objid] = nv(g)
+                        nv(g)
+                    end
+                    post_arg_id = get!(objid_to_vertex, post_objid) do
+                        add_vertex!(g)
+                        objid_to_vertex[post_objid] = nv(g)
+                        nv(g)
+                    end
+                    add_edge!(g, pre_arg_id, post_arg_id)
                 end
             elseif category == :data_annotation && kind == :start
                 id::NamedTuple
-                objid = id.objectid
-                name = id.name
-                arg_names[objid] = String(name)
+                name = String(id.name)
+                obj = id.objectid::LoggedMutableObject
+                objid = obj.objid
+                objid_to_name[objid] = name
+                if obj.kind == :task
+                    # N.B. We don't need the object vertex,
+                    # since we'll just render it as a task
+                    push!(dtasks_to_patch, objid)
+                else
+                    get!(objid_to_vertex, objid) do
+                        add_vertex!(g)
+                        objid_to_vertex[objid] = nv(g)
+                        nv(g)
+                    end
+                end
+            elseif category == :finish && kind == :finish
+                if haskey(logs[w], :tasktochunk)
+                    tid_chunk = logs[w][:tasktochunk][idx]
+                    if tid_chunk !== nothing
+                        tid, chunk_obj = tid_chunk::Pair{Int,LoggedMutableObject}
+                        chunk_id = chunk_obj.objid
+                        v = get!(objid_to_vertex, chunk_id) do
+                            add_vertex!(g)
+                            objid_to_vertex[chunk_id] = nv(g)
+                            nv(g)
+                        end
+                        add_edge!(g, tid_to_vertex[tid], v)
+                    end
+                end
             end
         end
     end
-    tids_sorted = map(first, sort(collect(tid_to_vertex); by=last))
-    task_procs = Processor[tid_to_proc[tid] for tid in tids_sorted]
+
+    # Process DTasks-to-Thunk mappings
+    for uid in dtasks_to_patch
+        if haskey(uid_to_tid, uid)
+            tid = uid_to_tid[uid]
+            v = get!(tid_to_vertex, tid) do
+                add_vertex!(g)
+                tid_to_vertex[tid] = nv(g)
+                nv(g)
+            end
+
+            # Fixup any missing tid data
+            if haskey(objid_to_name, uid)
+                tid_to_name[tid] = objid_to_name[uid]
+            end
+        end
+    end
+
+    # Auto-assign names
+    for (tid, name) in tid_to_auto_name
+        if !haskey(tid_to_name, tid)
+            tid_to_name[tid] = name
+        end
+    end
+
+    # Create reverse mappings
+    vertex_to_tid = Dict{Int,Int}(v=>k for (k,v) in tid_to_vertex)
+    vertex_to_objid = Dict{Int,UInt}(v=>k for (k,v) in objid_to_vertex)
 
     # Find all connected and disconnected vertices
     if !disconnected
@@ -110,31 +225,7 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
         con_vs = vertices(g)
     end
 
-    # Assign colors
-    labels = task_names
-    all_fns = unique(map(label->first(split(label, " ")), labels[con_vs]))
-    all_procs = unique(task_procs)
-
-    if color_by == :fn
-        _colors = [name_to_color(all_fns[i], colors) for i in 1:length(all_fns)]
-        colors = Dict(v=>_colors[findfirst(fn->occursin(fn, labels[v]), all_fns)] for v in con_vs)
-    elseif color_by == :proc
-        _colors = [name_to_color(string(all_procs[i]), colors) for i in 1:length(all_procs)]
-        colors = Dict(v=>_colors[findfirst(proc->proc==task_procs[v], all_procs)] for v in con_vs)
-    else
-        throw(ArgumentError("Unknown `color_by` value: $color_by\nAllowed: :fn, :proc"))
-    end
-
-    str = is_directed(g) ? "digraph mygraph {\n" : "graph mygraph {\n"
-
-    # Add raw arguments
-    for (id, name) in arg_names
-        str *= "a$id [label=\"$name\", shape=box]\n"
-    end
-
     if times
-        vertex_to_tid = Dict{Int,Int}(v=>k for (k,v) in tid_to_vertex)
-
         # Determine per-worker start times
         worker_start_times = Dict{Int,UInt64}()
         for w in keys(logs)
@@ -169,68 +260,127 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
         end
     end
 
-    # Add argument moves
-    seen_moves = Set{Tuple{UInt,UInt}}()
-    for (tid, moves) in arg_moves
-        for (pos, (pre_objid, post_objid)) in moves
-            pre_objid == post_objid && continue
-            (pre_objid, post_objid) in seen_moves && continue
-            push!(seen_moves, (pre_objid, post_objid))
-            move_str = "a$pre_objid -> a$post_objid [label=\"move\"]\n"
-            str *= move_str
-        end
+    # Get the set of all unique task and object IDs
+    all_tids = collect(keys(tid_to_vertex))
+    all_objids = collect(keys(objid_to_vertex))
+
+    # Assign colors
+    if color_by == :fn
+        all_fns = unique(map(label->first(split(label, " ")), values(tid_to_name)))
+        _colors = [name_to_color(all_fns[i], colors) for i in 1:length(all_fns)]
+        colors = Dict(tid=>_colors[findfirst(fn->occursin(fn, tid_to_name[tid]), all_fns)] for tid in all_tids)
+    elseif color_by == :proc
+        all_procs = unique(values(tid_to_proc))
+        _colors = [name_to_color(string(all_procs[i]), colors) for i in 1:length(all_procs)]
+        colors = Dict(tid=>_colors[findfirst(proc->proc==tid_to_proc[tid], all_procs)] for tid in all_tids)
+    else
+        throw(ArgumentError("Unknown `color_by` value: $color_by\nAllowed: :fn, :proc"))
     end
 
-    # Add tasks
-    for v in con_vs
+    str = is_directed(g) ? "digraph mygraph {\n" : "graph mygraph {\n"
+
+    # Add task vertices
+    for tid in all_tids
+        v = tid_to_vertex[tid]
         if !disconnected && (v in discon_vs)
             continue
         end
-        label = labels[v]
-        color = colors[v]
-        proc = task_procs[v]
-        proc_str = '(' * Dagger.short_name(task_procs[v]) * ')'
-        label_str = "$label\\n$proc_str"
+        label_str = tid_to_name[tid]
+        if haskey(tid_to_auto_name, tid) && tid_to_name[tid] != tid_to_auto_name[tid]
+            label_str *= "\\nTask: $(tid_to_auto_name[tid])"
+        end
+        color = colors[tid]
+        proc = tid_to_proc[tid]
+        label_str *= "\\n($(Dagger.short_name(tid_to_proc[tid])))"
         if times
-            tid = vertex_to_tid[v]
             start_time = pretty_time(start_times[tid]; digits=times_digits)
             finish_time = pretty_time(finish_times[tid]; digits=times_digits)
-            label_str *= "\\n[+$start_time -> +$finish_time]"
+            diff_time = pretty_time(finish_times[tid] - start_times[tid]; digits=times_digits)
+            label_str *= "\\n[+$start_time -> +$finish_time (diff: $diff_time)]"
         end
-        str *= "v$v [label=\"$label_str\", color=\"$color\", penwidth=2.0]\n"
+        label_str = sanitize_label(label_str)
+        str *= "v$v [label=\"$label_str\", shape=box, color=\"$color\", penwidth=2.0]\n"
     end
 
-    # Add task dependencies
+    # Add object vertices
+    for objid in all_objids
+        objid_v = objid_to_vertex[objid]
+        if !disconnected && !(objid_v in con_vs)
+            continue
+        end
+        if objid in dtasks_to_patch || haskey(uid_to_tid, objid)
+            # DTask, skip it
+            continue
+        end
+        # Object
+        if haskey(objid_to_name, objid)
+            label = sanitize_label(objid_to_name[objid])
+            label *= "\\nData: $(repr(objid))"
+        else
+            label = "Data: $(repr(objid))"
+        end
+        str *= "a$objid_v [label=\"$label\", shape=oval]\n"
+    end
+
+    # Add task argument move edges
+    seen_moves = Set{Tuple{UInt,UInt}}()
+    for (tid, moves) in task_arg_moves
+        for (pos, (pre_objid, post_objid)) in moves
+            pre_objid == post_objid && continue
+            (pre_objid, post_objid) in seen_moves && continue
+            push!(seen_moves, (pre_objid, post_objid))
+            pre_objid_v = objid_to_vertex[pre_objid]
+            post_objid_v = objid_to_vertex[post_objid]
+            move_str = "a$pre_objid_v -> a$post_objid_v [label=\"move\"]\n"
+            str *= move_str
+        end
+    end
+
+    # Add task-to-task (syncdep) dependency edges
     edge_sep = is_directed(g) ? "->" : "--"
     for edge in edges(g)
+        if !haskey(vertex_to_tid, src(edge)) || !haskey(vertex_to_tid, dst(edge))
+            continue
+        end
+        if !disconnected && !(src(edge) in con_vs) || !(dst(edge) in con_vs)
+            continue
+        end
         # FIXME: Label syncdeps with associated arguments and datadeps directions
         str *= "v$(src(edge)) $edge_sep v$(dst(edge)) [label=\"syncdep\"]\n"
     end
 
-    # Add task arguments
-    con_args = Vector{UInt}(collect(keys(arg_names)))
-    for moves in values(arg_moves)
-        for (_, (pre_objid, post_objid)) in moves
-            push!(con_args, pre_objid)
-            push!(con_args, post_objid)
-        end
-    end
+    # Add task argument edges
     for (tid, args) in task_args
         haskey(tid_to_vertex, tid) || continue
-        id = tid_to_vertex[tid]
-        id in con_vs || continue
+        tid_v = tid_to_vertex[tid]
+        tid_v in con_vs || continue
         for (pos, arg) in args
-            if !disconnected && !(arg in con_args)
+            arg_v = objid_to_vertex[arg]
+            if !disconnected && !(arg_v in con_vs)
                 continue
             end
-            arg_str = pos isa Int ? "arg $pos" : "kwarg $pos"
-            str *= "a$arg $edge_sep v$id [label=\"$arg_str\"]\n"
+            arg_str = sanitize_label(pos isa Int ? "arg $pos" : "kwarg $pos")
+            str *= "a$arg_v $edge_sep v$tid_v [label=\"$arg_str\"]\n"
         end
     end
 
+    # Add task result edges
+    for (tid, result) in task_result
+        haskey(tid_to_vertex, tid) || continue
+        tid_v = tid_to_vertex[tid]
+        tid_v in con_vs || continue
+        result_v = objid_to_vertex[result]
+        if !disconnected && !(result_v in con_vs)
+            continue
+        end
+        str *= "v$tid_v $edge_sep a$result_v [label=\"result\"]\n"
+    end
+
+    # Generate the final graph
     str *= "}\n"
     gv = GraphViz.Graph(str)
     GraphViz.layout!(gv; engine=layout_engine)
+
     return gv
 end