address comments of Co Pilot AI

Author: afterincomparableyum

cpp/celeborn/client/ShuffleClient.cpp

@@ -359,7 +359,8 @@ std::unique_ptr<CelebornInputStream> ShuffleClientImpl::readPartition(
 void ShuffleClientImpl::excludeFailedFetchLocation(
     const std::string& hostAndFetchPort,
     const std::exception& e) {
-  if (pushReplicateEnabled_ && fetchExcludeWorkerOnFailureEnabled_) {
+  if (pushReplicateEnabled_ && fetchExcludeWorkerOnFailureEnabled_ &&
+      utils::isCriticalCauseForFetch(e)) {
     auto now = std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::steady_clock::now().time_since_epoch())
                    .count();

cpp/celeborn/client/reader/CelebornInputStream.cpp

@@ -17,6 +17,7 @@
 
 #include "celeborn/client/reader/CelebornInputStream.h"
 #include <lz4.h>
+#include <thread>
 #include "celeborn/client/compress/Decompressor.h"
 
 namespace celeborn {
@@ -226,8 +227,7 @@ std::shared_ptr<PartitionReader> CelebornInputStream::createReaderWithRetry(
 
       if (currentLocation->hasPeer() && !readSkewPartitionWithoutMapRange_) {
         if (fetchChunkRetryCnt_ % 2 == 0) {
-          std::this_thread::sleep_for(
-              std::chrono::milliseconds(retryWait_.count()));
+          std::this_thread::sleep_for(retryWait_);
         }
         LOG(WARNING) << "CreatePartitionReader failed " << fetchChunkRetryCnt_
                      << "/" << fetchChunkMaxRetry_ << " times for location "
@@ -244,8 +244,7 @@ std::shared_ptr<PartitionReader> CelebornInputStream::createReaderWithRetry(
                      << "/" << fetchChunkMaxRetry_ << " times for location "
                      << currentLocation->hostAndFetchPort()
                      << ", retry the same location. Error: " << e.what();
-        std::this_thread::sleep_for(
-            std::chrono::milliseconds(retryWait_.count()));
+        std::this_thread::sleep_for(retryWait_);
       }
     }
   }

cpp/celeborn/client/reader/CelebornInputStream.h

@@ -17,8 +17,6 @@
 
 #pragma once
 
-#include <thread>
-
 #include "celeborn/client/compress/Decompressor.h"
 #include "celeborn/client/reader/WorkerPartitionReader.h"
 #include "celeborn/conf/CelebornConf.h"

cpp/celeborn/client/tests/CMakeLists.txt

@@ -17,6 +17,7 @@ add_executable(
         celeborn_client_test
         WorkerPartitionReaderTest.cpp
         PushDataCallbackTest.cpp
+        CelebornInputStreamRetryTest.cpp
         PushStateTest.cpp
         ReviveManagerTest.cpp
         Lz4DecompressorTest.cpp