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 | <undefined> | 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 |