You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@celeborn.apache.org by "turboFei (via GitHub)" <gi...@apache.org> on 2025/03/19 07:00:44 UTC

[PR] Broadcast GetReducerFileGroupResponse [celeborn]

turboFei opened a new pull request, #3158:
URL: https://github.com/apache/celeborn/pull/3158

   <!--
   Thanks for sending a pull request!  Here are some tips for you:
     - Make sure the PR title start w/ a JIRA ticket, e.g. '[CELEBORN-XXXX] Your PR title ...'.
     - Be sure to keep the PR description updated to reflect all changes.
     - Please write your PR title to summarize what this PR proposes.
     - If possible, provide a concise example to reproduce the issue for a faster review.
   -->
   
   ### What changes were proposed in this pull request?
   
   
   
   ### Why are the changes needed?
   
   
   
   ### Does this PR introduce _any_ user-facing change?
   
   
   
   ### How was this patch tested?
   
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2013393825


##########
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala:
##########
@@ -337,14 +349,44 @@ class ReducePartitionCommitHandler(
                   shufflePushFailedBatches.getOrDefault(
                     shuffleId,
                     new util.HashMap[String, util.Set[PushFailedBatch]]()))
-              context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              val serializedMsg =
+                context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              if (getReducerFileGroupResponseBroadcastEnabled &&
+                serializedMsg.capacity() >= getReducerFileGroupResponseBroadcastMiniSize) {
+                val broadcastMsg = broadcastGetReducerFileGroup(shuffleId, returnedMsg)
+                if (broadcastMsg != returnedMsg) {
+                  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()}.")

Review Comment:
   > How much bandwidth can be saved for this RPC response? Maybe adding some logs here can be helpful.
   Hi @FMX Here is the log.
   
   



##########
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala:
##########
@@ -337,14 +349,44 @@ class ReducePartitionCommitHandler(
                   shufflePushFailedBatches.getOrDefault(
                     shuffleId,
                     new util.HashMap[String, util.Set[PushFailedBatch]]()))
-              context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              val serializedMsg =
+                context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              if (getReducerFileGroupResponseBroadcastEnabled &&
+                serializedMsg.capacity() >= getReducerFileGroupResponseBroadcastMiniSize) {
+                val broadcastMsg = broadcastGetReducerFileGroup(shuffleId, returnedMsg)
+                if (broadcastMsg != returnedMsg) {
+                  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()}.")

Review Comment:
   > How much bandwidth can be saved for this RPC response? Maybe adding some logs here can be helpful.
   
   
   Hi @FMX Here is the log.
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#issuecomment-2748753948

   will rebase the code to resolve conflicts


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004744350


##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java:
##########
@@ -462,4 +470,121 @@ public static void addSparkListener(SparkListener listener) {
       sparkContext.addSparkListener(listener);
     }
   }
+
+  /**
+   * A [[KeyLock]] whose key is a shuffle id to ensure there is only one thread accessing the
+   * broadcast belonging to the shuffle id at a time.
+   */
+  private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();
+
+  @VisibleForTesting
+  public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
+
+  @VisibleForTesting
+  public static Map<Integer, Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>>
+      getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
+
+  public static byte[] serializeGetReducerFileGroupResponse(
+      Integer shuffleId, GetReducerFileGroupResponse response) {
+    SparkContext sparkContext = SparkContext$.MODULE$.getActive().getOrElse(null);
+    if (sparkContext == null) {
+      LOG.error("Can not get active SparkContext.");
+      return null;
+    }
+
+    return shuffleBroadcastLock.withLock(
+        shuffleId,
+        () -> {
+          Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>
+              cachedSerializeGetReducerFileGroupResponse =
+                  getReducerFileGroupResponseBroadcasts.get(shuffleId);
+          if (cachedSerializeGetReducerFileGroupResponse != null) {
+            return cachedSerializeGetReducerFileGroupResponse._2;
+          }
+
+          try {
+            LOG.info("Broadcasting GetReducerFileGroupResponse for shuffle: {}", shuffleId);
+            Broadcast<GetReducerFileGroupResponse> broadcast =
+                sparkContext.broadcast(
+                    response,
+                    scala.reflect.ClassManifestFactory.fromClass(
+                        GetReducerFileGroupResponse.class));
+
+            CompressionCodec codec = CompressionCodec.createCodec(sparkContext.conf());
+            // Using `org.apache.commons.io.output.ByteArrayOutputStream` instead of the standard
+            // one
+            // This implementation doesn't reallocate the whole memory block but allocates
+            // additional buffers. This way no buffers need to be garbage collected and
+            // the contents don't have to be copied to the new buffer.
+            org.apache.commons.io.output.ByteArrayOutputStream out =
+                new org.apache.commons.io.output.ByteArrayOutputStream();
+            try (ObjectOutputStream oos =
+                new ObjectOutputStream(codec.compressedOutputStream(out))) {
+              oos.writeObject(broadcast);
+            }
+            byte[] _serializeResult = out.toByteArray();
+            getReducerFileGroupResponseBroadcasts.put(
+                shuffleId, Tuple2.apply(broadcast, _serializeResult));
+            getReducerFileGroupResponseBroadcastNum.incrementAndGet();
+            return _serializeResult;
+          } catch (Throwable e) {
+            LOG.error(
+                "Failed to serialize GetReducerFileGroupResponse for shuffle: " + shuffleId, e);
+            return null;
+          }
+        });
+  }
+
+  public static GetReducerFileGroupResponse deserializeGetReducerFileGroupResponse(
+      Integer shuffleId, byte[] bytes) {
+    SparkContext sparkContext = SparkContext$.MODULE$.getActive().getOrElse(null);
+    if (sparkContext == null) {
+      LOG.error("Can not get active SparkContext.");
+      return null;
+    }
+
+    return shuffleBroadcastLock.withLock(
+        shuffleId,
+        () -> {
+          GetReducerFileGroupResponse response = null;
+          LOG.info(
+              "Deserializing GetReducerFileGroupResponse broadcast for shuffle: {}", shuffleId);
+
+          try {
+            CompressionCodec codec = CompressionCodec.createCodec(sparkContext.conf());
+            try (ObjectInputStream objIn =
+                new ObjectInputStream(
+                    codec.compressedInputStream(new ByteArrayInputStream(bytes)))) {
+              Broadcast<GetReducerFileGroupResponse> broadcast =
+                  (Broadcast<GetReducerFileGroupResponse>) objIn.readObject();
+              response = broadcast.value();
+            }
+          } catch (Throwable e) {
+            LOG.error(
+                "Failed to deserialize GetReducerFileGroupResponse for shuffle: " + shuffleId, e);
+          }
+          return response;
+        });
+  }
+
+  public static void invalidateSerializedGetReducerFileGroupResponse(Integer shuffleId) {
+    shuffleBroadcastLock.withLock(
+        shuffleId,
+        () -> {
+          try {
+            Tuple2<Broadcast<GetReducerFileGroupResponse>, byte[]>
+                cachedSerializeGetReducerFileGroupResponse =
+                    getReducerFileGroupResponseBroadcasts.remove(shuffleId);
+            if (cachedSerializeGetReducerFileGroupResponse != null) {
+              cachedSerializeGetReducerFileGroupResponse._1().destroy();
+            }
+          } catch (Throwable e) {
+            LOG.error(
+                "Failed to invalidate serialized GetReducerFileGroupResponse for shuffle: "
+                    + shuffleId,
+                e);
+          }
+          return null;
+        });
+  }

Review Comment:
   Refer Spark `MapOutputTracker` https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/MapOutputTracker.scala#L393
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004744738


##########
common/src/main/scala/org/apache/celeborn/common/util/KeyLock.scala:
##########
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.util
+
+import java.util.concurrent.{Callable, ConcurrentHashMap}
+
+/**
+ * This class is copied from Apache Spark.
+ * A special locking mechanism to provide locking with a given key. By providing the same key
+ * (identity is tested using the `equals` method), we ensure there is only one `func` running at
+ * the same time.
+ *
+ * @tparam K the type of key to identify a lock. This type must implement `equals` and `hashCode`
+ *           correctly as it will be the key type of an internal Map.
+ */
+class KeyLock[K] {
+
+  private val lockMap = new ConcurrentHashMap[K, AnyRef]()
+
+  private def acquireLock(key: K): Unit = {
+    while (true) {
+      val lock = lockMap.putIfAbsent(key, new Object)
+      if (lock == null) return
+      lock.synchronized {
+        while (lockMap.get(key) eq lock) {
+          lock.wait()
+        }
+      }
+    }
+  }
+
+  private def releaseLock(key: K): Unit = {
+    val lock = lockMap.remove(key)
+    lock.synchronized {
+      lock.notifyAll()
+    }
+  }
+
+  /**
+   * Run `func` under a lock identified by the given key. Multiple calls with the same key
+   * (identity is tested using the `equals` method) will be locked properly to ensure there is only
+   * one `func` running at the same time.
+   */
+  def withLock[T](key: K)(func: Callable[T]): T = {
+    if (key == null) {
+      throw new NullPointerException("key must not be null")
+    }
+    acquireLock(key)
+    try {
+      func.call()
+    } finally {
+      releaseLock(key)
+    }
+  }
+}

Review Comment:
   copied from https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/util/KeyLock.scala for spark-2



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "cxzl25 (via GitHub)" <gi...@apache.org>.
cxzl25 commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2013870101


##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java:
##########
@@ -476,4 +489,121 @@ public static boolean isCelebornSkewShuffleOrChildShuffle(int appShuffleId) {
       return false;
     }
   }
+
+  /**
+   * A [[KeyLock]] whose key is a shuffle id to ensure there is only one thread accessing the
+   * broadcast belonging to the shuffle id at a time.
+   */
+  private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();

Review Comment:
   ```suggestion
     private static final KeyLock<Integer> shuffleBroadcastLock = new KeyLock<>();
   ```



##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java:
##########
@@ -476,4 +489,121 @@ public static boolean isCelebornSkewShuffleOrChildShuffle(int appShuffleId) {
       return false;
     }
   }
+
+  /**
+   * A [[KeyLock]] whose key is a shuffle id to ensure there is only one thread accessing the
+   * broadcast belonging to the shuffle id at a time.
+   */
+  private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();
+
+  @VisibleForTesting
+  public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
+
+  @VisibleForTesting
+  public static Map<Integer, Tuple2<Broadcast<TransportMessage>, byte[]>>
+      getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
+
+  public static byte[] serializeGetReducerFileGroupResponse(
+      Integer shuffleId, GetReducerFileGroupResponse response) {
+    SparkContext sparkContext = SparkContext$.MODULE$.getActive().getOrElse(null);
+    if (sparkContext == null) {
+      LOG.error("Can not get active SparkContext.");
+      return null;
+    }
+
+    return shuffleBroadcastLock.withLock(
+        shuffleId,
+        () -> {
+          Tuple2<Broadcast<TransportMessage>, byte[]> cachedSerializeGetReducerFileGroupResponse =
+              getReducerFileGroupResponseBroadcasts.get(shuffleId);
+          if (cachedSerializeGetReducerFileGroupResponse != null) {
+            return cachedSerializeGetReducerFileGroupResponse._2;
+          }
+
+          try {
+            LOG.info("Broadcasting GetReducerFileGroupResponse for shuffle: {}", shuffleId);
+            TransportMessage transportMessage =
+                (TransportMessage) Utils.toTransportMessage(response);
+            Broadcast<TransportMessage> broadcast =
+                sparkContext.broadcast(
+                    transportMessage,
+                    scala.reflect.ClassManifestFactory.fromClass(TransportMessage.class));
+
+            CompressionCodec codec = CompressionCodec$.MODULE$.createCodec(sparkContext.conf());
+            // Using `org.apache.commons.io.output.ByteArrayOutputStream` instead of the standard
+            // one
+            // This implementation doesn't reallocate the whole memory block but allocates
+            // additional buffers. This way no buffers need to be garbage collected and
+            // the contents don't have to be copied to the new buffer.
+            org.apache.commons.io.output.ByteArrayOutputStream out =
+                new org.apache.commons.io.output.ByteArrayOutputStream();
+            try (ObjectOutputStream oos =
+                new ObjectOutputStream(codec.compressedOutputStream(out))) {
+              oos.writeObject(broadcast);
+            }
+            byte[] _serializeResult = out.toByteArray();
+            getReducerFileGroupResponseBroadcasts.put(
+                shuffleId, Tuple2.apply(broadcast, _serializeResult));
+            getReducerFileGroupResponseBroadcastNum.incrementAndGet();
+            return _serializeResult;
+          } catch (Throwable e) {
+            LOG.error(
+                "Failed to serialize GetReducerFileGroupResponse for shuffle: " + shuffleId, e);

Review Comment:
   ```suggestion
               LOG.error("Failed to serialize GetReducerFileGroupResponse for shuffle: {}", shuffleId, e);
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei closed pull request #3158: [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted
URL: https://github.com/apache/celeborn/pull/3158


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#issuecomment-2739510342

   ![image](https://github.com/user-attachments/assets/a82c9183-bf19-43b1-9394-ece4548a4082)
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004904313


##########
client-spark/spark-3-shaded/pom.xml:
##########
@@ -77,6 +77,7 @@
               <include>io.netty:*</include>
               <include>org.apache.commons:commons-lang3</include>
               <include>org.roaringbitmap:RoaringBitmap</include>
+              <include>commons-io:commons-io</include>

Review Comment:
   The org.apache.commons package is relocated.
   https://github.com/apache/celeborn/blob/9e8f9f6b193a94ca6e9467aeda6e7172f588eb70/client-spark/spark-3-shaded/pom.xml#L63-L64



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004897639


##########
client-spark/spark-3-shaded/pom.xml:
##########
@@ -77,6 +77,7 @@
               <include>io.netty:*</include>
               <include>org.apache.commons:commons-lang3</include>
               <include>org.roaringbitmap:RoaringBitmap</include>
+              <include>commons-io:commons-io</include>

Review Comment:
    
   
   
   To fix the class not found:
   ```
   rn.ApplicationMaster$$anon$2.run(ApplicationMaster.scala:732)
   Caused by: java.lang.NoClassDefFoundError: org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream
   	at org.apache.spark.shuffle.celeborn.SparkShuffleManager.<init>(SparkShuffleManager.java:113)
   	... 20 more
   Caused by: java.lang.ClassNotFoundException: org.apache.celeborn.shaded.org.apache.commons.io.output.ByteArrayOutputStream
   	at java.net.URLClassLoader.findClass(URLClassLoader.java:387)
   	at java.lang.ClassLoader.loadClass(ClassLoader.java:419)
   	at sun.misc.Launcher$AppClassLoader.loadClass(Launcher.java:365)
   	at java.lang.ClassLoader.loadClass(ClassLoader.java:352)
   	... 21 more
   ```
   
   
   ```
   jar -tf client-spark/spark-3-shaded/target/celeborn-client-spark-3-shaded_2.12-0.6.0-SNAPSHOT.jar|grep org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream
   org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream.class
   ```
   
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004742371


##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -5208,6 +5212,27 @@ object CelebornConf extends Logging {
       .intConf
       .createWithDefault(2000)
 
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED =
+    buildConf("celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.enabled")
+      .categories("client")
+      .doc(
+        "Whether to leverage Spark broadcast mechanism to send the GetReducerFileGroupResponse. " +
+          "If the response size is large and Spark executor number is large, the Spark driver network " +
+          "may be exhausted because each executor will pull the response from the driver. With broadcasting " +
+          "GetReducerFileGroupResponse, it prevents the driver from being the bottleneck in sending out multiple " +
+          "copies of the GetReducerFileGroupResponse (one per executor).")
+      .version("0.6.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE =

Review Comment:
   same with https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1718
   
   ```
     private[spark] val SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST =
       ConfigBuilder("spark.shuffle.mapOutput.minSizeForBroadcast")
         .doc("The size at which we use Broadcast to send the map output statuses to the executors.")
         .version("2.0.0")
         .bytesConf(ByteUnit.BYTE)
         .createWithDefaultString("512k")
   ```



##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -5208,6 +5212,27 @@ object CelebornConf extends Logging {
       .intConf
       .createWithDefault(2000)
 
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED =
+    buildConf("celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.enabled")
+      .categories("client")
+      .doc(
+        "Whether to leverage Spark broadcast mechanism to send the GetReducerFileGroupResponse. " +
+          "If the response size is large and Spark executor number is large, the Spark driver network " +
+          "may be exhausted because each executor will pull the response from the driver. With broadcasting " +
+          "GetReducerFileGroupResponse, it prevents the driver from being the bottleneck in sending out multiple " +
+          "copies of the GetReducerFileGroupResponse (one per executor).")
+      .version("0.6.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE =

Review Comment:
   same default value with https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1718
   
   ```
     private[spark] val SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST =
       ConfigBuilder("spark.shuffle.mapOutput.minSizeForBroadcast")
         .doc("The size at which we use Broadcast to send the map output statuses to the executors.")
         .version("2.0.0")
         .bytesConf(ByteUnit.BYTE)
         .createWithDefaultString("512k")
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004904313


##########
client-spark/spark-3-shaded/pom.xml:
##########
@@ -77,6 +77,7 @@
               <include>io.netty:*</include>
               <include>org.apache.commons:commons-lang3</include>
               <include>org.roaringbitmap:RoaringBitmap</include>
+              <include>commons-io:commons-io</include>

Review Comment:
   https://github.com/apache/celeborn/blob/9e8f9f6b193a94ca6e9467aeda6e7172f588eb70/client-spark/spark-3-shaded/pom.xml#L63-L64



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2013394762


##########
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala:
##########
@@ -337,14 +349,44 @@ class ReducePartitionCommitHandler(
                   shufflePushFailedBatches.getOrDefault(
                     shuffleId,
                     new util.HashMap[String, util.Set[PushFailedBatch]]()))
-              context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              val serializedMsg =
+                context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              if (getReducerFileGroupResponseBroadcastEnabled &&
+                serializedMsg.capacity() >= getReducerFileGroupResponseBroadcastMiniSize) {
+                val broadcastMsg = broadcastGetReducerFileGroup(shuffleId, returnedMsg)
+                if (broadcastMsg != returnedMsg) {
+                  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()}.")

Review Comment:
   Here is an example:
   ![image](https://github.com/user-attachments/assets/efc4d75c-f592-489b-8fde-34440f938b98)
   
   @FMX 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [WIP] Broadcast GetReducerFileGroupResponse [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei closed pull request #3158: [WIP] Broadcast GetReducerFileGroupResponse
URL: https://github.com/apache/celeborn/pull/3158


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "FMX (via GitHub)" <gi...@apache.org>.
FMX commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2011699248


##########
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala:
##########
@@ -337,14 +349,44 @@ class ReducePartitionCommitHandler(
                   shufflePushFailedBatches.getOrDefault(
                     shuffleId,
                     new util.HashMap[String, util.Set[PushFailedBatch]]()))
-              context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              val serializedMsg =
+                context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
+
+              if (getReducerFileGroupResponseBroadcastEnabled &&
+                serializedMsg.capacity() >= getReducerFileGroupResponseBroadcastMiniSize) {
+                val broadcastMsg = broadcastGetReducerFileGroup(shuffleId, returnedMsg)

Review Comment:
   How much bandwidth can be saved for this RPC response? Maybe adding some logs here can be helpful.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004742371


##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -5208,6 +5212,27 @@ object CelebornConf extends Logging {
       .intConf
       .createWithDefault(2000)
 
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED =
+    buildConf("celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.enabled")
+      .categories("client")
+      .doc(
+        "Whether to leverage Spark broadcast mechanism to send the GetReducerFileGroupResponse. " +
+          "If the response size is large and Spark executor number is large, the Spark driver network " +
+          "may be exhausted because each executor will pull the response from the driver. With broadcasting " +
+          "GetReducerFileGroupResponse, it prevents the driver from being the bottleneck in sending out multiple " +
+          "copies of the GetReducerFileGroupResponse (one per executor).")
+      .version("0.6.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE =

Review Comment:
   same default value with https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1718
   
   ```
     private[spark] val SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST =
       ConfigBuilder("spark.shuffle.mapOutput.minSizeForBroadcast")
         .doc("The size at which we use Broadcast to send the map output statuses to the executors.")
         .version("2.0.0")
         .bytesConf(ByteUnit.BYTE)
         .createWithDefaultString("512k")
   ```



##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -5208,6 +5212,27 @@ object CelebornConf extends Logging {
       .intConf
       .createWithDefault(2000)
 
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_ENABLED =
+    buildConf("celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.enabled")
+      .categories("client")
+      .doc(
+        "Whether to leverage Spark broadcast mechanism to send the GetReducerFileGroupResponse. " +
+          "If the response size is large and Spark executor number is large, the Spark driver network " +
+          "may be exhausted because each executor will pull the response from the driver. With broadcasting " +
+          "GetReducerFileGroupResponse, it prevents the driver from being the bottleneck in sending out multiple " +
+          "copies of the GetReducerFileGroupResponse (one per executor).")
+      .version("0.6.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  val CLIENT_SHUFFLE_GET_REDUCER_FILE_GROUP_BROADCAST_MINI_SIZE =
+    buildConf("celeborn.client.spark.shuffle.getReducerFileGroup.broadcast.miniSize")
+      .categories("client")
+      .doc("The size at which we use Broadcast to send the GetReducerFileGroupResponse to the executors.")
+      .version("0.6.0")
+      .bytesConf(ByteUnit.BYTE)
+      .createWithDefaultString("512k")

Review Comment:
   same default value with https://github.com/apache/spark/blob/8d260084b8a50ff59a127c7292c0cdb6737981b0/core/src/main/scala/org/apache/spark/internal/config/package.scala#L1718
   
   ```
     private[spark] val SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST =
       ConfigBuilder("spark.shuffle.mapOutput.minSizeForBroadcast")
         .doc("The size at which we use Broadcast to send the map output statuses to the executors.")
         .version("2.0.0")
         .bytesConf(ByteUnit.BYTE)
         .createWithDefaultString("512k")
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#issuecomment-2769766128

   thanks, merged to main


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2014205885


##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java:
##########
@@ -476,4 +489,121 @@ public static boolean isCelebornSkewShuffleOrChildShuffle(int appShuffleId) {
       return false;
     }
   }
+
+  /**
+   * A [[KeyLock]] whose key is a shuffle id to ensure there is only one thread accessing the
+   * broadcast belonging to the shuffle id at a time.
+   */
+  private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();

Review Comment:
   thanks, adopted for both spark3 and spark2



##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java:
##########
@@ -476,4 +489,121 @@ public static boolean isCelebornSkewShuffleOrChildShuffle(int appShuffleId) {
       return false;
     }
   }
+
+  /**
+   * A [[KeyLock]] whose key is a shuffle id to ensure there is only one thread accessing the
+   * broadcast belonging to the shuffle id at a time.
+   */
+  private static KeyLock<Integer> shuffleBroadcastLock = new KeyLock();
+
+  @VisibleForTesting
+  public static AtomicInteger getReducerFileGroupResponseBroadcastNum = new AtomicInteger();
+
+  @VisibleForTesting
+  public static Map<Integer, Tuple2<Broadcast<TransportMessage>, byte[]>>
+      getReducerFileGroupResponseBroadcasts = JavaUtils.newConcurrentHashMap();
+
+  public static byte[] serializeGetReducerFileGroupResponse(
+      Integer shuffleId, GetReducerFileGroupResponse response) {
+    SparkContext sparkContext = SparkContext$.MODULE$.getActive().getOrElse(null);
+    if (sparkContext == null) {
+      LOG.error("Can not get active SparkContext.");
+      return null;
+    }
+
+    return shuffleBroadcastLock.withLock(
+        shuffleId,
+        () -> {
+          Tuple2<Broadcast<TransportMessage>, byte[]> cachedSerializeGetReducerFileGroupResponse =
+              getReducerFileGroupResponseBroadcasts.get(shuffleId);
+          if (cachedSerializeGetReducerFileGroupResponse != null) {
+            return cachedSerializeGetReducerFileGroupResponse._2;
+          }
+
+          try {
+            LOG.info("Broadcasting GetReducerFileGroupResponse for shuffle: {}", shuffleId);
+            TransportMessage transportMessage =
+                (TransportMessage) Utils.toTransportMessage(response);
+            Broadcast<TransportMessage> broadcast =
+                sparkContext.broadcast(
+                    transportMessage,
+                    scala.reflect.ClassManifestFactory.fromClass(TransportMessage.class));
+
+            CompressionCodec codec = CompressionCodec$.MODULE$.createCodec(sparkContext.conf());
+            // Using `org.apache.commons.io.output.ByteArrayOutputStream` instead of the standard
+            // one
+            // This implementation doesn't reallocate the whole memory block but allocates
+            // additional buffers. This way no buffers need to be garbage collected and
+            // the contents don't have to be copied to the new buffer.
+            org.apache.commons.io.output.ByteArrayOutputStream out =
+                new org.apache.commons.io.output.ByteArrayOutputStream();
+            try (ObjectOutputStream oos =
+                new ObjectOutputStream(codec.compressedOutputStream(out))) {
+              oos.writeObject(broadcast);
+            }
+            byte[] _serializeResult = out.toByteArray();
+            getReducerFileGroupResponseBroadcasts.put(
+                shuffleId, Tuple2.apply(broadcast, _serializeResult));
+            getReducerFileGroupResponseBroadcastNum.incrementAndGet();
+            return _serializeResult;
+          } catch (Throwable e) {
+            LOG.error(
+                "Failed to serialize GetReducerFileGroupResponse for shuffle: " + shuffleId, e);

Review Comment:
   thanks, adopted for both spark3 and spark2



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#issuecomment-2769763999

   Thanks all for your suggestions and review, merging


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on code in PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#discussion_r2004897639


##########
client-spark/spark-3-shaded/pom.xml:
##########
@@ -77,6 +77,7 @@
               <include>io.netty:*</include>
               <include>org.apache.commons:commons-lang3</include>
               <include>org.roaringbitmap:RoaringBitmap</include>
+              <include>commons-io:commons-io</include>

Review Comment:
    
   
   
   To fix the class not found:
   ```
   Caused by: java.lang.NoClassDefFoundError: org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream
   	at org.apache.spark.shuffle.celeborn.SparkShuffleManager.<init>(SparkShuffleManager.java:113)
   	... 20 more
   Caused by: java.lang.ClassNotFoundException: org.apache.celeborn.shaded.org.apache.commons.io.output.ByteArrayOutputStream
   	at java.net.URLClassLoader.findClass(URLClassLoader.java:387)
   	at java.lang.ClassLoader.loadClass(ClassLoader.java:419)
   	at sun.misc.Launcher$AppClassLoader.loadClass(Launcher.java:365)
   	at java.lang.ClassLoader.loadClass(ClassLoader.java:352)
   ```
   
   
   ```
   jar -tf client-spark/spark-3-shaded/target/celeborn-client-spark-3-shaded_2.12-0.6.0-SNAPSHOT.jar|grep org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream
   org/apache/celeborn/shaded/org/apache/commons/io/output/ByteArrayOutputStream.class
   ```
   
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [CELEBORN-1921] Broadcast large GetReducerFileGroupResponse to prevent Spark driver network exhausted [celeborn]

Posted by "turboFei (via GitHub)" <gi...@apache.org>.
turboFei commented on PR #3158:
URL: https://github.com/apache/celeborn/pull/3158#issuecomment-2741316510

   I have finished the cluster testing.
   
   ![image](https://github.com/user-attachments/assets/7938b79f-6f01-4349-9c8b-1981594196e8)
   
   It works fine.
   
   cc @FMX @RexXiong 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@celeborn.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org