dongjoon-hyun commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4247343518
##########
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:
Putting this listener first also puts it ahead of the close listeners that
the writer thread registers lazily (a DSv2/JDBC/S3 reader opened on its first
pull). Before this PR those ran first (LIFO), and closing the stream was what
unblocked a writer stuck in a socket read that ignores `Thread.interrupt`.
Now `stopAndAwaitExit` waits without bound first, logging every 10s, and the
close listener never gets to run. Since the task reaper is off by default, the
task never completes and the executor slot is lost.
The UAF fix needs this ordering, so I'm not suggesting a revert. Could we at
least document this hang mode (and the reaper as the escape hatch) in the
scaladoc and the warning message, or add a configurable bound after which the
listener fails loudly?
##########
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:
When another thread already holds the invocation, `claimedInvocation` is
`false`, so `invokeListeners` returns immediately. The holder only drains
`onFailureCallbacks`, which means neither the pre-completion stop listener nor
any completion listener runs.
For example, a kill arrives and an interrupt listener fails.
`markInterrupted` then calls `markTaskFailed` on the kill thread, which runs
the failure listeners. If the task thread reaches `markTaskCompleted(None)`
meanwhile, it returns without running anything. `Task.run` finishes, and the
executor's `cleanUpAllAllocatedMemory` frees the pages while the pipelined
writer may still be copying rows (SPARK-33277). If the writer adds a listener
later, it takes over and waits for itself.
This predates the PR, but the PR now depends on the stop listener always
running on the completing thread. Should `markTaskCompleted` wait for the
current holder and then run the completion stacks, or should the holder also
drain the completion stacks once `completed` is set?
##########
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:
The claim in `markTaskCompleted` closes the takeover window at completion
start, but the failure path still opens one. When a listener throws, the catch
block (L370) sets `listenerInvocationThread = None` before `markTaskFailed(e)`
and only reclaims it in the `finally`.
Suppose a task has two pipelined writers, W1 and W2. W2's stop listener (or
any pre-completion listener registered after W1's) throws while W1 is still
running. If W1 then lazily adds a listener in that window, it sees `completed`
and an empty slot, becomes the invocation thread, and pops its own stop
listener. `stopAndAwaitExit` then waits on W1's own `exited` latch forever. The
task thread also keeps popping listeners at the same time, which breaks the
sequential guarantee.
Could the catch block keep this thread registered while it calls
`markTaskFailed` (e.g. let the failure invocation reuse the already-claimed
slot) instead of clearing it?
##########
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:
"regardless of when or from which thread that listener was added" claims
more than the implementation does. A pre-completion listener added after
completion has started runs only before the *remaining* completion listeners.
The new test `a pre-completion listener added during completion runs next`
asserts `Seq("B", "P", "A")`, so `B` runs before `P`.
Maybe: "before every pending listener added with
`addTaskCompletionListener`. Completion listeners that have already run are not
affected." That makes it clear a pre-completion listener must be registered
before completion starts to protect everything.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -128,6 +134,20 @@ private[spark] class TaskContextImpl(
this
}
+ override private[spark] def addTaskPreCompletionListener(
Review Comment:
nit: this copies `addTaskCompletionListener` with only the stack changed
(the comment even says "Same as `addTaskCompletionListener`"). A private helper
such as `addCompletionListenerTo(stack, listener)` used by both would keep the
add/invoke protocol in one place.
##########
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
+ // completion listener added here would run before them.
+ context.addTaskCompletionListener(makeLoggingListener("B"))
+
+ context.markTaskCompleted(None)
+ assert(invocationOrder === Seq("P2", "P1", "B", "A"))
+ // Called immediately once the task has completed, like completion
listeners.
+ context.addTaskPreCompletionListener(makeLoggingListener("P3"))
+ assert(invocationOrder === Seq("P2", "P1", "B", "A", "P3"))
+ }
+
+ test("SPARK-60113: completion listeners run even if a pre-completion
listener fails") {
+ val context = TaskContext.empty()
+ val listener = mock(classOf[TaskCompletionListener])
+ context.addTaskCompletionListener(listener)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ verify(listener, times(1)).onTaskCompletion(any())
+ assert(context.isFailed())
+ }
+
+ test("SPARK-60113: failures of pre-completion and completion listeners are
both reported") {
+ val context = TaskContext.empty()
+ context.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("post")
+ })
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ assert(e.getMessage.contains("post"))
+ }
+
+ test("SPARK-60113: a pre-completion listener added during completion runs
next") {
+ 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.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += "B"
+ // Added from a second thread while this thread runs the completion
listeners, so the
+ // call returns without running it; this thread must still run it.
+ val thread = new Thread(() => {
+ context.addTaskPreCompletionListener(makeLoggingListener("P"))
+ })
+ thread.start()
+ thread.join()
+ assert(invocationOrder === Seq("B"))
+ }
+ })
+
+ context.markTaskCompleted(None)
+ assert(invocationOrder === Seq("B", "P", "A"))
+ }
+
+ test("SPARK-60113: a listener added as completion starts does not take over
the listeners") {
Review Comment:
This depends on the busy-spin landing in the few-instruction window between
`completed = true` and the claim, without forcing that interleaving. If the
`markTaskCompleted` claim is reverted, the adder almost always arrives after
the task thread has registered itself, so I expect all 500 iterations to still
pass.
Could the test force the interleaving deterministically, for example with a
`TaskContextImpl` test subclass or hook that blocks between setting `completed`
and invoking the listeners? A case for the failure-path window (inline comment
on L322) would also help.
##########
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:
With `NonFatal(t)` widened to `Throwable`, a fatal error such as an OOM now
goes through `writer.setException` and `shutdownOutput` (both may allocate near
OOM) before it is rethrown. `PipelinedWriterTask` then logs it at ERROR, while
`ReaderIterator` rethrows the same error on the task thread, so one fatal error
is reported from two threads.
Is the `setException` path needed for fatal errors, or would rethrowing them
directly (as before) and letting `PipelinedWriterTask` log them be enough? If
the goal is to unblock the reader, a short note on the expected double report
would help.
##########
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:
`CountDownLatch.await` checks `Thread.interrupted()` before it checks the
count, so when the writer was claimed, a kill interrupt that is already set is
consumed here even if the writer has already exited. If the writer was never
claimed, the early `return` at L247 leaves the flag set.
As a result, the listeners that run next see the interrupt flag set or not
depending on a race. The comment above says the flag would make their blocking
channel and lock operations fail and skip cleanup, which is the case this is
trying to avoid.
If the intent is to never pass the interrupt on, `Thread.interrupted()`
could be called on the early-return path as well. Otherwise, could the comment
describe the actual behavior?
--
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]
