This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 20d8cce7e51d [SPARK-57931][CORE] Restore worker channel blocking mode
after pipelined Python UDF execution
20d8cce7e51d is described below
commit 20d8cce7e51daa861e03428794d1b1c6d33a2cf1
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Mon Jul 6 00:00:17 2026 -0700
[SPARK-57931][CORE] Restore worker channel blocking mode after pipelined
Python UDF execution
### What changes were proposed in this pull request?
The pipelined Python UDF path (SPARK-56642) switches a borrowed worker's
`SocketChannel` from non-blocking to blocking mode in `createPipelinedDataIn`
and never restores it. With worker reuse enabled (the default), the worker is
returned to the idle pool with its channel still in blocking mode.
This changes `PythonWorkerFactory.create()` to normalize a reused daemon
worker's channel back to non-blocking before calling `refresh()`, so a pooled
worker is always handed to the next task in the same (non-blocking) mode as a
freshly created one — restoring the invariant that a daemon worker taken from
the pool is non-blocking.
### Why are the changes needed?
`PythonWorker.refresh()` only opens a selector when the channel is
non-blocking. A pooled worker left in blocking mode therefore comes back with a
null `selector` / `selectionKey`. Code on the non-pipelined (single-threaded
NIO selector) path dereferences `worker.selector` / `worker.selectionKey`, so a
worker in this corrupted state would throw a `NullPointerException` there.
In the current OSS code this does not surface as an end-to-end failure,
because the worker-factory cache key (`PythonWorkersKey`) includes the worker
`envVars`, and the pipelined path adds `SPARK_PIPELINED_UDF=1` to `envVars`
before requesting a worker (`BasePythonRunner.compute`). Pipelined and
non-pipelined tasks therefore draw from separate idle pools: a worker left in
blocking mode only returns to the pipelined pool, and the next borrower from
that pool is again a pipelined task, [...]
That masking is fragile. It relies on the two pools staying disjoint via
`envVars`; it does not fix the broken invariant that a pooled daemon worker is
non-blocking. Any worker-management layer that pools or reuses workers across
that boundary — e.g. reusing a warmed worker regardless of whether the previous
task was pipelined — will hand a blocking-mode worker to selector-path code and
hit the `NullPointerException`. Fixing it at the pool boundary (`create()`)
restores the invariant [...]
The fix is applied in `create()` rather than in the pipelined path's
task-completion listener because the worker is released back to the pool from
the reader iterator when it reaches `END_OF_STREAM`, which runs *before* the
task-completion listener; a restore in the listener would therefore run after
the worker is already back in the pool. Normalizing at the single pool exit
point (`create()`) is correct regardless of that ordering.
### Does this PR introduce _any_ user-facing change?
No. This hardens an invariant in an opt-in feature
(`spark.python.udf.pipelined.enabled`) that has not been released yet; the
pool-isolation behavior above means it is not an observable failure in OSS
today.
### How was this patch tested?
New unit test in `PythonWorkerFactorySuite` that constructs a
`PythonWorker` over a loopback channel, puts it in the blocking state the
pipelined path leaves behind (where `refresh()` opens no selector), and asserts
`refreshNonBlocking()` — the method `create()` uses when reusing a pooled
worker — normalizes the channel back to non-blocking and re-opens the selector.
The test uses a mock channel rather than a real daemon worker so it runs in the
core module's test environment (which h [...]
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
Closes #56995 from viirya/fix-pipelined-udf-channel-mode.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../org/apache/spark/api/python/PythonRunner.scala | 5 ++++
.../spark/api/python/PythonWorkerFactory.scala | 22 ++++++++++++++-
.../api/python/PythonWorkerFactorySuite.scala | 32 ++++++++++++++++++++++
3 files changed, 58 insertions(+), 1 deletion(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 923df591cc2e..697545044e0e 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -413,6 +413,11 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
handle: Option[ProcessHandle],
context: TaskContext): DataInputStream = {
// Switch the channel to blocking mode for true full-duplex I/O.
+ // The channel is left in blocking mode after the task completes; with
worker reuse
+ // enabled the worker is returned to the idle pool, so
PythonWorkerFactory.create()
+ // normalizes it back to non-blocking before handing it to the next task
+ // (SPARK-57931). Without that, a later task on the non-pipelined selector
path would
+ // NPE on worker.selector because refresh() only opens a selector in
non-blocking mode.
// Must close the selector first because configureBlocking() fails
// if the channel is registered with a selector
(IllegalBlockingModeException).
if (worker.selectionKey != null) {
diff --git
a/core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala
index 350818e18cb9..c9f83bca62cd 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala
@@ -65,6 +65,21 @@ case class PythonWorker(channel: SocketChannel) {
this
}
+ /**
+ * Normalizes the channel to non-blocking mode and then refreshes. Used when
a worker is
+ * taken from the idle pool: the pipelined Python UDF path (SPARK-56642) may
have switched
+ * the channel to blocking mode without restoring it, and refresh() only
opens a selector
+ * for a non-blocking channel. Restoring non-blocking mode here keeps the
invariant that a
+ * pooled daemon worker is handed out non-blocking, so selector-path code
does not hit a
+ * null selector / selectionKey (SPARK-57931).
+ */
+ def refreshNonBlocking(): this.type = synchronized {
+ if (channel.isBlocking) {
+ channel.configureBlocking(false)
+ }
+ refresh()
+ }
+
def stop(): Unit = synchronized {
closeSelector()
Option(channel).foreach(_.close())
@@ -138,7 +153,12 @@ private[spark] class PythonWorkerFactory(
daemonWorkers.get(worker).foreach { workerHandle =>
if (workerHandle.isAlive()) {
try {
- return (worker.refresh(), Some(workerHandle))
+ // A daemon worker is always handed out in non-blocking mode
(see
+ // createWorker below). The pipelined Python UDF path
temporarily switches
+ // the channel to blocking mode and does not restore it, so
normalize here
+ // before reuse; otherwise refresh() would not open a selector
and the next
+ // task on the selector path would NPE on worker.selector
(SPARK-57931).
+ return (worker.refreshNonBlocking(), Some(workerHandle))
} catch {
case _: CancelledKeyException => /* pass */
}
diff --git
a/core/src/test/scala/org/apache/spark/api/python/PythonWorkerFactorySuite.scala
b/core/src/test/scala/org/apache/spark/api/python/PythonWorkerFactorySuite.scala
index 4f9dafb6cbea..d9315b715bcf 100644
---
a/core/src/test/scala/org/apache/spark/api/python/PythonWorkerFactorySuite.scala
+++
b/core/src/test/scala/org/apache/spark/api/python/PythonWorkerFactorySuite.scala
@@ -80,6 +80,38 @@ class PythonWorkerFactorySuite extends SparkFunSuite with
SharedSparkContext {
}
}
+ test("SPARK-57931: refreshNonBlocking normalizes a blocking channel and
opens a selector") {
+ // The pipelined Python UDF path switches a borrowed worker's channel to
blocking mode
+ // and does not restore it, so a worker returned to the idle pool can be
left blocking.
+ // create() calls refreshNonBlocking() when reusing a pooled worker; it
must return the
+ // channel to non-blocking mode and re-open the selector, otherwise a
later task on the
+ // selector path would NPE on a null worker.selector / worker.selectionKey.
+ val channel = java.nio.channels.SocketChannel.open()
+ try {
+ val worker = PythonWorker(channel)
+
+ // Simulate the state the pipelined path leaves behind: a blocking
channel whose
+ // refresh() opens no selector.
+ channel.configureBlocking(true)
+ worker.refresh()
+ assert(worker.selector === null, "precondition: blocking channel has no
selector")
+ assert(worker.selectionKey === null)
+
+ // Reuse normalization must restore non-blocking mode with a live
selector.
+ worker.refreshNonBlocking()
+ assert(!worker.channel.isBlocking, "channel should be normalized to
non-blocking")
+ assert(worker.selector != null, "a non-blocking channel should have a
selector")
+ assert(worker.selectionKey != null, "a non-blocking channel should have
a selection key")
+
+ // refreshNonBlocking on an already non-blocking channel is a no-op for
the mode.
+ worker.refreshNonBlocking()
+ assert(!worker.channel.isBlocking)
+ assert(worker.selector != null)
+ } finally {
+ channel.close()
+ }
+ }
+
test("idle worker pool is bounded when idleWorkerMaxPoolSize is set") {
sc.conf.set("spark.python.factory.idleWorkerMaxPoolSize", "2")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
