You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ka...@apache.org on 2026/08/14 04:13:54 UTC
(spark) branch branch-4.1 updated: [SPARK-58751][SS][PYTHON] Stop leaking Python workers on TransformWithState init failure
This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new 8f5e2f5e558c [SPARK-58751][SS][PYTHON] Stop leaking Python workers on TransformWithState init failure
8f5e2f5e558c is described below
commit 8f5e2f5e558c16e87ef0c95378f050169ca31b27
Author: jon-gao-db <24...@users.noreply.github.com>
AuthorDate: Fri Aug 14 13:12:39 2026 +0900
[SPARK-58751][SS][PYTHON] Stop leaking Python workers on TransformWithState init failure
### What changes were proposed in this pull request?
This patch makes the driver-side `TransformWithStateInPySpark` pre-initialization runner clean up on every lifecycle path. It moves `init()` and `process()` into `Utils.tryWithSafeFinally`, guaranteeing that `stop()` runs when initialization or processing fails. It also preserves the original initialization failure when cleanup itself throws by relying on `tryWithSafeFinally`'s suppressed-exception behavior.
The runner's state-server daemon thread is nullable until the end of initialization, so `stop()` now checks for null before interrupting it. A six-case regression suite covers initialization failure cleanup, repeated failures, suppressed cleanup errors, process failure wrapping, the success path, and stopping before state-server startup.
JIRA: https://issues.apache.org/jira/browse/SPARK-58751
### Why are the changes needed?
`StreamingPythonRunner.init()` creates the Python worker before it finishes initialization. If initialization fails afterward, the previous code never reached `stop()`, leaking the worker and its associated resources for the driver's lifetime. Repeated streaming restarts can accumulate these leaked resources and eventually prevent new isolated workers from starting.
### Does this PR introduce _any_ user-facing change?
No. Successful execution behavior is unchanged. This only restores cleanup on existing failure paths and preserves the original initialization error instead of allowing cleanup failures or a null-thread error to obscure it.
### How was this patch tested?
Added `TransformWithStateInPySparkPreInitCleanupSuite` with six tests and no real Python worker dependency. Static validation passed with `git diff --check` and the changed-file line-length check.
The focused Spark test could not run because this checkout does not have `sbt` installed and its launcher download was unavailable. A Maven `test-compile` fallback also could not resolve dependencies because the configured Maven mirrors were unreachable from the environment.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (Codex CLI).
Closes #57941 from jon-gao-db/fix/transform-with-state-pre-init-cleanup.
Lead-authored-by: jon-gao-db <24...@users.noreply.github.com>
Co-authored-by: Jonathan Gao <jo...@hotmail.com>
Signed-off-by: Jungtaek Lim <ka...@gmail.com>
(cherry picked from commit 3b39c59d34094a4aa9ae53bdeae05810f2987b9c)
Signed-off-by: Jungtaek Lim <ka...@gmail.com>
---
.../TransformWithStateInPySparkExec.scala | 28 ++--
.../TransformWithStateInPySparkPythonRunner.scala | 4 +-
...formWithStateInPySparkPreInitCleanupSuite.scala | 166 +++++++++++++++++++++
3 files changed, 186 insertions(+), 12 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
index c10d21933c2f..5196e2798f28 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala
@@ -157,17 +157,7 @@ case class TransformWithStateInPySparkExec(
groupingKeySchema,
driverProcessorHandle
)
- // runner initialization
- runner.init()
- try {
- // execute UDF on the python runner
- runner.process()
- } catch {
- case e: Throwable =>
- throw new SparkException("TransformWithStateInPySpark driver worker " +
- "exited unexpectedly (crashed)", e)
- }
- runner.stop()
+ TransformWithStateInPySparkExec.runPreInitRunner(runner)
val info = getStateInfo
val stateSchemaDir = stateSchemaDirPath()
@@ -408,6 +398,22 @@ case class TransformWithStateInPySparkExec(
// scalastyle:off argcount
object TransformWithStateInPySparkExec {
+ private[streaming] def runPreInitRunner(
+ runner: TransformWithStateInPySparkPythonPreInitRunner): Unit = {
+ Utils.tryWithSafeFinally {
+ runner.init()
+ try {
+ runner.process()
+ } catch {
+ case e: Throwable =>
+ throw new SparkException("TransformWithStateInPySpark driver worker " +
+ "exited unexpectedly (crashed)", e)
+ }
+ } {
+ runner.stop()
+ }
+ }
+
// Plan logical transformWithStateInPySpark for batch queries
def generateSparkPlanForBatchQueries(
functionExpr: Expression,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
index 42d4ad68c29a..b0bf851087e0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala
@@ -314,7 +314,9 @@ class TransformWithStateInPySparkPythonPreInitRunner(
override def stop(): Unit = {
super.stop()
closeServerSocketChannelSilently(stateServerSocket)
- daemonThread.interrupt()
+ if (daemonThread != null) {
+ daemonThread.interrupt()
+ }
}
private def startStateServer(): Unit = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala
new file mode 100644
index 000000000000..e56a8a671dd4
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala
@@ -0,0 +1,166 @@
+/*
+ * 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.sql.execution.python.streaming
+
+import java.io.{DataInputStream, DataOutputStream}
+import java.util.{ArrayList, HashMap}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
+
+import org.apache.spark.SparkException
+import org.apache.spark.api.python.{PythonFunction, SimplePythonFunction}
+import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.statefulprocessor.DriverStatefulProcessorHandleImpl
+import org.apache.spark.sql.streaming.TimeMode
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StructType
+
+class TransformWithStateInPySparkPreInitCleanupSuite extends SharedSparkSession {
+
+ private val groupingKeySchema = StructType(Nil)
+
+ private def newPythonFunction(): PythonFunction = {
+ new SimplePythonFunction(
+ command = Seq.empty[Byte],
+ envVars = new HashMap[String, String](),
+ pythonIncludes = new ArrayList[String](),
+ pythonExec = "python3",
+ pythonVer = "3",
+ broadcastVars = null,
+ accumulator = null)
+ }
+
+ private def newDriverHandle(): DriverStatefulProcessorHandleImpl = {
+ new DriverStatefulProcessorHandleImpl(TimeMode.None(), null)
+ }
+
+ private class StubPreInitRunner(
+ failInitWith: Option[Throwable] = None,
+ failProcessWith: Option[Throwable] = None,
+ failStopWith: Option[Throwable] = None)
+ extends TransformWithStateInPySparkPythonPreInitRunner(
+ newPythonFunction(),
+ "pyspark.sql.streaming.transform_with_state_driver_worker",
+ groupingKeySchema,
+ newDriverHandle()) {
+
+ val workerAlive = new AtomicBoolean(false)
+ val initCount = new AtomicInteger(0)
+ val processCount = new AtomicInteger(0)
+ val stopCount = new AtomicInteger(0)
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ initCount.incrementAndGet()
+ workerAlive.set(true)
+ failInitWith.foreach(error => throw error)
+ (null, null)
+ }
+
+ override def process(): Unit = {
+ processCount.incrementAndGet()
+ failProcessWith.foreach(error => throw error)
+ }
+
+ override def stop(): Unit = {
+ stopCount.incrementAndGet()
+ workerAlive.set(false)
+ failStopWith.foreach(error => throw error)
+ }
+ }
+
+ private def runPreInit(runner: TransformWithStateInPySparkPythonPreInitRunner): Unit = {
+ TransformWithStateInPySparkExec.runPreInitRunner(runner)
+ }
+
+ test("init failure after worker creation still stops the runner") {
+ val initFailure = new RuntimeException("init failed")
+ val runner = new StubPreInitRunner(failInitWith = Some(initFailure))
+
+ val thrown = intercept[Throwable] {
+ runPreInit(runner)
+ }
+
+ assert(thrown eq initFailure)
+ assert(runner.initCount.get() === 1)
+ assert(runner.stopCount.get() === 1)
+ assert(!runner.workerAlive.get())
+ assert(runner.processCount.get() === 0)
+ }
+
+ test("repeated init failures do not accumulate live workers") {
+ val runners = (1 to 20).map { _ =>
+ val runner = new StubPreInitRunner(failInitWith = Some(new RuntimeException("init failed")))
+ intercept[Throwable] {
+ runPreInit(runner)
+ }
+ runner
+ }
+
+ assert(runners.forall(_.initCount.get() === 1))
+ assert(runners.count(_.workerAlive.get()) === 0)
+ }
+
+ test("a stop failure does not mask the original init failure") {
+ val initFailure = new RuntimeException("init failed")
+ val stopFailure = new IllegalStateException("cleanup failed")
+ val runner =
+ new StubPreInitRunner(failInitWith = Some(initFailure), failStopWith = Some(stopFailure))
+
+ val thrown = intercept[Throwable] {
+ runPreInit(runner)
+ }
+
+ assert(thrown eq initFailure)
+ assert(thrown.getSuppressed.contains(stopFailure))
+ assert(runner.stopCount.get() === 1)
+ assert(!runner.workerAlive.get())
+ }
+
+ test("process failure is wrapped but still stops the runner") {
+ val processFailure = new RuntimeException("worker crashed")
+ val runner = new StubPreInitRunner(failProcessWith = Some(processFailure))
+
+ val thrown = intercept[SparkException] {
+ runPreInit(runner)
+ }
+
+ assert(thrown.getMessage.contains("exited unexpectedly (crashed)"))
+ assert(thrown.getCause eq processFailure)
+ assert(runner.stopCount.get() === 1)
+ assert(!runner.workerAlive.get())
+ }
+
+ test("success path stops the runner exactly once") {
+ val runner = new StubPreInitRunner()
+
+ runPreInit(runner)
+
+ assert(runner.initCount.get() === 1)
+ assert(runner.processCount.get() === 1)
+ assert(runner.stopCount.get() === 1)
+ assert(!runner.workerAlive.get())
+ }
+
+ test("stop before startStateServer ran does not throw NPE") {
+ val runner = new TransformWithStateInPySparkPythonPreInitRunner(
+ newPythonFunction(),
+ "pyspark.sql.streaming.transform_with_state_driver_worker",
+ groupingKeySchema,
+ newDriverHandle())
+
+ runner.stop()
+ runner.stop()
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org