logging/Sch/GraphVizExt: Various :graphviz improvements

Author: jpsamaroo

ext/GraphVizExt.jl

@@ -65,18 +65,19 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                 tid_to_vertex[tid] = nv(g)
                 push!(task_names, taskname)
                 for dep in deps
+                    haskey(tid_to_vertex, dep) || continue
                     add_edge!(g, tid_to_vertex[dep], nv(g))
                 end
-                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)
-                end
             elseif category == :compute && kind == :start
                 id::NamedTuple
                 tid = id.thunk_id
                 proc = id.processor
                 tid_to_proc[tid] = proc
             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)
+                end
                 if haskey(logs[w], :taskargmoves)
                     move_info = logs[w][:taskargmoves][idx]
                     move_info === nothing && continue
@@ -88,7 +89,7 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                 id::NamedTuple
                 objid = id.objectid
                 name = id.name
-                arg_names[objid] = name
+                arg_names[objid] = String(name)
             end
         end
     end
@@ -163,8 +164,12 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
     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
@@ -205,6 +210,7 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
         end
     end
     for (tid, args) in task_args
+        haskey(tid_to_vertex, tid) || continue
         id = tid_to_vertex[tid]
         id in con_vs || continue
         for (pos, arg) in args

src/datadeps.jl

@@ -413,7 +413,7 @@ function generate_slot!(state::DataDepsState, dest_space, data)
         w = only(unique(map(get_parent, collect(processors(dest_space))))).pid
         ctx = Sch.eager_context()
         id = rand(Int)
-        timespan_start(ctx, :move, (;thunk_id=0, id, processor=to_proc), (;f=nothing, data))
+        timespan_start(ctx, :move, (;thunk_id=0, id, position=0, processor=to_proc), (;f=nothing, data))
         dest_space_args[data] = remotecall_fetch(w, from_proc, to_proc, data) do from_proc, to_proc, data
             data_converted = move(from_proc, to_proc, data)
             data_chunk = tochunk(data_converted, to_proc)
@@ -422,7 +422,7 @@ function generate_slot!(state::DataDepsState, dest_space, data)
             @assert orig_space != memory_space(data_chunk) "space preserved! $orig_space != $(memory_space(data_chunk)) ($(typeof(data)) vs. $(typeof(data_chunk))), spaces ($orig_space -> $dest_space)"
             return data_chunk
         end
-        timespan_finish(ctx, :move, (;thunk_id=0, id, processor=to_proc), (;f=nothing, data=dest_space_args[data]))
+        timespan_finish(ctx, :move, (;thunk_id=0, id, position=0, processor=to_proc), (;f=nothing, data=dest_space_args[data]))
     end
     return dest_space_args[data]
 end

src/sch/Sch.jl

@@ -1049,13 +1049,21 @@ function fire_tasks!(ctx, thunks::Vector{<:Tuple}, (gproc, proc), state)
 
         ids = Int[0]
         data = Any[thunk.f]
-        positions = Union{Symbol,Nothing}[]
+        positions = Union{Symbol,Int}[0]
+        arg_ctr = 1
         for (idx, pos_x) in enumerate(thunk.inputs)
             pos, x = pos_x
             x = unwrap_weak_checked(x)
             push!(ids, istask(x) ? x.id : -idx)
             push!(data, istask(x) ? state.cache[x] : x)
-            push!(positions, pos)
+            if pos !== nothing
+                # Keyword arg
+                push!(positions, pos)
+            else
+                # Positional arg
+                push!(positions, arg_ctr)
+                arg_ctr += 1
+            end
         end
         toptions = thunk.options !== nothing ? thunk.options : ThunkOptions()
         options = merge(ctx.options, toptions)
@@ -1537,14 +1545,14 @@ function do_task(to_proc, task_desc)
     # Initiate data transfers for function and arguments
     transfer_time = Threads.Atomic{UInt64}(0)
     transfer_size = Threads.Atomic{UInt64}(0)
-    _data, _ids = if meta
-        (Any[first(data)], Int[first(ids)]) # always fetch function
+    _data, _ids, _positions = if meta
+        (Any[first(data)], Int[first(ids)], Union{Symbol,Int}[first(positions)]) # always fetch function
     else
-        (data, ids)
+        (data, ids, positions)
     end
-    fetch_tasks = map(Iterators.zip(_data,_ids)) do (x, id)
+    fetch_tasks = map(Iterators.zip(_data, _ids, _positions)) do (x, id, position)
         Threads.@spawn begin
-            timespan_start(ctx, :move, (;thunk_id, id, processor=to_proc), (;f, data=x))
+            timespan_start(ctx, :move, (;thunk_id, id, position, processor=to_proc), (;f, data=x))
             #= FIXME: This isn't valid if x is written to
             x = if x isa Chunk
                 value = lock(TASK_SYNC) do
@@ -1587,11 +1595,13 @@ function do_task(to_proc, task_desc)
                 end
             else
             =#
-            x = @invokelatest move(to_proc, x)
+            new_x = @invokelatest move(to_proc, x)
             #end
-            @dagdebug thunk_id :move "Moved argument $id to $to_proc: $(typeof(x))"
-            timespan_finish(ctx, :move, (;thunk_id, id, processor=to_proc), (;f, data=x); tasks=[Base.current_task()])
-            return x
+            if new_x !== x
+                @dagdebug thunk_id :move "Moved argument $position to $to_proc: $(typeof(x)) -> $(typeof(new_x))"
+            end
+            timespan_finish(ctx, :move, (;thunk_id, id, position, processor=to_proc), (;f, data=new_x); tasks=[Base.current_task()])
+            return new_x
         end
     end
     fetched = Any[]
@@ -1607,8 +1617,8 @@ function do_task(to_proc, task_desc)
     fetched_args = Any[]
     fetched_kwargs = Pair{Symbol,Any}[]
     for (idx, x) in enumerate(fetched)
-        pos = positions[idx]
-        if pos === nothing
+        pos = positions[idx+1]
+        if pos isa Int
             push!(fetched_args, x)
         else
             push!(fetched_kwargs, pos => x)

src/submission.jl

@@ -38,6 +38,7 @@ function eager_submit_internal!(ctx, state, task, tid, payload; uid_to_tid=Dict{
     end::Union{Vector{Any},Nothing}
     lock(Sch.EAGER_ID_MAP) do id_map
         for (idx, (pos, arg)) in enumerate(args)
+            # FIXME: Switch to Union{Symbol,Int} to preserve positional information
             pos::Union{Symbol,Nothing}
             newarg = if arg isa DTask
                 arg_uid = arg.uid

src/thunk.jl

@@ -362,6 +362,8 @@ function replace_broadcast(fn::Symbol)
     end
     return fn
 end
+# For TaskNames logging event
+Base.nameof(::ExpandedBroadcast{F}) where F = Symbol('.', nameof(F))
 
 to_namedtuple(;kwargs...) = (;kwargs...)
 

src/utils/logging-events.jl

@@ -117,21 +117,20 @@ end
 Records the raw (mutable) arguments of each submitted task.
 """
 struct TaskArguments end
-function (::TaskArguments)(ev::Event{:start})
-    if ev.category == :add_thunk
+(::TaskArguments)(ev::Event{:start}) = nothing
+function (ta::TaskArguments)(ev::Event{:finish})
+    if ev.category == :move
         args = Pair{Union{Symbol,Int},UInt}[]
-        for (idx, (pos, arg)) in enumerate(ev.timeline.args)
-            pos_idx = pos === nothing ? idx : pos
-            arg = Dagger.unwrap_weak_checked(arg)
-            if ismutable(arg)
-                push!(args, pos_idx => objectid(arg))
-            end
+        thunk_id = ev.id.thunk_id::Int
+        pos = ev.id.position::Union{Symbol,Int}
+        arg = ev.timeline.data
+        if ismutable(arg)
+            push!(args, pos => objectid(arg))
         end
-        return ev.id.thunk_id => args
+        return thunk_id => args
     end
     return
 end
-(ta::TaskArguments)(ev::Event{:finish}) = nothing
 
 """
     TaskArgumentMoves
@@ -148,8 +147,10 @@ function (ta::TaskArgumentMoves)(ev::Event{:start})
     if ev.category == :move
         data = ev.timeline.data
         if ismutable(data)
-            d = get!(Dict{Union{Int,Symbol},UInt}, ta.pre_move_args, ev.id.thunk_id)
-            d[ev.id.id] = objectid(data)
+            thunk_id = ev.id.thunk_id::Int
+            position = ev.id.position::Union{Symbol,Int}
+            d = get!(Dict{Union{Int,Symbol},UInt}, ta.pre_move_args, thunk_id)
+            d[position] = objectid(data)
         end
     end
     return
@@ -158,16 +159,18 @@ function (ta::TaskArgumentMoves)(ev::Event{:finish})
     if ev.category == :move
         post_data = ev.timeline.data
         if ismutable(post_data)
-            if haskey(ta.pre_move_args, ev.id.thunk_id)
-                d = ta.pre_move_args[ev.id.thunk_id]
-                if haskey(d, ev.id.id)
-                    pre_data = d[ev.id.id]
-                    return ev.id.thunk_id, ev.id.id, pre_data, objectid(post_data)
+            thunk_id = ev.id.thunk_id::Int
+            position = ev.id.position::Union{Symbol,Int}
+            if haskey(ta.pre_move_args, thunk_id)
+                d = ta.pre_move_args[thunk_id]
+                if haskey(d, position)
+                    pre_data = d[position]
+                    return thunk_id, position, pre_data, objectid(post_data)
                 else
-                    @warn "No TID $(ev.id.thunk_id), ID $(ev.id.id)"
+                    @warn "No TID $(thunk_id), Position $(position)"
                 end
             else
-                @warn "No TID $(ev.id.thunk_id)"
+                @warn "No TID $(thunk_id)"
             end
         end
     end