Separate HTML generating code from listener

Author: Andrew Or

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

@@ -105,6 +105,9 @@ function toggleDagViz(forJob) {
  *   div#dag-viz-graph >
  *   svg >
  *   g#cluster_stage_[stageId]
+ *
+ * Note that the input metadata is populated by o.a.s.ui.UIUtils.showDagViz.
+ * Any changes in the input format here must be reflected there.
  */
 function renderDagViz(forJob) {
 

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

@@ -23,6 +23,7 @@ import java.util.{Locale, Date}
 import scala.xml.{Node, Text}
 
 import org.apache.spark.Logging
+import org.apache.spark.ui.viz.VizGraph
 
 /** Utility functions for generating XML pages with spark content. */
 private[spark] object UIUtils extends Logging {
@@ -329,4 +330,42 @@ private[spark] object UIUtils extends Logging {
       <div class="bar bar-running" style={startWidth}></div>
     </div>
   }
+
+  /** Return a "DAG visualization" DOM element that expands into a visualization for a stage. */
+  def showDagVizForStage(stageId: Int, graph: Option[VizGraph]): Seq[Node] = {
+    showDagViz(graph.toSeq, forJob = false)
+  }
+
+  /** Return a "DAG visualization" DOM element that expands into a visualization for a job. */
+  def showDagVizForJob(jobId: Int, graphs: Seq[VizGraph]): Seq[Node] = {
+    showDagViz(graphs, forJob = true)
+  }
+
+  /**
+   * Return a "DAG visualization" DOM element that expands into a visualization on the UI.
+   *
+   * This populates metadata necessary for generating the visualization on the front-end in
+   * a format that is expected by spark-dag-viz.js. Any changes in the format here must be
+   * reflected there.
+   */
+  private def showDagViz(graphs: Seq[VizGraph], forJob: Boolean): Seq[Node] = {
+    <div>
+      <span class="expand-dag-viz" onclick={s"toggleDagViz($forJob);"}>
+        <span class="expand-dag-viz-arrow arrow-closed"></span>
+        <strong>DAG visualization</strong>
+      </span>
+      <div id="dag-viz-graph"></div>
+      <div id="dag-viz-metadata">
+        {
+          graphs.map { g =>
+            <div class="stage-metadata" stageId={g.rootScope.id} style="display:none">
+              <div class="dot-file">{VizGraph.makeDotFile(g, forJob)}</div>
+              { g.incomingEdges.map { e => <div class="incoming-edge">{e.fromId},{e.toId}</div> } }
+              { g.outgoingEdges.map { e => <div class="outgoing-edge">{e.fromId},{e.toId}</div> } }
+            </div>
+          }
+        }
+      </div>
+    </div>
+  }
 }

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

@@ -309,7 +309,7 @@ private[ui] class JobPage(parent: JobsTab) extends WebUIPage("job") {
       content ++= makeTimeline(activeStages ++ completedStages ++ failedStages,
           executorListener.executorIdToData, appStartTime)
 
-      content ++= vizListener.showVizElementForJob(jobId)
+      content ++= UIUtils.showDagVizForJob(jobId, vizListener.getVizGraphsForJob(jobId))
 
       if (shouldShowActiveStages) {
         content ++= <h4 id="active">Active Stages ({activeStages.size})</h4> ++

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

@@ -171,6 +171,8 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
           </div>
         </div>
 
+      val dagViz = UIUtils.showDagVizForStage(stageId, vizListener.getVizGraphForStage(stageId))
+
       val accumulableHeaders: Seq[String] = Seq("Accumulable", "Value")
       def accumulableRow(acc: AccumulableInfo): Elem =
         <tr><td>{acc.name}</td><td>{acc.value}</td></tr>
@@ -436,7 +438,7 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
       val content =
         summary ++
         showAdditionalMetrics ++
-        vizListener.showVizElementForStage(stageId) ++
+        dagViz ++
         <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/viz/VisualizationListener.scala

@@ -36,47 +36,15 @@ private[ui] class VisualizationListener(conf: SparkConf) extends SparkListener {
   private val retainedStages =
     conf.getInt("spark.ui.retainedStages", SparkUI.DEFAULT_RETAINED_STAGES)
 
-  /** Construct a "DAG 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 "DAG visualization" DOM element that expands into a visualization for a job. */
-  def showVizElementForJob(jobId: Int): Seq[Node] = {
-    showVizElement(getVizGraphsForJob(jobId), forJob = true)
-  }
-
-  /** Construct a "DAG 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-dag-viz" onclick={s"toggleDagViz($forJob);"}>
-        <span class="expand-dag-viz-arrow arrow-closed"></span>
-        <strong>DAG visualization</strong>
-      </span>
-      <div id="dag-viz-graph"></div>
-      <div id="dag-viz-metadata">
-        {
-          graphs.map { g =>
-            <div class="stage-metadata" stageId={g.rootScope.id} style="display:none">
-              <div class="dot-file">{VizGraph.makeDotFile(g, forJob)}</div>
-              { g.incomingEdges.map { e => <div class="incoming-edge">{e.fromId},{e.toId}</div> } }
-              { g.outgoingEdges.map { e => <div class="outgoing-edge">{e.fromId},{e.toId}</div> } }
-            </div>
-          }
-        }
-      </div>
-    </div>
-  }
-
   /** Return the graph metadata for the given stage, or None if no such information exists. */
-  private def getVizGraphsForJob(jobId: Int): Seq[VizGraph] = {
+  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] = {
+  def getVizGraphForStage(stageId: Int): Option[VizGraph] = {
     stageIdToGraph.get(stageId)
   }