path: root/core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala

/*
 * 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.ui.jobs

import org.scalatest.FunSuite
import org.scalatest.Matchers

import org.apache.spark._
import org.apache.spark.{LocalSparkContext, SparkConf, Success}
import org.apache.spark.executor._
import org.apache.spark.scheduler._
import org.apache.spark.util.Utils

class JobProgressListenerSuite extends FunSuite with LocalSparkContext with Matchers {

  val jobSubmissionTime = 1421191042750L
  val jobCompletionTime = 1421191296660L

  private def createStageStartEvent(stageId: Int) = {
    val stageInfo = new StageInfo(stageId, 0, stageId.toString, 0, null, "")
    SparkListenerStageSubmitted(stageInfo)
  }

  private def createStageEndEvent(stageId: Int, failed: Boolean = false) = {
    val stageInfo = new StageInfo(stageId, 0, stageId.toString, 0, null, "")
    if (failed) {
      stageInfo.failureReason = Some("Failed!")
    }
    SparkListenerStageCompleted(stageInfo)
  }

  private def createJobStartEvent(jobId: Int, stageIds: Seq[Int]) = {
    val stageInfos = stageIds.map { stageId =>
      new StageInfo(stageId, 0, stageId.toString, 0, null, "")
    }
    SparkListenerJobStart(jobId, jobSubmissionTime, stageInfos)
  }

  private def createJobEndEvent(jobId: Int, failed: Boolean = false) = {
    val result = if (failed) JobFailed(new Exception("dummy failure")) else JobSucceeded
    SparkListenerJobEnd(jobId, jobCompletionTime, result)
  }

  private def runJob(listener: SparkListener, jobId: Int, shouldFail: Boolean = false) {
    val stagesThatWontBeRun = jobId * 200 to jobId * 200 + 10
    val stageIds = jobId * 100 to jobId * 100 + 50
    listener.onJobStart(createJobStartEvent(jobId, stageIds ++ stagesThatWontBeRun))
    for (stageId <- stageIds) {
      listener.onStageSubmitted(createStageStartEvent(stageId))
      listener.onStageCompleted(createStageEndEvent(stageId, failed = stageId % 2 == 0))
    }
    listener.onJobEnd(createJobEndEvent(jobId, shouldFail))
  }

  private def assertActiveJobsStateIsEmpty(listener: JobProgressListener) {
    listener.getSizesOfActiveStateTrackingCollections.foreach { case (fieldName, size) =>
      assert(size === 0, s"$fieldName was not empty")
    }
  }

  test("test LRU eviction of stages") {
    val conf = new SparkConf()
    conf.set("spark.ui.retainedStages", 5.toString)
    val listener = new JobProgressListener(conf)

    for (i <- 1 to 50) {
      listener.onStageSubmitted(createStageStartEvent(i))
      listener.onStageCompleted(createStageEndEvent(i))
    }
    assertActiveJobsStateIsEmpty(listener)

    listener.completedStages.size should be (5)
    listener.completedStages.map(_.stageId).toSet should be (Set(50, 49, 48, 47, 46))
  }

  test("test LRU eviction of jobs") {
    val conf = new SparkConf()
    conf.set("spark.ui.retainedStages", 5.toString)
    conf.set("spark.ui.retainedJobs", 5.toString)
    val listener = new JobProgressListener(conf)

    // Run a bunch of jobs to get the listener into a state where we've exceeded both the
    // job and stage retention limits:
    for (jobId <- 1 to 10) {
      runJob(listener, jobId, shouldFail = false)
    }
    for (jobId <- 200 to 210) {
      runJob(listener, jobId, shouldFail = true)
    }
    assertActiveJobsStateIsEmpty(listener)
    // Snapshot the sizes of various soft- and hard-size-limited collections:
    val softLimitSizes = listener.getSizesOfSoftSizeLimitedCollections
    val hardLimitSizes = listener.getSizesOfHardSizeLimitedCollections
    // Run some more jobs:
    for (jobId <- 11 to 50) {
      runJob(listener, jobId, shouldFail = false)
      // We shouldn't exceed the hard / soft limit sizes after the jobs have finished:
      listener.getSizesOfSoftSizeLimitedCollections should be (softLimitSizes)
      listener.getSizesOfHardSizeLimitedCollections should be (hardLimitSizes)
    }

    listener.completedJobs.size should be (5)
    listener.completedJobs.map(_.jobId).toSet should be (Set(50, 49, 48, 47, 46))

    for (jobId <- 51 to 100) {
      runJob(listener, jobId, shouldFail = true)
      // We shouldn't exceed the hard / soft limit sizes after the jobs have finished:
      listener.getSizesOfSoftSizeLimitedCollections should be (softLimitSizes)
      listener.getSizesOfHardSizeLimitedCollections should be (hardLimitSizes)
    }
    assertActiveJobsStateIsEmpty(listener)

    // Completed and failed jobs each their own size limits, so this should still be the same:
    listener.completedJobs.size should be (5)
    listener.completedJobs.map(_.jobId).toSet should be (Set(50, 49, 48, 47, 46))
    listener.failedJobs.size should be (5)
    listener.failedJobs.map(_.jobId).toSet should be (Set(100, 99, 98, 97, 96))
  }

  test("test executor id to summary") {
    val conf = new SparkConf()
    val listener = new JobProgressListener(conf)
    val taskMetrics = new TaskMetrics()
    val shuffleReadMetrics = new ShuffleReadMetrics()
    assert(listener.stageIdToData.size === 0)

    // finish this task, should get updated shuffleRead
    shuffleReadMetrics.incRemoteBytesRead(1000)
    taskMetrics.setShuffleReadMetrics(Some(shuffleReadMetrics))
    var taskInfo = new TaskInfo(1234L, 0, 1, 0L, "exe-1", "host1", TaskLocality.NODE_LOCAL, false)
    taskInfo.finishTime = 1
    var task = new ShuffleMapTask(0)
    val taskType = Utils.getFormattedClassName(task)
    listener.onTaskEnd(
      SparkListenerTaskEnd(task.stageId, 0, taskType, Success, taskInfo, taskMetrics))
    assert(listener.stageIdToData.getOrElse((0, 0), fail())
      .executorSummary.getOrElse("exe-1", fail()).shuffleRead === 1000)

    // finish a task with unknown executor-id, nothing should happen
    taskInfo =
      new TaskInfo(1234L, 0, 1, 1000L, "exe-unknown", "host1", TaskLocality.NODE_LOCAL, true)
    taskInfo.finishTime = 1
    task = new ShuffleMapTask(0)
    listener.onTaskEnd(
      SparkListenerTaskEnd(task.stageId, 0, taskType, Success, taskInfo, taskMetrics))
    assert(listener.stageIdToData.size === 1)

    // finish this task, should get updated duration
    taskInfo = new TaskInfo(1235L, 0, 1, 0L, "exe-1", "host1", TaskLocality.NODE_LOCAL, false)
    taskInfo.finishTime = 1
    task = new ShuffleMapTask(0)
    listener.onTaskEnd(
      SparkListenerTaskEnd(task.stageId, 0, taskType, Success, taskInfo, taskMetrics))
    assert(listener.stageIdToData.getOrElse((0, 0), fail())
      .executorSummary.getOrElse("exe-1", fail()).shuffleRead === 2000)

    // finish this task, should get updated duration
    taskInfo = new TaskInfo(1236L, 0, 2, 0L, "exe-2", "host1", TaskLocality.NODE_LOCAL, false)
    taskInfo.finishTime = 1
    task = new ShuffleMapTask(0)
    listener.onTaskEnd(
      SparkListenerTaskEnd(task.stageId, 0, taskType, Success, taskInfo, taskMetrics))
    assert(listener.stageIdToData.getOrElse((0, 0), fail())
      .executorSummary.getOrElse("exe-2", fail()).shuffleRead === 1000)
  }

  test("test task success vs failure counting for different task end reasons") {
    val conf = new SparkConf()
    val listener = new JobProgressListener(conf)
    val metrics = new TaskMetrics()
    val taskInfo = new TaskInfo(1234L, 0, 3, 0L, "exe-1", "host1", TaskLocality.NODE_LOCAL, false)
    taskInfo.finishTime = 1
    val task = new ShuffleMapTask(0)
    val taskType = Utils.getFormattedClassName(task)

    // Go through all the failure cases to make sure we are counting them as failures.
    val taskFailedReasons = Seq(
      Resubmitted,
      new FetchFailed(null, 0, 0, 0, "ignored"),
      ExceptionFailure("Exception", "description", null, null, None),
      TaskResultLost,
      TaskKilled,
      ExecutorLostFailure("0"),
      UnknownReason)
    var failCount = 0
    for (reason <- taskFailedReasons) {
      listener.onTaskEnd(
        SparkListenerTaskEnd(task.stageId, 0, taskType, reason, taskInfo, metrics))
      failCount += 1
      assert(listener.stageIdToData((task.stageId, 0)).numCompleteTasks === 0)
      assert(listener.stageIdToData((task.stageId, 0)).numFailedTasks === failCount)
    }

    // Make sure we count success as success.
    listener.onTaskEnd(
      SparkListenerTaskEnd(task.stageId, 1, taskType, Success, taskInfo, metrics))
    assert(listener.stageIdToData((task.stageId, 1)).numCompleteTasks === 1)
    assert(listener.stageIdToData((task.stageId, 0)).numFailedTasks === failCount)
  }

  test("test update metrics") {
    val conf = new SparkConf()
    val listener = new JobProgressListener(conf)

    val taskType = Utils.getFormattedClassName(new ShuffleMapTask(0))
    val execId = "exe-1"

    def makeTaskMetrics(base: Int) = {
      val taskMetrics = new TaskMetrics()
      val shuffleReadMetrics = new ShuffleReadMetrics()
      val shuffleWriteMetrics = new ShuffleWriteMetrics()
      taskMetrics.setShuffleReadMetrics(Some(shuffleReadMetrics))
      taskMetrics.shuffleWriteMetrics = Some(shuffleWriteMetrics)
      shuffleReadMetrics.incRemoteBytesRead(base + 1)
      shuffleReadMetrics.incRemoteBlocksFetched(base + 2)
      shuffleWriteMetrics.incShuffleBytesWritten(base + 3)
      taskMetrics.setExecutorRunTime(base + 4)
      taskMetrics.incDiskBytesSpilled(base + 5)
      taskMetrics.incMemoryBytesSpilled(base + 6)
      val inputMetrics = new InputMetrics(DataReadMethod.Hadoop)
      taskMetrics.setInputMetrics(Some(inputMetrics))
      inputMetrics.incBytesRead(base + 7)
      val outputMetrics = new OutputMetrics(DataWriteMethod.Hadoop)
      taskMetrics.outputMetrics = Some(outputMetrics)
      outputMetrics.setBytesWritten(base + 8)
      taskMetrics
    }

    def makeTaskInfo(taskId: Long, finishTime: Int = 0) = {
      val taskInfo = new TaskInfo(taskId, 0, 1, 0L, execId, "host1", TaskLocality.NODE_LOCAL,
        false)
      taskInfo.finishTime = finishTime
      taskInfo
    }

    listener.onTaskStart(SparkListenerTaskStart(0, 0, makeTaskInfo(1234L)))
    listener.onTaskStart(SparkListenerTaskStart(0, 0, makeTaskInfo(1235L)))
    listener.onTaskStart(SparkListenerTaskStart(1, 0, makeTaskInfo(1236L)))
    listener.onTaskStart(SparkListenerTaskStart(1, 0, makeTaskInfo(1237L)))

    listener.onExecutorMetricsUpdate(SparkListenerExecutorMetricsUpdate(execId, Array(
      (1234L, 0, 0, makeTaskMetrics(0)),
      (1235L, 0, 0, makeTaskMetrics(100)),
      (1236L, 1, 0, makeTaskMetrics(200)))))

    var stage0Data = listener.stageIdToData.get((0, 0)).get
    var stage1Data = listener.stageIdToData.get((1, 0)).get
    assert(stage0Data.shuffleReadBytes == 102)
    assert(stage1Data.shuffleReadBytes == 201)
    assert(stage0Data.shuffleWriteBytes == 106)
    assert(stage1Data.shuffleWriteBytes == 203)
    assert(stage0Data.executorRunTime == 108)
    assert(stage1Data.executorRunTime == 204)
    assert(stage0Data.diskBytesSpilled == 110)
    assert(stage1Data.diskBytesSpilled == 205)
    assert(stage0Data.memoryBytesSpilled == 112)
    assert(stage1Data.memoryBytesSpilled == 206)
    assert(stage0Data.inputBytes == 114)
    assert(stage1Data.inputBytes == 207)
    assert(stage0Data.outputBytes == 116)
    assert(stage1Data.outputBytes == 208)
    assert(stage0Data.taskData.get(1234L).get.taskMetrics.get.shuffleReadMetrics.get
      .totalBlocksFetched == 2)
    assert(stage0Data.taskData.get(1235L).get.taskMetrics.get.shuffleReadMetrics.get
      .totalBlocksFetched == 102)
    assert(stage1Data.taskData.get(1236L).get.taskMetrics.get.shuffleReadMetrics.get
      .totalBlocksFetched == 202)

    // task that was included in a heartbeat
    listener.onTaskEnd(SparkListenerTaskEnd(0, 0, taskType, Success, makeTaskInfo(1234L, 1),
      makeTaskMetrics(300)))
    // task that wasn't included in a heartbeat
    listener.onTaskEnd(SparkListenerTaskEnd(1, 0, taskType, Success, makeTaskInfo(1237L, 1),
      makeTaskMetrics(400)))

    stage0Data = listener.stageIdToData.get((0, 0)).get
    stage1Data = listener.stageIdToData.get((1, 0)).get
    assert(stage0Data.shuffleReadBytes == 402)
    assert(stage1Data.shuffleReadBytes == 602)
    assert(stage0Data.shuffleWriteBytes == 406)
    assert(stage1Data.shuffleWriteBytes == 606)
    assert(stage0Data.executorRunTime == 408)
    assert(stage1Data.executorRunTime == 608)
    assert(stage0Data.diskBytesSpilled == 410)
    assert(stage1Data.diskBytesSpilled == 610)
    assert(stage0Data.memoryBytesSpilled == 412)
    assert(stage1Data.memoryBytesSpilled == 612)
    assert(stage0Data.inputBytes == 414)
    assert(stage1Data.inputBytes == 614)
    assert(stage0Data.outputBytes == 416)
    assert(stage1Data.outputBytes == 616)
    assert(stage0Data.taskData.get(1234L).get.taskMetrics.get.shuffleReadMetrics.get
      .totalBlocksFetched == 302)
    assert(stage1Data.taskData.get(1237L).get.taskMetrics.get.shuffleReadMetrics.get
      .totalBlocksFetched == 402)
  }
}