remove target segment desc cache when disconnect (kvcache-ai#1624)

Co-authored-by: youxiao <youxiao@huawei.com>

Author: ascend-direct-dev

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

@@ -84,10 +84,10 @@ class AscendDirectTransport : public Transport {
 
     void processSliceList(const std::vector<Slice *> &slice_list);
 
-    void connectAndTransfer(const std::string &target_adxl_engine_name,
+    void connectAndTransfer(const std::string &target_seg_name,
+                            const std::string &target_adxl_engine_name,
                             adxl::TransferOp operation,
-                            const std::vector<Slice *> &slice_list,
-                            int32_t times = 0);
+                            const std::vector<Slice *> &slice_list);
 
     void localCopy(TransferRequest::OpCode opcode,
                    const std::vector<Slice *> &slice_list);
@@ -112,9 +112,10 @@ 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, bool force = false);
+                   int32_t timeout_in_millis);
 
-    void TransferWithAsync(const std::string &target_adxl_engine_name,
+    void TransferWithAsync(const std::string &target_seg_name,
+                           const std::string &target_adxl_engine_name,
                            adxl::TransferOp operation,
                            const std::vector<Slice *> &slice_list,
                            const std::vector<adxl::TransferOpDesc> &op_descs);

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

@@ -728,6 +728,7 @@ void AscendDirectTransport::queryThread() {
                     slice->markFailed();
                 }
                 it = pending_batches.erase(it);
+                metadata_->removeSegmentDesc(target_segment_desc->name);
                 disconnect(target_adxl_engine_name, connect_timeout_);
             } else if (task_status == adxl::TransferStatus::COMPLETED) {
                 auto now = getCurrentTimeInNano();
@@ -750,6 +751,7 @@ void AscendDirectTransport::queryThread() {
                     for (auto &slice : slice_list) {
                         slice->markFailed();
                     }
+                    metadata_->removeSegmentDesc(target_segment_desc->name);
                     disconnect(target_adxl_engine_name, connect_timeout_);
                     it = pending_batches.erase(it);
                 } else {
@@ -815,12 +817,14 @@ void AscendDirectTransport::processSliceList(
                 << "us";
         return;
     }
-    return connectAndTransfer(target_adxl_engine_name, operation, slice_list);
+    return connectAndTransfer(target_segment_desc->name,
+                              target_adxl_engine_name, operation, slice_list);
 }
 
 void AscendDirectTransport::connectAndTransfer(
+    const std::string &target_seg_name,
     const std::string &target_adxl_engine_name, adxl::TransferOp operation,
-    const std::vector<Slice *> &slice_list, int32_t times) {
+    const std::vector<Slice *> &slice_list) {
     if (!auto_connect_) {
         int ret = checkAndConnect(target_adxl_engine_name);
         if (ret != 0) {
@@ -842,8 +846,8 @@ void AscendDirectTransport::connectAndTransfer(
         op_descs.emplace_back(op_desc);
     }
     if (use_async_transfer_) {
-        return TransferWithAsync(target_adxl_engine_name, operation, slice_list,
-                                 op_descs);
+        return TransferWithAsync(target_seg_name, target_adxl_engine_name,
+                                 operation, slice_list, op_descs);
     }
     auto status = adxl_->TransferSync(target_adxl_engine_name.c_str(),
                                       operation, op_descs, transfer_timeout_);
@@ -854,6 +858,7 @@ void AscendDirectTransport::connectAndTransfer(
                        .count()
                 << " us";
         if (use_short_connection_) {
+            metadata_->removeSegmentDesc(target_seg_name);
             disconnect(target_adxl_engine_name, connect_timeout_);
         }
         for (auto &slice : slice_list) {
@@ -877,11 +882,13 @@ void AscendDirectTransport::connectAndTransfer(
         // set small timeout to just release local res.
         LOG(INFO) << "transfer failed and disconnect to:"
                   << target_adxl_engine_name;
+        metadata_->removeSegmentDesc(target_seg_name);
         disconnect(target_adxl_engine_name, kDefaultDisconnectTime);
     }
 }
 
 void AscendDirectTransport::TransferWithAsync(
+    const std::string &target_seg_name,
     const std::string &target_adxl_engine_name, adxl::TransferOp operation,
     const std::vector<Slice *> &slice_list,
     const std::vector<adxl::TransferOpDesc> &op_descs) {
@@ -933,6 +940,7 @@ void AscendDirectTransport::TransferWithAsync(
         }
         // the connection is probably broken.
         // set small timeout to just release local res.
+        metadata_->removeSegmentDesc(target_seg_name);
         disconnect(target_adxl_engine_name, kDefaultDisconnectTime);
     }
 #endif
@@ -1166,8 +1174,7 @@ int AscendDirectTransport::checkAndConnect(
 }
 
 int AscendDirectTransport::disconnect(
-    const std::string &target_adxl_engine_name, int32_t timeout_in_millis,
-    bool force) {
+    const std::string &target_adxl_engine_name, int32_t timeout_in_millis) {
     if (auto_connect_) {
         auto status = adxl_->Disconnect(target_adxl_engine_name.c_str(),
                                         timeout_in_millis);
@@ -1186,18 +1193,15 @@ int AscendDirectTransport::disconnect(
                   << " is not connected.";
         return 0;
     }
-    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
-                       << ", errmsg: " << aclGetRecentErrMsg();
-            connected_segments_.erase(it);
-            return -1;
-        }
-    }
+    auto status =
+        adxl_->Disconnect(target_adxl_engine_name.c_str(), timeout_in_millis);
     connected_segments_.erase(it);
+    if (status != adxl::SUCCESS) {
+        LOG(ERROR) << "Failed to disconnect to: " << target_adxl_engine_name
+                   << ", status: " << status
+                   << ", errmsg: " << aclGetRecentErrMsg();
+        return -1;
+    }
     return 0;
 }
 }  // namespace mooncake