feat(driver): track and report driver exit state to AppMaster

core/raydp-main/src/main/java/org/apache/spark/deploy/raydp/RayAppMasterUtils.java

@@ -43,6 +43,7 @@ public static ActorHandle<RayAppMaster> createAppMaster(
               .substring(SparkOnRayConfigs.SPARK_MASTER_ACTOR_RESOURCE_PREFIX.length() + 1);
       creator.setResource(resourceName, resource.getValue());
     }
+    creator.setMaxTaskRetries(3);
 
     return creator.remote();
   }
@@ -57,6 +58,16 @@ public static Map<String, String> getRestartedExecutors(
     return handle.task(RayAppMaster::getRestartedExecutors).remote().get();
   }
 
+  public static boolean finishApplication(
+      ActorHandle<RayAppMaster> handle,
+      String appId,
+      String stateName,
+      int exitCode,
+      String diagnostics) {
+    return handle.task(RayAppMaster::finishApplication, appId, stateName, exitCode, diagnostics)
+            .remote().get();
+  }
+
   public static void stopAppMaster(
       ActorHandle<RayAppMaster> handle) {
     handle.task(RayAppMaster::stop).remote().get();

core/raydp-main/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala

@@ -50,6 +50,7 @@ import org.apache.ivy.plugins.resolver.{ChainResolver, FileSystemResolver, IBibl
 
 import org.apache.spark._
 import org.apache.spark.api.r.RUtils
+import org.apache.spark.deploy.raydp.{DriverAppMasterReporter, DriverExitState}
 import org.apache.spark.deploy.rest._
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config._
@@ -1011,6 +1012,19 @@ object SparkSubmit extends CommandLineUtils with Logging {
 
   private val CLASS_NOT_FOUND_EXIT_STATUS = 101
 
+  private def finalizeDriverTermination(): Unit = {
+    val snapshot = DriverExitState.current()
+    if (DriverExitState.isTerminal(snapshot.state)) {
+      DriverAppMasterReporter.tryReportAndCleanup()
+    }
+  }
+
+  private def describeFailure(t: Throwable): String = {
+    val message = Option(t.getMessage).filter(_.nonEmpty)
+      .getOrElse("No additional diagnostics available.")
+    s"${t.getClass.getName}: $message"
+  }
+
   // Following constants are visible for testing.
   private[deploy] val YARN_CLUSTER_SUBMIT_CLASS =
     "org.apache.spark.deploy.yarn.YarnClusterApplication"
@@ -1020,6 +1034,19 @@ object SparkSubmit extends CommandLineUtils with Logging {
     "org.apache.spark.deploy.k8s.submit.KubernetesClientApplication"
 
   override def main(args: Array[String]): Unit = {
+    DriverExitState.reset()
+    DriverAppMasterReporter.reset()
+    val originalExitFn = exitFn
+    exitFn = (exitCode: Int) => {
+      if (exitCode == 0) {
+        DriverExitState.trySetFinished()
+      } else {
+        DriverExitState.trySetFailed(exitCode, s"SparkSubmit exited with status $exitCode")
+      }
+      finalizeDriverTermination()
+      originalExitFn(exitCode)
+    }
+
     val submit = new SparkSubmit() {
       self =>
 
@@ -1050,7 +1077,18 @@ object SparkSubmit extends CommandLineUtils with Logging {
 
     }
 
-    submit.doSubmit(args)
+    try {
+      submit.doSubmit(args)
+      DriverExitState.trySetFinished()
+      finalizeDriverTermination()
+    } catch {
+      case t: Throwable =>
+        DriverExitState.trySetFailed(1, describeFailure(t))
+        finalizeDriverTermination()
+        throw t
+    } finally {
+      exitFn = originalExitFn
+    }
   }
 
   /**

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/ApplicationInfo.scala

@@ -53,6 +53,8 @@ private[spark] class ApplicationInfo(
   var removedExecutors: ArrayBuffer[ExecutorDesc] = _
   var coresGranted: Int = _
   var endTime: Long = _
+  var exitCode: Int = _
+  var diagnostics: String = _
   private var nextExecutorId: Int = _
   // this only count those registered executors and minus removed executors
   private var registeredExecutors: Int = 0
@@ -65,6 +67,8 @@ private[spark] class ApplicationInfo(
     addressToExecutorId = new HashMap[RpcAddress, String]
     executorIdToHandler = new HashMap[String, ActorHandle[RayDPExecutor]]
     endTime = -1L
+    exitCode = 0
+    diagnostics = null
     nextExecutorId = 0
     removedExecutors = new ArrayBuffer[ExecutorDesc]
   }
@@ -165,9 +169,21 @@ private[spark] class ApplicationInfo(
 
   def resetRetryCount(): Unit = _retryCount = 0
 
+  def finish(endState: ApplicationState.Value, endExitCode: Int, endDiagnostics: String): Boolean =
+    synchronized {
+      if (isFinished) {
+        false
+      } else {
+        state = endState
+        exitCode = endExitCode
+        diagnostics = endDiagnostics
+        endTime = System.currentTimeMillis()
+        true
+      }
+    }
+
   def markFinished(endState: ApplicationState.Value): Unit = {
-    state = endState
-    endTime = System.currentTimeMillis()
+    finish(endState, 0, null)
   }
 
   def isFinished: Boolean = {

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/DriverAppMasterReporter.scala

@@ -0,0 +1,108 @@
+/*
+ * 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.deploy.raydp
+
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.control.NonFatal
+
+import io.ray.api.ActorHandle
+
+import org.apache.spark.internal.Logging
+
+object DriverAppMasterReporter extends Logging {
+
+  private val reported = new AtomicBoolean(false)
+
+  private var appId: String = null
+  private var masterHandle: ActorHandle[RayAppMaster] = null
+
+  def reset(): Unit = synchronized {
+    reported.set(false)
+    appId = null
+    masterHandle = null
+  }
+
+  def bind(appId: String): Unit = synchronized {
+    if (!reported.get()) {
+      this.appId = appId
+    }
+  }
+
+  def bindMasterHandle(masterHandle: ActorHandle[RayAppMaster]): Unit = synchronized {
+    if (!reported.get()) {
+      this.masterHandle = masterHandle
+    }
+  }
+
+  def tryReportAndCleanup(): Boolean = {
+    val snapshot = DriverExitState.current()
+    if (!DriverExitState.isTerminal(snapshot.state)) {
+      logDebug(s"Skip AppMaster report because driver state is not terminal: ${snapshot.state}")
+      false
+    } else {
+      val binding = synchronized {
+        if (reported.get()) None
+        else Some((appId, masterHandle))
+      }
+      binding match {
+        case None => false
+        case Some((currentAppId, currentMasterHandle)) =>
+          if (currentAppId == null || currentMasterHandle == null) {
+            logWarning("Skip reporting terminal application state because AppMaster binding " +
+              "is incomplete.")
+            false
+          } else {
+            try {
+              val accepted = RayAppMasterUtils.finishApplication(
+                currentMasterHandle,
+                currentAppId,
+                snapshot.state.toString,
+                snapshot.exitCode,
+                snapshot.diagnostics)
+
+              if (!accepted) {
+                logWarning("Terminal application state report was not accepted by AppMaster; " +
+                  "keeping reporter state for a later retry.")
+                false
+              } else {
+                reported.set(true)
+                if (currentMasterHandle != null) {
+                  try {
+                    RayAppMasterUtils.stopAppMaster(currentMasterHandle)
+                  } catch {
+                    case NonFatal(e) =>
+                      logWarning("Failed to stop AppMaster during driver cleanup", e)
+                  }
+                }
+                synchronized {
+                  appId = null
+                  masterHandle = null
+                }
+                true
+              }
+            } catch {
+              case NonFatal(e) =>
+                logWarning("Failed to report terminal application state to AppMaster", e)
+                false
+            }
+          }
+      }
+    }
+  }
+}

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/DriverExitState.scala

@@ -0,0 +1,76 @@
+/*
+ * 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.deploy.raydp
+
+object DriverExitState {
+
+  case class Snapshot(state: ApplicationState.Value, exitCode: Int, diagnostics: String)
+
+  private var snapshot = Snapshot(ApplicationState.UNKNOWN, 0, null)
+
+  def reset(): Unit = synchronized {
+    snapshot = Snapshot(ApplicationState.UNKNOWN, 0, null)
+  }
+
+  def current(): Snapshot = synchronized {
+    snapshot
+  }
+
+  def isTerminal(state: ApplicationState.Value): Boolean = {
+    state == ApplicationState.FINISHED ||
+      state == ApplicationState.FAILED ||
+      state == ApplicationState.KILLED
+  }
+
+  def trySetFinished(): Boolean = synchronized {
+    trySet(ApplicationState.FINISHED, 0, null)
+  }
+
+  def trySetFailed(exitCode: Int, diagnostics: String): Boolean = synchronized {
+    trySet(ApplicationState.FAILED, normalizedFailureCode(exitCode), diagnostics)
+  }
+
+  def trySetKilled(exitCode: Int, diagnostics: String): Boolean = synchronized {
+    val normalizedExitCode = if (exitCode == 0) {
+      143
+    } else {
+      exitCode
+    }
+    trySet(ApplicationState.KILLED, normalizedExitCode, diagnostics)
+  }
+
+  private def trySet(
+      state: ApplicationState.Value,
+      exitCode: Int,
+      diagnostics: String): Boolean = {
+    if (isTerminal(snapshot.state)) {
+      false
+    } else {
+      snapshot = Snapshot(state, exitCode, diagnostics)
+      true
+    }
+  }
+
+  private def normalizedFailureCode(exitCode: Int): Int = {
+    if (exitCode == 0) {
+      1
+    } else {
+      exitCode
+    }
+  }
+}

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/Messages.scala

@@ -26,7 +26,11 @@ case class RegisterApplication(appDescription: ApplicationDescription, driver: R
 
 case class RegisteredApplication(appId: String, master: RpcEndpointRef) extends RayDPDeployMessage
 
-case class UnregisterApplication(appId: String) extends RayDPDeployMessage
+case class FinishApplication(
+    appId: String,
+    state: ApplicationState.Value,
+    exitCode: Int,
+    diagnostics: String) extends RayDPDeployMessage
 
 case class RegisterExecutor(executorId: String, nodeIp: String) extends RayDPDeployMessage
 

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/RayAppMaster.scala

@@ -83,6 +83,15 @@ class RayAppMaster(host: String,
 
   def getRestartedExecutors(): java.util.Map[String, String] = restartedExecutors.asJava
 
+  def finishApplication(
+      appId: String,
+      stateName: String,
+      exitCode: Int,
+      diagnostics: String): Boolean = {
+    endpoint.askSync[Boolean](
+      FinishApplication(appId, ApplicationState.withName(stateName), exitCode, diagnostics))
+  }
+
   /**
    * This is used to represent the Spark on Ray cluster URL.
    */
@@ -141,13 +150,13 @@ class RayAppMaster(host: String,
         logInfo("Registered app " + appDescription.name + " with ID " + app.id)
         driver.send(RegisteredApplication(app.id, self))
         schedule()
-
-      case UnregisterApplication(appId) =>
-        assert(appInfo != null && appInfo.id == appId)
-        appInfo.markFinished(ApplicationState.FINISHED)
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
+      case FinishApplication(appId, state, exitCode, diagnostics) =>
+        assert(appInfo != null && appInfo.id == appId)
+        context.reply(appInfo.finish(state, exitCode, diagnostics))
+
       case RegisterExecutor(executorId, executorIp) =>
         val success = appInfo.registerExecutor(executorId)
         if (success) {