sarutak opened a new pull request, #57041: URL: https://github.com/apache/spark/pull/57041
### What changes were proposed in this pull request? This PR fixes an issue that interrupt statuses set by `queryExecutionThread.interrupt()` in [QueryExecution.interruptAndAwaitExecutionThreadTermination()](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecution.scala#L503) are unexpectedly cleared. This can be the root cause of the flakiness of the test `python foreachBatch process: process terminates after query is stopped` in `SparkConnectSessionHolderSuite`. https://github.com/apache/spark/actions/runs/28604689132/job/84825842409 ``` [info] - python foreachBatch process: process terminates after query is stopped *** FAILED *** (1 minute, 46 seconds) [info] java.util.concurrent.TimeoutException: Stream Execution thread for stream foreachBatch_termination_test_q1_3866433486358 [id = e5635fe9-1633-4de9-a0d7-1f47a267e44a, runId = 5c01eacc-9822-4b9a-8390-e47267a2a6d0] failed to stop within 30000 milliseconds (specified by spark.sql.streaming.stopTimeout). See the cause on what was being executed in the streaming query thread. [info] at org.apache.spark.sql.execution.streaming.runtime.StreamExecution.interruptAndAwaitExecutionThreadTermination(StreamExecution.scala:510) [info] at org.apache.spark.sql.execution.streaming.runtime.MicroBatchExecution.stop(MicroBatchExecution.scala:475) [info] at org.apache.spark.sql.execution.streaming.runtime.StreamingQueryWrapper.stop(StreamingQueryWrapper.scala:61) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.runPythonForeachBatchTerminationTestBody(SparkConnectSessionHolderSuite.scala:387) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.$anonfun$new$16(SparkConnectSessionHolderSuite.scala:450) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.$anonfun$awaitTestBodyInNewThread$1(SparkConnectSessionHolderSuite.scala:281) [info] at java.base/java.lang.Thread.run(Thread.java:840) [info] Cause: org.apache.spark.SparkException: The stream thread was last executing: ``` When a streaming query runs with micro-batch mode, `markMicroBatchExecutionStart()` is invoked [here](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1236) every batch, and it eventually leads to `FileContextBasedCheckpointFileManager.renameTempFile()`. In this method, [fc.rename()](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointFileManager.scala#L391) is invoked where `fc` is an instance of `org.apache.hadoop.fs.FileContext`. Surprisingly, `fc.rename()` seems to clear interrupt statuses. You can confirm that by modifying `CheckpointFileManager.scala` like as follows: ``` @@ -388,7 +388,18 @@ class FileContextBasedCheckpointFileManager(path: Path, hadoopConf: Configuratio override def renameTempFile(srcPath: Path, dstPath: Path, overwriteIfPossible: Boolean): Unit = { import Options.Rename._ + + try { + // Receive interrupt here. + Thread.sleep(1000) + } catch { + case e: Throwable => + // Restore the interrupt status cleared by Thread.sleep. + Thread.currentThread.interrupt + } + println("Interrupt status before fc.rename is called: " + Thread.currentThread.isInterrupted) fc.rename(srcPath, dstPath, if (overwriteIfPossible) OVERWRITE else NONE) + println("Interrupt status after fc.rename is called: " + Thread.currentThread.isInterrupted) ``` And then, run the flaky test. ``` $ build/sbt 'testOnly org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite -- -z "python foreachBatch proce ss: process terminates after query is stopped"' ... Interrupt status before fc.rename is called: true Interrupt status after fc.rename is called: false ... ``` `QueryExecution.interruptAndAwaitExecutionThreadTermination()` is invoked from [MicroBatchExecution.stop()](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L475) so in the test above, if an interrupt status set through [query1.stop()](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionHolderSuite.scala#L387) is unexpectedly cleared, `TimeoutException` is thrown here because first some batches take over 30 seconds (This is because `foreach_batch_worker.py` hangs for several tens of seconds after exit from [process()](https://github.com/apache/spark/blob/4995e6f68c424de361371ecfa6fa0cd87c83617f/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py#L85) and before [write_int()](https://github.com/apache/spark/blob/4995e6f68c424d e361371ecfa6fa0cd87c83617f/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py#L86) is invoked. The reason is not clear for now but it's a separate issue). You can reproduce this issue by modifying `MicroBatchExecution.scala` as well like as follows. ``` override def stop(): Unit = { // Set the state to TERMINATED so that the batching thread knows that it was interrupted // intentionally + Thread.sleep(500) state.set(TERMINATED) if (queryExecutionThread.isAlive) { sparkSession.sparkContext.cancelJobGroup(runId.toString, ``` The `Thread.sleep` inserted is to ensure `state.set(TERMINATED)` is invoked after micro-batches start. And then, run the flaky test. ``` $ build/sbt 'testOnly org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite -- -z "python foreachBatch process: process terminates after query is stopped"' ... [info] - python foreachBatch process: process terminates after query is stopped *** FAILED *** (2 minutes, 12 seconds) [info] java.util.concurrent.TimeoutException: Stream Execution thread for stream foreachBatch_termination_test_q1_516328 7691187916 [id = 3df450c6-5410-49fa-bc3b-0b938830fcac, runId = 757cb7bf-5619-4afd-b82f-b8a3e63d2521] failed to stop within 30000 milliseconds (specified by spark.sql.streaming.stopTimeout). See the cause on what was being executed in the stream ing query thread. [info] at org.apache.spark.sql.execution.streaming.runtime.StreamExecution.interruptAndAwaitExecutionThreadTermination(S treamExecution.scala:510) [info] at org.apache.spark.sql.execution.streaming.runtime.MicroBatchExecution.stop(MicroBatchExecution.scala:476) [info] at org.apache.spark.sql.execution.streaming.runtime.StreamingQueryWrapper.stop(StreamingQueryWrapper.scala:61) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.runPythonForeachBatchTerminationTestBody(S parkConnectSessionHolderSuite.scala:387) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.$anonfun$new$16(SparkConnectSessionHolderS uite.scala:450) [info] at org.apache.spark.sql.connect.service.SparkConnectSessionHolderSuite.$anonfun$awaitTestBodyInNewThread$1(SparkC onnectSessionHolderSuite.scala:281) [info] at java.base/java.lang.Thread.run(Thread.java:840) [info] Cause: org.apache.spark.SparkException: The stream thread was last executing: ... [info] Run completed in 2 minutes, 15 seconds. [info] Total number of tests run: 1 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 0, failed 1, canceled 0, ignored 0, pending 0 [info] *** 1 TEST FAILED *** ``` The reason why `FileContext.rename()` clears interrupt statuses is not clear for now but this PR proposes to restore the interrupt status as a workaround. ### Why are the changes needed? To allow `StreamExecution.stop` can stop micro-batch streaming query by interruption expectedly. ### Does this PR introduce _any_ user-facing change? Yes. `StreamExecution.stop` will work as expected. ### How was this patch tested? Confirmed `python foreachBatch process: process terminates after query is stopped` in `SparkConnectSessionHolderSuite` passes even if `Thread.sleep`s are inserted like above. ``` [info] - python foreachBatch process: process terminates after query is stopped (18 seconds, 765 milliseconds) [info] Passed: Total 0, Failed 0, Errors 0, Passed 0 [info] No tests to run for sql / Test / testOnly [info] Run completed in 22 seconds, 214 milliseconds. [info] Total number of tests run: 1 [info] Suites: completed 1, aborted 0 [info] Tests: succeeded 1, failed 0, canceled 0, ignored 0, pending 0 [info] All tests passed. ``` ### Was this patch authored or co-authored using generative AI tooling? No. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
