SteNicholas commented on code in PR #3724:
URL: https://github.com/apache/celeborn/pull/3724#discussion_r3400519726
##########
cpp/celeborn/client/writer/PushMergedDataCallback.cpp:
##########
@@ -118,117 +135,128 @@ void PushMergedDataCallback::onSuccess(
<< groupedBatchId_ << ".";
if (response->remainingSize() > 0) {
- // Parse PbPushMergedDataSplitPartitionInfo from TransportMessage
- auto transportMsg = std::make_unique<protocol::TransportMessage>(
- response->readToReadOnlyBuffer(response->remainingSize()));
- PbPushMergedDataSplitPartitionInfo partitionInfo;
- if (!partitionInfo.ParseFromString(transportMsg->payload())) {
- pushState_->setException(std::make_unique<std::runtime_error>(
- "Failed to parse PbPushMergedDataSplitPartitionInfo"));
- return;
- }
+ try {
+ // Parse PbPushMergedDataSplitPartitionInfo from TransportMessage
+ auto transportMsg = std::make_unique<protocol::TransportMessage>(
+ response->readToReadOnlyBuffer(response->remainingSize()));
+ PbPushMergedDataSplitPartitionInfo partitionInfo;
+ if (!partitionInfo.ParseFromString(transportMsg->payload())) {
+ pushState_->setException(std::make_unique<std::runtime_error>(
+ "Failed to parse PbPushMergedDataSplitPartitionInfo"));
+ return;
+ }
- CELEBORN_CHECK_EQ(
- partitionInfo.statuscodes_size(),
- partitionInfo.splitpartitionindexes_size(),
- "Mismatched sizes: statuscodes {} vs splitpartitionindexes {}",
- partitionInfo.statuscodes_size(),
- partitionInfo.splitpartitionindexes_size());
- const int numBatches = static_cast<int>(batches_.size());
- for (int i = 0; i < partitionInfo.splitpartitionindexes_size(); i++) {
- int partitionIndex = partitionInfo.splitpartitionindexes(i);
- CELEBORN_CHECK_GE(partitionIndex, 0);
- CELEBORN_CHECK_LT(
- partitionIndex,
- numBatches,
- "Partition index {} out of range [0, {})",
- partitionIndex,
- numBatches);
- int statusCode = partitionInfo.statuscodes(i);
+ CELEBORN_CHECK_EQ(
+ partitionInfo.statuscodes_size(),
+ partitionInfo.splitpartitionindexes_size(),
+ "Mismatched sizes: statuscodes {} vs splitpartitionindexes {}",
+ partitionInfo.statuscodes_size(),
+ partitionInfo.splitpartitionindexes_size());
+ const int numBatches = static_cast<int>(batches_.size());
+ for (int i = 0; i < partitionInfo.splitpartitionindexes_size(); i++)
{
+ int partitionIndex = partitionInfo.splitpartitionindexes(i);
+ CELEBORN_CHECK_GE(partitionIndex, 0);
+ CELEBORN_CHECK_LT(
+ partitionIndex,
+ numBatches,
+ "Partition index {} out of range [0, {})",
+ partitionIndex,
+ numBatches);
+ int statusCode = partitionInfo.statuscodes(i);
- if (statusCode ==
- static_cast<int>(protocol::StatusCode::SOFT_SPLIT)) {
- int partitionId = partitionIds_[partitionIndex];
- if (!ShuffleClientImpl::newerPartitionLocationExists(
- sharedClient->getPartitionLocationMap(shuffleId_).value(),
+ if (statusCode ==
+ static_cast<int>(protocol::StatusCode::SOFT_SPLIT)) {
+ int partitionId = partitionIds_[partitionIndex];
+ if (!ShuffleClientImpl::newerPartitionLocationExists(
+
sharedClient->getPartitionLocationMap(shuffleId_).value(),
+ partitionId,
+ batches_[partitionIndex].loc->epoch)) {
+ auto reviveRequest = std::make_shared<protocol::ReviveRequest>(
+ shuffleId_,
+ mapId_,
+ attemptId_,
partitionId,
- batches_[partitionIndex].loc->epoch)) {
+ batches_[partitionIndex].loc->epoch,
+ batches_[partitionIndex].loc,
+ protocol::StatusCode::SOFT_SPLIT);
+ sharedClient->addRequestToReviveManager(reviveRequest);
+ }
+ }
+ }
+
+ // For any HARD_SPLIT partitions, need to resubmit
+ std::vector<DataBatch> batchesToRetry;
+ std::vector<std::shared_ptr<protocol::ReviveRequest>> reviveRequests;
+ for (int i = 0; i < partitionInfo.splitpartitionindexes_size(); i++)
{
+ int partitionIndex = partitionInfo.splitpartitionindexes(i);
+ CELEBORN_DCHECK_GE(partitionIndex, 0);
+ CELEBORN_DCHECK_LT(partitionIndex, numBatches);
+ int statusCode = partitionInfo.statuscodes(i);
+ if (statusCode ==
+ static_cast<int>(protocol::StatusCode::HARD_SPLIT)) {
Review Comment:
Branch rebased onto `main`; this change now lives in `294f89377` (rebased
equivalent of `f22033679`, identical diff).
##########
cpp/celeborn/client/ShuffleClient.cpp:
##########
@@ -858,6 +904,12 @@ bool ShuffleClientImpl::cleanupShuffle(int shuffleId) {
return true;
}
+void ShuffleClientImpl::shutdown() {
Review Comment:
Branch rebased onto `main`; this change now lives in `294f89377` (rebased
equivalent of `f22033679`, identical diff).
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
