HyukjinKwon commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4255354335
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +196,110 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, Claim)
+ }
+
+ /**
+ * @param mode `Nested` if this thread is already invoking listeners, i.e. a
listener failed. The
+ * failure listeners then run as part of that invocation, which
this thread keeps.
+ */
+ private def markTaskFailed(error: Throwable, mode: InvocationMode): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, mode)
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ // The completion listeners always run on this thread, and this method
returns only after they
+ // have run. If any other thread could run them, it could 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. And if this method
returned before they
+ // ran, the task could release its memory while that writer still uses it.
So the listener
+ // invocation is claimed together with setting `completed`; if another
thread holds it (e.g. a
+ // kill thread running failure listeners), that thread hands it over to
this one when done.
+ //
+ // That wait has no time bound: if the other thread's listeners block, so
does this task. It
+ // logs a warning with that thread's stack trace periodically. If such a
task is killed,
+ // `spark.task.reaper.enabled` with a positive
`spark.task.reaper.killTimeout` bounds how long
Review Comment:
**Nit (P3):** This bound does not hold for the example holder named above, a
kill thread running failure listeners. With `spark.task.reaper.enabled`,
`Executor.killTask` runs the kill on a `TaskReaper`, and `TaskReaper.run` calls
`taskRunner.kill()` synchronously before it enters its `killTimeout` loop.
`kill()` goes through `Task.kill` to `markInterrupted`, and if an interrupt
listener throws, `markTaskFailed` runs the failure listeners on that same
reaper thread. If one of them blocks, `kill()` never returns,
`timeoutExceeded()` is never checked, and this thread stays in
`awaitListenerInvocationHandover` with no JVM kill. The bound does hold when
the holder is some other thread, and for the writer-exit wait in
`startPipelinedWriter`, where the task thread holds the invocation. Could the
comment say that the reaper bounds this wait only when the blocked listeners do
not run inside the reaper's own kill call, and that otherwise only the periodic
warning shows it?
**Verification:**
- **Inspection:** The comment's recovery statement agrees with the
TaskReaper sequence, where kill() runs synchronously before the killTimeout
loop, and with markInterrupted running failure listeners on the kill thread
when an interrupt listener throws.
##########
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, I checked 7ab7ae2e6af. `stopAndAwaitExit` no longer waits when it
runs on the thread recorded in `PipelinedWriterRunning`, and that branch now
logs at ERROR and clears the self-interrupt; "stopping on the writer thread
itself does not wait for the writer" exercises it through a real `FutureTask`.
The handoff that led there is closed at the source too: `markTaskCompleted` now
waits for the holder and takes the invocation over directly, and a failing
listener keeps the invocation while the failure listeners run, so the guard is
only a safety net now.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4246906867","thread_id":"inline:4246906867","verdict_sha256":"698cadbb5ec60306a643aac1495c6da95c02eb7786ffec3bc65dd5edfc06c375"}
-->
##########
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:
Thanks, verified. The new test drives a real `PipelinedWriterRunnable`
against a socket pair and checks both `writer.exception` and the EOF on the
Python side. Narrowing the branch back to `NonFatal` would leave
`writer.exception` empty, and dropping the rethrow would fail the
`intercept[LinkageError]`, so both regressions are caught.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4246906877","thread_id":"inline:4246906877","verdict_sha256":"698cadbb5ec60306a643aac1495c6da95c02eb7786ffec3bc65dd5edfc06c375"}
-->
##########
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, the comment now matches the asserted order.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4246906884","thread_id":"inline:4246906884","verdict_sha256":"698cadbb5ec60306a643aac1495c6da95c02eb7786ffec3bc65dd5edfc06c375"}
-->
--
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]
