[SPARK-52675][ML][CONNECT] Interrupt hanging ML handlers in tests

### What changes were proposed in this pull request?

Interrupt hanging ML handlers in tests:

Recently some ML connect tests hangs randomly.
We need to add code to interrupt hanging handler threads and print the stack trace, for debuggability.

### Why are the changes needed?

Recently some ML connect tests hangs randomly.
We need to add code to interrupt hanging handler threads and print the stack trace, for debuggability.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Manually.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes #51364 from WeichenXu123/ml-connect-hang-debugger.

Authored-by: Weichen Xu <weichen.xu@databricks.com>
Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

Author: WeichenXu123

python/pyspark/testing/connectutils.py

@@ -166,6 +166,9 @@ def master(cls):
 
     @classmethod
     def setUpClass(cls):
+        # This environment variable is for interrupting hanging ML-handler and making the
+        # tests fail fast.
+        os.environ["SPARK_CONNECT_ML_HANDLER_INTERRUPTION_TIMEOUT_MINUTES"] = "5"
         cls.spark = (
             PySparkSession.builder.config(conf=cls.conf())
             .appName(cls.__name__)

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala

@@ -17,11 +17,14 @@
 
 package org.apache.spark.sql.connect.ml
 
+import java.io.{PrintWriter, StringWriter}
 import java.lang.ThreadLocal
+import java.util.concurrent.ConcurrentHashMap
 
 import scala.collection.mutable
 import scala.jdk.CollectionConverters.CollectionHasAsScala
 
+import org.apache.spark.SparkException
 import org.apache.spark.connect.proto
 import org.apache.spark.internal.Logging
 import org.apache.spark.ml.{Estimator, EstimatorUtils, Model, Transformer}
@@ -121,6 +124,9 @@ private[connect] object MLHandler extends Logging {
     override def initialValue: SessionHolder = null
   }
 
+  // A map of thread-id -> handler execution start time (UNIX timestamp)
+  val handlerExecutionStartTimeMap = new ConcurrentHashMap[Long, Long]()
+
   private val allowlistedMLClasses = {
     val transformerClasses = MLUtils.loadOperators(classOf[Transformer])
     val estimatorClasses = MLUtils.loadOperators(classOf[Estimator[_]])
@@ -150,7 +156,43 @@ private[connect] object MLHandler extends Logging {
     }
   }
 
-  def handleMlCommand(
+  def startHangingHandlerReaper(): Unit = {
+    val handlerInterruptionTimeoutMinutes = {
+      try {
+        val envValue = System.getenv("SPARK_CONNECT_ML_HANDLER_INTERRUPTION_TIMEOUT_MINUTES")
+        if (envValue != null) {
+          envValue.toInt
+        } else { 0 }
+      } catch {
+        case _: Exception => 0
+      }
+    }
+
+    if (handlerInterruptionTimeoutMinutes > 0) {
+      val handlerInterruptionTimeoutMillis = handlerInterruptionTimeoutMinutes * 60 * 1000
+      val thread = new Thread(() => {
+        while (true) {
+          handlerExecutionStartTimeMap.forEach { (threadId, startTime) =>
+            val execTime = System.currentTimeMillis() - startTime
+            if (execTime > handlerInterruptionTimeoutMillis) {
+              for (t <- Thread.getAllStackTraces().keySet().asScala) {
+                if (t.getId() == threadId) {
+                  t.interrupt()
+                }
+              }
+            }
+          }
+          Thread.sleep(60 * 1000)
+        }
+      })
+      thread.setDaemon(true)
+      thread.start()
+    }
+  }
+
+  startHangingHandlerReaper()
+
+  def _handleMlCommand(
       sessionHolder: SessionHolder,
       mlCommand: proto.MlCommand): proto.MlCommandResult = {
 
@@ -410,6 +452,39 @@ private[connect] object MLHandler extends Logging {
     }
   }
 
+  def wrapHandler(
+      originHandler: () => Any,
+      reqProto: com.google.protobuf.GeneratedMessage): Any = {
+    val threadId = Thread.currentThread().getId
+    val startTime = System.currentTimeMillis()
+    handlerExecutionStartTimeMap.put(threadId, startTime)
+    try {
+      originHandler()
+    } catch {
+      case e: InterruptedException =>
+        val stackTrace = {
+          val sw = new StringWriter()
+          val pw = new PrintWriter(sw)
+          e.printStackTrace(pw)
+          sw.toString
+        }
+        val execTime = (System.currentTimeMillis() - startTime) / (60 * 1000)
+        throw SparkException.internalError(
+          s"The Spark Connect ML handler thread is interrupted after executing for " +
+            s"$execTime minutes.\nThe request proto message is:\n${reqProto.toString}\n, " +
+            s"the current stack trace is:\n$stackTrace\n")
+    } finally {
+      handlerExecutionStartTimeMap.remove(threadId)
+    }
+  }
+
+  def handleMlCommand(
+      sessionHolder: SessionHolder,
+      mlCommand: proto.MlCommand): proto.MlCommandResult = {
+    wrapHandler(() => _handleMlCommand(sessionHolder, mlCommand), mlCommand)
+      .asInstanceOf[proto.MlCommandResult]
+  }
+
   private def createModelSummary(
       sessionHolder: SessionHolder,
       createSummaryCmd: proto.MlCommand.CreateSummary): proto.MlCommandResult =
@@ -431,7 +506,9 @@ private[connect] object MLHandler extends Logging {
         .build()
     }
 
-  def transformMLRelation(relation: proto.MlRelation, sessionHolder: SessionHolder): DataFrame = {
+  def _transformMLRelation(
+      relation: proto.MlRelation,
+      sessionHolder: SessionHolder): DataFrame = {
     relation.getMlTypeCase match {
       // Ml transform
       case proto.MlRelation.MlTypeCase.TRANSFORM =>
@@ -487,4 +564,9 @@ private[connect] object MLHandler extends Logging {
       case other => throw MlUnsupportedException(s"$other not supported")
     }
   }
+
+  def transformMLRelation(relation: proto.MlRelation, sessionHolder: SessionHolder): DataFrame = {
+    wrapHandler(() => _transformMLRelation(relation, sessionHolder), relation)
+      .asInstanceOf[DataFrame]
+  }
 }