add

Author: turboFei

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/ApplicationOperation.scala

@@ -91,10 +91,17 @@ object ApplicationState extends Enumeration {
   type ApplicationState = Value
   val PENDING, RUNNING, FINISHED, KILLED, FAILED, ZOMBIE, NOT_FOUND, UNKNOWN = Value
 
-  def isFailed(state: ApplicationState): Boolean = state match {
+  def isFailed(
+      state: ApplicationState,
+      appOperation: Option[ApplicationOperation]): Boolean = state match {
     case FAILED => true
     case KILLED => true
-    case NOT_FOUND => true
+    case NOT_FOUND => appOperation match {
+        // For YARN and Kubernetes operations, if the application is not found, treat it as failed
+        // to prevent mistakenly set unsuccessful applications to the finished state.
+        case Some(_: YarnApplicationOperation) | Some(_: KubernetesApplicationOperation) => true
+        case _ => false
+      }
     case _ => false
   }
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala

@@ -136,7 +136,7 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
           val shouldDelete = cleanupDriverPodStrategy match {
             case NONE => false
             case ALL => true
-            case COMPLETED => !ApplicationState.isFailed(notification.getValue)
+            case COMPLETED => !ApplicationState.isFailed(notification.getValue, Some(this))
           }
           if (shouldDelete) {
             deletePod(kubernetesInfo, removed.podName.orNull, appLabel)
@@ -573,8 +573,9 @@ object KubernetesApplicationOperation extends Logging {
       case Some(containerAppState) => containerAppState
       case None => podAppState
     }
-    val applicationError =
-      if (ApplicationState.isFailed(applicationState)) {
+    val applicationError = {
+      // here the applicationState could not been NOT_FOUND, safe to use None ApplicationOperation
+      if (ApplicationState.isFailed(applicationState, appOperation = None)) {
         val errorMap = containerStatusToBuildAppState.map { cs =>
           Map(
             "Pod" -> podName,
@@ -588,6 +589,7 @@ object KubernetesApplicationOperation extends Logging {
       } else {
         None
       }
+    }
     applicationState -> applicationError
   }
 

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KyuubiApplicationManager.scala

@@ -103,6 +103,11 @@ class KyuubiApplicationManager(metadataManager: Option[MetadataManager])
     operations.find(_.isInstanceOf[KubernetesApplicationOperation])
       .map(_.asInstanceOf[KubernetesApplicationOperation])
   }
+
+  private[kyuubi] def getApplicationOperation(appMgrInfo: ApplicationManagerInfo)
+      : Option[ApplicationOperation] = {
+    operations.find(_.isSupported(appMgrInfo))
+  }
 }
 
 object KyuubiApplicationManager {

kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala

@@ -25,7 +25,7 @@ import com.google.common.annotations.VisibleForTesting
 
 import org.apache.kyuubi.{KyuubiException, KyuubiSQLException, Utils}
 import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.engine.{ApplicationInfo, ApplicationState, KillResponse, ProcBuilder}
+import org.apache.kyuubi.engine.{ApplicationInfo, ApplicationOperation, ApplicationState, KillResponse, ProcBuilder}
 import org.apache.kyuubi.engine.spark.SparkBatchProcessBuilder
 import org.apache.kyuubi.metrics.MetricsConstants.OPERATION_OPEN
 import org.apache.kyuubi.metrics.MetricsSystem
@@ -99,6 +99,8 @@ class BatchJobSubmission(
       getOperationLog)
   }
 
+  private lazy val appOperation = applicationManager.getApplicationOperation(builder.appMgrInfo())
+
   def startupProcessAlive: Boolean =
     builder.processLaunched && Option(builder.process).exists(_.isAlive)
 
@@ -275,7 +277,7 @@ class BatchJobSubmission(
     try {
       info(s"Submitting $batchType batch[$batchId] job:\n$builder")
       val process = builder.start
-      while (process.isAlive && !applicationFailed(_applicationInfo)) {
+      while (process.isAlive && !applicationFailed(_applicationInfo, appOperation)) {
         doUpdateApplicationInfoMetadataIfNeeded()
         process.waitFor(applicationCheckInterval, TimeUnit.MILLISECONDS)
       }
@@ -284,7 +286,7 @@ class BatchJobSubmission(
         doUpdateApplicationInfoMetadataIfNeeded()
       }
 
-      if (applicationFailed(_applicationInfo)) {
+      if (applicationFailed(_applicationInfo, appOperation)) {
         Utils.terminateProcess(process, applicationStartupDestroyTimeout)
         throw new KyuubiException(s"Batch job failed: ${_applicationInfo}")
       }
@@ -332,7 +334,7 @@ class BatchJobSubmission(
       info(s"The $batchType batch[$batchId] job: $appId not found, assume that it has finished.")
       return
     }
-    if (applicationFailed(_applicationInfo)) {
+    if (applicationFailed(_applicationInfo, appOperation)) {
       throw new KyuubiException(s"$batchType batch[$batchId] job failed: ${_applicationInfo}")
     }
     updateBatchMetadata()
@@ -341,7 +343,7 @@ class BatchJobSubmission(
       Thread.sleep(applicationCheckInterval)
       updateApplicationInfoMetadataIfNeeded()
     }
-    if (applicationFailed(_applicationInfo)) {
+    if (applicationFailed(_applicationInfo, appOperation)) {
       throw new KyuubiException(s"$batchType batch[$batchId] job failed: ${_applicationInfo}")
     }
   }
@@ -445,8 +447,12 @@ class BatchJobSubmission(
 }
 
 object BatchJobSubmission {
-  def applicationFailed(applicationStatus: Option[ApplicationInfo]): Boolean = {
-    applicationStatus.map(_.state).exists(ApplicationState.isFailed)
+  def applicationFailed(
+      applicationStatus: Option[ApplicationInfo],
+      appOperation: Option[ApplicationOperation]): Boolean = {
+    applicationStatus.map(_.state).exists { state =>
+      ApplicationState.isFailed(state, appOperation)
+    }
   }
 
   def applicationTerminated(applicationStatus: Option[ApplicationInfo]): Boolean = {

kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala

@@ -129,8 +129,9 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
       metadata: Metadata,
       batchAppStatus: Option[ApplicationInfo]): Batch = {
     batchAppStatus.map { appStatus =>
+      val appOp = sessionManager.applicationManager.getApplicationOperation(metadata.appMgrInfo)
       val currentBatchState =
-        if (BatchJobSubmission.applicationFailed(batchAppStatus)) {
+        if (BatchJobSubmission.applicationFailed(batchAppStatus, appOp)) {
           OperationState.ERROR.toString
         } else if (BatchJobSubmission.applicationTerminated(batchAppStatus)) {
           OperationState.FINISHED.toString