From a59aebb3ded272771988793d9d7c56dbebae6d24 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Tue, 8 Sep 2026 17:06:28 +0800 Subject: [PATCH 1/9] =?UTF-8?q?feat(transfer-engine):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E8=B0=83=E5=BA=A6=E5=AD=90=E7=B3=BB=E7=BB=9F=EF=BC=8C=E6=94=AF?= =?UTF-8?q?=E6=8C=81=E7=AD=96=E7=95=A5=E5=8C=96=E5=87=86=E5=85=A5=E4=B8=8E?= =?UTF-8?q?=E5=88=86=E5=B1=82=E4=BB=BB=E5=8A=A1=E9=80=89=E6=8B=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 引入 SchedulerCore 实现调度循环，支持任务提交、取消、状态查询；提供可插拔策略：意图解析、分层准入、限界调度预算、分层任务选择；为 RDMA 和 UB 传输添加 scheduledTransferLength 接口；新增调度策略与核心的单元测试。 --- .../src/scheduler_policy.cpp | 189 +++++++++++++ .../src/transfer_engine_scheduling.cpp | 38 +++ .../kunpeng_transport/ub_transport.cpp | 23 ++ .../rdma_transport/rdma_transport.cpp | 20 ++ mooncake-transfer-engine/tests/CMakeLists.txt | 10 + .../tests/scheduler_core_test.cpp | 250 ++++++++++++++++++ .../tests/scheduler_policy_test.cpp | 130 +++++++++ 7 files changed, 660 insertions(+) create mode 100644 mooncake-transfer-engine/src/scheduler_policy.cpp create mode 100644 mooncake-transfer-engine/src/transfer_engine_scheduling.cpp create mode 100644 mooncake-transfer-engine/tests/scheduler_core_test.cpp create mode 100644 mooncake-transfer-engine/tests/scheduler_policy_test.cpp diff --git a/mooncake-transfer-engine/src/scheduler_policy.cpp b/mooncake-transfer-engine/src/scheduler_policy.cpp new file mode 100644 index 0000000000..6595036eb0 --- /dev/null +++ b/mooncake-transfer-engine/src/scheduler_policy.cpp @@ -0,0 +1,189 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#include "scheduler/scheduler_policy.h" + +#include +#include +#include +#include +#include + +namespace mooncake::scheduling { +namespace { + +class HierarchicalAdmissionPolicy final : public AdmissionPolicy { + public: + AdmissionDecision admit(uint64_t bytes, const RuntimeSnapshot& runtime, + const SchedulerConfig& config) const override { + if (bytes > config.max_outstanding_bytes || + bytes > config.max_tenant_outstanding_bytes) + return AdmissionDecision::REJECT; + if (runtime.outstanding_tasks >= config.max_outstanding_tasks || + bytes > config.max_outstanding_bytes - runtime.outstanding_bytes || + runtime.tenant_outstanding_tasks >= + config.max_tenant_outstanding_tasks || + bytes > config.max_tenant_outstanding_bytes - + runtime.tenant_outstanding_bytes) + return AdmissionDecision::DEFER; + return AdmissionDecision::ACCEPT; + } +}; + +class BoundedDispatchBudgetPolicy final : public DispatchBudgetPolicy { + public: + uint64_t budget(const ReadyTaskView& task, uint64_t fair_bytes, + const RuntimeSnapshot& runtime, + const SchedulerConfig& config) const override { + uint64_t limit = config.max_inflight_bytes; + if (task.qos.traffic_class != TrafficClass::HIGH) + limit -= config.reserved_high_bytes; + if (runtime.inflight_bytes >= limit) return 0; + return std::min({task.remaining_bytes, config.quantum_bytes, fair_bytes, + limit - runtime.inflight_bytes}); + } +}; + +class IntentQoSPolicy final : public QoSResolutionPolicy { + public: + ResolvedQoS resolve(const SchedulingHint& hint, + const SchedulerConfig& config) const override { + ResolvedQoS qos; + switch (hint.intent) { + case TaskIntent::CONTROL: + case TaskIntent::FOREGROUND_GET: + case TaskIntent::P2D_TRANSFER: + qos.traffic_class = TrafficClass::HIGH; + break; + case TaskIntent::BACKGROUND_PUT: + case TaskIntent::MIGRATION: + case TaskIntent::CHECKPOINT: + qos.traffic_class = TrafficClass::LOW; + break; + default: + break; + } + qos.priority_rank = hint.requested_priority.value_or(0); + auto it = config.tenant_weights.find(hint.tenant_id); + if (it != config.tenant_weights.end()) qos.weight = it->second; + return qos; + } +}; + +class HierarchicalTaskPolicy final : public TaskSelectionPolicy { + public: + explicit HierarchicalTaskPolicy(SchedulerConfig config) + : config_(std::move(config)) {} + + TaskSelection select(const std::vector& ready, + uint64_t now_ns) override { + using TenantTasks = + std::map>; + std::array groups; + for (const auto& task : ready) { + groups[static_cast(task.qos.traffic_class)] + [task.requested->tenant_id] + .push_back(&task); + } + // Drop inactive domains rather than accumulating credit while idle. + for (size_t c = 0; c < groups.size(); ++c) { + if (groups[c].empty()) class_deficit_[c] = 0; + for (auto it = tenant_deficit_[c].begin(); + it != tenant_deficit_[c].end();) { + if (!groups[c].count(it->first)) + it = tenant_deficit_[c].erase(it); + else + ++it; + } + } + for (size_t attempt = 0; attempt < 3; ++attempt) { + const size_t c = class_cursor_; + if (groups[c].empty()) { + class_cursor_ = (c + 1) % 3; + continue; + } + if (!class_deficit_[c]) { + class_deficit_[c] = + config_.quantum_bytes * config_.class_weights[c]; + } + auto tenant = groups[c].lower_bound(tenant_cursor_[c]); + if (tenant == groups[c].end()) tenant = groups[c].begin(); + auto& deficit = tenant_deficit_[c][tenant->first]; + if (!deficit) { + deficit = + config_.quantum_bytes * tenant->second.front()->qos.weight; + } + auto key = [&](const ReadyTaskView* task) { + uint64_t age = + now_ns > task->enqueue_ns ? now_ns - task->enqueue_ns : 0; + uint64_t bonus = + config_.aging_interval_ns + ? std::min(age / config_.aging_interval_ns, + config_.max_aging_bonus) + : 0; + int64_t priority = + int64_t(task->qos.priority_rank) - int64_t(bonus); + return std::make_tuple( + priority, + task->requested->deadline_ns.value_or( + std::numeric_limits::max()), + task->sequence); + }; + auto task = *std::min_element( + tenant->second.begin(), tenant->second.end(), + [&](auto* a, auto* b) { return key(a) < key(b); }); + selected_class_ = c; + selected_tenant_ = tenant->first; + auto next = std::next(tenant); + next_tenant_ = next == groups[c].end() ? groups[c].begin()->first + : next->first; + return {task->key, std::min(class_deficit_[c], deficit)}; + } + return {}; + } + + void accepted(const ReadyTaskView&, uint64_t bytes) override { + auto c = selected_class_; + auto& tenant = tenant_deficit_[c][selected_tenant_]; + class_deficit_[c] -= bytes; + tenant -= bytes; + tenant_cursor_[c] = tenant ? selected_tenant_ : next_tenant_; + if (!class_deficit_[c]) class_cursor_ = (c + 1) % 3; + } + + private: + SchedulerConfig config_; + std::array class_deficit_{}; + std::array, 3> tenant_deficit_; + std::array tenant_cursor_; + size_t class_cursor_{0}; + size_t selected_class_{0}; + std::string selected_tenant_; + std::string next_tenant_; +}; + +} // namespace + +std::unique_ptr makeIntentQoSPolicy() { + return std::make_unique(); +} + +std::unique_ptr makeHierarchicalAdmissionPolicy() { + return std::make_unique(); +} + +std::unique_ptr makeBoundedDispatchBudgetPolicy() { + return std::make_unique(); +} + +std::unique_ptr makeHierarchicalTaskPolicy( + const SchedulerConfig& config) { + return std::make_unique(config); +} + +SchedulerPolicySet makeDefaultPolicySet(const SchedulerConfig& config) { + return {makeIntentQoSPolicy(), makeHierarchicalAdmissionPolicy(), + makeHierarchicalTaskPolicy(config), + makeBoundedDispatchBudgetPolicy()}; +} + +} // namespace mooncake::scheduling diff --git a/mooncake-transfer-engine/src/transfer_engine_scheduling.cpp b/mooncake-transfer-engine/src/transfer_engine_scheduling.cpp new file mode 100644 index 0000000000..b5c57b934c --- /dev/null +++ b/mooncake-transfer-engine/src/transfer_engine_scheduling.cpp @@ -0,0 +1,38 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#include "transfer_engine.h" +#include "transfer_engine_impl.h" + +namespace mooncake { + +Status TransferEngine::configureScheduling( + const scheduling::SchedulerConfig& config) { + if (use_tent_ || !impl_) + return Status::NotImplemented( + "Scheduling requires initialized legacy TE"); + return impl_->configureScheduling(config); +} + +Status TransferEngine::submitScheduledTransfer( + BatchID batch_id, const std::vector& entries) { + if (use_tent_) { + std::vector requests; + requests.reserve(entries.size()); + for (const auto& entry : entries) requests.push_back(entry.request); + // TENT retains its existing QoS/admission path until its scheduling + // adapter can consume the common hint contract. + return submitTransfer(batch_id, requests); + } + if (!impl_) + return Status::InvalidArgument("Transfer Engine is not initialized"); + return impl_->submitScheduledTransfer(batch_id, entries); +} + +Status TransferEngine::cancelTransfer(BatchID batch_id, size_t task_id) { + if (use_tent_ || !impl_) + return Status::NotImplemented( + "Scheduled cancellation requires legacy TE"); + return impl_->cancelTransfer(batch_id, task_id); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 6016f828d1..d926eeae8e 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -728,6 +728,17 @@ Status UbTransport::submitTransferTask( << "UbTransport: Address not registered by any device(s) " << source_addr; if (staged_request) cleanupStagingForTask(&task, true); + if (task.scheduled) { + // Earlier watermark batches may already be in flight. + // Complete only the unposted suffix here; the scheduler + // retains the grant until that prefix also drains. + __sync_fetch_and_add(&task.slice_count, 1); + slice->markFailed(); + for (auto& pending : slices_to_post) + for (auto* unposted : pending.second) + unposted->markFailed(); + slices_to_post.clear(); + } return Status::AddressNotRegistered( "UbTransport: not registered by any device(s), " "address: " + @@ -820,6 +831,18 @@ Status UbTransport::getTransferStatus(BatchID batch_id, size_t task_id, return Status::OK(); } +Status UbTransport::scheduledTransferLength(const TransferRequest& request, + uint32_t max_slices, + size_t& length) { + const uint64_t block = globalConfig().slice_size; + if (!block || !max_slices) + return Status::InvalidArgument("Invalid UB slice budget"); + length = request.length; + if (uint64_t(max_slices) <= UINT64_MAX / block) + length = std::min(length, block * max_slices); + return Status::OK(); +} + Transport::SegmentID UbTransport::getSegmentID( const std::string& segment_name) { return metadata_->getSegmentID(segment_name); diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index fc96e4505f..f26c561b56 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -166,6 +166,26 @@ RdmaTransport::RdmaTransport() { } } +Status RdmaTransport::scheduledTransferLength(const TransferRequest &request, + uint32_t max_slices, + size_t &length) { + auto local = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); + auto target = metadata_->getSegmentDescByID(request.target_id); + if (!local || !target || !max_slices || !globalConfig().slice_size) + return Status::InvalidArgument("Invalid RDMA scheduling range"); + SliceLengthCalculator calculator{request, globalConfig().slice_size, + globalConfig().fragment_limit, local.get(), + target.get()}; + length = 0; + for (uint32_t count = 0; count < max_slices && length < request.length; + ++count) { + auto bytes = calculator.calculate(length); + if (!bytes) return Status::InvalidArgument("Empty RDMA slice"); + length += bytes; + } + return Status::OK(); +} + RdmaTransport::~RdmaTransport() { #ifdef CONFIG_USE_BATCH_DESC_SET for (auto &entry : batch_desc_set_) delete entry.second; diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index f68f55e26f..faf2a81452 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -1,5 +1,15 @@ set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}") +add_executable(scheduler_policy_test scheduler_policy_test.cpp) +target_link_libraries(scheduler_policy_test PRIVATE transfer_engine gtest + gtest_main) +add_test(NAME scheduler_policy_test COMMAND scheduler_policy_test) + +add_executable(scheduler_core_test scheduler_core_test.cpp) +target_link_libraries(scheduler_core_test PRIVATE transfer_engine gtest + gtest_main) +add_test(NAME scheduler_core_test COMMAND scheduler_core_test) + if(USE_HIP) file(GLOB TEST_SOURCES "*.cpp") hipify_files(TEST_SOURCES) diff --git a/mooncake-transfer-engine/tests/scheduler_core_test.cpp b/mooncake-transfer-engine/tests/scheduler_core_test.cpp new file mode 100644 index 0000000000..9ee9cca791 --- /dev/null +++ b/mooncake-transfer-engine/tests/scheduler_core_test.cpp @@ -0,0 +1,250 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#include + +#include +#include +#include + +#include "scheduler/scheduler_core.h" + +namespace mooncake::scheduling { +namespace { + +class FakeTransport final : public Transport { + public: + bool complete_inline{true}; + bool reject{false}; + std::vector ranges; + std::mutex mutex; + std::deque pending; + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::NotImplemented("Use prepared tasks"); + } + Status scheduledTransferLength(const TransferRequest& request, uint32_t, + size_t& length) override { + length = request.length; + return Status::OK(); + } + Status submitTransferTask( + const std::vector& tasks) override { + std::lock_guard lock(mutex); + for (auto* task : tasks) { + ranges.push_back(*task->request); + if (reject) return Status::InvalidArgument("Injected rejection"); + auto* slice = getSliceCache().allocate(); + slice->task = task; + slice->length = task->request->length; + task->slice_list.push_back(slice); + ++task->slice_count; + if (complete_inline) + slice->markSuccess(); + else + pending.push_back(slice); + } + return Status::OK(); + } + void drain() { + std::lock_guard lock(mutex); + complete_inline = true; + while (!pending.empty()) { + pending.front()->markSuccess(); + pending.pop_front(); + } + } + size_t submitted() { + std::lock_guard lock(mutex); + return ranges.size(); + } + Status getTransferStatus(BatchID, size_t, TransferStatus&) override { + return Status::NotImplemented("Scheduler aggregates completion"); + } + 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 "fake"; } +}; + +class SchedulerCoreTest : public ::testing::Test { + protected: + FakeTransport transport; + Transport::BatchDesc batch; + std::unique_ptr scheduler; + std::array data{}; + + void SetUp() override { + batch.id = reinterpret_cast(&batch); + batch.batch_size = 16; + batch.task_list.reserve(batch.batch_size); + SchedulerConfig config; + config.quantum_bytes = 16; + config.max_inflight_bytes = 16; + config.reserved_high_bytes = 0; + config.max_outstanding_tasks = 4; + config.max_outstanding_bytes = 128; + config.max_deferred_tasks = 0; + scheduler = std::make_unique( + config, [&](const auto&, Transport*& selected) { + selected = &transport; + return Status::OK(); + }); + } + void TearDown() override { + transport.drain(); + scheduler.reset(); + } + ScheduledTransferRequest request(size_t bytes) { + return { + {Transport::TransferRequest::WRITE, data.data(), 1, 1000, bytes}, + {}}; + } + template + bool wait(Predicate predicate) { + auto limit = std::chrono::steady_clock::now() + std::chrono::seconds(2); + do { + if (predicate()) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } while (std::chrono::steady_clock::now() < limit); + return false; + } + bool finished(Transport::TransferStatusEnum expected) { + return wait([&] { + Transport::TransferStatus status; + return scheduler->batchStatus(batch.id, status).ok() && + status.s == expected; + }); + } +}; + +TEST_F(SchedulerCoreTest, OwnsRequestsAndAdvancesNonoverlappingRanges) { + { + std::vector temporary{request(70)}; + ASSERT_TRUE(scheduler->submit(batch.id, temporary).ok()); + temporary[0].request.source = nullptr; + } + ASSERT_TRUE(finished(Transport::COMPLETED)); + Transport::TransferStatus status; + ASSERT_TRUE(scheduler->status(batch.id, 0, status).ok()); + EXPECT_EQ(status.transferred_bytes, 70); + std::lock_guard lock(transport.mutex); + ASSERT_EQ(transport.ranges.size(), 5); + size_t offset = 0; + for (const auto& range : transport.ranges) { + EXPECT_EQ(range.source, data.data() + offset); + EXPECT_EQ(range.target_offset, 1000 + offset); + EXPECT_LE(range.length, 16); + offset += range.length; + } + EXPECT_EQ(offset, 70); +} + +TEST_F(SchedulerCoreTest, RunningCancelWaitsForDrainAndStopsFurtherGrants) { + transport.complete_inline = false; + ASSERT_TRUE(scheduler->submit(batch.id, {request(80)}).ok()); + ASSERT_TRUE(wait([&] { return transport.submitted() == 1; })); + ASSERT_TRUE(scheduler->cancel(batch.id, 0).ok()); + EXPECT_FALSE(scheduler->release(batch.id).ok()); + Transport::TransferStatus status; + ASSERT_TRUE(scheduler->status(batch.id, 0, status).ok()); + EXPECT_EQ(status.s, Transport::WAITING); + transport.drain(); + ASSERT_TRUE(finished(Transport::CANCELED)); + EXPECT_EQ(transport.submitted(), 1); + ASSERT_TRUE(scheduler->status(batch.id, 0, status).ok()); + EXPECT_EQ(status.transferred_bytes, 16); + EXPECT_TRUE(scheduler->release(batch.id).ok()); + EXPECT_TRUE(batch.task_list[0].is_finished); + EXPECT_TRUE(batch.has_failure.load()); +} + +TEST_F(SchedulerCoreTest, AdmissionFailureIsAtomicAndCapacityIsReclaimed) { + transport.complete_inline = false; + ASSERT_TRUE(scheduler->submit(batch.id, {request(80)}).ok()); + EXPECT_FALSE(scheduler->submit(batch.id, {request(64), request(16)}).ok()); + EXPECT_EQ(batch.task_list.size(), 1); + transport.drain(); + ASSERT_TRUE(finished(Transport::COMPLETED)); + ASSERT_TRUE(scheduler->submit(batch.id, {request(64)}).ok()); + ASSERT_TRUE(finished(Transport::COMPLETED)); + EXPECT_EQ(batch.task_list.size(), 2); +} + +TEST_F(SchedulerCoreTest, SynchronousFailureDoesNotLeakInflightCredits) { + transport.reject = true; + ASSERT_TRUE(scheduler->submit(batch.id, {request(64)}).ok()); + ASSERT_TRUE(finished(Transport::FAILED)); + { + std::lock_guard lock(transport.mutex); + transport.reject = false; + } + ASSERT_TRUE(scheduler->submit(batch.id, {request(64)}).ok()); + ASSERT_TRUE(wait([&] { + Transport::TransferStatus status; + return scheduler->status(batch.id, 1, status).ok() && + status.s == Transport::COMPLETED; + })); +} + +TEST_F(SchedulerCoreTest, ZeroLengthCompletesWithoutPhysicalSubmission) { + ASSERT_TRUE(scheduler->submit(batch.id, {request(0)}).ok()); + ASSERT_TRUE(finished(Transport::COMPLETED)); + EXPECT_EQ(transport.submitted(), 0); +} + +TEST_F(SchedulerCoreTest, HighTaskCompetesAtNextGrantBoundary) { + transport.complete_inline = false; + auto low = request(64); + low.hint.intent = TaskIntent::BACKGROUND_PUT; + ASSERT_TRUE(scheduler->submit(batch.id, {low}).ok()); + ASSERT_TRUE(wait([&] { return transport.submitted() == 1; })); + auto high = request(8); + high.request.target_offset = 2000; + high.hint.intent = TaskIntent::FOREGROUND_GET; + ASSERT_TRUE(scheduler->submit(batch.id, {high}).ok()); + Transport::TransferStatus status; + ASSERT_TRUE(scheduler->status(batch.id, 1, status).ok()); + EXPECT_EQ(status.s, Transport::PENDING); + transport.drain(); + ASSERT_TRUE(finished(Transport::COMPLETED)); + std::lock_guard lock(transport.mutex); + ASSERT_GE(transport.ranges.size(), 2); + EXPECT_EQ(transport.ranges[1].target_offset, 2000); +} + +TEST_F(SchedulerCoreTest, DeferredTaskIsReadmittedAfterQuotaRelease) { + scheduler.reset(); + SchedulerConfig config; + config.max_outstanding_tasks = 1; + config.max_outstanding_bytes = 128; + config.max_deferred_tasks = 1; + config.max_deferred_bytes = 128; + scheduler = std::make_unique( + config, [&](const auto&, Transport*& selected) { + selected = &transport; + return Status::OK(); + }); + transport.complete_inline = false; + ASSERT_TRUE(scheduler->submit(batch.id, {request(32), request(32)}).ok()); + ASSERT_TRUE(wait([&] { return transport.submitted() == 1; })); + Transport::TransferStatus status; + ASSERT_TRUE(scheduler->status(batch.id, 1, status).ok()); + EXPECT_EQ(status.s, Transport::PENDING); + EXPECT_FALSE(scheduler->submit(batch.id, {request(1)}).ok()); + transport.drain(); + ASSERT_TRUE(finished(Transport::COMPLETED)); + EXPECT_EQ(transport.submitted(), 2); +} + +} // namespace +} // namespace mooncake::scheduling diff --git a/mooncake-transfer-engine/tests/scheduler_policy_test.cpp b/mooncake-transfer-engine/tests/scheduler_policy_test.cpp new file mode 100644 index 0000000000..f0095f17ea --- /dev/null +++ b/mooncake-transfer-engine/tests/scheduler_policy_test.cpp @@ -0,0 +1,130 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#include + +#include + +#include "scheduler/scheduler_policy.h" + +namespace mooncake::scheduling { + +TEST(SchedulerPolicy, IntentAndExplicitZeroAreIndependent) { + auto policy = makeIntentQoSPolicy(); + SchedulerConfig config; + SchedulingHint hint; + hint.intent = TaskIntent::BACKGROUND_PUT; + hint.requested_priority = 0; + auto qos = policy->resolve(hint, config); + EXPECT_EQ(qos.traffic_class, TrafficClass::LOW); + EXPECT_EQ(qos.priority_rank, 0); + hint.requested_priority = std::numeric_limits::min(); + EXPECT_EQ(policy->resolve(hint, config).priority_rank, + std::numeric_limits::min()); + hint.intent = TaskIntent::FOREGROUND_GET; + EXPECT_EQ(policy->resolve(hint, config).traffic_class, TrafficClass::HIGH); +} + +TEST(SchedulerPolicy, HierarchicalByteShares) { + SchedulerConfig config; + config.quantum_bytes = 1024; + config.class_weights = {4, 2, 1}; + auto policy = makeHierarchicalTaskPolicy(config); + std::array hints; + hints[0].tenant_id = "A"; + hints[1].tenant_id = "B"; + hints[2].tenant_id = "C"; + hints[3].tenant_id = "D"; + std::vector ready{ + {1, &hints[0], {TrafficClass::HIGH, 0, 2}, 1024, 0, 1}, + {2, &hints[1], {TrafficClass::HIGH, 0, 1}, 1024, 0, 2}, + {3, &hints[2], {TrafficClass::MEDIUM, 0, 1}, 1024, 0, 3}, + {4, &hints[3], {TrafficClass::LOW, 0, 1}, 1024, 0, 4}}; + std::array bytes{}; + for (int grant = 0; grant < 210; ++grant) { + auto selected = policy->select(ready, 0); + ASSERT_GE(selected.available_bytes, 1024); + ASSERT_GE(selected.key, 1); + ASSERT_LE(selected.key, 4); + bytes[selected.key - 1] += 1024; + policy->accepted(ready[selected.key - 1], 1024); + } + EXPECT_EQ(bytes[0], 80 * 1024); + EXPECT_EQ(bytes[1], 40 * 1024); + EXPECT_EQ(bytes[2], 60 * 1024); + EXPECT_EQ(bytes[3], 30 * 1024); +} + +TEST(SchedulerPolicy, PriorityDeadlineAndStableSequence) { + SchedulerConfig config; + auto policy = makeHierarchicalTaskPolicy(config); + SchedulingHint a, b; + a.deadline_ns = 100; + b.deadline_ns = 50; + std::vector ready{ + {1, &a, {TrafficClass::MEDIUM, -1, 1}, 16, 0, 2}, + {2, &b, {TrafficClass::MEDIUM, 0, 1}, 16, 0, 1}}; + EXPECT_EQ(policy->select(ready, 0).key, 1); + ready[1].qos.priority_rank = -1; + EXPECT_EQ(policy->select(ready, 0).key, 2); + b.deadline_ns = 100; + EXPECT_EQ(policy->select(ready, 0).key, 2); +} + +TEST(SchedulerPolicy, AgingDoesNotOverflowOrChangeClass) { + SchedulerConfig config; + config.aging_interval_ns = 1; + auto policy = makeHierarchicalTaskPolicy(config); + SchedulingHint a, b; + std::vector ready{ + {1, + &a, + {TrafficClass::LOW, std::numeric_limits::min(), 1}, + 16, + 0, + 2}, + {2, &b, {TrafficClass::LOW, -10, 1}, 16, 100, 1}}; + EXPECT_EQ(policy->select(ready, 100).key, 1); +} + +TEST(SchedulerPolicy, RejectedDispatchDoesNotConsumeCredit) { + SchedulerConfig config; + auto policy = makeHierarchicalTaskPolicy(config); + SchedulingHint hint; + std::vector ready{{1, &hint, {}, 128, 0, 1}}; + auto first = policy->select(ready, 0); + auto retry = policy->select(ready, 0); + EXPECT_EQ(first.key, retry.key); + EXPECT_EQ(first.available_bytes, retry.available_bytes); + policy->accepted(ready[0], 64); + EXPECT_EQ(policy->select(ready, 0).available_bytes, + first.available_bytes - 64); +} + +TEST(SchedulerPolicy, AdmissionDistinguishesPermanentAndTransientLimits) { + SchedulerConfig config; + config.max_outstanding_bytes = 1024; + config.max_tenant_outstanding_bytes = 512; + auto policy = makeHierarchicalAdmissionPolicy(); + RuntimeSnapshot runtime; + EXPECT_EQ(policy->admit(513, runtime, config), AdmissionDecision::REJECT); + runtime.tenant_outstanding_bytes = 500; + EXPECT_EQ(policy->admit(16, runtime, config), AdmissionDecision::DEFER); + runtime.tenant_outstanding_bytes = 0; + EXPECT_EQ(policy->admit(16, runtime, config), AdmissionDecision::ACCEPT); +} + +TEST(SchedulerPolicy, HighReservationCannotBeSpentByLow) { + SchedulerConfig config; + config.max_inflight_bytes = 1024; + config.reserved_high_bytes = 256; + auto policy = makeBoundedDispatchBudgetPolicy(); + SchedulingHint hint; + ReadyTaskView task{1, &hint, {TrafficClass::LOW, 0, 1}, 1024, 0, 1}; + RuntimeSnapshot runtime; + runtime.inflight_bytes = 768; + EXPECT_EQ(policy->budget(task, 1024, runtime, config), 0); + task.qos.traffic_class = TrafficClass::HIGH; + EXPECT_EQ(policy->budget(task, 1024, runtime, config), 256); +} + +} // namespace mooncake::scheduling From 87196643e8931bf6ac1741c322a0c8b0706c3761 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Tue, 8 Sep 2026 17:07:11 +0800 Subject: [PATCH 2/9] =?UTF-8?q?feat(scheduler):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E8=B0=83=E5=BA=A6=E5=99=A8=E6=A0=B8=E5=BF=83=E5=B9=B6=E9=9B=86?= =?UTF-8?q?=E6=88=90=E5=88=B0=E5=A4=9A=E4=BC=A0=E8=BE=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 实现 SchedulerCore 调度核心，支持 QoS 分类、准入控制、分片调度与批量任务生命周期管理。 集成到 MultiTransport，新增 configureScheduling、submitScheduledTransfer 等接口， 并适配 freeBatchID、submitTransfer、getTransferStatus 等方法以处理调度器场景。 --- .../src/multi_transport.cpp | 57 +++ .../src/scheduler_core.cpp | 469 ++++++++++++++++++ 2 files changed, 526 insertions(+) create mode 100644 mooncake-transfer-engine/src/scheduler_core.cpp diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 2c8d53988b..8e757216ca 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -83,6 +83,46 @@ MultiTransport::MultiTransport(std::shared_ptr metadata, MultiTransport::~MultiTransport() {} +Status MultiTransport::configureScheduling( + const scheduling::SchedulerConfig& config) { + if (scheduler_) return Status::InvalidArgument("Scheduler already configured"); + auto status = scheduling::SchedulerCore::validate(config); + if (!status.ok()) return status; + scheduler_ = std::make_unique( + config, [this](const TransferRequest& request, Transport*& transport) { + auto status = selectTransport(request, transport); + if (!status.ok()) return status; + // These transports support byte-addressed ranges and publish slice + // completion from their workers. Other protocols need an adapter + // capability implementation before enabling range scheduling. + for (const char* proto : {"rdma", "tcp", "ub"}) { + auto it = transport_map_.find(proto); + if (it != transport_map_.end() && it->second.get() == transport) + return Status::OK(); + } + return Status::NotSupportedTransport( + "Transport does not support scheduled byte ranges"); + }); + return Status::OK(); +} + +Status MultiTransport::submitScheduledTransfer( + BatchID batch_id, const std::vector& entries) { + if (!scheduler_) { + std::vector requests; + requests.reserve(entries.size()); + for (const auto& entry : entries) requests.push_back(entry.request); + return submitTransfer(batch_id, requests); + } + return scheduler_->submit(batch_id, entries); +} + +Status MultiTransport::cancelTransfer(BatchID batch_id, size_t task_id) { + if (!scheduler_) + return Status::InvalidArgument("Scheduling is not enabled"); + return scheduler_->cancel(batch_id, task_id); +} + MultiTransport::BatchID MultiTransport::allocateBatchID(size_t batch_size) { auto batch_desc = new BatchDesc(); if (!batch_desc) return ERR_MEMORY; @@ -99,6 +139,10 @@ MultiTransport::BatchID MultiTransport::allocateBatchID(size_t batch_size) { } Status MultiTransport::freeBatchID(BatchID batch_id) { + if (scheduler_) { + auto status = scheduler_->release(batch_id); + if (!status.ok()) return status; + } 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++) { @@ -117,6 +161,12 @@ Status MultiTransport::freeBatchID(BatchID batch_id) { Status MultiTransport::submitTransfer( BatchID batch_id, const std::vector& entries) { + if (scheduler_) { + std::vector scheduled; + scheduled.reserve(entries.size()); + for (const auto& entry : entries) scheduled.push_back({entry, {}}); + return scheduler_->submit(batch_id, scheduled); + } auto& batch_desc = *((BatchDesc*)(batch_id)); if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { return Status::TooManyRequests( @@ -160,6 +210,9 @@ Status MultiTransport::submitTransfer( Status MultiTransport::mp_submitTransfer( BatchID batch_id, const std::vector& entries, std::string& proto) { + if (scheduler_) + return Status::NotImplemented( + "Explicit multi-protocol scheduling is not supported"); auto& batch_desc = *((BatchDesc*)(batch_id)); if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { return Status::TooManyRequests( @@ -202,6 +255,8 @@ Status MultiTransport::mp_submitTransfer( Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus& status) { + if (scheduler_ && scheduler_->owns(batch_id)) + return scheduler_->status(batch_id, task_id, status); auto& batch_desc = *((BatchDesc*)(batch_id)); const size_t task_count = batch_desc.task_list.size(); if (task_id >= task_count) { @@ -270,6 +325,8 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, Status MultiTransport::getBatchTransferStatus(BatchID batch_id, TransferStatus& status) { + if (scheduler_ && scheduler_->owns(batch_id)) + return scheduler_->batchStatus(batch_id, status); auto& batch_desc = *((BatchDesc*)(batch_id)); const size_t task_count = batch_desc.task_list.size(); status.transferred_bytes = 0; diff --git a/mooncake-transfer-engine/src/scheduler_core.cpp b/mooncake-transfer-engine/src/scheduler_core.cpp new file mode 100644 index 0000000000..f7c283d162 --- /dev/null +++ b/mooncake-transfer-engine/src/scheduler_core.cpp @@ -0,0 +1,469 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#include "scheduler/scheduler_core.h" + +#include +#include +#include +#include +#include + +#include "config.h" + +namespace mooncake::scheduling { + +uint64_t SchedulerCore::nowNs() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); +} + +Status SchedulerCore::validate(const SchedulerConfig& config) { + if (!config.max_outstanding_tasks || !config.max_outstanding_bytes || + !config.max_inflight_bytes || !config.quantum_bytes || + !config.max_slices || !config.max_tenant_outstanding_tasks || + !config.max_tenant_outstanding_bytes || + config.max_deferred_bytes > std::numeric_limits::max() - + config.max_outstanding_bytes || + config.max_deferred_tasks > + std::numeric_limits::max() - config.max_outstanding_tasks || + config.reserved_high_bytes >= config.max_inflight_bytes) { + return Status::InvalidArgument("Invalid scheduler capacity"); + } + auto valid_weight = [&](uint32_t weight) { + return weight && config.quantum_bytes <= + std::numeric_limits::max() / weight; + }; + for (auto weight : config.class_weights) { + if (!valid_weight(weight)) + return Status::InvalidArgument("Invalid scheduler class weight"); + } + for (const auto& entry : config.tenant_weights) { + if (!valid_weight(entry.second)) + return Status::InvalidArgument("Invalid scheduler tenant weight"); + } + return Status::OK(); +} + +SchedulerCore::SchedulerCore(SchedulerConfig config, Route route) + : SchedulerCore(config, std::move(route), makeDefaultPolicySet(config)) {} + +SchedulerCore::SchedulerCore(SchedulerConfig config, Route route, + SchedulerPolicySet policies) + : config_(std::move(config)), + route_(std::move(route)), + qos_(std::move(policies.qos)), + selection_(std::move(policies.task_selection)), + admission_(std::move(policies.admission)), + budget_(std::move(policies.dispatch_budget)) { + if (!validate(config_).ok() || !route_ || !qos_ || !selection_ || + !admission_ || !budget_) + throw std::invalid_argument( + "Invalid scheduler configuration or policies"); + worker_ = std::thread(&SchedulerCore::run, this); +} + +SchedulerCore::~SchedulerCore() { + { + std::lock_guard lock(mutex_); + stopping_ = true; + for (auto& batch : batches_) + for (auto& task : batch.second.tasks) task->canceling = true; + } + wake_.notify_one(); + worker_.join(); +} + +bool SchedulerCore::terminal(const Task& task) { + return task.state == Transport::COMPLETED || + task.state == Transport::FAILED || task.state == Transport::CANCELED; +} + +ReadyTaskView SchedulerCore::view(const Task& task) const { + auto qos = task.qos; + auto deadline = task.input.hint.deadline_ns; + auto now = nowNs(); + if (config_.deadline_aware && deadline && + (*deadline <= now || + *deadline - now <= config_.deadline_promotion_window_ns)) { + if (qos.traffic_class == TrafficClass::LOW) + qos.traffic_class = TrafficClass::MEDIUM; + else if (qos.traffic_class == TrafficClass::MEDIUM) + qos.traffic_class = TrafficClass::HIGH; + } + return {task.key, + &task.input.hint, + qos, + task.input.request.length - task.completed, + task.enqueue_ns, + task.sequence}; +} + +Status SchedulerCore::submit( + BatchID id, const std::vector& requests) { + std::lock_guard lock(mutex_); + if (stopping_) return Status::BatchBusy("Scheduler is stopping"); + auto& public_batch = Transport::toBatchDesc(id); + auto existing = batches_.find(id); + size_t count = + existing == batches_.end() ? 0 : existing->second.tasks.size(); + if (public_batch.task_list.size() != count) { + return Status::InvalidArgument( + "Cannot mix scheduled and unscheduled tasks in a batch"); + } + if (count > public_batch.batch_size || + requests.size() > public_batch.batch_size - count) { + return Status::TooManyRequests("Scheduler task capacity exhausted"); + } + uint64_t bytes = 0; + auto simulated = snapshot(""); + auto tenants = tenant_usage_; + size_t deferred_tasks = deferred_tasks_; + uint64_t deferred_bytes = deferred_bytes_; + std::vector> prepared; + for (const auto& entry : requests) { + const auto& request = entry.request; + if (next_key_ == std::numeric_limits::max() || + next_sequence_ == std::numeric_limits::max()) + return Status::TooManyRequests( + "Scheduler identity space exhausted"); + if (request.length > config_.max_outstanding_bytes || + request.length > config_.max_tenant_outstanding_bytes) { + return Status::InvalidArgument("Task exceeds scheduler byte limit"); + } + if (request.length && + (!request.source || + request.length - 1 > + std::numeric_limits::max() - + reinterpret_cast(request.source) || + request.length - 1 > std::numeric_limits::max() - + request.target_offset)) { + return Status::InvalidArgument("Invalid transfer address range"); + } + auto task = std::make_unique(); + task->input = entry; + if (task->input.hint.tenant_id.empty()) + task->input.hint.tenant_id = "default"; + task->key = next_key_++; + if (task->input.hint.request_id.empty()) + task->input.hint.request_id = std::to_string(task->key); + task->qos = qos_->resolve(task->input.hint, config_); + task->enqueue_ns = nowNs(); + task->sequence = next_sequence_++; + task->public_id = count + prepared.size(); + auto result = route_(request, task->transport); + if (!result.ok()) return result; + if (!task->transport) + return Status::InvalidArgument( + "Scheduler route returned no transport"); + auto& tenant = tenants[task->input.hint.tenant_id]; + simulated.tenant_outstanding_tasks = tenant.first; + simulated.tenant_outstanding_bytes = tenant.second; + auto decision = admission_->admit(request.length, simulated, config_); + if (decision == AdmissionDecision::REJECT) + return Status::InvalidArgument("Scheduler admission rejected task"); + task->admitted = decision == AdmissionDecision::ACCEPT; + if (task->admitted) { + ++simulated.outstanding_tasks; + simulated.outstanding_bytes += request.length; + ++tenant.first; + tenant.second += request.length; + } else { + if (deferred_tasks >= config_.max_deferred_tasks || + request.length > config_.max_deferred_bytes - deferred_bytes) + return Status::TooManyRequests( + "Scheduler deferred queue is full"); + ++deferred_tasks; + deferred_bytes += request.length; + } + if (request.length > std::numeric_limits::max() - bytes) + return Status::InvalidArgument("Batch byte count overflows"); + bytes += request.length; + prepared.push_back(std::move(task)); + } + if (prepared.empty()) return Status::OK(); + auto& batch = batches_[id]; + public_batch.task_list.resize(count + prepared.size()); + public_batch.is_finished.store(false, std::memory_order_release); + for (auto& task : prepared) { + public_batch.task_list[task->public_id].batch_id = id; + if (task->admitted) { + admit(*task); + } else { + ++deferred_tasks_; + deferred_bytes_ += task->input.request.length; + } + batch.tasks.push_back(std::move(task)); + } + outstanding_tasks_ += requests.size(); + outstanding_bytes_ += bytes; + wake_.notify_one(); + return Status::OK(); +} + +void SchedulerCore::finish(Task& task, Transport::TransferStatusEnum state) { + task.state = state; + --outstanding_tasks_; + outstanding_bytes_ -= task.input.request.length; + if (task.admitted) { + --admitted_tasks_; + admitted_bytes_ -= task.input.request.length; + auto tenant = tenant_usage_.find(task.input.hint.tenant_id); + --tenant->second.first; + tenant->second.second -= task.input.request.length; + if (!tenant->second.first) tenant_usage_.erase(tenant); + } else { + --deferred_tasks_; + deferred_bytes_ -= task.input.request.length; + } +} + +RuntimeSnapshot SchedulerCore::snapshot(const std::string& tenant) const { + RuntimeSnapshot result; + result.outstanding_tasks = admitted_tasks_; + result.outstanding_bytes = admitted_bytes_; + result.inflight_bytes = inflight_bytes_; + auto it = tenant_usage_.find(tenant); + if (it != tenant_usage_.end()) { + result.tenant_outstanding_tasks = it->second.first; + result.tenant_outstanding_bytes = it->second.second; + } + return result; +} + +void SchedulerCore::admit(Task& task) { + task.admitted = true; + ++admitted_tasks_; + admitted_bytes_ += task.input.request.length; + auto& tenant = tenant_usage_[task.input.hint.tenant_id]; + ++tenant.first; + tenant.second += task.input.request.length; +} + +void SchedulerCore::publishBatch(BatchID id, const Batch& batch) { + auto& public_batch = Transport::toBatchDesc(id); + bool done = true; + bool failed = false; + uint64_t bytes = 0; + for (const auto& task : batch.tasks) { + auto& public_task = public_batch.task_list[task->public_id]; + public_task.transferred_bytes = task->completed; + public_task.is_finished = terminal(*task); + bytes += task->completed; + done = done && terminal(*task); + failed = failed || task->state == Transport::FAILED || + task->state == Transport::CANCELED; + } + public_batch.finished_transfer_bytes.store(bytes, + std::memory_order_relaxed); + if (done) { + public_batch.has_failure.store(failed, std::memory_order_relaxed); +#ifdef USE_EVENT_DRIVEN_COMPLETION + std::lock_guard lock(public_batch.completion_mutex); +#endif + public_batch.is_finished.store(true, std::memory_order_release); +#ifdef USE_EVENT_DRIVEN_COMPLETION + public_batch.completion_cv.notify_all(); +#endif + } +} + +void SchedulerCore::poll() { + for (auto& batch : batches_) { + for (auto& ptr : batch.second.tasks) { + auto& task = *ptr; + if (terminal(task)) continue; + if (task.grant) { + auto& physical = task.grant->task_list[0]; + // Failure/timeout is not proof that DMA has drained. Keep the + // grant and its credits until every submitted slice is done. + auto done = __atomic_load_n( + &physical.scheduled_completed_slices, __ATOMIC_ACQUIRE); + auto slice_count = + __atomic_load_n(&physical.slice_count, __ATOMIC_ACQUIRE); + if (done != slice_count) continue; + auto failed_slices = __atomic_load_n( + &physical.failed_slice_count, __ATOMIC_ACQUIRE); + auto transferred = __atomic_load_n(&physical.transferred_bytes, + __ATOMIC_ACQUIRE); + bool failed = task.dispatch_failed || failed_slices || + transferred != task.range.length; + inflight_bytes_ -= task.range.length; + task.completed += transferred; + task.grant.reset(); + if (failed) { + finish(task, Transport::FAILED); + continue; + } + task.state = Transport::PENDING; + task.enqueue_ns = nowNs(); + task.sequence = next_sequence_++; + } + if (task.canceling) + finish(task, Transport::CANCELED); + else if (task.completed == task.input.request.length) + finish(task, Transport::COMPLETED); + else if (!task.admitted && + admission_->admit(task.input.request.length, + snapshot(task.input.hint.tenant_id), + config_) == AdmissionDecision::ACCEPT) { + --deferred_tasks_; + deferred_bytes_ -= task.input.request.length; + admit(task); + } + } + publishBatch(batch.first, batch.second); + } +} + +bool SchedulerCore::dispatch() { + std::vector ready; + std::unordered_map tasks; + for (auto& batch : batches_) { + for (auto& ptr : batch.second.tasks) { + auto& task = *ptr; + if (task.state != Transport::PENDING || task.canceling || + !task.admitted) + continue; + if (!budget_->budget(view(task), config_.quantum_bytes, + snapshot(task.input.hint.tenant_id), config_)) + continue; + ready.push_back(view(task)); + tasks[task.key] = &task; + } + } + auto selected = selection_->select(ready, nowNs()); + if (!selected.key) return false; + auto found = tasks.find(selected.key); + if (found == tasks.end() || !selected.available_bytes) return false; + auto& task = *found->second; + uint64_t slice_size = globalConfig().slice_size; + if (!slice_size) { + finish(task, Transport::FAILED); + return true; + } + uint64_t descriptor_bytes = + slice_size > std::numeric_limits::max() / config_.max_slices + ? std::numeric_limits::max() + : slice_size * config_.max_slices; + uint64_t bytes = + std::min(budget_->budget(view(task), selected.available_bytes, + snapshot(task.input.hint.tenant_id), config_), + descriptor_bytes); + uint64_t hard_limit = config_.max_inflight_bytes; + if (task.qos.traffic_class != TrafficClass::HIGH) + hard_limit -= config_.reserved_high_bytes; + if (inflight_bytes_ >= hard_limit) return false; + bytes = std::min({bytes, selected.available_bytes, + uint64_t(task.input.request.length - task.completed), + hard_limit - inflight_bytes_}); + if (!bytes) return false; + task.range = task.input.request; + task.range.source = reinterpret_cast( + reinterpret_cast(task.range.source) + task.completed); + task.range.target_offset += task.completed; + task.range.length = bytes; + size_t accepted_length = 0; + auto bounded = task.transport->scheduledTransferLength( + task.range, config_.max_slices, accepted_length); + if (!bounded.ok() || !accepted_length || accepted_length > bytes) { + finish(task, Transport::FAILED); + return true; + } + bytes = accepted_length; + task.range.length = bytes; + task.grant = std::make_unique(); + task.grant->id = reinterpret_cast(task.grant.get()); + task.grant->batch_size = 1; + task.grant->context = nullptr; + task.grant->task_list.resize(1); + auto& physical = task.grant->task_list[0]; + physical.batch_id = task.grant->id; + physical.transport_ = task.transport; + physical.request = &task.range; + physical.scheduled = true; + inflight_bytes_ += bytes; + task.state = Transport::WAITING; + auto result = task.transport->submitTransferTask({&physical}); + task.dispatch_failed = !result.ok(); + // A legacy transport may return an error after posting a prefix. Do not + // replay that range or release its memory/credits before the prefix drains. + if (result.ok() || + __atomic_load_n(&physical.slice_count, __ATOMIC_ACQUIRE) != 0) + selection_->accepted(view(task), bytes); + return true; +} + +void SchedulerCore::run() { + std::unique_lock lock(mutex_); + for (;;) { + poll(); + if (stopping_ && !outstanding_tasks_) break; + if (!stopping_) { + while (dispatch()) { + } + } + // Completion counters are grant-level aggregation; bounded polling + // also makes progress when callers never query their batch status. + wake_.wait_for(lock, std::chrono::microseconds(100)); + } +} + +bool SchedulerCore::owns(BatchID id) { + std::lock_guard lock(mutex_); + return batches_.count(id) != 0; +} + +Status SchedulerCore::status(BatchID id, size_t index, TransferStatus& result) { + std::lock_guard lock(mutex_); + auto it = batches_.find(id); + if (it == batches_.end() || index >= it->second.tasks.size()) + return Status::InvalidArgument("Scheduled task ID out of range"); + const auto& task = *it->second.tasks[index]; + result = {task.state, size_t(task.completed)}; + return Status::OK(); +} + +Status SchedulerCore::batchStatus(BatchID id, TransferStatus& result) { + std::lock_guard lock(mutex_); + auto it = batches_.find(id); + if (it == batches_.end()) + return Status::InvalidArgument("Unknown scheduled batch"); + result = {Transport::COMPLETED, 0}; + bool pending = false, failed = false, canceled = false; + for (const auto& task : it->second.tasks) { + result.transferred_bytes += task->completed; + pending |= !terminal(*task); + failed |= task->state == Transport::FAILED; + canceled |= task->state == Transport::CANCELED; + } + result.s = pending ? Transport::WAITING + : failed ? Transport::FAILED + : canceled ? Transport::CANCELED + : Transport::COMPLETED; + return Status::OK(); +} + +Status SchedulerCore::cancel(BatchID id, size_t index) { + std::lock_guard lock(mutex_); + auto it = batches_.find(id); + if (it == batches_.end() || index >= it->second.tasks.size()) + return Status::InvalidArgument("Scheduled task ID out of range"); + it->second.tasks[index]->canceling = true; + wake_.notify_one(); + return Status::OK(); +} + +Status SchedulerCore::release(BatchID id) { + std::lock_guard lock(mutex_); + auto it = batches_.find(id); + if (it == batches_.end()) return Status::OK(); + for (const auto& task : it->second.tasks) + if (!terminal(*task)) + return Status::BatchBusy("Scheduled batch is busy"); + publishBatch(id, it->second); + batches_.erase(it); + return Status::OK(); +} + +} // namespace mooncake::scheduling From 9d7d6132fe2f2418c91cdc99b2a4fed6cae87438 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Tue, 8 Sep 2026 17:08:09 +0800 Subject: [PATCH 3/9] =?UTF-8?q?feat(scheduler):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E4=BC=A0=E8=BE=93=E4=BB=BB=E5=8A=A1=E8=B0=83=E5=BA=A6=E5=99=A8?= =?UTF-8?q?=EF=BC=8C=E6=94=AF=E6=8C=81=E4=BC=98=E5=85=88=E7=BA=A7=E5=92=8C?= =?UTF-8?q?QoS=E6=84=9F=E7=9F=A5=E7=9A=84=E8=B0=83=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 引入 SchedulerCore 核心调度器及可插拔策略（QoS 解析、准入、选路、预算）， 为 TransferEngine 提供 configureScheduling、submitScheduledTransfer 和 cancelTransfer 接口。 各 Transport 适配 scheduledTransferLength 以支持分片调度， 并提供基准测试验证调度器对前台时延和总吞吐的影响。 --- .../benchmark/scheduler/README.md | 58 ++++ .../scheduler/scheduler_benchmark.cpp | 312 ++++++++++++++++++ .../include/CMakeLists.txt | 5 + .../include/multi_transport.h | 9 + .../include/scheduler/scheduler_core.h | 102 ++++++ .../include/scheduler/scheduler_policy.h | 117 +++++++ .../include/scheduler/scheduling_hint.h | 35 ++ .../include/transfer_engine.h | 5 + .../include/transfer_engine_impl.h | 16 + .../kunpeng_transport/ub_transport.h | 4 + .../transport/rdma_transport/rdma_transport.h | 4 + .../transport/tcp_transport/tcp_transport.h | 8 + .../include/transport/transport.h | 17 + 13 files changed, 692 insertions(+) create mode 100644 mooncake-transfer-engine/benchmark/scheduler/README.md create mode 100644 mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp create mode 100644 mooncake-transfer-engine/include/scheduler/scheduler_core.h create mode 100644 mooncake-transfer-engine/include/scheduler/scheduler_policy.h create mode 100644 mooncake-transfer-engine/include/scheduler/scheduling_hint.h diff --git a/mooncake-transfer-engine/benchmark/scheduler/README.md b/mooncake-transfer-engine/benchmark/scheduler/README.md new file mode 100644 index 0000000000..5750c3920e --- /dev/null +++ b/mooncake-transfer-engine/benchmark/scheduler/README.md @@ -0,0 +1,58 @@ +# Classic TE Scheduler Benchmark + +This standalone benchmark compares the original classic Transfer Engine +submission path with the task scheduler. It does not change or link against +the existing `tebench` implementation. + +The workload follows the scheduler evaluation design: + +- 4 KiB and 64 KiB foreground reads with `FOREGROUND_GET` intent; +- 8 MiB migration and 64 MiB checkpoint background writes; +- 30 seconds of warmup, 120 seconds of measurement, and five repetitions; +- median foreground P99 and aggregate throughput comparison. + +## Build + +Configure this directory as an independent CMake project: + +```bash +cmake -S mooncake-transfer-engine/benchmark/scheduler \ + -B build/scheduler-benchmark +cmake --build build/scheduler-benchmark -j +``` + +The executable and comparison script are produced in +`build/scheduler-benchmark/`. + +## Run + +Start the existing classic `tebench` target with at least a 1 GiB buffer: + +```bash +./tebench --backend=classic --seg_type=DRAM --total_buffer_size=1073741824 +``` + +Copy the printed segment name. Run the baseline and scheduled cases from the +standalone benchmark build directory: + +```bash +rm -f baseline.jsonl scheduled.jsonl + +./scheduler_benchmark --target_seg_name= \ + --scheduling=false --output_jsonl=baseline.jsonl + +./scheduler_benchmark --target_seg_name= \ + --scheduling=true --output_jsonl=scheduled.jsonl + +python3 compare_results.py baseline.jsonl scheduled.jsonl +``` + +The comparison passes when the median P99 of both foreground classes improves +by at least 20% and median aggregate throughput retains at least 95% of the +baseline. Override the gates with `--min-p99-reduction` and +`--min-throughput-retention`. + +Use the same target, transport configuration, CPU affinity, registered memory, +and connection setup for both runs. The scheduler parameters can be scanned +with `--scheduler_quantum_bytes`, `--scheduler_max_inflight_bytes`, +`--scheduler_reserved_high_bytes`, and `--scheduler_max_slices`. diff --git a/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp new file mode 100644 index 0000000000..4f1f0cebed --- /dev/null +++ b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp @@ -0,0 +1,312 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "scheduler/scheduler_core.h" +#include "scheduler/scheduler_policy.h" +#include "transfer_engine.h" + +DEFINE_string(target_seg_name, "", "Target segment printed by tebench"); +DEFINE_string(metadata_conn_string, "P2PHANDSHAKE", + "Metadata connection string"); +DEFINE_string(local_server_name, mooncake::getHostname(), + "Local server name used by the initiator"); +DEFINE_bool(scheduling, false, "Submit through the classic TE scheduler"); +DEFINE_uint64(buffer_size, 1ULL << 30, + "Registered initiator buffer size in bytes"); +DEFINE_int32(foreground_threads, 1, + "Worker count for each foreground request size"); +DEFINE_int32(background_threads, 3, + "Worker count for each background request size"); +DEFINE_int32(warmup_seconds, 30, "Warmup duration for each repetition"); +DEFINE_int32(duration_seconds, 120, "Measurement duration per repetition"); +DEFINE_int32(repetitions, 5, "Number of independent repetitions"); +DEFINE_uint64(foreground_deadline_us, 0, + "Optional relative deadline for foreground requests"); +DEFINE_string(output_jsonl, "scheduler-benchmark.jsonl", + "File to which one JSON record per repetition is appended"); +DEFINE_uint64(scheduler_quantum_bytes, 1ULL << 20, "Scheduler byte quantum"); +DEFINE_uint64(scheduler_max_inflight_bytes, 16ULL << 20, + "Scheduler maximum in-flight bytes"); +DEFINE_uint64(scheduler_reserved_high_bytes, 1ULL << 20, + "Scheduler bytes reserved for HIGH traffic"); +DEFINE_uint32(scheduler_max_slices, 32, + "Scheduler maximum transport slices per grant"); + +namespace { + +using Clock = std::chrono::steady_clock; + +struct TrafficClass { + const char* name; + size_t block_size; + int threads; + mooncake::TransferRequest::OpCode opcode; + mooncake::TaskIntent intent; + const char* tenant; +}; + +struct WorkerResult { + std::vector latency_us; + uint64_t bytes{0}; + double duration_seconds{0}; +}; + +void check(const mooncake::Status& status, const char* operation) { + LOG_ASSERT(status.ok()) << operation << " failed: " << status.ToString(); +} + +double percentile(std::vector samples, double value) { + if (samples.empty()) return 0.0; + std::sort(samples.begin(), samples.end()); + const double rank = value / 100.0 * (samples.size() - 1); + const size_t lower = static_cast(rank); + const size_t upper = std::min(lower + 1, samples.size() - 1); + const double fraction = rank - lower; + return samples[lower] * (1.0 - fraction) + samples[upper] * fraction; +} + +double runTransfer(mooncake::TransferEngine& engine, mooncake::SegmentID target, + void* local, uint64_t remote_offset, + const TrafficClass& traffic) { + const auto batch = engine.allocateBatchID(1); + mooncake::TransferRequest request; + request.opcode = traffic.opcode; + request.source = local; + request.target_id = target; + request.target_offset = remote_offset; + request.length = traffic.block_size; + + const auto started = Clock::now(); + if (FLAGS_scheduling) { + mooncake::SchedulingHint hint; + hint.tenant_id = traffic.tenant; + hint.intent = traffic.intent; + if (FLAGS_foreground_deadline_us != 0 && + traffic.intent == mooncake::TaskIntent::FOREGROUND_GET) { + hint.deadline_ns = + std::chrono::duration_cast( + Clock::now().time_since_epoch()) + .count() + + FLAGS_foreground_deadline_us * 1000ULL; + } + check(engine.submitScheduledTransfer(batch, {{request, hint}}), + "submitScheduledTransfer"); + } else { + check(engine.submitTransfer(batch, {request}), "submitTransfer"); + } + + while (true) { + mooncake::TransferStatus status; + check(engine.getTransferStatus(batch, 0, status), "getTransferStatus"); + if (status.s == mooncake::TransferStatusEnum::COMPLETED) break; + LOG_ASSERT(status.s != mooncake::TransferStatusEnum::FAILED && + status.s != mooncake::TransferStatusEnum::TIMEOUT && + status.s != mooncake::TransferStatusEnum::CANCELED) + << "transfer reached terminal status " << status.s; + std::this_thread::yield(); + } + const double elapsed_us = + std::chrono::duration(Clock::now() - started) + .count(); + check(engine.freeBatchID(batch), "freeBatchID"); + return elapsed_us; +} + +void runWorker(mooncake::TransferEngine& engine, mooncake::SegmentID target, + uint8_t* local_base, uint64_t remote_base, size_t worker_index, + size_t address_stride, const TrafficClass& traffic, + WorkerResult& result) { + void* local = local_base + worker_index * address_stride; + const uint64_t remote = remote_base + worker_index * address_stride; + auto until = Clock::now() + std::chrono::seconds(FLAGS_warmup_seconds); + while (Clock::now() < until) + runTransfer(engine, target, local, remote, traffic); + + const auto started = Clock::now(); + until = started + std::chrono::seconds(FLAGS_duration_seconds); + while (Clock::now() < until) { + result.latency_us.push_back( + runTransfer(engine, target, local, remote, traffic)); + result.bytes += traffic.block_size; + } + result.duration_seconds = + std::chrono::duration(Clock::now() - started).count(); +} + +void appendRecord(const std::vector& classes, + const std::vector>& results, + int repetition) { + std::ofstream output(FLAGS_output_jsonl, std::ios::app); + LOG_ASSERT(output) << "cannot open " << FLAGS_output_jsonl; + + double aggregate_throughput = 0.0; + struct ClassMetrics { + double p99_us; + double throughput_gbps; + uint64_t operations; + }; + std::vector metrics; + for (size_t class_index = 0; class_index < classes.size(); ++class_index) { + std::vector latencies; + uint64_t bytes = 0; + uint64_t operations = 0; + double throughput = 0.0; + for (const auto& worker : results[class_index]) { + latencies.insert(latencies.end(), worker.latency_us.begin(), + worker.latency_us.end()); + bytes += worker.bytes; + operations += worker.latency_us.size(); + if (worker.duration_seconds > 0) + throughput += worker.bytes / 1e9 / worker.duration_seconds; + } + aggregate_throughput += throughput; + metrics.push_back( + {percentile(std::move(latencies), 99.0), throughput, operations}); + } + + output << std::fixed << std::setprecision(6) + << "{\"schema_version\":1,\"scheduling\":" + << (FLAGS_scheduling ? "true" : "false") + << ",\"repetition\":" << repetition + << ",\"aggregate_throughput_gbps\":" << aggregate_throughput + << ",\"classes\":["; + for (size_t i = 0; i < classes.size(); ++i) { + if (i != 0) output << ','; + output << "{\"name\":\"" << classes[i].name + << "\",\"threads\":" << classes[i].threads + << ",\"block_size\":" << classes[i].block_size + << ",\"operations\":" << metrics[i].operations + << ",\"p99_us\":" << metrics[i].p99_us + << ",\"throughput_gbps\":" << metrics[i].throughput_gbps << '}'; + } + output << "]}\n"; + LOG_ASSERT(output) << "failed to write " << FLAGS_output_jsonl; + + std::cout << "repetition=" << repetition << " scheduling=" << std::boolalpha + << FLAGS_scheduling + << " aggregate_throughput=" << aggregate_throughput << " GB/s" + << std::endl; + for (size_t i = 0; i < classes.size(); ++i) + std::cout << " " << classes[i].name << ": p99=" << metrics[i].p99_us + << " us throughput=" << metrics[i].throughput_gbps << " GB/s" + << std::endl; +} + +} // namespace + +int main(int argc, char** argv) { + gflags::SetUsageMessage( + "Classic Transfer Engine scheduler A/B benchmark initiator"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + + LOG_ASSERT(!FLAGS_target_seg_name.empty()) + << "--target_seg_name is required"; + LOG_ASSERT(FLAGS_foreground_threads > 0 && FLAGS_background_threads > 0); + LOG_ASSERT(FLAGS_foreground_threads <= 128 && + FLAGS_background_threads <= 128) + << "thread counts must not exceed 128 per traffic class"; + LOG_ASSERT(FLAGS_warmup_seconds >= 0 && FLAGS_duration_seconds > 0 && + FLAGS_repetitions > 0); + + std::vector classes = { + {"foreground-4k", 4ULL << 10, FLAGS_foreground_threads, + mooncake::TransferRequest::READ, mooncake::TaskIntent::FOREGROUND_GET, + "online"}, + {"foreground-64k", 64ULL << 10, FLAGS_foreground_threads, + mooncake::TransferRequest::READ, mooncake::TaskIntent::FOREGROUND_GET, + "online"}, + {"migration-8m", 8ULL << 20, FLAGS_background_threads, + mooncake::TransferRequest::WRITE, mooncake::TaskIntent::MIGRATION, + "batch"}, + {"checkpoint-64m", 64ULL << 20, FLAGS_background_threads, + mooncake::TransferRequest::WRITE, mooncake::TaskIntent::CHECKPOINT, + "batch"}, + }; + + size_t worker_count = 0; + for (const auto& traffic : classes) worker_count += traffic.threads; + const size_t address_stride = 64ULL << 20; + const size_t required_buffer = worker_count * address_stride; + LOG_ASSERT(FLAGS_buffer_size <= std::numeric_limits::max() && + FLAGS_buffer_size % 4096 == 0) + << "--buffer_size must fit in size_t and be a multiple of 4096"; + LOG_ASSERT(FLAGS_buffer_size >= required_buffer) + << "--buffer_size must be at least " << required_buffer; + + auto* local_buffer = static_cast( + std::aligned_alloc(4096, static_cast(FLAGS_buffer_size))); + LOG_ASSERT(local_buffer != nullptr) << "failed to allocate local buffer"; + + mooncake::TransferEngine engine(true); + LOG_ASSERT( + engine.init(FLAGS_metadata_conn_string, FLAGS_local_server_name) == 0) + << "TransferEngine initialization failed"; + LOG_ASSERT(engine.registerLocalMemory(local_buffer, FLAGS_buffer_size) == 0) + << "local memory registration failed"; + + if (FLAGS_scheduling) { + mooncake::scheduling::SchedulerConfig config; + config.quantum_bytes = FLAGS_scheduler_quantum_bytes; + config.max_inflight_bytes = FLAGS_scheduler_max_inflight_bytes; + config.reserved_high_bytes = FLAGS_scheduler_reserved_high_bytes; + config.max_slices = FLAGS_scheduler_max_slices; + check(engine.configureScheduling(config), "configureScheduling"); + } + + const auto target = engine.openSegment(FLAGS_target_seg_name); + auto segment = engine.getMetadata()->getSegmentDescByID(target); + LOG_ASSERT(segment && !segment->buffers.empty()) + << "target segment has no registered buffers"; + const auto& remote_buffer = segment->buffers.front(); + LOG_ASSERT(remote_buffer.length >= required_buffer) + << "target buffer must be at least " << required_buffer << " bytes"; + + for (int repetition = 1; repetition <= FLAGS_repetitions; ++repetition) { + std::vector> results; + results.reserve(classes.size()); + for (const auto& traffic : classes) + results.emplace_back(traffic.threads); + + std::vector workers; + size_t worker_index = 0; + for (size_t class_index = 0; class_index < classes.size(); + ++class_index) { + for (int class_worker = 0; + class_worker < classes[class_index].threads; + ++class_worker, ++worker_index) { + workers.emplace_back( + runWorker, std::ref(engine), target, local_buffer, + remote_buffer.addr, worker_index, address_stride, + std::cref(classes[class_index]), + std::ref(results[class_index][class_worker])); + } + } + for (auto& worker : workers) worker.join(); + appendRecord(classes, results, repetition); + } + + engine.closeSegment(target); + engine.unregisterLocalMemory(local_buffer); + engine.freeEngine(); + std::free(local_buffer); + return 0; +} diff --git a/mooncake-transfer-engine/include/CMakeLists.txt b/mooncake-transfer-engine/include/CMakeLists.txt index 7d5444d452..0a089507f6 100644 --- a/mooncake-transfer-engine/include/CMakeLists.txt +++ b/mooncake-transfer-engine/include/CMakeLists.txt @@ -10,6 +10,11 @@ install(FILES transfer_metadata.h DESTINATION include) install(FILES ub_allocator.h DESTINATION include) install(FILES common/base/status.h DESTINATION include/common/base) install(FILES transport/transport.h DESTINATION include/transport) +install(FILES + scheduler/scheduling_hint.h + scheduler/scheduler_policy.h + scheduler/scheduler_core.h + DESTINATION include/scheduler) # Device API headers (header-only, consumed by EP kernel) install(FILES transport/device/device_transport.h DESTINATION include/transport/device) diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index 541a7391ea..7f81e60dac 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -18,6 +18,7 @@ #include #include "transport/transport.h" +#include "scheduler/scheduler_core.h" namespace mooncake { class TransferEngineImplTestPeer; @@ -43,6 +44,12 @@ class MultiTransport { Status submitTransfer(BatchID batch_id, const std::vector &entries); + // Configure before submitting any batches. Scheduling is opt-in. + Status configureScheduling(const scheduling::SchedulerConfig &config); + Status submitScheduledTransfer( + BatchID batch_id, const std::vector &entries); + Status cancelTransfer(BatchID batch_id, size_t task_id); + #ifdef ENABLE_MULTI_PROTOCOL Status mp_submitTransfer(BatchID batch_id, const std::vector &entries, @@ -86,6 +93,8 @@ class MultiTransport { std::map> transport_map_; RWSpinlock batch_desc_lock_; std::unordered_map> batch_desc_set_; + // Destroy first: drain scheduled work while transports still exist. + std::unique_ptr scheduler_; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/scheduler/scheduler_core.h b/mooncake-transfer-engine/include/scheduler/scheduler_core.h new file mode 100644 index 0000000000..b4e9489073 --- /dev/null +++ b/mooncake-transfer-engine/include/scheduler/scheduler_core.h @@ -0,0 +1,102 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#pragma once + +#include +#include +#include +#include +#include + +#include "scheduler/scheduler_policy.h" +#include "transport/transport.h" + +namespace mooncake { + +struct ScheduledTransferRequest { + Transport::TransferRequest request; + SchedulingHint hint; +}; + +namespace scheduling { + +// Each grant has a private physical batch. Public batch/task identities never +// enter a transport, so physical completion cannot finish a logical task early. +class SchedulerCore { + public: + using BatchID = Transport::BatchID; + using TransferStatus = Transport::TransferStatus; + using Route = + std::function; + + SchedulerCore(SchedulerConfig config, Route route); + SchedulerCore(SchedulerConfig config, Route route, + SchedulerPolicySet policies); + ~SchedulerCore(); + static Status validate(const SchedulerConfig& config); + + Status submit(BatchID batch, + const std::vector& requests); + Status status(BatchID batch, size_t task, TransferStatus& result); + Status batchStatus(BatchID batch, TransferStatus& result); + Status cancel(BatchID batch, size_t task); + Status release(BatchID batch); + bool owns(BatchID batch); + + private: + struct Task { + uint64_t key; + size_t public_id; + ScheduledTransferRequest input; + ResolvedQoS qos; + uint64_t enqueue_ns; + uint64_t sequence; + uint64_t completed{0}; + Transport::TransferStatusEnum state{Transport::PENDING}; + bool canceling{false}; + bool dispatch_failed{false}; + bool admitted{false}; + Transport* transport{nullptr}; + Transport::TransferRequest range; + std::unique_ptr grant; + }; + struct Batch { + std::vector> tasks; + }; + + static uint64_t nowNs(); + static bool terminal(const Task& task); + ReadyTaskView view(const Task& task) const; + void run(); + void poll(); + bool dispatch(); + void finish(Task& task, Transport::TransferStatusEnum state); + void publishBatch(BatchID id, const Batch& batch); + RuntimeSnapshot snapshot(const std::string& tenant) const; + void admit(Task& task); + + SchedulerConfig config_; + Route route_; + std::unique_ptr qos_; + std::unique_ptr selection_; + std::unique_ptr admission_; + std::unique_ptr budget_; + std::mutex mutex_; + std::condition_variable wake_; + std::unordered_map batches_; + uint64_t next_key_{1}; + uint64_t next_sequence_{1}; + size_t outstanding_tasks_{0}; + uint64_t outstanding_bytes_{0}; + size_t admitted_tasks_{0}; + uint64_t admitted_bytes_{0}; + size_t deferred_tasks_{0}; + uint64_t deferred_bytes_{0}; + std::map> tenant_usage_; + uint64_t inflight_bytes_{0}; + bool stopping_{false}; + std::thread worker_; +}; + +} // namespace scheduling +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/scheduler/scheduler_policy.h b/mooncake-transfer-engine/include/scheduler/scheduler_policy.h new file mode 100644 index 0000000000..a38d5b6383 --- /dev/null +++ b/mooncake-transfer-engine/include/scheduler/scheduler_policy.h @@ -0,0 +1,117 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "scheduler/scheduling_hint.h" + +namespace mooncake::scheduling { + +enum class TrafficClass : uint8_t { HIGH, MEDIUM, LOW }; + +struct SchedulerConfig { + size_t max_outstanding_tasks{1024}; + uint64_t max_outstanding_bytes{1ULL << 30}; + size_t max_deferred_tasks{256}; + uint64_t max_deferred_bytes{256ULL << 20}; + size_t max_tenant_outstanding_tasks{1024}; + uint64_t max_tenant_outstanding_bytes{1ULL << 30}; + uint64_t max_inflight_bytes{16ULL << 20}; + uint64_t reserved_high_bytes{1ULL << 20}; + uint64_t quantum_bytes{1ULL << 20}; + uint32_t max_slices{32}; + uint64_t aging_interval_ns{10000000}; + uint32_t max_aging_bonus{1024}; + bool deadline_aware{true}; + uint64_t deadline_promotion_window_ns{1000000}; + std::array class_weights{8, 4, 1}; + std::map tenant_weights; +}; + +struct ResolvedQoS { + TrafficClass traffic_class{TrafficClass::MEDIUM}; + int32_t priority_rank{0}; + uint32_t weight{1}; +}; + +struct ReadyTaskView { + uint64_t key; + const SchedulingHint* requested; + ResolvedQoS qos; + uint64_t remaining_bytes; + uint64_t enqueue_ns; + uint64_t sequence; +}; + +struct TaskSelection { + uint64_t key{0}; + uint64_t available_bytes{0}; +}; + +enum class AdmissionDecision { ACCEPT, DEFER, REJECT }; + +struct RuntimeSnapshot { + size_t outstanding_tasks{0}; + uint64_t outstanding_bytes{0}; + size_t tenant_outstanding_tasks{0}; + uint64_t tenant_outstanding_bytes{0}; + uint64_t inflight_bytes{0}; +}; + +class AdmissionPolicy { + public: + virtual ~AdmissionPolicy() = default; + virtual AdmissionDecision admit(uint64_t bytes, + const RuntimeSnapshot& runtime, + const SchedulerConfig& config) const = 0; +}; + +class DispatchBudgetPolicy { + public: + virtual ~DispatchBudgetPolicy() = default; + virtual uint64_t budget(const ReadyTaskView& task, uint64_t fair_bytes, + const RuntimeSnapshot& runtime, + const SchedulerConfig& config) const = 0; +}; + +class QoSResolutionPolicy { + public: + virtual ~QoSResolutionPolicy() = default; + virtual ResolvedQoS resolve(const SchedulingHint& hint, + const SchedulerConfig& config) const = 0; +}; + +// A selection is only charged after transport acceptance. The policy owns +// fairness state, while the scheduler owns task and resource state. +class TaskSelectionPolicy { + public: + virtual ~TaskSelectionPolicy() = default; + virtual TaskSelection select(const std::vector& ready, + uint64_t now_ns) = 0; + virtual void accepted(const ReadyTaskView& task, uint64_t bytes) = 0; +}; + +// Construct a complete policy set before enabling scheduling. No registry +// lookup or configuration parsing occurs in the scheduling loop. +struct SchedulerPolicySet { + std::unique_ptr qos; + std::unique_ptr admission; + std::unique_ptr task_selection; + std::unique_ptr dispatch_budget; +}; + +SchedulerPolicySet makeDefaultPolicySet(const SchedulerConfig& config); + +std::unique_ptr makeIntentQoSPolicy(); +std::unique_ptr makeHierarchicalAdmissionPolicy(); +std::unique_ptr makeBoundedDispatchBudgetPolicy(); +std::unique_ptr makeHierarchicalTaskPolicy( + const SchedulerConfig& config); + +} // namespace mooncake::scheduling diff --git a/mooncake-transfer-engine/include/scheduler/scheduling_hint.h b/mooncake-transfer-engine/include/scheduler/scheduling_hint.h new file mode 100644 index 0000000000..755da0a843 --- /dev/null +++ b/mooncake-transfer-engine/include/scheduler/scheduling_hint.h @@ -0,0 +1,35 @@ +// Copyright 2026 Mooncake Authors +// Licensed under the Apache License, Version 2.0. +#pragma once + +#include +#include +#include + +namespace mooncake { + +enum class TaskIntent : uint8_t { + UNSPEC = 0, + CONTROL, + FOREGROUND_GET, + P2D_TRANSFER, + BACKGROUND_PUT, + PREFETCH, + MIGRATION, + CHECKPOINT, + WEIGHT_LOADING, +}; + +// deadline_ns is an absolute deadline in this TE process's steady clock. +// Smaller numeric priorities are more urgent; zero is an explicit priority. +struct SchedulingHint { + std::string request_id; + uint64_t generation{0}; + std::string tenant_id{"default"}; + std::optional requested_priority; + TaskIntent intent{TaskIntent::UNSPEC}; + std::optional deadline_ns; + bool allow_degrade{false}; +}; + +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index bd822ab3f9..e5936bc08f 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -122,6 +122,11 @@ class TransferEngine { Status submitTransfer(BatchID batch_id, const std::vector& entries); + Status configureScheduling(const scheduling::SchedulerConfig& config); + Status submitScheduledTransfer( + BatchID batch_id, const std::vector& entries); + Status cancelTransfer(BatchID batch_id, size_t task_id); + Status submitTransferWithNotify(BatchID batch_id, const std::vector& entries, TransferMetadata::NotifyDesc notify_msg); diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index c7508c89e5..2ec142a6c1 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -141,6 +141,22 @@ class TransferEngineImpl { return s; } + Status configureScheduling(const scheduling::SchedulerConfig& config) { + if (!multi_transports_) + return Status::InvalidArgument("Initialize TE before scheduling"); + return multi_transports_->configureScheduling(config); + } + + Status submitScheduledTransfer( + BatchID batch_id, + const std::vector& entries) { + return multi_transports_->submitScheduledTransfer(batch_id, entries); + } + + Status cancelTransfer(BatchID batch_id, size_t task_id) { + return multi_transports_->cancelTransfer(batch_id, task_id); + } + Status submitTransferWithNotify(BatchID batch_id, const std::vector& entries, TransferMetadata::NotifyDesc notify_msg) { diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h index 1ba12c0815..bae5571db0 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h @@ -83,6 +83,10 @@ class UbTransport : public Transport { Status submitTransferTask( const std::vector& task_list) override; + Status scheduledTransferLength(const TransferRequest& request, + uint32_t max_slices, + size_t& length) override; + Status getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus& status) override; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h index d5d8b3b5e5..095d87e0e3 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -92,6 +92,10 @@ class RdmaTransport : public Transport { Status submitTransferTask( const std::vector &task_list) override; + Status scheduledTransferLength(const TransferRequest &request, + uint32_t max_slices, + size_t &length) override; + Status getTransferStatus(BatchID batch_id, std::vector &status); diff --git a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h index be096e96df..0b969cfb4e 100644 --- a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h +++ b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h @@ -75,6 +75,14 @@ class TcpTransport : public Transport { Status submitTransferTask( const std::vector &task_list) override; + Status scheduledTransferLength(const TransferRequest &request, + uint32_t max_slices, + size_t &length) override { + if (!max_slices) return Status::InvalidArgument("Empty slice budget"); + length = request.length; + return Status::OK(); + } + Status getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status) override; diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index ed4124feba..1db62c2f9b 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -205,6 +205,9 @@ class Transport { __atomic_fetch_add(&task->success_slice_count, 1, __ATOMIC_RELAXED); check_batch_completion(false); + if (task->scheduled) + __atomic_fetch_add(&task->scheduled_completed_slices, 1, + __ATOMIC_RELEASE); } void markFailed() { @@ -212,6 +215,9 @@ class Transport { __atomic_fetch_add(&task->failed_slice_count, 1, __ATOMIC_RELAXED); check_batch_completion(true); + if (task->scheduled) + __atomic_fetch_add(&task->scheduled_completed_slices, 1, + __ATOMIC_RELEASE); } volatile int64_t ts; @@ -322,6 +328,9 @@ class Transport { }; struct TransferTask { + // Publish only after completion has finished accessing batch state. + bool scheduled = false; + uint64_t scheduled_completed_slices = 0; volatile uint64_t slice_count = 0; volatile uint64_t success_slice_count = 0; volatile uint64_t failed_slice_count = 0; @@ -404,6 +413,14 @@ class Transport { "Transport::submitTransferTask is not implemented"); } + // Bound a proposed byte range by physical descriptor capacity. Only + // transports with worker-published completion may opt into scheduling. + virtual Status scheduledTransferLength(const TransferRequest &request, + uint32_t max_slices, + size_t &length) { + return Status::NotImplemented("Transport has no scheduling adapter"); + } + /// @brief Get the status of a submitted transfer. This function shall not /// be called again after completion. /// @return Return 1 on completed (either success or failure); 0 if still in From 076158f20949d8d9d3a3595480c0a229b399c740 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Tue, 8 Sep 2026 17:08:58 +0800 Subject: [PATCH 4/9] =?UTF-8?q?feat(store):=20=E6=94=AF=E6=8C=81=E9=80=9A?= =?UTF-8?q?=E8=BF=87=E6=93=8D=E4=BD=9C=E9=80=89=E9=A1=B9=E4=BC=A0=E9=80=92?= =?UTF-8?q?=E8=B0=83=E5=BA=A6=E6=8F=90=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TransferSubmitter 的 submit、submit_batch 等方法现在接受 OperationOptions， 其中包含可选的 SchedulingHint。未提供时，根据操作类型自动设置任务意图。 submitTransfer 改为使用 submitScheduledTransfer 传递调度信息。 同时添加了调度器的基准测试和结果比较脚本，用于验证调度效果。 --- mooncake-store/include/operation_options.h | 15 +++ mooncake-store/include/transfer_task.h | 14 ++- mooncake-store/src/transfer_task.cpp | 41 +++++--- .../benchmark/scheduler/CMakeLists.txt | 15 +++ .../benchmark/scheduler/compare_results.py | 96 +++++++++++++++++++ 5 files changed, 163 insertions(+), 18 deletions(-) create mode 100644 mooncake-store/include/operation_options.h create mode 100644 mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt create mode 100644 mooncake-transfer-engine/benchmark/scheduler/compare_results.py diff --git a/mooncake-store/include/operation_options.h b/mooncake-store/include/operation_options.h new file mode 100644 index 0000000000..80bcc51324 --- /dev/null +++ b/mooncake-store/include/operation_options.h @@ -0,0 +1,15 @@ +#pragma once + +#include + +#include "scheduler/scheduling_hint.h" + +namespace mooncake { + +// Options shared by Store operations that may use Transfer Engine. Local +// memcpy and storage paths ignore scheduling because they use separate pools. +struct OperationOptions { + std::optional scheduling; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h index 52f44e412f..90113fbfab 100644 --- a/mooncake-store/include/transfer_task.h +++ b/mooncake-store/include/transfer_task.h @@ -21,6 +21,7 @@ #include "rpc_types.h" #include "storage_backend.h" #include "client_metric.h" +#include "operation_options.h" #ifdef USE_NOF #include "spdk/spdk_wrapper.h" #endif @@ -564,7 +565,8 @@ class TransferSubmitter { std::optional submit(const Replica::Descriptor& replica, std::vector& slices, TransferRequest::OpCode op_code, - void* ptr = nullptr, size_t size = 0); + void* ptr = nullptr, size_t size = 0, + const OperationOptions& options = {}); /** * @brief Submit a range read: read [src_offset, src_offset+size) from @@ -577,7 +579,7 @@ class TransferSubmitter { std::optional submit_batch( const std::vector& replicas, std::vector>& all_slices, - TransferRequest::OpCode op_code); + TransferRequest::OpCode op_code, const OperationOptions& options = {}); std::optional submit_batch_get_offload_object( const std::string& transfer_engine_addr, @@ -664,11 +666,12 @@ class TransferSubmitter { std::optional submitTransferEngineOperation( const AllocatedBuffer::Descriptor& handle, const std::vector& slices, const TransferRequest::OpCode op_code, - uint64_t src_offset = 0); + uint64_t src_offset = 0, const OperationOptions& options = {}); std::optional submitMemoryReadOperation( const AllocatedBuffer::Descriptor& handle, - const std::vector& slices, uint64_t src_offset); + const std::vector& slices, uint64_t src_offset, + const OperationOptions& options = {}); std::optional submitFileReadOperation( const Replica::Descriptor& replica, std::vector& slices, @@ -681,7 +684,8 @@ class TransferSubmitter { TransferRequest::OpCode op); std::optional submitTransfer( - std::vector& requests); + std::vector& requests, + const SchedulingHint& hint = {}); }; } // namespace mooncake diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp index 3e9dff18d9..d1e33f6046 100644 --- a/mooncake-store/src/transfer_task.cpp +++ b/mooncake-store/src/transfer_task.cpp @@ -1001,7 +1001,8 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine, std::optional TransferSubmitter::submit( const Replica::Descriptor& replica, std::vector& slices, - TransferRequest::OpCode op_code, void* ptr, size_t size) { + TransferRequest::OpCode op_code, void* ptr, size_t size, + const OperationOptions& options) { std::optional future; if (replica.is_memory_replica()) { @@ -1013,7 +1014,7 @@ std::optional TransferSubmitter::submit( } if (op_code == TransferRequest::READ) { - future = submitMemoryReadOperation(handle, slices, 0); + future = submitMemoryReadOperation(handle, slices, 0, options); } else { TransferStrategy strategy = selectStrategy(handle, slices); @@ -1022,8 +1023,8 @@ std::optional TransferSubmitter::submit( future = submitMemcpyOperation(handle, slices, op_code); break; case TransferStrategy::TRANSFER_ENGINE: - future = - submitTransferEngineOperation(handle, slices, op_code); + future = submitTransferEngineOperation(handle, slices, + op_code, 0, options); break; default: LOG(ERROR) << "Unknown transfer strategy: " << strategy; @@ -1059,7 +1060,7 @@ std::optional TransferSubmitter::submit( std::optional TransferSubmitter::submit_batch( const std::vector& replicas, std::vector>& all_slices, - TransferRequest::OpCode op_code) { + TransferRequest::OpCode op_code, const OperationOptions& options) { std::optional future; std::vector requests; for (size_t i = 0; i < replicas.size(); ++i) { @@ -1088,7 +1089,12 @@ std::optional TransferSubmitter::submit_batch( offset += slice.size; } } - future = submitTransfer(requests); + auto hint = options.scheduling.value_or(SchedulingHint{}); + if (!options.scheduling) + hint.intent = op_code == TransferRequest::READ + ? TaskIntent::FOREGROUND_GET + : TaskIntent::BACKGROUND_PUT; + future = submitTransfer(requests, hint); // Update metrics on successful submission if (future.has_value()) { for (auto& slices : all_slices) { @@ -1223,7 +1229,7 @@ std::optional TransferSubmitter::submitMemcpyOperation( } std::optional TransferSubmitter::submitTransfer( - std::vector& requests) { + std::vector& requests, const SchedulingHint& hint) { // Allocate batch ID const size_t batch_size = requests.size(); BatchID batch_id = engine_.allocateBatchID(batch_size); @@ -1233,7 +1239,10 @@ std::optional TransferSubmitter::submitTransfer( } // Submit transfer - Status s = engine_.submitTransfer(batch_id, requests); + std::vector scheduled; + scheduled.reserve(requests.size()); + for (const auto& request : requests) scheduled.push_back({request, hint}); + Status s = engine_.submitScheduledTransfer(batch_id, scheduled); if (!s.ok()) { LOG(ERROR) << "Failed to submit all transfers, error code is " << s.code(); @@ -1259,7 +1268,8 @@ std::optional TransferSubmitter::submitTransfer( std::optional TransferSubmitter::submitTransferEngineOperation( const AllocatedBuffer::Descriptor& handle, const std::vector& slices, - const TransferRequest::OpCode op_code, uint64_t src_offset) { + const TransferRequest::OpCode op_code, uint64_t src_offset, + const OperationOptions& options) { if (handle.transport_endpoint_.empty()) { LOG(ERROR) << "Transport endpoint is empty for handle with address " << handle.buffer_address_; @@ -1293,12 +1303,17 @@ std::optional TransferSubmitter::submitTransferEngineOperation( offset += slice.size; requests.emplace_back(request); } - return submitTransfer(requests); + auto hint = options.scheduling.value_or(SchedulingHint{}); + if (!options.scheduling) + hint.intent = op_code == TransferRequest::READ + ? TaskIntent::FOREGROUND_GET + : TaskIntent::BACKGROUND_PUT; + return submitTransfer(requests, hint); } std::optional TransferSubmitter::submitMemoryReadOperation( const AllocatedBuffer::Descriptor& handle, const std::vector& slices, - uint64_t src_offset) { + uint64_t src_offset, const OperationOptions& options) { TransferStrategy strategy = selectStrategy(handle, slices); if (strategy == TransferStrategy::LOCAL_MEMCPY) { @@ -1306,8 +1321,8 @@ std::optional TransferSubmitter::submitMemoryReadOperation( src_offset); } if (strategy == TransferStrategy::TRANSFER_ENGINE) { - return submitTransferEngineOperation(handle, slices, - TransferRequest::READ, src_offset); + return submitTransferEngineOperation( + handle, slices, TransferRequest::READ, src_offset, options); } LOG(ERROR) << "Read only supports LOCAL_MEMCPY or TRANSFER_ENGINE, got: " diff --git a/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt b/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt new file mode 100644 index 0000000000..e38d8c51de --- /dev/null +++ b/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt @@ -0,0 +1,15 @@ +cmake_minimum_required(VERSION 3.16) +project(mooncake-scheduler-benchmark LANGUAGES CXX) + +set(BUILD_BENCHMARK OFF CACHE BOOL "" FORCE) +set(BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) +set(BUILD_UNIT_TESTS OFF CACHE BOOL "" FORCE) + +add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../.. + ${CMAKE_CURRENT_BINARY_DIR}/transfer-engine) + +add_executable(scheduler_benchmark scheduler_benchmark.cpp) +target_compile_features(scheduler_benchmark PRIVATE cxx_std_20) +target_link_libraries(scheduler_benchmark PRIVATE transfer_engine) + +configure_file(compare_results.py compare_results.py COPYONLY) diff --git a/mooncake-transfer-engine/benchmark/scheduler/compare_results.py b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py new file mode 100644 index 0000000000..d2851fc078 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Compare scheduler_benchmark JSONL output from baseline and scheduled runs.""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +from pathlib import Path + + +def load(path: Path, expected_scheduling: bool) -> list[dict]: + records = [] + with path.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + if not line.strip(): + continue + record = json.loads(line) + if record.get("schema_version") != 1: + raise ValueError(f"{path}:{line_number}: unsupported schema") + if record.get("scheduling") is not expected_scheduling: + raise ValueError( + f"{path}:{line_number}: unexpected scheduling mode" + ) + records.append(record) + if not records: + raise ValueError(f"{path}: no records") + return records + + +def median(records: list[dict], key: str) -> float: + return statistics.median(float(record[key]) for record in records) + + +def class_p99(records: list[dict]) -> dict[str, float]: + values: dict[str, list[float]] = {} + for record in records: + for traffic in record["classes"]: + values.setdefault(traffic["name"], []).append(traffic["p99_us"]) + return {name: statistics.median(samples) for name, samples in values.items()} + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("baseline", type=Path) + parser.add_argument("scheduled", type=Path) + parser.add_argument("--min-p99-reduction", type=float, default=20.0) + parser.add_argument("--min-throughput-retention", type=float, default=95.0) + args = parser.parse_args() + + try: + baseline = load(args.baseline, False) + scheduled = load(args.scheduled, True) + if len(baseline) != len(scheduled): + raise ValueError("run counts differ") + baseline_p99 = class_p99(baseline) + scheduled_p99 = class_p99(scheduled) + if baseline_p99.keys() != scheduled_p99.keys(): + raise ValueError("traffic class sets differ") + + baseline_bw = median(baseline, "aggregate_throughput_gbps") + scheduled_bw = median(scheduled, "aggregate_throughput_gbps") + if baseline_bw <= 0: + raise ValueError("baseline throughput must be positive") + retention = scheduled_bw / baseline_bw * 100.0 + passed = retention >= args.min_throughput_retention + + print("class baseline p99 scheduled p99 reduction") + for name in baseline_p99: + if baseline_p99[name] <= 0: + raise ValueError(f"baseline P99 for {name} must be positive") + reduction = ( + (baseline_p99[name] - scheduled_p99[name]) + / baseline_p99[name] + * 100.0 + ) + print( + f"{name:<21} {baseline_p99[name]:>10.2f} us" + f" {scheduled_p99[name]:>12.2f} us {reduction:>10.2f}%" + ) + if name.startswith("foreground-"): + passed = passed and reduction >= args.min_p99_reduction + print( + f"aggregate throughput: {baseline_bw:.6f} -> " + f"{scheduled_bw:.6f} GB/s ({retention:.2f}% retained)" + ) + print("result: " + ("PASS" if passed else "FAIL")) + return 0 if passed else 1 + except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) From 01bd087d868f6369b9125bd4d17a133149cc028b Mon Sep 17 00:00:00 2001 From: yuanhao Date: Wed, 9 Sep 2026 16:33:11 +0800 Subject: [PATCH 5/9] =?UTF-8?q?build(transfer-engine):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8DFindSpDiag.cmake=E7=9A=84=E5=8C=85=E5=90=AB=E8=B7=AF?= =?UTF-8?q?=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mooncake-transfer-engine/src/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index ed0c4c52ad..2ea755cb4c 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -1,5 +1,5 @@ file(GLOB ENGINE_SOURCES "*.cpp") -include(${CMAKE_SOURCE_DIR}/mooncake-common/FindSpDiag.cmake) +include(${CMAKE_CURRENT_LIST_DIR}/../../mooncake-common/FindSpDiag.cmake) add_subdirectory(common) add_subdirectory(transport) From 351c4fc4303ecb119b5e351ad1fbda70e7976dac Mon Sep 17 00:00:00 2001 From: yuanhao Date: Wed, 9 Sep 2026 16:42:47 +0800 Subject: [PATCH 6/9] =?UTF-8?q?build(scheduler=5Fbenchmark):=20=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=E7=BC=BA=E5=A4=B1=E7=9A=84=20include=20=E8=B7=AF?= =?UTF-8?q?=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt b/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt index e38d8c51de..f67c886e31 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt +++ b/mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt @@ -10,6 +10,8 @@ add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../.. add_executable(scheduler_benchmark scheduler_benchmark.cpp) target_compile_features(scheduler_benchmark PRIVATE cxx_std_20) +target_include_directories( + scheduler_benchmark PRIVATE ${CMAKE_CURRENT_LIST_DIR}/../../include) target_link_libraries(scheduler_benchmark PRIVATE transfer_engine) configure_file(compare_results.py compare_results.py COPYONLY) From 78e1eb60ef9cd2a18b9c2d1d66be4b098a276d02 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Thu, 10 Sep 2026 11:11:01 +0800 Subject: [PATCH 7/9] =?UTF-8?q?feat(scheduler):=20=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E5=8F=8C=E8=A7=92=E8=89=B2=E5=8F=8A=E5=8D=8F=E8=AE=AE=E9=80=89?= =?UTF-8?q?=E6=8B=A9=EF=BC=8C=E4=BC=98=E5=8C=96NUMA=E7=BB=91=E5=AE=9A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 扩展基准测试为目标/发起者双角色，取消对tebench的依赖。 添加--protocol、--device_name、--numa_node参数，支持多设备列表及拓扑生成。 内存分配改用numa_alloc，目标模式添加信号处理。 更新README文档，提供完整运行示例。 --- .../benchmark/scheduler/README.md | 50 ++++++++--- .../benchmark/scheduler/compare_results.py | 8 ++ .../scheduler/scheduler_benchmark.cpp | 90 ++++++++++++++++--- 3 files changed, 124 insertions(+), 24 deletions(-) diff --git a/mooncake-transfer-engine/benchmark/scheduler/README.md b/mooncake-transfer-engine/benchmark/scheduler/README.md index 5750c3920e..0a33ac1402 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/README.md +++ b/mooncake-transfer-engine/benchmark/scheduler/README.md @@ -1,8 +1,9 @@ # Classic TE Scheduler Benchmark This standalone benchmark compares the original classic Transfer Engine -submission path with the task scheduler. It does not change or link against -the existing `tebench` implementation. +submission path with the task scheduler. The same executable provides target +and initiator roles. Both roles explicitly install `ub`, `rdma`, or `tcp`, so +the selected protocol cannot be replaced by auto-discovery. The workload follows the scheduler evaluation design: @@ -17,7 +18,12 @@ Configure this directory as an independent CMake project: ```bash cmake -S mooncake-transfer-engine/benchmark/scheduler \ - -B build/scheduler-benchmark + -B build/scheduler-benchmark \ + -DCMAKE_BUILD_TYPE=Release \ + -DENABLE_DEBUG_SYMBOLS=OFF \ + -DUSE_UB=ON \ + -DURMA_INCLUDE_DIR=/usr/include \ + -DURMA_LIBRARY=/usr/lib64/liburma.so cmake --build build/scheduler-benchmark -j ``` @@ -26,23 +32,39 @@ The executable and comparison script are produced in ## Run -Start the existing classic `tebench` target with at least a 1 GiB buffer: +Start the target on the server that owns the remote memory: ```bash -./tebench --backend=classic --seg_type=DRAM --total_buffer_size=1073741824 +./scheduler_benchmark \ + --mode=target \ + --protocol=ub \ + --device_name=urma0 \ + --numa_node=0 ``` -Copy the printed segment name. Run the baseline and scheduled cases from the -standalone benchmark build directory: +Copy the printed `TARGET_SEGMENT=:` value. Run the baseline and +scheduled cases on the initiator with its local UB device: ```bash rm -f baseline.jsonl scheduled.jsonl -./scheduler_benchmark --target_seg_name= \ - --scheduling=false --output_jsonl=baseline.jsonl - -./scheduler_benchmark --target_seg_name= \ - --scheduling=true --output_jsonl=scheduled.jsonl +./scheduler_benchmark \ + --mode=initiator \ + --protocol=ub \ + --device_name=urma0 \ + --numa_node=0 \ + --target_seg_name= \ + --scheduling=false \ + --output_jsonl=baseline.jsonl + +./scheduler_benchmark \ + --mode=initiator \ + --protocol=ub \ + --device_name=urma0 \ + --numa_node=0 \ + --target_seg_name= \ + --scheduling=true \ + --output_jsonl=scheduled.jsonl python3 compare_results.py baseline.jsonl scheduled.jsonl ``` @@ -56,3 +78,7 @@ Use the same target, transport configuration, CPU affinity, registered memory, and connection setup for both runs. The scheduler parameters can be scanned with `--scheduler_quantum_bytes`, `--scheduler_max_inflight_bytes`, `--scheduler_reserved_high_bytes`, and `--scheduler_max_slices`. + +For multiple UB devices, pass a comma-separated list such as +`--device_name=urma0,urma1`. The NUMA location and device topology are generated +from `--numa_node` and `--device_name` on each host. diff --git a/mooncake-transfer-engine/benchmark/scheduler/compare_results.py b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py index d2851fc078..5922cc9c87 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/compare_results.py +++ b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py @@ -54,6 +54,13 @@ def main() -> int: scheduled = load(args.scheduled, True) if len(baseline) != len(scheduled): raise ValueError("run counts differ") + baseline_protocols = {record["protocol"] for record in baseline} + scheduled_protocols = {record["protocol"] for record in scheduled} + if ( + len(baseline_protocols) != 1 + or baseline_protocols != scheduled_protocols + ): + raise ValueError("baseline and scheduled protocols differ") baseline_p99 = class_p99(baseline) scheduled_p99 = class_p99(scheduled) if baseline_p99.keys() != scheduled_p99.keys(): @@ -66,6 +73,7 @@ def main() -> int: retention = scheduled_bw / baseline_bw * 100.0 passed = retention >= args.min_throughput_retention + print(f"protocol: {next(iter(baseline_protocols))}") print("class baseline p99 scheduled p99 reduction") for name in baseline_p99: if baseline_p99[name] <= 0: diff --git a/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp index 4f1f0cebed..1ef1dd8d9f 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp +++ b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp @@ -3,11 +3,12 @@ #include #include +#include #include #include +#include #include -#include #include #include #include @@ -23,11 +24,16 @@ #include "scheduler/scheduler_policy.h" #include "transfer_engine.h" -DEFINE_string(target_seg_name, "", "Target segment printed by tebench"); +DEFINE_string(mode, "initiator", "Benchmark role: target or initiator"); +DEFINE_string(protocol, "ub", "Transport protocol: ub, rdma, or tcp"); +DEFINE_string(device_name, "urma0", + "Comma-separated devices for UB/RDMA, e.g. urma0,urma1"); +DEFINE_int32(numa_node, 0, "NUMA node used for the registered buffer"); +DEFINE_string(target_seg_name, "", "Segment printed by the target process"); DEFINE_string(metadata_conn_string, "P2PHANDSHAKE", "Metadata connection string"); DEFINE_string(local_server_name, mooncake::getHostname(), - "Local server name used by the initiator"); + "Local server name used for P2P discovery"); DEFINE_bool(scheduling, false, "Submit through the classic TE scheduler"); DEFINE_uint64(buffer_size, 1ULL << 30, "Registered initiator buffer size in bytes"); @@ -53,6 +59,7 @@ DEFINE_uint32(scheduler_max_slices, 32, namespace { using Clock = std::chrono::steady_clock; +volatile std::sig_atomic_t target_running = 1; struct TrafficClass { const char* name; @@ -73,6 +80,42 @@ void check(const mooncake::Status& status, const char* operation) { LOG_ASSERT(status.ok()) << operation << " failed: " << status.ToString(); } +void stopTarget(int) { target_running = 0; } + +std::string topologyJson() { + std::string devices; + size_t begin = 0; + while (begin < FLAGS_device_name.size()) { + const size_t comma = FLAGS_device_name.find(',', begin); + const std::string device = + FLAGS_device_name.substr(begin, comma - begin); + LOG_ASSERT(!device.empty()) << "--device_name contains an empty entry"; + if (!devices.empty()) devices += ','; + devices += "\"" + device + "\""; + if (comma == std::string::npos) break; + begin = comma + 1; + } + LOG_ASSERT(!devices.empty()) + << "--device_name is required for " << FLAGS_protocol; + return "{\"cpu:" + std::to_string(FLAGS_numa_node) + "\":[[" + devices + + "],[]]}"; +} + +void installSelectedTransport(mooncake::TransferEngine& engine) { + LOG_ASSERT(FLAGS_protocol == "ub" || FLAGS_protocol == "rdma" || + FLAGS_protocol == "tcp") + << "--protocol must be ub, rdma, or tcp"; + if (FLAGS_protocol == "tcp") { + LOG_ASSERT(engine.installTransport("tcp", nullptr) != nullptr) + << "TCP Transport installation failed"; + return; + } + std::string topology = topologyJson(); + void* args[] = {topology.data(), nullptr}; + LOG_ASSERT(engine.installTransport(FLAGS_protocol, args) != nullptr) + << FLAGS_protocol << " Transport installation failed"; +} + double percentile(std::vector samples, double value) { if (samples.empty()) return 0.0; std::sort(samples.begin(), samples.end()); @@ -185,6 +228,7 @@ void appendRecord(const std::vector& classes, output << std::fixed << std::setprecision(6) << "{\"schema_version\":1,\"scheduling\":" << (FLAGS_scheduling ? "true" : "false") + << ",\"protocol\":\"" << FLAGS_protocol << "\"" << ",\"repetition\":" << repetition << ",\"aggregate_throughput_gbps\":" << aggregate_throughput << ",\"classes\":["; @@ -213,13 +257,16 @@ void appendRecord(const std::vector& classes, } // namespace int main(int argc, char** argv) { - gflags::SetUsageMessage( - "Classic Transfer Engine scheduler A/B benchmark initiator"); + gflags::SetUsageMessage("Classic Transfer Engine scheduler A/B benchmark"); gflags::ParseCommandLineFlags(&argc, &argv, true); google::InitGoogleLogging(argv[0]); - LOG_ASSERT(!FLAGS_target_seg_name.empty()) - << "--target_seg_name is required"; + LOG_ASSERT(FLAGS_mode == "target" || FLAGS_mode == "initiator") + << "--mode must be target or initiator"; + if (FLAGS_mode == "initiator") + LOG_ASSERT(!FLAGS_target_seg_name.empty()) + << "--target_seg_name is required in initiator mode"; + LOG_ASSERT(FLAGS_numa_node >= 0) << "--numa_node must be non-negative"; LOG_ASSERT(FLAGS_foreground_threads > 0 && FLAGS_background_threads > 0); LOG_ASSERT(FLAGS_foreground_threads <= 128 && FLAGS_background_threads <= 128) @@ -252,17 +299,33 @@ int main(int argc, char** argv) { LOG_ASSERT(FLAGS_buffer_size >= required_buffer) << "--buffer_size must be at least " << required_buffer; - auto* local_buffer = static_cast( - std::aligned_alloc(4096, static_cast(FLAGS_buffer_size))); + auto* local_buffer = static_cast(numa_alloc_onnode( + static_cast(FLAGS_buffer_size), FLAGS_numa_node)); LOG_ASSERT(local_buffer != nullptr) << "failed to allocate local buffer"; - mooncake::TransferEngine engine(true); + mooncake::TransferEngine engine(false); LOG_ASSERT( engine.init(FLAGS_metadata_conn_string, FLAGS_local_server_name) == 0) << "TransferEngine initialization failed"; - LOG_ASSERT(engine.registerLocalMemory(local_buffer, FLAGS_buffer_size) == 0) + installSelectedTransport(engine); + const std::string location = "cpu:" + std::to_string(FLAGS_numa_node); + LOG_ASSERT(engine.registerLocalMemory(local_buffer, FLAGS_buffer_size, + location) == 0) << "local memory registration failed"; + if (FLAGS_mode == "target") { + std::signal(SIGINT, stopTarget); + std::signal(SIGTERM, stopTarget); + std::cout << "TARGET_SEGMENT=" << engine.getLocalIpAndPort() + << " PROTOCOL=" << FLAGS_protocol << std::endl; + while (target_running) + std::this_thread::sleep_for(std::chrono::seconds(1)); + engine.unregisterLocalMemory(local_buffer); + engine.freeEngine(); + numa_free(local_buffer, FLAGS_buffer_size); + return 0; + } + if (FLAGS_scheduling) { mooncake::scheduling::SchedulerConfig config; config.quantum_bytes = FLAGS_scheduler_quantum_bytes; @@ -276,6 +339,9 @@ int main(int argc, char** argv) { auto segment = engine.getMetadata()->getSegmentDescByID(target); LOG_ASSERT(segment && !segment->buffers.empty()) << "target segment has no registered buffers"; + LOG_ASSERT(segment->protocol == FLAGS_protocol) + << "target protocol is " << segment->protocol << ", expected " + << FLAGS_protocol; const auto& remote_buffer = segment->buffers.front(); LOG_ASSERT(remote_buffer.length >= required_buffer) << "target buffer must be at least " << required_buffer << " bytes"; @@ -307,6 +373,6 @@ int main(int argc, char** argv) { engine.closeSegment(target); engine.unregisterLocalMemory(local_buffer); engine.freeEngine(); - std::free(local_buffer); + numa_free(local_buffer, FLAGS_buffer_size); return 0; } From 5e4d855f4d70232c6c4555270be2774a3605bcb9 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Thu, 10 Sep 2026 17:38:46 +0800 Subject: [PATCH 8/9] =?UTF-8?q?feat(scheduler):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E7=B1=BB=E5=88=AB=E6=9D=83=E9=87=8D=E5=92=8C=E5=BB=B6=E8=BF=9F?= =?UTF-8?q?=E5=88=86=E5=B8=83=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 支持通过 --scheduler_class_weights 配置 HIGH:MEDIUM:LOW 权重，影响调度器决策。 JSONL 输出升级至 schema v2，新增每个流量类别的 avg/min/p50/p99/p999/p9999/max 延迟与吞吐量字段。 对比脚本更新以计算各类别的 Avg 减少、P99 减少和带宽保留率，并支持 class_weights 验证。 --- .../benchmark/scheduler/README.md | 13 +- .../benchmark/scheduler/compare_results.py | 92 ++++++++++---- .../scheduler/scheduler_benchmark.cpp | 116 +++++++++++++++--- 3 files changed, 180 insertions(+), 41 deletions(-) diff --git a/mooncake-transfer-engine/benchmark/scheduler/README.md b/mooncake-transfer-engine/benchmark/scheduler/README.md index 0a33ac1402..7ef37cf9ed 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/README.md +++ b/mooncake-transfer-engine/benchmark/scheduler/README.md @@ -64,6 +64,7 @@ rm -f baseline.jsonl scheduled.jsonl --numa_node=0 \ --target_seg_name= \ --scheduling=true \ + --scheduler_class_weights=8:4:1 \ --output_jsonl=scheduled.jsonl python3 compare_results.py baseline.jsonl scheduled.jsonl @@ -74,10 +75,20 @@ by at least 20% and median aggregate throughput retains at least 95% of the baseline. Override the gates with `--min-p99-reduction` and `--min-throughput-retention`. +Each repetition records `avg_us`, `min_us`, `p50_us`, `p99_us`, `p999_us`, +`p9999_us`, `max_us`, operation count, and bandwidth for every traffic class. +The comparison report shows the median Avg reduction, P99 reduction, and +bandwidth retention for each class, followed by aggregate bandwidth retention. +Remove older JSONL files before running because schema version 2 adds these +latency fields and output uses append mode. + Use the same target, transport configuration, CPU affinity, registered memory, and connection setup for both runs. The scheduler parameters can be scanned with `--scheduler_quantum_bytes`, `--scheduler_max_inflight_bytes`, -`--scheduler_reserved_high_bytes`, and `--scheduler_max_slices`. +`--scheduler_reserved_high_bytes`, `--scheduler_max_slices`, and +`--scheduler_class_weights`. Class weights use `HIGH:MEDIUM:LOW` order, must be +positive integers, and default to `8:4:1`. They affect only runs with +`--scheduling=true` and are recorded in every JSONL result. For multiple UB devices, pass a comma-separated list such as `--device_name=urma0,urma1`. The NUMA location and device topology are generated diff --git a/mooncake-transfer-engine/benchmark/scheduler/compare_results.py b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py index 5922cc9c87..1106dd10a4 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/compare_results.py +++ b/mooncake-transfer-engine/benchmark/scheduler/compare_results.py @@ -17,7 +17,7 @@ def load(path: Path, expected_scheduling: bool) -> list[dict]: if not line.strip(): continue record = json.loads(line) - if record.get("schema_version") != 1: + if record.get("schema_version") != 2: raise ValueError(f"{path}:{line_number}: unsupported schema") if record.get("scheduling") is not expected_scheduling: raise ValueError( @@ -33,12 +33,35 @@ def median(records: list[dict], key: str) -> float: return statistics.median(float(record[key]) for record in records) -def class_p99(records: list[dict]) -> dict[str, float]: - values: dict[str, list[float]] = {} +def class_medians(records: list[dict]) -> dict[str, dict[str, float]]: + values: dict[str, dict[str, list[float]]] = {} for record in records: for traffic in record["classes"]: - values.setdefault(traffic["name"], []).append(traffic["p99_us"]) - return {name: statistics.median(samples) for name, samples in values.items()} + samples = values.setdefault( + traffic["name"], + {"avg_us": [], "p99_us": [], "throughput_gbps": []}, + ) + for metric in samples: + samples[metric].append(float(traffic[metric])) + return { + name: { + metric: statistics.median(samples) + for metric, samples in metrics.items() + } + for name, metrics in values.items() + } + + +def reduction(before: float, after: float, description: str) -> float: + if before <= 0: + raise ValueError(f"baseline {description} must be positive") + return (before - after) / before * 100.0 + + +def retention(before: float, after: float, description: str) -> float: + if before <= 0: + raise ValueError(f"baseline {description} must be positive") + return after / before * 100.0 def main() -> int: @@ -61,37 +84,58 @@ def main() -> int: or baseline_protocols != scheduled_protocols ): raise ValueError("baseline and scheduled protocols differ") - baseline_p99 = class_p99(baseline) - scheduled_p99 = class_p99(scheduled) - if baseline_p99.keys() != scheduled_p99.keys(): + baseline_classes = class_medians(baseline) + scheduled_classes = class_medians(scheduled) + if baseline_classes.keys() != scheduled_classes.keys(): raise ValueError("traffic class sets differ") baseline_bw = median(baseline, "aggregate_throughput_gbps") scheduled_bw = median(scheduled, "aggregate_throughput_gbps") - if baseline_bw <= 0: - raise ValueError("baseline throughput must be positive") - retention = scheduled_bw / baseline_bw * 100.0 - passed = retention >= args.min_throughput_retention + total_retention = retention( + baseline_bw, scheduled_bw, "aggregate throughput" + ) + passed = total_retention >= args.min_throughput_retention print(f"protocol: {next(iter(baseline_protocols))}") - print("class baseline p99 scheduled p99 reduction") - for name in baseline_p99: - if baseline_p99[name] <= 0: - raise ValueError(f"baseline P99 for {name} must be positive") - reduction = ( - (baseline_p99[name] - scheduled_p99[name]) - / baseline_p99[name] - * 100.0 + for name, baseline_metrics in baseline_classes.items(): + scheduled_metrics = scheduled_classes[name] + avg_reduction = reduction( + baseline_metrics["avg_us"], + scheduled_metrics["avg_us"], + f"Avg for {name}", + ) + p99_reduction = reduction( + baseline_metrics["p99_us"], + scheduled_metrics["p99_us"], + f"P99 for {name}", + ) + bandwidth_retention = retention( + baseline_metrics["throughput_gbps"], + scheduled_metrics["throughput_gbps"], + f"bandwidth for {name}", + ) + print(name) + print( + f" Avg: {baseline_metrics['avg_us']:.2f} -> " + f"{scheduled_metrics['avg_us']:.2f} us " + f"({avg_reduction:.2f}% reduction)" + ) + print( + f" P99: {baseline_metrics['p99_us']:.2f} -> " + f"{scheduled_metrics['p99_us']:.2f} us " + f"({p99_reduction:.2f}% reduction)" ) print( - f"{name:<21} {baseline_p99[name]:>10.2f} us" - f" {scheduled_p99[name]:>12.2f} us {reduction:>10.2f}%" + " Bandwidth: " + f"{baseline_metrics['throughput_gbps']:.6f} -> " + f"{scheduled_metrics['throughput_gbps']:.6f} GB/s " + f"({bandwidth_retention:.2f}% retained)" ) if name.startswith("foreground-"): - passed = passed and reduction >= args.min_p99_reduction + passed = passed and p99_reduction >= args.min_p99_reduction print( f"aggregate throughput: {baseline_bw:.6f} -> " - f"{scheduled_bw:.6f} GB/s ({retention:.2f}% retained)" + f"{scheduled_bw:.6f} GB/s ({total_retention:.2f}% retained)" ) print("result: " + ("PASS" if passed else "FAIL")) return 0 if passed else 1 diff --git a/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp index 1ef1dd8d9f..4cfa3f3e34 100644 --- a/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp +++ b/mooncake-transfer-engine/benchmark/scheduler/scheduler_benchmark.cpp @@ -6,6 +6,8 @@ #include #include +#include +#include #include #include #include @@ -14,7 +16,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -55,6 +59,8 @@ DEFINE_uint64(scheduler_reserved_high_bytes, 1ULL << 20, "Scheduler bytes reserved for HIGH traffic"); DEFINE_uint32(scheduler_max_slices, 32, "Scheduler maximum transport slices per grant"); +DEFINE_string(scheduler_class_weights, "8:4:1", + "Scheduler class weights in HIGH:MEDIUM:LOW order"); namespace { @@ -76,12 +82,57 @@ struct WorkerResult { double duration_seconds{0}; }; +struct LatencyMetrics { + double average_us{0}; + double min_us{0}; + double p50_us{0}; + double p99_us{0}; + double p999_us{0}; + double p9999_us{0}; + double max_us{0}; +}; + void check(const mooncake::Status& status, const char* operation) { LOG_ASSERT(status.ok()) << operation << " failed: " << status.ToString(); } void stopTarget(int) { target_running = 0; } +std::array parseClassWeights() { + const std::string_view text = FLAGS_scheduler_class_weights; + const size_t first_separator = text.find(':'); + const size_t second_separator = + first_separator == std::string_view::npos + ? std::string_view::npos + : text.find(':', first_separator + 1); + LOG_ASSERT(first_separator != std::string_view::npos && + second_separator != std::string_view::npos && + text.find(':', second_separator + 1) == std::string_view::npos) + << "--scheduler_class_weights must be HIGH:MEDIUM:LOW, for example " + "8:4:1"; + + const std::array tokens = { + text.substr(0, first_separator), + text.substr(first_separator + 1, + second_separator - first_separator - 1), + text.substr(second_separator + 1), + }; + std::array weights{}; + for (size_t i = 0; i < tokens.size(); ++i) { + uint64_t value = 0; + const auto parsed = std::from_chars( + tokens[i].data(), tokens[i].data() + tokens[i].size(), value); + LOG_ASSERT(parsed.ec == std::errc() && + parsed.ptr == tokens[i].data() + tokens[i].size() && + value > 0 && + value <= std::numeric_limits::max()) + << "--scheduler_class_weights entries must be positive uint32 " + "values"; + weights[i] = static_cast(value); + } + return weights; +} + std::string topologyJson() { std::string devices; size_t begin = 0; @@ -116,14 +167,28 @@ void installSelectedTransport(mooncake::TransferEngine& engine) { << FLAGS_protocol << " Transport installation failed"; } -double percentile(std::vector samples, double value) { - if (samples.empty()) return 0.0; - std::sort(samples.begin(), samples.end()); - const double rank = value / 100.0 * (samples.size() - 1); +double percentile(const std::vector& sorted, double value) { + if (sorted.empty()) return 0.0; + const double rank = value / 100.0 * (sorted.size() - 1); const size_t lower = static_cast(rank); - const size_t upper = std::min(lower + 1, samples.size() - 1); + const size_t upper = std::min(lower + 1, sorted.size() - 1); const double fraction = rank - lower; - return samples[lower] * (1.0 - fraction) + samples[upper] * fraction; + return sorted[lower] * (1.0 - fraction) + sorted[upper] * fraction; +} + +LatencyMetrics summarizeLatency(std::vector samples) { + if (samples.empty()) return {}; + const double sum = std::accumulate(samples.begin(), samples.end(), 0.0); + std::sort(samples.begin(), samples.end()); + return { + sum / samples.size(), + samples.front(), + percentile(samples, 50.0), + percentile(samples, 99.0), + percentile(samples, 99.9), + percentile(samples, 99.99), + samples.back(), + }; } double runTransfer(mooncake::TransferEngine& engine, mooncake::SegmentID target, @@ -196,40 +261,42 @@ void runWorker(mooncake::TransferEngine& engine, mooncake::SegmentID target, void appendRecord(const std::vector& classes, const std::vector>& results, + const std::array& class_weights, int repetition) { std::ofstream output(FLAGS_output_jsonl, std::ios::app); LOG_ASSERT(output) << "cannot open " << FLAGS_output_jsonl; double aggregate_throughput = 0.0; struct ClassMetrics { - double p99_us; + LatencyMetrics latency; double throughput_gbps; uint64_t operations; }; std::vector metrics; for (size_t class_index = 0; class_index < classes.size(); ++class_index) { std::vector latencies; - uint64_t bytes = 0; uint64_t operations = 0; double throughput = 0.0; for (const auto& worker : results[class_index]) { latencies.insert(latencies.end(), worker.latency_us.begin(), worker.latency_us.end()); - bytes += worker.bytes; operations += worker.latency_us.size(); if (worker.duration_seconds > 0) throughput += worker.bytes / 1e9 / worker.duration_seconds; } aggregate_throughput += throughput; metrics.push_back( - {percentile(std::move(latencies), 99.0), throughput, operations}); + {summarizeLatency(std::move(latencies)), throughput, operations}); } output << std::fixed << std::setprecision(6) - << "{\"schema_version\":1,\"scheduling\":" + << "{\"schema_version\":2,\"scheduling\":" << (FLAGS_scheduling ? "true" : "false") << ",\"protocol\":\"" << FLAGS_protocol << "\"" << ",\"repetition\":" << repetition + << ",\"class_weights\":{\"high\":" << class_weights[0] + << ",\"medium\":" << class_weights[1] + << ",\"low\":" << class_weights[2] << '}' << ",\"aggregate_throughput_gbps\":" << aggregate_throughput << ",\"classes\":["; for (size_t i = 0; i < classes.size(); ++i) { @@ -238,7 +305,13 @@ void appendRecord(const std::vector& classes, << "\",\"threads\":" << classes[i].threads << ",\"block_size\":" << classes[i].block_size << ",\"operations\":" << metrics[i].operations - << ",\"p99_us\":" << metrics[i].p99_us + << ",\"avg_us\":" << metrics[i].latency.average_us + << ",\"min_us\":" << metrics[i].latency.min_us + << ",\"p50_us\":" << metrics[i].latency.p50_us + << ",\"p99_us\":" << metrics[i].latency.p99_us + << ",\"p999_us\":" << metrics[i].latency.p999_us + << ",\"p9999_us\":" << metrics[i].latency.p9999_us + << ",\"max_us\":" << metrics[i].latency.max_us << ",\"throughput_gbps\":" << metrics[i].throughput_gbps << '}'; } output << "]}\n"; @@ -246,12 +319,21 @@ void appendRecord(const std::vector& classes, std::cout << "repetition=" << repetition << " scheduling=" << std::boolalpha << FLAGS_scheduling + << " class_weights=" << class_weights[0] << ':' + << class_weights[1] << ':' << class_weights[2] << " aggregate_throughput=" << aggregate_throughput << " GB/s" << std::endl; for (size_t i = 0; i < classes.size(); ++i) - std::cout << " " << classes[i].name << ": p99=" << metrics[i].p99_us - << " us throughput=" << metrics[i].throughput_gbps << " GB/s" - << std::endl; + std::cout << " " << classes[i].name + << ": avg=" << metrics[i].latency.average_us + << " us min=" << metrics[i].latency.min_us + << " us p50=" << metrics[i].latency.p50_us + << " us p99=" << metrics[i].latency.p99_us + << " us p999=" << metrics[i].latency.p999_us + << " us p9999=" << metrics[i].latency.p9999_us + << " us max=" << metrics[i].latency.max_us + << " us throughput=" << metrics[i].throughput_gbps + << " GB/s operations=" << metrics[i].operations << std::endl; } } // namespace @@ -273,6 +355,7 @@ int main(int argc, char** argv) { << "thread counts must not exceed 128 per traffic class"; LOG_ASSERT(FLAGS_warmup_seconds >= 0 && FLAGS_duration_seconds > 0 && FLAGS_repetitions > 0); + const auto class_weights = parseClassWeights(); std::vector classes = { {"foreground-4k", 4ULL << 10, FLAGS_foreground_threads, @@ -332,6 +415,7 @@ int main(int argc, char** argv) { config.max_inflight_bytes = FLAGS_scheduler_max_inflight_bytes; config.reserved_high_bytes = FLAGS_scheduler_reserved_high_bytes; config.max_slices = FLAGS_scheduler_max_slices; + config.class_weights = class_weights; check(engine.configureScheduling(config), "configureScheduling"); } @@ -367,7 +451,7 @@ int main(int argc, char** argv) { } } for (auto& worker : workers) worker.join(); - appendRecord(classes, results, repetition); + appendRecord(classes, results, class_weights, repetition); } engine.closeSegment(target); From 04d334cca58b41de87c0a1b6fabc9233d70854d7 Mon Sep 17 00:00:00 2001 From: yuanhao Date: Thu, 10 Sep 2026 20:23:50 +0800 Subject: [PATCH 9/9] =?UTF-8?q?feat(scheduling):=20=E6=B7=BB=E5=8A=A0?= =?UTF-8?q?=E4=BC=A0=E8=BE=93=E5=BC=95=E6=93=8E=E8=B0=83=E5=BA=A6=E9=85=8D?= =?UTF-8?q?=E7=BD=AE=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为 Client 和 RealClient 增加 configureScheduling 接口，支持通过命令行参数（如 --scheduling、--scheduler_quantum_bytes 等）启用并配置传输引擎调度策略，并解析类权重字符串。 --- mooncake-store/include/client_service.h | 3 + mooncake-store/include/real_client.h | 3 + mooncake-store/src/client_service.cpp | 7 +++ mooncake-store/src/real_client.cpp | 7 +++ mooncake-store/src/real_client_main.cpp | 75 +++++++++++++++++++++++++ 5 files changed, 95 insertions(+) diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 631e4328bd..fc77348269 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -95,6 +95,9 @@ class Client { std::map labels = {}, const std::string& tenant_id = "default"); + Status configureScheduling( + const scheduling::SchedulerConfig& scheduler_config); + /** * @brief Retrieves data for a given key * @param object_key Key to retrieve diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 6e3a9be8c3..da21b7241b 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -79,6 +79,9 @@ class RealClient : public PyClient { // Factory to create shared instances and auto-register to ResourceTracker static std::shared_ptr create(); + Status configureScheduling( + const scheduling::SchedulerConfig &scheduler_config); + int setup_real( const std::string &local_hostname, const std::string &metadata_server, size_t global_segment_size = 1024 * 1024 * 16, diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 0dca231eca..0726a460a7 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -892,6 +892,13 @@ ErrorCode Client::InitTransferEngine( return ErrorCode::OK; } +Status Client::configureScheduling( + const scheduling::SchedulerConfig& scheduler_config) { + if (!transfer_engine_) + return Status::InvalidArgument("Transfer engine is not initialized"); + return transfer_engine_->configureScheduling(scheduler_config); +} + void Client::InitTransferSubmitter() { // Initialize TransferSubmitter after transfer engine is ready // Keep using logical local_hostname for name-based behaviors; endpoint is diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index e26b6ea96a..1b23c4a214 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -696,6 +696,13 @@ tl::expected RealClient::setup_ascend_internal( return {}; } +Status RealClient::configureScheduling( + const scheduling::SchedulerConfig &scheduler_config) { + if (!client_) + return Status::InvalidArgument("Real client is not initialized"); + return client_->configureScheduling(scheduler_config); +} + tl::expected RealClient::setup_internal( const std::string &local_hostname, const std::string &metadata_server, size_t global_segment_size, size_t local_buffer_size, diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 4b4137e1c0..3f3852a04a 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -1,6 +1,8 @@ #include #include +#include #include +#include #include #include "client_service.h" @@ -8,6 +10,7 @@ #include "config.h" #include "mooncake_logging.h" #include "real_client.h" +#include "scheduler/scheduler_policy.h" using namespace mooncake; @@ -23,6 +26,16 @@ DEFINE_string(global_segment_size, "4 GB", "Size of global segment"); DEFINE_string(local_buffer_size, "0", "Size of local buffer (e.g., 16MB, 1GB)"); DEFINE_int32(threads, 1, "Number of threads for client service"); DEFINE_string(tenant_id, "default", "Tenant identifier"); +DEFINE_bool(scheduling, false, "Enable Transfer Engine scheduling"); +DEFINE_uint64(scheduler_quantum_bytes, 1ULL << 20, "Scheduler byte quantum"); +DEFINE_uint64(scheduler_max_inflight_bytes, 16ULL << 20, + "Scheduler maximum in-flight bytes"); +DEFINE_uint64(scheduler_reserved_high_bytes, 1ULL << 20, + "Scheduler bytes reserved for HIGH traffic"); +DEFINE_uint32(scheduler_max_slices, 32, + "Scheduler maximum transport slices per grant"); +DEFINE_string(scheduler_class_weights, "8:4:1", + "Scheduler class weights in HIGH:MEDIUM:LOW order"); DEFINE_bool(enable_offload, false, "Enable offload availability"); DEFINE_bool(start_offload_rpc_server, true, "Expose TCP RPC for disk-tier reads " @@ -36,6 +49,40 @@ DEFINE_int32(offload_rpc_thread_num, 8, DECLARE_bool(enable_http_server); DECLARE_int32(http_port); +namespace { + +std::array parseClassWeights() { + const std::string_view text = FLAGS_scheduler_class_weights; + const size_t first_separator = text.find(':'); + const size_t second_separator = + first_separator == std::string_view::npos + ? std::string_view::npos + : text.find(':', first_separator + 1); + LOG_ASSERT(first_separator != std::string_view::npos && + second_separator != std::string_view::npos && + text.find(':', second_separator + 1) == std::string_view::npos) + << "--scheduler_class_weights must be HIGH:MEDIUM:LOW, for example " + "8:4:1"; + + const std::array tokens = { + text.substr(0, first_separator), + text.substr(first_separator + 1, + second_separator - first_separator - 1), + text.substr(second_separator + 1), + }; + std::array weights{}; + for (size_t i = 0; i < tokens.size(); ++i) { + const auto value = mooncake::parseFromString(tokens[i]); + LOG_ASSERT(value.has_value() && *value > 0) + << "--scheduler_class_weights entries must be positive uint32 " + "values"; + weights[i] = *value; + } + return weights; +} + +} // namespace + namespace mooncake { void RegisterClientRpcService(coro_rpc::coro_rpc_server &server, RealClient &real_client) { @@ -143,6 +190,34 @@ int main(int argc, char *argv[]) { return -1; } + if (FLAGS_scheduling) { + mooncake::scheduling::SchedulerConfig scheduler_config; + scheduler_config.quantum_bytes = FLAGS_scheduler_quantum_bytes; + scheduler_config.max_inflight_bytes = + FLAGS_scheduler_max_inflight_bytes; + scheduler_config.reserved_high_bytes = + FLAGS_scheduler_reserved_high_bytes; + scheduler_config.max_slices = FLAGS_scheduler_max_slices; + scheduler_config.class_weights = parseClassWeights(); + const auto status = + client_inst->configureScheduling(scheduler_config); + if (!status.ok()) { + LOG(ERROR) << "Failed to configure Transfer Engine scheduling: " + << status.ToString(); + return -1; + } + LOG(INFO) << "Transfer Engine scheduling enabled with class weights " + << scheduler_config.class_weights[0] << ':' + << scheduler_config.class_weights[1] << ':' + << scheduler_config.class_weights[2] + << ", quantum_bytes=" << scheduler_config.quantum_bytes + << ", max_inflight_bytes=" + << scheduler_config.max_inflight_bytes + << ", reserved_high_bytes=" + << scheduler_config.reserved_high_bytes + << ", max_slices=" << scheduler_config.max_slices; + } + if (client_inst->start_dummy_client_monitor()) { LOG(FATAL) << "Failed to start dummy client monitor thread"; return -1;