ut

turboFei · turboFei · commit ea274ca0bf72 · 2025-03-19T16:59:19.000-07:00
diff --git a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.stream.Collectors;
 
@@ -359,6 +360,9 @@ public static void addSparkListener(SparkListener listener) {
    */
   private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();
 
+  @VisibleForTesting
+  public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
+
   protected static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
       getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
 
@@ -401,6 +405,7 @@ public static byte[] serializeGetReducerFileGroupResponse(
             byte[] _serializeResult = out.toByteArray();
             getReducerFileGroupResponseBroadcasts.put(
                 shuffleId, Tuple2.apply(broadcast, _serializeResult));
+            getReducerFileGroupResponseBroadcastNum.incrementAndGet();
             return _serializeResult;
           } catch (Throwable e) {
             logger.error(
@@ -430,10 +435,6 @@ public static GetReducerFileGroupResponse deserializeGetReducerFileGroupResponse
               Broadcast<GetReducerFileGroupResponse> broadcast =
                   (Broadcast<GetReducerFileGroupResponse>) objIn.readObject();
               response = broadcast.value();
-              logger.info(
-                  "Shuffle {} broadcast GetReducerFileGroupResponse size={}.",
-                  shuffleId,
-                  bytes.length);
             }
           } catch (Throwable e) {
             logger.error(
diff --git a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -25,6 +25,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.LongAdder;
 import java.util.stream.Collectors;
 
@@ -476,6 +477,9 @@ public static void addSparkListener(SparkListener listener) {
    */
   private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();
 
+  @VisibleForTesting
+  public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
+
   protected static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
       getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
 
@@ -520,6 +524,7 @@ public static byte[] serializeGetReducerFileGroupResponse(
             byte[] _serializeResult = out.toByteArray();
             getReducerFileGroupResponseBroadcasts.put(
                 shuffleId, Tuple2.apply(broadcast, _serializeResult));
+            getReducerFileGroupResponseBroadcastNum.incrementAndGet();
             return _serializeResult;
           } catch (Throwable e) {
             LOG.error(
@@ -552,10 +557,6 @@ public static GetReducerFileGroupResponse deserializeGetReducerFileGroupResponse
               Broadcast<GetReducerFileGroupResponse> broadcast =
                   (Broadcast<GetReducerFileGroupResponse>) objIn.readObject();
               response = broadcast.value();
-              LOG.info(
-                  "Shuffle {} broadcast GetReducerFileGroupResponse size={}.",
-                  shuffleId,
-                  bytes.length);
             }
           } catch (Throwable e) {
             LOG.error(
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -308,7 +308,7 @@ public abstract int getShuffleId(
   public static void registerDeserializeReducerFileGroupResponseFunction(
       BiFunction<Integer, byte[], ControlMessages.GetReducerFileGroupResponse> function) {
     if (!deserializeReducerFileGroupResponseFunction.isPresent()) {
-      deserializeReducerFileGroupResponseFunction = Optional.of(function);
+      deserializeReducerFileGroupResponseFunction = Optional.ofNullable(function);
     }
   }
 
diff --git a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
@@ -79,6 +79,10 @@ class ReducePartitionCommitHandler(
   private val rpcCacheConcurrencyLevel = conf.clientRpcCacheConcurrencyLevel
   private val rpcCacheExpireTime = conf.clientRpcCacheExpireTime
 
+  private val getReducerFileGroupResponseBroadcastEnabled = conf.getReducerFileGroupBroadcastEnabled
+  private val getReducerFileGroupResponseBroadcastMiniSize =
+    conf.getReducerFileGroupBroadcastMiniSize
+
   // noinspection UnstableApiUsage
   private val getReducerFileGroupRpcCache: Cache[Int, ByteBuffer] = CacheBuilder.newBuilder()
     .concurrencyLevel(rpcCacheConcurrencyLevel)
@@ -327,7 +331,7 @@ class ReducePartitionCommitHandler(
           getMapperAttempts(shuffleId))
 
         // only check whether broadcast enabled for the UTs
-        if (conf.getReducerFileGroupBroadcastEnabled) {
+        if (getReducerFileGroupResponseBroadcastEnabled) {
           response = broadcastGetReducerFileGroup(shuffleId, response)
         }
 
@@ -349,11 +353,17 @@ class ReducePartitionCommitHandler(
               val serializedMsg =
                 context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
 
-              if (conf.getReducerFileGroupBroadcastEnabled &&
-                serializedMsg.capacity() >= conf.getReducerFileGroupBroadcastMiniSize) {
+              if (getReducerFileGroupResponseBroadcastEnabled &&
+                serializedMsg.capacity() >= getReducerFileGroupResponseBroadcastMiniSize) {
                 val broadcastMsg = broadcastGetReducerFileGroup(shuffleId, returnedMsg)
                 if (broadcastMsg != returnedMsg) {
-                  context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(broadcastMsg)
+                  val serializedBroadcastMsg =
+                    context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(broadcastMsg)
+                  logInfo(s"Shuffle $shuffleId GetReducerFileGroupResponse size" +
+                    s" ${serializedMsg.capacity()} reached the broadcast threshold" +
+                    s" $getReducerFileGroupResponseBroadcastMiniSize," +
+                    s" the broadcast response size is ${serializedBroadcastMsg.capacity()}.")
+                  serializedBroadcastMsg
                 } else {
                   serializedMsg
                 }
diff --git a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -1207,12 +1207,14 @@ object ControlMessages extends Logging {
           case (uniqueId, pushFailedBatchSet) =>
             (uniqueId, PbSerDeUtils.fromPbPushFailedBatchSet(pushFailedBatchSet))
         }.toMap.asJava
+        val broadcast = pbGetReducerFileGroupResponse.getBroadcast.toByteArray
         GetReducerFileGroupResponse(
           Utils.toStatusCode(pbGetReducerFileGroupResponse.getStatus),
           fileGroup,
           attempts,
           partitionIds,
-          pushFailedBatches)
+          pushFailedBatches,
+          broadcast)
 
       case GET_SHUFFLE_ID_VALUE =>
         message.getParsedPayload()
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornHashSuite.scala
@@ -18,11 +18,13 @@
 package org.apache.celeborn.tests.spark
 
 import org.apache.spark.SparkConf
+import org.apache.spark.shuffle.celeborn.SparkUtils
 import org.apache.spark.sql.SparkSession
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.protocol.ShuffleMode
 
 class CelebornHashSuite extends AnyFunSuite
@@ -64,4 +66,40 @@ class CelebornHashSuite extends AnyFunSuite
 
     celebornSparkSession.stop()
   }
+
+  test("celeborn spark integration test - GetReducerFileGroupResponse broadcast") {
+    SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
+    val sparkConf = new SparkConf().setAppName("celeborn-demo").setMaster("local[2]")
+      .set(
+        s"spark.${CelebornConf.CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED.key}",
+        "true")
+      .set(
+        s"spark.${CelebornConf.CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE.key}",
+        "0")
+    val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+    val combineResult = combine(sparkSession)
+    val groupbyResult = groupBy(sparkSession)
+    val repartitionResult = repartition(sparkSession)
+    val sqlResult = runsql(sparkSession)
+
+    Thread.sleep(3000L)
+    sparkSession.stop()
+
+    val celebornSparkSession = SparkSession.builder()
+      .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+      .getOrCreate()
+    val celebornCombineResult = combine(celebornSparkSession)
+    val celebornGroupbyResult = groupBy(celebornSparkSession)
+    val celebornRepartitionResult = repartition(celebornSparkSession)
+    val celebornSqlResult = runsql(celebornSparkSession)
+
+    assert(combineResult.equals(celebornCombineResult))
+    assert(groupbyResult.equals(celebornGroupbyResult))
+    assert(repartitionResult.equals(celebornRepartitionResult))
+    assert(combineResult.equals(celebornCombineResult))
+    assert(sqlResult.equals(celebornSqlResult))
+    assert(SparkUtils.getReducerFileGroupResponseBroadcastNum.get() > 0)
+
+    celebornSparkSession.stop()
+  }
 }
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornSortSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornSortSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.celeborn.tests.spark
 
 import org.apache.spark.SparkConf
+import org.apache.spark.shuffle.celeborn.SparkUtils
 import org.apache.spark.sql.SparkSession
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
@@ -66,4 +67,42 @@ class CelebornSortSuite extends AnyFunSuite
 
     celebornSparkSession.stop()
   }
+
+  test("celeborn spark integration test - GetReducerFileGroupResponse broadcast") {
+    SparkUtils.getReducerFileGroupResponseBroadcastNum.set(0)
+    val sparkConf = new SparkConf().setAppName("celeborn-demo").setMaster("local[2]")
+      .set(s"spark.${CelebornConf.CLIENT_PUSH_SORT_RANDOMIZE_PARTITION_ENABLED.key}", "false")
+      .set(
+        s"spark.${CelebornConf.CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED.key}",
+        "true")
+      .set(
+        s"spark.${CelebornConf.CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE.key}",
+        "0")
+
+    val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+    val combineResult = combine(sparkSession)
+    val groupbyResult = groupBy(sparkSession)
+    val repartitionResult = repartition(sparkSession)
+    val sqlResult = runsql(sparkSession)
+
+    Thread.sleep(3000L)
+    sparkSession.stop()
+
+    val celebornSparkSession = SparkSession.builder()
+      .config(updateSparkConf(sparkConf, ShuffleMode.SORT))
+      .getOrCreate()
+    val celebornCombineResult = combine(celebornSparkSession)
+    val celebornGroupbyResult = groupBy(celebornSparkSession)
+    val celebornRepartitionResult = repartition(celebornSparkSession)
+    val celebornSqlResult = runsql(celebornSparkSession)
+
+    assert(combineResult.equals(celebornCombineResult))
+    assert(groupbyResult.equals(celebornGroupbyResult))
+    assert(repartitionResult.equals(celebornRepartitionResult))
+    assert(combineResult.equals(celebornCombineResult))
+    assert(sqlResult.equals(celebornSqlResult))
+    assert(SparkUtils.getReducerFileGroupResponseBroadcastNum.get() > 0)
+
+    celebornSparkSession.stop()
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -308,7 +308,7 @@ public abstract int getShuffleId(`
`308`	`308`	`public static void registerDeserializeReducerFileGroupResponseFunction(`
`309`	`309`	`BiFunction<Integer, byte[], ControlMessages.GetReducerFileGroupResponse> function) {`
`310`	`310`	`if (!deserializeReducerFileGroupResponseFunction.isPresent()) {`
`311`		`- deserializeReducerFileGroupResponseFunction = Optional.of(function);`
	`311`	`+ deserializeReducerFileGroupResponseFunction = Optional.ofNullable(function);`
`312`	`312`	`}`
`313`	`313`	`}`
`314`	`314`