[SPARK-49502][CORE] Avoid NPE in SparkEnv.get.shuffleManager.unregisterShuffle

[cxzl25]

What changes were proposed in this pull request?

This PR aims to avoid NPE in SparkEnv.get.shuffleManager.unregisterShuffle.

Why are the changes needed?

After SPARK-45762, the shuffle manager is initialized after the block manager, which means that when the driver cleans up the shuffle, the shuffle manager may not have been initialized yet, causing NPE.

24/09/03 20:09:51,668 [dispatcher-Executor] INFO BlockManager: Initialized BlockManager: BlockManagerId(168, x, 25467, None)
24/09/03 20:09:51,684 [block-manager-storage-async-thread-pool-2] ERROR BlockManagerStorageEndpoint: Error in removing shuffle 29
java.lang.NullPointerException
        at org.apache.spark.storage.BlockManagerStorageEndpoint$$anonfun$receiveAndReply$1.$anonfun$applyOrElse$3(BlockManagerStorageEndpoint.scala:61)
        at scala.runtime.java8.JFunction0$mcZ$sp.apply(JFunction0$mcZ$sp.java:23)
        at org.apache.spark.storage.BlockManagerStorageEndpoint.$anonfun$doAsync$1(BlockManagerStorageEndpoint.scala:89)
        at scala.concurrent.Future$.$anonfun$apply$1(Future.scala:659)
        at scala.util.Success.$anonfun$map$1(Try.scala:255)
        at scala.util.Success.map(Try.scala:213)
        at scala.concurrent.Future.$anonfun$map$1(Future.scala:292)
        at scala.concurrent.impl.Promise.liftedTree1$1(Promise.scala:33)
        at scala.concurrent.impl.Promise.$anonfun$transform$1(Promise.scala:33)
        at scala.concurrent.impl.CallbackRunnable.run(Promise.scala:64)
        at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1142)
        at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:617)
        at java.lang.Thread.run(Thread.java:745)

Does this PR introduce any user-facing change?

No

How was this patch tested?

GA

Was this patch authored or co-authored using generative AI tooling?

No

[mridulm]

How is unregister shuffle getting called ? Shuffle registration requires SparkContext to have been initialized - which will include shuffle manager.

[cxzl25]

How is unregister shuffle getting called ?

The Executor is added in the case of dynamic allocation. It first initializes the BlockManager and registers it with the Driver, and then initializes the ShuffleManager.

spark/core/src/main/scala/org/apache/spark/executor/Executor.scala

Line 162 in 339d1c9

env.blockManager.initialize(conf.getAppId)

spark/core/src/main/scala/org/apache/spark/executor/Executor.scala

Line 355 in 339d1c9

env.initializeShuffleManager()

The Driver will send removeShuffle RPC to all registered BlockManagers. At this time, there are some newly added Executors that have not yet initialized the ShuffleManager.

spark/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

Lines 469 to 474 in 339d1c9

	val removeShuffleFromExecutorsFutures = blockManagerInfo.values.map { bm =>
	bm.storageEndpoint.ask[Boolean](removeMsg).recover {
	// use false as default value means no shuffle data were removed
	handleBlockRemovalFailure("shuffle", shuffleId.toString, bm.blockManagerId, false)
	}
	}.toSeq

RemoveShuffle RPC has two code paths that will trigger.

spark/core/src/main/scala/org/apache/spark/ContextCleaner.scala

Lines 235 to 241 in 339d1c9

	def doCleanupShuffle(shuffleId: Int, blocking: Boolean): Unit = {
	try {
	if (mapOutputTrackerMaster.containsShuffle(shuffleId)) {
	logDebug("Cleaning shuffle " + shuffleId)
	// Shuffle must be removed before it's unregistered from the output tracker
	// to find blocks served by the shuffle service on deallocated executors
	shuffleDriverComponents.removeShuffle(shuffleId, blocking)

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala

Lines 177 to 183 in 339d1c9

	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)

[mridulm]

Ah ! This is happening at the executor, not driver - thanks for clarifying.
Let me take a look at the detailed comment tomorrow (a bit late for me :-) )

[yaooqinn]

Shall we add a debug-level log here?

[mridulm]

Looks good to me.
Will leave it open for a bit, in case other reviewers of #43627 want to take a look as well.

[mridulm]

Merged to master.
Thanks for fixing this @cxzl25 !
Thanks for the review @yaooqinn :-)