[SPARK-38270][SQL] Spark SQL CLI's AM should keep same exit code with client side

### What changes were proposed in this pull request?
Current Spark SQL CLI alway use  shutdown hook to stop SparkSQLEnv
```
 // Clean up after we exit
    ShutdownHookManager.addShutdownHook { () => SparkSQLEnv.stop() }
```

but use process ret to close client side jvm
```
while (line != null) {
  if (!line.startsWith("--")) {
    if (prefix.nonEmpty) {
      prefix += '\n'
     }
    if (line.trim().endsWith(";") && !line.trim().endsWith("\\;")) {
      line = prefix + line
      ret = cli.processLine(line, true)
      prefix = ""
      currentPrompt = promptWithCurrentDB
    } else {
      prefix = prefix + line
      currentPrompt = continuedPromptWithDBSpaces
    }
  }
  line = reader.readLine(currentPrompt + "> ")
}

sessionState.close()

System.exit(ret)
}
```

```
    if (sessionState.execString != null) {
      exitCode = cli.processLine(sessionState.execString)
      System.exit(exitCode)
    }

    try {
      if (sessionState.fileName != null) {
        exitCode = cli.processFile(sessionState.fileName)
        System.exit(exitCode)
      }
    } catch {
      case e: FileNotFoundException =>
        logError(s"Could not open input file for reading. (${e.getMessage})")
        exitCode = 3
        System.exit(exitCode)
    }

```

This always cause client side exit code not consistent with AM.

IN this pr I prefer to pass the exit code to `SparkContext.stop()` method to pass a clear client side status to AM side in client mode.

In this pr, I add a new `stop` method in `SchedulerBackend`
```
def stop(exitCode: Int): Unit = stop()
```
So we don't need to implement it for all kinds of scheduler backend.
In this pr, we only handle the case me meet for `YarnClientSchedulerBackend`, then we can only implement `stop(exitCode: Int)` for this class, then for yarn client mode, it can work now.

I think this can benefit many similar case in future.

### Why are the changes needed?
Keep client side status consistent  with AM side

### Does this PR introduce _any_ user-facing change?
With this pr, client side status will be same with am side

### How was this patch tested?
MT

Closes apache#35594 from AngersZhuuuu/SPARK-38270.

Lead-authored-by: Angerszhuuuu <[email protected]>
Co-authored-by: AngersZhuuuu <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

Author: AngersZhuuuu

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -2068,7 +2068,20 @@ class SparkContext(config: SparkConf) extends Logging {
   /**
    * Shut down the SparkContext.
    */
-  def stop(): Unit = {
+  def stop(): Unit = stop(0)
+
+  /**
+   * Shut down the SparkContext with exit code that will passed to scheduler backend.
+   * In client mode, client side may call `SparkContext.stop()` to clean up but exit with
+   * code not equal to 0. This behavior cause resource scheduler such as `ApplicationMaster`
+   * exit with success status but client side exited with failed status. Spark can call
+   * this method to stop SparkContext and pass client side correct exit code to scheduler backend.
+   * Then scheduler backend should send the exit code to corresponding resource scheduler
+   * to keep consistent.
+   *
+   * @param exitCode Specified exit code that will passed to scheduler backend in client mode.
+   */
+  def stop(exitCode: Int): Unit = {
     if (LiveListenerBus.withinListenerThread.value) {
       throw new SparkException(s"Cannot stop SparkContext within listener bus thread.")
     }
@@ -2101,7 +2114,7 @@ class SparkContext(config: SparkConf) extends Logging {
     }
     if (_dagScheduler != null) {
       Utils.tryLogNonFatalError {
-        _dagScheduler.stop()
+        _dagScheduler.stop(exitCode)
       }
       _dagScheduler = null
     }

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -2889,7 +2889,7 @@ private[spark] class DAGScheduler(
     listenerBus.post(SparkListenerJobEnd(job.jobId, clock.getTimeMillis(), JobSucceeded))
   }
 
-  def stop(): Unit = {
+  def stop(exitCode: Int = 0): Unit = {
     Utils.tryLogNonFatalError {
       messageScheduler.shutdownNow()
     }
@@ -2900,7 +2900,7 @@ private[spark] class DAGScheduler(
       eventProcessLoop.stop()
     }
     Utils.tryLogNonFatalError {
-      taskScheduler.stop()
+      taskScheduler.stop(exitCode)
     }
   }
 

core/src/main/scala/org/apache/spark/scheduler/SchedulerBackend.scala

@@ -30,6 +30,7 @@ private[spark] trait SchedulerBackend {
 
   def start(): Unit
   def stop(): Unit
+  def stop(exitCode: Int): Unit = stop()
   /**
    * Update the current offers and schedule tasks
    */

core/src/main/scala/org/apache/spark/scheduler/TaskScheduler.scala

@@ -49,7 +49,7 @@ private[spark] trait TaskScheduler {
   def postStartHook(): Unit = { }
 
   // Disconnect from the cluster.
-  def stop(): Unit
+  def stop(exitCode: Int = 0): Unit
 
   // Submit a sequence of tasks to run.
   def submitTasks(taskSet: TaskSet): Unit

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -971,13 +971,13 @@ private[spark] class TaskSchedulerImpl(
     }
   }
 
-  override def stop(): Unit = {
+  override def stop(exitCode: Int = 0): Unit = {
     Utils.tryLogNonFatalError {
       speculationScheduler.shutdown()
     }
     if (backend != null) {
       Utils.tryLogNonFatalError {
-        backend.stop()
+        backend.stop(exitCode)
       }
     }
     if (taskResultGetter != null) {

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

@@ -151,7 +151,7 @@ private[spark] object CoarseGrainedClusterMessages {
   case class KillExecutors(executorIds: Seq[String]) extends CoarseGrainedClusterMessage
 
   // Used internally by executors to shut themselves down.
-  case object Shutdown extends CoarseGrainedClusterMessage
+  case class Shutdown(exitCode: Int = 0) extends CoarseGrainedClusterMessage
 
   // The message to check if `CoarseGrainedSchedulerBackend` thinks the executor is alive or not.
   case class IsExecutorAlive(executorId: String) extends CoarseGrainedClusterMessage

core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala

@@ -176,7 +176,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
     override def schedulingMode: SchedulingMode = SchedulingMode.FIFO
     override def rootPool: Pool = new Pool("", schedulingMode, 0, 0)
     override def start() = {}
-    override def stop() = {}
+    override def stop(exitCode: Int) = {}
     override def executorHeartbeatReceived(
         execId: String,
         accumUpdates: Array[(Long, Seq[AccumulatorV2[_, _]])],
@@ -846,7 +846,7 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti
       override def schedulingMode: SchedulingMode = SchedulingMode.FIFO
       override def rootPool: Pool = new Pool("", schedulingMode, 0, 0)
       override def start(): Unit = {}
-      override def stop(): Unit = {}
+      override def stop(exitCode: Int): Unit = {}
       override def submitTasks(taskSet: TaskSet): Unit = {
         taskSets += taskSet
       }

core/src/test/scala/org/apache/spark/scheduler/ExternalClusterManagerSuite.scala

@@ -80,7 +80,7 @@ private class DummyTaskScheduler extends TaskScheduler {
   override def schedulingMode: SchedulingMode = SchedulingMode.FIFO
   override def rootPool: Pool = new Pool("", schedulingMode, 0, 0)
   override def start(): Unit = {}
-  override def stop(): Unit = {}
+  override def stop(exitCode: Int): Unit = {}
   override def submitTasks(taskSet: TaskSet): Unit = {}
   override def cancelTasks(stageId: Int, interruptThread: Boolean): Unit = {}
   override def killTaskAttempt(

project/MimaExcludes.scala

@@ -100,7 +100,15 @@ object MimaExcludes {
 
     // [SPARK-37935][SQL] Eliminate separate error sub-classes fields
     ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.SparkException.this"),
-    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.AnalysisException.this")
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.sql.AnalysisException.this"),
+
+    // [SPARK-38270][SQL] Spark SQL CLI's AM should keep same exit code with client side
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.productPrefix"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.productArity"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.productElement"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.productIterator"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.canEqual"),
+    ProblemFilters.exclude[FinalMethodProblem]("org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages#Shutdown.toString")
   )
 
   // Defulat exclude rules

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

@@ -790,6 +790,8 @@ private[spark] class ApplicationMaster(
   private class AMEndpoint(override val rpcEnv: RpcEnv, driver: RpcEndpointRef)
     extends RpcEndpoint with Logging {
     @volatile private var shutdown = false
+    @volatile private var exitCode = 0
+
     private val clientModeTreatDisconnectAsFailed =
       sparkConf.get(AM_CLIENT_MODE_TREAT_DISCONNECT_AS_FAILED)
 
@@ -810,7 +812,9 @@ private[spark] class ApplicationMaster(
       case UpdateDelegationTokens(tokens) =>
         SparkHadoopUtil.get.addDelegationTokens(tokens, sparkConf)
 
-      case Shutdown => shutdown = true
+      case Shutdown(code) =>
+        exitCode = code
+        shutdown = true
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
@@ -854,8 +858,13 @@ private[spark] class ApplicationMaster(
       // This avoids potentially reporting incorrect exit codes if the driver fails
       if (!(isClusterMode || sparkConf.get(YARN_UNMANAGED_AM))) {
         if (shutdown || !clientModeTreatDisconnectAsFailed) {
-          logInfo(s"Driver terminated or disconnected! Shutting down. $remoteAddress")
-          finish(FinalApplicationStatus.SUCCEEDED, ApplicationMaster.EXIT_SUCCESS)
+          if (exitCode == 0) {
+            logInfo(s"Driver terminated or disconnected! Shutting down. $remoteAddress")
+            finish(FinalApplicationStatus.SUCCEEDED, ApplicationMaster.EXIT_SUCCESS)
+          } else {
+            logError(s"Driver terminated with exit code ${exitCode}! Shutting down. $remoteAddress")
+            finish(FinalApplicationStatus.FAILED, exitCode)
+          }
         } else {
           logError(s"Application Master lost connection with driver! Shutting down. $remoteAddress")
           finish(FinalApplicationStatus.FAILED, ApplicationMaster.EXIT_DISCONNECTED)