Exited spark-submit process should not block batch submit queue

apache · Jan 30, 2024 · 148ebf8 · 148ebf8
1 parent 3f993f4
commit 148ebf8
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/ProcBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/ProcBuilder.scala
@@ -24,7 +24,6 @@ import java.nio.file.{Files, Path, Paths}
 
 import scala.collection.JavaConverters._
 
-import com.google.common.annotations.VisibleForTesting
 import com.google.common.collect.EvictingQueue
 import org.apache.commons.lang3.StringUtils.containsIgnoreCase
 
@@ -166,8 +165,7 @@ trait ProcBuilder {
   // Visible for test
   @volatile private[kyuubi] var logCaptureThreadReleased: Boolean = true
   private var logCaptureThread: Thread = _
-  private var process: Process = _
-  @VisibleForTesting
+  @volatile private[kyuubi] var process: Process = _
   @volatile private[kyuubi] var processLaunched: Boolean = _
 
   private[kyuubi] lazy val engineLog: File = ProcBuilder.synchronized {

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala
@@ -100,6 +100,9 @@ class BatchJobSubmission(
       getOperationLog)
   }
 
+  def startupProcessAlive: Boolean =
+    builder.processLaunched && Option(builder.process).exists(_.isAlive)
+
   override def currentApplicationInfo(): Option[ApplicationInfo] = {
     if (isTerminal(state) && _applicationInfo.map(_.state).exists(ApplicationState.isTerminated)) {
       return _applicationInfo

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiBatchService.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiBatchService.scala
@@ -94,8 +94,13 @@ class KyuubiBatchService(
                   metadata.appState match {
                     // app that is not submitted to resource manager
                     case None | Some(ApplicationState.NOT_FOUND) => false
-                    // app that is pending in resource manager
-                    case Some(ApplicationState.PENDING) => false
+                    // app that is pending in resource manager while the local startup
+                    // process is alive. For example, in Spark YARN cluster mode, if set
+                    // spark.yarn.submit.waitAppCompletion=false, the local spark-submit
+                    // process exits immediately once Application goes ACCEPTED status,
+                    // even no resource could be allocated for the AM container.
+                    case Some(ApplicationState.PENDING) if batchSession.startupProcessAlive =>
+                      false
                     // not sure, added for safe
                     case Some(ApplicationState.UNKNOWN) => false
                     case _ => true

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiBatchSession.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiBatchSession.scala
@@ -111,6 +111,8 @@ class KyuubiBatchSession(
       batchArgs,
       metadata)
 
+  def startupProcessAlive: Boolean = batchJobSubmissionOp.startupProcessAlive
+
   private def waitMetadataRequestsRetryCompletion(): Unit = {
     val batchId = batchJobSubmissionOp.batchId
     sessionManager.getMetadataRequestsRetryRef(batchId).foreach {