Repository navigation
Conversation
…exit before task completion cleanup
In pipelined Python UDF mode (`spark.python.udf.pipelined.enabled`), the listener that stops the writer thread now waits until the writer has actually exited, and it runs before every task completion listener.
1. Add `TaskContext.addTaskPreCompletionListener` (`private[spark]`, abstract; `BarrierTaskContext` delegates). `TaskContextImpl.markTaskCompleted` runs these listeners on its own thread before it marks the task completed, outside the listener invocation used for the other listeners. Only then does it set `completed` (in the same `synchronized` block that finds no more pre-completion listeners) and run the completion listeners as before. While the pre-completion listeners run, a listener that another thread adds sees the task as not completed and is only queued, so that thread can never start running the listeners. A pre-completion listener added after completion is handed to `addTaskCompletionListener`, so it runs like a completion listener added at that point. A failing one marks the task failed, and its failure is reported together with those of the completion listeners. Tasks without pre-completion listeners behave as before.
2. The writer is wrapped in `BasePythonRunner.PipelinedWriterTask`. Its state (not started, running on a thread, or finished) is kept in a single field guarded by the task, and a latch is counted down on every transition to finished. The stop listener is added with `addTaskPreCompletionListener`, and the writer is submitted only if the task has not completed by then. The listener interrupts the writer thread directly and waits on that latch. If it runs before the pool has started the writer, it marks the writer finished, so the writer never runs.
3. The wait has no time bound, the same as the `WriterThread.join()` that the writer thread used before SPARK-44705 (added for SPARK-33277). While it waits, it logs a warning every 10 seconds that names the writer thread (with its stack trace the first time), and it interrupts the writer again, in case code it called consumed the interrupt. Like `Task.kill`, it does not interrupt while `context.interruptible()` is false, which is while some thread of the task, not necessarily the writer, creates a resource through `TaskContext.createResourceUninterruptibly`; the writer is interrupted at a later warning then. A failure to interrupt the writer or to log the wait does not end the wait; it is thrown once the writer has exited. An interrupt of the listener thread that arrives during the wait does not stop it and is not passed on to the completion listeners that run next (the kill is already recorded in the `TaskContext`); an interrupt that was pending before the wait is kept. If the listener ever ran on the writer thread itself, it would not wait, since the writer can only exit after it returns, and would log an error.
4. The writer writes through a `CloseShieldOutputStream` (added by SPARK-60305) around its `DirectByteBufferOutputStream`. The Arrow writers write through `Channels.newChannel(dataOut)`, an interruptible channel. When the stop listener interrupts the writer during such a write, the JDK closes that channel and the stream under it on the interrupting thread, and the buffer was freed while the writer still copied into it. Now the buffer is freed only in the writer's own `finally`.
5. A fatal error of the writer (e.g. `OutOfMemoryError`), including one while it sets up its buffer, is passed to the reader through `writer.setException` with a socket `shutdownOutput()`, and it is logged. Previously the thread pool's `FutureTask` kept it, and the reader waited for the Python worker.
The existing listener does:
```scala
writerFuture.cancel(true)
try { writerFuture.get() } catch { case _: CancellationException | ... => }
```
Its comment says this waits for the writer to exit before the listeners that run after it (registered earlier, LIFO) free the memory backing the input rows. But once `cancel(true)` succeeds, `FutureTask.get()` throws `CancellationException` right away and does not wait for the interrupted runnable to return. So the writer can still be serializing a row, or adding it to a `HybridRowQueue` page, when the queue's close listener frees that page. That is a use-after-free, and with off-heap memory it can crash the executor. A `HybridQueue.createNewQueue` racing with `close()` can also add a page that nobody closes until `cleanUpAllAllocatedMemory`.
A task completes while the writer is still running when the output is not fully consumed, for example with a `LIMIT` or when the task fails.
Ordering the writer listener in the completion listener stack is not enough. Listeners run in reverse order of registration, and the writer thread itself adds completion listeners lazily while it pulls the upstream iterator, e.g. `ParquetReaderCallback.initIfNotAlready` of the DSv2 Parquet reader (which closes the vectorized reader and its off-heap column vectors) or the `UnsafeExternalSorter` that `ExternalAppendOnlyUnsafeRowArray` creates once it spills. Those would run before the writer listener. Running the stop listener inside the usual listener invocation is not enough either: `TaskContextImpl` can let another thread run the pending listeners, e.g. when a kill thread is running failure listeners as the task completes (SPARK-60222). If that thread were the writer, it would run its own stop listener. Running the pre-completion listeners before the task is marked completed avoids both.
The wait cannot be bounded safely: if the listener gave up, the listeners that run next would still free the memory, and throwing from the listener does not stop them either. Even a writer blocked in upstream I/O that ignores interrupts (e.g. a slow JDBC fetch) adds the row it gets next to the `HybridRowQueue`, whose `InMemoryRowQueue` writes through the page's base object captured at construction, so a freed page would be corrupted even on-heap (`HeapMemoryAllocator` pools pages for reuse). The price is that task completion is held for as long as such I/O blocks, including I/O that only a close listener of the upstream reader would unblock, since that listener now runs after the wait; `spark.task.reaper.enabled` with a positive `spark.task.reaper.killTimeout` bounds this for killed tasks, by killing the executor JVM, except in local mode, where the reaper only logs an error. This is documented in the scaladoc and the warning. Usually the writer exits soon after the interrupt: `channel.write` fails with `ClosedByInterruptException`, and otherwise the writer loop checks the interrupt flag between `writeNextInputToStream` calls, each of which can pull and serialize a whole batch (e.g. for the Arrow runners).
The pipelined mode was added in 4.3.0 by SPARK-56642, so branch-4.3 and branch-4.x have the same issue. The same interrupt-closes-the-buffer problem also affects the default (non-pipelined) Arrow UDF path, where `Task.kill` interrupts the task thread while it serializes into `ReaderInputStream`'s buffer; SPARK-60305 fixed that separately, since it needs backports to 4.0 through 4.3, and added the `CloseShieldOutputStream` that this change also uses.
No. This fixes a bug in an opt-in mode that is not released yet (4.3.0). For tasks with no pre-completion listener, listener behavior is unchanged.
New tests in `BasePythonRunnerSuite`:
- Task completion stays blocked until a running writer that ignores interrupts exits. Both a cleanup listener added on the task thread before the writer starts and one that the writer thread adds after the stop listener see that the writer has exited. The tests first wait until task completion has reached the writer listener.
- A listener that the writer thread adds after task completion has started is queued and runs after the writer exits.
- Stopping interrupts a running writer right away; an interrupt of the stopping thread, whether it arrives during the wait or while stopping, does not stop it and is not passed on; an interrupt pending before stopping is kept.
- The writer is not interrupted while it creates a resource uninterruptibly, and is interrupted afterwards.
- A writer stopped before it starts never runs, also when it is stopped twice, when it is started after the task completed, or when it is started while another thread runs the completion listeners. Stopping on the writer thread itself does not wait.
- A failure to interrupt the writer does not end the wait, and is thrown once the writer has exited.
- A writer that consumed the first interrupt is interrupted again.
- An interrupted Arrow-style channel write in a later buffer leaves the buffer usable, through a real `PipelinedWriterRunnable` (with a socket pair standing in for the Python worker).
- A warning that fails on every tick does not skip the repeated interrupt.
- A fatal error of a real `PipelinedWriterRunnable` is rethrown, logged, and passed to the reader, and the Python side sees EOF.
New tests in `TaskContextSuite`:
- Pre-completion listener ordering and failure handling, and a pre-completion listener added after completion running like a completion listener, or failing the task and throwing if it fails. Failures of pre-completion and completion listeners are merged into one exception.
- While a pre-completion listener runs, listeners added from another thread are only queued and run on the completing thread afterwards, and failure listeners triggered from another thread still run.
Before the fix, a test with the old listener code showed `markTaskCompleted` returning while the writer was still running. Each of the following changes, reverted on its own, makes its test fail: the cancel-then-`get()` wait, adding the stop listener as a pre-completion listener, running the pre-completion listeners before setting `completed`, the close shield, the immediate interrupt, skipping the interrupt during uninterruptible resource creation, the repeated interrupt, clearing and keeping the interrupt status, and passing fatal errors to the reader.
Also ran the whole `core` test suite, the `sql` suites that use `markTaskFailed`, `markTaskCompleted` or `markInterrupted` (`JdbcTaskInterruptSuite`, `SupportsArchiveFormatSuite`, `JsonSuite`, `ParquetStorageFilterSuite`, `ValueStateSuite`, `RocksDBStateStoreSuite`, `RocksDBStateStoreLockHardeningSuite`, `RocksDBCheckpointFailureInjectionSuite`), and `pyspark.sql.tests.pandas.test_pipelined_udf`.
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
Closes apache#59332 from viirya/SPARK-60113-wait-pipelined-writer.
Authored-by: Liang-Chi Hsieh <viirya@gmail.com>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>
(cherry picked from commit ca1f5c5)
dongjoon-hyun
approved these changes
Oct 20, 2026
dongjoon-hyun
left a comment
dongjoon-hyun
left a comment
Member
There was a problem hiding this comment.
+1, LGTM (Pending CIs).
viirya
added a commit
that referenced
this pull request
Oct 20, 2026
…r to exit before task completion cleanup ### What changes were proposed in this pull request? This is a backport of #59332 to `branch-4.3`. In pipelined Python UDF mode (`spark.python.udf.pipelined.enabled`), the listener that stops the writer thread now waits until the writer has actually exited, and it runs before every task completion listener: `TaskContextImpl.markTaskCompleted` runs the new pre-completion listeners on its own thread before it marks the task completed. See #59332 for the full description. The cherry-pick of the merged commit (ca1f5c5) applied cleanly except for the imports of `BasePythonRunnerSuite`: `branch-4.3` does not have the SPARK-59773 worker metrics tests that use some of master's imports, so only the imports that the backported tests use are added. `CloseShieldOutputStream` (SPARK-60305) and the per-buffer `DataOutputStream` of SPARK-60343, which this change builds on, are already in `branch-4.3`. ### Why are the changes needed? The pipelined mode was added in 4.3.0 by SPARK-56642, so `branch-4.3` has the same use-after-free: once `cancel(true)` succeeds, the old listener's `Future.get()` returns without waiting for the writer, so the listeners that run after it can free memory the writer still uses, e.g. the pages of the `HybridRowQueue`. ### Does this PR introduce _any_ user-facing change? No. This fixes a bug in the opt-in pipelined mode (default `false`). For tasks with no pre-completion listener, listener behavior is unchanged. ### How was this patch tested? On this branch, with JDK 17: - `build/sbt 'core/testOnly *BasePythonRunnerSuite *TaskContextSuite *CloseShieldOutputStreamSuite *BarrierTaskContextSuite *ExecutorSuite *StreamingShuffleSuite'` - `build/sbt core/scalastyle core/Test/scalastyle core/mimaReportBinaryIssues` (scalastyle passes; MiMa reports no problem from this change) - `python/run-tests --testnames pyspark.sql.tests.pandas.test_pipelined_udf` after `build/sbt -Phive package` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) This pull request and its description were written by Isaac. Closes #59647 from viirya/SPARK-60113-4.3. Authored-by: Liang-Chi Hsieh <viirya@gmail.com> Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>
Member
Author
|
Merge Summary:
Posted by |
Member
Author
|
Thank you @dongjoon-hyun ! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This is a backport of #59332 to
branch-4.3.In pipelined Python UDF mode (
spark.python.udf.pipelined.enabled), the listener that stops the writer thread now waits until the writer has actually exited, and it runs before every task completion listener:TaskContextImpl.markTaskCompletedruns the new pre-completion listeners on its own thread before it marks the task completed. See #59332 for the full description.The cherry-pick of the merged commit (ca1f5c5) applied cleanly except for the imports of
BasePythonRunnerSuite:branch-4.3does not have the SPARK-59773 worker metrics tests that use some of master's imports, so only the imports that the backported tests use are added.CloseShieldOutputStream(SPARK-60305) and the per-bufferDataOutputStreamof SPARK-60343, which this change builds on, are already inbranch-4.3.Why are the changes needed?
The pipelined mode was added in 4.3.0 by SPARK-56642, so
branch-4.3has the same use-after-free: oncecancel(true)succeeds, the old listener'sFuture.get()returns without waiting for the writer, so the listeners that run after it can free memory the writer still uses, e.g. the pages of theHybridRowQueue.Does this PR introduce any user-facing change?
No. This fixes a bug in the opt-in pipelined mode (default
false). For tasks with no pre-completion listener, listener behavior is unchanged.How was this patch tested?
On this branch, with JDK 17:
build/sbt 'core/testOnly *BasePythonRunnerSuite *TaskContextSuite *CloseShieldOutputStreamSuite *BarrierTaskContextSuite *ExecutorSuite *StreamingShuffleSuite'build/sbt core/scalastyle core/Test/scalastyle core/mimaReportBinaryIssues(scalastyle passes; MiMa reports no problem from this change)python/run-tests --testnames pyspark.sql.tests.pandas.test_pipelined_udfafterbuild/sbt -Phive packageWas this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.