Adapt to adxl connection auto release feature

Author: youxiao

mooncake-common/common.cmake

@@ -160,14 +160,15 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT)
     file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
   endif()
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
-  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
   set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include")
   add_compile_options(-Wno-ignored-qualifiers)
   include_directories(/usr/local/include /usr/include ${ASCEND_INCLUDE_DIR})
-  link_directories(${ASCEND_LIB_DIR} ${ASCEND_DEVLIB_DIR})
+  link_directories(${ASCEND_LIB_DIR})
 endif()
 
 if (USE_ASCEND)
+  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
+  link_directories(${ASCEND_DEVLIB_DIR})
   add_compile_definitions(USE_ASCEND)
 endif()
 

mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h

@@ -81,6 +81,11 @@ class AscendDirectTransport : public Transport {
 
     void processSliceList(const std::vector<Slice *> &slice_list);
 
+    void connectAndTransfer(const std::string &target_adxl_engine_name,
+                            adxl::TransferOp operation,
+                            const std::vector<Slice *> &slice_list,
+                            int32_t times = 0);
+
     void localCopy(TransferRequest::OpCode opcode,
                    const std::vector<Slice *> &slice_list);
 
@@ -105,7 +110,7 @@ class AscendDirectTransport : public Transport {
     int checkAndConnect(const std::string &target_adxl_engine_name);
 
     int disconnect(const std::string &target_adxl_engine_name,
-                   int32_t timeout_in_millis);
+                   int32_t timeout_in_millis, bool force = false);
 
     std::atomic_bool running_;
     std::unique_ptr<adxl::AdxlEngine> adxl_;
@@ -131,6 +136,8 @@ class AscendDirectTransport : public Transport {
     bool use_buffer_pool_{false};
 
     int32_t base_port_ = 20000;
+
+    std::unordered_set<SegmentID> need_update_metadata_segs_;
 };
 
 }  // namespace mooncake

mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp

@@ -36,7 +36,8 @@
 namespace mooncake {
 namespace {
 constexpr size_t kMemcpyBatchLimit = 4096;
-}
+constexpr int32_t kMaxAdxlConnectRetries = 3;
+}  // namespace
 AscendDirectTransport::AscendDirectTransport() : running_(false) {}
 
 AscendDirectTransport::~AscendDirectTransport() {
@@ -78,6 +79,7 @@ AscendDirectTransport::~AscendDirectTransport() {
         }
     }
     addr_to_mem_handle_.clear();
+    adxl_->Finalize();
 }
 
 int AscendDirectTransport::install(std::string &local_server_name,
@@ -547,8 +549,9 @@ void AscendDirectTransport::processSliceList(
     if (slice_list.empty()) {
         return;
     }
-    auto target_segment_desc =
-        metadata_->getSegmentDescByID(slice_list[0]->target_id);
+    auto it = need_update_metadata_segs_.find(slice_list[0]->target_id);
+    auto target_segment_desc = metadata_->getSegmentDescByID(
+        slice_list[0]->target_id, (it != need_update_metadata_segs_.end()));
     if (!target_segment_desc) {
         LOG(ERROR) << "Cannot find segment descriptor for target_id: "
                    << slice_list[0]->target_id;
@@ -557,6 +560,9 @@ void AscendDirectTransport::processSliceList(
         }
         return;
     }
+    if (it != need_update_metadata_segs_.end()) {
+        need_update_metadata_segs_.erase(it);
+    }
     auto target_adxl_engine_name =
         (target_segment_desc->rank_info.hostIp + ":" +
          std::to_string(target_segment_desc->rank_info.hostPort));
@@ -582,10 +588,14 @@ void AscendDirectTransport::processSliceList(
                   << "us";
         return;
     }
+    return connectAndTransfer(target_adxl_engine_name, operation, slice_list);
+}
+
+void AscendDirectTransport::connectAndTransfer(
+    const std::string &target_adxl_engine_name, adxl::TransferOp operation,
+    const std::vector<Slice *> &slice_list, int32_t times) {
     int ret = checkAndConnect(target_adxl_engine_name);
     if (ret != 0) {
-        LOG(ERROR) << "Failed to connect to segment: "
-                   << target_segment_desc->name;
         for (auto &slice : slice_list) {
             slice->markFailed();
         }
@@ -613,6 +623,14 @@ void AscendDirectTransport::processSliceList(
                          std::chrono::steady_clock::now() - start)
                          .count()
                   << " us";
+    } else if (status == adxl::NOT_CONNECTED) {
+        LOG(INFO) << "Connection reset by backend, retry times:" << times;
+        disconnect(target_adxl_engine_name, 0, true);
+        if (times < kMaxAdxlConnectRetries) {
+            return connectAndTransfer(target_adxl_engine_name, operation,
+                                      slice_list, times + 1);
+        }
+        return;
     } else {
         if (status == adxl::TIMEOUT) {
             LOG(ERROR) << "Transfer timeout to: " << target_adxl_engine_name
@@ -628,6 +646,7 @@ void AscendDirectTransport::processSliceList(
         // the connection is probably broken.
         // set small timeout to just release local res.
         disconnect(target_adxl_engine_name, 10);
+        need_update_metadata_segs_.emplace(slice_list[0]->target_id);
     }
 }
 
@@ -844,21 +863,24 @@ int AscendDirectTransport::checkAndConnect(
 }
 
 int AscendDirectTransport::disconnect(
-    const std::string &target_adxl_engine_name, int32_t timeout_in_millis) {
+    const std::string &target_adxl_engine_name, int32_t timeout_in_millis,
+    bool force) {
     std::lock_guard<std::mutex> lock(connection_mutex_);
     auto it = connected_segments_.find(target_adxl_engine_name);
     if (it == connected_segments_.end()) {
         LOG(INFO) << "Target adxl engine: " << target_adxl_engine_name
                   << " is not connected.";
         return 0;
     }
-    auto status =
-        adxl_->Disconnect(target_adxl_engine_name.c_str(), timeout_in_millis);
-    if (status != adxl::SUCCESS) {
-        LOG(ERROR) << "Failed to disconnect to: " << target_adxl_engine_name
-                   << ", status: " << status;
-        connected_segments_.erase(target_adxl_engine_name);
-        return -1;
+    if (!force) {
+        auto status = adxl_->Disconnect(target_adxl_engine_name.c_str(),
+                                        timeout_in_millis);
+        if (status != adxl::SUCCESS) {
+            LOG(ERROR) << "Failed to disconnect to: " << target_adxl_engine_name
+                       << ", status: " << status;
+            connected_segments_.erase(target_adxl_engine_name);
+            return -1;
+        }
     }
     connected_segments_.erase(target_adxl_engine_name);
     return 0;