[SPARK-31486] [CORE] spark.submit.waitAppCompletion flag to control spark-submit exit in Standalone Cluster Mode

core/src/main/scala/org/apache/spark/deploy/Client.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.deploy
 
+import java.util.concurrent.TimeUnit
+
 import scala.collection.mutable.HashSet
 import scala.concurrent.ExecutionContext
 import scala.reflect.ClassTag
@@ -61,6 +63,11 @@ private class ClientEndpoint(
 
    private val lostMasters = new HashSet[RpcAddress]
    private var activeMasterEndpoint: RpcEndpointRef = null
+   private val waitAppCompletion = conf.getBoolean("spark.standalone.submit.waitAppCompletion",
+     false)
+   private val REPORT_DRIVER_STATUS_INTERVAL = 10000
+   private var submittedDriverID = ""
+
 
   private def getProperty(key: String, conf: SparkConf): Option[String] = {
     sys.props.get(key).orElse(conf.getOption(key))
@@ -105,6 +112,10 @@ private class ClientEndpoint(
         asyncSendToMasterAndForwardReply[SubmitDriverResponse](
           RequestSubmitDriver(driverDescription))
 
+        forwardMessageThread.scheduleAtFixedRate(() => Utils.tryLogNonFatalError {
+          MonitorDriverStatus()
+        }, 0, REPORT_DRIVER_STATUS_INTERVAL, TimeUnit.MILLISECONDS)
+
       case "kill" =>
         val driverId = driverArgs.driverId
         asyncSendToMasterAndForwardReply[KillDriverResponse](RequestKillDriver(driverId))
@@ -123,16 +134,24 @@ private class ClientEndpoint(
       }(forwardMessageExecutionContext)
     }
   }
+  private def MonitorDriverStatus(): Unit = {
+    if (submittedDriverID != "") {
+      asyncSendToMasterAndForwardReply[DriverStatusResponse](RequestDriverStatus(submittedDriverID))
+    }
+  }
 
-  /* Find out driver status then exit the JVM */
+  /**
+   * Find out driver status then exit the JVM. If the waitAppCompletion is set to true, monitors
+   * the application until it finishes, fails or is killed.
+   */
   def pollAndReportStatus(driverId: String): Unit = {
     // Since ClientEndpoint is the only RpcEndpoint in the process, blocking the event loop thread
     // is fine.
     logInfo("... waiting before polling master for driver state")
     Thread.sleep(5000)
     logInfo("... polling master for driver state")
     val statusResponse =
-      activeMasterEndpoint.askSync[DriverStatusResponse](RequestDriverStatus(driverId))
+        activeMasterEndpoint.askSync[DriverStatusResponse](RequestDriverStatus(driverId))
     if (statusResponse.found) {
       logInfo(s"State of $driverId is ${statusResponse.state.get}")
       // Worker node, if present
@@ -148,20 +167,27 @@ private class ClientEndpoint(
           e.printStackTrace()
           System.exit(-1)
         case _ =>
-          System.exit(0)
+          if (!waitAppCompletion) {
+            logInfo(s"spark-submit not configured to wait for completion, " +
+              s"exiting spark-submit JVM.")
+            System.exit(0)
+            } else {
+              logInfo(s"spark-submit is configured to wait for completion, " +
+                s"continue monitoring driver status.")
+            }
+        }
+      } else {
+        logError(s"ERROR: Cluster master did not recognize $driverId")
+        System.exit(-1)
       }
-    } else {
-      logError(s"ERROR: Cluster master did not recognize $driverId")
-      System.exit(-1)
     }
-  }
-
   override def receive: PartialFunction[Any, Unit] = {
 
     case SubmitDriverResponse(master, success, driverId, message) =>
       logInfo(message)
       if (success) {
         activeMasterEndpoint = master
+        submittedDriverID = driverId.get
         pollAndReportStatus(driverId.get)
       } else if (!Utils.responseFromBackup(message)) {
         System.exit(-1)
@@ -176,6 +202,22 @@ private class ClientEndpoint(
       } else if (!Utils.responseFromBackup(message)) {
         System.exit(-1)
       }
+
+    case DriverStatusResponse(found, state, _, _, _) =>
+      if (found) {
+        state.get match {
+          case DriverState.FINISHED | DriverState.FAILED |
+               DriverState.ERROR | DriverState.KILLED =>
+            logInfo(s"State of driver $submittedDriverID is ${state.get}, " +
+              s"exiting spark-submit JVM.")
+            System.exit(0)
+          case _ =>
+            logDebug(s"State of driver $submittedDriverID is ${state.get}, " +
+              s"continue monitoring driver status.")
+        }
+      } else {
+        System.exit(-1)
+      }
   }
 
   override def onDisconnected(remoteAddress: RpcAddress): Unit = {

docs/spark-standalone.md

@@ -374,6 +374,25 @@ To run an interactive Spark shell against the cluster, run the following command
 
 You can also pass an option `--total-executor-cores <numCores>` to control the number of cores that spark-shell uses on the cluster.
 
+# Client Properties
+
+Spark applications supports the following configuration properties specific to standalone mode: 
+
+<table class="table">
+  <tr><th style="width:21%">Property Name</th><th>Default Value</th><th>Meaning</th><th>Since Version</th></tr>
+  <tr>
+  <td><code>spark.standalone.submit.waitAppCompletion</code></td>
+  <td><code>false</code></td>
+  <td>
+  In standalone cluster mode, controls whether the client waits to exit until the application completes.
+  If set to <code>true</code>, the client process will stay alive polling the driver's status.
+  Otherwise, the client process will exit after submission.
+  </td>
+  <td>3.1.0</td>
+  </tr>
+</table>
+
+
 # Launching Spark Applications
 
 The [`spark-submit` script](submitting-applications.html) provides the most straightforward way to