viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4250752636
##########
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:
Thanks, confirmed. Fixed in 1053b61d411 on both sides. `stopAndAwaitExit`
now has the same-thread guard you suggested: `PipelinedWriterTask` records the
thread running the writer, and if the stop listener runs on that thread, it
logs a warning and returns after `cancel(true)` instead of waiting. The new
test "stopping on the writer thread itself does not wait for the writer" covers
it.
The handoff that led there is also closed in `TaskContextImpl` (see
dongjoon's comments on L322 and L208). When another thread holds the listener
invocation, `markTaskCompleted` now waits for it and gets the invocation handed
over directly, instead of returning. A failing listener no longer releases the
invocation while the failure listeners run. So the guard is a safety net rather
than the main mechanism.
##########
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:
Added "a fatal error of the pipelined writer is passed to the reader". It
runs a real `PipelinedWriterRunnable` against a socket pair, with a `Writer`
whose `writeNextInputToStream` throws a `LinkageError`, and checks that the
error is rethrown, that `writer.exception` holds it, and that the Python side
reads EOF. It fails if the branch is narrowed back to `NonFatal` or if the
rethrow is removed.
##########
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:
Thanks, reworded as suggested.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -267,13 +313,14 @@ private[spark] class TaskContextImpl(
}
def getNextListenerOrDeregisterThread(): Option[T] = synchronized {
- if (listeners.empty()) {
- // We have executed all listeners that have been added so far.
Deregister this thread as the
- // callback invocation thread.
- listenerInvocationThread = None
- None
- } else {
- Some(listeners.pop())
+ listenerStacks.find(!_.empty()) match {
+ case Some(listeners) =>
+ Some(listeners.pop())
+ case None =>
+ // We have executed all listeners that have been added so far.
Deregister this thread as
+ // the callback invocation thread.
+ listenerInvocationThread = None
Review Comment:
Done as you suggested. When a listener throws, the catch block no longer
clears `listenerInvocationThread`. It calls `markTaskFailed` in a nested mode,
which runs the failure listeners as part of the current invocation and keeps
the registration when they are done. Since nothing else cleared it, the old
reclaim in `finally` is gone, and so is the window in which two threads could
pop listeners at once. In the old code, `markTaskFailed` reclaimed the
invocation for the failure listeners right away, so the windows were the gaps
around that reclaim rather than the whole time failure listeners run. The new
test "a failing listener does not let another thread take over the listeners"
therefore uses contention: two pre-completion listeners throw while another
thread keeps adding listeners, repeated 200 times, and the pending listener
must run on the task thread. With the old catch block restored, it fails on the
first run.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -177,29 +197,48 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ val claimedInvocation = synchronized {
if (completed) return
completed = true
+ // Claim the listener invocation together with setting `completed`,
unless another thread is
+ // already invoking listeners. Otherwise a thread that adds a completion
listener in between
+ // would see `completed` and run the listeners itself. That thread can
be one the listeners
+ // are meant to stop and wait for, e.g. a pipelined Python UDF writer
that adds listeners
+ // lazily while it pulls its input, which would then wait for itself.
+ if (listenerInvocationThread.isEmpty) {
Review Comment:
Went with the first option. If another thread holds the invocation,
`markTaskCompleted` registers itself as the next invoker and waits. The holder
hands the invocation to it directly when it finishes, so no third thread (such
as the writer) can take it in between. The completing thread then runs the
pre-completion and completion listeners, and `markTaskCompleted` returns only
after they have run, so `Task.run` no longer reaches
`cleanUpAllAllocatedMemory` first. If `markTaskCompleted` is called from a
listener on the thread that already holds the invocation, it runs the
completion listeners as part of that invocation instead of waiting. An
interrupt during the wait is restored afterwards.
The one risk is a holder that itself waits for the task thread. Only threads
running failure listeners can hold the invocation at that point (e.g. the kill
thread after an interrupt listener failed), and a failure listener that blocks
on the task thread seems unlikely. The new test "completion waits for another
thread invoking listeners and runs them" reproduces your example: a failing
interrupt listener on a kill thread, with a failure listener that blocks.
`markTaskCompleted` stays blocked, a listener added from a third thread does
not take over, and once the failure listener returns, all completion listeners
run on the completing thread.
##########
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)
Review Comment:
Agreed, and documented. The scaladoc of `startPipelinedWriter` now says that
a writer blocked in upstream I/O that ignores interrupts holds task completion
until that I/O returns, including I/O that only the upstream reader's close
listener would unblock, and that `spark.task.reaper.enabled` bounds how long a
killed task can stay there. The 10-second warning says the same. I didn't add a
configurable bound: when it expires, failing the listener would not stop the
listeners after it from freeing memory under the writer, so the bound would
only trade the hang for the use-after-free.
For writers that are blocked because an interrupt was consumed
(HyukjinKwon's question in the review summary), the wait now interrupts the
writer again on every warning tick, as long as the writer is still running.
##########
core/src/main/scala/org/apache/spark/TaskContext.scala:
##########
@@ -147,6 +147,21 @@ abstract class TaskContext extends Serializable {
})
}
+ /**
+ * Adds a listener to be executed on task completion before every listener
added with
+ * `addTaskCompletionListener`, regardless of when or from which thread that
listener was
Review Comment:
Thanks, reworded along those lines: "before every pending listener added
with `addTaskCompletionListener`, from any thread and even if that listener was
added later. Completion listeners that have already run are not affected, so to
run before all of them, the listener must be added before the task completes."
##########
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:
Good catch. The wait now leaves the caller's interrupt status as it found
it. It reads and clears the flag before waiting and restores it afterwards.
Interrupts that arrive during the wait are still not passed on. Both paths
(writer never claimed, writer claimed) now leave a pending interrupt in place,
and the comment says so. The new test "an interrupt pending before stopping the
writer is kept" covers both paths.
--
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]
