From 27c410079e0755bfdcb51668ac9c55dde0c206d8 Mon Sep 17 00:00:00 2001
From: Mateusz Jakub Fila <mateusz.jakub.fila@cern.ch>
Date: Mon, 19 Aug 2024 08:40:35 +0200
Subject: [PATCH] add show_logs for graphviz

---
 docs/src/logging-visualization.md |  1 +
 ext/GraphVizExt.jl                | 43 +++++++++++++++++++++++++++----
 2 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/docs/src/logging-visualization.md b/docs/src/logging-visualization.md
index 928f816ae..f6717c0ba 100644
--- a/docs/src/logging-visualization.md
+++ b/docs/src/logging-visualization.md
@@ -18,6 +18,7 @@ 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, :graphviz)` to write a Graphviz dot graph of executed tasks and their dependencies
 - `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:
diff --git a/ext/GraphVizExt.jl b/ext/GraphVizExt.jl
index 1b41f9fa1..3169d2859 100644
--- a/ext/GraphVizExt.jl
+++ b/ext/GraphVizExt.jl
@@ -38,6 +38,30 @@ tab20_colors = [
 ]
 _default_colors = tab20_colors
 
+"""
+    Dagger.show_logs(io::IO, logs::Dict, ::Val{:graphviz}; disconnected=false,
+                       color_by=:fn, times::Bool=true, times_digits::Integer=3)
+
+Write a graph of the task dependencies and data dependencies in `logs` to dot format
+
+Requires the following events enabled in `enable_logging!`: `taskdeps`, `tasknames`, `taskargs`, `taskargmoves`
+
+Options:
+- `disconnected`: If `true`, render disconnected vertices (tasks or arguments without upstream/downstream dependencies)
+- `color_by`: How to color tasks; if `:fn`, then color by unique function name, if `:proc`, then color by unique processor
+- `times`: If `true`, annotate each task with its start and finish times
+- `times_digits`: Number of digits to display in the time annotations
+- `colors`: A list of colors to use for coloring tasks
+- `name_to_color`: A function that maps task names to colors
+"""
+function Dagger.show_logs(io::IO, logs::Dict, ::Val{:graphviz}; disconnected=false,
+                          color_by=:fn, times::Bool=true, times_digits::Integer=3,
+                          colors=_default_colors, name_to_color=_name_to_color)
+    dot = logs_to_dot(logs; disconnected, times, times_digits,
+                      color_by, colors, name_to_color)
+    println(io, dot)
+end
+
 """
     Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                        color_by=:fn, layout_engine="dot",
@@ -50,14 +74,26 @@ Requires the `all_task_deps` event enabled in `enable_logging!`
 Options:
 - `disconnected`: If `true`, render disconnected vertices (tasks or arguments without upstream/downstream dependencies)
 - `color_by`: How to color tasks; if `:fn`, then color by unique function name, if `:proc`, then color by unique processor
-- `layout_engine`: The layout engine to use for GraphViz
+- `layout_engine`: The layout engine to use for GraphViz rendering
 - `times`: If `true`, annotate each task with its start and finish times
 - `times_digits`: Number of digits to display in the time annotations
+- `colors`: A list of colors to use for coloring tasks
+- `name_to_color`: A function that maps task names to colors
 """
 function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
                             color_by=:fn, layout_engine="dot",
                             times::Bool=true, times_digits::Integer=3,
                             colors=_default_colors, name_to_color=_name_to_color)
+    dot = logs_to_dot(logs; disconnected, times, times_digits,
+                            color_by, colors, name_to_color)
+    gv = GraphViz.Graph(dot)
+    GraphViz.layout!(gv; engine=layout_engine)
+    return gv
+end
+
+function logs_to_dot(logs::Dict; disconnected=false, color_by=:fn,
+                     times::Bool=true, times_digits::Integer=3,
+                     colors=_default_colors, name_to_color=_name_to_color)
     # Lookup all relevant task/argument dependencies and values in logs
     g = SimpleDiGraph()
 
@@ -397,10 +433,7 @@ function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
 
     # Generate the final graph
     str *= "}\n"
-    gv = GraphViz.Graph(str)
-    GraphViz.layout!(gv; engine=layout_engine)
-
-    return gv
+    return str
 end
 
 end