From f35768af30be72c30d129d2822e3d9c742c124f5 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:30:36 +0800 Subject: [PATCH 1/2] [TENT] Make UB batch submission all-or-nothing UbTransport::submitTransferTasks() handed tasks to UbWorkers one at a time. When a later task was rejected, the earlier tasks of the same call were already queued and could only be asked to cancel, so a failed submit could still run part of the batch. UbWorkers::submit() also skipped any slice that could not be marked queued and reported success for the rest of the task. Add UbWorkers::submitBatch(). It checks every task and slice before changing any state, then marks and queues all slices under the queue lock. If a slice changes state, the workers stop, or the queue cannot grow, nothing is published, and any task whose slices were already marked is cancelled so no slice is left stuck in the queued state. submit() now goes through submitBatch(), and submitTransferTasks() submits the whole request list in one call before adding the tasks to the sub-batch. Arm the progress callback only after the tasks are in the sub-batch. A very small task could complete before it was queryable and its notification was lost; notify once after arming to cover that window. Co-authored-by: Le1zyCatt <2490162471@qq.com> Signed-off-by: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> --- .../tent/include/tent/transport/ub/workers.h | 2 + .../tent/src/transport/ub/ub_transport.cpp | 29 +++--- .../tent/src/transport/ub/workers.cpp | 98 +++++++++++++++---- .../tent/tests/ub_native_data_path_test.cpp | 29 ++++++ 4 files changed, 126 insertions(+), 32 deletions(-) diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h index a2e9031fe8..7f3fc4e1dd 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -147,6 +147,8 @@ class UbWorkers final { Status start(); Status stop(); + Status submitBatch(const std::vector& tasks, + uint64_t device_mask = ~0ULL); Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL); Status cancel(const UbTask::Ptr& task); diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index 652b18c6d3..c590931c95 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -705,6 +706,7 @@ Status UbTransport::submitTransferTasks( const auto notify_progress = ub_batch->notify_progress; const auto progress_batch_id = ub_batch->progress_batch_id; + auto progress_armed = std::make_shared>(false); std::vector new_tasks; new_tasks.reserve(request_list.size()); for (const auto& request : request_list) { @@ -717,9 +719,12 @@ Status UbTransport::submitTransferTasks( "transports/ub/max_slices_per_task" LOC_MARK); } auto task = ub::UbTask::create( - request, - [notify_progress, progress_batch_id](const TransferStatus&) { - if (notify_progress) notify_progress(progress_batch_id); + request, [notify_progress, progress_batch_id, + progress_armed](const TransferStatus&) { + if (progress_armed->load(std::memory_order_acquire) && + notify_progress) { + notify_progress(progress_batch_id); + } }); size_t offset = 0; while (offset < request.length) { @@ -741,15 +746,15 @@ Status UbTransport::submitTransferTasks( new_tasks.push_back(std::move(task)); } - for (auto& task : new_tasks) { - ub_batch->task_list.push_back(task); - auto status = impl_->workers->submit(task, ub_batch->device_mask); - if (!status.ok()) { - for (auto& queued : new_tasks) { - if (queued) queued->requestCancellation(); - } - return status; - } + CHECK_STATUS(impl_->workers->submitBatch(new_tasks, ub_batch->device_mask)); + ub_batch->task_list.insert(ub_batch->task_list.end(), new_tasks.begin(), + new_tasks.end()); + progress_armed->store(true, std::memory_order_release); + // A very small task can finish between worker publication and the batch + // append above. Its callback intentionally stays disarmed until the task + // is queryable; this notification closes that race. + if (!new_tasks.empty() && notify_progress) { + notify_progress(progress_batch_id); } return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp index effe77fff2..37a263298c 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -275,44 +276,101 @@ Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { if (!task) { return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK); } + return submitBatch({task}, device_mask); +} + +Status UbWorkers::submitBatch(const std::vector& tasks, + uint64_t device_mask) { if (!accepting_.load(std::memory_order_acquire)) { return Status::InvalidArgument( "UB workers are not accepting work" LOC_MARK); } - const auto snapshot = task->snapshot(); - if (!snapshot.sealed) { - return Status::InvalidArgument( - "UB task must be sealed before submission" LOC_MARK); - } - const int priority = - std::clamp(task->request().priority, static_cast(PRIO_HIGH), - static_cast(PRIO_LOW)); std::vector pending; - for (const auto& slice : task->slices()) { - if (slice && slice->markQueued()) { + std::unordered_set unique_tasks; + for (const auto& task : tasks) { + if (!task) { + return Status::InvalidArgument( + "Cannot submit a null UB task" LOC_MARK); + } + if (!unique_tasks.insert(task.get()).second) { + return Status::InvalidArgument( + "Cannot submit the same UB task twice in one batch" LOC_MARK); + } + const auto snapshot = task->snapshot(); + if (!snapshot.sealed) { + return Status::InvalidArgument( + "UB task must be sealed before submission" LOC_MARK); + } + const int priority = + std::clamp(task->request().priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW)); + for (const auto& slice : task->slices()) { + if (!slice || slice->snapshot().state != UbSliceState::kInitial) { + return Status::InvalidArgument( + "UB task contains a slice that cannot be " + "submitted" LOC_MARK); + } pending.push_back(PendingSlice{task, slice, device_mask, priority, task->request().target_id, task->request().opcode}); } } - bool stopped_during_submit = false; + + Status failure = Status::OK(); + bool notify = false; + size_t marked_count = 0; { std::lock_guard lock(queue_mutex_); if (!accepting_.load(std::memory_order_relaxed)) { - stopped_during_submit = true; - } else { - for (auto& item : pending) { - queues_[item.priority].push_back(std::move(item)); + failure = Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + + if (failure.ok()) { + for (const auto& item : pending) { + if (!item.slice->markQueued()) { + failure = Status::InvalidArgument( + "UB task changed while it was being " + "submitted" LOC_MARK); + break; + } + ++marked_count; + } + } + + if (failure.ok()) { + std::array inserted_counts{}; + try { + for (const auto& item : pending) { + queues_[item.priority].push_back(item); + ++inserted_counts[item.priority]; + } + notify = !pending.empty(); + } catch (const std::bad_alloc&) { + for (size_t priority = 0; priority < queues_.size(); + ++priority) { + while (inserted_counts[priority] != 0) { + queues_[priority].pop_back(); + --inserted_counts[priority]; + } + } + failure = Status::InternalError( + "Unable to enqueue UB task batch" LOC_MARK); } } } - if (stopped_during_submit) { - for (auto& item : pending) item.slice->requestCancellation(); - return Status::InvalidArgument( - "UB workers stopped during submission" LOC_MARK); + + if (!failure.ok()) { + // A markQueued() race or allocation failure may occur after earlier + // slices changed state. None of them are visible to posting workers; + // cancel the private tasks so no slice remains orphaned in kQueued. + for (size_t index = 0; index < marked_count; ++index) { + (void)pending[index].task->requestCancellation(); + } + return failure; } - if (!pending.empty()) queue_cv_.notify_all(); + if (notify) queue_cv_.notify_all(); return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp index 4dc2694de6..447c4f47dd 100644 --- a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -547,6 +547,30 @@ TEST(UbNativeDataPathTest, (void)endpoints.retire(endpoint); }); ASSERT_TRUE(workers.start().ok()); + + Request invalid_batch_request{}; + invalid_batch_request.opcode = Request::WRITE; + invalid_batch_request.source = source.data(); + invalid_batch_request.target_id = LOCAL_SEGMENT_ID; + invalid_batch_request.target_offset = + reinterpret_cast(target.data()); + invalid_batch_request.length = 16; + auto valid_task = UbTask::create(invalid_batch_request); + auto valid_slice = valid_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, 0, 1}); + ASSERT_NE(valid_slice, nullptr); + ASSERT_TRUE(valid_task->seal()); + auto unsealed_task = UbTask::create(invalid_batch_request); + ASSERT_NE(unsealed_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, + 0, 1}), + nullptr); + + EXPECT_FALSE(workers.submitBatch({valid_task, unsealed_task}).ok()); + EXPECT_EQ(workers.queuedCount(), 0U); + EXPECT_EQ(workers.inflightCount(), 0U); + EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kInitial); + adapter->holdNextCompletion(); adapter->dropNextQuiescedCompletion(); @@ -913,6 +937,11 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { request.target_id = LOCAL_SEGMENT_ID; request.target_offset = reinterpret_cast(target.data()); request.length = source.size(); + auto invalid_request = request; + invalid_request.length = 0; + EXPECT_FALSE( + transport.submitTransferTasks(batch, {request, invalid_request}).ok()); + EXPECT_EQ(batch->size(), 0U); ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); TransferStatus transfer{PENDING, 0}; From 11226e38e44a5e7f04649cc6c317f5c0662b7293 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:38:13 +0800 Subject: [PATCH 2/2] [TENT] Handle UB batch enqueue exceptions atomically Roll back partial enqueues for allocation and length errors. Cancel marked tasks before constructing the returned error status. Inject a second-enqueue failure to verify worker rollback and preserve existing transport sub-batch task IDs after a failed submit. Signed-off-by: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> --- .../tent/include/tent/transport/ub/workers.h | 5 ++ .../tent/src/transport/ub/workers.cpp | 52 +++++++++++-------- .../tent/tests/ub_native_data_path_test.cpp | 46 +++++++++++++--- 3 files changed, 75 insertions(+), 28 deletions(-) diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h index 7f3fc4e1dd..ecae50cf15 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -159,6 +159,11 @@ class UbWorkers final { [[nodiscard]] size_t inflightCount() const; private: + friend class UbWorkersTestPeer; + // Scoped to the submitting thread so transport tests can exercise the + // real workers without changing their lifecycle or racing posting lanes. + static thread_local bool fail_second_enqueue_for_test_; + struct PendingSlice { // Keep the task alive until every queued/in-flight slice reaches a // terminal state. Callers may release the sub-batch handle before the diff --git a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp index 37a263298c..62a6cfe85c 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -279,8 +280,12 @@ Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { return submitBatch({task}, device_mask); } +thread_local bool UbWorkers::fail_second_enqueue_for_test_ = false; + Status UbWorkers::submitBatch(const std::vector& tasks, uint64_t device_mask) { + const bool fail_second_enqueue = + std::exchange(fail_second_enqueue_for_test_, false); if (!accepting_.load(std::memory_order_acquire)) { return Status::InvalidArgument( "UB workers are not accepting work" LOC_MARK); @@ -317,37 +322,37 @@ Status UbWorkers::submitBatch(const std::vector& tasks, } } - Status failure = Status::OK(); - bool notify = false; + bool enqueue_failed = false; size_t marked_count = 0; { std::lock_guard lock(queue_mutex_); if (!accepting_.load(std::memory_order_relaxed)) { - failure = Status::InvalidArgument( + return Status::InvalidArgument( "UB workers stopped during submission" LOC_MARK); } - if (failure.ok()) { - for (const auto& item : pending) { - if (!item.slice->markQueued()) { - failure = Status::InvalidArgument( - "UB task changed while it was being " - "submitted" LOC_MARK); - break; - } - ++marked_count; - } + for (const auto& item : pending) { + if (!item.slice->markQueued()) break; + ++marked_count; } - if (failure.ok()) { + if (marked_count == pending.size()) { std::array inserted_counts{}; try { - for (const auto& item : pending) { + for (size_t index = 0; index < pending.size(); ++index) { + if (fail_second_enqueue && index == 1) { + throw std::length_error("Injected UB enqueue failure"); + } + const auto& item = pending[index]; queues_[item.priority].push_back(item); ++inserted_counts[item.priority]; } - notify = !pending.empty(); } catch (const std::bad_alloc&) { + enqueue_failed = true; + } catch (const std::length_error&) { + enqueue_failed = true; + } + if (enqueue_failed) { for (size_t priority = 0; priority < queues_.size(); ++priority) { while (inserted_counts[priority] != 0) { @@ -355,22 +360,25 @@ Status UbWorkers::submitBatch(const std::vector& tasks, --inserted_counts[priority]; } } - failure = Status::InternalError( - "Unable to enqueue UB task batch" LOC_MARK); } } } - if (!failure.ok()) { - // A markQueued() race or allocation failure may occur after earlier + if (enqueue_failed || marked_count != pending.size()) { + // A markQueued() failure or enqueue exception may occur after earlier // slices changed state. None of them are visible to posting workers; // cancel the private tasks so no slice remains orphaned in kQueued. for (size_t index = 0; index < marked_count; ++index) { (void)pending[index].task->requestCancellation(); } - return failure; + if (enqueue_failed) { + return Status::InternalError( + "Unable to enqueue UB task batch" LOC_MARK); + } + return Status::InvalidArgument( + "UB task changed while it was being submitted" LOC_MARK); } - if (notify) queue_cv_.notify_all(); + if (!pending.empty()) queue_cv_.notify_all(); return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp index 447c4f47dd..101dc23232 100644 --- a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -27,6 +27,14 @@ #include "tent/transport/ub/workers.h" namespace mooncake::tent::ub { + +class UbWorkersTestPeer { + public: + static void failSecondEnqueue() { + UbWorkers::fail_second_enqueue_for_test_ = true; + } +}; + namespace { class FakeContext final : public Context { @@ -571,6 +579,22 @@ TEST(UbNativeDataPathTest, EXPECT_EQ(workers.inflightCount(), 0U); EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kInitial); + auto rollback_task = UbTask::create(invalid_batch_request); + auto rollback_slice = rollback_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, 0, 1}); + ASSERT_NE(rollback_slice, nullptr); + ASSERT_TRUE(rollback_task->seal()); + + // Both slices are marked queued and the first is inserted before the + // second enqueue throws. + UbWorkersTestPeer::failSecondEnqueue(); + EXPECT_TRUE( + workers.submitBatch({valid_task, rollback_task}).IsInternalError()); + EXPECT_EQ(workers.queuedCount(), 0U); + EXPECT_EQ(workers.inflightCount(), 0U); + EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kCanceled); + EXPECT_EQ(rollback_slice->snapshot().state, UbSliceState::kCanceled); + adapter->holdNextCompletion(); adapter->dropNextQuiescedCompletion(); @@ -930,18 +954,13 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { EXPECT_TRUE(transport.freeSubBatch(rejected_batch).ok()); Transport::SubBatchRef batch = nullptr; - ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok()); + ASSERT_TRUE(transport.allocateSubBatch(batch, 3).ok()); Request request{}; request.opcode = Request::WRITE; request.source = source.data(); request.target_id = LOCAL_SEGMENT_ID; request.target_offset = reinterpret_cast(target.data()); request.length = source.size(); - auto invalid_request = request; - invalid_request.length = 0; - EXPECT_FALSE( - transport.submitTransferTasks(batch, {request, invalid_request}).ok()); - EXPECT_EQ(batch->size(), 0U); ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); TransferStatus transfer{PENDING, 0}; @@ -957,10 +976,25 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { EXPECT_EQ(source, target); EXPECT_GT(transport.getEstimatedBandwidth(), 0.0); + ASSERT_EQ(batch->size(), 1U); + auto single_slice_request = request; + single_slice_request.length = 16; + UbWorkersTestPeer::failSecondEnqueue(); + EXPECT_TRUE(transport + .submitTransferTasks( + batch, {single_slice_request, single_slice_request}) + .IsInternalError()); + EXPECT_EQ(batch->size(), 1U); + transfer = TransferStatus{PENDING, 0}; + ASSERT_TRUE(transport.getTransferStatus(batch, 0, transfer).ok()); + EXPECT_EQ(transfer.s, COMPLETED); + EXPECT_EQ(transfer.transferred_bytes, source.size()); + source.fill(0); Request read_request = request; read_request.opcode = Request::READ; ASSERT_TRUE(transport.submitTransferTasks(batch, {read_request}).ok()); + EXPECT_EQ(batch->size(), 2U); transfer = TransferStatus{PENDING, 0}; const auto read_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);