From 63cd958b4439aaf28581bdde3762e2413577a4bc Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Tue, 7 Jul 2015 11:47:01 -0700
Subject: [PATCH 01/10] [SPARK-8167] Make tasks that fail from YARN preemption
 not fail job.

The architecture is that, in YARN mode, if the driver detects that an executor
has disconnected, it asks the ApplicationMaster why the executor died. If the
ApplicationMaster is aware that the executor died because of preemption, all
tasks associated with that executor are not marked as failed. The executor
is still removed from the driver's list of available executors, however.
---
 .../org/apache/spark/TaskEndReason.scala      |  12 +
 .../spark/scheduler/ExecutorLossReason.scala  |  18 +-
 .../org/apache/spark/scheduler/Pool.scala     |   4 +-
 .../apache/spark/scheduler/Schedulable.scala  |   2 +-
 .../spark/scheduler/TaskSchedulerImpl.scala   |   8 +-
 .../spark/scheduler/TaskSetManager.scala      |  19 +-
 .../cluster/CoarseGrainedClusterMessage.scala |   8 +-
 .../CoarseGrainedSchedulerBackend.scala       |  18 +-
 .../cluster/SparkDeploySchedulerBackend.scala |   6 +-
 .../cluster/YarnSchedulerBackend.scala        |  44 +++-
 .../cluster/YarnSchedulerBackend.scala.orig   | 209 ++++++++++++++++++
 .../mesos/CoarseMesosSchedulerBackend.scala   |   4 +-
 .../cluster/mesos/MesosSchedulerBackend.scala |   2 +-
 .../spark/scheduler/TaskSetManagerSuite.scala |  10 +-
 .../spark/deploy/yarn/ApplicationMaster.scala |   9 +
 .../spark/deploy/yarn/YarnAllocator.scala     |  51 ++++-
 16 files changed, 383 insertions(+), 41 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig

diff --git a/core/src/main/scala/org/apache/spark/TaskEndReason.scala b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
index 48fd3e7e23d52..a6641b167cf7a 100644
--- a/core/src/main/scala/org/apache/spark/TaskEndReason.scala
+++ b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
@@ -166,6 +166,18 @@ case class ExecutorLostFailure(execId: String) extends TaskFailedReason {
   override def toErrorString: String = s"ExecutorLostFailure (executor ${execId} lost)"
 }
 
+/**
+ * :: DeveloperApi ::
+ * The task failed because the executor that it was running on was prematurely terminated. The
+ * executor is forcibly exited but the exit should be considered as part of normal cluster
+ * behavior.
+ */
+@DeveloperApi
+case class ExecutorForTaskExited(taskId: Long, execId: String, exitReason: String, exitCode: Int) extends TaskFailedReason {
+  override def toErrorString: String = s"Task with ID $taskId had its executor ${execId} exit normally with exit code "
+    s"$exitCode, due to the following reason: $exitReason."
+}
+
 /**
  * :: DeveloperApi ::
  * We don't know why the task ended -- for example, because of a ClassNotFound exception when
diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
index 2bc43a9186449..94e40f65541af 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
@@ -23,13 +23,25 @@ import org.apache.spark.executor.ExecutorExitCode
  * Represents an explanation for a executor or whole slave failing or exiting.
  */
 private[spark]
-class ExecutorLossReason(val message: String) {
+class ExecutorLossReason(val message: String) extends Serializable {
   override def toString: String = message
 }
 
+private[spark] case class ExecutorExitedAbnormally(val exitCode: Int, reason: String)
+  extends ExecutorLossReason(reason) {
+}
+
+private[spark] object ExecutorExitedAbnormally {
+  def apply(exitCode: Int): ExecutorExitedAbnormally = ExecutorExitedAbnormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
+}
+
 private[spark]
-case class ExecutorExited(val exitCode: Int)
-  extends ExecutorLossReason(ExecutorExitCode.explainExitCode(exitCode)) {
+case class ExecutorExitedNormally(val exitCode: Int, reason: String)
+  extends ExecutorLossReason(reason) {
+}
+
+private[spark] object ExecutorExitedNormally {
+  def apply(exitCode: Int): ExecutorExitedNormally = ExecutorExitedNormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
 }
 
 private[spark]
diff --git a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
index 174b73221afc0..7f2e721df3588 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Pool.scala
@@ -83,8 +83,8 @@ private[spark] class Pool(
     null
   }
 
-  override def executorLost(executorId: String, host: String) {
-    schedulableQueue.foreach(_.executorLost(executorId, host))
+  override def executorLost(executorId: String, host: String, reason: ExecutorLossReason) {
+    schedulableQueue.foreach(_.executorLost(executorId, host, reason))
   }
 
   override def checkSpeculatableTasks(): Boolean = {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala b/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
index a87ef030e69c2..ab00bc8f0bf4e 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Schedulable.scala
@@ -42,7 +42,7 @@ private[spark] trait Schedulable {
   def addSchedulable(schedulable: Schedulable): Unit
   def removeSchedulable(schedulable: Schedulable): Unit
   def getSchedulableByName(name: String): Schedulable
-  def executorLost(executorId: String, host: String): Unit
+  def executorLost(executorId: String, host: String, reason: ExecutorLossReason): Unit
   def checkSpeculatableTasks(): Boolean
   def getSortedTaskSetQueue: ArrayBuffer[TaskSetManager]
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index ed3dde0fc3055..14350371f1ba6 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -315,7 +315,7 @@ private[spark] class TaskSchedulerImpl(
           // We lost this entire executor, so remember that it's gone
           val execId = taskIdToExecutorId(tid)
           if (activeExecutorIds.contains(execId)) {
-            removeExecutor(execId)
+            removeExecutor(execId, SlaveLost(s"Task $tid was lost, so marking the executor as lost as well."))
             failedExecutor = Some(execId)
           }
         }
@@ -446,7 +446,7 @@ private[spark] class TaskSchedulerImpl(
       if (activeExecutorIds.contains(executorId)) {
         val hostPort = executorIdToHost(executorId)
         logError("Lost executor %s on %s: %s".format(executorId, hostPort, reason))
-        removeExecutor(executorId)
+        removeExecutor(executorId, reason)
         failedExecutor = Some(executorId)
       } else {
          // We may get multiple executorLost() calls with different loss reasons. For example, one
@@ -464,7 +464,7 @@ private[spark] class TaskSchedulerImpl(
   }
 
   /** Remove an executor from all our data structures and mark it as lost */
-  private def removeExecutor(executorId: String) {
+  private def removeExecutor(executorId: String, reason: ExecutorLossReason) {
     activeExecutorIds -= executorId
     val host = executorIdToHost(executorId)
     val execs = executorsByHost.getOrElse(host, new HashSet)
@@ -479,7 +479,7 @@ private[spark] class TaskSchedulerImpl(
       }
     }
     executorIdToHost -= executorId
-    rootPool.executorLost(executorId, host)
+    rootPool.executorLost(executorId, host, reason)
   }
 
   def executorAdded(execId: String, host: String) {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 82455b0426a5d..6a58ffff3f22b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -707,6 +707,9 @@ private[spark] class TaskSetManager(
             s"${ef.className} (${ef.description}) [duplicate $dupCount]")
         }
 
+      case e: ExecutorForTaskExited =>
+        logWarning(s"Task $tid failed because while it was being computed, its executor exited normally." +
+          s" Not marking the task as failed.")
       case e: TaskFailedReason =>  // TaskResultLost, TaskKilled, and others
         logWarning(failureReason)
 
@@ -718,7 +721,7 @@ private[spark] class TaskSetManager(
       put(info.executorId, clock.getTimeMillis())
     sched.dagScheduler.taskEnded(tasks(index), reason, null, null, info, taskMetrics)
     addPendingTask(index)
-    if (!isZombie && state != TaskState.KILLED && !reason.isInstanceOf[TaskCommitDenied]) {
+    if (!isZombie && state != TaskState.KILLED && shouldTaskFailureEventuallyFailJob(reason)) {
       // If a task failed because its attempt to commit was denied, do not count this failure
       // towards failing the stage. This is intended to prevent spurious stage failures in cases
       // where many speculative tasks are launched and denied to commit.
@@ -735,6 +738,10 @@ private[spark] class TaskSetManager(
     maybeFinishTaskSet()
   }
 
+  private def shouldTaskFailureEventuallyFailJob(reason: TaskEndReason): Boolean = {
+    !reason.isInstanceOf[TaskCommitDenied] && !reason.isInstanceOf[ExecutorForTaskExited]
+  }
+
   def abort(message: String): Unit = sched.synchronized {
     // TODO: Kill running tasks if we were not terminated due to a Mesos error
     sched.dagScheduler.taskSetFailed(taskSet, message)
@@ -774,7 +781,7 @@ private[spark] class TaskSetManager(
   }
 
   /** Called by TaskScheduler when an executor is lost so we can re-enqueue our tasks */
-  override def executorLost(execId: String, host: String) {
+  override def executorLost(execId: String, host: String, reason: ExecutorLossReason) {
     logInfo("Re-queueing tasks for " + execId + " from TaskSet " + taskSet.id)
 
     // Re-enqueue pending tasks for this host based on the status of the cluster. Note
@@ -805,9 +812,13 @@ private[spark] class TaskSetManager(
         }
       }
     }
-    // Also re-enqueue any tasks that were running on the node
     for ((tid, info) <- taskInfos if info.running && info.executorId == execId) {
-      handleFailedTask(tid, TaskState.FAILED, ExecutorLostFailure(execId))
+      // Also re-enqueue any tasks that were running on the node
+      val executorFailureReason = reason match {
+        case exited: ExecutorExitedNormally => ExecutorForTaskExited(tid, execId, exited.reason, exited.exitCode)
+        case default => ExecutorLostFailure(execId)
+      }
+      handleFailedTask(tid, TaskState.FAILED, executorFailureReason)
     }
     // recalculate valid locality levels and waits when executor is lost
     recomputeLocality()
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
index 4be1eda2e9291..e8cc4b135cf0f 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
@@ -21,6 +21,7 @@ import java.nio.ByteBuffer
 
 import org.apache.spark.TaskState.TaskState
 import org.apache.spark.rpc.RpcEndpointRef
+import org.apache.spark.scheduler.{SlaveLost, ExecutorLossReason}
 import org.apache.spark.util.{SerializableBuffer, Utils}
 
 private[spark] sealed trait CoarseGrainedClusterMessage extends Serializable
@@ -70,7 +71,7 @@ private[spark] object CoarseGrainedClusterMessages {
 
   case object StopExecutors extends CoarseGrainedClusterMessage
 
-  case class RemoveExecutor(executorId: String, reason: String) extends CoarseGrainedClusterMessage
+  case class RemoveExecutor(executorId: String, reason: ExecutorLossReason)
 
   case class SetupDriver(driver: RpcEndpointRef) extends CoarseGrainedClusterMessage
 
@@ -88,6 +89,11 @@ private[spark] object CoarseGrainedClusterMessages {
   // This includes executors already pending or running
   case class RequestExecutors(requestedTotal: Int) extends CoarseGrainedClusterMessage
 
+  // Check if an executor was force-killed but for a normal reason
+  // This could be the case if e.g. the cluster manager supports killing an executor to move
+  // it elsewhere or to kill an executor in order to free resources
+  case class GetExecutorLossReason(executorId: String)
+
   case class KillExecutors(executorIds: Seq[String]) extends CoarseGrainedClusterMessage
 
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 7c7f70d8a193b..83cae7e3c02ed 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -76,7 +76,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
 
     override protected def log = CoarseGrainedSchedulerBackend.this.log
 
-    private val addressToExecutorId = new HashMap[RpcAddress, String]
+    protected val addressToExecutorId = new HashMap[RpcAddress, String]
 
     private val reviveThread =
       ThreadUtils.newDaemonSingleThreadScheduledExecutor("driver-revive-thread")
@@ -175,8 +175,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
-      addressToExecutorId.get(remoteAddress).foreach(removeExecutor(_,
-        "remote Rpc client disassociated"))
+      addressToExecutorId.get(remoteAddress).foreach(removeExecutor(_, SlaveLost("remote Rpc client disassociated")))
     }
 
     // Make fake resource offers on just one executor
@@ -214,7 +213,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     // Remove a disconnected slave from the cluster
-    def removeExecutor(executorId: String, reason: String): Unit = {
+    def removeExecutor(executorId: String, reason: ExecutorLossReason): Unit = {
       executorDataMap.get(executorId) match {
         case Some(executorInfo) =>
           // This must be synchronized because variables mutated
@@ -226,9 +225,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
           }
           totalCoreCount.addAndGet(-executorInfo.totalCores)
           totalRegisteredExecutors.addAndGet(-1)
-          scheduler.executorLost(executorId, SlaveLost(reason))
+          scheduler.executorLost(executorId, reason)
           listenerBus.post(
-            SparkListenerExecutorRemoved(System.currentTimeMillis(), executorId, reason))
+            SparkListenerExecutorRemoved(System.currentTimeMillis(), executorId, reason.toString))
         case None => logError(s"Asked to remove non-existent executor $executorId")
       }
     }
@@ -250,10 +249,11 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     // TODO (prashant) send conf instead of properties
-    driverEndpoint = rpcEnv.setupEndpoint(
-      CoarseGrainedSchedulerBackend.ENDPOINT_NAME, new DriverEndpoint(rpcEnv, properties))
+    driverEndpoint = rpcEnv.setupEndpoint(CoarseGrainedSchedulerBackend.ENDPOINT_NAME, createDriverEndpoint(properties))
   }
 
+  protected def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = new DriverEndpoint(rpcEnv, properties)
+
   def stopExecutors() {
     try {
       if (driverEndpoint != null) {
@@ -291,7 +291,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
   }
 
   // Called by subclasses when notified of a lost worker
-  def removeExecutor(executorId: String, reason: String) {
+  def removeExecutor(executorId: String, reason: ExecutorLossReason) {
     try {
       driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, reason))
     } catch {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
index 687ae9620460f..2825ef0871381 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
@@ -23,7 +23,7 @@ import org.apache.spark.rpc.RpcAddress
 import org.apache.spark.{Logging, SparkConf, SparkContext, SparkEnv}
 import org.apache.spark.deploy.{ApplicationDescription, Command}
 import org.apache.spark.deploy.client.{AppClient, AppClientListener}
-import org.apache.spark.scheduler.{ExecutorExited, ExecutorLossReason, SlaveLost, TaskSchedulerImpl}
+import org.apache.spark.scheduler.{ExecutorExitedNormally, ExecutorLossReason, SlaveLost, TaskSchedulerImpl}
 import org.apache.spark.util.Utils
 
 private[spark] class SparkDeploySchedulerBackend(
@@ -135,11 +135,11 @@ private[spark] class SparkDeploySchedulerBackend(
 
   override def executorRemoved(fullId: String, message: String, exitStatus: Option[Int]) {
     val reason: ExecutorLossReason = exitStatus match {
-      case Some(code) => ExecutorExited(code)
+      case Some(code) => ExecutorExitedNormally(code)
       case None => SlaveLost(message)
     }
     logInfo("Executor %s removed: %s".format(fullId, message))
-    removeExecutor(fullId.split("/")(1), reason.toString)
+    removeExecutor(fullId.split("/")(1), reason)
   }
 
   override def sufficientResourcesRegistered(): Boolean = {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index bc67abb5df446..71c2bf7eab261 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -17,12 +17,15 @@
 
 package org.apache.spark.scheduler.cluster
 
+import java.util.Properties
+
+import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{Future, ExecutionContext}
 
 import org.apache.spark.{Logging, SparkContext}
 import org.apache.spark.rpc._
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
-import org.apache.spark.scheduler.TaskSchedulerImpl
+import org.apache.spark.scheduler._
 import org.apache.spark.ui.JettyUtils
 import org.apache.spark.util.{ThreadUtils, RpcUtils}
 
@@ -47,6 +50,7 @@ private[spark] abstract class YarnSchedulerBackend(
     YarnSchedulerBackend.ENDPOINT_NAME, new YarnSchedulerEndpoint(rpcEnv))
 
   private implicit val askTimeout = RpcUtils.askRpcTimeout(sc.conf)
+  private val executorDisconnectedHandlerPool = ThreadUtils.newDaemonSingleThreadExecutor("executor-disconnected-handler")
 
   /**
    * Request executors from the ApplicationMaster by specifying the total number desired.
@@ -89,6 +93,36 @@ private[spark] abstract class YarnSchedulerBackend(
     }
   }
 
+  /**
+   * Override the DriverEndpoint to add extra logic for the case when an executor is disconnected.
+   * We should check the cluster manager and find if the loss of the executor was caused by YARN
+   * force killing it due to preemption.
+   */
+  class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
+      extends DriverEndpoint(rpcEnv, sparkProperties) {
+
+    override def onDisconnected(rpcAddress: RpcAddress): Unit = {
+      addressToExecutorId.get(rpcAddress).foreach({ executorId =>
+        executorDisconnectedHandlerPool.submit(new Runnable() {
+          override def run(): Unit = {
+            val wasExecutorForceKilledNormally =
+              yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
+            wasExecutorForceKilledNormally match {
+              case Some(killReason) =>
+                driverEndpoint.send(RemoveExecutor(executorId, killReason))
+              case None =>
+                driverEndpoint.send(RemoveExecutor(executorId, SlaveLost("Executor was terminated for an unknown reason.")))
+            }
+          }
+        })
+      })
+    }
+  }
+
+  override def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = {
+    new YarnDriverEndpoint(rpcEnv, properties)
+  }
+
   /**
    * An [[RpcEndpoint]] that communicates with the ApplicationMaster.
    */
@@ -141,6 +175,13 @@ private[spark] abstract class YarnSchedulerBackend(
             context.reply(false)
         }
 
+      case c: GetExecutorLossReason =>
+        amEndpoint match {
+          case Some(am) => context.reply(am.askWithRetry[ExecutorLossReason](c))
+          case None =>
+            logWarning("Attempted to check if an executor exited normally before the AM has registered!")
+        }
+
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
@@ -155,6 +196,7 @@ private[spark] abstract class YarnSchedulerBackend(
   }
 }
 
+
 private[spark] object YarnSchedulerBackend {
   val ENDPOINT_NAME = "YarnScheduler"
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig
new file mode 100644
index 0000000000000..fb1c3af55bcc0
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig
@@ -0,0 +1,209 @@
+/*
+ * 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.scheduler.cluster
+
+import java.util.Properties
+
+import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.{Future, ExecutionContext}
+
+import org.apache.spark.{Logging, SparkContext}
+import org.apache.spark.rpc._
+import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
+import org.apache.spark.scheduler._
+import org.apache.spark.ui.JettyUtils
+import org.apache.spark.util.{ThreadUtils, RpcUtils}
+
+import scala.util.control.NonFatal
+
+/**
+ * Abstract Yarn scheduler backend that contains common logic
+ * between the client and cluster Yarn scheduler backends.
+ */
+private[spark] abstract class YarnSchedulerBackend(
+    scheduler: TaskSchedulerImpl,
+    sc: SparkContext)
+  extends CoarseGrainedSchedulerBackend(scheduler, sc.env.rpcEnv) {
+
+  if (conf.getOption("spark.scheduler.minRegisteredResourcesRatio").isEmpty) {
+    minRegisteredRatio = 0.8
+  }
+
+  protected var totalExpectedExecutors = 0
+
+  private val yarnSchedulerEndpoint = rpcEnv.setupEndpoint(
+    YarnSchedulerBackend.ENDPOINT_NAME, new YarnSchedulerEndpoint(rpcEnv))
+
+<<<<<<< Updated upstream
+  private implicit val askTimeout = RpcUtils.askRpcTimeout(sc.conf)
+||||||| merged common ancestors
+  private implicit val askTimeout = RpcUtils.askTimeout(sc.conf)
+=======
+  private val executorDisconnectedHandlerPool = ThreadUtils.newDaemonSingleThreadExecutor("executor-disconnected-handler")
+
+  private implicit val askTimeout = RpcUtils.askTimeout(sc.conf)
+>>>>>>> Stashed changes
+
+  /**
+   * Request executors from the ApplicationMaster by specifying the total number desired.
+   * This includes executors already pending or running.
+   */
+  override def doRequestTotalExecutors(requestedTotal: Int): Boolean = {
+    yarnSchedulerEndpoint.askWithRetry[Boolean](RequestExecutors(requestedTotal))
+  }
+
+  /**
+   * Request that the ApplicationMaster kill the specified executors.
+   */
+  override def doKillExecutors(executorIds: Seq[String]): Boolean = {
+    yarnSchedulerEndpoint.askWithRetry[Boolean](KillExecutors(executorIds))
+  }
+
+  override def sufficientResourcesRegistered(): Boolean = {
+    totalRegisteredExecutors.get() >= totalExpectedExecutors * minRegisteredRatio
+  }
+
+  /**
+   * Add filters to the SparkUI.
+   */
+  private def addWebUIFilter(
+      filterName: String,
+      filterParams: Map[String, String],
+      proxyBase: String): Unit = {
+    if (proxyBase != null && proxyBase.nonEmpty) {
+      System.setProperty("spark.ui.proxyBase", proxyBase)
+    }
+
+    val hasFilter =
+      filterName != null && filterName.nonEmpty &&
+      filterParams != null && filterParams.nonEmpty
+    if (hasFilter) {
+      logInfo(s"Add WebUI Filter. $filterName, $filterParams, $proxyBase")
+      conf.set("spark.ui.filters", filterName)
+      filterParams.foreach { case (k, v) => conf.set(s"spark.$filterName.param.$k", v) }
+      scheduler.sc.ui.foreach { ui => JettyUtils.addFilters(ui.getHandlers, conf) }
+    }
+  }
+
+  /**
+   * Override the DriverEndpoint to add extra logic for the case when an executor is disconnected.
+   * We should check the cluster manager and find if the loss of the executor was caused by YARN
+   * force killing it due to preemption.
+   */
+  class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
+      extends DriverEndpoint(rpcEnv, sparkProperties) {
+
+    override def onDisconnected(rpcAddress: RpcAddress): Unit = {
+      addressToExecutorId.get(rpcAddress).foreach({ executorId =>
+        executorDisconnectedHandlerPool.submit(new Runnable() {
+          override def run(): Unit = {
+            val wasExecutorForceKilledNormally =
+              yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
+            wasExecutorForceKilledNormally match {
+              case Some(killReason) =>
+                driverEndpoint.send(RemoveExecutor(executorId, killReason))
+              case None =>
+                driverEndpoint.send(RemoveExecutor(executorId, SlaveLost("Executor was terminated for an unknown reason.")))
+            }
+          }
+        })
+      })
+    }
+  }
+
+  override def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = {
+    new YarnDriverEndpoint(rpcEnv, properties)
+  }
+
+  /**
+   * An [[RpcEndpoint]] that communicates with the ApplicationMaster.
+   */
+  private class YarnSchedulerEndpoint(override val rpcEnv: RpcEnv)
+    extends ThreadSafeRpcEndpoint with Logging {
+    private var amEndpoint: Option[RpcEndpointRef] = None
+
+    private val askAmThreadPool =
+      ThreadUtils.newDaemonCachedThreadPool("yarn-scheduler-ask-am-thread-pool")
+    implicit val askAmExecutor = ExecutionContext.fromExecutor(askAmThreadPool)
+
+    override def receive: PartialFunction[Any, Unit] = {
+      case RegisterClusterManager(am) =>
+        logInfo(s"ApplicationMaster registered as $am")
+        amEndpoint = Some(am)
+
+      case AddWebUIFilter(filterName, filterParams, proxyBase) =>
+        addWebUIFilter(filterName, filterParams, proxyBase)
+
+    }
+
+    override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
+      case r: RequestExecutors =>
+        amEndpoint match {
+          case Some(am) =>
+            Future {
+              context.reply(am.askWithRetry[Boolean](r))
+            } onFailure {
+              case NonFatal(e) =>
+                logError(s"Sending $r to AM was unsuccessful", e)
+                context.sendFailure(e)
+            }
+          case None =>
+            logWarning("Attempted to request executors before the AM has registered!")
+            context.reply(false)
+        }
+
+      case k: KillExecutors =>
+        amEndpoint match {
+          case Some(am) =>
+            Future {
+              context.reply(am.askWithRetry[Boolean](k))
+            } onFailure {
+              case NonFatal(e) =>
+                logError(s"Sending $k to AM was unsuccessful", e)
+                context.sendFailure(e)
+            }
+          case None =>
+            logWarning("Attempted to kill executors before the AM has registered!")
+            context.reply(false)
+        }
+
+      case c: GetExecutorLossReason =>
+        amEndpoint match {
+          case Some(am) => context.reply(am.askWithRetry[ExecutorLossReason](c))
+          case None =>
+            logWarning("Attempted to check if an executor exited normally before the AM has registered!")
+        }
+
+    }
+
+    override def onDisconnected(remoteAddress: RpcAddress): Unit = {
+      if (amEndpoint.exists(_.address == remoteAddress)) {
+        logWarning(s"ApplicationMaster has disassociated: $remoteAddress")
+      }
+    }
+
+    override def onStop(): Unit = {
+      askAmThreadPool.shutdownNow()
+    }
+  }
+}
+
+
+private[spark] object YarnSchedulerBackend {
+  val ENDPOINT_NAME = "YarnScheduler"
+}
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
index b68f8c7685eba..b59e156101e34 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
@@ -27,7 +27,7 @@ import org.apache.mesos.{Scheduler => MScheduler, _}
 import org.apache.mesos.Protos.{TaskInfo => MesosTaskInfo, _}
 import org.apache.spark.{SparkContext, SparkEnv, SparkException, TaskState}
 import org.apache.spark.rpc.RpcAddress
-import org.apache.spark.scheduler.TaskSchedulerImpl
+import org.apache.spark.scheduler.{SlaveLost, TaskSchedulerImpl}
 import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
 import org.apache.spark.util.Utils
 
@@ -268,7 +268,7 @@ private[spark] class CoarseMesosSchedulerBackend(
       if (slaveIdsWithExecutors.contains(slaveId.getValue)) {
         // Note that the slave ID corresponds to the executor ID on that slave
         slaveIdsWithExecutors -= slaveId.getValue
-        removeExecutor(slaveId.getValue, "Mesos slave lost")
+        removeExecutor(slaveId.getValue, SlaveLost("Mesos slave lost"))
       }
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
index d72e2af456e15..a20c829275fe7 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
@@ -357,7 +357,7 @@ private[spark] class MesosSchedulerBackend(
                             slaveId: SlaveID, status: Int) {
     logInfo("Executor lost: %s, marking slave %s as lost".format(executorId.getValue,
                                                                  slaveId.getValue))
-    recordSlaveLost(d, slaveId, ExecutorExited(status))
+    recordSlaveLost(d, slaveId, ExecutorExitedNormally(status))
   }
 
   override def killTask(taskId: Long, executorId: String, interruptThread: Boolean): Unit = {
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
index 0060f3396dcde..3a3254cd4b3a6 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
@@ -330,7 +330,7 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
 
     // Now mark host2 as dead
     sched.removeExecutor("exec2")
-    manager.executorLost("exec2", "host2")
+    manager.executorLost("exec2", "host2", SlaveLost())
 
     // nothing should be chosen
     assert(manager.resourceOffer("exec1", "host1", ANY) === None)
@@ -500,10 +500,10 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
       Array(PROCESS_LOCAL, NODE_LOCAL, NO_PREF, RACK_LOCAL, ANY)))
     // test if the valid locality is recomputed when the executor is lost
     sched.removeExecutor("execC")
-    manager.executorLost("execC", "host2")
+    manager.executorLost("execC", "host2", SlaveLost())
     assert(manager.myLocalityLevels.sameElements(Array(NODE_LOCAL, NO_PREF, ANY)))
     sched.removeExecutor("execD")
-    manager.executorLost("execD", "host1")
+    manager.executorLost("execD", "host1", SlaveLost())
     assert(manager.myLocalityLevels.sameElements(Array(NO_PREF, ANY)))
   }
 
@@ -717,8 +717,8 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     assert(manager.resourceOffer("execB.2", "host2", ANY) !== None)
     sched.removeExecutor("execA")
     sched.removeExecutor("execB.2")
-    manager.executorLost("execA", "host1")
-    manager.executorLost("execB.2", "host2")
+    manager.executorLost("execA", "host1", SlaveLost())
+    manager.executorLost("execB.2", "host2", SlaveLost())
     clock.advance(LOCALITY_WAIT_MS * 4)
     sched.addExecutor("execC", "host3")
     manager.executorAdded()
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index 83dafa4a125d2..09dc87169587b 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -576,6 +576,15 @@ private[spark] class ApplicationMaster(
           case None => logWarning("Container allocator is not ready to kill executors yet.")
         }
         context.reply(true)
+
+      case GetExecutorLossReason(executorId) =>
+        val executorExitStatus = Option(allocator) match {
+          case Some(a) => Some(a.getExecutorLossReason(executorId))
+          case None =>
+            logWarning("Container allocator was not ready to report on executor status.")
+            None
+        }
+        context.reply(executorExitStatus)
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index 940873fbd046c..f91089b0ae588 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -21,6 +21,8 @@ import java.util.Collections
 import java.util.concurrent._
 import java.util.regex.Pattern
 
+import org.apache.spark.scheduler.{ExecutorExitedAbnormally, ExecutorExitedNormally, ExecutorLossReason}
+
 import scala.collection.JavaConversions._
 import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet}
 
@@ -62,6 +64,8 @@ private[yarn] class YarnAllocator(
 
   import YarnAllocator._
 
+  private val NORMAL_CONTAINER_EXIT_STATUS = ExecutorExitedNormally(0, "Executor exited normally.")
+
   // RackResolver logs an INFO message whenever it resolves a rack, which is way too often.
   if (Logger.getLogger(classOf[RackResolver]).getLevel == null) {
     Logger.getLogger(classOf[RackResolver]).setLevel(Level.WARN)
@@ -88,6 +92,16 @@ private[yarn] class YarnAllocator(
   // Visible for testing.
   private[yarn] val executorIdToContainer = new HashMap[String, Container]
 
+  // Maintain container exit statuses, for any executors that exit with non-zero codes.
+  // Used to report to the driver how a given executor was terminated.
+  // For example, the driver will want to ignore preempted containers
+  // and thus should not be considered as a failure counting against
+  // the job's task failure count.
+
+  // Store only non-zero exit codes to save on space for the expected-majority case where executors
+  // terminate normally.
+  private val completedNonZeroContainerExitCodes = new HashMap[ContainerId, ExecutorLossReason]
+
   // Executor memory in MB.
   protected val executorMemory = args.executorMemory
   // Additional memory overhead.
@@ -172,6 +186,11 @@ private[yarn] class YarnAllocator(
     }
   }
 
+  def getExecutorLossReason(executorId: String): ExecutorLossReason = synchronized {
+    allocateResources()
+    completedNonZeroContainerExitCodes.getOrElse(executorIdToContainer(executorId).getId, NORMAL_CONTAINER_EXIT_STATUS)
+  }
+
   /**
    * Request resources such that, if YARN gives us all we ask for, we'll have a number of containers
    * equal to maxExecutors.
@@ -399,21 +418,43 @@ private[yarn] class YarnAllocator(
         // Hadoop 2.2.X added a ContainerExitStatus we should switch to use
         // there are some exit status' we shouldn't necessarily count against us, but for
         // now I think its ok as none of the containers are expected to exit
+        var isExecutorNonZeroExitNormal = true
+        var containerExitReason = "Container exited for an unknown reason."
         if (completedContainer.getExitStatus == ContainerExitStatus.PREEMPTED) {
-          logInfo("Container preempted: " + containerId)
+          containerExitReason = s"Container $containerId was preempted."
+          logInfo(containerExitReason)
         } else if (completedContainer.getExitStatus == -103) { // vmem limit exceeded
-          logWarning(memLimitExceededLogMessage(
+          // Should probably still count these towards task failures
+          isExecutorNonZeroExitNormal = false
+          containerExitReason = memLimitExceededLogMessage(
             completedContainer.getDiagnostics,
-            VMEM_EXCEEDED_PATTERN))
+            VMEM_EXCEEDED_PATTERN)
+          logWarning(containerExitReason)
         } else if (completedContainer.getExitStatus == -104) { // pmem limit exceeded
-          logWarning(memLimitExceededLogMessage(
+          // Should probably still count these towards task failures
+          isExecutorNonZeroExitNormal = false
+          containerExitReason = memLimitExceededLogMessage(
             completedContainer.getDiagnostics,
-            PMEM_EXCEEDED_PATTERN))
+            PMEM_EXCEEDED_PATTERN)
+          logWarning(containerExitReason)
         } else if (completedContainer.getExitStatus != 0) {
           logInfo("Container marked as failed: " + containerId +
             ". Exit status: " + completedContainer.getExitStatus +
             ". Diagnostics: " + completedContainer.getDiagnostics)
           numExecutorsFailed += 1
+          isExecutorNonZeroExitNormal = false
+          containerExitReason = s"Container $containerId exited abnormally, and was marked as failed."
+        }
+
+        if (completedContainer.getExitStatus() != 0) {
+          val exitStatus = {
+            if (isExecutorNonZeroExitNormal) {
+              ExecutorExitedNormally(completedContainer.getExitStatus(), containerExitReason)
+            } else {
+              ExecutorExitedAbnormally(completedContainer.getExitStatus(), containerExitReason)
+            }
+          }
+          completedNonZeroContainerExitCodes.put(containerId, exitStatus)
         }
       }
 

From 1100ed8d1cb3f504729580af672bf6d6b3eeb624 Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Wed, 8 Jul 2015 11:34:55 -0700
Subject: [PATCH 02/10] Handle multiple onDisconnected events since executor
 removal doesn't happen immediately.

Also begin unit tests.
---
 .../cluster/YarnSchedulerBackend.scala        | 48 +++++++++++++------
 .../spark/deploy/yarn/ApplicationMaster.scala |  4 +-
 .../spark/deploy/yarn/YarnAllocator.scala     | 35 ++++++++------
 .../deploy/yarn/YarnAllocatorSuite.scala      | 19 ++++++++
 4 files changed, 74 insertions(+), 32 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index 71c2bf7eab261..95f40c4b1e990 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -19,6 +19,7 @@ package org.apache.spark.scheduler.cluster
 
 import java.util.Properties
 
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{Future, ExecutionContext}
 
@@ -50,7 +51,6 @@ private[spark] abstract class YarnSchedulerBackend(
     YarnSchedulerBackend.ENDPOINT_NAME, new YarnSchedulerEndpoint(rpcEnv))
 
   private implicit val askTimeout = RpcUtils.askRpcTimeout(sc.conf)
-  private val executorDisconnectedHandlerPool = ThreadUtils.newDaemonSingleThreadExecutor("executor-disconnected-handler")
 
   /**
    * Request executors from the ApplicationMaster by specifying the total number desired.
@@ -98,23 +98,34 @@ private[spark] abstract class YarnSchedulerBackend(
    * We should check the cluster manager and find if the loss of the executor was caused by YARN
    * force killing it due to preemption.
    */
-  class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
+  private class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
       extends DriverEndpoint(rpcEnv, sparkProperties) {
 
+    private val pendingDisconnectedExecutors = new mutable.HashSet[String]
+    private val handleDisconnectedExecutorThreadPool =
+      ThreadUtils.newDaemonCachedThreadPool("yarn-driver-endpoint-handle-disconnected-executor-thread-pool")
+
     override def onDisconnected(rpcAddress: RpcAddress): Unit = {
       addressToExecutorId.get(rpcAddress).foreach({ executorId =>
-        executorDisconnectedHandlerPool.submit(new Runnable() {
-          override def run(): Unit = {
-            val wasExecutorForceKilledNormally =
-              yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
-            wasExecutorForceKilledNormally match {
-              case Some(killReason) =>
-                driverEndpoint.send(RemoveExecutor(executorId, killReason))
-              case None =>
-                driverEndpoint.send(RemoveExecutor(executorId, SlaveLost("Executor was terminated for an unknown reason.")))
-            }
+        pendingDisconnectedExecutors.synchronized {
+          if (!pendingDisconnectedExecutors.contains(executorId)) {
+            pendingDisconnectedExecutors.add(executorId)
+            handleDisconnectedExecutorThreadPool.submit(new Runnable() {
+              override def run(): Unit = {
+                val executorLossReason = yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
+                executorLossReason match {
+                  case Some(reason) => driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, reason))
+                  case None =>
+                    logWarning(s"Attempted to get executor loss reason for $rpcAddress but got no response. Marking as slave lost.")
+                    driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, SlaveLost()))
+                }
+                pendingDisconnectedExecutors.synchronized {
+                  pendingDisconnectedExecutors.remove(executorId)
+                }
+              }
+            })
           }
-        })
+        }
       })
     }
   }
@@ -177,11 +188,18 @@ private[spark] abstract class YarnSchedulerBackend(
 
       case c: GetExecutorLossReason =>
         amEndpoint match {
-          case Some(am) => context.reply(am.askWithRetry[ExecutorLossReason](c))
+          case Some(am) =>
+            Future {
+              context.reply(am.askWithRetry[Option[ExecutorLossReason]](c))
+            } onFailure {
+              case NonFatal(e) =>
+                logError(s"Finding the executor loss reason was unsuccessful", e)
+                context.sendFailure(e)
+            }
           case None =>
             logWarning("Attempted to check if an executor exited normally before the AM has registered!")
+            context.reply(None)
         }
-
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index 09dc87169587b..e95d2682cce9f 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -578,13 +578,13 @@ private[spark] class ApplicationMaster(
         context.reply(true)
 
       case GetExecutorLossReason(executorId) =>
-        val executorExitStatus = Option(allocator) match {
+        val executorLossReason = Option(allocator) match {
           case Some(a) => Some(a.getExecutorLossReason(executorId))
           case None =>
             logWarning("Container allocator was not ready to report on executor status.")
             None
         }
-        context.reply(executorExitStatus)
+        context.reply(executorLossReason)
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index f91089b0ae588..2b1f347b323cf 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -101,6 +101,7 @@ private[yarn] class YarnAllocator(
   // Store only non-zero exit codes to save on space for the expected-majority case where executors
   // terminate normally.
   private val completedNonZeroContainerExitCodes = new HashMap[ContainerId, ExecutorLossReason]
+  private val killedExecutors = new HashSet[String]
 
   // Executor memory in MB.
   protected val executorMemory = args.executorMemory
@@ -181,6 +182,7 @@ private[yarn] class YarnAllocator(
       val container = executorIdToContainer.remove(executorId).get
       internalReleaseContainer(container)
       numExecutorsRunning -= 1
+      killedExecutors += executorId
     } else {
       logWarning(s"Attempted to kill unknown executor $executorId!")
     }
@@ -188,7 +190,11 @@ private[yarn] class YarnAllocator(
 
   def getExecutorLossReason(executorId: String): ExecutorLossReason = synchronized {
     allocateResources()
-    completedNonZeroContainerExitCodes.getOrElse(executorIdToContainer(executorId).getId, NORMAL_CONTAINER_EXIT_STATUS)
+    if (killedExecutors.contains(executorId)) {
+        NORMAL_CONTAINER_EXIT_STATUS
+    } else {
+      completedNonZeroContainerExitCodes.getOrElse(executorIdToContainer(executorId).getId, NORMAL_CONTAINER_EXIT_STATUS)
+    }
   }
 
   /**
@@ -418,43 +424,42 @@ private[yarn] class YarnAllocator(
         // Hadoop 2.2.X added a ContainerExitStatus we should switch to use
         // there are some exit status' we shouldn't necessarily count against us, but for
         // now I think its ok as none of the containers are expected to exit
-        var isExecutorNonZeroExitNormal = true
+        var isExecutorNonZeroExitNormal = false
         var containerExitReason = "Container exited for an unknown reason."
-        if (completedContainer.getExitStatus == ContainerExitStatus.PREEMPTED) {
+        val exitStatus = completedContainer.getExitStatus
+        if (exitStatus == ContainerExitStatus.PREEMPTED) {
+          isExecutorNonZeroExitNormal = true
           containerExitReason = s"Container $containerId was preempted."
           logInfo(containerExitReason)
-        } else if (completedContainer.getExitStatus == -103) { // vmem limit exceeded
+        } else if (exitStatus == -103) { // vmem limit exceeded
           // Should probably still count these towards task failures
-          isExecutorNonZeroExitNormal = false
           containerExitReason = memLimitExceededLogMessage(
             completedContainer.getDiagnostics,
             VMEM_EXCEEDED_PATTERN)
           logWarning(containerExitReason)
-        } else if (completedContainer.getExitStatus == -104) { // pmem limit exceeded
+        } else if (exitStatus == -104) { // pmem limit exceeded
           // Should probably still count these towards task failures
-          isExecutorNonZeroExitNormal = false
           containerExitReason = memLimitExceededLogMessage(
             completedContainer.getDiagnostics,
             PMEM_EXCEEDED_PATTERN)
           logWarning(containerExitReason)
-        } else if (completedContainer.getExitStatus != 0) {
+        } else if (exitStatus != 0) {
           logInfo("Container marked as failed: " + containerId +
             ". Exit status: " + completedContainer.getExitStatus +
             ". Diagnostics: " + completedContainer.getDiagnostics)
           numExecutorsFailed += 1
-          isExecutorNonZeroExitNormal = false
-          containerExitReason = s"Container $containerId exited abnormally, and was marked as failed."
+          containerExitReason = s"Container $containerId exited abnormally with exit status $exitStatus, and was marked as failed."
         }
 
-        if (completedContainer.getExitStatus() != 0) {
-          val exitStatus = {
+        if (exitStatus != 0) {
+          val exitReason = {
             if (isExecutorNonZeroExitNormal) {
-              ExecutorExitedNormally(completedContainer.getExitStatus(), containerExitReason)
+              ExecutorExitedNormally(completedContainer.getExitStatus, containerExitReason)
             } else {
-              ExecutorExitedAbnormally(completedContainer.getExitStatus(), containerExitReason)
+              ExecutorExitedAbnormally(completedContainer.getExitStatus, containerExitReason)
             }
           }
-          completedNonZeroContainerExitCodes.put(containerId, exitStatus)
+          completedNonZeroContainerExitCodes.put(containerId, exitReason)
         }
       }
 
diff --git a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
index 7509000771d94..8df8d13cea911 100644
--- a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
+++ b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
@@ -242,4 +242,23 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     assert(pmemMsg.contains("2.1 MB of 2 GB physical memory used."))
   }
 
+  test("Getting executor loss reason should depend on how container was terminated") {
+    val container1 = createContainer("host1")
+    val container2 = createContainer("host2")
+    val container3 = createContainer("host3")
+    val container4 = createContainer("host4")
+    val container5 = createContainer("host5")
+
+    val handler = createAllocator(5)
+    handler.handleAllocatedContainers(Array(container1, container2, container3, container4, container5))
+    val preemptedContainerStatus = ContainerStatus.newInstance(container1.getId(), ContainerState.COMPLETE, "Preempted", ContainerExitStatus.PREEMPTED)
+    val vmemLimitExceededContainerStatus = ContainerStatus.newInstance(container2.getId(), ContainerState.COMPLETE, "Vmem limit exceeded", -103)
+    val pmemLimitExceededContainerStatus = ContainerStatus.newInstance(container3.getId(), ContainerState.COMPLETE, "pmem limit exceeded", -104)
+    val unknownErrorContainerSTatus = ContainerStatus.newInstance(container4.getId(), ContainerState.COMPLETE, "Unknown error", 123)
+
+    val killedExecutorId =  handler.executorIdToContainer.filter(_._2.getId().equals(container5)).iterator().next()._1
+    handler.killExecutor(killedExecutorId)
+    assert(handler.getExecutorLossReason(killedExecutorId).isInstanceOf[ExecutorExitedNormally])
+  }
+
 }

From 08d2bdd058856b6b1ecb514a8e1066bfacf7518c Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Wed, 8 Jul 2015 20:52:49 -0700
Subject: [PATCH 03/10] Completing unit tests.

---
 .../spark/scheduler/TaskSetManagerSuite.scala | 23 +++++++
 .../deploy/yarn/YarnAllocatorSuite.scala      | 66 ++++++++++++++-----
 2 files changed, 73 insertions(+), 16 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
index 3a3254cd4b3a6..a2ace5bece17b 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala
@@ -507,6 +507,29 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     assert(manager.myLocalityLevels.sameElements(Array(NO_PREF, ANY)))
   }
 
+  test("Executors are added but exit normally while running tasks") {
+    sc = new SparkContext("local", "test")
+    val sched = new FakeTaskScheduler(sc)
+    val taskSet = FakeTask.createTaskSet(4,
+      Seq(TaskLocation("host1", "execA")),
+      Seq(TaskLocation("host1", "execB")),
+      Seq(TaskLocation("host2", "execC")),
+      Seq())
+    val manager = new TaskSetManager(sched, taskSet, 1, new ManualClock)
+    sched.addExecutor("execA", "host1")
+    manager.executorAdded()
+    sched.addExecutor("execC", "host2")
+    manager.executorAdded()
+    assert(manager.resourceOffer("exec1", "host1", ANY).isDefined)
+    sched.removeExecutor("execA")
+    manager.executorLost("execA", "host1", ExecutorExitedNormally(143, "Normal termination"))
+    assert(!sched.taskSetsFailed.contains(taskSet.id))
+    assert(manager.resourceOffer("execC", "host2", ANY).isDefined)
+    sched.removeExecutor("execC")
+    manager.executorLost("execC", "host2", ExecutorExitedAbnormally(1, "Abnormal termination"))
+    assert(sched.taskSetsFailed.contains(taskSet.id))
+  }
+
   test("test RACK_LOCAL tasks") {
     // Assign host1 to rack1
     FakeRackUtil.assignHostToRack("host1", "rack1")
diff --git a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
index 8df8d13cea911..c4c56deaaec80 100644
--- a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
+++ b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
@@ -18,10 +18,13 @@
 package org.apache.spark.deploy.yarn
 
 import java.util.{Arrays, List => JList}
+import org.mockito.Mockito._
+import scala.collection.JavaConversions._
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.CommonConfigurationKeysPublic
 import org.apache.hadoop.net.DNSToSwitchMapping
+import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse
 import org.apache.hadoop.yarn.api.records._
 import org.apache.hadoop.yarn.client.api.AMRMClient
 import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
@@ -30,7 +33,7 @@ import org.apache.spark.{SecurityManager, SparkFunSuite}
 import org.apache.spark.SparkConf
 import org.apache.spark.deploy.yarn.YarnSparkHadoopUtil._
 import org.apache.spark.deploy.yarn.YarnAllocator._
-import org.apache.spark.scheduler.SplitInfo
+import org.apache.spark.scheduler.{ExecutorExitedAbnormally, ExecutorExitedNormally, SplitInfo}
 
 import org.scalatest.{BeforeAndAfterEach, Matchers}
 
@@ -69,7 +72,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
   var containerNum = 0
 
   override def beforeEach() {
-    rmClient = AMRMClient.createAMRMClient()
+    rmClient = spy(AMRMClient.createAMRMClient())
     rmClient.init(conf)
     rmClient.start()
   }
@@ -243,22 +246,53 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
   }
 
   test("Getting executor loss reason should depend on how container was terminated") {
-    val container1 = createContainer("host1")
-    val container2 = createContainer("host2")
-    val container3 = createContainer("host3")
-    val container4 = createContainer("host4")
-    val container5 = createContainer("host5")
-
-    val handler = createAllocator(5)
-    handler.handleAllocatedContainers(Array(container1, container2, container3, container4, container5))
-    val preemptedContainerStatus = ContainerStatus.newInstance(container1.getId(), ContainerState.COMPLETE, "Preempted", ContainerExitStatus.PREEMPTED)
-    val vmemLimitExceededContainerStatus = ContainerStatus.newInstance(container2.getId(), ContainerState.COMPLETE, "Vmem limit exceeded", -103)
-    val pmemLimitExceededContainerStatus = ContainerStatus.newInstance(container3.getId(), ContainerState.COMPLETE, "pmem limit exceeded", -104)
-    val unknownErrorContainerSTatus = ContainerStatus.newInstance(container4.getId(), ContainerState.COMPLETE, "Unknown error", 123)
-
-    val killedExecutorId =  handler.executorIdToContainer.filter(_._2.getId().equals(container5)).iterator().next()._1
+    val preemptedContainer = createContainer("host1")
+    val vmemExceededContainer = createContainer("host2")
+    val pmemExceededContainer = createContainer("host3")
+    val unknownErrorContainer= createContainer("host4")
+    val killedContainer = createContainer("host5")
+    val normalExitContainer = createContainer("host6")
+
+    val containersToStatusAndExpectedLossReasons = Map(
+      preemptedContainer ->
+          (ContainerStatus.newInstance(preemptedContainer.getId(), ContainerState.COMPLETE, "Preempted", ContainerExitStatus.PREEMPTED),
+          classOf[ExecutorExitedNormally]),
+      vmemExceededContainer ->
+          (ContainerStatus.newInstance(vmemExceededContainer.getId(), ContainerState.COMPLETE, "Vmem limit exceeded", -103),
+          classOf[ExecutorExitedAbnormally]),
+      pmemExceededContainer ->
+          (ContainerStatus.newInstance(pmemExceededContainer.getId(), ContainerState.COMPLETE, "pmem limit exceeded", -104),
+          classOf[ExecutorExitedAbnormally]),
+      unknownErrorContainer ->
+          (ContainerStatus.newInstance(unknownErrorContainer.getId(), ContainerState.COMPLETE, "Unknown error", 123),
+          classOf[ExecutorExitedAbnormally]),
+      normalExitContainer ->
+        (ContainerStatus.newInstance(normalExitContainer.getId(), ContainerState.COMPLETE, "Container exited normally", 0),
+            classOf[ExecutorExitedNormally])
+    )
+
+    val handler = createAllocator(6)
+    val mockAllocateResponse = mock(classOf[AllocateResponse])
+    handler.requestTotalExecutors(6)
+    handler.updateResourceRequests()
+    handler.handleAllocatedContainers(containersToStatusAndExpectedLossReasons.keys.toSeq ++ Seq(killedContainer))
+    doReturn(mockAllocateResponse).when(rmClient).allocate(0.1f)
+    when(mockAllocateResponse.getAllocatedContainers).thenReturn(Seq())
+    when(mockAllocateResponse.getCompletedContainersStatuses)
+        .thenReturn(containersToStatusAndExpectedLossReasons.values.map(_._1).toSeq)
+        .thenReturn(Seq())
+
+    val killedExecutorId = getExecutorIdForContainer(handler, killedContainer)
     handler.killExecutor(killedExecutorId)
     assert(handler.getExecutorLossReason(killedExecutorId).isInstanceOf[ExecutorExitedNormally])
+    containersToStatusAndExpectedLossReasons.foreach({ testContainer =>
+      val executorId = getExecutorIdForContainer(handler, testContainer._1)
+      assert(testContainer._2._2.isInstance(handler.getExecutorLossReason(executorId)))
+    })
   }
 
+  private def getExecutorIdForContainer(handler: YarnAllocator, container: Container): String = {
+    handler.executorIdToContainer.filter(_._2.equals(container)).iterator.next()._1
+  }
 }
+

From 62e43428a3a89dda395bd6c8d47066ae460165a3 Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 9 Jul 2015 11:07:25 -0700
Subject: [PATCH 04/10] Removing rogue .orig file

---
 .../cluster/YarnSchedulerBackend.scala.orig   | 209 ------------------
 1 file changed, 209 deletions(-)
 delete mode 100644 core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig
deleted file mode 100644
index fb1c3af55bcc0..0000000000000
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala.orig
+++ /dev/null
@@ -1,209 +0,0 @@
-/*
- * 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.scheduler.cluster
-
-import java.util.Properties
-
-import scala.collection.mutable.ArrayBuffer
-import scala.concurrent.{Future, ExecutionContext}
-
-import org.apache.spark.{Logging, SparkContext}
-import org.apache.spark.rpc._
-import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
-import org.apache.spark.scheduler._
-import org.apache.spark.ui.JettyUtils
-import org.apache.spark.util.{ThreadUtils, RpcUtils}
-
-import scala.util.control.NonFatal
-
-/**
- * Abstract Yarn scheduler backend that contains common logic
- * between the client and cluster Yarn scheduler backends.
- */
-private[spark] abstract class YarnSchedulerBackend(
-    scheduler: TaskSchedulerImpl,
-    sc: SparkContext)
-  extends CoarseGrainedSchedulerBackend(scheduler, sc.env.rpcEnv) {
-
-  if (conf.getOption("spark.scheduler.minRegisteredResourcesRatio").isEmpty) {
-    minRegisteredRatio = 0.8
-  }
-
-  protected var totalExpectedExecutors = 0
-
-  private val yarnSchedulerEndpoint = rpcEnv.setupEndpoint(
-    YarnSchedulerBackend.ENDPOINT_NAME, new YarnSchedulerEndpoint(rpcEnv))
-
-<<<<<<< Updated upstream
-  private implicit val askTimeout = RpcUtils.askRpcTimeout(sc.conf)
-||||||| merged common ancestors
-  private implicit val askTimeout = RpcUtils.askTimeout(sc.conf)
-=======
-  private val executorDisconnectedHandlerPool = ThreadUtils.newDaemonSingleThreadExecutor("executor-disconnected-handler")
-
-  private implicit val askTimeout = RpcUtils.askTimeout(sc.conf)
->>>>>>> Stashed changes
-
-  /**
-   * Request executors from the ApplicationMaster by specifying the total number desired.
-   * This includes executors already pending or running.
-   */
-  override def doRequestTotalExecutors(requestedTotal: Int): Boolean = {
-    yarnSchedulerEndpoint.askWithRetry[Boolean](RequestExecutors(requestedTotal))
-  }
-
-  /**
-   * Request that the ApplicationMaster kill the specified executors.
-   */
-  override def doKillExecutors(executorIds: Seq[String]): Boolean = {
-    yarnSchedulerEndpoint.askWithRetry[Boolean](KillExecutors(executorIds))
-  }
-
-  override def sufficientResourcesRegistered(): Boolean = {
-    totalRegisteredExecutors.get() >= totalExpectedExecutors * minRegisteredRatio
-  }
-
-  /**
-   * Add filters to the SparkUI.
-   */
-  private def addWebUIFilter(
-      filterName: String,
-      filterParams: Map[String, String],
-      proxyBase: String): Unit = {
-    if (proxyBase != null && proxyBase.nonEmpty) {
-      System.setProperty("spark.ui.proxyBase", proxyBase)
-    }
-
-    val hasFilter =
-      filterName != null && filterName.nonEmpty &&
-      filterParams != null && filterParams.nonEmpty
-    if (hasFilter) {
-      logInfo(s"Add WebUI Filter. $filterName, $filterParams, $proxyBase")
-      conf.set("spark.ui.filters", filterName)
-      filterParams.foreach { case (k, v) => conf.set(s"spark.$filterName.param.$k", v) }
-      scheduler.sc.ui.foreach { ui => JettyUtils.addFilters(ui.getHandlers, conf) }
-    }
-  }
-
-  /**
-   * Override the DriverEndpoint to add extra logic for the case when an executor is disconnected.
-   * We should check the cluster manager and find if the loss of the executor was caused by YARN
-   * force killing it due to preemption.
-   */
-  class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
-      extends DriverEndpoint(rpcEnv, sparkProperties) {
-
-    override def onDisconnected(rpcAddress: RpcAddress): Unit = {
-      addressToExecutorId.get(rpcAddress).foreach({ executorId =>
-        executorDisconnectedHandlerPool.submit(new Runnable() {
-          override def run(): Unit = {
-            val wasExecutorForceKilledNormally =
-              yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
-            wasExecutorForceKilledNormally match {
-              case Some(killReason) =>
-                driverEndpoint.send(RemoveExecutor(executorId, killReason))
-              case None =>
-                driverEndpoint.send(RemoveExecutor(executorId, SlaveLost("Executor was terminated for an unknown reason.")))
-            }
-          }
-        })
-      })
-    }
-  }
-
-  override def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = {
-    new YarnDriverEndpoint(rpcEnv, properties)
-  }
-
-  /**
-   * An [[RpcEndpoint]] that communicates with the ApplicationMaster.
-   */
-  private class YarnSchedulerEndpoint(override val rpcEnv: RpcEnv)
-    extends ThreadSafeRpcEndpoint with Logging {
-    private var amEndpoint: Option[RpcEndpointRef] = None
-
-    private val askAmThreadPool =
-      ThreadUtils.newDaemonCachedThreadPool("yarn-scheduler-ask-am-thread-pool")
-    implicit val askAmExecutor = ExecutionContext.fromExecutor(askAmThreadPool)
-
-    override def receive: PartialFunction[Any, Unit] = {
-      case RegisterClusterManager(am) =>
-        logInfo(s"ApplicationMaster registered as $am")
-        amEndpoint = Some(am)
-
-      case AddWebUIFilter(filterName, filterParams, proxyBase) =>
-        addWebUIFilter(filterName, filterParams, proxyBase)
-
-    }
-
-    override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
-      case r: RequestExecutors =>
-        amEndpoint match {
-          case Some(am) =>
-            Future {
-              context.reply(am.askWithRetry[Boolean](r))
-            } onFailure {
-              case NonFatal(e) =>
-                logError(s"Sending $r to AM was unsuccessful", e)
-                context.sendFailure(e)
-            }
-          case None =>
-            logWarning("Attempted to request executors before the AM has registered!")
-            context.reply(false)
-        }
-
-      case k: KillExecutors =>
-        amEndpoint match {
-          case Some(am) =>
-            Future {
-              context.reply(am.askWithRetry[Boolean](k))
-            } onFailure {
-              case NonFatal(e) =>
-                logError(s"Sending $k to AM was unsuccessful", e)
-                context.sendFailure(e)
-            }
-          case None =>
-            logWarning("Attempted to kill executors before the AM has registered!")
-            context.reply(false)
-        }
-
-      case c: GetExecutorLossReason =>
-        amEndpoint match {
-          case Some(am) => context.reply(am.askWithRetry[ExecutorLossReason](c))
-          case None =>
-            logWarning("Attempted to check if an executor exited normally before the AM has registered!")
-        }
-
-    }
-
-    override def onDisconnected(remoteAddress: RpcAddress): Unit = {
-      if (amEndpoint.exists(_.address == remoteAddress)) {
-        logWarning(s"ApplicationMaster has disassociated: $remoteAddress")
-      }
-    }
-
-    override def onStop(): Unit = {
-      askAmThreadPool.shutdownNow()
-    }
-  }
-}
-
-
-private[spark] object YarnSchedulerBackend {
-  val ENDPOINT_NAME = "YarnScheduler"
-}

From bd1056ec2554158e1a7eb3509732d89334552fe8 Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 9 Jul 2015 11:17:52 -0700
Subject: [PATCH 05/10] Put back extends CoarseGrainedClusterMessage

---
 .../spark/scheduler/cluster/CoarseGrainedClusterMessage.scala   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
index e8cc4b135cf0f..faedf6568f5d7 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
@@ -71,7 +71,7 @@ private[spark] object CoarseGrainedClusterMessages {
 
   case object StopExecutors extends CoarseGrainedClusterMessage
 
-  case class RemoveExecutor(executorId: String, reason: ExecutorLossReason)
+  case class RemoveExecutor(executorId: String, reason: ExecutorLossReason) extends CoarseGrainedClusterMessage
 
   case class SetupDriver(driver: RpcEndpointRef) extends CoarseGrainedClusterMessage
 

From 2bbdcfd94c66e82b825977676c76d1cb05c1d9a9 Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 9 Jul 2015 11:19:36 -0700
Subject: [PATCH 06/10] Removing unnecessary import

---
 .../apache/spark/scheduler/cluster/YarnSchedulerBackend.scala   | 2 --
 1 file changed, 2 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index 95f40c4b1e990..018d7dc1a87b2 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.scheduler.cluster
 
-import java.util.Properties
-
 import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.{Future, ExecutionContext}

From a697c7cc38ac417eb6a7fed658acd4782824d8fa Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 9 Jul 2015 11:21:28 -0700
Subject: [PATCH 07/10] Organizing more imports

---
 .../spark/scheduler/cluster/YarnSchedulerBackend.scala       | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index 018d7dc1a87b2..d72395299a827 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -17,8 +17,7 @@
 
 package org.apache.spark.scheduler.cluster
 
-import scala.collection.mutable
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable.{ArrayBuffer, HashSet}
 import scala.concurrent.{Future, ExecutionContext}
 
 import org.apache.spark.{Logging, SparkContext}
@@ -99,7 +98,7 @@ private[spark] abstract class YarnSchedulerBackend(
   private class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
       extends DriverEndpoint(rpcEnv, sparkProperties) {
 
-    private val pendingDisconnectedExecutors = new mutable.HashSet[String]
+    private val pendingDisconnectedExecutors = new HashSet[String]
     private val handleDisconnectedExecutorThreadPool =
       ThreadUtils.newDaemonCachedThreadPool("yarn-driver-endpoint-handle-disconnected-executor-thread-pool")
 

From e791cd6378d052ce65d8835d43790b75b8c7c0cd Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Tue, 4 Aug 2015 13:51:56 -0700
Subject: [PATCH 08/10] Fixing a compiler error caused by the merge from
 master.

Also I changed the semantics of YarnAllocator.getExecutorLossReason() to
better "clean up" the completed executor map. Also
processCompletedContainers() now always adds an executor reason to the
completed executor exit reason map regardless of exit status as it is
expected for the client to always call getExecutorLossReason().
---
 .../mesos/CoarseMesosSchedulerBackend.scala   |  2 +-
 .../spark/deploy/yarn/YarnAllocator.scala     | 40 +++++++++----------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
index 826dd5f364bb2..5bcb2202e3381 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
@@ -359,7 +359,7 @@ private[spark] class CoarseMesosSchedulerBackend(
         if (slaveIdToTaskId.contains(slaveId)) {
           val taskId: Int = slaveIdToTaskId.get(slaveId)
           taskIdToSlaveId.remove(taskId)
-          removeExecutor(sparkExecutorId(slaveId, taskId.toString), reason)
+          removeExecutor(sparkExecutorId(slaveId, taskId.toString), SlaveLost(reason))
         }
         // TODO: This assumes one Spark executor per Mesos slave,
         // which may no longer be true after SPARK-5095
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index 6d77b41370665..c09d802e9f019 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -39,7 +39,6 @@ import org.apache.log4j.{Level, Logger}
 import org.apache.spark.{Logging, SecurityManager, SparkConf}
 import org.apache.spark.deploy.yarn.YarnSparkHadoopUtil._
 import org.apache.spark.rpc.RpcEndpointRef
-import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._
 
 /**
@@ -68,7 +67,7 @@ private[yarn] class YarnAllocator(
 
   import YarnAllocator._
 
-  private val NORMAL_CONTAINER_EXIT_STATUS = ExecutorExitedNormally(0, "Executor exited normally.")
+  private val UNKNOWN_CONTAINER_EXIT_STATUS = ExecutorExitedAbnormally(-1, "Executor exited for an unknown reason.")
 
   // RackResolver logs an INFO message whenever it resolves a rack, which is way too often.
   if (Logger.getLogger(classOf[RackResolver]).getLevel == null) {
@@ -98,6 +97,7 @@ private[yarn] class YarnAllocator(
 
   private var numUnexpectedContainerRelease = 0L
   private val containerIdToExecutorId = new HashMap[ContainerId, String]
+  private val completedExecutorExitReasons = new HashMap[String, ExecutorLossReason]
 
   // Executor memory in MB.
   protected val executorMemory = args.executorMemory
@@ -198,19 +198,19 @@ private[yarn] class YarnAllocator(
       containerIdToExecutorId.remove(container.getId)
       internalReleaseContainer(container)
       numExecutorsRunning -= 1
-      killedExecutors += executorId
     } else {
       logWarning(s"Attempted to kill unknown executor $executorId!")
     }
   }
 
+  /**
+   * Gets the executor loss reason for a disconnected executor.
+   * Note that this method is expected to be called exactly once per executor ID.
+   */
   def getExecutorLossReason(executorId: String): ExecutorLossReason = synchronized {
     allocateResources()
-    if (killedExecutors.contains(executorId)) {
-        NORMAL_CONTAINER_EXIT_STATUS
-    } else {
-      completedNonZeroContainerExitCodes.getOrElse(executorIdToContainer(executorId).getId, NORMAL_CONTAINER_EXIT_STATUS)
-    }
+    // Expect to be asked for a loss reason once and exactly once.
+    completedExecutorExitReasons.remove(executorId).getOrElse(executorId, UNKNOWN_CONTAINER_EXIT_STATUS)
   }
 
   /**
@@ -437,7 +437,7 @@ private[yarn] class YarnAllocator(
     for (completedContainer <- completedContainers) {
       val containerId = completedContainer.getContainerId
       val alreadyReleased = releasedContainers.remove(containerId)
-      if (!alreadyReleased) {
+      val exitReason = if (!alreadyReleased) {
         // Decrement the number of executors running. The next iteration of
         // the ApplicationMaster's reporting thread will take care of allocating.
         numExecutorsRunning -= 1
@@ -475,16 +475,16 @@ private[yarn] class YarnAllocator(
           containerExitReason = s"Container $containerId exited abnormally with exit status $exitStatus, and was marked as failed."
         }
 
-        if (exitStatus != 0) {
-          val exitReason = {
-            if (isExecutorNonZeroExitNormal) {
-              ExecutorExitedNormally(completedContainer.getExitStatus, containerExitReason)
-            } else {
-              ExecutorExitedAbnormally(completedContainer.getExitStatus, containerExitReason)
-            }
-          }
-          completedNonZeroContainerExitCodes.put(containerId, exitReason)
+        if (exitStatus == 0) {
+          ExecutorExitedNormally(0, s"Executor for container $containerId exited normally.")
+        } else if (isExecutorNonZeroExitNormal) {
+          ExecutorExitedNormally(completedContainer.getExitStatus, containerExitReason)
+        } else {
+          ExecutorExitedAbnormally(completedContainer.getExitStatus, containerExitReason)
         }
+      } else {
+        ExecutorExitedNormally(completedContainer.getExitStatus,
+          s"Container $containerId exited from explicit termination request.")
       }
 
       if (allocatedContainerToHostMap.containsKey(containerId)) {
@@ -503,13 +503,13 @@ private[yarn] class YarnAllocator(
 
       containerIdToExecutorId.remove(containerId).foreach { eid =>
         executorIdToContainer.remove(eid)
+        completedExecutorExitReasons.put(eid, exitReason)
 
         if (!alreadyReleased) {
           // The executor could have gone away (like no route to host, node failure, etc)
           // Notify backend about the failure of the executor
           numUnexpectedContainerRelease += 1
-          driverRef.send(RemoveExecutor(eid,
-            s"Yarn deallocated the executor $eid (container $containerId)"))
+          driverRef.send(RemoveExecutor(eid, exitReason))
         }
       }
     }

From aa69b6fc2fb2db7c98a81bebdf02f23fd18d2b52 Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 6 Aug 2015 14:32:45 -0700
Subject: [PATCH 09/10] Adding a stronger assertion, and fixing compiler error

---
 .../scala/org/apache/spark/deploy/yarn/YarnAllocator.scala     | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index c09d802e9f019..f5c1d6ab8f40a 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -210,7 +210,8 @@ private[yarn] class YarnAllocator(
   def getExecutorLossReason(executorId: String): ExecutorLossReason = synchronized {
     allocateResources()
     // Expect to be asked for a loss reason once and exactly once.
-    completedExecutorExitReasons.remove(executorId).getOrElse(executorId, UNKNOWN_CONTAINER_EXIT_STATUS)
+    assert(completedExecutorExitReasons.contains(executorId))
+    completedExecutorExitReasons.remove(executorId).getOrElse(UNKNOWN_CONTAINER_EXIT_STATUS)
   }
 
   /**

From e3e827fa1dd5afcedff5edd293aee001d404bb5b Mon Sep 17 00:00:00 2001
From: mcheah <mcheah@palantir.com>
Date: Thu, 6 Aug 2015 17:40:24 -0700
Subject: [PATCH 10/10] Fix all the scalastyle errors

---
 .../scala/org/apache/spark/TaskEndReason.scala   | 11 ++++++++---
 .../spark/scheduler/ExecutorLossReason.scala     |  8 ++++++--
 .../spark/scheduler/TaskSchedulerImpl.scala      |  3 ++-
 .../apache/spark/scheduler/TaskSetManager.scala  |  7 ++++---
 .../cluster/CoarseGrainedClusterMessage.scala    |  7 ++++---
 .../cluster/CoarseGrainedSchedulerBackend.scala  | 11 ++++++++---
 .../scheduler/cluster/YarnSchedulerBackend.scala | 16 ++++++++++------
 .../apache/spark/deploy/yarn/YarnAllocator.scala |  6 ++++--
 .../spark/deploy/yarn/YarnAllocatorSuite.scala   |  1 -
 9 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/TaskEndReason.scala b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
index a6641b167cf7a..03b1f92211af7 100644
--- a/core/src/main/scala/org/apache/spark/TaskEndReason.scala
+++ b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
@@ -173,9 +173,14 @@ case class ExecutorLostFailure(execId: String) extends TaskFailedReason {
  * behavior.
  */
 @DeveloperApi
-case class ExecutorForTaskExited(taskId: Long, execId: String, exitReason: String, exitCode: Int) extends TaskFailedReason {
-  override def toErrorString: String = s"Task with ID $taskId had its executor ${execId} exit normally with exit code "
-    s"$exitCode, due to the following reason: $exitReason."
+case class ExecutorForTaskExited(
+    taskId: Long,
+    execId: String,
+    exitReason: String,
+    exitCode: Int)
+    extends TaskFailedReason {
+  override def toErrorString: String = s"Task with ID $taskId had its executor ${execId}" +
+    s" exit normally with exit code $exitCode, due to the following reason: $exitReason."
 }
 
 /**
diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
index 94e40f65541af..fac9e928947d2 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
@@ -32,7 +32,9 @@ private[spark] case class ExecutorExitedAbnormally(val exitCode: Int, reason: St
 }
 
 private[spark] object ExecutorExitedAbnormally {
-  def apply(exitCode: Int): ExecutorExitedAbnormally = ExecutorExitedAbnormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
+  def apply(exitCode: Int): ExecutorExitedAbnormally = {
+    ExecutorExitedAbnormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
+  }
 }
 
 private[spark]
@@ -41,7 +43,9 @@ case class ExecutorExitedNormally(val exitCode: Int, reason: String)
 }
 
 private[spark] object ExecutorExitedNormally {
-  def apply(exitCode: Int): ExecutorExitedNormally = ExecutorExitedNormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
+  def apply(exitCode: Int): ExecutorExitedNormally = {
+    ExecutorExitedNormally(exitCode, ExecutorExitCode.explainExitCode(exitCode))
+  }
 }
 
 private[spark]
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 86672b9ff50f2..e2eeba07d881a 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -332,7 +332,8 @@ private[spark] class TaskSchedulerImpl(
           // We lost this entire executor, so remember that it's gone
           val execId = taskIdToExecutorId(tid)
           if (activeExecutorIds.contains(execId)) {
-            removeExecutor(execId, SlaveLost(s"Task $tid was lost, so marking the executor as lost as well."))
+            removeExecutor(execId, SlaveLost(s"Task $tid was lost, so" +
+              s" marking the executor as lost as well."))
             failedExecutor = Some(execId)
           }
         }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 6a58ffff3f22b..e5c263749d085 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -708,8 +708,8 @@ private[spark] class TaskSetManager(
         }
 
       case e: ExecutorForTaskExited =>
-        logWarning(s"Task $tid failed because while it was being computed, its executor exited normally." +
-          s" Not marking the task as failed.")
+        logWarning(s"Task $tid failed because while it was being computed, its executor" +
+          s" exited normally. Not marking the task as failed.")
       case e: TaskFailedReason =>  // TaskResultLost, TaskKilled, and others
         logWarning(failureReason)
 
@@ -815,7 +815,8 @@ private[spark] class TaskSetManager(
     for ((tid, info) <- taskInfos if info.running && info.executorId == execId) {
       // Also re-enqueue any tasks that were running on the node
       val executorFailureReason = reason match {
-        case exited: ExecutorExitedNormally => ExecutorForTaskExited(tid, execId, exited.reason, exited.exitCode)
+        case exited: ExecutorExitedNormally =>
+          ExecutorForTaskExited(tid, execId, exited.reason, exited.exitCode)
         case default => ExecutorLostFailure(execId)
       }
       handleFailedTask(tid, TaskState.FAILED, executorFailureReason)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
index 6c9067e8bae77..b45d18008daf4 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
@@ -71,7 +71,8 @@ private[spark] object CoarseGrainedClusterMessages {
 
   case object StopExecutors extends CoarseGrainedClusterMessage
 
-  case class RemoveExecutor(executorId: String, reason: ExecutorLossReason) extends CoarseGrainedClusterMessage
+  case class RemoveExecutor(executorId: String, reason: ExecutorLossReason)
+    extends CoarseGrainedClusterMessage
 
   case class SetupDriver(driver: RpcEndpointRef) extends CoarseGrainedClusterMessage
 
@@ -94,8 +95,8 @@ private[spark] object CoarseGrainedClusterMessages {
     extends CoarseGrainedClusterMessage
 
   // Check if an executor was force-killed but for a normal reason
-  // This could be the case if e.g. the cluster manager supports killing an executor to move
-  // it elsewhere or to kill an executor in order to free resources
+  // This could be the case if e.g. the cluster manager supports killing an executor to
+  // move it elsewhere or to kill an executor in order to free resources
   case class GetExecutorLossReason(executorId: String)
 
   case class KillExecutors(executorIds: Seq[String]) extends CoarseGrainedClusterMessage
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 9cda77c063f0f..a67f9a95960fb 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -185,7 +185,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
-      addressToExecutorId.get(remoteAddress).foreach(removeExecutor(_, SlaveLost("remote Rpc client disassociated")))
+      addressToExecutorId
+        .get(remoteAddress)
+        .foreach(removeExecutor(_, SlaveLost("remote Rpc client disassociated")))
     }
 
     // Make fake resource offers on just one executor
@@ -262,10 +264,13 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     // TODO (prashant) send conf instead of properties
-    driverEndpoint = rpcEnv.setupEndpoint(CoarseGrainedSchedulerBackend.ENDPOINT_NAME, createDriverEndpoint(properties))
+    driverEndpoint = rpcEnv.setupEndpoint(
+        CoarseGrainedSchedulerBackend.ENDPOINT_NAME, createDriverEndpoint(properties))
   }
 
-  protected def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = new DriverEndpoint(rpcEnv, properties)
+  protected def createDriverEndpoint(
+      properties: ArrayBuffer[(String, String)]): DriverEndpoint
+    = new DriverEndpoint(rpcEnv, properties)
 
   def stopExecutors() {
     try {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index 3756a0948524c..dfcee1567f798 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -101,7 +101,7 @@ private[spark] abstract class YarnSchedulerBackend(
 
     private val pendingDisconnectedExecutors = new HashSet[String]
     private val handleDisconnectedExecutorThreadPool =
-      ThreadUtils.newDaemonCachedThreadPool("yarn-driver-endpoint-handle-disconnected-executor-thread-pool")
+      ThreadUtils.newDaemonCachedThreadPool("yarn-driver-handle-lost-executor-thread-pool")
 
     override def onDisconnected(rpcAddress: RpcAddress): Unit = {
       addressToExecutorId.get(rpcAddress).foreach({ executorId =>
@@ -110,11 +110,15 @@ private[spark] abstract class YarnSchedulerBackend(
             pendingDisconnectedExecutors.add(executorId)
             handleDisconnectedExecutorThreadPool.submit(new Runnable() {
               override def run(): Unit = {
-                val executorLossReason = yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](GetExecutorLossReason(executorId))
+                val executorLossReason =
+                  yarnSchedulerEndpoint.askWithRetry[Option[ExecutorLossReason]](
+                      GetExecutorLossReason(executorId))
                 executorLossReason match {
-                  case Some(reason) => driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, reason))
+                  case Some(reason) =>
+                    driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, reason))
                   case None =>
-                    logWarning(s"Attempted to get executor loss reason for $rpcAddress but got no response. Marking as slave lost.")
+                    logWarning(s"Attempted to get executor loss reason" +
+                      s" for $rpcAddress but got no response. Marking as slave lost.")
                     driverEndpoint.askWithRetry[Boolean](RemoveExecutor(executorId, SlaveLost()))
                 }
                 pendingDisconnectedExecutors.synchronized {
@@ -197,7 +201,8 @@ private[spark] abstract class YarnSchedulerBackend(
                 context.sendFailure(e)
             }
           case None =>
-            logWarning("Attempted to check if an executor exited normally before the AM has registered!")
+            logWarning("Attempted to check if an executor exited normally" +
+              " before the AM has registered!")
             context.reply(None)
         }
     }
@@ -214,7 +219,6 @@ private[spark] abstract class YarnSchedulerBackend(
   }
 }
 
-
 private[spark] object YarnSchedulerBackend {
   val ENDPOINT_NAME = "YarnScheduler"
 }
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index f5c1d6ab8f40a..8734faba10554 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -67,7 +67,8 @@ private[yarn] class YarnAllocator(
 
   import YarnAllocator._
 
-  private val UNKNOWN_CONTAINER_EXIT_STATUS = ExecutorExitedAbnormally(-1, "Executor exited for an unknown reason.")
+  private val UNKNOWN_CONTAINER_EXIT_STATUS =
+    ExecutorExitedAbnormally(-1, "Executor exited for an unknown reason.")
 
   // RackResolver logs an INFO message whenever it resolves a rack, which is way too often.
   if (Logger.getLogger(classOf[RackResolver]).getLevel == null) {
@@ -473,7 +474,8 @@ private[yarn] class YarnAllocator(
             ". Exit status: " + completedContainer.getExitStatus +
             ". Diagnostics: " + completedContainer.getDiagnostics)
           numExecutorsFailed += 1
-          containerExitReason = s"Container $containerId exited abnormally with exit status $exitStatus, and was marked as failed."
+          containerExitReason = s"Container $containerId exited abnormally with exit" +
+            s" status $exitStatus, and was marked as failed."
         }
 
         if (exitStatus == 0) {
diff --git a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
index 0b661fb8e6351..7173a62be2712 100644
--- a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
+++ b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
@@ -273,4 +273,3 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     assert(pmemMsg.contains("2.1 MB of 2 GB physical memory used."))
   }
 }
-