viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4238719125
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
Review Comment:
Good catch, confirmed. `DataSourceRDD.compute` creates the reader inside
`flatMap`, so `ParquetReaderCallback.initIfNotAlready` runs on the writer
thread and its `closeCurrent()` lands above the writer listener in the stack.
Added `TaskContext.addTaskPreCompletionListener` (`private[spark]`, default
falls back to `addTaskCompletionListener`). `TaskContextImpl.invokeListeners`
now always takes the next listener from the pre-completion stack first, so
these run before every completion listener, including ones added later or from
other threads, and also when a pre-completion listener is added while another
thread is already running the completion listeners. The writer listener uses
it. The test now adds a cleanup listener from the writer thread itself, and it
fails if the writer listener is switched back to an ordinary completion
listener.
As far as I can tell, the `WriterThread` listener before SPARK-44705 had the
same gap, since it was also registered before the writer thread started pulling
upstream.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerTask.stopAndAwaitExit(writerFuture, taskIdentifier(context))
Review Comment:
Fixed in two places. `PipelinedWriterRunnable` now also calls
`writer.setException(t)` and `shutdownOutput()` for fatal errors before
rethrowing them, so the reader fails with the cause instead of waiting for
Python. `PipelinedWriterTask.run` logs anything that escapes the writer. A
`future.get()` after the latch would not be enough here: once `cancel(true)`
has succeeded, the future's outcome is no longer available.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerTask.stopAndAwaitExit(writerFuture, taskIdentifier(context))
+ }
+ }
+
+ /**
+ * Wraps a pipelined writer so that a task completion listener can wait
until the writer has
+ * either exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable) 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()
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_], taskName: String): Unit = {
+ // Interrupts the writer thread if the writer is running. This unblocks
channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException),
and the writer
+ // loop checks the interrupt flag between rows, so the wait below is
normally short.
Review Comment:
Fixed the comment: it now says that the interrupt is seen only between
`writeNextInputToStream` calls, each of which can pull and serialize a whole
batch. Will fix the PR description as well.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerTask.stopAndAwaitExit(writerFuture, taskIdentifier(context))
+ }
+ }
+
+ /**
+ * Wraps a pipelined writer so that a task completion listener can wait
until the writer has
+ * either exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable) 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()
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_], taskName: String): Unit = {
+ // Interrupts the writer thread if the writer is running. This unblocks
channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException),
and the writer
+ // loop checks the interrupt flag between rows, so the wait below is
normally short.
+ future.cancel(true)
+ if (claimed.compareAndSet(false, true)) {
+ // The writer never started and now never will.
+ return
+ }
+ // Wait without a time bound: giving up would let the listeners that run
next free memory
+ // the writer may still be reading, and failing this listener would not
stop them either.
+ // A killed task that stays stuck here is handled by the task reaper, if
enabled.
Review Comment:
The trade-off is real, and I documented it in the comment. I still kept the
wait unbounded, because in these examples the wait does buy safety. In
`BatchEvalPython` / `ArrowEvalPython`, the row the writer gets next from the
blocked upstream read goes to `queue.add` before the loop checks the interrupt
again. `InMemoryRowQueue` writes through the page's base object that it
captured at construction. So even for on-heap memory, a freed page (at least
the page size, which `HeapMemoryAllocator` pools and hands to other consumers)
would be silently corrupted, and off-heap it would crash. I couldn't find a
reliable way to tell when giving up would be safe. The `WriterThread.join()`
before SPARK-44705 was also unbounded, for every Python UDF.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerTask.stopAndAwaitExit(writerFuture, taskIdentifier(context))
+ }
+ }
+
+ /**
+ * Wraps a pipelined writer so that a task completion listener can wait
until the writer has
+ * either exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable) 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()
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_], taskName: String): Unit = {
+ // Interrupts the writer thread if the writer is running. This unblocks
channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException),
and the writer
+ // loop checks the interrupt flag between rows, so the wait below is
normally short.
+ future.cancel(true)
+ if (claimed.compareAndSet(false, true)) {
+ // The writer never started and now never will.
+ return
+ }
+ // Wait without a time bound: giving up would let the listeners that run
next free memory
+ // the writer may still be reading, and failing this listener would not
stop them either.
+ // A killed task that stays stuck here is handled by the task reaper, if
enabled.
+ val startNs = System.nanoTime()
Review Comment:
Following your next comment, the wait no longer restores the interrupt
status, while `Uninterruptibles.awaitUninterruptibly` always restores it. So I
kept the loop, which is now just a catch-and-continue without the outer
`try`/`finally`.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,86 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while a task completion listener waits for a pipelined
writer to exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Listeners registered earlier run later (LIFO) and free memory backing the
input rows, e.g.
+ * the off-heap pages of a HybridRowQueue or of an off-heap vectorized
reader. If the writer
+ * were still serializing or buffering such a row at that point, it would be
a use-after-free
+ * that can crash the executor (SPARK-33277). `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)
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskCompletionListener[Unit] { _ =>
+ writerTask.stopAndAwaitExit(writerFuture, taskIdentifier(context))
+ }
+ }
+
+ /**
+ * Wraps a pipelined writer so that a task completion listener can wait
until the writer has
+ * either exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(writer: Runnable) 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()
+ } finally {
+ exited.countDown()
+ }
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_], taskName: String): Unit = {
+ // Interrupts the writer thread if the writer is running. This unblocks
channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException),
and the writer
+ // loop checks the interrupt flag between rows, so the wait below is
normally short.
+ future.cancel(true)
+ if (claimed.compareAndSet(false, true)) {
+ // The writer never started and now never will.
+ return
+ }
+ // Wait without a time bound: giving up would let the listeners that run
next free memory
+ // the writer may still be reading, and failing this listener would not
stop them either.
+ // A killed task that stays stuck here is handled by the task reaper, if
enabled.
+ val startNs = System.nanoTime()
+ var interrupted = false
+ try {
+ var done = false
+ while (!done) {
+ try {
+ done = exited.await(pipelinedWriterExitWarnIntervalMs,
TimeUnit.MILLISECONDS)
+ if (!done) {
+ val waitedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
startNs)
+ logWarning(log"Still waiting for the pipelined Python writer of
" +
+ log"${MDC(TASK_NAME, taskName)} to exit after " +
+ log"${MDC(TOTAL_TIME, waitedMs)} ms; task completion cleanup
is blocked " +
+ log"until it does.")
+ }
+ } catch {
+ case _: InterruptedException => interrupted = true
+ }
+ }
+ } finally {
+ if (interrupted) {
Review Comment:
Not intended. Changed it to not restore the flag, with a comment: the kill
is already recorded in the `TaskContext`, and the old listener also swallowed
the interrupt (`get()` threw `InterruptedException`). The test now asserts that
a later completion listener does not see the flag.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +50,105 @@ class BasePythonRunnerSuite extends SparkFunSuite {
}
}
+ /**
+ * A writer that, once started, keeps running until `release` is counted
down, ignoring
+ * interrupts the way a writer in the middle of copying a row into a
HybridRowQueue page does.
+ */
+ private class UninterruptibleWriter extends Runnable {
+ val started = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ @volatile var exited = false
+
+ override def run(): Unit = {
+ started.countDown()
+ try {
+ var done = false
+ while (!done) {
+ try {
+ done = release.await(1, TimeUnit.MINUTES)
+ } catch {
+ case _: InterruptedException => // keep going, like a row copy in
progress
+ }
+ }
+ } finally {
+ exited = true
+ }
+ }
+ }
+
+ test("SPARK-60113: task completion waits for the pipelined writer to exit") {
+ val context = TaskContext.empty()
+ val writer = new UninterruptibleWriter
+ // Registered before the writer, so it runs after the writer's listener
(LIFO), like the
+ // listener that closes the HybridRowQueue and frees its pages.
+ @volatile var exitedBeforeCleanup: Option[Boolean] = None
+ context.addTaskCompletionListener[Unit] { _ => exitedBeforeCleanup =
Some(writer.exited) }
+ BasePythonRunner.startPipelinedWriter(writer, context)
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+
+ val completion = new Thread(() => context.markTaskCompleted(None))
+ completion.setDaemon(true)
+ completion.start()
+ try {
+ // Task completion must stay blocked while the writer is still running.
+ completion.join(500)
Review Comment:
Fixed. The tests add a pre-completion listener after `startPipelinedWriter`,
which runs right before the writer listener and counts down a latch. The tests
wait for that latch before checking `isAlive`, and the interrupt test also
interrupts only after it.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +50,105 @@ class BasePythonRunnerSuite extends SparkFunSuite {
}
}
+ /**
+ * A writer that, once started, keeps running until `release` is counted
down, ignoring
+ * interrupts the way a writer in the middle of copying a row into a
HybridRowQueue page does.
+ */
+ private class UninterruptibleWriter extends Runnable {
+ val started = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ @volatile var exited = false
+
+ override def run(): Unit = {
+ started.countDown()
+ try {
+ var done = false
+ while (!done) {
+ try {
+ done = release.await(1, TimeUnit.MINUTES)
+ } catch {
+ case _: InterruptedException => // keep going, like a row copy in
progress
+ }
+ }
+ } finally {
+ exited = true
+ }
+ }
+ }
+
+ test("SPARK-60113: task completion waits for the pipelined writer to exit") {
+ val context = TaskContext.empty()
+ val writer = new UninterruptibleWriter
+ // Registered before the writer, so it runs after the writer's listener
(LIFO), like the
+ // listener that closes the HybridRowQueue and frees its pages.
+ @volatile var exitedBeforeCleanup: Option[Boolean] = None
+ context.addTaskCompletionListener[Unit] { _ => exitedBeforeCleanup =
Some(writer.exited) }
+ BasePythonRunner.startPipelinedWriter(writer, context)
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+
+ val completion = new Thread(() => context.markTaskCompleted(None))
+ completion.setDaemon(true)
+ completion.start()
+ try {
+ // Task completion must stay blocked while the writer is still running.
+ completion.join(500)
+ assert(completion.isAlive)
+ assert(exitedBeforeCleanup.isEmpty)
+ } finally {
+ writer.release.countDown()
+ }
+ completion.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!completion.isAlive)
+ assert(exitedBeforeCleanup === Some(true))
+ }
+
+ test("SPARK-60113: task completion waits for the pipelined writer even if
interrupted") {
+ val context = TaskContext.empty()
+ val writer = new UninterruptibleWriter
+ @volatile var exitedBeforeCleanup: Option[Boolean] = None
+ context.addTaskCompletionListener[Unit] { _ => exitedBeforeCleanup =
Some(writer.exited) }
+ BasePythonRunner.startPipelinedWriter(writer, context)
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+
+ @volatile var interruptedAfterCompletion = false
+ val completion = new Thread(() => {
+ context.markTaskCompleted(None)
+ interruptedAfterCompletion = Thread.currentThread().isInterrupted
+ })
+ completion.setDaemon(true)
+ completion.start()
+ try {
+ completion.join(200)
+ // E.g. a task kill interrupting the task thread while it runs the
listeners.
+ completion.interrupt()
+ completion.join(500)
+ assert(completion.isAlive)
+ assert(exitedBeforeCleanup.isEmpty)
+ } finally {
+ writer.release.countDown()
+ }
+ completion.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!completion.isAlive)
+ assert(exitedBeforeCleanup === Some(true))
+ // The interrupt is preserved for the code that runs after the listeners.
+ assert(interruptedAfterCompletion)
+ }
+
+ test("SPARK-60113: a pipelined writer stopped before it starts never runs") {
+ val writer = new UninterruptibleWriter
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(writer)
+ // A future the pool has not started yet: stopping must not wait for it.
+ val future = new FutureTask[Unit](writerTask, ())
Review Comment:
Added "stopping waits for a writer that started after the interrupt": the
writer is claimed by `run()` on a thread whose interrupt flag is already set,
and `stopAndAwaitExit` must block until it exits.
--
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]
