You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2023/03/27 08:18:59 UTC

[incubator-celeborn] branch main updated: [CELEBORN-479][PERF] Refactor DataPushQueue.takePushTask to avoid busy wait (#1386)

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

zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new cb19ed1c6 [CELEBORN-479][PERF] Refactor DataPushQueue.takePushTask to avoid busy wait (#1386)
cb19ed1c6 is described below

commit cb19ed1c66a8ece834fc6aea10b8b57627d2eec6
Author: Keyong Zhou <zh...@apache.org>
AuthorDate: Mon Mar 27 16:18:55 2023 +0800

    [CELEBORN-479][PERF] Refactor DataPushQueue.takePushTask to avoid busy wait (#1386)
---
 .../celeborn/client/write/DataPushQueue.java       | 48 ++++++++++++++--------
 .../apache/celeborn/client/write/DataPusher.java   | 11 ++---
 .../celeborn/client/write/DataPushQueueSuitJ.java  | 17 ++++----
 .../common/write/InFlightRequestTracker.java       |  9 ----
 .../apache/celeborn/common/write/PushState.java    |  4 +-
 .../org/apache/celeborn/common/CelebornConf.scala  |  9 ++++
 docs/configuration/client.md                       |  1 +
 7 files changed, 59 insertions(+), 40 deletions(-)

diff --git a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
index 5d202bf6a..40e7626c5 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPushQueue.java
@@ -18,10 +18,7 @@
 package org.apache.celeborn.client.write;
 
 import java.io.IOException;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 
@@ -31,6 +28,7 @@ import org.slf4j.LoggerFactory;
 import org.apache.celeborn.client.ShuffleClient;
 import org.apache.celeborn.common.CelebornConf;
 import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.ExceptionUtils;
 import org.apache.celeborn.common.util.Utils;
 import org.apache.celeborn.common.write.PushState;
 
@@ -55,7 +53,7 @@ public class DataPushQueue {
   private final int numMappers;
   private final int numPartitions;
   private final ShuffleClient client;
-  private final Set<String> reachLimitWorker = new HashSet<>();
+  private final long takeTaskWaitTimeMs;
 
   public DataPushQueue(
       CelebornConf conf,
@@ -76,17 +74,23 @@ public class DataPushQueue {
     final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
     this.pushState = client.getPushState(mapKey);
     this.maxInFlight = conf.pushMaxReqsInFlight();
+    this.takeTaskWaitTimeMs = conf.pushTakeTaskWaitTimeMs();
     final int capacity = conf.pushQueueCapacity();
     workingQueue = new LinkedBlockingQueue<>(capacity);
   }
 
   /*
-   * Now, `takePushTask` is only used by one thread,
+   * Now, `takePushTasks` is only used by one thread,
    * so it is not thread-safe.
    * */
-  public PushTask takePushTask() throws IOException {
+  public ArrayList<PushTask> takePushTasks() throws IOException {
+    ArrayList<PushTask> tasks = new ArrayList<>();
+    HashMap<String, Integer> workerCapacity = new HashMap<>();
     while (dataPusher.stillRunning()) {
-      reachLimitWorker.clear();
+      // clear() here is necessary since inflight pushes might change after sleeping
+      // takeTaskWaitTimeMs
+      // in last loop
+      workerCapacity.clear();
       Iterator<PushTask> iterator = workingQueue.iterator();
       while (iterator.hasNext()) {
         PushTask task = iterator.next();
@@ -95,19 +99,29 @@ public class DataPushQueue {
             client.getPartitionLocation(appId, shuffleId, numMappers, numPartitions);
         if (partitionLocationMap != null) {
           PartitionLocation loc = partitionLocationMap.get(partitionId);
-          if (!reachLimitWorker.contains(loc.hostAndPushPort())) {
-            boolean reachLimit = pushState.reachLimit(loc.hostAndPushPort(), maxInFlight);
-            if (!reachLimit) {
-              iterator.remove();
-              return task;
-            } else {
-              reachLimitWorker.add(loc.hostAndPushPort());
-            }
+          Integer oldCapacity = workerCapacity.get(loc.hostAndPushPort());
+          if (oldCapacity == null) {
+            oldCapacity = maxInFlight - pushState.inflightPushes(loc.hostAndPushPort());
+            workerCapacity.put(loc.hostAndPushPort(), oldCapacity);
+          }
+          if (oldCapacity > 0) {
+            iterator.remove();
+            tasks.add(task);
+            workerCapacity.put(loc.hostAndPushPort(), oldCapacity - 1);
           }
         } else {
-          return task;
+          tasks.add(task);
         }
       }
+      if (!tasks.isEmpty()) {
+        return tasks;
+      }
+      try {
+        // Reaching here means no available tasks can be pushed to any worker, wait for a while
+        Thread.sleep(takeTaskWaitTimeMs);
+      } catch (InterruptedException ie) {
+        ExceptionUtils.wrapAndThrowIOException(ie);
+      }
     }
     return null;
   }
diff --git a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
index 48b0a2163..3fdaaafa9 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
@@ -18,6 +18,7 @@
 package org.apache.celeborn.client.write;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Objects;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -121,12 +122,12 @@ public class DataPusher {
           public void run() {
             while (stillRunning()) {
               try {
-                PushTask task = dataPushQueue.takePushTask();
-                if (Objects.isNull(task)) {
-                  continue;
+                ArrayList<PushTask> tasks = dataPushQueue.takePushTasks();
+                for (int i = 0; i < tasks.size(); i++) {
+                  PushTask task = tasks.get(i);
+                  pushData(task);
+                  reclaimTask(task);
                 }
-                pushData(task);
-                reclaimTask(task);
               } catch (CelebornIOException e) {
                 exceptionRef.set(e);
               } catch (InterruptedException | IOException e) {
diff --git a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
index 9311a3440..29b66a33f 100644
--- a/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
+++ b/client/src/test/java/org/apache/celeborn/client/write/DataPushQueueSuitJ.java
@@ -129,13 +129,16 @@ public class DataPushQueueSuitJ {
             () -> {
               while (running.get()) {
                 try {
-                  PushTask task = dataPushQueue.takePushTask();
-                  byte[] buffer = task.getBuffer();
-                  int partitionId = task.getPartitionId();
-                  tarWorkerData.get(partitionId % numWorker).add(bytesToInt(buffer));
-                  pushState.removeBatch(
-                      partitionBatchIdMap.get(partitionId),
-                      reducePartitionMap.get(partitionId).hostAndPushPort());
+                  ArrayList<PushTask> tasks = dataPushQueue.takePushTasks();
+                  for (int i = 0; i < tasks.size(); i++) {
+                    PushTask task = tasks.get(i);
+                    byte[] buffer = task.getBuffer();
+                    int partitionId = task.getPartitionId();
+                    tarWorkerData.get(partitionId % numWorker).add(bytesToInt(buffer));
+                    pushState.removeBatch(
+                        partitionBatchIdMap.get(partitionId),
+                        reducePartitionMap.get(partitionId).hostAndPushPort());
+                  }
                 } catch (IOException e) {
                   throw new RuntimeException(e);
                 }
diff --git a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
index aec6bd7bf..76ffc33ff 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
@@ -160,15 +160,6 @@ public class InFlightRequestTracker {
     return times <= 0;
   }
 
-  public boolean reachLimit(String hostAndPushPort, int maxInFlight) throws IOException {
-    if (pushState.exception.get() != null) {
-      throw pushState.exception.get();
-    }
-
-    Set<Integer> batchIdSet = getBatchIdSetByAddressPair(hostAndPushPort);
-    return batchIdSet.size() > maxInFlight;
-  }
-
   protected int nextBatchId() {
     return batchId.incrementAndGet();
   }
diff --git a/common/src/main/java/org/apache/celeborn/common/write/PushState.java b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
index e8fb9a2e9..873ffc8ad 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
@@ -90,7 +90,7 @@ public class PushState {
     return inFlightRequestTracker.limitZeroInFlight();
   }
 
-  public boolean reachLimit(String hostAndPushPort, int maxInFlight) throws IOException {
-    return inFlightRequestTracker.reachLimit(hostAndPushPort, maxInFlight);
+  public int inflightPushes(String hostAndPushPort) {
+    return inFlightRequestTracker.getBatchIdSetByAddressPair(hostAndPushPort).size();
   }
 }
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 27a3cf6c1..7a8506f3a 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -683,6 +683,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
     }
   def pushLimitInFlightSleepDeltaMs: Long = get(PUSH_LIMIT_IN_FLIGHT_SLEEP_INTERVAL)
   def pushSplitPartitionThreads: Int = get(PUSH_SPLIT_PARTITION_THREADS)
+  def pushTakeTaskWaitTimeMs: Long = get(PUSH_TAKE_TASK_WAIT_TIME)
   def partitionSplitMode: PartitionSplitMode = PartitionSplitMode.valueOf(get(PARTITION_SPLIT_MODE))
   def partitionSplitThreshold: Long = get(PARTITION_SPLIT_THRESHOLD)
   def batchHandleChangePartitionEnabled: Boolean = get(BATCH_HANDLE_CHANGE_PARTITION_ENABLED)
@@ -2344,6 +2345,14 @@ object CelebornConf extends Logging {
       .intConf
       .createWithDefault(8)
 
+  val PUSH_TAKE_TASK_WAIT_TIME: ConfigEntry[Long] =
+    buildConf("celeborn.push.takeTaskWaitTime")
+      .categories("client")
+      .doc("Wait time if no task available to push to worker.")
+      .version("0.3.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createWithDefaultString("50ms")
+
   val PARTITION_SPLIT_THRESHOLD: ConfigEntry[Long] =
     buildConf("celeborn.shuffle.partitionSplit.threshold")
       .withAlternative("rss.partition.split.threshold")
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index ab97f568a..b3bc3a32c 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -45,6 +45,7 @@ license: |
 | celeborn.push.sortMemory.threshold | 64m | When SortBasedPusher use memory over the threshold, will trigger push data. | 0.2.0 | 
 | celeborn.push.splitPartition.threads | 8 | Thread number to process shuffle split request in shuffle client. | 0.2.0 | 
 | celeborn.push.stageEnd.timeout | &lt;undefined&gt; | Timeout for waiting StageEnd. Default value should be `celeborn.rpc.askTimeout * (celeborn.rpc.requestCommitFiles.maxRetries + 1)`. | 0.2.0 | 
+| celeborn.push.takeTaskWaitTime | 50ms | Wait time if no task available to push to worker. | 0.3.0 | 
 | celeborn.rpc.cache.concurrencyLevel | 32 | The number of write locks to update rpc cache. | 0.2.0 | 
 | celeborn.rpc.cache.expireTime | 15s | The time before a cache item is removed. | 0.2.0 | 
 | celeborn.rpc.cache.size | 256 | The max cache items count for rpc cache. | 0.2.0 |