LuciferYang commented on code in PR #58518:
URL: https://github.com/apache/spark/pull/58518#discussion_r3955413682
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -298,15 +305,23 @@ object SQLExecution extends Logging {
event.duration = endTime - startTime
event.qe = queryExecution
event.executionFailure = ex
+ // Snapshot the `@volatile` `dagScheduler` once and share it
across both reads below.
+ // `SparkContext.stop()` nulls `dagScheduler` before it stops
the listener bus, so a
+ // query unwinding here while the context tears down would
otherwise NPE. As this runs
+ // in a `finally`, that NPE would replace the query's real
failure and skip the event
+ // post and observation completion below, so tolerate an
already-stopped context.
+ val dagSchedulerOpt = Option(sc.dagScheduler)
if (Utils.isTesting) {
import scala.jdk.CollectionConverters._
- event.jobIds =
Option(sc.dagScheduler.activeQueryToJobs.get(executionId))
+ // Only runs under `Utils.isTesting`; hits the same teardown
race as the cleanup.
+ event.jobIds = dagSchedulerOpt
+ .flatMap(ds => Option(ds.activeQueryToJobs.get(executionId)))
.map(_.asScala.map(_.jobId).toSet)
.getOrElse(Set.empty)
}
// Clean up jobs tracked by DAGScheduler for this query
execution.
- sc.dagScheduler.cleanupQueryJobs(executionId)
+ dagSchedulerOpt.foreach(_.cleanupQueryJobs(executionId))
Review Comment:
This round guards the two statements we know can throw; it does not
establish the invariant that nothing in the `finally` can leave a waiter hung.
The last statement, `observationManager.tryComplete` at :330, is still
unguarded, and it is the only one whose failure blocks a waiter:
`Observation.get` waits on `awaitResult(future, Duration.Inf)` and never
returns if the promise is not completed. The next statement someone adds
between :262 and :326 that can fail during teardown, which is exactly how the
shuffle cleanup got here, hangs that waiter again, and neither of the two tests
that assert the observation completes would catch it.
Moving `tryComplete` into its own `finally` makes it run whatever the block
above throws. `promise.tryComplete` is idempotent, so calling it on an
already-completed observation is a no-op.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -261,25 +261,32 @@ object SQLExecution extends Logging {
}
if (queryExecution.shuffleCleanupMode != DoNotCleanup
&& isExecutedPlanAvailable) {
- val shuffleIds = queryExecution.executedPlan match {
- case command: V2CommandExec =>
- command.children.flatMap(extractShuffleIds)
- case dataWritingCommand: DataWritingCommandExec =>
- extractShuffleIds(dataWritingCommand.child)
- case plan =>
- extractShuffleIds(plan)
- }
- shuffleIds.foreach { shuffleId =>
- queryExecution.shuffleCleanupMode match {
- case RemoveShuffleFiles =>
- // Same as what we do in
ContextCleaner.doCleanupShuffle, but do not
- // unregister the shuffle on MapOutputTracker, so that
stage retries would be
- // triggered.
- // Set blocking to Utils.isTesting to deflake unit tests.
- sc.shuffleDriverComponents.removeShuffle(shuffleId,
Utils.isTesting)
- case SkipMigration =>
-
SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId)
- case _ => // this should not happen
+ // Best-effort shuffle cleanup. This runs in a `finally` while
the query may be
+ // unwinding as the `SparkContext` is torn down:
`removeShuffle` reaches a stopped
+ // `BlockManagerMaster` and `SparkEnv.get` is null once
`SparkContext.stop()` has
+ // stopped `SparkEnv`, either of which would throw. As this
runs before the event
+ // post and observation completion below, an escaping failure
would replace the
+ // query's real exception and hang observation waiters, so log
and swallow it.
+ Utils.tryLogNonFatalError {
Review Comment:
`Utils.tryLogNonFatalError` logs `Uncaught exception in thread <name>`,
which names neither the shuffle cleanup nor the execution id nor the shuffle
id. When `askSync` itself throws, that line is the only trace that shuffle
files were left behind: the `logWarning` carrying the shuffle id inside
`removeShuffle` only fires when the removal future fails, which this path never
reaches. Bounding disk usage is the whole reason an operator turns
`spark.sql.classic.shuffleDependency.fileCleanup.enabled` on, and here the
query returns fine while the files stay on disk.
A local `catch { case NonFatal(e) => logWarning(...) }` carrying the
execution id would say which execution leaked the files, and `removeShuffle`
already logs the analogous leak at that level. The existing
`tryLogNonFatalError` uses under sql/core cover `outputStream.close()` and
cache-table cleanup, where a failure genuinely needs no follow-up.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala:
##########
@@ -423,6 +425,120 @@ class SQLExecutionSuite extends SparkFunSuite with
SQLConfHelper {
spark.stop()
}
}
+
+ /**
+ * Runs `f` with `spark`'s `dagScheduler` nulled out, standing in for a
`SparkContext` that has
+ * already been stopped. `SparkContext.stop()` nulls `_dagScheduler` before
it stops the listener
+ * bus, so a query really can unwind through `withNewExecutionId`'s
`finally` in this state.
+ */
+ private def withStoppedDagScheduler[T](spark: SparkSession)(f: => T): T = {
+ val sc = spark.sparkContext
+ val savedDagScheduler = sc.dagScheduler
+ sc.dagScheduler = null
+ try {
+ f
+ } finally {
+ sc.dagScheduler = savedDagScheduler
+ }
+ }
+
+ /**
+ * Runs `f` with `spark`'s `BlockManagerMaster.driverEndpoint` nulled out,
standing in for a
+ * `SparkContext` whose `SparkEnv` has been stopped. `SparkContext.stop()`
stops `SparkEnv` (which
+ * nulls that endpoint) after nulling `dagScheduler`, so the shuffle cleanup
in
+ * `withNewExecutionId`'s `finally` -- which runs before the `dagScheduler`
cleanup -- really can
+ * hit a stopped `BlockManagerMaster` and NPE while a query unwinds during
teardown.
+ */
+ private def withStoppedBlockManagerMaster[T](spark: SparkSession)(f: => T):
T = {
+ val master = spark.sparkContext.env.blockManager.master
+ val savedEndpoint = master.driverEndpoint
+ master.driverEndpoint = null
+ try {
+ f
+ } finally {
+ master.driverEndpoint = savedEndpoint
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId surfaces the body's failure when the
SparkContext " +
+ "has been stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ val qe = spark.range(1, 10).queryExecution
+ val bodyFailure = new IllegalStateException("body failed")
+ // Without the null guards, the `finally` dereferences `dagScheduler`
while it is null:
+ // under `Utils.isTesting` the `activeQueryToJobs` read throws first,
and the
+ // `cleanupQueryJobs` cleanup would do the same. Because the NPE is
thrown from a `finally`,
+ // it replaces `bodyFailure` entirely -- destroying the only record of
why the query failed.
+ val thrown = intercept[IllegalStateException] {
+ withStoppedDagScheduler(spark) {
+ SQLExecution.withNewExecutionId(qe) {
+ throw bodyFailure
+ }
+ }
+ }
+ assert(thrown eq bodyFailure)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId completes normally when the
SparkContext has been " +
+ "stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Attach an observation to pin down that it is completed. `tryComplete`
runs at the end of
+ // the `finally`, after the guarded cleanup, so a cleanup failure would
skip it and leave the
+ // observation uncompleted. Assert on `future.isCompleted`
(non-blocking) so a regression
+ // fails fast, rather than calling `get`, which would block until the
suite timeout.
+ val observation = new Observation("obs")
+ val df = spark.range(1, 10).observe(observation, count(lit(1)).as("cnt"))
+ val qe = df.queryExecution
+ withStoppedDagScheduler(spark) {
+ assert(SQLExecution.withNewExecutionId(qe)("result") === "result")
+ }
+ assert(observation.future.isCompleted)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId tolerates shuffle cleanup failing when
the " +
+ "SparkContext is stopping") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Disable AQE so the shuffle id is materialized from the plan below
without running a job;
+ // the `finally`'s shuffle cleanup then actually calls `removeShuffle`
for it.
+ spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, false)
+ val observation = new Observation("obs")
+ val df = spark.range(0, 4).repartition(2).observe(observation,
count(lit(1)).as("cnt"))
+ val qe = df.queryExecution
+ // Guard against a silent no-op: the cleanup only calls `removeShuffle`
if the plan yields a
+ // shuffle id, so assert there is one to clean up (accessing it also
materializes it).
+ val shuffleIds = qe.executedPlan.collect { case e: ShuffleExchangeLike
=> e.shuffleId }
+ assert(shuffleIds.nonEmpty)
Review Comment:
The `assert(shuffleIds.nonEmpty)` at :519 does not rule out the no-op run it
is written against. What decides whether the block at :270 runs is
`shuffleCleanupMode != DoNotCleanup` at :262, and that mode is fixed when the
`QueryExecution` is constructed, from
`spark.sql.classic.shuffleDependency.fileCleanup.enabled`, whose default is
`Utils.isTesting`. The test never asserts it, so "it can't silently pass
without reaching `removeShuffle`" does not follow: run this suite without
`spark.testing` / `SPARK_TESTING` set, which is what a single-test run from an
IDE does, and the mode falls back to `DoNotCleanup`, the whole block is
skipped, and both asserts still hold.
Pin that conf to `true` in a `withSQLConf` together with the AQE switch,
which my other comment asks to move there anyway, before `df` is built since
the mode is captured at `QueryExecution` construction, and add
`assert(qe.shuffleCleanupMode == RemoveShuffleFiles)`. No new import needed.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -261,25 +261,32 @@ object SQLExecution extends Logging {
}
if (queryExecution.shuffleCleanupMode != DoNotCleanup
&& isExecutedPlanAvailable) {
- val shuffleIds = queryExecution.executedPlan match {
- case command: V2CommandExec =>
- command.children.flatMap(extractShuffleIds)
- case dataWritingCommand: DataWritingCommandExec =>
- extractShuffleIds(dataWritingCommand.child)
- case plan =>
- extractShuffleIds(plan)
- }
- shuffleIds.foreach { shuffleId =>
- queryExecution.shuffleCleanupMode match {
- case RemoveShuffleFiles =>
- // Same as what we do in
ContextCleaner.doCleanupShuffle, but do not
- // unregister the shuffle on MapOutputTracker, so that
stage retries would be
- // triggered.
- // Set blocking to Utils.isTesting to deflake unit tests.
- sc.shuffleDriverComponents.removeShuffle(shuffleId,
Utils.isTesting)
- case SkipMigration =>
-
SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId)
- case _ => // this should not happen
+ // Best-effort shuffle cleanup. This runs in a `finally` while
the query may be
+ // unwinding as the `SparkContext` is torn down:
`removeShuffle` reaches a stopped
+ // `BlockManagerMaster` and `SparkEnv.get` is null once
`SparkContext.stop()` has
+ // stopped `SparkEnv`, either of which would throw. As this
runs before the event
+ // post and observation completion below, an escaping failure
would replace the
+ // query's real exception and hang observation waiters, so log
and swallow it.
+ Utils.tryLogNonFatalError {
+ val shuffleIds = queryExecution.executedPlan match {
+ case command: V2CommandExec =>
+ command.children.flatMap(extractShuffleIds)
+ case dataWritingCommand: DataWritingCommandExec =>
+ extractShuffleIds(dataWritingCommand.child)
+ case plan =>
+ extractShuffleIds(plan)
+ }
+ shuffleIds.foreach { shuffleId =>
Review Comment:
`tryLogNonFatalError` wraps the shuffle-id extraction plus the whole
`foreach`, while `removeShuffle` sends one `askSync` per shuffle id. So the
first id that throws ends the loop, and every remaining id never reaches
`removeShuffle` at all, leaving its files on disk. The loop already stopped at
the first throw before this PR, but the exception used to escape; now that the
block is best-effort, "the first failure abandons the rest, and nothing says
so" is its actual behavior.
Keep the outer guard's scope as it is, since `executedPlan` at :271
re-throws the failure its `LazyTry` cached; swap its handler as my other
comment asks, and add an inner one per id inside the `foreach`. That also lets
the log carry the shuffle id. The cost is that under a systemic failure such as
teardown the log goes from one line to one per id.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala:
##########
@@ -423,6 +425,120 @@ class SQLExecutionSuite extends SparkFunSuite with
SQLConfHelper {
spark.stop()
}
}
+
+ /**
+ * Runs `f` with `spark`'s `dagScheduler` nulled out, standing in for a
`SparkContext` that has
+ * already been stopped. `SparkContext.stop()` nulls `_dagScheduler` before
it stops the listener
+ * bus, so a query really can unwind through `withNewExecutionId`'s
`finally` in this state.
+ */
+ private def withStoppedDagScheduler[T](spark: SparkSession)(f: => T): T = {
+ val sc = spark.sparkContext
+ val savedDagScheduler = sc.dagScheduler
+ sc.dagScheduler = null
+ try {
+ f
+ } finally {
+ sc.dagScheduler = savedDagScheduler
+ }
+ }
+
+ /**
+ * Runs `f` with `spark`'s `BlockManagerMaster.driverEndpoint` nulled out,
standing in for a
+ * `SparkContext` whose `SparkEnv` has been stopped. `SparkContext.stop()`
stops `SparkEnv` (which
+ * nulls that endpoint) after nulling `dagScheduler`, so the shuffle cleanup
in
+ * `withNewExecutionId`'s `finally` -- which runs before the `dagScheduler`
cleanup -- really can
+ * hit a stopped `BlockManagerMaster` and NPE while a query unwinds during
teardown.
+ */
+ private def withStoppedBlockManagerMaster[T](spark: SparkSession)(f: => T):
T = {
+ val master = spark.sparkContext.env.blockManager.master
+ val savedEndpoint = master.driverEndpoint
+ master.driverEndpoint = null
+ try {
+ f
+ } finally {
+ master.driverEndpoint = savedEndpoint
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId surfaces the body's failure when the
SparkContext " +
+ "has been stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ val qe = spark.range(1, 10).queryExecution
+ val bodyFailure = new IllegalStateException("body failed")
+ // Without the null guards, the `finally` dereferences `dagScheduler`
while it is null:
+ // under `Utils.isTesting` the `activeQueryToJobs` read throws first,
and the
+ // `cleanupQueryJobs` cleanup would do the same. Because the NPE is
thrown from a `finally`,
+ // it replaces `bodyFailure` entirely -- destroying the only record of
why the query failed.
+ val thrown = intercept[IllegalStateException] {
+ withStoppedDagScheduler(spark) {
+ SQLExecution.withNewExecutionId(qe) {
+ throw bodyFailure
+ }
+ }
+ }
+ assert(thrown eq bodyFailure)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId completes normally when the
SparkContext has been " +
+ "stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Attach an observation to pin down that it is completed. `tryComplete`
runs at the end of
+ // the `finally`, after the guarded cleanup, so a cleanup failure would
skip it and leave the
+ // observation uncompleted. Assert on `future.isCompleted`
(non-blocking) so a regression
+ // fails fast, rather than calling `get`, which would block until the
suite timeout.
+ val observation = new Observation("obs")
+ val df = spark.range(1, 10).observe(observation, count(lit(1)).as("cnt"))
+ val qe = df.queryExecution
+ withStoppedDagScheduler(spark) {
+ assert(SQLExecution.withNewExecutionId(qe)("result") === "result")
+ }
+ assert(observation.future.isCompleted)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId tolerates shuffle cleanup failing when
the " +
+ "SparkContext is stopping") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Disable AQE so the shuffle id is materialized from the plan below
without running a job;
+ // the `finally`'s shuffle cleanup then actually calls `removeShuffle`
for it.
+ spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, false)
Review Comment:
The `spark.conf.set` at :512 turns AQE off and never restores it. Two tests
above, at :330 and :408, set config through `withSQLConf`, which this suite
already mixes in via `SQLConfHelper`. Nothing leaks today because the test
calls `spark.stop()` in its `finally`, but the session comes from
`getOrCreate()`, so the day these three tests share one session, AQE stays off
for whatever runs after them, showing up as changed behavior rather than a
failure.
Switch it to `withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
{ ... }`, and let the block span from `val observation` through the last
assert: AQE is read at planning time, so wrapping only the `intercept` leaves
`qe.executedPlan` as an `AdaptiveSparkPlanExec` and the assert at :519 fails.
The `fileCleanup` conf from my other comment belongs in the same block.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -261,25 +261,32 @@ object SQLExecution extends Logging {
}
if (queryExecution.shuffleCleanupMode != DoNotCleanup
&& isExecutedPlanAvailable) {
- val shuffleIds = queryExecution.executedPlan match {
- case command: V2CommandExec =>
- command.children.flatMap(extractShuffleIds)
- case dataWritingCommand: DataWritingCommandExec =>
- extractShuffleIds(dataWritingCommand.child)
- case plan =>
- extractShuffleIds(plan)
- }
- shuffleIds.foreach { shuffleId =>
- queryExecution.shuffleCleanupMode match {
- case RemoveShuffleFiles =>
- // Same as what we do in
ContextCleaner.doCleanupShuffle, but do not
- // unregister the shuffle on MapOutputTracker, so that
stage retries would be
- // triggered.
- // Set blocking to Utils.isTesting to deflake unit tests.
- sc.shuffleDriverComponents.removeShuffle(shuffleId,
Utils.isTesting)
- case SkipMigration =>
-
SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId)
- case _ => // this should not happen
+ // Best-effort shuffle cleanup. This runs in a `finally` while
the query may be
+ // unwinding as the `SparkContext` is torn down:
`removeShuffle` reaches a stopped
+ // `BlockManagerMaster` and `SparkEnv.get` is null once
`SparkContext.stop()` has
+ // stopped `SparkEnv`, either of which would throw. As this
runs before the event
+ // post and observation completion below, an escaping failure
would replace the
+ // query's real exception and hang observation waiters, so log
and swallow it.
+ Utils.tryLogNonFatalError {
+ val shuffleIds = queryExecution.executedPlan match {
+ case command: V2CommandExec =>
+ command.children.flatMap(extractShuffleIds)
+ case dataWritingCommand: DataWritingCommandExec =>
+ extractShuffleIds(dataWritingCommand.child)
+ case plan =>
+ extractShuffleIds(plan)
+ }
+ shuffleIds.foreach { shuffleId =>
+ queryExecution.shuffleCleanupMode match {
+ case RemoveShuffleFiles =>
+ // Same as ContextCleaner.doCleanupShuffle, but do not
unregister the
+ // shuffle on MapOutputTracker so that stage retries
would be triggered.
+ // Set blocking to Utils.isTesting to deflake unit
tests.
+ sc.shuffleDriverComponents.removeShuffle(shuffleId,
Utils.isTesting)
Review Comment:
I agree with the conclusion: letting a genuine interrupt propagate is right.
But "nothing in this path interrupts the driver thread anyway" does not hold,
so it should not go into the PR description as a justification. Spark Connect's
`ExecuteThreadRunner.interrupt()` (:111) interrupts exactly the thread running
`withNewExecutionId`, reachable from `InterruptRequest`, `ReleaseExecute`, and
session close. The state machine only CASes to `completed` after `handlePlan`
returns, so the whole `finally` sits inside the interruptible window, and
Connect itself turns this cleanup on (`SparkConnectPlanExecution.scala:61`).
Structured Streaming's `queryExecutionThread` is a second such path.
So when the interrupt lands while `removeShuffle` is parked on the `askSync`
reply, the `InterruptedException` passes through `tryLogNonFatalError` and
replaces the query's own exception. That was true before this PR too, so it is
not a regression and needs no code change here. Only the reasoning needs
restating.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -298,15 +305,23 @@ object SQLExecution extends Logging {
event.duration = endTime - startTime
event.qe = queryExecution
event.executionFailure = ex
+ // Snapshot the `@volatile` `dagScheduler` once and share it
across both reads below.
+ // `SparkContext.stop()` nulls `dagScheduler` before it stops
the listener bus, so a
+ // query unwinding here while the context tears down would
otherwise NPE. As this runs
+ // in a `finally`, that NPE would replace the query's real
failure and skip the event
+ // post and observation completion below, so tolerate an
already-stopped context.
+ val dagSchedulerOpt = Option(sc.dagScheduler)
Review Comment:
The same mechanism is spelled out over and over in this PR: six lines above
:270, five above :313, plus two helper scaladocs and three inline comments in
the test (:469, :490, :522). All of them say that the `finally` runs while the
`SparkContext` is being torn down, that an escaping failure replaces the
query's own exception, and that it skips the event post and the observation
completion. The change itself is ten lines.
Keeping the block above :270 and reducing the one above :313 to the half
that is not a repeat, that `SparkContext.stop()` nulls `dagScheduler` before it
stops the listener bus, would cover the production side. The test comments only
need to say why the simulated state really occurs. Separately, `the cleanup` at
:316 now has two candidates, the shuffle cleanup at :270 and the job cleanup at
:324, so it is worth naming which.
--
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]
