ivoson commented on code in PR #58518:
URL: https://github.com/apache/spark/pull/58518#discussion_r3965947211
##########
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:
Done. `observationManager.tryComplete` now runs in its own inner `finally`
wrapping the end-event post and the shuffle/job cleanup, so it completes the
observation no matter what that block throws (`tryComplete` is idempotent). The
two teardown tests assert completion via `future.isCompleted`, so a regression
fails fast instead of blocking a waiter to the suite timeout.
##########
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:
Done. The AQE switch and
`spark.sql.classic.shuffleDependency.fileCleanup.enabled -> true` are now in
one `withSQLConf` before the `QueryExecution` is built, plus
`assert(qe.shuffleCleanupMode == RemoveShuffleFiles)`, so the cleanup block is
reached regardless of `SPARK_TESTING`.
##########
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:
Done. Replaced the outer `tryLogNonFatalError` with local `catch { case
NonFatal(e) => logWarning(...) }` handlers carrying the execution id and
shuffle id, so a leak names which execution/shuffle left files on disk --
matching the level `removeShuffle` already logs at.
##########
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:
Done. Kept the outer guard (so `executedPlan`'s cached `LazyTry` re-throw is
still handled) and added a per-id `catch NonFatal` inside the `foreach`, so one
shuffle's failure no longer abandons the rest and each warning carries its
shuffle id. As you note, under a systemic failure the log now goes to one line
per id.
##########
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:
You're right -- thanks for the correction. I've dropped that claim:
Connect's `ExecuteThreadRunner.interrupt()` and Structured Streaming's
`queryExecutionThread` both interrupt the thread running `withNewExecutionId`,
so the `finally` sits in the interruptible window. The description and cleanup
scaladoc now state only that letting a genuine `InterruptedException` propagate
is intentional (it isn't `NonFatal`) and that this is pre-existing behavior,
not a regression -- no code change.
--
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]
