[SPARK-39601][YARN] AllocationFailure should not be treated as exitCausedByApp when driver is shutting down

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala

@@ -815,6 +815,7 @@ private[spark] class ApplicationMaster(
       case Shutdown(code) =>
         exitCode = code
         shutdown = true
+        allocator.setShutdown(true)
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala

@@ -199,6 +199,8 @@ private[yarn] class YarnAllocator(
     }
   }
 
+  @volatile private var shutdown = false
+
   // The default profile is always present so we need to initialize the datastructures keyed by
   // ResourceProfile id to ensure its present if things start running before a request for
   // executors could add it. This approach is easier then going and special casing everywhere.
@@ -215,6 +217,8 @@ private[yarn] class YarnAllocator(
 
   initDefaultProfile()
 
+  def setShutdown(shutdown: Boolean): Unit = this.shutdown = shutdown
+
   def getNumExecutorsRunning: Int = synchronized {
     runningExecutorsPerResourceProfileId.values.map(_.size).sum
   }
@@ -835,6 +839,8 @@ private[yarn] class YarnAllocator(
         // now I think its ok as none of the containers are expected to exit.
         val exitStatus = completedContainer.getExitStatus
         val (exitCausedByApp, containerExitReason) = exitStatus match {
+          case _ if shutdown =>
+            (false, s"Executor for container $containerId exited after Application shutdown.")
           case ContainerExitStatus.SUCCESS =>
             (false, s"Executor for container $containerId exited because of a YARN event (e.g., " +
               "preemption) and not because of an error in the running job.")

resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnClientSchedulerBackend.scala

@@ -162,7 +162,7 @@ private[spark] class YarnClientSchedulerBackend(
    */
   override def stop(exitCode: Int): Unit = {
     assert(client != null, "Attempted to stop this scheduler before starting it!")
-    yarnSchedulerEndpoint.handleClientModeDriverStop(exitCode)
+    yarnSchedulerEndpoint.signalDriverStop(exitCode)
     if (monitorThread != null) {
       monitorThread.stopMonitor()
     }

resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnClusterSchedulerBackend.scala

@@ -35,6 +35,10 @@ private[spark] class YarnClusterSchedulerBackend(
     startBindings()
   }
 
+  override def stop(exitCode: Int): Unit = {
+    yarnSchedulerEndpoint.signalDriverStop(exitCode)
+  }
+
   override def getDriverLogUrls: Option[Map[String, String]] = {
     YarnContainerInfoHelper.getLogUrls(sc.hadoopConfiguration, container = None)
   }

resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala

@@ -319,7 +319,7 @@ private[spark] abstract class YarnSchedulerBackend(
       removeExecutorMessage.foreach { message => driverEndpoint.send(message) }
     }
 
-    private[cluster] def handleClientModeDriverStop(exitCode: Int): Unit = {
+    private[cluster] def signalDriverStop(exitCode: Int): Unit = {
       amEndpoint match {
         case Some(am) =>
           am.send(Shutdown(exitCode))

resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala

@@ -693,6 +693,28 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers {
       .updateBlacklist(hosts.slice(10, 11).asJava, Collections.emptyList())
   }
 
+  test("SPARK-39601 YarnAllocator should not count executor failure after shutdown") {
+    val (handler, _) = createAllocator()
+    handler.updateResourceRequests()
+    handler.getNumExecutorsFailed should be(0)
+
+    val failedBeforeShutdown = createContainer("host1")
+    val failedAfterShutdown = createContainer("host2")
+    handler.handleAllocatedContainers(Seq(failedBeforeShutdown, failedAfterShutdown))
+
+    val failedBeforeShutdownStatus = ContainerStatus.newInstance(
+      failedBeforeShutdown.getId, ContainerState.COMPLETE, "Failed", -1)
+    val failedAfterShutdownStatus = ContainerStatus.newInstance(
+      failedAfterShutdown.getId, ContainerState.COMPLETE, "Failed", -1)
+
+    handler.processCompletedContainers(Seq(failedBeforeShutdownStatus))
+    handler.getNumExecutorsFailed should be(1)
+
+    handler.setShutdown(true)
+    handler.processCompletedContainers(Seq(failedAfterShutdownStatus))
+    handler.getNumExecutorsFailed should be(1)
+  }
+
   test("SPARK-28577#YarnAllocator.resource.memory should include offHeapSize " +
     "when offHeapEnabled is true.") {
     val originalOffHeapEnabled = sparkConf.get(MEMORY_OFFHEAP_ENABLED)