This is an automated email from the ASF dual-hosted git repository.
rexxiong pushed a commit to branch branch-0.3
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/branch-0.3 by this push:
new 5dd44fc58 [CELEBORN-656] Batch revive RPCs in client to avoid too many
requests
5dd44fc58 is described below
commit 5dd44fc58575d8b16001cda32c040e1d9b78a8f1
Author: zky.zhoukeyong <[email protected]>
AuthorDate: Tue Jun 27 22:11:04 2023 +0800
[CELEBORN-656] Batch revive RPCs in client to avoid too many requests
### What changes were proposed in this pull request?
This PR batches revive requests and periodically send to LifecycleManager
to reduce number or RPC requests.
To be more detailed. This PR changes Revive message to support multiple
unique partitions, and also passes a set unique mapIds for checking MapEnd.
Each time ShuffleClientImpl wants to revive, it adds a ReviveRquest to
ReviveManager and wait for result. ReviveManager batches revive requests and
periodically send to LifecycleManager (deduplicated by partitionId).
LifecycleManager constructs ChangeLocationsCallContext and after all locations
are notified, it replies to ShuffleClientImpl.
### Why are the changes needed?
In my test 3T TPCDS q23a with 3 Celeborn workers, when kill a worker, the
LifecycleManger will receive 4.8w Revive requests:
```
[emr-usermaster-1-1 logs]$ cat
spark-emr-user-org.apache.spark.sql.hive.thriftserver.HiveThriftServer2-1-master-1-1.c-fa08904e94c028d1.out.1
|grep -i revive |wc -l
64364
```
After this PR, number of ReviveBatch requests reduces to 708:
```
[emr-usermaster-1-1 logs]$ cat
spark-emr-user-org.apache.spark.sql.hive.thriftserver.HiveThriftServer2-1-master-1-1.c-fa08904e94c028d1.out
|grep -i revive |wc -l
2573
```
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Manual test. I have tested:
1. Disable graceful shutdown, kill one worker, job succeeds
2. Disable graceful shutdown, kill two workers successively, job fails as
expected
3. Enable graceful shutdown, restart two workers successively, job succeeds
4. Enable graceful shutdown, restart two workers successively, then kill
the third one, job succeeds
Closes #1588 from waitinfuture/656-2.
Lead-authored-by: zky.zhoukeyong <[email protected]>
Co-authored-by: Keyong Zhou <[email protected]>
Co-authored-by: Keyong Zhou <[email protected]>
Signed-off-by: Shuang <[email protected]>
(cherry picked from commit 57b0e815cf861ee975a70c0d9744cb04a3db2393)
Signed-off-by: Shuang <[email protected]>
---
.../flink/readclient/FlinkShuffleClientImpl.java | 32 +-
.../org/apache/celeborn/client/ReviveManager.java | 128 +++++++
.../apache/celeborn/client/ShuffleClientImpl.java | 423 +++++++++++++++------
.../celeborn/client/ChangePartitionManager.scala | 8 +-
.../apache/celeborn/client/LifecycleManager.scala | 118 +++---
.../client/RequestLocationCallContext.scala | 36 +-
.../celeborn/client/ShuffleClientHelper.scala | 9 +-
.../celeborn/client/WorkerStatusTracker.scala | 5 +-
.../network/client/TransportResponseHandler.java | 6 +
.../celeborn/common/protocol/ReviveRequest.java | 49 +++
.../common/protocol/message/StatusCode.java | 3 +-
common/src/main/proto/TransportMessages.proto | 27 +-
.../org/apache/celeborn/common/CelebornConf.scala | 19 +
.../common/protocol/message/ControlMessages.scala | 49 ++-
.../org/apache/celeborn/common/util/Utils.scala | 2 +
docs/configuration/client.md | 2 +
.../celeborn/tests/spark/RetryReviveTest.scala | 3 +-
.../service/deploy/worker/PushDataHandler.scala | 7 +-
18 files changed, 689 insertions(+), 237 deletions(-)
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
index c1e93687f..7d2b11e0c 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
@@ -19,7 +19,7 @@ package org.apache.celeborn.plugin.flink.readclient;
import java.io.IOException;
import java.nio.ByteBuffer;
-import java.util.Optional;
+import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
@@ -49,7 +49,9 @@ import
org.apache.celeborn.common.network.protocol.RegionFinish;
import org.apache.celeborn.common.network.protocol.RegionStart;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbChangeLocationPartitionInfo;
import org.apache.celeborn.common.protocol.PbChangeLocationResponse;
+import org.apache.celeborn.common.protocol.ReviveRequest;
import org.apache.celeborn.common.protocol.TransportModuleConstants;
import org.apache.celeborn.common.protocol.message.ControlMessages;
import org.apache.celeborn.common.protocol.message.StatusCode;
@@ -376,22 +378,30 @@ public class FlinkShuffleClientImpl extends
ShuffleClientImpl {
if (regionStartResponse.hasRemaining()
&& regionStartResponse.get() ==
StatusCode.HARD_SPLIT.getValue()) {
// if split then revive
+ Set<Integer> mapIds = new HashSet<>();
+ mapIds.add(mapId);
+ List<ReviveRequest> requests = new ArrayList<>();
+ ReviveRequest req =
+ new ReviveRequest(
+ shuffleId,
+ mapId,
+ attemptId,
+ location.getId(),
+ location.getEpoch(),
+ location,
+ StatusCode.HARD_SPLIT);
+ requests.add(req);
PbChangeLocationResponse response =
driverRssMetaService.askSync(
- ControlMessages.Revive$.MODULE$.apply(
- shuffleId,
- mapId,
- attemptId,
- location.getId(),
- location.getEpoch(),
- location,
- StatusCode.HARD_SPLIT),
+ ControlMessages.Revive$.MODULE$.apply(shuffleId, mapIds,
requests),
conf.clientRpcRequestPartitionLocationRpcAskTimeout(),
ClassTag$.MODULE$.apply(PbChangeLocationResponse.class));
// per partitionKey only serve single PartitionLocation in Client
Cache.
- StatusCode respStatus = Utils.toStatusCode(response.getStatus());
+ PbChangeLocationPartitionInfo partitionInfo =
response.getPartitionInfo(0);
+ StatusCode respStatus =
Utils.toStatusCode(partitionInfo.getStatus());
if (StatusCode.SUCCESS.equals(respStatus)) {
- return
Optional.of(PbSerDeUtils.fromPbPartitionLocation(response.getLocation()));
+ return Optional.of(
+
PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition()));
} else {
// throw exception
logger.error(
diff --git a/client/src/main/java/org/apache/celeborn/client/ReviveManager.java
b/client/src/main/java/org/apache/celeborn/client/ReviveManager.java
new file mode 100644
index 000000000..6d94f4755
--- /dev/null
+++ b/client/src/main/java/org/apache/celeborn/client/ReviveManager.java
@@ -0,0 +1,128 @@
+/*
+ * 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.client;
+
+import java.util.*;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.ReviveRequest;
+import org.apache.celeborn.common.protocol.message.StatusCode;
+import org.apache.celeborn.common.util.ThreadUtils;
+
+class ReviveManager {
+ private static final Logger logger =
LoggerFactory.getLogger(ShuffleClientImpl.class);
+
+ LinkedBlockingQueue<ReviveRequest> requestQueue = new
LinkedBlockingQueue<>();
+ private final long interval;
+ private final int batchSize;
+ ShuffleClientImpl shuffleClient;
+ private ScheduledExecutorService batchReviveRequestScheduler =
+
ThreadUtils.newDaemonSingleThreadScheduledExecutor("batch-revive-scheduler");
+
+ public ReviveManager(ShuffleClientImpl shuffleClient, CelebornConf conf) {
+ this.shuffleClient = shuffleClient;
+ this.interval = conf.clientPushReviveInterval();
+ this.batchSize = conf.clientPushReviveBatchSize();
+
+ batchReviveRequestScheduler.scheduleAtFixedRate(
+ () -> {
+ Map<Integer, Set<ReviveRequest>> shuffleMap = new HashMap<>();
+ do {
+ ArrayList<ReviveRequest> batchRequests = new ArrayList<>();
+ requestQueue.drainTo(batchRequests, batchSize);
+ for (ReviveRequest req : batchRequests) {
+ Set<ReviveRequest> set =
+ shuffleMap.computeIfAbsent(req.shuffleId, id -> new
HashSet<>());
+ set.add(req);
+ }
+ for (Map.Entry<Integer, Set<ReviveRequest>> shuffleEntry :
shuffleMap.entrySet()) {
+ // Call reviveBatch for requests in the same (appId, shuffleId)
+ int shuffleId = shuffleEntry.getKey();
+ Set<ReviveRequest> requests = shuffleEntry.getValue();
+ Set<Integer> mapIds = new HashSet<>();
+ ArrayList<ReviveRequest> filteredRequests = new ArrayList<>();
+ Map<Integer, ReviveRequest> requestsToSend = new HashMap<>();
+
+ Map<Integer, PartitionLocation> partitionMap =
+ shuffleClient.reducePartitionMap.get(shuffleId);
+ // Insert request that is not MapperEnded and with the max epoch
+ // into requestsToSend
+ Iterator<ReviveRequest> iter = requests.iterator();
+ while (iter.hasNext()) {
+ ReviveRequest req = iter.next();
+ if (shuffleClient.newerPartitionLocationExists(
+ partitionMap, req.partitionId, req.epoch, false)
+ || shuffleClient.mapperEnded(shuffleId, req.mapId)) {
+ req.reviveStatus = StatusCode.SUCCESS.getValue();
+ } else {
+ filteredRequests.add(req);
+ mapIds.add(req.mapId);
+ PartitionLocation loc = req.loc;
+ if (!requestsToSend.containsKey(req.partitionId)
+ || requestsToSend.get(req.partitionId).epoch <
req.epoch) {
+ requestsToSend.put(req.partitionId, req);
+ }
+ }
+ }
+
+ if (!requestsToSend.isEmpty()) {
+ // Call reviveBatch. Return null means Exception caught or
+ // SHUFFLE_NOT_REGISTERED
+ Map<Integer, Integer> results =
+ shuffleClient.reviveBatch(shuffleId, mapIds,
requestsToSend.values());
+ if (results == null) {
+ for (ReviveRequest req : filteredRequests) {
+ req.reviveStatus = StatusCode.REVIVE_FAILED.getValue();
+ }
+ } else {
+ for (ReviveRequest req : filteredRequests) {
+ if (shuffleClient.mapperEnded(shuffleId, req.mapId)) {
+ req.reviveStatus = StatusCode.SUCCESS.getValue();
+ } else {
+ req.reviveStatus = results.get(req.partitionId);
+ }
+ }
+ }
+ }
+ }
+ // break the loop if remaining requests is less than half of
+ // `celeborn.client.push.revive.batchSize`
+ } while (requestQueue.size() > batchSize / 2);
+ },
+ interval,
+ interval,
+ TimeUnit.MILLISECONDS);
+ }
+
+ public void addRequest(ReviveRequest request) {
+ shuffleClient.excludeWorkerByCause(request.cause, request.loc);
+ // This sync is necessary to ensure the add action is atomic
+ try {
+ requestQueue.put(request);
+ } catch (InterruptedException e) {
+ logger.error("Exception when put into requests!", e);
+ }
+ }
+}
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 61e42714e..70dd9cb59 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -87,7 +87,7 @@ public class ShuffleClientImpl extends ShuffleClient {
protected final int BATCH_HEADER_SIZE = 4 * 4;
// key: shuffleId, value: (partitionId, PartitionLocation)
- private final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>>
reducePartitionMap =
+ final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>>
reducePartitionMap =
JavaUtils.newConcurrentHashMap();
// key: shuffleId, value: Set(mapId)
@@ -112,7 +112,7 @@ public class ShuffleClientImpl extends ShuffleClient {
protected final String appUniqueId;
- ThreadLocal<Compressor> compressorThreadLocal =
+ private ThreadLocal<Compressor> compressorThreadLocal =
new ThreadLocal<Compressor>() {
@Override
protected Compressor initialValue() {
@@ -120,6 +120,8 @@ public class ShuffleClientImpl extends ShuffleClient {
}
};
+ private final ReviveManager reviveManager;
+
protected static class ReduceFileGroups {
public Map<Integer, Set<PartitionLocation>> partitionGroups;
public int[] mapAttempts;
@@ -187,6 +189,7 @@ public class ShuffleClientImpl extends ShuffleClient {
partitionSplitPool =
ThreadUtils.newDaemonCachedThreadPool(
"celeborn-shuffle-split", pushSplitPartitionThreads, 60);
+ reviveManager = new ReviveManager(this, conf);
logger.info("Created ShuffleClientImpl, appUniqueId: {}", appUniqueId);
}
@@ -207,20 +210,32 @@ public class ShuffleClientImpl extends ShuffleClient {
private void submitRetryPushData(
int shuffleId,
- int mapId,
- int attemptId,
byte[] body,
int batchId,
- PartitionLocation loc,
RpcResponseCallback wrappedCallback,
PushState pushState,
- StatusCode cause,
- int remainReviveTimes) {
+ ReviveRequest request,
+ int remainReviveTimes,
+ long dueTime) {
+ int mapId = request.mapId;
+ int attemptId = request.attemptId;
+ PartitionLocation loc = request.loc;
+ StatusCode cause = request.cause;
int partitionId = loc.getId();
- if (!revive(shuffleId, mapId, attemptId, partitionId, loc.getEpoch(), loc,
cause)) {
- wrappedCallback.onFailure(
- new CelebornIOException(cause + " then revive but " +
StatusCode.REVIVE_FAILED));
- } else if (mapperEnded(shuffleId, mapId)) {
+ long reviveWaitTime = dueTime - System.currentTimeMillis();
+ final long delta = 50;
+ long accumulatedTime = 0;
+ while (request.reviveStatus == StatusCode.REVIVE_INITIALIZED.getValue()
+ && accumulatedTime <= reviveWaitTime) {
+ try {
+ Thread.sleep(delta);
+ accumulatedTime += delta;
+ } catch (InterruptedException e) {
+ logger.error("Interrupted while waiting for Revive result!");
+ Thread.currentThread().interrupt();
+ }
+ }
+ if (mapperEnded(shuffleId, mapId)) {
logger.debug(
"Revive for push data success, but the mapper already ended for
shuffle {} map {} attempt {} partition {} batch {} location {}.",
shuffleId,
@@ -230,6 +245,19 @@ public class ShuffleClientImpl extends ShuffleClient {
batchId,
loc);
pushState.removeBatch(batchId, loc.hostAndPushPort());
+ } else if (request.reviveStatus != StatusCode.SUCCESS.getValue()) {
+ wrappedCallback.onFailure(
+ new CelebornIOException(
+ cause
+ + " then revive but "
+ + StatusCode.REVIVE_FAILED
+ + ", revive status "
+ + request.reviveStatus
+ + "("
+ + Utils.toStatusCode(request.reviveStatus)
+ + ")"
+ + ", old location: "
+ + request.loc));
} else {
PartitionLocation newLoc =
reducePartitionMap.get(shuffleId).get(partitionId);
logger.info(
@@ -273,6 +301,24 @@ public class ShuffleClientImpl extends ShuffleClient {
}
}
+ public ReviveRequest[] addAndGetReviveRequests(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ ArrayList<DataBatches.DataBatch> batches,
+ StatusCode cause) {
+ ReviveRequest[] reviveRequests = new ReviveRequest[batches.size()];
+ for (int i = 0; i < batches.size(); i++) {
+ DataBatches.DataBatch batch = batches.get(i);
+ PartitionLocation loc = batch.loc;
+ ReviveRequest reviveRequest =
+ new ReviveRequest(shuffleId, mapId, attemptId, loc.getId(),
loc.getEpoch(), loc, cause);
+ reviveManager.addRequest(reviveRequest);
+ reviveRequests[i] = reviveRequest;
+ }
+ return reviveRequests;
+ }
+
private void submitRetryPushMergedData(
PushState pushState,
int shuffleId,
@@ -281,49 +327,89 @@ public class ShuffleClientImpl extends ShuffleClient {
ArrayList<DataBatches.DataBatch> batches,
StatusCode cause,
Integer oldGroupedBatchId,
- int remainReviveTimes) {
+ ReviveRequest[] reviveRequests,
+ int remainReviveTimes,
+ long reviveResponseDueTime) {
HashMap<Pair<String, String>, DataBatches> newDataBatchesMap = new
HashMap<>();
ArrayList<DataBatches.DataBatch> reviveFailedBatchesMap = new
ArrayList<>();
- for (DataBatches.DataBatch batch : batches) {
- int partitionId = batch.loc.getId();
- if (!revive(
- shuffleId, mapId, attemptId, partitionId, batch.loc.getEpoch(),
batch.loc, cause)) {
- if (remainReviveTimes > 0) {
- reviveFailedBatchesMap.add(batch);
+ long reviveWaitTime = reviveResponseDueTime - System.currentTimeMillis();
+ final long delta = 50;
+ long accumulatedTime = 0;
+ int index = 0;
+ while (index < reviveRequests.length && accumulatedTime <= reviveWaitTime)
{
+ ReviveRequest request = reviveRequests[index];
+ DataBatches.DataBatch batch = batches.get(index);
+ if (request.reviveStatus != StatusCode.REVIVE_INITIALIZED.getValue()) {
+ if (mapperEnded(shuffleId, mapId)) {
+ logger.debug(
+ "Revive for push merged data success, but the mapper already
ended for shuffle {} map {} attempt {} partition {} batch {}.",
+ shuffleId,
+ mapId,
+ attemptId,
+ request.partitionId,
+ oldGroupedBatchId);
+ } else if (request.reviveStatus == StatusCode.SUCCESS.getValue()) {
+ PartitionLocation newLoc =
reducePartitionMap.get(shuffleId).get(request.partitionId);
+ DataBatches newDataBatches =
+ newDataBatchesMap.computeIfAbsent(genAddressPair(newLoc), (s) ->
new DataBatches());
+ newDataBatches.addDataBatch(newLoc, batch.batchId, batch.body);
} else {
- String errorMsg =
- String.format(
- "Revive failed while pushing merged for shuffle %d map %d
attempt %d partition %d batch %d location %s.",
- shuffleId, mapId, attemptId, partitionId, oldGroupedBatchId,
batch.loc);
- pushState.exception.compareAndSet(
- null,
- new CelebornIOException(
- errorMsg,
- new CelebornIOException(cause + " then revive but " +
StatusCode.REVIVE_FAILED)));
- return;
+ if (remainReviveTimes > 0) {
+ reviveFailedBatchesMap.add(batch);
+ } else {
+ String errorMsg =
+ String.format(
+ "Revive failed while pushing merged for shuffle %d map %d
attempt %d partition %d batch %d location %s.",
+ shuffleId, mapId, attemptId, request.partitionId,
oldGroupedBatchId, batch.loc);
+ pushState.exception.compareAndSet(
+ null,
+ new CelebornIOException(
+ errorMsg,
+ new CelebornIOException(
+ cause
+ + " then revive but "
+ + request.reviveStatus
+ + "("
+ + Utils.toStatusCode(request.reviveStatus)
+ + ")")));
+ return;
+ }
}
- } else if (mapperEnded(shuffleId, mapId)) {
- logger.debug(
- "Revive for push merged data success, but the mapper already ended
for shuffle {} map {} attempt {} partition {} batch {}.",
- shuffleId,
- mapId,
- attemptId,
- partitionId,
- oldGroupedBatchId);
+ index++;
} else {
- PartitionLocation newLoc =
reducePartitionMap.get(shuffleId).get(partitionId);
- logger.info(
- "Revive for push merged data success, new location for shuffle {}
map {} attempt {} partition {} batch {} is location {}.",
- shuffleId,
- mapId,
- attemptId,
- partitionId,
- oldGroupedBatchId,
- newLoc);
- DataBatches newDataBatches =
- newDataBatchesMap.computeIfAbsent(genAddressPair(newLoc), (s) ->
new DataBatches());
- newDataBatches.addDataBatch(newLoc, batch.batchId, batch.body);
+ try {
+ Thread.sleep(delta);
+ } catch (InterruptedException e) {
+ logger.error("Interrupted while waiting for Revive result!");
+ Thread.currentThread().interrupt();
+ }
+ accumulatedTime += delta;
+ }
+ }
+
+ for (int i = index; i < reviveRequests.length; i++) {
+ ReviveRequest request = reviveRequests[index];
+ DataBatches.DataBatch batch = batches.get(i);
+ if (remainReviveTimes > 0) {
+ reviveFailedBatchesMap.add(batch);
+ } else {
+ String errorMsg =
+ String.format(
+ "Revive failed while pushing merged for shuffle %d map %d
attempt %d partition %d batch %d location %s.",
+ shuffleId, mapId, attemptId, request.partitionId,
oldGroupedBatchId, batch.loc);
+ pushState.exception.compareAndSet(
+ null,
+ new CelebornIOException(
+ errorMsg,
+ new CelebornIOException(
+ cause
+ + " then revive but "
+ + request.reviveStatus
+ + "("
+ + Utils.toStatusCode(request.reviveStatus)
+ + ")")));
+ return;
}
}
@@ -342,6 +428,8 @@ public class ShuffleClientImpl extends ShuffleClient {
if (reviveFailedBatchesMap.isEmpty()) {
pushState.removeBatch(oldGroupedBatchId,
batches.get(0).loc.hostAndPushPort());
} else {
+ ReviveRequest[] requests =
+ addAndGetReviveRequests(shuffleId, mapId, attemptId,
reviveFailedBatchesMap, cause);
pushDataRetryPool.submit(
() ->
submitRetryPushMergedData(
@@ -352,7 +440,12 @@ public class ShuffleClientImpl extends ShuffleClient {
reviveFailedBatchesMap,
cause,
oldGroupedBatchId,
- remainReviveTimes - 1));
+ requests,
+ remainReviveTimes - 1,
+ System.currentTimeMillis()
+ + conf.clientRpcRequestPartitionLocationRpcAskTimeout()
+ .duration()
+ .toMillis()));
}
}
@@ -494,35 +587,39 @@ public class ShuffleClientImpl extends ShuffleClient {
}
}
- private boolean waitRevivedLocation(
- ConcurrentHashMap<Integer, PartitionLocation> map, int partitionId, int
epoch) {
- PartitionLocation currentLocation = map.get(partitionId);
+ /**
+ * check if a newer PartitionLocation(with larger epoch) exists in local
cache
+ *
+ * @param shuffleMap
+ * @param partitionId
+ * @param epoch
+ * @param wait wheter to wait for some time for a newer PartitionLocation
+ * @return
+ */
+ boolean newerPartitionLocationExists(
+ Map<Integer, PartitionLocation> shuffleMap, int partitionId, int epoch,
boolean wait) {
+ PartitionLocation currentLocation = shuffleMap.get(partitionId);
if (currentLocation != null && currentLocation.getEpoch() > epoch) {
return true;
- }
-
- long sleepTimeMs = RND.nextInt(50);
- if (sleepTimeMs > 30) {
- try {
- TimeUnit.MILLISECONDS.sleep(sleepTimeMs);
- } catch (InterruptedException e) {
- logger.error("Waiting revived location was interrupted.", e);
- Thread.currentThread().interrupt();
+ } else if (wait) {
+ long sleepTimeMs = RND.nextInt(50);
+ if (sleepTimeMs > 30) {
+ try {
+ TimeUnit.MILLISECONDS.sleep(sleepTimeMs);
+ } catch (InterruptedException e) {
+ logger.error("Waiting revived location was interrupted.", e);
+ Thread.currentThread().interrupt();
+ }
}
- }
- currentLocation = map.get(partitionId);
- return currentLocation != null && currentLocation.getEpoch() > epoch;
+ currentLocation = shuffleMap.get(partitionId);
+ return currentLocation != null && currentLocation.getEpoch() > epoch;
+ } else {
+ return false;
+ }
}
- private boolean revive(
- int shuffleId,
- int mapId,
- int attemptId,
- int partitionId,
- int epoch,
- PartitionLocation oldLocation,
- StatusCode cause) {
+ void excludeWorkerByCause(StatusCode cause, PartitionLocation oldLocation) {
// Add ShuffleClient side blacklist
if (shuffleClientPushBlacklistEnabled && oldLocation != null) {
if (cause == StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_MASTER) {
@@ -539,18 +636,26 @@ public class ShuffleClientImpl extends ShuffleClient {
blacklist.add(oldLocation.getPeer().hostAndPushPort());
}
}
+ }
+
+ private boolean revive(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ int epoch,
+ PartitionLocation oldLocation,
+ StatusCode cause) {
+ excludeWorkerByCause(cause, oldLocation);
+
+ Set<Integer> mapIds = new HashSet<>();
+ mapIds.add(mapId);
+ List<ReviveRequest> requests = new ArrayList<>();
+ ReviveRequest req =
+ new ReviveRequest(shuffleId, mapId, attemptId, partitionId, epoch,
oldLocation, cause);
+ requests.add(req);
+ Map<Integer, Integer> results = reviveBatch(shuffleId, mapIds, requests);
- ConcurrentHashMap<Integer, PartitionLocation> map =
reducePartitionMap.get(shuffleId);
- if (waitRevivedLocation(map, partitionId, epoch)) {
- logger.debug(
- "Revive already success for shuffle {} map {} attempt {} partition
{} epoch {}, just return true(Assume revive successfully).",
- shuffleId,
- mapId,
- attemptId,
- partitionId,
- epoch);
- return true;
- }
if (mapperEnded(shuffleId, mapId)) {
logger.debug(
"Revive success, but the mapper ended for shuffle {} map {} attempt
{} partition {}, just return true(Assume revive successfully).",
@@ -559,48 +664,80 @@ public class ShuffleClientImpl extends ShuffleClient {
attemptId,
partitionId);
return true;
+ } else if (results == null || results.get(partitionId) !=
StatusCode.SUCCESS.getValue()) {
+ return false;
+ } else {
+ return true;
}
+ }
+
+ /** @return partitionId -> StatusCode#getValue */
+ Map<Integer, Integer> reviveBatch(
+ int shuffleId, Set<Integer> mapIds, Collection<ReviveRequest> requests) {
+ // partitionId -> StatusCode#getValue
+ Map<Integer, Integer> results = new HashMap<>();
+ // Local cached map of (partitionId -> PartitionLocation)
+ ConcurrentHashMap<Integer, PartitionLocation> partitionLocationMap =
+ reducePartitionMap.get(shuffleId);
+
+ Map<Integer, PartitionLocation> oldLocMap = new HashMap<>();
+ Iterator<ReviveRequest> iter = requests.iterator();
+ while (iter.hasNext()) {
+ ReviveRequest req = iter.next();
+ oldLocMap.put(req.partitionId, req.loc);
+ }
try {
PbChangeLocationResponse response =
driverRssMetaService.askSync(
- Revive$.MODULE$.apply(
- shuffleId, mapId, attemptId, partitionId, epoch,
oldLocation, cause),
+ Revive$.MODULE$.apply(shuffleId, mapIds, requests),
conf.clientRpcRequestPartitionLocationRpcAskTimeout(),
ClassTag$.MODULE$.apply(PbChangeLocationResponse.class));
- // per partitionKey only serve single PartitionLocation in Client Cache.
- StatusCode respStatus = Utils.toStatusCode(response.getStatus());
- if (response.getAvailable()) {
- blacklist.remove(oldLocation.hostAndPushPort());
- }
- if (StatusCode.SUCCESS.equals(respStatus)) {
- PartitionLocation newLocation =
- PbSerDeUtils.fromPbPartitionLocation(response.getLocation());
- map.put(partitionId, newLocation);
- blacklist.remove(newLocation.hostAndPushPort());
- return true;
- } else if (StatusCode.MAP_ENDED.equals(respStatus)) {
- logger.debug(
- "Revive success, but the mapper ended for shuffle {} map {}
attempt {} partition {}, just return true(Assume revive successfully).",
- shuffleId,
- mapId,
- attemptId,
- partitionId);
+
+ for (int i = 0; i < response.getEndedMapIdCount(); i++) {
+ int mapId = response.getEndedMapId(i);
mapperEndMap.computeIfAbsent(shuffleId, (id) ->
ConcurrentHashMap.newKeySet()).add(mapId);
- return true;
- } else {
- return false;
}
+
+ for (int i = 0; i < response.getPartitionInfoCount(); i++) {
+ PbChangeLocationPartitionInfo partitionInfo =
response.getPartitionInfo(i);
+ int partitionId = partitionInfo.getPartitionId();
+ int statusCode = partitionInfo.getStatus();
+ if (partitionInfo.getOldAvailable()) {
+ blacklist.remove(oldLocMap.get(partitionId).hostAndPushPort());
+ }
+
+ if (StatusCode.SUCCESS.getValue() == statusCode) {
+ PartitionLocation loc =
+
PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition());
+ partitionLocationMap.put(partitionId, loc);
+ blacklist.remove(loc.hostAndPushPort());
+ } else if (StatusCode.STAGE_ENDED.getValue() == statusCode) {
+ stageEnded(shuffleId);
+ return results;
+ } else if (StatusCode.SHUFFLE_NOT_REGISTERED.getValue() == statusCode)
{
+ logger.error("SHUFFLE_NOT_REGISTERED!");
+ return null;
+ }
+ results.put(partitionId, statusCode);
+ }
+
+ return results;
} catch (Exception e) {
+ StringBuilder partitionIds = new StringBuilder();
+ StringBuilder epochs = new StringBuilder();
+ requests.forEach(
+ (req) -> {
+ partitionIds.append(req.partitionId).append(",");
+ epochs.append(req.epoch).append(",");
+ });
logger.error(
- "Exception raised while reviving for shuffle {} map {} attempt {}
partition {} epoch {}.",
+ "Exception raised while reviving for shuffle {} partitionIds {}
epochs {}.",
shuffleId,
- mapId,
- attemptId,
- partitionId,
- epoch,
+ partitionIds,
+ epochs,
e);
- return false;
+ return null;
}
}
@@ -769,19 +906,32 @@ public class ShuffleClientImpl extends ShuffleClient {
attemptId,
partitionId,
nextBatchId);
+ ReviveRequest reviveRequest =
+ new ReviveRequest(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ loc.getEpoch(),
+ loc,
+ StatusCode.HARD_SPLIT);
+ reviveManager.addRequest(reviveRequest);
+ long dueTime =
+ System.currentTimeMillis()
+ +
conf.clientRpcRequestPartitionLocationRpcAskTimeout()
+ .duration()
+ .toMillis();
pushDataRetryPool.submit(
() ->
submitRetryPushData(
shuffleId,
- mapId,
- attemptId,
body,
nextBatchId,
- loc,
this,
pushState,
- StatusCode.HARD_SPLIT,
- remainReviveTimes));
+ reviveRequest,
+ remainReviveTimes,
+ dueTime));
} else if (reason ==
StatusCode.PUSH_DATA_SUCCESS_MASTER_CONGESTED.getValue()) {
logger.debug(
"Push data to {} master congestion required for shuffle
{} map {} attempt {} partition {} batch {}.",
@@ -846,19 +996,26 @@ public class ShuffleClientImpl extends ShuffleClient {
// async retry push data
if (!mapperEnded(shuffleId, mapId)) {
remainReviveTimes = remainReviveTimes - 1;
+ ReviveRequest reviveRequest =
+ new ReviveRequest(
+ shuffleId, mapId, attemptId, partitionId,
loc.getEpoch(), loc, cause);
+ reviveManager.addRequest(reviveRequest);
+ long dueTime =
+ System.currentTimeMillis()
+ + conf.clientRpcRequestPartitionLocationRpcAskTimeout()
+ .duration()
+ .toMillis();
pushDataRetryPool.submit(
() ->
submitRetryPushData(
shuffleId,
- mapId,
- attemptId,
body,
nextBatchId,
- loc,
this,
pushState,
- cause,
- remainReviveTimes));
+ reviveRequest,
+ remainReviveTimes,
+ dueTime));
} else {
pushState.removeBatch(nextBatchId, loc.hostAndPushPort());
logger.info(
@@ -1134,6 +1291,10 @@ public class ShuffleClientImpl extends ShuffleClient {
Arrays.toString(partitionIds),
groupedBatchId,
Arrays.toString(batchIds));
+
+ ReviveRequest[] requests =
+ addAndGetReviveRequests(
+ shuffleId, mapId, attemptId, batches,
StatusCode.HARD_SPLIT);
pushDataRetryPool.submit(
() ->
submitRetryPushMergedData(
@@ -1144,7 +1305,12 @@ public class ShuffleClientImpl extends ShuffleClient {
batches,
StatusCode.HARD_SPLIT,
groupedBatchId,
- remainReviveTimes));
+ requests,
+ remainReviveTimes,
+ System.currentTimeMillis()
+ +
conf.clientRpcRequestPartitionLocationRpcAskTimeout()
+ .duration()
+ .toMillis()));
} else if (reason ==
StatusCode.PUSH_DATA_SUCCESS_MASTER_CONGESTED.getValue()) {
logger.debug(
"Push merged data to {} master congestion required for
shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}.",
@@ -1208,6 +1374,8 @@ public class ShuffleClientImpl extends ShuffleClient {
remainReviveTimes,
e);
if (!mapperEnded(shuffleId, mapId)) {
+ ReviveRequest[] requests =
+ addAndGetReviveRequests(shuffleId, mapId, attemptId,
batches, cause);
pushDataRetryPool.submit(
() ->
submitRetryPushMergedData(
@@ -1218,7 +1386,12 @@ public class ShuffleClientImpl extends ShuffleClient {
batches,
cause,
groupedBatchId,
- remainReviveTimes - 1));
+ requests,
+ remainReviveTimes - 1,
+ System.currentTimeMillis()
+ +
conf.clientRpcRequestPartitionLocationRpcAskTimeout()
+ .duration()
+ .toMillis()));
} else {
pushState.removeBatch(groupedBatchId, hostPort);
logger.info(
@@ -1461,7 +1634,7 @@ public class ShuffleClientImpl extends ShuffleClient {
driverRssMetaService = endpointRef;
}
- protected boolean mapperEnded(int shuffleId, int mapId) {
+ boolean mapperEnded(int shuffleId, int mapId) {
return (mapperEndMap.containsKey(shuffleId) &&
mapperEndMap.get(shuffleId).contains(mapId))
|| stageEnded(shuffleId);
}
diff --git
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
index 557cc3239..164d041d9 100644
---
a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
@@ -161,6 +161,7 @@ class ChangePartitionManager(
// Else register and allocate for it.
getLatestPartition(shuffleId, partitionId, oldEpoch).foreach {
latestLoc =>
context.reply(
+ partitionId,
StatusCode.SUCCESS,
Some(latestLoc),
lifecycleManager.workerStatusTracker.workerAvailable(oldPartition))
@@ -224,12 +225,12 @@ class ChangePartitionManager(
location -> Option(requestsMap.remove(location.getId))
}
}.foreach { case (newLocation, requests) =>
- requests.foreach(_.asScala.foreach { req =>
+ requests.map(_.asScala.toList.foreach(req =>
req.context.reply(
+ req.partitionId,
StatusCode.SUCCESS,
Option(newLocation),
-
lifecycleManager.workerStatusTracker.workerAvailable(req.oldPartition))
- })
+
lifecycleManager.workerStatusTracker.workerAvailable(req.oldPartition))))
}
}
@@ -245,6 +246,7 @@ class ChangePartitionManager(
}.foreach { requests =>
requests.map(_.asScala.toList.foreach(req =>
req.context.reply(
+ req.partitionId,
status,
None,
lifecycleManager.workerStatusTracker.workerAvailable(req.oldPartition))))
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index 840a85f27..fece54cf2 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -235,29 +235,33 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
case pb: PbRevive =>
val shuffleId = pb.getShuffleId
- val mapId = pb.getMapId
- val attemptId = pb.getAttemptId
- val partitionId = pb.getPartitionId
- val epoch = pb.getEpoch
- val oldPartition =
- if (pb.hasOldPartition) {
- PbSerDeUtils.fromPbPartitionLocation(pb.getOldPartition)
+ val mapIds = pb.getMapIdList
+ val partitionInfos = pb.getPartitionInfoList
+
+ val partitionIds = new util.ArrayList[Integer]()
+ val epochs = new util.ArrayList[Integer]()
+ val oldPartitions = new util.ArrayList[PartitionLocation]()
+ val causes = new util.ArrayList[StatusCode]()
+ (0 until partitionInfos.size()).foreach { idx =>
+ val info = partitionInfos.get(idx)
+ partitionIds.add(info.getPartitionId)
+ epochs.add(info.getEpoch)
+ if (info.hasPartition) {
+
oldPartitions.add(PbSerDeUtils.fromPbPartitionLocation(info.getPartition))
} else {
- null
+ oldPartitions.add(null)
}
- val cause = Utils.toStatusCode(pb.getStatus)
- logTrace(s"Received Revive request, " +
- s"$shuffleId, $mapId, $attemptId ,$partitionId," +
- s" $epoch, $oldPartition, $cause.")
+ causes.add(Utils.toStatusCode(info.getStatus))
+ }
+ logWarning(s"Received Revive request, number of partitions
${partitionIds.size()}")
handleRevive(
context,
shuffleId,
- mapId,
- attemptId,
- partitionId,
- epoch,
- oldPartition,
- cause)
+ mapIds,
+ partitionIds,
+ epochs,
+ oldPartitions,
+ causes)
case pb: PbPartitionSplit =>
val shuffleId = pb.getShuffleId
@@ -267,7 +271,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
logTrace(s"Received split request, " +
s"$shuffleId, $partitionId, $epoch, $oldPartition")
changePartitionManager.handleRequestPartitionLocation(
- ChangeLocationCallContext(context),
+ ChangeLocationsCallContext(context, 1),
shuffleId,
partitionId,
epoch,
@@ -487,51 +491,51 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
private def handleRevive(
context: RpcCallContext,
shuffleId: Int,
- mapId: Int,
- attemptId: Int,
- partitionId: Int,
- oldEpoch: Int,
- oldPartition: PartitionLocation,
- cause: StatusCode): Unit = {
+ mapIds: util.List[Integer],
+ partitionIds: util.List[Integer],
+ oldEpochs: util.List[Integer],
+ oldPartitions: util.List[PartitionLocation],
+ causes: util.List[StatusCode]): Unit = {
+ val contextWrapper =
+ ChangeLocationsCallContext(context, partitionIds.size())
// If shuffle not registered, reply ShuffleNotRegistered and return
if (!registeredShuffle.contains(shuffleId)) {
logError(s"[handleRevive] shuffle $shuffleId not registered!")
- context.reply(ChangeLocationResponse(
+ contextWrapper.reply(
+ -1,
StatusCode.SHUFFLE_NOT_REGISTERED,
None,
- workerStatusTracker.workerAvailable(oldPartition)))
+ false)
return
}
- if (getPartitionType(shuffleId) == PartitionType.MAP) {
- logError(s"[handleRevive] shuffle $shuffleId revived filed, because map
partition don't support revive!")
- context.reply(ChangeLocationResponse(
- StatusCode.REVIVE_FAILED,
+ if (commitManager.isStageEnd(shuffleId)) {
+ logError(s"[handleRevive] shuffle $shuffleId stage ended!")
+ contextWrapper.reply(
+ -1,
+ StatusCode.STAGE_ENDED,
None,
- workerStatusTracker.workerAvailable(oldPartition)))
+ false)
return
}
- if (commitManager.isMapperEnded(shuffleId, mapId)) {
- logWarning(s"[handleRevive] Mapper ended, mapId $mapId, current
attemptId $attemptId, " +
- s"ended attemptId
${commitManager.getMapperAttempts(shuffleId)(mapId)}, shuffleId $shuffleId.")
- context.reply(ChangeLocationResponse(
- StatusCode.MAP_ENDED,
- None,
- workerStatusTracker.workerAvailable(oldPartition)))
- return
+ mapIds.asScala.foreach { mapId =>
+ if (commitManager.isMapperEnded(shuffleId, mapId)) {
+ logWarning(s"[handleRevive] Mapper ended, mapId $mapId, ended
attemptId ${commitManager.getMapperAttempts(
+ shuffleId)(mapId)}, shuffleId $shuffleId")
+ contextWrapper.markMapperEnd(mapId)
+ }
}
- logWarning(s"Do Revive for shuffle shuffleId $shuffleId, " +
- s"oldPartition: $oldPartition, cause: $cause")
-
- changePartitionManager.handleRequestPartitionLocation(
- ChangeLocationCallContext(context),
- shuffleId,
- partitionId,
- oldEpoch,
- oldPartition,
- Some(cause))
+ (0 until partitionIds.size()).foreach { idx =>
+ changePartitionManager.handleRequestPartitionLocation(
+ contextWrapper,
+ shuffleId,
+ partitionIds.get(idx),
+ oldEpochs.get(idx),
+ oldPartitions.get(idx),
+ Some(causes.get(idx)))
+ }
}
private def handleMapperEnd(
@@ -724,9 +728,15 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
logWarning(s"Cannot find workInfo for $shuffleId from previous
success workResource:" +
s" ${destroyWorkerInfo.readableAddress()}, init according to
partition info")
try {
- destroyWorkerInfo.endpoint = rpcEnv.setupEndpointRef(
- RpcAddress.apply(destroyWorkerInfo.host,
destroyWorkerInfo.rpcPort),
- WORKER_EP)
+ if (workerStatusTracker.workerAvailable(destroyWorkerInfo)) {
+ destroyWorkerInfo.endpoint = rpcEnv.setupEndpointRef(
+ RpcAddress.apply(destroyWorkerInfo.host,
destroyWorkerInfo.rpcPort),
+ WORKER_EP)
+ } else {
+ logInfo(
+ s"${destroyWorkerInfo.toUniqueId()} is unavailable, set
destroyWorkerInfo to null")
+ destroyWorkerInfo = null
+ }
} catch {
case t: Throwable =>
logError(
@@ -767,7 +777,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
}
val msg = ReleaseSlots(appUniqueId, shuffleId, workerIds,
workerSlotsPerDisk)
requestMasterReleaseSlots(msg)
- logInfo(s"Released slots for reserve buffer failed workers " +
+ logDebug(s"Released slots for reserve buffer failed workers " +
s"${workerIds.asScala.mkString(",")}" +
s"${slots.asScala.mkString(",")}" +
s"shuffleId $shuffleId")
}
diff --git
a/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
b/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
index fa3a9eb95..d5e6c4c36 100644
---
a/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/RequestLocationCallContext.scala
@@ -17,6 +17,10 @@
package org.apache.celeborn.client
+import java.util
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.common.protocol.PartitionLocation
import
org.apache.celeborn.common.protocol.message.ControlMessages.{ChangeLocationResponse,
RegisterShuffleResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
@@ -24,25 +28,47 @@ import org.apache.celeborn.common.rpc.RpcCallContext
trait RequestLocationCallContext {
def reply(
+ partitionId: Int,
status: StatusCode,
partitionLocationOpt: Option[PartitionLocation],
- excluded: Boolean): Unit
+ available: Boolean): Unit
}
-case class ChangeLocationCallContext(context: RpcCallContext) extends
RequestLocationCallContext {
+case class ChangeLocationsCallContext(
+ context: RpcCallContext,
+ partitionCount: Int)
+ extends RequestLocationCallContext with Logging {
+ val endedMapIds = new util.HashSet[Integer]()
+ val newLocs =
+ new ConcurrentHashMap[Integer, (StatusCode, Boolean,
PartitionLocation)](partitionCount)
+
+ def markMapperEnd(mapId: Int): Unit = this.synchronized {
+ endedMapIds.add(mapId)
+ }
+
override def reply(
+ partitionId: Int,
status: StatusCode,
partitionLocationOpt: Option[PartitionLocation],
- excluded: Boolean): Unit = {
- context.reply(ChangeLocationResponse(status, partitionLocationOpt,
excluded))
+ available: Boolean): Unit = this.synchronized {
+ if (newLocs.containsKey(partitionId)) {
+ logError(s"PartitionId $partitionId already exists!")
+ }
+ newLocs.put(partitionId, (status, available,
partitionLocationOpt.getOrElse(null)))
+
+ if (newLocs.size() == partitionCount || StatusCode.SHUFFLE_NOT_REGISTERED
== status
+ || StatusCode.STAGE_ENDED == status) {
+ context.reply(ChangeLocationResponse(endedMapIds, newLocs))
+ }
}
}
case class ApplyNewLocationCallContext(context: RpcCallContext) extends
RequestLocationCallContext {
override def reply(
+ partitionId: Int,
status: StatusCode,
partitionLocationOpt: Option[PartitionLocation],
- excluded: Boolean = false): Unit = {
+ available: Boolean): Unit = {
partitionLocationOpt match {
case Some(partitionLocation) =>
context.reply(RegisterShuffleResponse(status,
Array(partitionLocation)))
diff --git
a/client/src/main/scala/org/apache/celeborn/client/ShuffleClientHelper.scala
b/client/src/main/scala/org/apache/celeborn/client/ShuffleClientHelper.scala
index d6b9be376..f8a69468c 100644
--- a/client/src/main/scala/org/apache/celeborn/client/ShuffleClientHelper.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/ShuffleClientHelper.scala
@@ -42,9 +42,14 @@ object ShuffleClientHelper extends Logging {
req,
conf.clientRpcRequestPartitionLocationRpcAskTimeout).onComplete {
case Success(resp) =>
- val respStatus = Utils.toStatusCode(resp.getStatus)
+ val partitionInfo = resp.getPartitionInfo(0)
+ val respStatus = Utils.toStatusCode(partitionInfo.getStatus)
if (respStatus == StatusCode.SUCCESS) {
- shuffleLocs.put(partitionId,
PbSerDeUtils.fromPbPartitionLocation(resp.getLocation))
+ shuffleLocs.put(
+ partitionId,
+ PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition))
+ } else if (respStatus == StatusCode.STAGE_ENDED) {
+ logInfo(s"Stage ended for $shuffleId")
} else {
logInfo(s"split failed for $respStatus, " +
s"shuffle file can be larger than expected, try split again");
diff --git
a/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
b/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
index 113b7d779..74cdeca2a 100644
--- a/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/WorkerStatusTracker.scala
@@ -153,7 +153,7 @@ class WorkerStatusTracker(
def handleHeartbeatResponse(res: HeartbeatFromApplicationResponse): Unit = {
if (res.statusCode == StatusCode.SUCCESS) {
- logDebug(s"Received Blacklist from Master, blacklist: ${res.blacklist} "
+
+ logInfo(s"Received Blacklist from Master, blacklist: ${res.blacklist} " +
s"unknown workers: ${res.unknownWorkers}, shutdown workers:
${res.shuttingWorkers}")
val current = System.currentTimeMillis()
@@ -200,7 +200,8 @@ class WorkerStatusTracker(
}
}
- logDebug(s"Current blacklist $blacklist")
+ logInfo(s"Current blacklist $blacklist, Current shuttingDown
${shuttingWorkers.asScala.map(
+ _.readableAddress()).mkString("\n")}")
}
}
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
index 7a2f8e61e..01b225dc8 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
@@ -143,6 +143,12 @@ public class TransportResponseHandler extends
MessageHandler<ResponseMessage> {
if (info.channelFuture != null) {
info.channelFuture.cancel(true);
}
+ logger.info(
+ "Fail expire fetch request {},{},{},{}",
+ entry.getKey().streamId,
+ entry.getKey().chunkIndex,
+ entry.getKey().offset,
+ entry.getKey().len);
info.callback.onFailure(
entry.getKey().chunkIndex, new
CelebornIOException(StatusCode.FETCH_DATA_TIMEOUT));
info.channelFuture = null;
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/ReviveRequest.java
b/common/src/main/java/org/apache/celeborn/common/protocol/ReviveRequest.java
new file mode 100644
index 000000000..0d001adfa
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/protocol/ReviveRequest.java
@@ -0,0 +1,49 @@
+/*
+ * 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.protocol;
+
+import org.apache.celeborn.common.protocol.message.StatusCode;
+
+public class ReviveRequest {
+ public int shuffleId;
+ public int mapId;
+ public int attemptId;
+ public int partitionId;
+ public int epoch;
+ public PartitionLocation loc;
+ public StatusCode cause;
+ public volatile int reviveStatus;
+
+ public ReviveRequest(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ int epoch,
+ PartitionLocation loc,
+ StatusCode cause) {
+ this.shuffleId = shuffleId;
+ this.mapId = mapId;
+ this.attemptId = attemptId;
+ this.partitionId = partitionId;
+ this.epoch = epoch;
+ this.loc = loc;
+ this.cause = cause;
+ reviveStatus = StatusCode.REVIVE_INITIALIZED.getValue();
+ }
+}
diff --git
a/common/src/main/java/org/apache/celeborn/common/protocol/message/StatusCode.java
b/common/src/main/java/org/apache/celeborn/common/protocol/message/StatusCode.java
index 688f0def8..5ad47bed8 100644
---
a/common/src/main/java/org/apache/celeborn/common/protocol/message/StatusCode.java
+++
b/common/src/main/java/org/apache/celeborn/common/protocol/message/StatusCode.java
@@ -77,7 +77,8 @@ public enum StatusCode {
PUSH_DATA_MASTER_BLACKLISTED(44),
PUSH_DATA_SLAVE_BLACKLISTED(45),
- FETCH_DATA_TIMEOUT(46);
+ FETCH_DATA_TIMEOUT(46),
+ REVIVE_INITIALIZED(47);
private final byte value;
diff --git a/common/src/main/proto/TransportMessages.proto
b/common/src/main/proto/TransportMessages.proto
index 731c6a81a..c3ebb92a5 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -205,20 +205,29 @@ message PbRequestSlotsResponse {
map<string, PbWorkerResource> workerResource = 2;
}
+message PbRevivePartitionInfo {
+ int32 partitionId = 1;
+ int32 epoch = 2;
+ PbPartitionLocation partition = 3;
+ int32 status = 4;
+}
+
message PbRevive {
int32 shuffleId = 1;
- int32 mapId = 2;
- int32 attemptId = 3;
- int32 partitionId = 4;
- int32 epoch = 5;
- PbPartitionLocation oldPartition = 6;
- int32 status = 7;
+ repeated int32 mapId = 2;
+ repeated PbRevivePartitionInfo partitionInfo = 3;
+}
+
+message PbChangeLocationPartitionInfo {
+ int32 partitionId = 1;
+ int32 status = 2;
+ PbPartitionLocation partition = 3;
+ bool oldAvailable = 4;
}
message PbChangeLocationResponse {
- int32 status = 1;
- PbPartitionLocation location = 2;
- bool available = 3;
+ repeated int32 endedMapId = 1;
+ repeated PbChangeLocationPartitionInfo partitionInfo = 2;
}
message PbPartitionSplit {
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 8f5a4a752..c86e9da7f 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -750,6 +750,8 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable
with Logging with Se
def clientPushQueueCapacity: Int = get(CLIENT_PUSH_QUEUE_CAPACITY)
def clientPushMaxReqsInFlight: Int = get(CLIENT_PUSH_MAX_REQS_IN_FLIGHT)
def clientPushMaxReviveTimes: Int = get(CLIENT_PUSH_MAX_REVIVE_TIMES)
+ def clientPushReviveInterval: Long = get(CLIENT_PUSH_REVIVE_INTERVAL)
+ def clientPushReviveBatchSize: Int = get(CLIENT_PUSH_REVIVE_BATCHSIZE)
def clientPushSortMemoryThreshold: Long =
get(CLIENT_PUSH_SORT_MEMORY_THRESHOLD)
def clientPushSortPipelineEnabled: Boolean =
get(CLIENT_PUSH_SORT_PIPELINE_ENABLED)
def clientPushSortRandomizePartitionIdEnabled: Boolean =
@@ -2674,6 +2676,23 @@ object CelebornConf extends Logging {
.intConf
.createWithDefault(5)
+ val CLIENT_PUSH_REVIVE_INTERVAL: ConfigEntry[Long] =
+ buildConf("celeborn.client.push.revive.interval")
+ .categories("client")
+ .version("0.3.0")
+ .doc("Interval for client to trigger Revive to LifecycleManager. The
number of partitions in one Revive " +
+ "request is `celeborn.client.push.revive.batchSize`.")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("100ms")
+
+ val CLIENT_PUSH_REVIVE_BATCHSIZE: ConfigEntry[Int] =
+ buildConf("celeborn.client.push.revive.batchSize")
+ .categories("client")
+ .version("0.3.0")
+ .doc("Max number of partitions in one Revive request.")
+ .intConf
+ .createWithDefault(2048)
+
val CLIENT_PUSH_BLACKLIST_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.client.push.blacklist.enabled")
.categories("client")
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
index 9bbbda7c4..ef42e6603 100644
---
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
@@ -19,6 +19,7 @@ package org.apache.celeborn.common.protocol.message
import java.util
import java.util.UUID
+import java.util.concurrent.atomic.AtomicLong
import scala.collection.JavaConverters._
@@ -186,23 +187,23 @@ object ControlMessages extends Logging {
object Revive {
def apply(
shuffleId: Int,
- mapId: Int,
- attemptId: Int,
- partitionId: Int,
- epoch: Int,
- oldPartition: PartitionLocation,
- cause: StatusCode): PbRevive = {
+ mapIds: util.Set[Integer],
+ reviveRequests: util.Collection[ReviveRequest]): PbRevive = {
val builder = PbRevive.newBuilder()
- builder
.setShuffleId(shuffleId)
- .setMapId(mapId)
- .setAttemptId(attemptId)
- .setPartitionId(partitionId)
- .setEpoch(epoch)
- .setStatus(cause.getValue)
- if (oldPartition != null) {
-
builder.setOldPartition(PbSerDeUtils.toPbPartitionLocation(oldPartition))
+ .addAllMapId(mapIds)
+
+ reviveRequests.asScala.foreach { req =>
+ val partitionInfoBuilder = PbRevivePartitionInfo.newBuilder()
+ .setPartitionId(req.partitionId)
+ .setEpoch(req.epoch)
+ .setStatus(req.cause.getValue)
+ if (req.loc != null) {
+
partitionInfoBuilder.setPartition(PbSerDeUtils.toPbPartitionLocation(req.loc))
+ }
+ builder.addPartitionInfo(partitionInfoBuilder.build())
}
+
builder.build()
}
}
@@ -223,14 +224,20 @@ object ControlMessages extends Logging {
object ChangeLocationResponse {
def apply(
- status: StatusCode,
- partitionLocationOpt: Option[PartitionLocation],
- available: Boolean): PbChangeLocationResponse = {
+ mapIds: util.Set[Integer],
+ newLocs: util.Map[Integer, (StatusCode, Boolean, PartitionLocation)])
+ : PbChangeLocationResponse = {
val builder = PbChangeLocationResponse.newBuilder()
- builder.setStatus(status.getValue)
- .setAvailable(available)
- partitionLocationOpt.foreach { partitionLocation =>
-
builder.setLocation(PbSerDeUtils.toPbPartitionLocation(partitionLocation))
+ builder.addAllEndedMapId(mapIds)
+ newLocs.asScala.foreach { case (partitionId, (status, available, loc)) =>
+ val pbChangeLocationPartitionInfoBuilder =
PbChangeLocationPartitionInfo.newBuilder()
+ .setPartitionId(partitionId)
+ .setStatus(status.getValue)
+ .setOldAvailable(available)
+ if (loc != null) {
+
pbChangeLocationPartitionInfoBuilder.setPartition(PbSerDeUtils.toPbPartitionLocation(loc))
+ }
+ builder.addPartitionInfo(pbChangeLocationPartitionInfoBuilder.build())
}
builder.build()
}
diff --git a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
index 247d77657..7b1163e54 100644
--- a/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/util/Utils.scala
@@ -883,6 +883,8 @@ object Utils extends Logging {
StatusCode.PUSH_DATA_SLAVE_BLACKLISTED
case 46 =>
StatusCode.FETCH_DATA_TIMEOUT
+ case 47 =>
+ StatusCode.REVIVE_INITIALIZED
case _ =>
null
}
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index badec4ae9..9d9e891e8 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -47,6 +47,8 @@ license: |
| celeborn.client.push.queue.capacity | 512 | Push buffer queue size for a
task. The maximum memory is `celeborn.push.buffer.max.size` *
`celeborn.push.queue.capacity`, default: 64KiB * 512 = 32MiB | 0.3.0 |
| celeborn.client.push.replicate.enabled | false | When true, Celeborn worker
will replicate shuffle data to another Celeborn worker asynchronously to ensure
the pushed shuffle data won't be lost after the node failure. It's recommended
to set `false` when `HDFS` is enabled in `celeborn.storage.activeTypes`. |
0.3.0 |
| celeborn.client.push.retry.threads | 8 | Thread number to process shuffle
re-send push data requests. | 0.3.0 |
+| celeborn.client.push.revive.batchSize | 2048 | Max number of partitions in
one Revive request. | 0.3.0 |
+| celeborn.client.push.revive.interval | 100ms | Interval for client to
trigger Revive to LifecycleManager. The number of partitions in one Revive
request is `celeborn.client.push.revive.batchSize`. | 0.3.0 |
| celeborn.client.push.revive.maxRetries | 5 | Max retry times for reviving
when celeborn push data failed. | 0.3.0 |
| celeborn.client.push.slowStart.initialSleepTime | 500ms | The initial sleep
time if the current max in flight requests is 0 | 0.3.0 |
| celeborn.client.push.slowStart.maxSleepTime | 2s | If
celeborn.client.push.limit.strategy is set to SLOWSTART, push side will take a
sleep strategy for each batch of requests, this controls the max sleep time if
the max in flight requests limit is 1 for a long time | 0.3.0 |
diff --git
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
index 791ac05df..b5c2c3599 100644
---
a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
+++
b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala
@@ -50,8 +50,9 @@ class RetryReviveTest extends AnyFunSuite
val ss = SparkSession.builder()
.config(updateSparkConf(sparkConf, ShuffleMode.HASH))
.getOrCreate()
- ss.sparkContext.parallelize(1 to 1000, 2)
+ val result = ss.sparkContext.parallelize(1 to 1000, 2)
.map { i => (i, Range(1, 1000).mkString(",")) }.groupByKey(16).collect()
+ assert(result.size == 1000)
ss.stop()
}
}
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
index 068afbc58..54f167e4c 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -178,8 +178,9 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
if (shuffleMapperAttempts.containsKey(shuffleKey)) {
if (-1 != shuffleMapperAttempts.get(shuffleKey).get(mapId)) {
// partition data has already been committed
- logInfo(s"Receive push data from speculative task(shuffle
$shuffleKey, map $mapId, " +
- s" attempt $attemptId), but this mapper has already been ended.")
+ logInfo(
+ s"[Case1] Receive push data from speculative task(shuffle
$shuffleKey, map $mapId, " +
+ s" attempt $attemptId), but this mapper has already been ended.")
callbackWithTimer.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.STAGE_ENDED.getValue)))
} else {
logInfo(
@@ -192,7 +193,7 @@ class PushDataHandler extends BaseMessageHandler with
Logging {
// If there is no shuffle key in shuffleMapperAttempts but there is
shuffle key
// in StorageManager. This partition should be HARD_SPLIT partition
and
// after worker restart, some task still push data to this
HARD_SPLIT partition.
- logInfo(s"Receive push data for committed hard split partition of " +
+ logInfo(s"[Case2] Receive push data for committed hard split
partition of " +
s"(shuffle $shuffleKey, map $mapId attempt $attemptId)")
callbackWithTimer.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
} else {
