From ce3866073cf20d8fbb9ac5581956cfaf16d6efbd Mon Sep 17 00:00:00 2001 From: saurabhdamle11 Date: Thu, 24 Sep 2026 13:53:31 -0700 Subject: [PATCH 1/3] [SPARK-59776][CORE][SQL] Release shuffle-tracked executors when SQL shuffle file cleanup removes shuffle files When SQL shuffle dependency file cleanup is enabled, SQLExecution removes shuffle files via ShuffleDriverComponents.removeShuffle directly, bypassing ContextCleaner, so the dynamic allocation ExecutorMonitor is never told and keeps executors alive for those shuffles (forever with the default infinite shuffle tracking timeout). Add CleanerListener.shuffleFilesRemoved (no-op by default) and ContextCleaner.notifyShuffleFilesRemoved, call it after a successful removal in SQLExecution, and have ExecutorMonitor stop tracking the shuffle. A separate callback is used instead of shuffleCleaned because the shuffle intentionally stays registered on the MapOutputTracker. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../org/apache/spark/ContextCleaner.scala | 20 ++++++++ .../scheduler/dynalloc/ExecutorMonitor.scala | 3 ++ .../dynalloc/ExecutorMonitorSuite.scala | 23 ++++++++++ .../spark/sql/execution/SQLExecution.scala | 3 ++ .../sql/execution/SQLExecutionSuite.scala | 46 ++++++++++++++++++- 5 files changed, 93 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 52d8507fb953a..d31b6fc8941b8 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -306,6 +306,21 @@ private[spark] class ContextCleaner( mgr != null && !mgr.usesStreamingShuffleOutputTracker && mgr.holdsShuffle(shuffleId) } + /** + * Notify listeners that the files of `shuffleId` were removed outside of this cleaner (e.g. by + * SQL shuffle dependency cleanup) while the shuffle stays registered on the MapOutputTracker. + * Unlike `doCleanupShuffle`, this fires `shuffleFilesRemoved` rather than `shuffleCleaned`, since + * the shuffle's map outputs are not unregistered. + */ + def notifyShuffleFilesRemoved(shuffleId: Int): Unit = { + try { + listeners.asScala.foreach(_.shuffleFilesRemoved(shuffleId)) + } catch { + case e: Exception => + logError(log"Error notifying removal of shuffle ${MDC(SHUFFLE_ID, shuffleId)}", e) + } + } + /** Perform broadcast cleanup. */ def doCleanupBroadcast(broadcastId: Long, blocking: Boolean): Unit = { try { @@ -376,6 +391,11 @@ private object ContextCleaner { private[spark] trait CleanerListener { def rddCleaned(rddId: Int): Unit def shuffleCleaned(shuffleId: Int): Unit + /** + * Called when a shuffle's files were removed but the shuffle is still registered on the + * MapOutputTracker (see `ContextCleaner.notifyShuffleFilesRemoved`). + */ + def shuffleFilesRemoved(shuffleId: Int): Unit = {} def broadcastCleaned(broadcastId: Long): Unit def accumCleaned(accId: Long): Unit def checkpointCleaned(rddId: Long): Unit diff --git a/core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala b/core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala index a98672fd7db4a..e30619964ab79 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala @@ -460,6 +460,9 @@ private[spark] class ExecutorMonitor( } } + // Executors no longer hold the shuffle's files, so they need not be kept around for it. + override def shuffleFilesRemoved(shuffleId: Int): Unit = shuffleCleaned(shuffleId) + override def broadcastCleaned(broadcastId: Long): Unit = { } override def accumCleaned(accId: Long): Unit = { } diff --git a/core/src/test/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitorSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitorSuite.scala index 4838255009028..dd4d949c167d2 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitorSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitorSuite.scala @@ -385,6 +385,29 @@ class ExecutorMonitorSuite extends SparkFunSuite { } + test("SPARK-59776: executor with removed shuffle files times out at the idle deadline") { + val bus = mockListenerBus() + conf.set(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, true).set(SHUFFLE_SERVICE_ENABLED, false) + monitor = new ExecutorMonitor(conf, client, bus, clock, allocationManagerSource()) + + val stage1 = stageInfo(1, shuffleId = 0) + val stage2 = stageInfo(2) + monitor.onExecutorAdded(SparkListenerExecutorAdded(clock.getTimeMillis(), "1", execInfo)) + monitor.onJobStart(SparkListenerJobStart(1, clock.getTimeMillis(), Seq(stage1, stage2))) + monitor.onTaskStart(SparkListenerTaskStart(1, 0, taskInfo("1", 1))) + monitor.onTaskEnd(SparkListenerTaskEnd(1, 0, "foo", Success, taskInfo("1", 1), + new ExecutorMetrics, null)) + monitor.onJobEnd(SparkListenerJobEnd(1, clock.getTimeMillis(), JobSucceeded)) + + // The executor holds shuffle 0's files, so only the shuffle timeout applies. + assert(monitor.timedOutExecutors(idleDeadline).isEmpty) + assert(monitor.timedOutExecutors(shuffleDeadline) === Seq("1")) + + // Files removed while the shuffle stays registered (SQL shuffle cleanup). + monitor.shuffleFilesRemoved(0) + assert(monitor.timedOutExecutors(idleDeadline) === Seq("1")) + } + test("SPARK-28839: Avoids NPE in context cleaner when shuffle service is on") { val bus = mockListenerBus() conf.set(DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED, true).set(SHUFFLE_SERVICE_ENABLED, true) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala index 295441296c506..e92ec69bfb04a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala @@ -152,6 +152,9 @@ object SQLExecution extends Logging { // MapOutputTracker so that stage retries would be triggered. Blocking is // Utils.isTesting to deflake unit tests. sc.shuffleDriverComponents.removeShuffle(shuffleId, Utils.isTesting) + // Let listeners such as the dynamic allocation ExecutorMonitor stop tracking the + // shuffle, as ContextCleaner will not report it cleaned while it stays registered. + sc.cleaner.foreach(_.notifyShuffleFilesRemoved(shuffleId)) } catch { case NonFatal(e) => logWarning(log"Failed to remove shuffle ${MDC(SHUFFLE_ID, shuffleId)} for " + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala index 870c91d2f0a0f..faeb470fe2265 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala @@ -18,14 +18,16 @@ package org.apache.spark.sql.execution import java.util.Locale -import java.util.concurrent.Executors +import java.util.concurrent.{ConcurrentLinkedQueue, Executors} import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} import scala.collection.parallel.immutable.ParRange import scala.concurrent.{ExecutionContext, Future} import scala.concurrent.duration._ +import scala.jdk.CollectionConverters._ -import org.apache.spark.{SparkConf, SparkContext, SparkEnv, SparkFunSuite} +import org.apache.spark.{CleanerListener, MapOutputTrackerMaster, SparkConf, SparkContext, SparkEnv, + SparkFunSuite} import org.apache.spark.launcher.SparkLauncher import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart} import org.apache.spark.sql.{Observation, Row, SparkSession} @@ -603,6 +605,46 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper { } } + test("SPARK-59776: RemoveShuffleFiles cleanup notifies cleaner listeners") { + val spark = SparkSession.builder().master("local[*]").appName("test").getOrCreate() + try { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key -> "true") { + val filesRemoved = new ConcurrentLinkedQueue[Int]() + val cleaned = new ConcurrentLinkedQueue[Int]() + spark.sparkContext.cleaner.get.attachListener(new CleanerListener { + override def rddCleaned(rddId: Int): Unit = {} + override def shuffleCleaned(shuffleId: Int): Unit = cleaned.add(shuffleId) + override def shuffleFilesRemoved(shuffleId: Int): Unit = filesRemoved.add(shuffleId) + override def broadcastCleaned(broadcastId: Long): Unit = {} + override def accumCleaned(accId: Long): Unit = {} + override def checkpointCleaned(rddId: Long): Unit = {} + }) + + val df = spark.range(0, 10).join(spark.range(0, 20), "id") + val qe = df.queryExecution + assert(qe.shuffleCleanupMode == RemoveShuffleFiles) + val shuffleIds = qe.executedPlan.collect { case e: ShuffleExchangeLike => e.shuffleId } + assert(shuffleIds.size > 1) + assert(df.collect().length == 10) + + // Every shuffle's files removal is reported, so e.g. the dynamic allocation + // ExecutorMonitor stops keeping executors alive for them. The shuffles stay registered on + // the MapOutputTracker, so they must not be reported as fully cleaned. + assert(filesRemoved.asScala.toSet == shuffleIds.toSet) + assert(cleaned.asScala.toSet.intersect(shuffleIds.toSet).isEmpty) + shuffleIds.foreach { id => + assert(SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + .containsShuffle(id)) + } + } + } finally { + spark.stop() + } + } + test("SPARK-59242: SkipMigration cleanup failure has mode-specific warning") { val spark = SparkSession.builder().master("local[*]").appName("test").getOrCreate() try { From 1b25e56f52f8ddafd628f58a9d9d64530f3ecd5b Mon Sep 17 00:00:00 2001 From: saurabhdamle11 Date: Thu, 24 Sep 2026 13:58:15 -0700 Subject: [PATCH 2/3] Trigger CI Co-Authored-By: Claude Opus 5.5 (1M context) From 887d8811e3086c178fd898b60fc8e7f05a2126da Mon Sep 17 00:00:00 2001 From: saurabhdamle11 Date: Fri, 25 Sep 2026 10:36:01 -0700 Subject: [PATCH 3/3] [SPARK-59776][SQL][TESTS] Drop MapOutputTracker registration assertion from RemoveShuffleFiles listener test In local mode the driver's BlockManagerStorageEndpoint handles RemoveShuffle and unregisters the shuffle from the MapOutputTrackerMaster, so the shuffle is not guaranteed to stay registered after SQL shuffle file cleanup. Keep only the listener notification assertions. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../spark/sql/execution/SQLExecutionSuite.scala | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala index faeb470fe2265..ed0b8333e84f2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala @@ -26,8 +26,7 @@ import scala.concurrent.{ExecutionContext, Future} import scala.concurrent.duration._ import scala.jdk.CollectionConverters._ -import org.apache.spark.{CleanerListener, MapOutputTrackerMaster, SparkConf, SparkContext, SparkEnv, - SparkFunSuite} +import org.apache.spark.{CleanerListener, SparkConf, SparkContext, SparkEnv, SparkFunSuite} import org.apache.spark.launcher.SparkLauncher import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart} import org.apache.spark.sql.{Observation, Row, SparkSession} @@ -631,14 +630,10 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper { assert(df.collect().length == 10) // Every shuffle's files removal is reported, so e.g. the dynamic allocation - // ExecutorMonitor stops keeping executors alive for them. The shuffles stay registered on - // the MapOutputTracker, so they must not be reported as fully cleaned. + // ExecutorMonitor stops keeping executors alive for them. The shuffles are not cleaned + // through ContextCleaner, so they must not be reported as fully cleaned. assert(filesRemoved.asScala.toSet == shuffleIds.toSet) assert(cleaned.asScala.toSet.intersect(shuffleIds.toSet).isEmpty) - shuffleIds.foreach { id => - assert(SparkEnv.get.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] - .containsShuffle(id)) - } } } finally { spark.stop()