fix(driver): force JVM exit to prevent hanging RayJobs

bin/raydp-submit

@@ -143,6 +143,12 @@ added_args+=("spark.driver.host=$RAY_DRIVER_NODE_IP")
 added_args+=("--conf")
 added_args+=("spark.driver.bindAddress=$RAY_DRIVER_NODE_IP")
 
+# JvmExitGuard: force JVM exit after Spark app ends if non-daemon threads keep it alive.
+# Set RAYDP_JVM_EXIT_TIMEOUT env var to override the default timeout (300s) or disable (0).
+if [ -n "$RAYDP_JVM_EXIT_TIMEOUT" ]; then
+  added_confs+=("-Draydp.jvm.exit.timeout=$RAYDP_JVM_EXIT_TIMEOUT")
+fi
+
 # Find the java binary
 if [ -n "${JAVA_HOME}" ]; then
   RUNNER="${JAVA_HOME}/bin/java"

core/raydp-main/src/main/java/org/apache/spark/deploy/raydp/ExternalShuffleServiceUtils.java

@@ -18,21 +18,31 @@
 package org.apache.spark.deploy.raydp;
 
 import java.util.List;
+import java.util.Optional;
 
 import io.ray.api.ActorHandle;
 import io.ray.api.Ray;
 
 public class ExternalShuffleServiceUtils {
+  private static String getShuffleServiceActorName(String node) {
+    return "raydp-shuffle-service-" + node.replace('.', '-');
+  }
+
   public static ActorHandle<RayExternalShuffleService> createShuffleService(
       String node, List<String> options) {
+    String actorName = getShuffleServiceActorName(node);
+    Optional<ActorHandle<RayExternalShuffleService>> existing = Ray.getActor(actorName);
+    if (existing.isPresent()) {
+      return existing.get();
+    }
+
     return Ray.actor(RayExternalShuffleService::new)
+              .setName(actorName)
               .setResource("node:" + node, 0.01)
-              .setJvmOptions(options).remote();
-  }
-
-  public static void startShuffleService(
-      ActorHandle<RayExternalShuffleService> handle) {
-    handle.task(RayExternalShuffleService::start).remote();
+              .setJvmOptions(options)
+              .setMaxRestarts(-1)
+              .setMaxTaskRetries(-1)
+              .remote();
   }
 
   public static void stopShuffleService(

core/raydp-main/src/main/java/org/apache/spark/deploy/raydp/RayAppMasterUtils.java

@@ -57,6 +57,16 @@ public static Map<String, String> getRestartedExecutors(
     return handle.task(RayAppMaster::getRestartedExecutors).remote().get();
   }
 
+  public static boolean finishApplication(
+      ActorHandle<RayAppMaster> handle,
+      String appId,
+      String stateName,
+      int exitCode,
+      String diagnostics) {
+    return handle.task(RayAppMaster::finishApplication, appId, stateName, exitCode, diagnostics)
+            .remote().get();
+  }
+
   public static void stopAppMaster(
       ActorHandle<RayAppMaster> handle) {
     handle.task(RayAppMaster::stop).remote().get();

core/raydp-main/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala

@@ -50,6 +50,7 @@ import org.apache.ivy.plugins.resolver.{ChainResolver, FileSystemResolver, IBibl
 
 import org.apache.spark._
 import org.apache.spark.api.r.RUtils
+import org.apache.spark.deploy.raydp.{DriverAppMasterReporter, DriverExitState, JvmExitGuard}
 import org.apache.spark.deploy.rest._
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config._
@@ -1011,6 +1012,20 @@ object SparkSubmit extends CommandLineUtils with Logging {
 
   private val CLASS_NOT_FOUND_EXIT_STATUS = 101
 
+  private def finalizeDriverTermination(): Unit = {
+    val snapshot = DriverExitState.current()
+    if (DriverExitState.isTerminal(snapshot.state)) {
+      DriverAppMasterReporter.tryReportAndCleanup()
+      JvmExitGuard.arm(snapshot.exitCode)
+    }
+  }
+
+  private def describeFailure(t: Throwable): String = {
+    val message = Option(t.getMessage).filter(_.nonEmpty)
+      .getOrElse("No additional diagnostics available.")
+    s"${t.getClass.getName}: $message"
+  }
+
   // Following constants are visible for testing.
   private[deploy] val YARN_CLUSTER_SUBMIT_CLASS =
     "org.apache.spark.deploy.yarn.YarnClusterApplication"
@@ -1020,6 +1035,19 @@ object SparkSubmit extends CommandLineUtils with Logging {
     "org.apache.spark.deploy.k8s.submit.KubernetesClientApplication"
 
   override def main(args: Array[String]): Unit = {
+    DriverExitState.reset()
+    DriverAppMasterReporter.reset()
+    val originalExitFn = exitFn
+    exitFn = (exitCode: Int) => {
+      if (exitCode == JvmExitGuard.EXIT_SUCCESS) {
+        DriverExitState.trySetFinished()
+      } else {
+        DriverExitState.trySetFailed(exitCode, s"SparkSubmit exited with status $exitCode")
+      }
+      finalizeDriverTermination()
+      originalExitFn(exitCode)
+    }
+
     val submit = new SparkSubmit() {
       self =>
 
@@ -1050,7 +1078,18 @@ object SparkSubmit extends CommandLineUtils with Logging {
 
     }
 
-    submit.doSubmit(args)
+    try {
+      submit.doSubmit(args)
+      DriverExitState.trySetFinished()
+      finalizeDriverTermination()
+    } catch {
+      case t: Throwable =>
+        DriverExitState.trySetFailed(JvmExitGuard.EXIT_APP_FAILED, describeFailure(t))
+        finalizeDriverTermination()
+        throw t
+    } finally {
+      exitFn = originalExitFn
+    }
   }
 
   /**

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/ApplicationInfo.scala

@@ -25,6 +25,7 @@ import io.ray.api.ActorHandle
 
 import org.apache.spark.executor.RayDPExecutor
 import org.apache.spark.internal.Logging
+import org.apache.spark.deploy.raydp.JvmExitGuard
 import org.apache.spark.raydp.RayExecutorUtils
 import org.apache.spark.resource.ResourceInformation
 import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef}
@@ -53,6 +54,8 @@ private[spark] class ApplicationInfo(
   var removedExecutors: ArrayBuffer[ExecutorDesc] = _
   var coresGranted: Int = _
   var endTime: Long = _
+  var exitCode: Int = _
+  var diagnostics: String = _
   private var nextExecutorId: Int = _
   // this only count those registered executors and minus removed executors
   private var registeredExecutors: Int = 0
@@ -65,6 +68,8 @@ private[spark] class ApplicationInfo(
     addressToExecutorId = new HashMap[RpcAddress, String]
     executorIdToHandler = new HashMap[String, ActorHandle[RayDPExecutor]]
     endTime = -1L
+    exitCode = JvmExitGuard.EXIT_SUCCESS
+    diagnostics = null
     nextExecutorId = 0
     removedExecutors = new ArrayBuffer[ExecutorDesc]
   }
@@ -165,9 +170,21 @@ private[spark] class ApplicationInfo(
 
   def resetRetryCount(): Unit = _retryCount = 0
 
+  def finish(endState: ApplicationState.Value, endExitCode: Int, endDiagnostics: String): Boolean =
+    synchronized {
+      if (isFinished) {
+        false
+      } else {
+        state = endState
+        exitCode = endExitCode
+        diagnostics = endDiagnostics
+        endTime = System.currentTimeMillis()
+        true
+      }
+    }
+
   def markFinished(endState: ApplicationState.Value): Unit = {
-    state = endState
-    endTime = System.currentTimeMillis()
+    finish(endState, JvmExitGuard.EXIT_SUCCESS, null)
   }
 
   def isFinished: Boolean = {

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/DriverAppMasterReporter.scala

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.deploy.raydp
+
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.control.NonFatal
+
+import io.ray.api.ActorHandle
+
+import org.apache.spark.internal.Logging
+
+object DriverAppMasterReporter extends Logging {
+
+  private val reported = new AtomicBoolean(false)
+
+  private var appId: String = null
+  private var masterHandle: ActorHandle[RayAppMaster] = null
+
+  def reset(): Unit = synchronized {
+    reported.set(false)
+    appId = null
+    masterHandle = null
+  }
+
+  def bind(appId: String): Unit = synchronized {
+    if (!reported.get()) {
+      this.appId = appId
+    }
+  }
+
+  def bindMasterHandle(masterHandle: ActorHandle[RayAppMaster]): Unit = synchronized {
+    if (!reported.get()) {
+      this.masterHandle = masterHandle
+    }
+  }
+
+  def tryReportAndCleanup(): Boolean = {
+    val snapshot = DriverExitState.current()
+    if (!DriverExitState.isTerminal(snapshot.state)) {
+      logDebug(s"Skip AppMaster report because driver state is not terminal: ${snapshot.state}")
+      false
+    } else {
+      val binding = synchronized {
+        if (!reported.compareAndSet(false, true)) None
+        else Some((appId, masterHandle))
+      }
+      binding match {
+        case None => false
+        case Some((currentAppId, currentMasterHandle)) =>
+          try {
+            if (currentAppId != null && currentMasterHandle != null) {
+              RayAppMasterUtils.finishApplication(
+                currentMasterHandle,
+                currentAppId,
+                snapshot.state.toString,
+                snapshot.exitCode,
+                snapshot.diagnostics)
+            } else {
+              logWarning("Skip reporting terminal application state because AppMaster binding " +
+                "is incomplete.")
+            }
+          } catch {
+            case NonFatal(e) =>
+              logWarning("Failed to report terminal application state to AppMaster", e)
+          } finally {
+            if (currentMasterHandle != null) {
+              try {
+                RayAppMasterUtils.stopAppMaster(currentMasterHandle)
+              } catch {
+                case NonFatal(e) =>
+                  logWarning("Failed to stop AppMaster during driver cleanup", e)
+              }
+            }
+            synchronized {
+              appId = null
+              masterHandle = null
+            }
+          }
+          true
+      }
+    }
+  }
+}

core/raydp-main/src/main/scala/org/apache/spark/deploy/raydp/DriverExitState.scala

@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.deploy.raydp
+
+object DriverExitState {
+
+  case class Snapshot(state: ApplicationState.Value, exitCode: Int, diagnostics: String)
+
+  private var snapshot = Snapshot(ApplicationState.UNKNOWN, JvmExitGuard.EXIT_SUCCESS, null)
+
+  def reset(): Unit = synchronized {
+    snapshot = Snapshot(ApplicationState.UNKNOWN, JvmExitGuard.EXIT_SUCCESS, null)
+  }
+
+  def current(): Snapshot = synchronized {
+    snapshot
+  }
+
+  def isTerminal(state: ApplicationState.Value): Boolean = {
+    state == ApplicationState.FINISHED ||
+      state == ApplicationState.FAILED ||
+      state == ApplicationState.KILLED
+  }
+
+  def trySetFinished(): Boolean = synchronized {
+    trySet(ApplicationState.FINISHED, JvmExitGuard.EXIT_SUCCESS, null)
+  }
+
+  def trySetFailed(exitCode: Int, diagnostics: String): Boolean = synchronized {
+    trySet(ApplicationState.FAILED, normalizedFailureCode(exitCode), diagnostics)
+  }
+
+  def trySetKilled(exitCode: Int, diagnostics: String): Boolean = synchronized {
+    val normalizedExitCode = if (exitCode == JvmExitGuard.EXIT_SUCCESS) {
+      JvmExitGuard.EXIT_KILLED
+    } else {
+      exitCode
+    }
+    trySet(ApplicationState.KILLED, normalizedExitCode, diagnostics)
+  }
+
+  private def trySet(
+      state: ApplicationState.Value,
+      exitCode: Int,
+      diagnostics: String): Boolean = {
+    if (isTerminal(snapshot.state)) {
+      false
+    } else {
+      snapshot = Snapshot(state, exitCode, diagnostics)
+      true
+    }
+  }
+
+  private def normalizedFailureCode(exitCode: Int): Int = {
+    if (exitCode == JvmExitGuard.EXIT_SUCCESS) {
+      JvmExitGuard.EXIT_APP_FAILED
+    } else {
+      exitCode
+    }
+  }
+}