Implement feature.

JoshRosen · Jun 5, 2024 · d6108ce · d6108ce
1 parent 2c0b148
commit d6108ce
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
@@ -1021,11 +1021,19 @@ private[spark] class SparkSubmit extends Logging {
         e
     }
 
+    var exitCode: Int = 1
     try {
       app.start(childArgs.toArray, sparkConf)
+      exitCode = 0
     } catch {
       case t: Throwable =>
-        throw findCause(t)
+        val cause = findCause(t)
+        cause match {
+          case e: SparkUserAppException =>
+            exitCode = e.exitCode
+          case _ =>
+        }
+        throw cause
     } finally {
       if (args.master.startsWith("k8s") && !isShell(args.primaryResource) &&
           !isSqlShell(args.mainClass) && !isThriftServer(args.mainClass) &&
@@ -1036,6 +1044,12 @@ private[spark] class SparkSubmit extends Logging {
           case e: Throwable => logError("Failed to close SparkContext", e)
         }
       }
+      if (sparkConf.get(SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT)) {
+        logInfo(
+          log"Calling System.exit() with exit code ${MDC(LogKeys.EXIT_CODE, exitCode)} " +
+          log"because main ${MDC(LogKeys.CONFIG, SUBMIT_CALL_SYSTEM_EXIT_ON_MAIN_EXIT.key)}=true")
+        exitFn(exitCode)
+      }
     }
   }