(celeborn) branch branch-0.4 updated: [CELEBORN-1463][FOLLOWUP] Respeact client/server threads num to avoid competitiveness

[et...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ethanfeng pushed a commit to branch branch-0.4
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/branch-0.4 by this push:
     new 4add7d2f0 [CELEBORN-1463][FOLLOWUP] Respeact client/server threads num to avoid competitiveness
4add7d2f0 is described below

commit 4add7d2f0df2af9b51d979d563449c8bc30b45b5
Author: binjie yang <ya...@shizhuang-inc.com>
AuthorDate: Tue Jun 18 10:50:36 2024 +0800

    [CELEBORN-1463][FOLLOWUP] Respeact client/server threads num to avoid competitiveness
    
    ### What changes were proposed in this pull request?
     #2570
    Respect client/server threads num when create transport client/server memory allocator.
    
    ### Why are the changes needed?
    Follow up, as title.
    
    Respeact client/server threads num to avoid competitiveness
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    No
    
    Closes #2571 from zwangsheng/CELEBORN-1463-follow-up.
    
    Authored-by: binjie yang <ya...@shizhuang-inc.com>
    Signed-off-by: mingji <fe...@alibaba-inc.com>
    (cherry picked from commit 565e39ac1bfd691ea9f165f12f1c2482af4bcfd5)
    Signed-off-by: mingji <fe...@alibaba-inc.com>
---
 .../common/network/client/TransportClientFactory.java      |  3 ++-
 .../celeborn/common/network/server/TransportServer.java    |  3 ++-
 .../apache/celeborn/common/network/util/NettyUtils.java    | 14 ++++++++++++--
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
index d24e9db8f..e5670a4a6 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
@@ -106,7 +106,8 @@ public class TransportClientFactory implements Closeable {
     logger.info("mode " + ioMode + " threads " + conf.clientThreads());
     this.workerGroup =
         NettyUtils.createEventLoop(ioMode, conf.clientThreads(), conf.getModuleName() + "-client");
-    this.pooledAllocator = NettyUtils.getPooledByteBufAllocator(conf, null, false);
+    this.pooledAllocator =
+        NettyUtils.getPooledByteBufAllocator(conf, null, false, conf.clientThreads());
   }
 
   /**
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java b/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
index 53f7b7169..24de6c876 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/TransportServer.java
@@ -93,7 +93,8 @@ public class TransportServer implements Closeable {
     EventLoopGroup workerGroup =
         NettyUtils.createEventLoop(ioMode, conf.serverThreads(), conf.getModuleName() + "-server");
 
-    PooledByteBufAllocator allocator = NettyUtils.getPooledByteBufAllocator(conf, source, true);
+    PooledByteBufAllocator allocator =
+        NettyUtils.getPooledByteBufAllocator(conf, source, true, conf.serverThreads());
 
     bootstrap =
         new ServerBootstrap()
diff --git a/common/src/main/java/org/apache/celeborn/common/network/util/NettyUtils.java b/common/src/main/java/org/apache/celeborn/common/network/util/NettyUtils.java
index 60eb44744..00c970eb3 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/util/NettyUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/util/NettyUtils.java
@@ -143,15 +143,25 @@ public class NettyUtils {
 
   public static PooledByteBufAllocator getPooledByteBufAllocator(
       TransportConf conf, AbstractSource source, boolean allowCache) {
+    return getPooledByteBufAllocator(conf, source, allowCache, 0);
+  }
+
+  public static PooledByteBufAllocator getPooledByteBufAllocator(
+      TransportConf conf, AbstractSource source, boolean allowCache, int coreNum) {
     if (conf.getCelebornConf().networkShareMemoryAllocator()) {
       return getSharedPooledByteBufAllocator(
           conf.getCelebornConf(),
           source,
           allowCache && conf.getCelebornConf().networkMemoryAllocatorAllowCache());
     }
+    int arenas;
+    if (coreNum != 0) {
+      arenas = coreNum;
+    } else {
+      arenas = conf.getCelebornConf().networkAllocatorArenas();
+    }
     PooledByteBufAllocator allocator =
-        createPooledByteBufAllocator(
-            conf.preferDirectBufs(), allowCache, conf.getCelebornConf().networkAllocatorArenas());
+        createPooledByteBufAllocator(conf.preferDirectBufs(), allowCache, arenas);
     if (source != null) {
       String poolName = "default-netty-pool";
       Map<String, String> labels = new HashMap<>();