[BLAZE-287][FOLLOWUP] BlazeCelebornShuffleWriter should use mapped sh…

…uffle id for rerunning stage of fetch failure (#712)
kwai · Dec 21, 2024 · cbeda0a · cbeda0a
1 parent 092c9a0
commit cbeda0a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/...a/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala b/...a/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleManager.scala
@@ -134,7 +134,10 @@ class BlazeCelebornShuffleManager(conf: SparkConf, isDriver: Boolean)
       .asInstanceOf[ShuffleClient]
 
     val celebornHandle = handle.asInstanceOf[CelebornShuffleHandle[_, _, _]]
-    val writer = new BlazeCelebornShuffleWriter(shuffleClient, context, celebornHandle, metrics)
+    val shuffleIdTracker = FieldUtils
+      .readField(celebornShuffleManager, "shuffleIdTracker", true)
+      .asInstanceOf[ExecutorShuffleIdTracker]
+    val writer = new BlazeCelebornShuffleWriter(shuffleClient, context, celebornHandle, metrics, shuffleIdTracker)
     writer.asInstanceOf[BlazeRssShuffleWriterBase[K, V]]
   }
 

diff --git a/...la/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala b/...la/org/apache/spark/sql/execution/blaze/shuffle/celeborn/BlazeCelebornShuffleWriter.scala
@@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.blaze.shuffle.celeborn
 import org.apache.celeborn.client.ShuffleClient
 import org.apache.spark.shuffle.ShuffleHandle
 import org.apache.spark.shuffle.ShuffleWriteMetricsReporter
-import org.apache.spark.shuffle.celeborn.CelebornShuffleHandle
+import org.apache.spark.shuffle.celeborn.{CelebornShuffleHandle, ExecutorShuffleIdTracker, SparkUtils}
 import org.apache.spark.sql.execution.blaze.shuffle.BlazeRssShuffleWriterBase
 import org.apache.spark.sql.execution.blaze.shuffle.RssPartitionWriterBase
 import org.apache.spark.TaskContext
@@ -29,7 +29,8 @@ class BlazeCelebornShuffleWriter[K, C](
     shuffleClient: ShuffleClient,
     taskContext: TaskContext,
     handle: CelebornShuffleHandle[K, _, C],
-    metrics: ShuffleWriteMetricsReporter)
+    metrics: ShuffleWriteMetricsReporter,
+    shuffleIdTracker: ExecutorShuffleIdTracker)
     extends BlazeRssShuffleWriterBase[K, C](metrics) {
 
   private val numMappers = handle.numMappers
@@ -41,9 +42,11 @@ class BlazeCelebornShuffleWriter[K, C](
       metrics: ShuffleWriteMetricsReporter,
       numPartitions: Int): RssPartitionWriterBase = {
 
+    val shuffleId = SparkUtils.celebornShuffleId(shuffleClient, handle, taskContext, true)
+    shuffleIdTracker.track(handle.shuffleId, shuffleId)
     new CelebornPartitionWriter(
       shuffleClient,
-      handle.shuffleId,
+      shuffleId,
       encodedAttemptId,
       numMappers,
       numPartitions,