Use driver <-> AM communication to send "remove executor" request.

Author: Marcelo Vanzin

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

@@ -108,6 +108,8 @@ private[spark] abstract class YarnSchedulerBackend(
       case AddWebUIFilter(filterName, filterParams, proxyBase) =>
         addWebUIFilter(filterName, filterParams, proxyBase)
 
+      case RemoveExecutor(executorId, reason) =>
+        removeExecutor(executorId, reason)
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {

yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala

@@ -77,8 +77,6 @@ private[spark] class ApplicationMaster(
   @volatile private var allocator: YarnAllocator = _
   private val allocatorLock = new Object()
 
-  @volatile private var backend: CoarseGrainedSchedulerBackend = _
-
   // Fields used in client mode.
   private var rpcEnv: RpcEnv = null
   private var amEndpoint: RpcEndpointRef = _
@@ -220,20 +218,22 @@ private[spark] class ApplicationMaster(
     }
   }
 
-  private def sparkContextInitialized(sc: SparkContext, backend: CoarseGrainedSchedulerBackend) = {
+  private def sparkContextInitialized(sc: SparkContext) = {
     sparkContextRef.synchronized {
       sparkContextRef.compareAndSet(null, sc)
       sparkContextRef.notifyAll()
     }
-    this.backend = backend
-    if (null != allocator) allocator.setScheduler(backend)
   }
 
   private def sparkContextStopped(sc: SparkContext) = {
     sparkContextRef.compareAndSet(sc, null)
   }
 
-  private def registerAM(_rpcEnv: RpcEnv, uiAddress: String, securityMgr: SecurityManager) = {
+  private def registerAM(
+      _rpcEnv: RpcEnv,
+      driverRef: RpcEndpointRef,
+      uiAddress: String,
+      securityMgr: SecurityManager) = {
     val sc = sparkContextRef.get()
 
     val appId = client.getAttemptId().getApplicationId().toString()
@@ -250,13 +250,13 @@ private[spark] class ApplicationMaster(
         RpcAddress(_sparkConf.get("spark.driver.host"), _sparkConf.get("spark.driver.port").toInt),
         CoarseGrainedSchedulerBackend.ENDPOINT_NAME)
     allocator = client.register(driverUrl,
+      driverRef,
       yarnConf,
       _sparkConf,
       if (sc != null) sc.preferredNodeLocationData else Map(),
       uiAddress,
       historyAddress,
       securityMgr)
-    if (null != backend) allocator.setScheduler(backend)
 
     allocator.allocateResources()
     reporterThread = launchReporterThread()
@@ -267,17 +267,20 @@ private[spark] class ApplicationMaster(
    *
    * In cluster mode, the AM and the driver belong to same process
    * so the AMEndpoint need not monitor lifecycle of the driver.
+   *
+   * @return A reference to the driver's RPC endpoint.
    */
   private def runAMEndpoint(
       host: String,
       port: String,
-      isClusterMode: Boolean): Unit = {
+      isClusterMode: Boolean): RpcEndpointRef = {
     val driverEndpoint = rpcEnv.setupEndpointRef(
       SparkEnv.driverActorSystemName,
       RpcAddress(host, port.toInt),
       YarnSchedulerBackend.ENDPOINT_NAME)
     amEndpoint =
       rpcEnv.setupEndpoint("YarnAM", new AMEndpoint(rpcEnv, driverEndpoint, isClusterMode))
+    driverEndpoint
   }
 
   private def runDriver(securityMgr: SecurityManager): Unit = {
@@ -295,21 +298,21 @@ private[spark] class ApplicationMaster(
         "Timed out waiting for SparkContext.")
     } else {
       rpcEnv = sc.env.rpcEnv
-      runAMEndpoint(
+      val driverRef = runAMEndpoint(
         sc.getConf.get("spark.driver.host"),
         sc.getConf.get("spark.driver.port"),
         isClusterMode = true)
-      registerAM(rpcEnv, sc.ui.map(_.appUIAddress).getOrElse(""), securityMgr)
+      registerAM(rpcEnv, driverRef, sc.ui.map(_.appUIAddress).getOrElse(""), securityMgr)
       userClassThread.join()
     }
   }
 
   private def runExecutorLauncher(securityMgr: SecurityManager): Unit = {
     val port = sparkConf.getInt("spark.yarn.am.port", 0)
     rpcEnv = RpcEnv.create("sparkYarnAM", Utils.localHostName, port, sparkConf, securityMgr)
-    waitForSparkDriver()
+    val driverRef = waitForSparkDriver()
     addAmIpFilter()
-    registerAM(rpcEnv, sparkConf.get("spark.driver.appUIAddress", ""), securityMgr)
+    registerAM(rpcEnv, driverRef, sparkConf.get("spark.driver.appUIAddress", ""), securityMgr)
 
     // In client mode the actor will stop the reporter thread.
     reporterThread.join()
@@ -433,7 +436,7 @@ private[spark] class ApplicationMaster(
     }
   }
 
-  private def waitForSparkDriver(): Unit = {
+  private def waitForSparkDriver(): RpcEndpointRef = {
     logInfo("Waiting for Spark driver to be reachable.")
     var driverUp = false
     val hostport = args.userArgs(0)
@@ -617,9 +620,8 @@ object ApplicationMaster extends Logging {
     }
   }
 
-  private[spark] def sparkContextInitialized(sc: SparkContext,
-      backend: CoarseGrainedSchedulerBackend): Unit = {
-    master.sparkContextInitialized(sc, backend)
+  private[spark] def sparkContextInitialized(sc: SparkContext): Unit = {
+    master.sparkContextInitialized(sc)
   }
 
   private[spark] def sparkContextStopped(sc: SparkContext): Boolean = {

yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala

@@ -36,7 +36,9 @@ 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._
 
 /**
  * YarnAllocator is charged with requesting containers from the YARN ResourceManager and deciding
@@ -53,6 +55,7 @@ import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
  */
 private[yarn] class YarnAllocator(
     driverUrl: String,
+    driverRef: RpcEndpointRef,
     conf: Configuration,
     sparkConf: SparkConf,
     amClient: AMRMClient[ContainerRequest],
@@ -90,7 +93,6 @@ private[yarn] class YarnAllocator(
   private[yarn] val executorIdToContainer = new HashMap[String, Container]
 
   private var numUnexpectedContainerRelease = 0L
-  private var backend: CoarseGrainedSchedulerBackend = _
   private val containerIdToExecutorId = new HashMap[ContainerId, String]
 
   // Executor memory in MB.
@@ -390,13 +392,8 @@ private[yarn] class YarnAllocator(
   private[yarn] def processCompletedContainers(completedContainers: Seq[ContainerStatus]): Unit = {
     for (completedContainer <- completedContainers) {
       val containerId = completedContainer.getContainerId
-
-      var needNotify = false
-      if (releasedContainers.contains(containerId)) {
-        // Already marked the container for release, so remove it from
-        // `releasedContainers`.
-        releasedContainers.remove(containerId)
-      } else {
+      val alreadyReleased = releasedContainers.remove(containerId)
+      if (!alreadyReleased) {
         // Decrement the number of executors running. The next iteration of
         // the ApplicationMaster's reporting thread will take care of allocating.
         numExecutorsRunning -= 1
@@ -423,7 +420,6 @@ private[yarn] class YarnAllocator(
             ". Diagnostics: " + completedContainer.getDiagnostics)
           numExecutorsFailed += 1
         }
-        needNotify = true
       }
 
       if (allocatedContainerToHostMap.containsKey(containerId)) {
@@ -440,13 +436,16 @@ private[yarn] class YarnAllocator(
         allocatedContainerToHostMap.remove(containerId)
       }
 
-      val executorIdOpt = containerIdToExecutorId.remove(containerId)
-      if (executorIdOpt.isDefined) executorIdToContainer.remove(executorIdOpt.get)
+      containerIdToExecutorId.remove(containerId).foreach { eid =>
+        executorIdToContainer.remove(eid)
 
-      if (needNotify && executorIdOpt.isDefined) {
-        // The executor could have gone away (like no route to host, node failure, etc)
-        // Notify backend about the failure of the executor
-        notifyBackend(executorIdOpt.get, containerId)
+        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)"))
+        }
       }
     }
   }
@@ -456,19 +455,8 @@ private[yarn] class YarnAllocator(
     amClient.releaseAssignedContainer(container.getId())
   }
 
-  private[yarn] def notifyBackend(executorId: String, containerId: ContainerId): Unit = {
-    numUnexpectedContainerRelease += 1
-    if (null != backend) {
-      backend.removeExecutor(executorId,
-        "Yarn deallocated the executor (" + executorId + ") container " + containerId)
-    }
-  }
-
   private[yarn] def getNumUnexpectedContainerRelease = numUnexpectedContainerRelease
 
-  private[yarn] def setScheduler(backend: CoarseGrainedSchedulerBackend): Unit = synchronized {
-    this.backend = backend
-  }
 }
 
 private object YarnAllocator {

yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnRMClient.scala

@@ -33,6 +33,7 @@ import org.apache.hadoop.yarn.util.ConverterUtils
 import org.apache.hadoop.yarn.webapp.util.WebAppUtils
 
 import org.apache.spark.{Logging, SecurityManager, SparkConf}
+import org.apache.spark.rpc.RpcEndpointRef
 import org.apache.spark.scheduler.SplitInfo
 import org.apache.spark.util.Utils
 
@@ -56,6 +57,7 @@ private[spark] class YarnRMClient(args: ApplicationMasterArguments) extends Logg
    */
   def register(
       driverUrl: String,
+      driverRef: RpcEndpointRef,
       conf: YarnConfiguration,
       sparkConf: SparkConf,
       preferredNodeLocations: Map[String, Set[SplitInfo]],
@@ -73,7 +75,8 @@ private[spark] class YarnRMClient(args: ApplicationMasterArguments) extends Logg
       amClient.registerApplicationMaster(Utils.localHostName(), 0, uiAddress)
       registered = true
     }
-    new YarnAllocator(driverUrl, conf, sparkConf, amClient, getAttemptId(), args, securityMgr)
+    new YarnAllocator(driverUrl, driverRef, conf, sparkConf, amClient, getAttemptId(), args,
+      securityMgr)
   }
 
   /**

yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnClusterScheduler.scala

@@ -29,8 +29,7 @@ private[spark] class YarnClusterScheduler(sc: SparkContext) extends YarnSchedule
   logInfo("Created YarnClusterScheduler")
 
   override def postStartHook() {
-    ApplicationMaster.sparkContextInitialized(sc,
-      this.backend.asInstanceOf[CoarseGrainedSchedulerBackend])
+    ApplicationMaster.sparkContextInitialized(sc)
     super.postStartHook()
     logInfo("YarnClusterScheduler.postStartHook done")
   }

yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala

@@ -26,14 +26,16 @@ import org.apache.hadoop.yarn.api.records._
 import org.apache.hadoop.yarn.client.api.AMRMClient
 import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
 
+import org.scalatest.{BeforeAndAfterEach, Matchers}
+import org.mockito.Mockito._
+
 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.rpc.RpcEndpointRef
 import org.apache.spark.scheduler.SplitInfo
 
-import org.scalatest.{BeforeAndAfterEach, Matchers}
-
 class MockResolver extends DNSToSwitchMapping {
 
   override def resolve(names: JList[String]): JList[String] = {
@@ -91,6 +93,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
       "--class", "SomeClass")
     new YarnAllocator(
       "not used",
+      mock(classOf[RpcEndpointRef]),
       conf,
       sparkConf,
       rmClient,