logging: Add TaskFunctionNames event

Author: jpsamaroo

ext/GraphVizExt.jl

@@ -84,7 +84,7 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
             if category == :add_thunk && kind == :start
                 id::NamedTuple
                 taskdeps = logs[w][:taskdeps][idx]::Pair{Int,Vector{Int}}
-                taskname = logs[w][:tasknames][idx]::String
+                taskname = logs[w][:taskfuncnames][idx]::String
                 tid, deps = taskdeps
                 v = get!(tid_to_vertex, tid) do
                     add_vertex!(g)

ext/JSON3Ext.jl

@@ -37,8 +37,7 @@ function logs_to_chrome_trace(logs::Dict)
             if !haskey(execution_logs, tid)
                 execution_logs[tid] = Dict{Symbol,Any}()
             end
-            taskname = logs[w][:tasknames][start_idx]
-            fname = first(split(taskname, ' '))
+            fname = logs[w][:taskfuncnames][start_idx]
             execution_logs[tid][:name] = fname
         end
     end
@@ -63,4 +62,4 @@ function Dagger.show_logs(io::IO, logs::Dict, ::Val{:chrome_trace})
     JSON3.write(io, logs_to_chrome_trace(logs))
 end
 
-end # module JSON3Ext
+end # module JSON3Ext

ext/PlotsExt.jl

@@ -35,8 +35,8 @@ function logs_to_df(logs::Dict, ::Val{:execution}; colors=_default_colors, name_
             kind = logs[w][:core][idx].kind::Symbol
             if category == :add_thunk && kind == :start
                 tid = logs[w][:id][idx].thunk_id::Int
-                if haskey(logs[w], :tasknames)
-                    fn_names[tid] = first(split(logs[w][:tasknames][idx]::String, ' '))
+                if haskey(logs[w], :taskfuncnames)
+                    fn_names[tid] = logs[w][:taskfuncnames][idx]::String
                 else
                     @warn "Task names missing from logs"
                     fn_names[tid] = "unknown"

src/utils/logging-events.jl

@@ -119,6 +119,24 @@ func_name(f::Function) = nameof(f)
 func_name(x) = repr(x)
 func_name(::Dagger.ExpandedBroadcast{F}) where F = Symbol('.', nameof(F))
 
+"""
+    TaskFunctionNames
+
+Records the function name of each task.
+"""
+struct TaskFunctionNames end
+function (::TaskFunctionNames)(ev::Event{:start})
+    if ev.category == :add_thunk
+        f = Dagger.chunktype(ev.timeline.f)
+        if hasproperty(f, :instance) && isdefined(f, :instance)
+            f = f.instance
+        end
+        return String(func_name(f))
+    end
+    return
+end
+(td::TaskFunctionNames)(ev::Event{:finish}) = nothing
+
 """
     TaskArguments
 

src/utils/logging.jl

@@ -10,6 +10,7 @@ Extra events:
 - `timeline::Bool`: Enables raw "timeline" values, which are event-specific; not recommended except for debugging
 - `all_task_deps::Bool`: Enables all task dependency-related logging
 - `tasknames::Bool`: Enables generating unique task names for each task
+- `taskfuncnames::Bool`: Enables reporting of task function names for each task
 - `taskdeps::Bool`: Enables reporting of upstream task dependencies (as task IDs) for each task argument
 - `taskargs::Bool`: Enables reporting of upstream non-task dependencies (as `objectid` hash) for each task argument
 - `taskargmoves::Bool`: Enables reporting of copies of upstream dependencies (as original and copy `objectid` hashes) for each task argument
@@ -22,6 +23,7 @@ function enable_logging!(;metrics::Bool=true,
                           timeline::Bool=false,
                           all_task_deps::Bool=false,
                           tasknames::Bool=true,
+                          taskfuncnames::Bool=false,
                           taskdeps::Bool=true,
                           taskargs::Bool=false,
                           taskargmoves::Bool=false,
@@ -36,7 +38,7 @@ function enable_logging!(;metrics::Bool=true,
         ml[:timeline] = TimespanLogging.Events.TimelineMetrics()
     end
     if all_task_deps
-        tasknames = true
+        taskfuncnames = true
         taskdeps = true
         taskargs = true
         taskargmoves = true
@@ -47,6 +49,9 @@ function enable_logging!(;metrics::Bool=true,
     if tasknames
         ml[:tasknames] = Dagger.Events.TaskNames()
     end
+    if taskfuncnames
+        ml[:taskfuncnames] = Dagger.Events.TaskFunctionNames()
+    end
     if taskdeps
         ml[:taskdeps] = Dagger.Events.TaskDependencies()
     end