[CELEBORN-2465][SPARK] Proactively unregister shuffles on SQL query completion

[yew1eb]

What changes were proposed in this pull request?

Add an opt-in configuration celeborn.client.spark.shuffleCleanup.enabled (default false). When enabled, a SparkListener registered on the SparkContext listener bus handles SparkListenerSQLExecutionEnd events and proactively unregisters all shuffles written by the completed SQL query (or Dataset action), instead of waiting for the driver GC to trigger Spark's ContextCleaner or for the application to end.

• New CelebornSQLExecutionEndCleaner listener collects shuffle ids from the query's final executed plan (AQE-aware) and cleans them up through the standard Spark cleanup entry point shuffleDriverComponents.removeShuffle, reusing the existing unregister pipeline with zero master/worker/protocol changes.
• New LifecycleManager.isAppShuffleRegistered guard so only actually registered shuffles are unregistered.

Why are the changes needed?

In long-running Spark applications (ThriftServer, Kyuubi, spark-sql, spark-connect-server), one application runs many SQL queries, each producing new shuffles. Shuffle unregister is only triggered by Spark's GC-driven ContextCleaner, which may be delayed indefinitely in long sessions due to DAG/query-plan reference chains (see CELEBORN-2244), so obsolete shuffle data piles up on Celeborn worker disks for the whole session lifetime.

Does this PR resolve a correctness bug?

No.

Does this PR introduce any user-facing change?

Yes, a new client configuration celeborn.client.spark.shuffleCleanup.enabled (default false, opt-in).

How was this patch tested?

add ut CelebornQueryEndShuffleCleanupSuite

[s0nskar]

cc: @saurabhd336

[venkata91]

Thanks for adding this feature. This is a very nice improvement to keep the shuffle sizes bounded proactively. I would also add spark-connect-server to the above list.

[venkata91]

Thanks for working on this. Bounding shuffle storage for long-running sessions is valuable. I ran the new suite plus a few extra scenarios locally (Spark 3.5 and 4.0) and found two blockers and some smaller items:

1. Spark 4 + AQE is a no-op (likely cause of the spark4 CI failures). Spark 4.0 wraps the final AQE plan in ResultQueryStageExec, a leaf QueryStageExec, so collect never reaches the exchanges. With -Pspark-4.0 I get 3/6 failures in CelebornQueryEndShuffleCleanupSuite (all the AQE cases). Matching case stage: QueryStageExec => extractShuffleIds(stage.plan) makes all 6 pass locally.

2. Silent empty results when the same DataFrame is executed again with celeborn.client.spark.stageRerun.enabled=false. CelebornShuffleDriverComponents doesn't override removeShuffle, so it goes through LocalDiskShuffleDriverComponents → BlockManagerMaster.removeShuffle → RemoveShuffle to every BlockManager, including the driver's. That also unregisters the shuffle from the driver's MapOutputTrackerMaster. Because the ShuffleDependency is still reachable from the DataFrame's executed plan, the same shuffle id is reused on the next execution. Combined with the LifecycleManager's delayed expiry, Spark and Celeborn can disagree on whether the shuffle exists. Celeborn treats unknown shuffles as empty (SHUFFLE_UNREGISTERED / UNKNOWN_APP_SHUFFLE_ID), so the result is wrong rather than an error. Repro (Spark 3.5, AQE off, stageRerun=false, cleanFailedShuffle=false, shuffle.expired.checkInterval=1s):

   val df = spark.sql("SELECT id % 10 AS k, COUNT(1) AS cnt FROM ta GROUP BY id % 10")
   df.collect()  // correct
   // wait until LifecycleManager expires the shuffle
   df.collect()  // correct
   df.collect()  // correct
   df.collect()  // returns an EMPTY result, no exception

   Two threads calling collect() on the same DataFrame concurrently also fail (CelebornIOException: createPartitionReader failed) with stageRerun disabled. With stageRerun enabled, all my sequential and concurrent runs were correct. Could we fail fast at init unless stageRerun.enabled is true, as fetch.cleanFailedShuffle already does, rather than only recommending it in the doc? Even with stageRerun, a reader that asks for the shuffle id after the mapping was removed gets UNKNOWN_APP_SHUFFLE_ID with success=true, which is read as empty. I couldn't trigger that race, but it may be worth a guard.

3. Scope question. Spark 4's built-in equivalent (SPARK-47764, spark.sql.shuffleDependency.fileCleanup.enabled) is scoped to Spark Connect / Thrift executions, which re-plan every statement, so DataFrame re-use can't happen. This PR applies to all SQL executions, including classic Dataset/PySpark patterns like df.toPandas(); df.collect(). Would a narrower scope (or a per-session switch) make sense? On Spark 4+ it might also be possible to rely on Spark's own cleanup, which calls shuffleManager.unregisterShuffle.

Happy to share the repro suite if useful. A few inline comments follow.

[venkata91]

On Spark 4.x the final AQE plan is wrapped in ResultQueryStageExec (also BroadcastQueryStageExec here), both leaf QueryStageExecs, so collect doesn't descend into them. Matching case stage: QueryStageExec => extractShuffleIds(stage.plan) fixed all AQE tests for me on -Pspark-4.0. TableCacheQueryStageExec.plan is an InMemoryTableScanExec, so cached plans still won't be touched.

[venkata91]

ShuffleExchangeExec.shuffleDependency is a lazy val. For exchanges that never materialized (e.g. an AQE query that failed or was cancelled before creating that stage), touching it runs child.execute(), allocates a new shuffle id and calls shuffleManager.registerShuffle, all on the listener-bus thread. Could we only collect ids from materialized stages (e.g. ShuffleQueryStageExec with isMaterialized)?

Also: collectWithSubqueries would cover scalar-subquery shuffles, and non-vanilla ShuffleExchangeLike implementations (e.g. Gluten's columnar exchange) aren't handled.

[venkata91]

If planning itself failed, qe.executedPlan will re-trigger planning here on the listener thread and throw again (Spark 4 guards this with isExecutedPlanAvailable). Nit: could the LifecycleManager be passed into the listener's constructor instead of SparkEnv.get.shuffleManager.asInstanceOf[...] on every event? Also maybe NonFatal instead of Throwable at L37.

[venkata91]

Nit: the PR description says celeborn.client.spark.sql.queryEndShuffleCleanup.enabled, but the key here is celeborn.client.spark.queryEndShuffleCleanup.enabled, while the val is CLIENT_SPARK_SQL_.... Please align them. Also see the top-level comment about enforcing stageRerun.enabled instead of only recommending it.

[venkata91]

Could we add cases for re-executing the same DataFrame (sequentially and concurrently, with stageRerun on/off) and for a failed AQE query? Nit: spark.stop() at the end of each test is redundant with afterEach.

[yew1eb]

@venkata91 Thanks for the thorough review and the repros — all findings confirmed and addressed in the latest commit. Summary of the changes:

Blocker 1 — Spark 4 + AQE no-op. Fixed. The extractor now matches generic QueryStageExec and descends into stage.plan, so Spark 4's ResultQueryStageExec leaf is reached; TableCacheQueryStageExec (plan = InMemoryTableScanExec) contains no shuffles and is unaffected. The suite now passes 10/10 locally on both Spark 3.5 and Spark 4.0.

Blocker 2 — silent empty results with stageRerun.enabled=false. Init now fails fast when cleanup is enabled but stage rerun is not, following the existing fetch.cleanFailedShuffle precedent. Your repro scenarios are covered by new tests: sequential and concurrent re-execution of the same DataFrame, both returning correct results with stage rerun on. The UNKNOWN_APP_SHUFFLE_ID race you mentioned is left as a follow-up since neither of us could trigger it.

Lazy shuffleDependency. Extraction is now side-effect free on the listener thread:

• Shuffle ids are read only from materialized ShuffleQueryStageExec (dependency already computed); raw exchanges inside an AdaptiveSparkPlanExec subtree are never touched.
• For failed non-AQE queries, exchanges with RangePartitioning output are skipped (reading them could launch a sampling job); the isAppShuffleRegistered filter drops anything never actually registered.
• Two subtleties surfaced while testing the multi-join case: AQE nests earlier stages inside later stages' plan subtrees, so traversal must continue below stage/exchange boundaries; and identical inputs (self-joins) canonicalize multiple exchanges into one.
• Scalar subqueries are covered via subqueriesAll. Non-vanilla ShuffleExchangeLike implementations (e.g. Gluten) remain intentionally unsupported to avoid coupling to third-party exchange internals.

Planning failures / nits. Failed queries are detected via SparkListenerSQLExecutionEnd.executionFailure; the residual re-planning risk of qe.executedPlan is contained by the NonFatal catch (no public isExecutedPlanAvailable exists on Spark 3.5/4.0–4.2). LifecycleManager is now constructor-injected, and Throwable was narrowed to NonFatal.

Config key. Aligned to celeborn.client.spark.sql.queryEndShuffleCleanup.enabled; docs regenerated via ConfigurationSuite. spark-connect-server added to the doc and PR description.

Tests. Added: failed AQE query, sequential/concurrent DataFrame re-execution, and the stage-rerun fail-fast case; removed the redundant spark.stop() calls.

Re: scope and SPARK-47764. Confirmed it shipped in 4.0.0 scoped to Spark Connect (shuffleCleanupMode defaults to DoNotCleanup); the classic and ThriftServer variants arrived only in 4.1 and 4.2 respectively, all opt-in. So this feature still covers ground Spark doesn't: all of Spark 3.x and classic/Thrift on 4.0. Both implementations share the same shuffleDriverComponents.removeShuffle entry point, so they coexist idempotently if a user enables both. That said, narrowing this listener's scope on Spark 4.1+/4.2+ in favor of the built-in cleanup could be a reasonable follow-up once those versions are the common denominator.

[venkata91]

@yew1eb thanks for the quick turnaround. I re-ran last round's repros against bd7f9b3 on Spark 3.5 and 4.0. The Spark 4 AQE fix works, and the sequential and concurrent re-execution cases are correct with stageRerun on.

I also added a few cases around the new traversal. Each finding below reproduces on both 3.5 and 4.0, and details are inline:

1. A failed non-AQE query can still launch Spark jobs on the listener bus thread.
2. With AQE, some shuffles are not cleaned at query end: scalar subqueries, finished child stages of a failed stage, and stages AQE drops during re-optimization. They fall back to GC cleanup, so this is a leak, not wrong results.
3. The stageRerun check only fails the first query.

Two other things:

• The PR description still mentions CelebornSQLExecutionEndCleaner and CelebornQueryEndShuffleCleanupSuite (renamed in 2096b7d). It could also mention that stageRerun.enabled is now required.
• Could we file a follow-up JIRA for the deferred UNKNOWN_APP_SHUFFLE_ID race and link it here?

I have a local repro suite and a prototype patch if either would help.

[venkata91]

In bd7f9b3, this still reads shuffleDependency for a failed non-AQE query when the outer exchange isn't RangePartitioning. That lazy val runs child.execute(), so a nested range exchange, subquery or broadcast underneath can still start a job.

Repro (AQE off): bad.union(good).collect(), where each side is repartitionByRange(4, x).groupBy(x % 10).count() and bad has a throwing UDF. After the query fails, the driver log shows Got job 2 (onOtherEvent at SparkListenerBus.scala:100), coming from CelebornShuffleCleanupListener.scala:103 -> shuffleDependency$lzycompute -> RangePartitioner.sketch. The exception then aborts the walk, so good's shuffle isn't cleaned either.

Some options:
(a) Skip raw exchanges when executionFailure is set and leave those shuffles to GC cleanup. This is simplest, and it's what Spark's own query-end cleanup effectively does, since it only uses AQE stage info.
(b) Only read the id when the exchange has already executed, e.g. by checking whether mapOutputStatisticsFuture / cachedShuffleRDD was initialized. This relies on Spark internals that differ across versions.
(c) Catch per exchange so that one failure doesn't stop the rest of the walk. This contains the damage, but the job still runs.

[venkata91]

With AQE on, scalar-subquery shuffles aren't found. The root is an AdaptiveSparkPlanExec, which is a leaf, so subqueriesAll is empty. The subqueries live in expressions of nodes inside executedPlan, and visit only follows children.

Repro (AQE on): SELECT COUNT(*) FROM ta WHERE id > (SELECT MAX(c) FROM (SELECT id % 10 k, COUNT(1) c FROM ta GROUP BY id % 10)) cleans only the main-query shuffle. The two subquery shuffles stay registered.

Some options:
(a) Also visit p.subqueries for each node inside visit. That covers both AQE and non-AQE, and the top-level subqueriesAll would no longer be needed.
(b) Read stages from AdaptiveSparkPlanExec.context.stageCache. AQE subqueries share the context, so their stages are in it too (see the next comment).

[venkata91]

Two cases with AQE are missed here:

• If a stage isn't materialized, we don't descend into it, so its materialized children are skipped. For example, if the UDF fails in the second map stage, the first stage's shuffle stays registered.
• Stages that AQE drops from the final plan are never visited. For example, ta JOIN (SELECT * FROM tb WHERE id < 0) b ON ... with broadcast off leaves both shuffles registered after the query ends.

Some options:
(a) Descend into the stage regardless of whether it is materialized, since each stage below is checked on its own. This fixes the first case only.
(b) Also collect the materialized ShuffleQueryStageExecs from adaptivePlan.context.stageCache. That holds every stage the query created; I checked it exists in 3.0 through 4.x. This covers both cases and is close to Spark 4's context.shuffleIds.

Minor: the inAdaptivePlan = false reset on L88 doesn't seem needed. Keeping true would guarantee that no raw exchange inside an AQE subtree is ever touched.

[venkata91]

This check runs lazily on the first registerShuffle, after lifecycleManager has already been assigned. With stageRerun=false, the first query fails, but later queries run normally with cleanup silently off. The cleanFailedShuffle and broadcast callbacks below are also never registered.

Some options:
(a) Move the check to the constructor (driver only, e.g. next to validateAttemptConfig), so the SparkContext fails to start.
(b) Keep it here, but validate before lifecycleManager is assigned, so every query fails consistently.
(c) Log a warning and disable cleanup instead of throwing.

(a) seems the most predictable to me. The fail-fast test would then intercept around getOrCreate().

[venkata91]

Nit: now that stageRerun is required whenever this is used, hasMapping is always true. We could either drop the parameter and check shuffleIdMapping only, or keep it if you expect other callers.

[venkata91]

Nit: shuffleCleanup.enabled could be read as general shuffle cleanup, which already happens via ContextCleaner. Some options: keep it as is, or use a name that says when it fires, e.g. celeborn.client.spark.sql.cleanShuffleOnQueryEnd.enabled. No strong opinion.

Separately, it might help to note the limits in the doc: it only applies to SQL/Dataset executions, non-vanilla exchanges such as Gluten aren't covered, and shuffles of failed or cancelled stages fall back to GC cleanup.