diff --git a/docs/source/design/tent/cpp-api.md b/docs/source/design/tent/cpp-api.md index fc68a22c62..3acead95ea 100644 --- a/docs/source/design/tent/cpp-api.md +++ b/docs/source/design/tent/cpp-api.md @@ -75,6 +75,7 @@ For users migrating from Transfer Engine, the following table shows how TE APIs | `allocateBatchID(batch_size)` | `allocateBatch(batch_size)` | Renamed | | `freeBatchID(batch_id)` | `freeBatch(batch_id)` | Renamed | | `submitTransfer(batch_id, entries)` | `submitTransfer(batch_id, request_list)` | `TransferRequest` → `Request` | +| *Not available* | `cancelTransfer(batch_id, task_id)` | TENT-only: best-effort cancellation for queued and RDMA tasks | | `submitTransferWithNotify(batch_id, entries, notify_msg)` | `submitTransfer(batch_id, request_list, notifi)` | Unified API with optional notification | | `getTransferStatus(batch_id, task_id, status)` | `getTransferStatus(batch_id, task_id, status)` | Same | | `getBatchTransferStatus(batch_id, status)` | `getTransferStatus(batch_id, status)` | Overloaded; single `TransferStatus` output = overall status | @@ -274,6 +275,24 @@ Queries the status of transfer requests. - `status` / `status_list` / `overall_status`: Output parameter(s) for status. - Return value: `Status::OK()` on success; otherwise a non-OK status. +#### TransferEngine::cancelTransfer + +```cpp +Status cancelTransfer(BatchID batch_id, size_t task_id); +``` + +Requests best-effort cancellation of one public task. A task still waiting in +the TENT admission queue becomes `CANCELED` without being dispatched. For RDMA, +workers suppress slices they observe before `ibv_post_send`; work already +posted to a QP is allowed to drain and may complete successfully. Consequently, +the API returning `OK` means the cancellation request was accepted, not that +the task is already terminal. Continue polling `getTransferStatus` before +calling `freeBatch`. + +Cancellation is idempotent. Merged public tasks share one physical transfer, +so canceling any alias cancels the shared task. Direct cancellation of staging +or non-RDMA transport work currently returns `Status::NotImplemented`. + #### TransferEngine::freeBatch ```cpp diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index d74cb087e4..61fc5580a3 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -134,6 +134,11 @@ class LocalTransferAdmissionQueue { Status complete(QueueOwnerId owner_id, TransferStatusEnum terminal_status); + // Cancel an owner that has not entered the dispatch window. Idempotent for + // an owner already canceled; dispatching owners must be canceled through + // their selected transport instead. + Status cancel(QueueOwnerId owner_id); + Status retireBatch(uint64_t batch_token); Status resolveOwner(uint64_t batch_token, size_t public_task_id, diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 1bd440545f..25f1fb6ea1 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -58,6 +58,7 @@ struct TaskInfo { std::string qp_pool; // Named QP pool (RFC #2568 step 3), "" = none Request request; bool staging{false}; + bool cancel_requested{false}; TransferStatusEnum status{TransferStatusEnum::PENDING}; volatile TransferStatusEnum staging_status{TransferStatusEnum::PENDING}; std::chrono::steady_clock::time_point start_time{}; // For latency tracking @@ -140,6 +141,8 @@ class TransferEngineImpl { const std::vector& request_list, const Notification& notifi); + Status cancelTransfer(BatchID batch_id, size_t task_id); + Status sendNotification(SegmentID target_id, const Notification& notifi); Status receiveNotification(std::vector& notifi_list); @@ -246,6 +249,8 @@ class TransferEngineImpl { Status finishQueuedOwner(QueueOwnerId owner_id, TransferStatusEnum terminal_status); + Status cancelQueuedOwner(QueueOwnerId owner_id); + Status retireQueueForBatch(Batch* batch); Status pollTaskStatus(Batch* batch, size_t task_id, diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index 7f1a2d3c2b..e82d8c1d89 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -63,14 +63,14 @@ class Transport { std::function notify_progress; }; - using SubBatchRef = SubBatch *; + using SubBatchRef = SubBatch*; public: Transport() = default; virtual ~Transport() = default; - virtual Status install(std::string &local_segment_name, + virtual Status install(std::string& local_segment_name, std::shared_ptr metadata, std::shared_ptr local_topology, std::shared_ptr conf = nullptr) { @@ -81,56 +81,67 @@ class Transport { virtual const Capabilities capabilities() const { return caps; } - virtual Status allocateSubBatch(SubBatchRef &batch, size_t max_size) { + virtual Status allocateSubBatch(SubBatchRef& batch, size_t max_size) { return Status::NotImplemented( "allocateSubBatch not implemented" LOC_MARK); } - virtual Status freeSubBatch(SubBatchRef &batch) { + virtual Status freeSubBatch(SubBatchRef& batch) { return Status::NotImplemented("freeSubBatch not implemented" LOC_MARK); } virtual Status submitTransferTasks( - SubBatchRef batch, const std::vector &request_list) { + SubBatchRef batch, const std::vector& request_list) { return Status::NotImplemented( "submitTransferTasks not implemented" LOC_MARK); } virtual Status getTransferStatus(SubBatchRef batch, int task_id, - TransferStatus &status) { + TransferStatus& status) { return Status::NotImplemented( "getTransferStatus not implemented" LOC_MARK); } - virtual Status allocateLocalMemory(void **addr, size_t size, - MemoryOptions &options) { + // Cancellation is best effort: implementations must prevent work that has + // not reached the device from being submitted, but work already posted to + // a device may still complete. Callers must continue polling until the + // task reaches a terminal state. + virtual bool supportsCancellation() const { return false; } + + virtual Status cancelTransferTask(SubBatchRef batch, int task_id) { + return Status::NotImplemented( + "cancelTransferTask not implemented" LOC_MARK); + } + + virtual Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions& options) { return Platform::getLoader().allocate(addr, size, options); } - virtual Status freeLocalMemory(void *addr, size_t size) { + virtual Status freeLocalMemory(void* addr, size_t size) { return Platform::getLoader().free(addr, size); } // Pre-registration warm-up that pins pages before NUMA probing. // Returns true if pages were successfully pinned (caller may skip // prefault). Default: no-op, returns false. - virtual bool warmupMemory(void *addr, size_t length) { return false; } + virtual bool warmupMemory(void* addr, size_t length) { return false; } - virtual Status addMemoryBuffer(BufferDesc &desc, - const MemoryOptions &options) { + virtual Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) { return Status::NotImplemented( "addMemoryBuffer not implemented" LOC_MARK); } - virtual Status addMemoryBuffer(std::vector &desc_list, - const MemoryOptions &options) { - for (auto &desc : desc_list) { + virtual Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) { + for (auto& desc : desc_list) { CHECK_STATUS(addMemoryBuffer(desc, options)); } return Status::OK(); } - virtual Status removeMemoryBuffer(BufferDesc &desc) { + virtual Status removeMemoryBuffer(BufferDesc& desc) { return Status::NotImplemented( "removeMemoryBuffer not implemented" LOC_MARK); } @@ -138,17 +149,17 @@ class Transport { virtual bool supportNotification() const { return false; } virtual Status sendNotification(SegmentID target_id, - const Notification ¬ify) { + const Notification& notify) { return Status::NotImplemented( "sendNotification not implemented" LOC_MARK); } - virtual Status receiveNotification(std::vector ¬ify_list) { + virtual Status receiveNotification(std::vector& notify_list) { return Status::NotImplemented( "receiveNotification not implemented" LOC_MARK); } - virtual const char *getName() const { return ""; } + virtual const char* getName() const { return ""; } protected: Capabilities caps; diff --git a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h index ff93e0ed02..52a9776427 100644 --- a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h +++ b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h @@ -174,6 +174,9 @@ void tent_free_notifs(tent_notifi_info* info); int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, size_t task_id, tent_status_t* status); +int tent_cancel_task(tent_engine_t engine, tent_batch_id_t batch_id, + size_t task_id); + int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* status); @@ -299,6 +302,11 @@ class TransferEngine { const std::vector& request_list, const Notification& notifi); + // Best-effort task cancellation. Work that has not reached the transport + // is prevented from being submitted. Device work already posted may still + // complete, so callers must continue polling for a terminal status. + Status cancelTransfer(BatchID batch_id, size_t task_id); + Status sendNotification(SegmentID target_id, const Notification& notifi); Status receiveNotification(std::vector& notifi_list); @@ -328,4 +336,4 @@ class TransferEngine { } // namespace mooncake #endif -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h index 3dc6deba4b..2e87b8aec8 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h @@ -79,6 +79,10 @@ class RdmaTransport : public Transport { virtual Status getTransferStatus(SubBatchRef batch, int task_id, TransferStatus& status); + bool supportsCancellation() const override { return true; } + + Status cancelTransferTask(SubBatchRef batch, int task_id) override; + virtual Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions& options); diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index b27a6b7163..8332a7c34a 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -57,10 +57,14 @@ struct RdmaTask { std::string qp_pool; volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; - volatile int success_slices; - volatile int resolved_slices; + std::atomic success_slices{0}; + std::atomic resolved_slices{0}; volatile TransferStatusEnum first_error = PENDING; + // Set by the control thread. Workers observe this flag before posting or + // retrying a slice. Already-posted WRs are allowed to drain normally. + std::atomic cancel_requested{false}; + // Reference counting for UAF protection std::atomic ref_count{0}; @@ -92,6 +96,9 @@ struct RdmaSlice { int qp_index = 0; int retry_count = 0; bool failed = false; + // True while DeviceSelector accounts this slice against source_dev_id. + // The worker clears it exactly once on completion, failure, or cancel. + bool quota_charged = false; uint64_t enqueue_ts = 0; uint64_t submit_ts = 0; // Non-owning pointer to the per-worker RailMonitor for this slice's @@ -111,15 +118,18 @@ static inline void updateSliceStatus(RdmaSlice* slice, if (!__sync_bool_compare_and_swap(&slice->word, PENDING, status)) return; if (status == COMPLETED) { __sync_fetch_and_add(&task->transferred_bytes, slice->length); - __sync_fetch_and_add(&task->success_slices, 1); + task->success_slices.fetch_add(1, std::memory_order_acq_rel); } else { __sync_bool_compare_and_swap(&task->first_error, PENDING, status); } - int resolved = __sync_add_and_fetch(&task->resolved_slices, 1); + int resolved = + task->resolved_slices.fetch_add(1, std::memory_order_acq_rel) + 1; if (resolved >= task->num_slices) { - TransferStatusEnum final_st = (task->success_slices == task->num_slices) - ? COMPLETED - : task->first_error; + TransferStatusEnum final_st = + (task->success_slices.load(std::memory_order_acquire) == + task->num_slices) + ? COMPLETED + : task->first_error; if (final_st == PENDING) final_st = FAILED; __sync_bool_compare_and_swap(&task->status_word, PENDING, final_st); } @@ -129,4 +139,4 @@ static inline void updateSliceStatus(RdmaSlice* slice, } // namespace tent } // namespace mooncake -#endif // TENT_SLICE_H \ No newline at end of file +#endif // TENT_SLICE_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h index c18d5de130..05750142f3 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -43,7 +43,7 @@ class Workers { using BoundedSliceQueue = BoundedMPSCQueue; public: - Workers(RdmaTransport *transport); + Workers(RdmaTransport* transport); ~Workers(); @@ -51,16 +51,17 @@ class Workers { Status stop(); - Status submit(RdmaSlice *slice); + Status submit(RdmaSlice* slice); - Status submit(RdmaSliceList &slice_list, int worker_id = -1); + Status submit(RdmaSliceList& slice_list, int worker_id = -1); - Status cancel(RdmaSliceList &slice_list); + Status cancel(RdmaTask* task); - DeviceSelector *getDeviceSelector() const { return device_selector_.get(); } + DeviceSelector* getDeviceSelector() const { return device_selector_.get(); } private: using Task = std::function; + struct WorkerContext; void workerThread(int thread_id); @@ -68,41 +69,45 @@ class Workers { void asyncPollCq(); + bool cancelUnpostedSlice(WorkerContext& worker, RdmaSlice* slice); + + void releaseSliceQuota(RdmaSlice* slice, double latency = 0.0); + void monitorThread(); - int handleContextEvents(std::shared_ptr &context); + int handleContextEvents(std::shared_ptr& context); - Status generatePostPath(RdmaSlice *slice); + Status generatePostPath(RdmaSlice* slice); private: struct RouteHint { // Owning reference to the segment snapshot; keeps all raw pointers // below valid for the lifetime of this hint. SegmentDescRef pin; - SegmentDesc *segment; - BufferDesc *buffer; - const Topology::MemEntry *topo_entry; - const Topology *topo; + SegmentDesc* segment; + BufferDesc* buffer; + const Topology::MemEntry* topo_entry; + const Topology* topo; std::string location; }; - Status getRouteHint(RouteHint &hint, SegmentID segment_id, uint64_t addr, + Status getRouteHint(RouteHint& hint, SegmentID segment_id, uint64_t addr, uint64_t length); - Status selectOptimalDevice(RouteHint &source, RouteHint &target, - RdmaSlice *slice); + Status selectOptimalDevice(RouteHint& source, RouteHint& target, + RdmaSlice* slice); - Status selectFallbackDevice(RouteHint &source, RouteHint &target, - RdmaSlice *slice); + Status selectFallbackDevice(RouteHint& source, RouteHint& target, + RdmaSlice* slice); - int getDeviceByFlatIndex(const RouteHint &hint, size_t flat_idx); + int getDeviceByFlatIndex(const RouteHint& hint, size_t flat_idx); - int getDeviceRank(const RouteHint &hint, int device_id); + int getDeviceRank(const RouteHint& hint, int device_id); void showLatencyInfo(); private: - RdmaTransport *transport_; + RdmaTransport* transport_; size_t num_workers_; std::thread monitor_; @@ -113,7 +118,7 @@ class Workers { SegmentID remote_segment_id; int remote_device_id; - bool operator==(const PostPath &rhs) const { + bool operator==(const PostPath& rhs) const { return local_device_id == rhs.local_device_id && remote_segment_id == rhs.remote_segment_id && remote_device_id == rhs.remote_device_id; @@ -121,7 +126,7 @@ class Workers { }; struct PostPathHash { - size_t operator()(const PostPath &postPath) const { + size_t operator()(const PostPath& postPath) const { size_t h1 = std::hash{}(postPath.local_device_id); size_t h2 = std::hash{}(postPath.remote_segment_id); size_t h3 = std::hash{}(postPath.remote_device_id); @@ -131,10 +136,10 @@ class Workers { std::shared_ptr getEndpoint(PostPath path); - void disableEndpoint(RdmaSlice *slice); + void disableEndpoint(RdmaSlice* slice); using GroupedRequests = - std::unordered_map, PostPathHash>; + std::unordered_map, PostPathHash>; struct PerfMetric { void add(double val) { samples.push_back(val); } @@ -191,7 +196,7 @@ class Workers { std::thread thread; BoundedSliceQueue queues[kNumPriorityLevels]; // Priority queues GroupedRequests requests; - std::unordered_set inflight_slice_set; + std::unordered_set inflight_slice_set; std::atomic inflight_slices = 0; std::mutex mutex; @@ -210,9 +215,9 @@ class Workers { }; // Promote timed-out low priority requests to higher priority queues - void promoteTimedOutRequests(WorkerContext &worker); + void promoteTimedOutRequests(WorkerContext& worker); - WorkerContext *worker_context_; + WorkerContext* worker_context_; uint64_t slice_timeout_ns_; uint64_t priority_promotion_timeout_ns_; // Timeout for priority promotion // Opt-in (issue #2528): when true, a promotion pass promotes exactly the diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index c1f2abe390..1a39544b11 100644 --- a/mooncake-transfer-engine/tent/src/python/pybind.cpp +++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp @@ -703,6 +703,15 @@ PYBIND11_MODULE(tent, m) { py::arg("batch_id"), py::arg("request_list"), py::arg("name"), py::arg("message")) + .def( + "cancel_transfer", + [](TransferEngine& self, uint64_t batch_id, size_t task_id) { + py::gil_scoped_release release; + auto s = self.cancelTransfer((BatchID)batch_id, task_id); + ThrowStatus(s, "cancel_transfer"); + }, + py::arg("batch_id"), py::arg("task_id")) + // --------------------------------------------------------------------- // notification send/receive // --------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 6eae246f95..ee0b3d2df8 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -367,6 +367,37 @@ Status LocalTransferAdmissionQueue::complete( return Status::OK(); } +Status LocalTransferAdmissionQueue::cancel(QueueOwnerId owner_id) { + if (owner_id == 0) { + return Status::InvalidArgument("invalid queue owner id" LOC_MARK); + } + auto owner_it = owners_.find(owner_id); + if (owner_it == owners_.end()) { + return Status::InvalidEntry("queue owner not found" LOC_MARK); + } + auto& owner = owner_it->second; + if (owner.state == QueueState::Terminal) { + return owner.terminal_status == TransferStatusEnum::CANCELED + ? Status::OK() + : Status::InvalidEntry( + "queue owner is already terminal" LOC_MARK); + } + if (owner.state != QueueState::Queued) { + return Status::InvalidEntry( + "queue owner is already dispatching" LOC_MARK); + } + + owner.state = QueueState::Terminal; + owner.terminal_status = TransferStatusEnum::CANCELED; + --outstanding_owners_; + outstanding_bytes_ -= owner.request.length; + if (owner.kind == QueueOwnerKind::User) { + --outstanding_user_owners_; + outstanding_user_bytes_ -= owner.request.length; + } + return Status::OK(); +} + Status LocalTransferAdmissionQueue::retireBatch(uint64_t batch_token) { if (batch_token == 0) { return Status::InvalidArgument("invalid batch token" LOC_MARK); diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 6bee9e27aa..d0a8201dc8 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -1651,6 +1651,25 @@ Status TransferEngineImpl::finishQueuedOwner( return Status::OK(); } +Status TransferEngineImpl::cancelQueuedOwner(QueueOwnerId owner_id) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + return Status::InvalidEntry("queued owner not found" LOC_MARK); + } + if (queued_it->second.in_dispatch_window) { + return Status::InvalidEntry( + "queued owner is already dispatching" LOC_MARK); + } + CHECK_STATUS(runtime_queue_->cancel(owner_id)); + for (const auto task_id : queued_it->second.public_task_ids) { + auto& task = queued_it->second.batch->task_list[task_id]; + task.cancel_requested = true; + task.status = CANCELED; + } + queued_owners_.erase(queued_it); + return Status::OK(); +} + Status TransferEngineImpl::retireQueueForBatch(Batch* batch) { if (!batch || batch->queue_token == 0) return Status::OK(); auto status = runtime_queue_->retireBatch(batch->queue_token); @@ -1899,6 +1918,75 @@ Status TransferEngineImpl::submitTransfer( QueueOwnerKind::User); } +Status TransferEngineImpl::cancelTransfer(BatchID batch_id, size_t task_id) { + if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK); + std::lock_guard lk(progress_mutex_); + if (!alive_batches_.count(batch_id)) { + return Status::InvalidArgument("Batch is not alive" LOC_MARK); + } + auto* batch = reinterpret_cast(batch_id); + if (task_id >= batch->task_list.size()) { + return Status::InvalidArgument("Invalid task ID" LOC_MARK); + } + + size_t owner_task_id = task_id; + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + QueueOwnerId owner_id = 0; + auto resolve_status = + runtime_queue_->resolveOwner(batch->queue_token, task_id, owner_id); + if (resolve_status.ok()) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + TransferStatusEnum public_status = PENDING; + CHECK_STATUS(runtime_queue_->getPublicStatus( + batch->queue_token, task_id, public_status)); + return public_status != PENDING + ? Status::OK() + : Status::InvalidEntry( + "queued owner metadata missing" LOC_MARK); + } + owner_task_id = queued_it->second.owner_task_id; + if (!queued_it->second.in_dispatch_window) { + CHECK_STATUS(cancelQueuedOwner(owner_id)); + CHECK_STATUS(refillDispatchWindow()); + notifyRuntimeQueueReady(); + return Status::OK(); + } + } + } + + auto& owner = batch->task_list[owner_task_id]; + if (owner.status != PENDING) return Status::OK(); + if (owner.staging) { + return Status::NotImplemented( + "staging transfer cancellation is not implemented" LOC_MARK); + } + if (owner.type == UNSPEC) { + owner.cancel_requested = true; + owner.status = CANCELED; + return Status::OK(); + } + auto& transport = transport_list_[owner.type]; + auto& sub_batch = batch->sub_batch[owner.type]; + if (!transport || !sub_batch) { + return Status::InvalidArgument("Transport not available" LOC_MARK); + } + if (!transport->supportsCancellation()) { + return Status::NotImplemented( + "selected transport does not support cancellation" LOC_MARK); + } + + CHECK_STATUS(transport->cancelTransferTask(sub_batch, owner.sub_task_id)); + // Merged public tasks share one physical transport task. Mark every alias + // so polling any of them cannot trigger failover after cancellation. + for (auto& task : batch->task_list) { + if (task.type == owner.type && task.sub_task_id == owner.sub_task_id) { + task.cancel_requested = true; + } + } + return Status::OK(); +} + Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { auto& task = batch->task_list[task_id]; auto prev_type = task.type; @@ -1968,7 +2056,8 @@ void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id, bool allow_failover) { auto& task = batch->task_list[task_id]; task.status = task_status.s; - if (!allow_failover || task_status.s != FAILED || task.type == UNSPEC) + if (!allow_failover || task.cancel_requested || task_status.s != FAILED || + task.type == UNSPEC) return; if (resubmitTransferTask(batch, task_id).ok()) { diff --git a/mooncake-transfer-engine/tent/src/transfer_engine.cpp b/mooncake-transfer-engine/tent/src/transfer_engine.cpp index 5628a45c93..8d495763f6 100644 --- a/mooncake-transfer-engine/tent/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/tent/src/transfer_engine.cpp @@ -140,6 +140,10 @@ Status TransferEngine::submitTransfer(BatchID batch_id, return impl_->submitTransfer(batch_id, request_list, notifi); } +Status TransferEngine::cancelTransfer(BatchID batch_id, size_t task_id) { + return impl_->cancelTransfer(batch_id, task_id); +} + Status TransferEngine::sendNotification(SegmentID target_id, const Notification& notifi) { return impl_->sendNotification(target_id, notifi); diff --git a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp index 150df8b8e6..f9c88258f0 100644 --- a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp +++ b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp @@ -308,7 +308,7 @@ void tent_free_notifs(tent_notifi_info* info) { int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, size_t task_id, tent_status_t* xfer_status) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(xfer_status); mooncake::tent::TransferStatus internal_status; auto status = @@ -322,10 +322,22 @@ int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, return 0; } +int tent_cancel_task(tent_engine_t engine, tent_batch_id_t batch_id, + size_t task_id) { + CHECK_POINTER(engine); + if (!batch_id) return -1; + auto status = CAST(engine)->cancelTransfer(batch_id, task_id); + if (!status.ok()) { + LOG(ERROR) << "tent_cancel_task: " << status.ToString(); + return -1; + } + return 0; +} + int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* xfer_status) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(xfer_status); mooncake::tent::TransferStatus internal_status; auto status = CAST(engine)->getTransferStatus(batch_id, internal_status); @@ -463,7 +475,7 @@ int tent_register_memory_batch_ex(tent_engine_t engine, void** addrs, int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* statuses, size_t* count) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(statuses); CHECK_POINTER(count); std::vector status_list; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp index 53a8bb4296..d3073a2c5f 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp @@ -292,7 +292,9 @@ Status DeviceSelector::release(int dev_id, uint64_t length, double latency) { auto& dev = it->second; dev.releaseInflight(length); - if (!smart_selection_enabled_) { + // Cancellation of an unposted slice must release its inflight charge but + // has no latency sample from which to learn bandwidth. + if (!smart_selection_enabled_ || latency <= 0.0) { return Status::OK(); } @@ -357,4 +359,4 @@ int DeviceSelector::getDevicePriority(int dev_id) const { } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index dc7ca187a9..a744a3616b 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -441,6 +441,10 @@ Status RdmaTransport::submitTransferTasks( task->num_slices = 0; task->status_word = PENDING; task->transferred_bytes = 0; + task->success_slices.store(0, std::memory_order_relaxed); + task->resolved_slices.store(0, std::memory_order_relaxed); + task->first_error = PENDING; + task->cancel_requested.store(false, std::memory_order_relaxed); task->ref(); // Batch holds a reference to the task const double merge_ratio = 0.25; @@ -492,6 +496,7 @@ Status RdmaTransport::submitTransferTasks( slice->length = length; slice->task = task; slice->retry_count = 0; + slice->quota_charged = false; slice->ep_weak_ptr.reset(); slice->word = PENDING; slice->next = nullptr; @@ -499,8 +504,10 @@ Status RdmaTransport::submitTransferTasks( slice->priority = request.priority; // Copy priority from request task->num_slices++; task->ref(); // Each slice holds a reference to the task - if (slice_idx < slice_dev_ids.size()) + if (slice_idx < slice_dev_ids.size()) { slice->source_dev_id = slice_dev_ids[slice_idx]; + slice->quota_charged = true; + } offset += length; int part_id = next_worker_idx % num_workers; auto& list = slice_lists[part_id]; @@ -536,6 +543,19 @@ Status RdmaTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::OK(); } +Status RdmaTransport::cancelTransferTask(SubBatchRef batch, int task_id) { + auto* rdma_batch = dynamic_cast(batch); + if (!rdma_batch) { + return Status::InvalidArgument("Invalid RDMA sub-batch" LOC_MARK); + } + if (task_id < 0 || task_id >= (int)rdma_batch->task_list.size()) { + return Status::InvalidArgument("Invalid task ID" LOC_MARK); + } + auto* task = rdma_batch->task_list[task_id]; + if (task->status_word != PENDING) return Status::OK(); + return workers_->cancel(task); +} + bool RdmaTransport::warmupMemory(void* addr, size_t length) { if (length < kMrWarmupMinBytes) return false; unsigned hwc = std::thread::hardware_concurrency(); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 22e95b66af..2f58f470b1 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -47,7 +47,10 @@ RailMonitor& getOrCreateRail( } // namespace Workers::Workers(RdmaTransport* transport) - : transport_(transport), num_workers_(0), running_(false) { + : transport_(transport), + num_workers_(0), + running_(false), + worker_context_(nullptr) { device_selector_ = std::make_unique(); device_selector_->loadTopology(transport_->local_topology_); auto& conf = transport_->conf_; @@ -254,8 +257,42 @@ Status Workers::submit(RdmaSlice* slice) { return submit(slice_list); } -Status Workers::cancel(RdmaSliceList& slice_list) { - return Status::NotImplemented("cancel not implemented" LOC_MARK); +Status Workers::cancel(RdmaTask* task) { + if (!task) return Status::InvalidArgument("Invalid RDMA task" LOC_MARK); + if (task->cancel_requested.exchange(true, std::memory_order_acq_rel)) { + return Status::OK(); + } + if (!running_.load(std::memory_order_acquire) || !worker_context_ || + !num_workers_) { + return Status::OK(); + } + // Wake every worker because one task may have slices distributed across + // several queues. Cancellation remains best effort for slices already + // posted to a QP; those drain through the normal CQ path. + for (size_t id = 0; id < num_workers_; ++id) { + auto& worker = worker_context_[id]; + std::lock_guard lock(worker.mutex); + if (worker.in_suspend) worker.cv.notify_all(); + } + return Status::OK(); +} + +bool Workers::cancelUnpostedSlice(WorkerContext& worker, RdmaSlice* slice) { + if (!slice || !slice->task || + !slice->task->cancel_requested.load(std::memory_order_acquire)) + return false; + if (slice->word == PENDING) { + releaseSliceQuota(slice); + updateSliceStatus(slice, CANCELED); + } + worker.inflight_slices.fetch_sub(1); + return true; +} + +void Workers::releaseSliceQuota(RdmaSlice* slice, double latency) { + if (!slice || !slice->quota_charged || !device_selector_) return; + device_selector_->release(slice->source_dev_id, slice->length, latency); + slice->quota_charged = false; } std::shared_ptr Workers::getEndpoint(Workers::PostPath path) { @@ -353,11 +390,23 @@ void Workers::asyncPostSend() { if (slice_list.num_slices == 0) continue; auto slice = slice_list.first; for (int id = 0; id < slice_list.num_slices; ++id) { + if (cancelUnpostedSlice(worker, slice)) { + slice = slice->next; + continue; + } auto status = generatePostPath(slice); if (!status.ok()) { LOG(ERROR) << "Failed to generate post path for slice " << slice << ": " << status.ToString(); - updateSliceStatus(slice, FAILED); + releaseSliceQuota(slice); + updateSliceStatus(slice, slice->task->cancel_requested.load( + std::memory_order_acquire) + ? CANCELED + : FAILED); + worker.inflight_slices.fetch_sub(1); + } else if (cancelUnpostedSlice(worker, slice)) { + slice = slice->next; + continue; } else { PostPath path{ .local_device_id = slice->source_dev_id, @@ -373,21 +422,32 @@ void Workers::asyncPostSend() { auto& path = entry.first; auto& slices = entry.second; if (slices.empty()) continue; + slices.erase(std::remove_if(slices.begin(), slices.end(), + [&](RdmaSlice* slice) { + return cancelUnpostedSlice(worker, + slice); + }), + slices.end()); + if (slices.empty()) continue; auto endpoint = getEndpoint(path); if (!endpoint) { std::vector clone; slices.swap(clone); for (auto slice : clone) { + if (cancelUnpostedSlice(worker, slice)) continue; slice->retry_count++; if (slice->retry_count >= transport_->params_->workers.max_retry_count) { LOG(WARNING) << "Slice " << slice << " failed: retry count exceeded"; disableEndpoint(slice); + releaseSliceQuota(slice); updateSliceStatus(slice, FAILED); } else { + releaseSliceQuota(slice); submit(slice); } + worker.inflight_slices.fetch_sub(1); } continue; } @@ -396,6 +456,13 @@ void Workers::asyncPostSend() { for (int id = 0; id < num_submitted; ++id) { auto slice = slices[id]; if (slice->failed) { + releaseSliceQuota(slice); + if (slice->task->cancel_requested.load( + std::memory_order_acquire)) { + updateSliceStatus(slice, CANCELED); + worker.inflight_slices.fetch_sub(1); + continue; + } slice->retry_count++; if (slice->retry_count >= transport_->params_->workers.max_retry_count) { @@ -406,14 +473,14 @@ void Workers::asyncPostSend() { } else { submit(slice); } + worker.inflight_slices.fetch_sub(1); } else { slice->submit_ts = getCurrentTimeInNano(); + worker.inflight_slice_set.insert(slice); } } if (num_submitted) { - worker.inflight_slice_set.insert(slices.begin(), - slices.begin() + num_submitted); slices.erase(slices.begin(), slices.begin() + num_submitted); } } @@ -527,10 +594,7 @@ void Workers::asyncPollCq() { (slice->submit_ts - slice->enqueue_ts) / 1000.0; double inflight_lat = (poll_ts - slice->submit_ts) / 1000.0; double overall_lat_sec = (poll_ts - slice->enqueue_ts) / 1e9; - if (slice->retry_count == 0) { - device_selector_->release(slice->source_dev_id, slice->length, - overall_lat_sec); - } + releaseSliceQuota(slice, overall_lat_sec); if (slice->word != PENDING) continue; if (!ep) { updateSliceStatus(slice, FAILED); @@ -558,7 +622,12 @@ void Workers::asyncPollCq() { } else { num_slices += ep->acknowledge(slice, PENDING); disableEndpoint(slice); - submit(slice); + if (slice->task->cancel_requested.load( + std::memory_order_acquire)) { + updateSliceStatus(slice, CANCELED); + } else { + submit(slice); + } } } else { num_slices += ep->acknowledge(slice, COMPLETED); @@ -754,6 +823,7 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, if (slice->source_dev_id < 0) { CHECK_STATUS(device_selector_->allocate( slice->length, source.buffer->location, slice->source_dev_id)); + slice->quota_charged = true; } if (slice->source_dev_id < 0) diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 231e46ea0d..2b630d7deb 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -27,8 +27,8 @@ target_include_directories(admission_queue_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME admission_queue_test COMMAND admission_queue_test) -# Reproducible hot-path microbenchmark; intentionally not registered with -# ctest. Run manually when changing deadline promotion partitioning. +# Reproducible hot-path microbenchmark; intentionally not registered with ctest. +# Run manually when changing deadline promotion partitioning. add_executable(deadline_promotion_bench deadline_promotion_bench.cpp ../src/runtime/admission_queue.cpp) target_link_libraries(deadline_promotion_bench PRIVATE tent_common) @@ -36,11 +36,18 @@ target_include_directories(deadline_promotion_bench PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_executable(promotion_policy_test promotion_policy_test.cpp) -target_link_libraries(promotion_policy_test PRIVATE tent_common gtest gtest_main) +target_link_libraries(promotion_policy_test PRIVATE tent_common gtest + gtest_main) target_include_directories(promotion_policy_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME promotion_policy_test COMMAND promotion_policy_test) +add_executable(rdma_cancel_test rdma_cancel_test.cpp ../src/runtime/slab.cpp) +target_link_libraries(rdma_cancel_test PRIVATE tent_common gtest gtest_main) +target_include_directories(rdma_cancel_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME rdma_cancel_test COMMAND rdma_cancel_test) + add_executable(thread_local_storage_test thread_local_storage_test.cpp) target_link_libraries(thread_local_storage_test PRIVATE gtest gtest_main) target_include_directories(thread_local_storage_test @@ -112,22 +119,22 @@ target_include_directories(tent_failover_test add_test(NAME tent_failover_test COMMAND tent_failover_test) add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp) -target_link_libraries(tent_endpoint_lifecycle_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main + tent_link_group) target_include_directories(tent_endpoint_lifecycle_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_endpoint_lifecycle_test COMMAND tent_endpoint_lifecycle_test) add_executable(tent_endpoint_store_test endpoint_store_test.cpp) -target_link_libraries(tent_endpoint_store_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_endpoint_store_test PRIVATE gtest gtest_main + tent_link_group) target_include_directories(tent_endpoint_store_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_endpoint_store_test COMMAND tent_endpoint_store_test) add_executable(tent_rdma_transport_test rdma_transport_test.cpp) -target_link_libraries(tent_rdma_transport_test - PRIVATE gtest gtest_main tent_link_group ibverbs) +target_link_libraries(tent_rdma_transport_test PRIVATE gtest gtest_main + tent_link_group ibverbs) target_include_directories(tent_rdma_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rdma_transport_test COMMAND tent_rdma_transport_test) @@ -135,9 +142,8 @@ add_test(NAME tent_rdma_transport_test COMMAND tent_rdma_transport_test) if(USE_HIP) find_package(HIP REQUIRED) add_executable(tent_rocm_platform_test rocm_platform_test.cpp) - target_link_libraries(tent_rocm_platform_test PRIVATE gtest gtest_main - tent_link_group - hip::host) + target_link_libraries(tent_rocm_platform_test + PRIVATE gtest gtest_main tent_link_group hip::host) target_include_directories(tent_rocm_platform_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rocm_platform_test COMMAND tent_rocm_platform_test) @@ -146,9 +152,8 @@ endif() if(USE_SUNRISE) add_executable(tent_sunrise_link_transport_test sunrise_link_transport_test.cpp) - target_link_libraries(tent_sunrise_link_transport_test PRIVATE gtest - gtest_main - tent_link_group) + target_link_libraries(tent_sunrise_link_transport_test + PRIVATE gtest gtest_main tent_link_group) target_include_directories( tent_sunrise_link_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include @@ -158,14 +163,14 @@ if(USE_SUNRISE) endif() add_executable(tent_fault_proxy_test fault_proxy_test.cpp) target_link_libraries(tent_fault_proxy_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_fault_proxy_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_fault_proxy_test COMMAND tent_fault_proxy_test) add_executable(tent_rail_monitor_test rail_monitor_test.cpp) target_link_libraries(tent_rail_monitor_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_rail_monitor_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rail_monitor_test COMMAND tent_rail_monitor_test) @@ -173,7 +178,7 @@ add_test(NAME tent_rail_monitor_test COMMAND tent_rail_monitor_test) # Transport Selector Unit Test add_executable(tent_transport_selector_test transport_selector_test.cpp) target_link_libraries(tent_transport_selector_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_transport_selector_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_transport_selector_test COMMAND tent_transport_selector_test) @@ -181,8 +186,8 @@ add_test(NAME tent_transport_selector_test COMMAND tent_transport_selector_test) # End-to-end failover test: drives real TransferEngineImpl with # FaultProxyTransport-wrapped fakes to exercise resubmitTransferTask. add_executable(tent_engine_failover_e2e_test engine_failover_e2e_test.cpp) -target_link_libraries(tent_engine_failover_e2e_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_engine_failover_e2e_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_engine_failover_e2e_test PRIVATE asio_shared) endif() @@ -194,8 +199,8 @@ add_test(NAME tent_engine_failover_e2e_test # Per-request transport_hint: validates submitTransfer parameter, routing, # disabled-transport rejection, out-of-range rejection, mixed-hint batches. add_executable(tent_transport_hint_test transport_hint_test.cpp) -target_link_libraries(tent_transport_hint_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_transport_hint_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_transport_hint_test PRIVATE asio_shared) endif() @@ -213,20 +218,18 @@ add_test(NAME tent_intent_type_test COMMAND tent_intent_type_test) # ProgressWorker skeleton test: covers default-off behavior, event-driven # progress without poll-failover, and freeBatch races (issue #2116). add_executable(tent_progress_worker_test progress_worker_test.cpp) -target_link_libraries(tent_progress_worker_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_progress_worker_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_progress_worker_test PRIVATE asio_shared) endif() target_include_directories(tent_progress_worker_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_progress_worker_test - COMMAND tent_progress_worker_test) +add_test(NAME tent_progress_worker_test COMMAND tent_progress_worker_test) -add_executable(tent_runtime_queue_dispatch_test - runtime_queue_dispatch_test.cpp) -target_link_libraries(tent_runtime_queue_dispatch_test - PRIVATE gtest gtest_main tent_link_group) +add_executable(tent_runtime_queue_dispatch_test runtime_queue_dispatch_test.cpp) +target_link_libraries(tent_runtime_queue_dispatch_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_runtime_queue_dispatch_test PRIVATE asio_shared) endif() @@ -247,9 +250,9 @@ if(USE_TPU) add_executable(tent_tpu_pjrt_shim_test tpu/tpu_pjrt_shim_test.cpp) add_dependencies(tent_tpu_pjrt_shim_test mock_tpu_pjrt) - target_link_libraries(tent_tpu_pjrt_shim_test - PRIVATE gtest gtest_main tent_link_group - ${CMAKE_DL_LIBS}) + target_link_libraries( + tent_tpu_pjrt_shim_test PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) target_include_directories(tent_tpu_pjrt_shim_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) target_compile_definitions( @@ -261,9 +264,9 @@ if(USE_TPU) # exactly-one-device-side guard) against the same mock adapter. add_executable(tent_tpu_transport_test tpu/tpu_transport_test.cpp) add_dependencies(tent_tpu_transport_test mock_tpu_pjrt) - target_link_libraries(tent_tpu_transport_test - PRIVATE gtest gtest_main tent_link_group - ${CMAKE_DL_LIBS}) + target_link_libraries( + tent_tpu_transport_test PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) target_include_directories(tent_tpu_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) target_compile_definitions( diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index 30620a40c0..660c76e417 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -280,6 +280,40 @@ TEST(AdmissionQueueTest, RequiresDispatchBeforeTerminalCompletion) { EXPECT_EQ(status.code(), Status::Code::kInvalidEntry); } +TEST(AdmissionQueueTest, CancelsQueuedOwnerAndReleasesAccounting) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 1u); + + EXPECT_TRUE(queue.cancel(admitted_ids[0]).ok()); + EXPECT_TRUE(queue.cancel(admitted_ids[0]).ok()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); + EXPECT_TRUE(queue.pickForDispatch(1, 16).empty()); + + TransferStatusEnum status = PENDING; + ASSERT_TRUE(queue.getPublicStatus(1, 0, status).ok()); + EXPECT_EQ(status, CANCELED); + EXPECT_TRUE(queue.retireBatch(1).ok()); +} + +TEST(AdmissionQueueTest, RejectsQueueCancelAfterDispatchStarts) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + auto picked = queue.pickForDispatch(1, 16); + ASSERT_EQ(picked.size(), 1u); + + EXPECT_TRUE(queue.cancel(picked[0]).IsInvalidEntry()); + EXPECT_EQ(queue.outstandingOwners(), 1u); + EXPECT_TRUE(queue.complete(picked[0], COMPLETED).ok()); +} + TEST(AdmissionQueueTest, RetainsTerminalStatusUntilBatchRetire) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; diff --git a/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp new file mode 100644 index 0000000000..30bca1d3bd --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp @@ -0,0 +1,77 @@ +// Copyright 2026 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 "tent/transport/rdma/slice.h" + +#include + +namespace mooncake { +namespace tent { +namespace { + +TEST(RdmaCancelTest, PartialCancellationWaitsForPostedSlices) { + RdmaTask task{}; + task.num_slices = 2; + task.status_word = PENDING; + task.transferred_bytes = 0; + task.success_slices.store(0); + task.resolved_slices.store(0); + task.first_error = PENDING; + // Two slice references plus the batch reference. updateSliceStatus drops + // one reference per resolved slice; the stack object is never deallocated. + task.ref_count.store(3); + + RdmaSlice canceled{}; + canceled.task = &task; + canceled.word = PENDING; + canceled.length = 4096; + RdmaSlice posted{}; + posted.task = &task; + posted.word = PENDING; + posted.length = 8192; + + updateSliceStatus(&canceled, CANCELED); + EXPECT_EQ(task.status_word, PENDING); + EXPECT_EQ(task.transferred_bytes, 0u); + + updateSliceStatus(&posted, COMPLETED); + EXPECT_EQ(task.status_word, CANCELED); + EXPECT_EQ(task.transferred_bytes, 8192u); + EXPECT_EQ(task.resolved_slices.load(), 2); +} + +TEST(RdmaCancelTest, FullyPostedTaskMayStillComplete) { + RdmaTask task{}; + task.num_slices = 1; + task.status_word = PENDING; + task.transferred_bytes = 0; + task.success_slices.store(0); + task.resolved_slices.store(0); + task.first_error = PENDING; + task.cancel_requested.store(true); + task.ref_count.store(2); + + RdmaSlice posted{}; + posted.task = &task; + posted.word = PENDING; + posted.length = 4096; + updateSliceStatus(&posted, COMPLETED); + + EXPECT_EQ(task.status_word, COMPLETED); + EXPECT_EQ(task.transferred_bytes, 4096u); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp index bc4c4971ed..197cca1246 100644 --- a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -81,6 +82,8 @@ std::shared_ptr makeRdmaConfig() { config->set("metadata_type", "p2p"); config->set("metadata_servers", "P2PHANDSHAKE"); config->set("transports/rdma/enable", true); + config->set("transports/rdma/num_lanes", 1); + config->set("transports/rdma/endpoint/max_qp_wr", 1); config->set("transports/tcp/enable", false); config->set("transports/shm/enable", false); return config; @@ -122,6 +125,9 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { if (!hasRdmaDevice()) GTEST_SKIP() << "no RDMA device detected"; constexpr size_t kDataLength = 4 * 1024 * 1024; + constexpr size_t kCancelTaskCount = 16; + constexpr size_t kCancelStride = 8 * 1024 * 1024; + constexpr size_t kBufferLength = kCancelTaskCount * kCancelStride; int ready_pipe[2]; int stop_pipe[2]; ASSERT_EQ(pipe(ready_pipe), 0); @@ -135,7 +141,7 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { TransferEngine server(makeRdmaConfig()); if (!server.available()) _exit(2); - std::vector buffer(kDataLength); + std::vector buffer(kBufferLength); if (!server.registerLocalMemory(buffer.data(), buffer.size()).ok()) _exit(3); @@ -172,7 +178,7 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { TransferEngine client(makeRdmaConfig()); ASSERT_TRUE(client.available()); - std::vector buffer(kDataLength * 2); + std::vector buffer(kBufferLength); for (size_t i = 0; i < kDataLength; ++i) { buffer[i] = static_cast((i * 31) & 0xff); } @@ -214,6 +220,45 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { std::memcmp(buffer.data(), buffer.data() + kDataLength, kDataLength), 0); + // Keep one QP/worker and one outstanding WR so the tail task remains in + // the worker's unposted set long enough to exercise real cancellation. + std::vector cancel_requests; + cancel_requests.reserve(kCancelTaskCount); + for (size_t i = 0; i < kCancelTaskCount; ++i) { + Request cancel_request{}; + cancel_request.opcode = Request::WRITE; + cancel_request.source = buffer.data() + i * kCancelStride; + cancel_request.target_id = segment; + cancel_request.target_offset = info.buffers[0].base + i * kCancelStride; + cancel_request.length = kDataLength; + cancel_request.transport_hint = RDMA; + cancel_requests.push_back(cancel_request); + } + + batch = client.allocateBatch(kCancelTaskCount); + ASSERT_TRUE(client.submitTransfer(batch, cancel_requests).ok()); + const size_t cancel_task_id = kCancelTaskCount - 1; + ASSERT_TRUE(client.cancelTransfer(batch, cancel_task_id).ok()); + ASSERT_TRUE(client.cancelTransfer(batch, cancel_task_id).ok()); + + std::vector statuses; + for (int poll = 0; poll < 10000; ++poll) { + ASSERT_TRUE(client.getTransferStatus(batch, statuses).ok()); + if (std::all_of(statuses.begin(), statuses.end(), + [](const TransferStatus& task_status) { + return task_status.s != TransferStatusEnum::PENDING; + })) + break; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(statuses.size(), kCancelTaskCount); + for (size_t i = 0; i < cancel_task_id; ++i) { + EXPECT_EQ(statuses[i].s, TransferStatusEnum::COMPLETED); + } + EXPECT_EQ(statuses[cancel_task_id].s, TransferStatusEnum::CANCELED); + EXPECT_LE(statuses[cancel_task_id].transferred_bytes, kDataLength); + ASSERT_TRUE(client.freeBatch(batch).ok()); + EXPECT_TRUE(client.closeSegment(segment).ok()); EXPECT_TRUE( client.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); diff --git a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp index 2a94c9e5ff..b164b642a3 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -60,6 +60,8 @@ class FakeTransport : public Transport { std::atomic submit_calls{0}; std::atomic status_calls{0}; + std::atomic cancel_calls{0}; + bool cancellation_supported{true}; Status install(std::string&, std::shared_ptr, std::shared_ptr, @@ -110,6 +112,23 @@ class FakeTransport : public Transport { return Status::OK(); } + bool supportsCancellation() const override { + return cancellation_supported; + } + + Status cancelTransferTask(SubBatchRef batch, int task_id) override { + auto* fake = static_cast(batch); + if (task_id < 0 || task_id >= (int)fake->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + if (!cancellation_supported) { + return Status::NotImplemented("cancel unsupported" LOC_MARK); + } + ++cancel_calls; + fake->statuses[task_id] = {TransferStatusEnum::CANCELED, 0}; + return Status::OK(); + } + Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override { desc.transports.push_back(self_type_); return Status::OK(); @@ -358,6 +377,113 @@ TEST(RuntimeQueueDispatch, KeepsDispatchWindowUntilOwnerIsTerminal) { engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); } +TEST(RuntimeQueueDispatch, CancelsQueuedOwnerWithoutDispatchingIt) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_rdma = std::make_shared( + RDMA, [&complete_first](const Request& request, int) { + if (!complete_first.load()) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x5a); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + ASSERT_EQ(fake_rdma->submit_calls.load(), 1); + + ASSERT_TRUE(engine.cancelTransfer(batch, 1).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 0); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::CANCELED); + + complete_first.store(true); + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CancelsDispatchedRdmaTaskIdempotently) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x6b); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + + ASSERT_TRUE(engine.cancelTransfer(batch, 0).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 1); + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::CANCELED); + ASSERT_TRUE(engine.cancelTransfer(batch, 0).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsCancellationForUnsupportedTransport) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + fake_rdma->cancellation_supported = false; + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x7c); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + + EXPECT_TRUE(engine.cancelTransfer(batch, 0).IsNotImplemented()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 0); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + TEST(RuntimeQueueDispatch, ProgressWorkerRefillsWindowFromTransportNotify) { auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); cfg->set("enable_progress_worker", true);