Reimplement rendering with dagre-d3 instead of viz.js

Before this commit, this patch relies on a JavaScript version of
GraphViz that was compiled from C. Even the minified version of
this resource was ~2.5M. The main motivation for switching away
from this library, however, is that this is a complete black box
of which we have absolutely no control. It is not at all extensible,
and if something breaks we will have a hard time understanding
why.

The new library, dagre-d3, is not perfect either. It does not
officially support clustering of nodes; for certain large graphs,
the clusters will have a lot of unnecessary whitespace. A few in
the dagre-d3 community are looking into a solution, but until then
we will have to live with this (minor) inconvenience.

Author: Andrew Or

core/src/main/resources/org/apache/spark/ui/static/d3.min.js

@@ -0,0 +1,98 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * Render a DAG that describes the RDDs for a given stage.
+ *
+ * Input: The content of a dot file, stored in the text of the "#viz-dot-file" element
+ * Output: An SVG that visualizes the DAG, stored in the "#viz-graph" element
+ *
+ * This relies on a custom implementation of dagre-d3, which can be found under
+ * http://github.com/andrewor14/dagre-d3/dist/dagre-d3.js. For more detail, please
+ * track the changes in that project after it was forked.
+ */
+function renderStageViz() {
+
+  // Parse the dot file and render it in an SVG
+  var dot = d3.select("#viz-dot-file").text();
+  var escaped_dot = dot
+    .replace(/&lt;/g, "<")
+    .replace(/&gt;/g, ">")
+    .replace(/&quot;/g, "\"");
+  var g = graphlibDot.read(escaped_dot);
+  var render = new dagreD3.render();
+  var svg = d3.select("#viz-graph");
+  svg.call(render, g);
+
+  // Set the appropriate SVG dimensions to ensure that all elements are displayed
+  var svgMargin = 20;
+  var boundingBox = svg.node().getBBox();
+  svg.style("width", (boundingBox.width + svgMargin) + "px");
+  svg.style("height", (boundingBox.height + svgMargin) + "px"); 
+
+  // Add style to clusters, nodes and edges
+  d3.selectAll("svg g.cluster rect")
+    .style("fill", "none")
+    .style("stroke", "#AADFFF")
+    .style("stroke-width", "4px")
+    .style("stroke-opacity", "0.5");
+  d3.selectAll("svg g.node rect")
+    .style("fill", "white")
+    .style("stroke", "black")
+    .style("stroke-width", "2px")
+    .style("fill-opacity", "0.8")
+    .style("stroke-opacity", "0.9");
+  d3.selectAll("svg g.edgePath path")
+    .style("stroke", "black")
+    .style("stroke-width", "2px");
+
+  // Add labels to clusters
+  d3.selectAll("svg g.cluster").each(function(cluster_data) {
+    var cluster = d3.select(this);
+    cluster.selectAll("rect").each(function(rect_data) {
+      var rect = d3.select(this);
+      // Shift the boxes up a little
+      rect.attr("y", toFloat(rect.attr("y")) - 10);
+      rect.attr("height", toFloat(rect.attr("height")) + 10);
+      var labelX = toFloat(rect.attr("x")) + toFloat(rect.attr("width")) - 5;
+      var labelY = toFloat(rect.attr("y")) + 15;
+      var labelText = cluster.attr("id").replace(/^cluster/, "").replace(/_.*$/, "");
+      cluster.append("text")
+        .attr("x", labelX)
+        .attr("y", labelY)
+        .attr("fill", "#AAAAAA")
+        .attr("font-size", "11px")
+        .attr("text-anchor", "end")
+        .text(labelText);
+    });
+  });
+
+  // We have shifted a few elements upwards, so we should fix the SVG views
+  var startX = -svgMargin;
+  var startY = -svgMargin;
+  var endX = toFloat(svg.style("width")) + svgMargin;
+  var endY = toFloat(svg.style("height")) + svgMargin;
+  var newViewBox = startX + " " + startY + " " + endX + " " + endY;
+  svg.attr("viewBox", newViewBox);
+
+}
+
+/* Helper method to convert attributes to numeric values. */
+function toFloat(f) {
+  return parseFloat(f.replace(/px$/, ""))
+}
+

core/src/main/resources/org/apache/spark/ui/static/dagre-d3.min.js

@@ -168,7 +168,10 @@ private[spark] object UIUtils extends Logging {
   }
 
   def vizHeaderNodes: Seq[Node] = {
-    <script src={prependBaseUri("/static/viz.js")}></script>
+    <script src={prependBaseUri("/static/d3.min.js")}></script>
+    <script src={prependBaseUri("/static/dagre-d3.min.js")}></script>
+    <script src={prependBaseUri("/static/graphlib-dot.min.js")}></script>
+    <script src={prependBaseUri("/static/spark-stage-viz.js")}></script>
   }
 
   /** Returns a spark page with correctly formatted headers */

core/src/main/resources/org/apache/spark/ui/static/graphlib-dot.min.js

@@ -41,19 +41,13 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
     if (graph.isEmpty) {
       return Seq.empty
     }
-    val viz = <div id="stage-viz">{VizGraph.makeDotFile(graph.get)}</div>
-    val script = {
-      <script type="text/javascript">
-        <xml:unparsed>
-          var dot = document.getElementById("stage-viz").innerHTML;
-          var dot = dot.replace(/&lt;/g, "<").replace(/&gt;/g, ">").replace(/&quot;/g, "\"");
-          console.log(dot);
-          var viz = Viz(dot, "svg", "dot");
-          document.getElementById("stage-viz").innerHTML = viz;
-        </xml:unparsed>
-      </script>
+    {
+      <div id="viz-dot-file" style="display:none">
+        {VizGraph.makeDotFile(graph.get)}
+      </div>
+      <svg id="viz-graph"></svg>
+      <script type="text/javascript">renderStageViz()</script>
     }
-    viz ++ script
   }
 
   def render(request: HttpServletRequest): Seq[Node] = {

core/src/main/resources/org/apache/spark/ui/static/spark-stage-viz.js

@@ -81,7 +81,10 @@ private[ui] object VizGraph {
         while (scopeIt.hasNext) {
           val scopeId = scopeIt.next()
           val scope = scopes.getOrElseUpdate(scopeId, new VizScope(scopeId))
-          scope.attachChildNode(node)
+          // Only attach this node to the innermost scope
+          if (!scopeIt.hasNext) {
+            scope.attachChildNode(node)
+          }
           // RDD scopes are hierarchical, with the outermost scopes ordered first
           // If there is not a previous scope, then this must be a root scope
           if (previousScope == null) {
@@ -106,48 +109,35 @@ private[ui] object VizGraph {
    */
   def makeDotFile(graph: VizGraph): String = {
     val dotFile = new StringBuilder
-    dotFile.append(
-      """
-        |digraph G {
-        |  node[fontsize="12", style="rounded, bold", shape="box"]
-        |  graph[labeljust="r", style="bold", color="#DDDDDD", fontsize="10"]
-      """.stripMargin.trim)
+    dotFile.append("digraph G {\n")
     //
     graph.rootScopes.foreach { scope =>
       dotFile.append(makeDotSubgraph(scope, "  "))
     }
     //
     graph.rootNodes.foreach { node =>
-      dotFile.append("  " + makeDotNode(node) + "\n")
+      dotFile.append(s"  ${makeDotNode(node)};\n")
     }
     //
     graph.edges.foreach { edge =>
-      dotFile.append("  " + edge.fromId + "->" + edge.toId + "\n")
+      dotFile.append(s"  ${edge.fromId}->${edge.toId};\n")
     }
     dotFile.append("}")
+    println(dotFile.toString())
     dotFile.toString()
   }
 
   /** */
   private def makeDotNode(node: VizNode): String = {
-    val dnode = new StringBuilder
-    dnode.append(node.id)
-    dnode.append(s""" [label="${node.name}"""")
-    if (node.isCached) {
-      dnode.append(s""", URL="/storage/rdd/?id=${node.id}", color="red"""")
-    }
-    dnode.append("]")
-    dnode.toString()
+    s"""${node.id} [label="${node.name}"]"""
   }
 
   /** */
   private def makeDotSubgraph(scope: VizScope, indent: String): String = {
     val subgraph = new StringBuilder
-    subgraph.append(indent + "subgraph cluster" + scope.id + " {\n")
-    subgraph.append(indent + "  label=\"" + scope.name + "\"\n")
-    subgraph.append(indent + "  fontcolor=\"#AAAAAA\"\n")
+    subgraph.append(indent + s"subgraph cluster${scope.id} {\n")
     scope.childrenNodes.foreach { node =>
-      subgraph.append(indent + "  " + makeDotNode(node) + "\n")
+      subgraph.append(indent + s"  ${makeDotNode(node)};\n")
     }
     scope.childrenScopes.foreach { cscope =>
       subgraph.append(makeDotSubgraph(cscope, indent + "  "))