Introduce visualization to the Job Page

This includes a generalization of the visualization previously
displayed on the stage page. More functionality is needed in
JavaScript to prevent the job visualization from looking too
cluttered. This is still WIP.

Author: Andrew Or

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

@@ -15,84 +15,70 @@
  * limitations under the License.
  */
 
-var stageVizIsRendered = false
+var VizDefaults = {
+  stageVizOffset: 160
+};
 
 /*
- * Render or remove the stage visualization on the UI.
+ * Show or hide the RDD DAG visualization on the UI.
  * This assumes that the visualization is stored in the "#viz-graph" element.
  */
-function toggleStageViz() {
+function toggleShowViz(forJob) {
   $(".expand-visualization-arrow").toggleClass('arrow-closed');
   $(".expand-visualization-arrow").toggleClass('arrow-open');
-  var shouldRender = $(".expand-visualization-arrow").hasClass("arrow-open");
-  if (shouldRender) {
-    // If the viz is already rendered, just show it
-    if (stageVizIsRendered) {
-      $("#viz-graph").show();
-    } else {
-      renderStageViz();
-      stageVizIsRendered = true;
+  var show = $(".expand-visualization-arrow").hasClass("arrow-open");
+  if (show) {
+    var shouldRender = d3.select("#viz-graph svg").empty();
+    if (shouldRender) {
+      renderViz();
+      styleViz(forJob);
     }
+    $("#viz-graph").show();
   } else {
-    // Instead of emptying the element once and for all, cache it for use
-    // again later in case we want to expand the visualization again
     $("#viz-graph").hide();
   }
 }
 
 /*
- * Render a DAG that describes the RDDs for a given stage.
+ * Render a DAG visualization that describes RDDs in a stage or a job.
  *
- * Input: The content of a dot file, stored in the text of the "#viz-dot-file" element
+ * Input: The contents of the relevant dot files, stored in the "#viz-dot-files" 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
  * https://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() {
+function renderViz() {
 
   // If there is not a dot file to render, report error
-  if (d3.select("#viz-dot-file").empty()) {
+  if (d3.select("#viz-dot-files").empty()) {
     d3.select("#viz-graph")
       .append("div")
       .text("No visualization information available for this stage.");
     return;
   }
 
-  // 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").append("svg");
-  svg.call(render, g);
+
+  // Each div in #viz-dot-files stores the content of one dot file
+  // Each dot file is used to generate the visualization for a stage
+  d3.selectAll("#viz-dot-files div").each(function(d, i) {
+    var div = d3.select(this);
+    var stageId = div.attr("name");
+    var dot = div.text();
+    var container = svg.append("g").attr("id", "graph_" + stageId);
+    // Move the container so it doesn't overlap with the existing ones
+    container.attr("transform", "translate(" + VizDefaults.stageVizOffset * i + ", 0)");
+    renderDot(dot, container);
+  });
 
   // 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);
@@ -103,12 +89,10 @@ function renderStageViz() {
       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(/_.*$/, "");
+      var labelText = cluster.attr("name").replace(/^cluster_/, "");
       cluster.append("text")
         .attr("x", labelX)
         .attr("y", labelY)
-        .attr("fill", "#AAAAAA")
-        .attr("font-size", "11px")
         .attr("text-anchor", "end")
         .text(labelText);
     });
@@ -123,6 +107,71 @@ function renderStageViz() {
   svg.attr("viewBox", newViewBox);
 }
 
+/*
+ *
+ */
+function renderDot(dot, container) {
+  var escaped_dot = dot
+    .replace(/&lt;/g, "<")
+    .replace(/&gt;/g, ">")
+    .replace(/&quot;/g, "\"");
+  var g = graphlibDot.read(escaped_dot);
+  var renderer = new dagreD3.render();
+  renderer(container, g);
+}
+
+/*
+ * Style the visualization we just rendered.
+ * We apply a different style depending on whether this is a stage or a job.
+ */
+function styleViz(forJob) {
+  if (forJob) {
+    styleJobViz();
+  } else {
+    styleStageViz();
+  }
+}
+
+function styleJobViz() {
+  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");
+  d3.selectAll("svg g.cluster text")
+    .attr("fill", "#AAAAAA")
+    .attr("font-size", "11px")
+}
+
+function styleStageViz() {
+  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");
+  d3.selectAll("svg g.cluster text")
+    .attr("fill", "#AAAAAA")
+    .attr("font-size", "11px")
+}
+
 /* 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/spark-stage-viz.js renamed to core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.js

@@ -47,6 +47,8 @@ class OrderedRDDFunctions[K : Ordering : ClassTag,
     self: RDD[P])
   extends Logging with Serializable
 {
+  // TODO: Don't forget to scope me later
+
   private val ordering = implicitly[Ordering[K]]
 
   /**

core/src/main/scala/org/apache/spark/rdd/OrderedRDDFunctions.scala

@@ -171,7 +171,7 @@ private[spark] object UIUtils extends Logging {
     <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>
+    <script src={prependBaseUri("/static/spark-dag-viz.js")}></script>
   }
 
   /** Returns a spark page with correctly formatted headers */

core/src/main/scala/org/apache/spark/ui/UIUtils.scala

@@ -29,6 +29,7 @@ import org.apache.spark.ui.{UIUtils, WebUIPage}
 /** Page showing statistics and stage list for a given job */
 private[ui] class JobPage(parent: JobsTab) extends WebUIPage("job") {
   private val listener = parent.listener
+  private val vizListener = parent.vizListener
 
   def render(request: HttpServletRequest): Seq[Node] = {
     listener.synchronized {
@@ -154,6 +155,7 @@ private[ui] class JobPage(parent: JobsTab) extends WebUIPage("job") {
         </div>
 
       var content = summary
+      content ++= vizListener.showVizElementForJob(jobId)
       if (shouldShowActiveStages) {
         content ++= <h4 id="active">Active Stages ({activeStages.size})</h4> ++
           activeStagesTable.toNodeSeq
@@ -174,7 +176,7 @@ private[ui] class JobPage(parent: JobsTab) extends WebUIPage("job") {
         content ++= <h4 id ="failed">Failed Stages ({failedStages.size})</h4> ++
           failedStagesTable.toNodeSeq
       }
-      UIUtils.headerSparkPage(s"Details for Job $jobId", content, parent)
+      UIUtils.headerSparkPage(s"Details for Job $jobId", content, parent, showVisualization = true)
     }
   }
 }

core/src/main/scala/org/apache/spark/ui/jobs/JobPage.scala

@@ -25,7 +25,9 @@ private[ui] class JobsTab(parent: SparkUI) extends SparkUITab(parent, "jobs") {
   val sc = parent.sc
   val killEnabled = parent.killEnabled
   def isFairScheduler: Boolean = listener.schedulingMode.exists(_ == SchedulingMode.FAIR)
+  // TODO: rename this listener
   val listener = parent.jobProgressListener
+  val vizListener = parent.visualizationListener
 
   attachPage(new AllJobsPage(this))
   attachPage(new JobPage(this))

core/src/main/scala/org/apache/spark/ui/jobs/JobsTab.scala

@@ -36,32 +36,6 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
   private val progressListener = parent.progressListener
   private val vizListener = parent.vizListener
 
-  /**
-   * Return a DOM element containing the "Show Visualization" toggle that, if enabled,
-   * renders the visualization for the stage. If there is no visualization information
-   * available for this stage, an appropriate message is displayed to the user.
-   */
-  private def showVizElement(stageId: Int): Seq[Node] = {
-    val graph = vizListener.getVizGraph(stageId)
-    <div>
-      <span class="expand-visualization" onclick="render();">
-        <span class="expand-visualization-arrow arrow-closed"></span>
-        <strong>Show Visualization</strong>
-      </span>
-      <div id="viz-graph"></div>
-      <script type="text/javascript">
-        {Unparsed(s"function render() { toggleStageViz(); }")}
-      </script>
-      {
-        if (graph.isDefined) {
-          <div id="viz-dot-file" style="display:none">
-            {VizGraph.makeDotFile(graph.get)}
-          </div>
-        }
-      }
-    </div>
-  }
-
   def render(request: HttpServletRequest): Seq[Node] = {
     progressListener.synchronized {
       val parameterId = request.getParameter("id")
@@ -462,7 +436,7 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
       val content =
         summary ++
         showAdditionalMetrics ++
-        showVizElement(stageId) ++
+        vizListener.showVizElementForStage(stageId) ++
         <h4>Summary Metrics for {numCompleted} Completed Tasks</h4> ++
         <div>{summaryTable.getOrElse("No tasks have reported metrics yet.")}</div> ++
         <h4>Aggregated Metrics by Executor</h4> ++ executorTable.toNodeSeq ++

core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala

@@ -18,39 +18,81 @@
 package org.apache.spark.ui.viz
 
 import scala.collection.mutable
+import scala.xml.{Node, Unparsed}
 
 import org.apache.spark.SparkConf
 import org.apache.spark.scheduler._
 import org.apache.spark.ui.SparkUI
 
 /**
- * A SparkListener that constructs a graph of the RDD DAG for each stage.
- * This graph will be used for rendering visualization in the UI later.
+ * A SparkListener that constructs RDD DAG visualization for the UI.
  */
 private[ui] class VisualizationListener(conf: SparkConf) extends SparkListener {
-
-  // A list of stage IDs to track the order in which stages are inserted
-  private val stageIds = new mutable.ArrayBuffer[Int]
-
-  // Stage ID -> graph metadata for the stage
+  private val jobIdToStageIds = new mutable.HashMap[Int, Seq[Int]]
   private val stageIdToGraph = new mutable.HashMap[Int, VizGraph]
+  private val stageIds = new mutable.ArrayBuffer[Int]
 
-  // How many stages to retain graph metadata for
+  // How many jobs or stages to retain graph metadata for
   private val retainedStages =
     conf.getInt("spark.ui.retainedStages", SparkUI.DEFAULT_RETAINED_STAGES)
 
+  /** Construct a "Show visualization" DOM element that expands into a visualization for a stage. */
+  def showVizElementForStage(stageId: Int): Seq[Node] = {
+    showVizElement(getVizGraphForStage(stageId).toSeq, forJob = false)
+  }
+
+  /** Construct a "Show visualization" DOM element that expands into a visualization for a job. */
+  def showVizElementForJob(jobId: Int): Seq[Node] = {
+    showVizElement(getVizGraphsForJob(jobId), forJob = true)
+  }
+
+  /** Construct a "Show visualization" DOM element that expands into a visualization on the UI. */
+  private def showVizElement(graphs: Seq[VizGraph], forJob: Boolean): Seq[Node] = {
+    <div>
+      <span class="expand-visualization" onclick="render();">
+        <span class="expand-visualization-arrow arrow-closed"></span>
+        <strong>Show visualization</strong>
+      </span>
+      <div id="viz-graph"></div>
+      <div id="viz-dot-files">
+        {
+          graphs.map { g =>
+            <div name={g.rootScope.id} style="display:none">
+              {VizGraph.makeDotFile(g, forJob)}
+            </div>
+          }
+        }
+      </div>
+    </div> ++
+    <script type="text/javascript">
+      {Unparsed(s"function render() { toggleShowViz($forJob); }")}
+    </script>
+  }
+
   /** Return the graph metadata for the given stage, or None if no such information exists. */
-  def getVizGraph(stageId: Int): Option[VizGraph] = {
+  private def getVizGraphsForJob(jobId: Int): Seq[VizGraph] = {
+    jobIdToStageIds.get(jobId)
+      .map { sids => sids.flatMap { sid => stageIdToGraph.get(sid) } }
+      .getOrElse { Seq.empty }
+  }
+
+  /** Return the graph metadata for the given stage, or None if no such information exists. */
+  private def getVizGraphForStage(stageId: Int): Option[VizGraph] = {
     stageIdToGraph.get(stageId)
   }
 
-  override def onStageCompleted(stageCompleted: SparkListenerStageCompleted): Unit = synchronized {
-    val stageId = stageCompleted.stageInfo.stageId
-    val rddInfos = stageCompleted.stageInfo.rddInfos
-    val vizGraph = VizGraph.makeVizGraph(rddInfos)
-    stageIdToGraph(stageId) = vizGraph
+  /** On job start, construct a VizGraph for each stage in the job for display later. */
+  override def onJobStart(jobStart: SparkListenerJobStart): Unit = synchronized {
+    val jobId = jobStart.jobId
+    val stageInfos = jobStart.stageInfos
+
+    stageInfos.foreach { stageInfo =>
+      stageIds += stageInfo.stageId
+      stageIdToGraph(stageInfo.stageId) = VizGraph.makeVizGraph(stageInfo)
+    }
+    jobIdToStageIds(jobId) = stageInfos.map(_.stageId)
 
-    // Remove metadata for old stages
+    // Remove graph metadata for old stages
     if (stageIds.size >= retainedStages) {
       val toRemove = math.max(retainedStages / 10, 1)
       stageIds.take(toRemove).foreach { id => stageIdToGraph.remove(id) }