diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index 541a7391e..42463bd53 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -30,6 +30,7 @@ class MultiTransport { using TransferRequest = Transport::TransferRequest; using TransferStatus = Transport::TransferStatus; using BatchDesc = Transport::BatchDesc; + using AbandonOutcome = Transport::AbandonOutcome; MultiTransport(std::shared_ptr metadata, std::string &local_server_name); @@ -54,6 +55,24 @@ class MultiTransport { Status getBatchTransferStatus(BatchID batch_id, TransferStatus &status); + // Drives all unfinished tasks and reports whether borrowed buffers and + // callback state can be released. retired is valid even on an error Status; + // transfer/query failure alone never authorizes resource reuse. + Status getBatchRetirementStatus(BatchID batch_id, bool &retired); + + // Same as above plus a safety verdict: `quiescent` is false when the + // terminal state came from a give-up (timeout/abort not confirmed), so the + // caller must not release the batch's borrowed buffers. + Status getBatchRetirementStatus(BatchID batch_id, bool &retired, + bool &quiescent); + + // Bounded attempt to stop a stalled batch. Only kQuiescentRecoverable + // publishes the terminal counts that let the caller free the batch; + // kQuiescentDeviceFault keeps the batch pending on purpose so the caller + // shuts the engine down instead of releasing buffers it cannot trust. + Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout, + AbandonOutcome &outcome); + Transport *installTransport(const std::string &proto, std::shared_ptr topo); diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 597204832..8c7c6d357 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -165,6 +165,21 @@ class TransferEngine { Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status); + // Drives all unfinished tasks and reports whether borrowed buffers and + // callback state can be released. retired is valid even on an error Status; + // transfer/query failure alone never authorizes resource reuse. + Status getBatchRetirementStatus(BatchID batch_id, bool& retired); + + // Same as above plus a safety verdict: `quiescent` is false when the + // terminal state came from a give-up, so borrowed buffers must not be + // released. Callers must not free a retired but non-quiescent batch. + Status getBatchRetirementStatus(BatchID batch_id, bool& retired, + bool& quiescent); + + // Bounded attempt to stop a stalled batch. See Transport::AbandonOutcome. + Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout, + Transport::AbandonOutcome& outcome); + Status getNicLoadStats(std::vector& stats) const; Transport* getTransport(const std::string& proto); diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index a741eeca7..0de0582d2 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -52,6 +52,7 @@ using SegmentHandle = Transport::SegmentHandle; using SegmentID = Transport::SegmentID; using BatchID = Transport::BatchID; using BufferEntry = Transport::BufferEntry; +using AbandonOutcome = Transport::AbandonOutcome; #ifdef ENABLE_MULTI_PROTOCOL using RegisteredBuffer = TransferEngine::RegisteredBuffer; @@ -318,6 +319,25 @@ class TransferEngineImpl { return result; } + // Drives unfinished classic tasks and reports whether borrowed buffers + // and completion callbacks may be released. The public TransferEngine + // facade forwards here, so this middle layer must not be skipped even + // though MultiTransport already implements the check. + Status getBatchRetirementStatus(BatchID batch_id, bool& retired) { + return multi_transports_->getBatchRetirementStatus(batch_id, retired); + } + + Status getBatchRetirementStatus(BatchID batch_id, bool& retired, + bool& quiescent) { + return multi_transports_->getBatchRetirementStatus(batch_id, retired, + quiescent); + } + + Status abandonBatch(BatchID batch_id, std::chrono::milliseconds timeout, + AbandonOutcome& outcome) { + return multi_transports_->abandonBatch(batch_id, timeout, outcome); + } + Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status, bool skip_metrics = false) { Status result = diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h index b6970814f..4175528e0 100644 --- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h +++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h @@ -50,6 +50,10 @@ class AscendDirectTransport : public Transport { Status getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status) override; + // False when any slice of the batch was failed by a give-up path whose + // in-flight ADXL operation could not be cancelled or drained. + bool isBatchQuiescent(BatchID batch_id) override; + int install(std::string &local_server_name, std::shared_ptr meta, std::shared_ptr topo) override; diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 100108218..5ac04c5c5 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -80,6 +81,16 @@ class Transport { FAILED }; + // Result of a bounded attempt to stop a stalled batch. kQuiescent* means + // the transport confirmed that it will not access the batch's borrowed + // buffers again; kNotQuiescent means no such proof exists and the caller + // must keep waiting (returning without it reintroduces a use-after-free). + enum class AbandonOutcome { + kQuiescentRecoverable, + kQuiescentDeviceFault, + kNotQuiescent, + }; + struct TransferStatus { TransferStatusEnum s; size_t transferred_bytes; @@ -178,6 +189,10 @@ class Transport { void *handle; int64_t start_time; int32_t engine_id; + // Set when the terminal state came from a timeout/abort that + // could not prove the transfer stopped. Such a slice must never + // be treated as safe to release. + bool unsafe_terminal; } ascend_direct; struct { uint64_t dest_addr; @@ -186,17 +201,19 @@ class Transport { public: void markSuccess() { + CompletionCallbackPin pin(task); status = Slice::SUCCESS; __atomic_fetch_add(&task->transferred_bytes, length, __ATOMIC_RELAXED); - __atomic_fetch_add(&task->success_slice_count, 1, __ATOMIC_RELAXED); + __atomic_fetch_add(&task->success_slice_count, 1, __ATOMIC_ACQ_REL); check_batch_completion(false); } void markFailed() { + CompletionCallbackPin pin(task); status = Slice::FAILED; - __atomic_fetch_add(&task->failed_slice_count, 1, __ATOMIC_RELAXED); + __atomic_fetch_add(&task->failed_slice_count, 1, __ATOMIC_ACQ_REL); check_batch_completion(true); } @@ -204,6 +221,22 @@ class Transport { volatile int64_t ts; private: + class CompletionCallbackPin { + public: + explicit CompletionCallbackPin(TransferTask *task) + : batch_(toBatchDesc(task->batch_id)) { + batch_.active_completion_callbacks.fetch_add( + 1, std::memory_order_acq_rel); + } + ~CompletionCallbackPin() { + batch_.active_completion_callbacks.fetch_sub( + 1, std::memory_order_release); + } + + private: + BatchDesc &batch_; + }; + inline void check_batch_completion(bool is_failed) { #ifdef USE_EVENT_DRIVEN_COMPLETION auto &batch_desc = toBatchDesc(task->batch_id); @@ -346,6 +379,9 @@ class Transport { void *context; // for transport implementers. int64_t start_timestamp; + // A terminal count does not mean its completion callback has returned. + std::atomic active_completion_callbacks{0}; + // Track batch progress and notifies waiters std::atomic has_failure{false}; std::atomic is_finished{ @@ -371,6 +407,18 @@ class Transport { /// @brief Free an allocated batch. virtual Status freeBatchID(BatchID batch_id); + /// @brief Stop a stalled batch and report whether its buffers are safe to + /// release. Implementations must only report kQuiescent* when they can + /// prove that no transfer access is in flight any more. + virtual Status abandonBatch(BatchID batch_id, + std::chrono::milliseconds timeout, + AbandonOutcome& outcome); + + /// @brief Whether a retired batch is provably idle (no in-flight access). + /// Transports that publish terminal state on give-up/timeout must override + /// this and return false for those batches. + virtual bool isBatchQuiescent(BatchID batch_id); + /// @brief Submit a batch of transfer requests to the batch. /// @return The number of successfully submitted transfers on success. If /// that number is less than nr, errno is set. diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 81cb0174f..bef87276e 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -96,6 +96,13 @@ MultiTransport::BatchID MultiTransport::allocateBatchID(size_t batch_size) { } Status MultiTransport::freeBatchID(BatchID batch_id) { + bool retired = false; + auto status = getBatchRetirementStatus(batch_id, retired); + if (!retired) { + return Status::BatchBusy("Batch tasks or callbacks are pending: " + + status.ToString()); + } + // A transfer/query error does not prevent cleanup after retirement. auto& batch_desc = *((BatchDesc*)(batch_id)); const size_t task_count = batch_desc.task_list.size(); for (size_t task_id = 0; task_id < task_count; task_id++) { @@ -254,7 +261,7 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, } else { status.s = Transport::TransferStatusEnum::COMPLETED; } - task.is_finished = true; + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); } else { if (checkSliceTimeout(task)) { status.s = Transport::TransferStatusEnum::TIMEOUT; @@ -265,6 +272,136 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, return Status::OK(); } +Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, + bool& retired) { + auto& batch = Transport::toBatchDesc(batch_id); + retired = true; + Status result = Status::OK(); + for (size_t index = 0; index < batch.task_list.size(); ++index) { + auto& task = batch.task_list[index]; + const uint64_t count = + __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE); + uint64_t succeeded = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + uint64_t failed = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + if (succeeded > count || failed != count - succeeded) { + TransferStatus status; + auto queried = getTransferStatus(batch_id, index, status); + if (result.ok() && !queried.ok()) { + result = queried; + } else if (result.ok() && + (status.s == Transport::TransferStatusEnum::FAILED || + status.s == Transport::TransferStatusEnum::TIMEOUT || + status.s == Transport::TransferStatusEnum::CANCELED)) { + result = + Status::Context("Transfer task failed, state=" + + std::to_string(static_cast(status.s))); + } + succeeded = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + failed = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + } + if (succeeded > count || failed != count - succeeded) retired = false; + if (result.ok() && failed != 0) { + result = Status::Context("Transfer task failed: " + + std::to_string(index)); + } + } + retired = retired && batch.active_completion_callbacks.load( + std::memory_order_acquire) == 0; + if (retired) { + for (auto& task : batch.task_list) { + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + } + } + return result; +} + +Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, + bool &retired, + bool &quiescent) { + const Status status = getBatchRetirementStatus(batch_id, retired); + quiescent = retired; + if (retired) { + auto &batch = Transport::toBatchDesc(batch_id); + for (auto &task : batch.task_list) { + if (task.transport_ == nullptr) continue; + if (!task.transport_->isBatchQuiescent(batch_id)) { + quiescent = false; + break; + } + } + } + return status; +} + +Status MultiTransport::abandonBatch(BatchID batch_id, + std::chrono::milliseconds timeout, + AbandonOutcome &outcome) { + auto &batch = Transport::toBatchDesc(batch_id); + outcome = AbandonOutcome::kQuiescentRecoverable; + + // Collect the transports that still own unfinished tasks of this batch. + std::vector pending; + for (auto &task : batch.task_list) { + const uint64_t count = + __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE); + const uint64_t succeeded = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + const uint64_t failed = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + if (succeeded + failed == count) continue; + if (task.transport_ == nullptr) { + // No transport to ask: the batch can never be proven quiescent. + outcome = AbandonOutcome::kNotQuiescent; + return Status::OK(); + } + if (std::find(pending.begin(), pending.end(), task.transport_) == + pending.end()) { + pending.push_back(task.transport_); + } + } + + Status result = Status::OK(); + for (auto *transport : pending) { + AbandonOutcome local = AbandonOutcome::kNotQuiescent; + const Status status = transport->abandonBatch(batch_id, timeout, local); + if (!status.ok()) { + // The transport cannot stop this batch: no safety proof. + outcome = AbandonOutcome::kNotQuiescent; + return status; + } + if (local == AbandonOutcome::kQuiescentDeviceFault) { + outcome = local; + } else if (local == AbandonOutcome::kNotQuiescent && + outcome == AbandonOutcome::kQuiescentRecoverable) { + outcome = local; + } + } + + if (outcome == AbandonOutcome::kQuiescentRecoverable) { + // The transport proved that the batch is stopped, so publishing the + // terminal counts is safe and lets the caller free the batch. + for (auto &task : batch.task_list) { + const uint64_t count = + __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE); + const uint64_t succeeded = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + const uint64_t failed = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + if (succeeded + failed < count) { + __atomic_add_fetch(&task.failed_slice_count, + count - succeeded - failed, + __ATOMIC_ACQ_REL); + } + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + } + } + return result; +} + Status MultiTransport::getBatchTransferStatus(BatchID batch_id, TransferStatus& status) { auto& batch_desc = *((BatchDesc*)(batch_id)); diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 3c31a4320..f688be41c 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -247,6 +247,22 @@ Status TransferEngine::getBatchTransferStatus(BatchID batch_id, return impl_->getBatchTransferStatus(batch_id, status); } +Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, + bool& retired) { + return impl_->getBatchRetirementStatus(batch_id, retired); +} + +Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, + bool& retired, bool& quiescent) { + return impl_->getBatchRetirementStatus(batch_id, retired, quiescent); +} + +Status TransferEngine::abandonBatch(BatchID batch_id, + std::chrono::milliseconds timeout, + Transport::AbandonOutcome& outcome) { + return impl_->abandonBatch(batch_id, timeout, outcome); +} + Status TransferEngine::getNicLoadStats(std::vector& stats) const { stats.clear(); return Status::OK(); @@ -756,6 +772,61 @@ Status TransferEngine::getBatchTransferStatus(BatchID batch_id, return impl_->getBatchTransferStatus(batch_id, status); } +Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, + bool& retired) { + if (!use_tent_) return impl_->getBatchRetirementStatus(batch_id, retired); + retired = false; + TransferStatus status; + auto queried = getBatchTransferStatus(batch_id, status); + if (!queried.ok()) return queried; + // TENT caches failure/timeout states without exposing transport quiescence. + // Those states cannot authorize reuse of borrowed buffers. + retired = status.s == TransferStatusEnum::COMPLETED; + if (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::TIMEOUT || + status.s == TransferStatusEnum::CANCELED) { + return Status::Context("Transfer batch failed, state=" + + std::to_string(static_cast(status.s))); + } + return Status::OK(); +} + +Status TransferEngine::abandonBatch(BatchID batch_id, + std::chrono::milliseconds timeout, + Transport::AbandonOutcome& outcome) { + if (!use_tent_) return impl_->abandonBatch(batch_id, timeout, outcome); + // TENT does not expose transport quiescence: a stalled TENT batch can never + // be abandoned safely, so keep it pending and let the caller wait. + (void)batch_id; + (void)timeout; + outcome = Transport::AbandonOutcome::kNotQuiescent; + return Status::NotImplemented("TENT does not expose batch quiescence"); +} + +Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, + bool& retired, + bool& quiescent) { + if (!use_tent_) { + return impl_->getBatchRetirementStatus(batch_id, retired, quiescent); + } + retired = false; + quiescent = false; + TransferStatus status; + auto queried = getBatchTransferStatus(batch_id, status); + if (!queried.ok()) return queried; + retired = status.s == TransferStatusEnum::COMPLETED; + // TENT only certifies COMPLETED; cached failure/timeout states are not + // proof that the transport stopped accessing the buffers. + quiescent = retired; + if (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::TIMEOUT || + status.s == TransferStatusEnum::CANCELED) { + return Status::Context("Transfer batch failed, state=" + + std::to_string(static_cast(status.s))); + } + return Status::OK(); +} + Status TransferEngine::getNicLoadStats(std::vector& stats) const { stats.clear(); if (use_tent_) { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp index 365505cbb..0a0ab37bc 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp @@ -283,6 +283,16 @@ Status AscendDirectTransport::submitTransferTask( return Status::OK(); } +bool AscendDirectTransport::isBatchQuiescent(BatchID batch_id) { + auto &batch_desc = *((BatchDesc *)(batch_id)); + for (auto &task : batch_desc.task_list) { + for (auto *slice : task.slice_list) { + if (slice->ascend_direct.unsafe_terminal) return false; + } + } + return true; +} + // actually not called, just use getTransferStatus in multi_transport Status AscendDirectTransport::getTransferStatus(BatchID batch_id, size_t task_id, diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp index 3455d1217..38e27018c 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp @@ -216,6 +216,11 @@ void AsyncTransferExecutor::markBatchFailed(QueryBatch& batch, LOG(ERROR) << reason; } for (auto* slice : batch.slices) { + // A submitted ADXL operation has no cancel path: its timeout/failure + // publication is a give-up, not proof that the buffers are idle. + if (slice->ascend_direct.handle != nullptr) { + slice->ascend_direct.unsafe_terminal = true; + } slice->markFailed(); } handleTaskFinished(); diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/local_copy_engine.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/local_copy_engine.cpp index e752f2e9e..e6517b4d3 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/local_copy_engine.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/local_copy_engine.cpp @@ -284,6 +284,9 @@ void LocalCopyEngine::CopyWithAsync(TransferRequest::OpCode opcode, LOG(ERROR) << "Memory copy failed, ret:" << ret; (void)aclrtStreamAbort(stream_); for (auto &slice : async_list) { + // The abort is not proven to have taken effect; the slice must not + // be treated as safe to release. + slice->ascend_direct.unsafe_terminal = true; slice->markFailed(); } } diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp index 6f08a3257..1e8462c71 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp @@ -700,6 +700,10 @@ void TransferExecutorBase::processSliceList( if (result.status == adxl::TIMEOUT) { LOG(ERROR) << "Transfer timeout to: " << target_adxl_engine_name << ", errmsg: " << aclGetRecentErrMsg(); + // The synchronous transfer timed out without a cancellation proof. + for (auto* slice : slice_list) { + slice->ascend_direct.unsafe_terminal = true; + } } else { LOG(ERROR) << "Transfer failed to: " << target_adxl_engine_name << ", errmsg: " << aclGetRecentErrMsg(); diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp index ef81cb6f3..11c6344a2 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp @@ -133,8 +133,10 @@ void HcclTransport::initiatorLoop(int deviceLogicId, int selfIdx) { << ", source_addr: " << slice->source_addr << ", dest_addr: " << slice->hccl.dest_addr << ", ret: " << ret; - slice->markFailed(); - slice->status = Slice::SliceStatus::FAILED; + // Keep the error private until the stream and all retries + // have retired. Publishing FAILED here lets callers free + // buffers that a sibling operation still accesses. + slice->status = Slice::FAILED; } } @@ -147,7 +149,7 @@ void HcclTransport::initiatorLoop(int deviceLogicId, int selfIdx) { << ", remote devicePhyId: " << remote_rank_info_.devicePhyId << ", ret: " << ret; for (auto slice : slice_list) { - slice->markFailed(); + slice->status = Slice::FAILED; } } auto addOpfence = std::chrono::high_resolution_clock::now(); @@ -160,9 +162,21 @@ void HcclTransport::initiatorLoop(int deviceLogicId, int selfIdx) { << local_rank_info_.devicePhyId << ", remote devicePhyId: " << remote_rank_info_.devicePhyId << ", ret: " << ret; - aclrtStreamAbort(stream); + const int abort_ret = aclrtStreamAbort(stream); + if (abort_ret) { + LOG(ERROR) << "aclrtStreamAbort failed, ret: " << abort_ret; + } // Synchronize to drop ACL_ERROR_RT_STREAM_ABORT error - aclrtSynchronizeStream(stream); + int drain_ret = aclrtSynchronizeStream(stream); + while (drain_ret != ACL_SUCCESS) { + // An error result alone cannot release borrowed buffers. + // Keep the batch pending until the stream actually drains. + LOG_EVERY_T(ERROR, 1) + << "Waiting for failed HCCL stream retirement, ret: " + << drain_ret; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + drain_ret = aclrtSynchronizeStream(stream); + } if (loop < transfer_max_retry_cnt) { LOG(INFO) << "Hccl transport failed. Retry ..."; @@ -181,7 +195,7 @@ void HcclTransport::initiatorLoop(int deviceLogicId, int selfIdx) { LOG(ERROR) << "Hccl transport failed after " << transfer_max_retry_cnt << " retries."; for (auto slice : slice_list) { - slice->markFailed(); + slice->status = Slice::FAILED; } } @@ -222,9 +236,10 @@ void HcclTransport::initiatorLoop(int deviceLogicId, int selfIdx) { } for (auto slice : slice_list) { - if (slice->status != Slice::SliceStatus::FAILED) { + if (slice->status == Slice::SliceStatus::FAILED) { + slice->markFailed(); + } else { slice->markSuccess(); - slice->task->transferred_bytes = slice->length; } } (void)start; @@ -538,16 +553,20 @@ Status HcclTransport::getTransferStatus(BatchID batch_id, size_t task_id, std::to_string(batch_id)); } auto &task = batch_desc.task_list[task_id]; - status.transferred_bytes = task.transferred_bytes; - uint64_t success_slice_count = task.success_slice_count; - uint64_t failed_slice_count = task.failed_slice_count; + uint64_t success_slice_count = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + uint64_t failed_slice_count = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + // Completion counters publish the preceding byte updates. + status.transferred_bytes = + __atomic_load_n(&task.transferred_bytes, __ATOMIC_RELAXED); if (success_slice_count + failed_slice_count == task.slice_count) { if (failed_slice_count) { status.s = TransferStatusEnum::FAILED; } else { status.s = TransferStatusEnum::COMPLETED; } - task.is_finished = true; + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); } else { status.s = TransferStatusEnum::WAITING; } diff --git a/mooncake-transfer-engine/src/transport/transport.cpp b/mooncake-transfer-engine/src/transport/transport.cpp index 00e8bf96f..0d19bb5b1 100644 --- a/mooncake-transfer-engine/src/transport/transport.cpp +++ b/mooncake-transfer-engine/src/transport/transport.cpp @@ -49,6 +49,10 @@ Status Transport::freeBatchID(BatchID batch_id) { "BatchID cannot be freed until all tasks are done"); } } + if (batch_desc.active_completion_callbacks.load( + std::memory_order_acquire) != 0) { + return Status::BatchBusy("Completion callbacks are pending"); + } delete &batch_desc; #ifdef CONFIG_USE_BATCH_DESC_SET RWSpinlock::WriteGuard guard(batch_desc_lock_); @@ -57,6 +61,22 @@ Status Transport::freeBatchID(BatchID batch_id) { return Status::OK(); } +bool Transport::isBatchQuiescent(BatchID batch_id) { + (void)batch_id; + // Legacy: terminal counters are the only signal available. + return true; +} + +Status Transport::abandonBatch(BatchID batch_id, + std::chrono::milliseconds timeout, + AbandonOutcome &outcome) { + (void)batch_id; + (void)timeout; + outcome = AbandonOutcome::kNotQuiescent; + return Status::NotImplemented( + "Transport::abandonBatch is not implemented"); +} + int Transport::install(std::string &local_server_name, std::shared_ptr meta, std::shared_ptr topo) { @@ -64,4 +84,4 @@ int Transport::install(std::string &local_server_name, metadata_ = meta; return 0; } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index e7a343b4a..f128ac00f 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -117,6 +117,14 @@ add_executable(dmabuf_export_test ${WORKSPACE}/dmabuf_export_test.cpp) target_link_libraries(dmabuf_export_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME dmabuf_export_test COMMAND dmabuf_export_test) +# Hardware-free unit tests for the batch retirement contract: the +# TransferEngineImpl forwarding layer, the MultiTransport drain/callback checks +# and the Transport completion-callback gate. No transport device is required. +add_executable(batch_retirement_test ${WORKSPACE}/batch_retirement_test.cpp) +target_link_libraries(batch_retirement_test PUBLIC transfer_engine gtest + gtest_main) +add_test(NAME batch_retirement_test COMMAND batch_retirement_test) + if(USE_MNNVL) add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp) target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest diff --git a/mooncake-transfer-engine/tests/batch_retirement_test.cpp b/mooncake-transfer-engine/tests/batch_retirement_test.cpp new file mode 100644 index 000000000..22e667c40 --- /dev/null +++ b/mooncake-transfer-engine/tests/batch_retirement_test.cpp @@ -0,0 +1,477 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include + +#include "multi_transport.h" +#include "transfer_engine.h" +#include "transfer_engine_impl.h" +#include "transport/transport.h" + +namespace mooncake { + +// Exposes the transport layer of a TransferEngineImpl so the retirement +// contract can be tested without starting a metadata/transport stack. +class TransferEngineImplTestPeer { + public: + static void setMultiTransport(TransferEngineImpl& engine, + std::shared_ptr transports) { + engine.multi_transports_ = std::move(transports); + } +}; + +// Hardware-free transport: it materialises the task/slice layout of a submitted +// batch and reports the transfer status the test asks for. No buffer is ever +// touched, so the retirement contract is testable on CPU-only runners. +class RetirementProbeTransport : public Transport { + public: + Status submitTransfer( + BatchID batch_id, + const std::vector& entries) override { + auto& batch = toBatchDesc(batch_id); + auto& task = batch.task_list.emplace_back(); + task.batch_id = batch_id; + task.transport_ = this; + for (const auto& entry : entries) { + auto* slice = getSliceCache().allocate(); + slice->source_addr = entry.source; + slice->length = entry.length; + slice->opcode = entry.opcode; + slice->target_id = entry.target_id; + slice->status = Slice::PENDING; + slice->ts = 0; + slice->task = &task; + task.slice_list.push_back(slice); + task.total_bytes += entry.length; + } + task.slice_count = task.slice_list.size(); + return Status::OK(); + } + + Status getTransferStatus(BatchID, size_t, TransferStatus& status) override { + ++query_count_; + if (!query_status_.ok()) return query_status_; + status.s = reported_status_; + status.transferred_bytes = 0; + return Status::OK(); + } + + void failQueries(Status status) { query_status_ = std::move(status); } + + void clearQueryFailure() { query_status_ = Status::OK(); } + + int queryCount() const { return query_count_; } + + int registerLocalMemory(void*, size_t, const std::string&, bool, + bool) override { + return 0; + } + + int unregisterLocalMemory(void*, bool) override { return 0; } + + int registerLocalMemoryBatch(const std::vector&, + const std::string&) override { + return 0; + } + + int unregisterLocalMemoryBatch(const std::vector&) override { + return 0; + } + + const char* getName() const override { return "retirement-probe"; } + + private: + Status query_status_; + TransferStatusEnum reported_status_ = TransferStatusEnum::WAITING; + int query_count_ = 0; +}; + +// Probe transport that reports an unsafe (give-up) terminal state. +class NonQuiescentProbeTransport : public RetirementProbeTransport { + public: + bool isBatchQuiescent(BatchID) override { return false; } +}; + +// Probe transport that can claim a bounded stop, so the abandon contract can +// be tested without hardware. +class AbandonProbeTransport : public RetirementProbeTransport { + public: + Status abandonBatch(BatchID, std::chrono::milliseconds, + AbandonOutcome& outcome) override { + ++abandon_calls_; + outcome = abandon_outcome_; + return abandon_status_; + } + + void setAbandonResult(Status status, AbandonOutcome outcome) { + abandon_status_ = std::move(status); + abandon_outcome_ = outcome; + } + + int abandonCalls() const { return abandon_calls_; } + + private: + Status abandon_status_; + AbandonOutcome abandon_outcome_ = AbandonOutcome::kQuiescentRecoverable; + int abandon_calls_ = 0; +}; + +namespace { + +std::vector makeRequests(size_t count) { + std::vector entries; + for (size_t i = 0; i < count; ++i) { + Transport::TransferRequest request{}; + request.opcode = Transport::TransferRequest::WRITE; + request.source = nullptr; // the probe never dereferences the buffer + request.target_id = 0; + request.target_offset = i * 64; + request.length = 64; + entries.push_back(request); + } + return entries; +} + +// Owns a MultiTransport plus the probe transport and one submitted batch. The +// batch handle stays opaque to the test: only public transport calls are used. +class SubmittedBatchFixture { + public: + explicit SubmittedBatchFixture(size_t request_count) + : local_server_name_("retirement-test") { + transports_ = + std::make_shared(nullptr, local_server_name_); + probe_ = std::make_shared(); + batch_id_ = transports_->allocateBatchID(1); + submitted_ = + probe_->submitTransfer(batch_id_, makeRequests(request_count)); + } + + MultiTransport& transports() { return *transports_; } + + RetirementProbeTransport& probe() { return *probe_; } + + Transport::BatchID batchId() const { return batch_id_; } + + bool submitted() const { return submitted_.ok(); } + + Transport::BatchDesc& batch() { return Transport::toBatchDesc(batch_id_); } + + Transport::Slice* slice(size_t index) { + return batch().task_list[0].slice_list[index]; + } + + private: + std::string local_server_name_; + std::shared_ptr transports_; + std::shared_ptr probe_; + Transport::BatchID batch_id_ = 0; + Status submitted_; +}; + +} // namespace + +// The public facade must reach the retirement check through +// TransferEngineImpl; this cannot compile while the middle layer lacks the +// forwarding member that the wrapper calls. +TEST(BatchRetirementTest, PublicApiReportsRetiredEmptyBatch) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + const BatchID batch_id = engine.allocateBatchID(1); + ASSERT_NE(batch_id, 0); + + bool retired = false; + const Status status = engine.getBatchRetirementStatus(batch_id, retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + + EXPECT_TRUE(engine.freeBatchID(batch_id).ok()); +} + +// Same forwarding path without a metadata server: an empty batch is retired and +// can be released immediately. +TEST(BatchRetirementTest, ImplForwardsRetirementStatus) { + std::string local_server_name = "retirement-impl-test"; + TransferEngineImpl engine(false); + auto transports = + std::make_shared(nullptr, local_server_name); + TransferEngineImplTestPeer::setMultiTransport(engine, transports); + + const BatchID batch_id = engine.allocateBatchID(1); + ASSERT_NE(batch_id, 0); + + bool retired = false; + const Status status = engine.getBatchRetirementStatus(batch_id, retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + + EXPECT_TRUE(engine.freeBatchID(batch_id).ok()); +} + +// Borrowed buffers stay owner-held until every slice of every task has reached +// a terminal state; cleanup is refused in the meantime. +TEST(BatchRetirementTest, PendingSlicesKeepBatchBusy) { + SubmittedBatchFixture fixture(2); + ASSERT_TRUE(fixture.submitted()); + + bool retired = true; + Status status = fixture.transports().getBatchRetirementStatus( + fixture.batchId(), retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(retired); + + const Status busy = fixture.transports().freeBatchID(fixture.batchId()); + EXPECT_EQ(busy.code(), Status::Code::kBatchBusy) << busy.ToString(); + + fixture.slice(0)->markSuccess(); + status = fixture.transports().getBatchRetirementStatus(fixture.batchId(), + retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(retired) << "one slice is still in flight"; + + fixture.slice(1)->markSuccess(); + status = fixture.transports().getBatchRetirementStatus(fixture.batchId(), + retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + EXPECT_TRUE(fixture.transports().freeBatchID(fixture.batchId()).ok()); +} + +// A terminal slice count is not enough: a completion callback that has not +// returned still keeps the batch alive. +TEST(BatchRetirementTest, PendingCompletionCallbackBlocksRetirement) { + SubmittedBatchFixture fixture(1); + ASSERT_TRUE(fixture.submitted()); + fixture.slice(0)->markSuccess(); + + fixture.batch().active_completion_callbacks.fetch_add( + 1, std::memory_order_acq_rel); + + bool retired = true; + Status status = fixture.transports().getBatchRetirementStatus( + fixture.batchId(), retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(retired); + + const Status busy = fixture.transports().freeBatchID(fixture.batchId()); + EXPECT_EQ(busy.code(), Status::Code::kBatchBusy) << busy.ToString(); + + fixture.batch().active_completion_callbacks.fetch_sub( + 1, std::memory_order_acq_rel); + status = fixture.transports().getBatchRetirementStatus(fixture.batchId(), + retired); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + EXPECT_TRUE(fixture.transports().freeBatchID(fixture.batchId()).ok()); +} + +// A failed transfer is still retired: the error is reported to the caller, yet +// the transport confirms that no access is in flight, so cleanup must succeed. +TEST(BatchRetirementTest, FailedTaskRetiresWithErrorStatus) { + SubmittedBatchFixture fixture(2); + ASSERT_TRUE(fixture.submitted()); + fixture.slice(0)->markFailed(); + fixture.slice(1)->markFailed(); + + bool retired = false; + const Status status = fixture.transports().getBatchRetirementStatus( + fixture.batchId(), retired); + EXPECT_FALSE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired) + << "terminal failure with no in-flight access is retired"; + + const Status freed = fixture.transports().freeBatchID(fixture.batchId()); + EXPECT_TRUE(freed.ok()) << freed.ToString(); +} + +// A query error alone never authorises buffer reuse; the batch stays owned. +TEST(BatchRetirementTest, QueryErrorDoesNotRetire) { + SubmittedBatchFixture fixture(2); + ASSERT_TRUE(fixture.submitted()); + fixture.probe().failQueries(Status::Socket("probe query failed")); + + bool retired = true; + const Status status = fixture.transports().getBatchRetirementStatus( + fixture.batchId(), retired); + EXPECT_FALSE(status.ok()) << status.ToString(); + EXPECT_FALSE(retired); + EXPECT_GT(fixture.probe().queryCount(), 0); + + const Status busy = fixture.transports().freeBatchID(fixture.batchId()); + EXPECT_EQ(busy.code(), Status::Code::kBatchBusy) << busy.ToString(); + + // Recovery path: once the probe answers again and the task finishes, the + // batch must become releasable. + fixture.probe().clearQueryFailure(); + fixture.slice(0)->markSuccess(); + fixture.slice(1)->markSuccess(); + retired = false; + EXPECT_TRUE(fixture.transports() + .getBatchRetirementStatus(fixture.batchId(), retired) + .ok()); + EXPECT_TRUE(retired); + EXPECT_TRUE(fixture.transports().freeBatchID(fixture.batchId()).ok()); +} + +// Transport-level cleanup (used by transports that own their BatchDesc) applies +// the same callback gate. +TEST(BatchRetirementTest, TransportCleanupWaitsForCallbacks) { + auto probe = std::make_shared(); + const BatchID batch_id = probe->allocateBatchID(1); + auto& batch = Transport::toBatchDesc(batch_id); + batch.active_completion_callbacks.fetch_add(1, std::memory_order_acq_rel); + + const Status busy = probe->freeBatchID(batch_id); + EXPECT_EQ(busy.code(), Status::Code::kBatchBusy) << busy.ToString(); + + batch.active_completion_callbacks.fetch_sub(1, std::memory_order_acq_rel); + EXPECT_TRUE(probe->freeBatchID(batch_id).ok()); +} + +// A recoverable abandon proves that no access is in flight, so retirement is +// published and the caller may free the batch. +TEST(BatchRetirementTest, AbandonRecoverablePublishesTerminalState) { + std::string local_server_name = "abandon-recoverable-test"; + MultiTransport transports(nullptr, local_server_name); + auto probe = std::make_shared(); + const BatchID batch_id = transports.allocateBatchID(1); + ASSERT_TRUE(probe->submitTransfer(batch_id, makeRequests(2)).ok()); + + Transport::AbandonOutcome outcome = Transport::AbandonOutcome::kNotQuiescent; + const Status status = transports.abandonBatch( + batch_id, std::chrono::milliseconds(50), outcome); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_EQ(outcome, Transport::AbandonOutcome::kQuiescentRecoverable); + EXPECT_EQ(probe->abandonCalls(), 1); + + // The abandoned task reports its failure status, but retirement proves the + // buffers are safe and the batch can be freed. + bool retired = false; + const Status retired_status = + transports.getBatchRetirementStatus(batch_id, retired); + EXPECT_FALSE(retired_status.ok()) << retired_status.ToString(); + EXPECT_TRUE(retired); + EXPECT_TRUE(transports.freeBatchID(batch_id).ok()); +} + +// A device fault stops the batch but proves nothing about buffer safety: the +// batch must stay pending so the caller shuts the engine down instead of +// releasing borrowed buffers. +TEST(BatchRetirementTest, AbandonDeviceFaultKeepsBatchPending) { + std::string local_server_name = "abandon-device-fault-test"; + MultiTransport transports(nullptr, local_server_name); + auto probe = std::make_shared(); + probe->setAbandonResult(Status::OK(), + Transport::AbandonOutcome::kQuiescentDeviceFault); + const BatchID batch_id = transports.allocateBatchID(1); + ASSERT_TRUE(probe->submitTransfer(batch_id, makeRequests(2)).ok()); + + Transport::AbandonOutcome outcome = Transport::AbandonOutcome::kNotQuiescent; + const Status status = transports.abandonBatch( + batch_id, std::chrono::milliseconds(50), outcome); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_EQ(outcome, Transport::AbandonOutcome::kQuiescentDeviceFault); + + bool retired = true; + EXPECT_TRUE(transports.getBatchRetirementStatus(batch_id, retired).ok()); + EXPECT_FALSE(retired); + EXPECT_EQ(transports.freeBatchID(batch_id).code(), + Status::Code::kBatchBusy); + + // Finish the fake transfer so the test does not leak the batch handle. + auto& batch = Transport::toBatchDesc(batch_id); + for (auto* slice : batch.task_list[0].slice_list) slice->markSuccess(); + EXPECT_TRUE(transports.getBatchRetirementStatus(batch_id, retired).ok()); + EXPECT_TRUE(retired); + EXPECT_TRUE(transports.freeBatchID(batch_id).ok()); +} + +// A transport without the capability cannot prove safety; the generic layer +// must report kNotQuiescent and keep the batch pending. +TEST(BatchRetirementTest, AbandonUnsupportedTransportKeepsBatchPending) { + std::string local_server_name = "abandon-unsupported-test"; + MultiTransport transports(nullptr, local_server_name); + auto probe = std::make_shared(); + const BatchID batch_id = transports.allocateBatchID(1); + ASSERT_TRUE(probe->submitTransfer(batch_id, makeRequests(2)).ok()); + + Transport::AbandonOutcome outcome = Transport::AbandonOutcome::kQuiescentRecoverable; + const Status status = transports.abandonBatch( + batch_id, std::chrono::milliseconds(50), outcome); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(outcome, Transport::AbandonOutcome::kNotQuiescent); + + bool retired = true; + EXPECT_TRUE(transports.getBatchRetirementStatus(batch_id, retired).ok()); + EXPECT_FALSE(retired); + + auto& batch = Transport::toBatchDesc(batch_id); + for (auto* slice : batch.task_list[0].slice_list) slice->markSuccess(); + EXPECT_TRUE(transports.getBatchRetirementStatus(batch_id, retired).ok()); + EXPECT_TRUE(transports.freeBatchID(batch_id).ok()); +} + +// A give-up terminal state retires the batch but must not be advertised as +// safe to release. +TEST(BatchRetirementTest, NonQuiescentRetirementIsReported) { + std::string local_server_name = "non-quiescent-test"; + MultiTransport transports(nullptr, local_server_name); + auto probe = std::make_shared(); + const BatchID batch_id = transports.allocateBatchID(1); + ASSERT_TRUE(probe->submitTransfer(batch_id, makeRequests(2)).ok()); + auto& batch = Transport::toBatchDesc(batch_id); + for (auto* slice : batch.task_list[0].slice_list) slice->markSuccess(); + + bool retired = false; + bool quiescent = true; + const Status status = + transports.getBatchRetirementStatus(batch_id, retired, quiescent); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + EXPECT_FALSE(quiescent); + EXPECT_TRUE(transports.freeBatchID(batch_id).ok()); +} + +// A confirmed stop is reported as quiescent. +TEST(BatchRetirementTest, QuiescentRetirementIsReported) { + std::string local_server_name = "quiescent-test"; + MultiTransport transports(nullptr, local_server_name); + auto probe = std::make_shared(); + const BatchID batch_id = transports.allocateBatchID(1); + ASSERT_TRUE(probe->submitTransfer(batch_id, makeRequests(2)).ok()); + auto& batch = Transport::toBatchDesc(batch_id); + for (auto* slice : batch.task_list[0].slice_list) slice->markSuccess(); + + bool retired = false; + bool quiescent = false; + const Status status = + transports.getBatchRetirementStatus(batch_id, retired, quiescent); + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(retired); + EXPECT_TRUE(quiescent); + EXPECT_TRUE(transports.freeBatchID(batch_id).ok()); +} + +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}