[SPARK-8167] Tasks that fail from YARN preemption should not fail job

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

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]

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 = {

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]
 }

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) {

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()

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) extends CoarseGrainedClusterMessage
 
   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
 
 }

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 {

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 = {

core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala

@@ -17,12 +17,13 @@
 
 package org.apache.spark.scheduler.cluster
 
+import scala.collection.mutable.{ArrayBuffer, HashSet}
 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}
 
@@ -89,6 +90,47 @@ 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.
+   */
+  private class YarnDriverEndpoint(rpcEnv: RpcEnv, sparkProperties: ArrayBuffer[(String, String)])
+      extends DriverEndpoint(rpcEnv, sparkProperties) {
+
+    private val pendingDisconnectedExecutors = new 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 =>
+        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)
+                }
+              }
+            })
+          }
+        }
+      })
+    }
+  }
+
+  override def createDriverEndpoint(properties: ArrayBuffer[(String, String)]): DriverEndpoint = {
+    new YarnDriverEndpoint(rpcEnv, properties)
+  }
+
   /**
    * An [[RpcEndpoint]] that communicates with the ApplicationMaster.
    */
@@ -141,6 +183,20 @@ private[spark] abstract class YarnSchedulerBackend(
             context.reply(false)
         }
 
+      case c: GetExecutorLossReason =>
+        amEndpoint match {
+          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 = {
@@ -155,6 +211,7 @@ private[spark] abstract class YarnSchedulerBackend(
   }
 }
 
+
 private[spark] object YarnSchedulerBackend {
   val ENDPOINT_NAME = "YarnScheduler"
 }

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"))
       }
     }
   }

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 = {