Merge pull request #570 from m-fila/chrome_trace_names

use user given task names in chrome trace

Author: jpsamaroo

ext/JSON3Ext.jl

@@ -10,6 +10,8 @@ using Dagger
 
 function logs_to_chrome_trace(logs::Dict)
     execution_logs = Dict{Int,Any}()
+    tid_to_uid = Dict{Int,UInt}()
+    uid_to_name = Dict{UInt,String}()
     add_unknown_procs_metadata = false
     Dagger.logs_event_pairs(logs) do w, start_idx, finish_idx
         category = logs[w][:core][start_idx].category
@@ -20,9 +22,9 @@ function logs_to_chrome_trace(logs::Dict)
             end
             t_start = logs[w][:core][start_idx].timestamp / 1e3 # us
             t_stop = logs[w][:core][finish_idx].timestamp / 1e3 # us
-            proc = logs[w][:id][start_idx].processor
             execution_logs[tid][:ts] = t_start
             execution_logs[tid][:dur] = t_stop - t_start
+            proc = logs[w][:id][start_idx].processor
             if proc isa Dagger.ThreadProc
                 execution_logs[tid][:pid] = proc.owner
                 execution_logs[tid][:tid] = proc.tid # thread id
@@ -37,14 +39,37 @@ function logs_to_chrome_trace(logs::Dict)
             if !haskey(execution_logs, tid)
                 execution_logs[tid] = Dict{Symbol,Any}()
             end
+            # auto name
             fname = logs[w][:taskfuncnames][start_idx]
             execution_logs[tid][:name] = fname
+            # uid-tid mapping for user task name
+            if haskey(logs[w], :taskuidtotid)
+                uid_tid = logs[w][:taskuidtotid][start_idx]
+                if uid_tid !== nothing
+                    uid, tid = uid_tid::Pair{UInt,Int}
+                    tid_to_uid[tid] = uid
+                end
+            end
+        elseif category == :data_annotation
+            # user task name
+            id = logs[w][:id][start_idx]::NamedTuple
+            name = String(id.name)
+            obj = id.objectid::Dagger.LoggedMutableObject
+            objid = obj.objid
+            uid_to_name[objid] = name
         end
     end
     events = Vector{Dict{Symbol,Any}}()
-    for (_, v) in execution_logs
+    for (tid, v) in execution_logs
         v[:ph] = "X"
         v[:cat] = "compute"
+        # replace auto name with user task name if present
+        if haskey(tid_to_uid, tid)
+            uid = tid_to_uid[tid]
+            if haskey(uid_to_name, uid)
+                v[:name] = uid_to_name[uid]
+            end
+        end
         push!(events, v)
     end
     if add_unknown_procs_metadata

src/utils/logging.jl

@@ -136,8 +136,7 @@ utilize for display purposes.
 function logs_annotate!(ctx::Context, arg, name::Union{String,Symbol})
     ismutable(arg) || throw(ArgumentError("Argument must be mutable to be annotated"))
     Dagger.TimespanLogging.timespan_start(ctx, :data_annotation, (;objectid=objectid_or_chunkid(arg), name), nothing)
-    # TODO: Remove redundant log event
-    Dagger.TimespanLogging.timespan_finish(ctx, :data_annotation, nothing, nothing)
+    Dagger.TimespanLogging.timespan_finish(ctx, :data_annotation, (;objectid=objectid_or_chunkid(arg), name), nothing)
 end
 logs_annotate!(arg, name::Union{String,Symbol}) =
     logs_annotate!(Dagger.Sch.eager_context(), arg, name)