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..ecae50cf15 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); @@ -157,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/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..62a6cfe85c 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include #include @@ -275,42 +277,106 @@ 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); +} + +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); } - 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; + + bool enqueue_failed = 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)); + return Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + + for (const auto& item : pending) { + if (!item.slice->markQueued()) break; + ++marked_count; + } + + if (marked_count == pending.size()) { + std::array inserted_counts{}; + try { + 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]; + } + } 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) { + queues_[priority].pop_back(); + --inserted_counts[priority]; + } + } } } } - if (stopped_during_submit) { - for (auto& item : pending) item.slice->requestCancellation(); + + 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(); + } + if (enqueue_failed) { + return Status::InternalError( + "Unable to enqueue UB task batch" LOC_MARK); + } return Status::InvalidArgument( - "UB workers stopped during submission" LOC_MARK); + "UB task changed while it was being submitted" LOC_MARK); } 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 4dc2694de6..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 { @@ -547,6 +555,46 @@ 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); + + 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(); @@ -906,7 +954,7 @@ 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(); @@ -928,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);