Task scheduling

mooncake-store/include/client_service.h

@@ -95,6 +95,9 @@ class Client {
         std::map<std::string, std::string> 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

mooncake-store/include/operation_options.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <optional>
+
+#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<SchedulingHint> scheduling;
+};
+
+}  // namespace mooncake

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<RealClient> 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,

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<TransferFuture> submit(const Replica::Descriptor& replica,
                                          std::vector<Slice>& 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<TransferFuture> submit_batch(
         const std::vector<Replica::Descriptor>& replicas,
         std::vector<std::vector<Slice>>& all_slices,
-        TransferRequest::OpCode op_code);
+        TransferRequest::OpCode op_code, const OperationOptions& options = {});
 
     std::optional<TransferFuture> submit_batch_get_offload_object(
         const std::string& transfer_engine_addr,
@@ -664,11 +666,12 @@ class TransferSubmitter {
     std::optional<TransferFuture> submitTransferEngineOperation(
         const AllocatedBuffer::Descriptor& handle,
         const std::vector<Slice>& slices, const TransferRequest::OpCode op_code,
-        uint64_t src_offset = 0);
+        uint64_t src_offset = 0, const OperationOptions& options = {});
 
     std::optional<TransferFuture> submitMemoryReadOperation(
         const AllocatedBuffer::Descriptor& handle,
-        const std::vector<Slice>& slices, uint64_t src_offset);
+        const std::vector<Slice>& slices, uint64_t src_offset,
+        const OperationOptions& options = {});
 
     std::optional<TransferFuture> submitFileReadOperation(
         const Replica::Descriptor& replica, std::vector<Slice>& slices,
@@ -681,7 +684,8 @@ class TransferSubmitter {
                                TransferRequest::OpCode op);
 
     std::optional<TransferFuture> submitTransfer(
-        std::vector<TransferRequest>& requests);
+        std::vector<TransferRequest>& requests,
+        const SchedulingHint& hint = {});
 };
 
 }  // namespace mooncake

mooncake-store/src/client_service.cpp

@@ -1183,6 +1183,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

mooncake-store/src/real_client.cpp

@@ -696,6 +696,13 @@ tl::expected<void, ErrorCode> 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<void, ErrorCode> RealClient::setup_internal(
     const std::string &local_hostname, const std::string &metadata_server,
     size_t global_segment_size, size_t local_buffer_size,

mooncake-store/src/real_client_main.cpp

@@ -1,13 +1,16 @@
 #include <gflags/gflags.h>
 #include <algorithm>
+#include <array>
 #include <csignal>
+#include <string_view>
 #include <ylt/coro_rpc/coro_rpc_server.hpp>
 
 #include "client_service.h"
 #include "common.h"
 #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<uint32_t, 3> 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<std::string_view, 3> tokens = {
+        text.substr(0, first_separator),
+        text.substr(first_separator + 1,
+                    second_separator - first_separator - 1),
+        text.substr(second_separator + 1),
+    };
+    std::array<uint32_t, 3> weights{};
+    for (size_t i = 0; i < tokens.size(); ++i) {
+        const auto value = mooncake::parseFromString<uint32_t>(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;

mooncake-store/src/transfer_task.cpp

@@ -1001,7 +1001,8 @@ TransferSubmitter::TransferSubmitter(TransferEngine& engine,
 
 std::optional<TransferFuture> TransferSubmitter::submit(
     const Replica::Descriptor& replica, std::vector<Slice>& slices,
-    TransferRequest::OpCode op_code, void* ptr, size_t size) {
+    TransferRequest::OpCode op_code, void* ptr, size_t size,
+    const OperationOptions& options) {
     std::optional<TransferFuture> future;
 
     if (replica.is_memory_replica()) {
@@ -1013,7 +1014,7 @@ std::optional<TransferFuture> 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<TransferFuture> 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<TransferFuture> TransferSubmitter::submit(
 std::optional<TransferFuture> TransferSubmitter::submit_batch(
     const std::vector<Replica::Descriptor>& replicas,
     std::vector<std::vector<Slice>>& all_slices,
-    TransferRequest::OpCode op_code) {
+    TransferRequest::OpCode op_code, const OperationOptions& options) {
     std::optional<TransferFuture> future;
     std::vector<TransferRequest> requests;
     for (size_t i = 0; i < replicas.size(); ++i) {
@@ -1088,7 +1089,12 @@ std::optional<TransferFuture> 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<TransferFuture> TransferSubmitter::submitMemcpyOperation(
 }
 
 std::optional<TransferFuture> TransferSubmitter::submitTransfer(
-    std::vector<TransferRequest>& requests) {
+    std::vector<TransferRequest>& 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<TransferFuture> TransferSubmitter::submitTransfer(
     }
 
     // Submit transfer
-    Status s = engine_.submitTransfer(batch_id, requests);
+    std::vector<ScheduledTransferRequest> 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<TransferFuture> TransferSubmitter::submitTransfer(
 
 std::optional<TransferFuture> TransferSubmitter::submitTransferEngineOperation(
     const AllocatedBuffer::Descriptor& handle, const std::vector<Slice>& 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,21 +1303,26 @@ std::optional<TransferFuture> 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<TransferFuture> TransferSubmitter::submitMemoryReadOperation(
     const AllocatedBuffer::Descriptor& handle, const std::vector<Slice>& slices,
-    uint64_t src_offset) {
+    uint64_t src_offset, const OperationOptions& options) {
     TransferStrategy strategy = selectStrategy(handle, slices);
 
     if (strategy == TransferStrategy::LOCAL_MEMCPY) {
         return submitMemcpyOperation(handle, slices, TransferRequest::READ,
                                      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: "

mooncake-transfer-engine/benchmark/scheduler/CMakeLists.txt

@@ -0,0 +1,17 @@
+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_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)