SteNicholas commented on code in PR #3724:
URL: https://github.com/apache/celeborn/pull/3724#discussion_r3400463477
##########
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:
Good catch — done in f22033679. Merged the resubmit branch to mirror Java:
the loop now resubmits any split partition whose status is not `SOFT_SPLIT`
(still reviving with `HARD_SPLIT`), instead of only `== HARD_SPLIT`, so an
unexpected non-SOFT status can no longer be silently dropped.
##########
cpp/celeborn/client/ShuffleClient.cpp:
##########
@@ -858,6 +904,12 @@ bool ShuffleClientImpl::cleanupShuffle(int shuffleId) {
return true;
}
+void ShuffleClientImpl::shutdown() {
Review Comment:
Done in f22033679 — `shutdown()` now also clears `fetchExcludedWorkers_`, so
both client-wide exclusion sets are reset like Java. The pools/managers (revive
manager, retry pool, client factory) are torn down by the destructor (RAII), so
shutdown() intentionally only clears the two exclusion sets; added a comment to
note that.
--
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]
