HyukjinKwon commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4246906867
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,97 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while task completion waits for a pipelined writer to
exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
pre-completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Task completion listeners free memory that the writer may still read or
write, e.g. the
+ * pages of the HybridRowQueue that buffers the input rows, or the off-heap
column vectors of a
+ * vectorized reader. If the writer were still running at that point, it
would be a
+ * use-after-free that can crash the executor (SPARK-33277). The listener is
added with
+ * `addTaskPreCompletionListener` so that it runs before all of them,
including the ones that
+ * the writer thread adds lazily while it pulls the upstream iterator, which
would run first if
+ * this were an ordinary completion listener. `Future.get()` cannot be used
for this wait:
+ * once `cancel(true)` succeeds it throws `CancellationException`
immediately, without waiting
+ * for the interrupted runnable to return.
+ */
+ private[python] def startPipelinedWriter(writer: Runnable, context:
TaskContext): Unit = {
+ val writerTask = new PipelinedWriterTask(writer, taskIdentifier(context))
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
+ writerTask.stopAndAwaitExit(writerFuture)
+ })
+ }
+
+ /**
+ * Wraps a pipelined writer so that task completion can wait until the
writer has either
+ * exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable, taskName: String)
+ extends Runnable {
+ // Set by whichever of run() and stopAndAwaitExit() gets here first. If
run() wins, the
+ // writer runs and stopAndAwaitExit() waits for it; otherwise run() skips
the writer.
+ private val claimed = new AtomicBoolean(false)
+ private val exited = new CountDownLatch(1)
+
+ override def run(): Unit = {
+ if (claimed.compareAndSet(false, true)) {
+ try {
+ writer.run()
+ } catch {
+ case t: Throwable =>
+ // The thread pool's FutureTask would keep this to itself, and the
result of a
+ // cancelled future cannot be retrieved, so log it here.
+ logError(log"Pipelined Python writer of ${MDC(TASK_NAME,
taskName)} failed", t)
+ throw t
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_]): Unit = {
Review Comment:
**Non-blocking (P2):** The new claim in `markTaskCompleted` closes the
window from the last round when the slot is empty. When another thread already
holds `listenerInvocationThread`, though, the task thread still returns without
running any completion listener, and whichever thread adds a listener next
inherits them. That can happen when `Task.kill` runs `markInterrupted` on the
kill thread and a `TaskInterruptListener` throws (`addTaskInterruptListener` is
public), so `markTaskFailed` runs the failure listeners on that thread. It can
also happen when a failure listener is added from another thread after the task
failed. If the pipelined writer, which nobody has stopped yet, then lazily adds
a completion listener (e.g. `ParquetReaderCallback.initIfNotAlready`), it
claims the slot and pops this pre-completion listener first. `cancel(true)`
interrupts the writer itself, the `claimed` CAS fails, and the loop waits on
`exited`, which only the `finally` of `run()` further up the same
stack can count down. The pool thread and the Python worker leak permanently,
and `worker.stop()`, the `HybridRowQueue` close, and the reader closes never
run. The old `cancel(true)` + `get()` returned at once on this path. The
listener-failure handoff you mentioned ends in the same wait if it is ever hit.
**Recommended change:** In PipelinedWriterTask, remember the thread that
claimed and is running the writer. In stopAndAwaitExit, after cancel(true), if
the current thread is that writer thread, log that task completion cleanup is
running on the writer thread and return without waiting, because the writer can
only exit after this listener returns. Add Python runner unit-test coverage
that invokes the stop listener on the writer's own thread while the writer is
still running.
**Why this works:** On the writer thread, the latch can only be counted down
by run()'s finally block further up the same stack, so waiting there can never
succeed. Skipping the wait on that one thread restores the merge-target
behavior for this handoff: the listener returns, the remaining cleanup runs,
and the writer exits on its self-interrupt. Every other thread keeps the
unbounded wait.
**Scope:** Add a same-thread guard to the pipelined writer stop listener and
a regression test for it.
**Compatibility:** The unbounded wait from the task thread and the
run()/stopAndAwaitExit claim protocol stay as they are.
**Risks:** On the self-invocation path, the remaining completion listeners
still run on the writer thread while it is in the middle of pulling input, as
in the merge target. The guard prevents the hang; it does not add a
use-after-free safety net for that narrow path.
**Constraints:** Keep the unbounded, interrupt-ignoring wait unchanged for
every thread other than the writer. Keep the claim protocol between run() and
stopAndAwaitExit, including the never-started case.
**Success:** stopAndAwaitExit never blocks when it is invoked on the thread
that is running this writer. Invoked from any other thread, it still interrupts
the writer and returns only after the writer has exited or is guaranteed never
to run.
**Verification:**
- **Regression:** Invoking the stop listener on the writer's own thread
while that writer is still running returns instead of hanging, and the writer
then finishes and counts down its exit latch.
- **Behavior:** Task completion from another thread still waits for a
running writer, skips a writer that never started, and waits for a writer
claimed after the interrupt.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -1336,9 +1409,9 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
// automatically closed by the JVM, which will cause Python to
receive
// EOF and the reader thread to get IOException.
Thread.currentThread().interrupt()
- case NonFatal(t) =>
+ case t: Throwable =>
Review Comment:
**Nit (P3):** With this change, a fatal error goes through `setException`,
`shutdownOutput()`, and the rethrow below, but no test runs this branch. `a
fatal error of the pipelined writer is logged` wraps a bare `Runnable` in
`PipelinedWriterTask`, so narrowing this back to `NonFatal(t)`, or dropping the
rethrow, would keep every test green. The reader would then again wait on
Python with no cause recorded, which is the hang raised earlier in review.
**Verification:**
- **Regression:** With the fatal-error branch narrowed back to NonFatal, or
with the rethrow removed, the new coverage fails.
##########
core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala:
##########
@@ -757,6 +757,113 @@ class TaskContextSuite extends SparkFunSuite with
BeforeAndAfter with LocalSpark
assert(invocationOrder === Seq("C", "B", "A", "D"))
}
+ test("SPARK-60113: pre-completion listeners run before all completion
listeners") {
+ val context = TaskContext.empty()
+ val invocationOrder = ArrayBuffer.empty[String]
+
+ def makeLoggingListener(id: String): TaskCompletionListener = new
TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += id
+ }
+ }
+ context.addTaskCompletionListener(makeLoggingListener("A"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P1"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P2"))
+ // Added after the pre-completion listeners, e.g. lazily by the thread
they stop. An ordinary
Review Comment:
**Nit (P3):** `B` is an ordinary completion listener added at this exact
point, and the assertion below expects it to run after `P2` and `P1`. Read
literally, this sentence says the opposite of what the test checks. Something
like: "Had P1 and P2 been added with addTaskCompletionListener, B would run
before them."
**Verification:**
- **Inspection:** The comment names the condition (P1 and P2 being ordinary
completion listeners) under which B would run before them.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,97 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while task completion waits for a pipelined writer to
exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
pre-completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Task completion listeners free memory that the writer may still read or
write, e.g. the
+ * pages of the HybridRowQueue that buffers the input rows, or the off-heap
column vectors of a
+ * vectorized reader. If the writer were still running at that point, it
would be a
+ * use-after-free that can crash the executor (SPARK-33277). The listener is
added with
+ * `addTaskPreCompletionListener` so that it runs before all of them,
including the ones that
+ * the writer thread adds lazily while it pulls the upstream iterator, which
would run first if
+ * this were an ordinary completion listener. `Future.get()` cannot be used
for this wait:
+ * once `cancel(true)` succeeds it throws `CancellationException`
immediately, without waiting
+ * for the interrupted runnable to return.
+ */
+ private[python] def startPipelinedWriter(writer: Runnable, context:
TaskContext): Unit = {
+ val writerTask = new PipelinedWriterTask(writer, taskIdentifier(context))
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
+ writerTask.stopAndAwaitExit(writerFuture)
+ })
+ }
+
+ /**
+ * Wraps a pipelined writer so that task completion can wait until the
writer has either
+ * exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable, taskName: String)
+ extends Runnable {
+ // Set by whichever of run() and stopAndAwaitExit() gets here first. If
run() wins, the
+ // writer runs and stopAndAwaitExit() waits for it; otherwise run() skips
the writer.
+ private val claimed = new AtomicBoolean(false)
+ private val exited = new CountDownLatch(1)
+
+ override def run(): Unit = {
+ if (claimed.compareAndSet(false, true)) {
+ try {
+ writer.run()
+ } catch {
+ case t: Throwable =>
+ // The thread pool's FutureTask would keep this to itself, and the
result of a
+ // cancelled future cannot be retrieved, so log it here.
+ logError(log"Pipelined Python writer of ${MDC(TASK_NAME,
taskName)} failed", t)
+ throw t
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_]): Unit = {
+ // Interrupts the writer thread if the writer is running. This unblocks
channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException).
Otherwise the writer
+ // sees the interrupt only between calls to writeNextInputToStream, each
of which can pull
+ // and serialize a whole batch of input rows (e.g. for the Arrow
runners).
+ future.cancel(true)
+ if (claimed.compareAndSet(false, true)) {
+ // The writer never started and now never will.
+ return
+ }
+ // Wait without a time bound. The completion listeners that run next
free memory that
+ // the writer may still access: even a writer blocked in upstream I/O
that ignores
+ // interrupts adds the row it gets next to the HybridRowQueue. Giving up
here would allow
+ // that use-after-free, and failing this listener would not stop those
listeners either.
+ // The price is that task completion is held for as long as such
upstream I/O blocks.
+ //
+ // An interrupt of this thread (e.g. a task kill) is not restored after
the wait: the kill
+ // is already recorded in the TaskContext, and the flag would make the
blocking channel
+ // and lock operations of the completion listeners that run next fail
and skip cleanup.
+ val startNs = System.nanoTime()
+ var done = false
+ while (!done) {
+ try {
+ done = exited.await(pipelinedWriterExitWarnIntervalMs,
TimeUnit.MILLISECONDS)
Review Comment:
Thanks, I checked 038c4b3420e. `markTaskCompleted` now takes
`listenerInvocationThread` in the same `synchronized` block that sets
`completed`, only that entry point pre-claims, and the new `TaskContextSuite`
and `BasePythonRunnerSuite` tests cover a listener added just as completion
starts and one added by the writer after completion has reached its stop
listener. I agree the listener-failure handoff you describe is fine to leave:
it needs a pre-completion listener that throws, and the stop listener does not
throw in practice.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4242065501","thread_id":"inline:4242065501","verdict_sha256":"6704d95e924fc29e87f8cb36a03662f4adde76842bc4b448ca5b970a0c19ca82"}
-->
--
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]
