Add render/show_logs to chrome_trace (#540)

Author: m-fila

Diff for: Project.toml

@@ -46,16 +46,19 @@ julia = "1.8"
 [extensions]
 GraphVizExt = "GraphViz"
 GraphVizSimpleExt = "Colors"
+JSON3Ext = "JSON3"
 PlotsExt = ["DataFrames", "Plots"]
 
 [extras]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [weakdeps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

Diff for: docs/src/logging-visualization.md

@@ -17,9 +17,13 @@ and `render_logs` take not a `Val` but a raw `Symbol`, which will be internally
 converted to a `Val` for dispatch purposes
 (i.e. `render_logs(logs::Dict, :myrenderer)` -> `render_logs(logs, Val{:myrenderer}())`).
 
+Built-in `IO` support exists for:
+- `show_logs(io, logs, :chrome_trace)` to write a task execution timeline in the chrome-trace format (view in [perfetto web UI](https://ui.perfetto.dev/) or `about:tracing` in a chrome-based browser)
+
 Built-in rendering support exists for:
 - `render_logs(logs, :graphviz)` to generate a graph diagram of executed tasks and their dependencies
 - `render_logs(logs, :plots_gantt)` to generate a Gantt chart of task execution across all processors
+- `render_logs(logs, :chrome_trace)` to generate a `String` with chrome-trace formatted task execution timeline
 
 The latter (`MultiEventLog`) allows for continuously rendering logs as they're
 generated, permitting real-time visualization of Dagger's operations. This

Diff for: ext/JSON3Ext.jl

@@ -0,0 +1,66 @@
+module JSON3Ext
+
+if isdefined(Base, :get_extension)
+    using JSON3
+else
+    using ..JSON3
+end
+
+using Dagger
+
+function logs_to_chrome_trace(logs::Dict)
+    execution_logs = Dict{Int,Any}()
+    add_unknown_procs_metadata = false
+    Dagger.logs_event_pairs(logs) do w, start_idx, finish_idx
+        category = logs[w][:core][start_idx].category
+        if category == :compute
+            tid = logs[w][:id][start_idx].thunk_id # thunk id
+            if !haskey(execution_logs, tid)
+                execution_logs[tid] = Dict{Symbol,Any}()
+            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
+            if proc isa Dagger.ThreadProc
+                execution_logs[tid][:pid] = proc.owner
+                execution_logs[tid][:tid] = proc.tid # thread id
+            else
+                @warn "Compute event for [$tid] executed on non-Dagger.ThreadProc processor. Assigning unknown pid and tid"
+                execution_logs[tid][:pid] = -1
+                execution_logs[tid][:tid] = -1
+                add_unknown_procs_metadata = true
+            end
+        elseif category == :add_thunk
+            tid = logs[w][:id][start_idx].thunk_id
+            if !haskey(execution_logs, tid)
+                execution_logs[tid] = Dict{Symbol,Any}()
+            end
+            taskname = logs[w][:tasknames][start_idx]
+            fname = first(split(taskname, ' '))
+            execution_logs[tid][:name] = fname
+        end
+    end
+    events = Vector{Dict{Symbol,Any}}()
+    for (_, v) in execution_logs
+        v[:ph] = "X"
+        v[:cat] = "compute"
+        push!(events, v)
+    end
+    if add_unknown_procs_metadata
+        push!(events, Dict(:name => "process_name", :ph => "M", :cat => "__metadata", :pid => -1, :args => Dict(:name => "Unknown")))
+        push!(events, Dict(:name => "thread_name", :ph => "M", :cat => "__metadata", :pid => -1, :tid => -1, :args => Dict(:name => "Unknown")))
+    end
+    return Dict(:traceEvents => events)
+end
+
+function Dagger.render_logs(logs::Dict, ::Val{:chrome_trace})
+    return JSON3.write(logs_to_chrome_trace(logs))
+end
+
+function Dagger.show_logs(io::IO, logs::Dict, ::Val{:chrome_trace})
+    JSON3.write(io, logs_to_chrome_trace(logs))
+end
+
+end # module JSON3Ext

Diff for: src/Dagger.jl

@@ -115,6 +115,9 @@ function __init__()
                 # Gantt chart HTTP server
                 include("ui/gantt-mux.jl")
             end
+            @require JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" begin
+                include(joinpath(dirname(@__DIR__), "ext", "JSON3Ext.jl"))
+            end
         end
         # TODO: Move to Pkg extensions
         @require ProfileSVG="132c30aa-f267-4189-9183-c8a63c7e05e6" begin

Diff for: test/Project.toml

@@ -5,6 +5,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"

Diff for: test/logging.jl

@@ -1,6 +1,6 @@
 import TimespanLogging
 import TimespanLogging: Timespan, Event, Events, LocalEventLog, MultiEventLog
-import Colors, GraphViz, DataFrames, Plots
+import Colors, GraphViz, DataFrames, Plots, JSON3
 
 @testset "Logging" begin
     @testset "LocalEventLog" begin
@@ -180,6 +180,9 @@ import Colors, GraphViz, DataFrames, Plots
             # PlotsExt
             @test Dagger.render_logs(logs, :plots_gantt) !== nothing
 
+            # JSON3Ext
+            @test Dagger.render_logs(logs, :chrome_trace) !== nothing
+            
             Dagger.disable_logging!()
         end
     end