wire up task summary

Author: squito

core/src/main/scala/org/apache/spark/status/api/v1/AllStagesResource.scala

@@ -146,9 +146,9 @@ private[v1] object AllStagesResource {
     )
   }
 
-  def taskMetricDistributions(allTaskData: Seq[TaskUIData], quantiles: Array[Double]): TaskMetricDistributions = {
+  def taskMetricDistributions(allTaskData: Iterable[TaskUIData], quantiles: Array[Double]): TaskMetricDistributions = {
 
-    val rawMetrics = allTaskData.flatMap{_.taskMetrics}
+    val rawMetrics = allTaskData.flatMap{_.taskMetrics}.toSeq
 
     def getMetric[T](data: Seq[T], f: T => Double): IndexedSeq[Double] =
       Distribution(data.map{d=> f(d)}).get.getQuantiles(quantiles)

core/src/main/scala/org/apache/spark/status/api/v1/JsonRootResource.scala

@@ -73,6 +73,11 @@ private[v1] class JsonRootResource extends UIRootFromServletContext {
   }
 
 
+  @Path("applications/{appId}/stages/{stageId: \\d+}/{attemptId: \\d+}/taskSummary")
+  def getStageAttemptTaskSummary(): StageTaskSummary = {
+    new StageTaskSummary(uiRoot)
+  }
+
   @Path("applications/{appId}/storage/rdd")
   def getRdds(): AllRDDResource = {
     new AllRDDResource(uiRoot)

core/src/main/scala/org/apache/spark/status/api/v1/StageTaskSummary.scala

@@ -0,0 +1,69 @@
+/*
+ * 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.
+ */
+package org.apache.spark.status.api.v1
+
+import javax.ws.rs.core.MediaType
+import javax.ws.rs._
+
+import org.apache.spark.SparkException
+import org.apache.spark.SparkException
+
+@Produces(Array(MediaType.APPLICATION_JSON))
+private[v1] class StageTaskSummary(uiRoot: UIRoot) {
+
+  @GET
+  def stageData(
+    @PathParam("appId") appId: String,
+    @PathParam("stageId") stageId: Int,
+    @PathParam("attemptId") attemptId: Int,
+    @DefaultValue("0.05,0.25,0.5,0.75,0.95") @QueryParam("quantiles") quantileString: String
+  ): TaskMetricDistributions = {
+    uiRoot.withSparkUI(appId) { ui =>
+      val listener = ui.stagesTab.listener
+      val stageAndStatus = AllStagesResource.stagesAndStatus(ui)
+      val oneStage = stageAndStatus.flatMap { case (status, stages) =>
+        val matched = stages.find { stage =>
+          stage.stageId == stageId && stage.attemptId == attemptId
+        }
+        matched.map { status -> _ }
+      }.headOption
+      oneStage match {
+        case Some((status, stageInfo)) =>
+          val stageUiData = listener.synchronized {
+            listener.stageIdToData.get((stageInfo.stageId, stageInfo.attemptId)).
+              getOrElse(throw new SparkException("failed to get full stage data for stage: " +
+              stageInfo.stageId + ":" + stageInfo.attemptId)
+              )
+          }
+          val quantiles = quantileString.split(",").map{_.toDouble}
+          println("quantiles = " + quantiles.mkString(","))
+          AllStagesResource.taskMetricDistributions(stageUiData.taskData.values, quantiles)
+        case None =>
+          val stageAttempts = stageAndStatus.flatMap{ case (status, stages) =>
+            stages.filter { _.stageId == stageId }.map{_.attemptId}
+          }
+          if (stageAttempts.isEmpty) {
+            throw new NotFoundException(s"unknown stage: $stageId")
+          } else {
+            throw new NotFoundException(s"unknown attempt for stage $stageId.  " +
+              s"Found attempts: ${stageAttempts.mkString("[", ",", "]")}")
+          }
+      }
+
+    }
+  }
+}