diff --git a/mooncake-transfer-engine/benchmark/CMakeLists.txt b/mooncake-transfer-engine/benchmark/CMakeLists.txt index c99fe974..01f60652 100644 --- a/mooncake-transfer-engine/benchmark/CMakeLists.txt +++ b/mooncake-transfer-engine/benchmark/CMakeLists.txt @@ -1,5 +1,5 @@ find_package(CUDAToolkit) -if (CUDAToolkit_FOUND) +if(CUDAToolkit_FOUND) add_definitions(-DUSE_CUDA) include_directories(${CUDAToolkit_INCLUDE_DIRS}) message(STATUS "CUDA: Enabled") @@ -10,3 +10,7 @@ endif() file(GLOB TEBENCH_SOURCES "*.cpp") add_executable(tebench ${TEBENCH_SOURCES}) target_link_libraries(tebench PUBLIC transfer_engine tent) + +# Set RPATH for finding libasio.so at runtime +set_target_properties(tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE + INSTALL_RPATH "$ORIGIN/../lib") diff --git a/mooncake-transfer-engine/benchmark/main.cpp b/mooncake-transfer-engine/benchmark/main.cpp index 72f7bac9..bb3b792b 100644 --- a/mooncake-transfer-engine/benchmark/main.cpp +++ b/mooncake-transfer-engine/benchmark/main.cpp @@ -20,8 +20,8 @@ using namespace mooncake::tent; -void processBatchSizes(BenchRunner& runner, size_t block_size, - size_t batch_size, int num_threads) { +int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, + int num_threads) { bool mixed_opcode = false; OpCode opcode = READ; if (XferBenchConfig::check_consistency || XferBenchConfig::op_type == "mix") @@ -37,16 +37,16 @@ void processBatchSizes(BenchRunner& runner, size_t block_size, XferBenchStats stats; std::mutex mutex; - runner.runInitiatorTasks([&](int thread_id) -> int { + int rc = runner.runInitiatorTasks([&](int thread_id) -> int { runner.pinThread(thread_id); auto max_block_size = XferBenchConfig::max_block_size; auto max_batch_size = XferBenchConfig::max_batch_size; - uint64_t local_addr = - runner.getLocalBufferBase(XferBenchConfig::local_gpu_id + thread_id, - max_block_size, max_batch_size); + auto local_gpu_offset = std::max(0, XferBenchConfig::local_gpu_id); + auto target_gpu_offset = std::max(0, XferBenchConfig::target_gpu_id); + uint64_t local_addr = runner.getLocalBufferBase( + local_gpu_offset + thread_id, max_block_size, max_batch_size); uint64_t target_addr = runner.getTargetBufferBase( - XferBenchConfig::target_gpu_id + thread_id, max_block_size, - max_batch_size); + target_gpu_offset + thread_id, max_block_size, max_batch_size); XferBenchTimer timer; while (timer.lap_us(false) < 1000000ull) { @@ -89,7 +89,9 @@ void processBatchSizes(BenchRunner& runner, size_t block_size, return 0; }); + if (rc != 0) return -1; printStats(block_size, batch_size, stats, num_threads); + return 0; } int main(int argc, char* argv[]) { @@ -113,21 +115,25 @@ int main(int argc, char* argv[]) { return runner->runTarget(); } printStatsHeader(); + bool interrupted = false; for (int num_threads = XferBenchConfig::start_num_threads; - num_threads <= XferBenchConfig::max_num_threads; num_threads *= 2) { + !interrupted && num_threads <= XferBenchConfig::max_num_threads; + num_threads *= 2) { runner->startInitiator(num_threads); for (size_t block_size = XferBenchConfig::start_block_size; - block_size <= XferBenchConfig::max_block_size; block_size *= 2) { + !interrupted && block_size <= XferBenchConfig::max_block_size; + block_size *= 2) { for (size_t batch_size = XferBenchConfig::start_batch_size; - batch_size <= XferBenchConfig::max_batch_size; + !interrupted && batch_size <= XferBenchConfig::max_batch_size; batch_size *= 2) { if (block_size * batch_size * num_threads > XferBenchConfig::total_buffer_size) { LOG(INFO) << "Skipped for block_size " << block_size << " batch_size " << batch_size; } else { - processBatchSizes(*runner, block_size, batch_size, - num_threads); + if (processBatchSizes(*runner, block_size, batch_size, + num_threads) != 0) + interrupted = true; } } } diff --git a/mooncake-transfer-engine/benchmark/te_backend.cpp b/mooncake-transfer-engine/benchmark/te_backend.cpp index b67beb6c..af380690 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.cpp +++ b/mooncake-transfer-engine/benchmark/te_backend.cpp @@ -128,13 +128,22 @@ int TEBenchRunner::allocateBuffers() { } #ifdef USE_CUDA } else if (XferBenchConfig::seg_type == "VRAM") { - int num_buffers = 0; - cudaGetDeviceCount(&num_buffers); + int gpu_count = 0; + cudaGetDeviceCount(&gpu_count); + int start_gpu = 0, num_buffers = gpu_count; + if (XferBenchConfig::local_gpu_id != -1) { + start_gpu = XferBenchConfig::local_gpu_id; + num_buffers = 1; + LOG_ASSERT(start_gpu >= 0 && start_gpu < gpu_count) + << "local_gpu_id " << start_gpu << " out of range [0, " + << gpu_count << ")"; + } pinned_buffer_list_.resize(num_buffers, nullptr); for (int i = 0; i < num_buffers; ++i) { - auto location = "cuda:" + std::to_string(i); + int gpu_id = start_gpu + i; + auto location = "cuda:" + std::to_string(gpu_id); pinned_buffer_list_[i] = - allocateMemoryPool(total_buffer_size, i, true); + allocateMemoryPool(total_buffer_size, gpu_id, true); engine_->registerLocalMemory(pinned_buffer_list_[i], total_buffer_size, location); } @@ -252,8 +261,8 @@ int TEBenchRunner::runInitiatorTasks( current_task_[id] = func; pending_ = (int)threads_.size(); cv_task_.notify_all(); - cv_done_.wait(lk, [&] { return g_te_running && pending_ == 0; }); - return 0; + cv_done_.wait(lk, [&] { return !g_te_running || pending_ == 0; }); + return g_te_running ? 0 : -1; } double TEBenchRunner::runSingleTransfer(uint64_t local_addr, @@ -294,4 +303,4 @@ double TEBenchRunner::runSingleTransfer(uint64_t local_addr, } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 8dc924e0..99e4b95a 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.cpp +++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp @@ -44,6 +44,7 @@ std::shared_ptr loadConfig() { config->set("local_segment_name", XferBenchConfig::seg_name); config->set("metadata_type", XferBenchConfig::metadata_type); config->set("metadata_servers", XferBenchConfig::metadata_url_list); + config->set("rpc_server_port", XferBenchConfig::rpc_server_port); // Configure transport types based on xport_type parameter if (!XferBenchConfig::xport_type.empty()) { @@ -79,58 +80,64 @@ static TransportType getTransportType(const std::string& xport_type) { } int TENTBenchRunner::allocateBuffers() { - auto total_buffer_size = XferBenchConfig::total_buffer_size; - if (XferBenchConfig::seg_type == "DRAM") { - int num_buffers = numa_num_configured_nodes(); - pinned_buffer_list_.resize(num_buffers, nullptr); - auto start_ts = getCurrentTimeInNano(); - for (int i = 0; i < num_buffers; ++i) { - if (!XferBenchConfig::xport_type.empty()) { - MemoryOptions options; - options.type = getTransportType(XferBenchConfig::xport_type); - options.location = "cpu:" + std::to_string(i); - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, options)); - } else { - auto location = "cpu:" + std::to_string(i); - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, location)); - } - } - auto allocated_ts = getCurrentTimeInNano(); - std::vector buffers_size; - buffers_size.resize(pinned_buffer_list_.size(), total_buffer_size); - CHECK_FAIL( - engine_->registerLocalMemory(pinned_buffer_list_, buffers_size)); - auto registered_ts = getCurrentTimeInNano(); - LOG(INFO) << "Allocated " << total_buffer_size * num_buffers - << " bytes DRAM buffers in " - << (allocated_ts - start_ts) / 1e6 << " ms, registered in " - << (registered_ts - allocated_ts) / 1e6 << " ms"; + const auto total_buffer_size = XferBenchConfig::total_buffer_size; + const auto& seg_type = XferBenchConfig::seg_type; + const auto& xport_type = XferBenchConfig::xport_type; + + // Resolve device prefix, start index, and buffer count per seg_type + std::string device_prefix; + int start_idx = 0, num_buffers = 0; + + if (seg_type == "DRAM") { + device_prefix = "cpu"; + num_buffers = numa_num_configured_nodes(); #ifdef USE_CUDA - } else if (XferBenchConfig::seg_type == "VRAM") { - int num_buffers = 0; - cudaGetDeviceCount(&num_buffers); - pinned_buffer_list_.resize(num_buffers, nullptr); - for (int i = 0; i < num_buffers; ++i) { - if (!XferBenchConfig::xport_type.empty()) { - MemoryOptions options; - options.type = getTransportType(XferBenchConfig::xport_type); - options.location = "cuda:" + std::to_string(i); - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, options)); - } else { - auto location = "cuda:" + std::to_string(i); - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, location)); - } - CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_[i], - total_buffer_size)); + } else if (seg_type == "VRAM") { + device_prefix = "cuda"; + int gpu_count = 0; + cudaGetDeviceCount(&gpu_count); + start_idx = 0; + num_buffers = gpu_count; + if (XferBenchConfig::local_gpu_id != -1) { + start_idx = XferBenchConfig::local_gpu_id; + num_buffers = 1; + LOG_ASSERT(start_idx >= 0 && start_idx < gpu_count) + << "local_gpu_id " << start_idx << " out of range [0, " + << gpu_count << ")"; } #endif } else { - LOG(ERROR) << "Unknown seg_type: " << XferBenchConfig::seg_type; + LOG(ERROR) << "Unknown seg_type: " << seg_type; + return -1; } + + // Allocate + pinned_buffer_list_.resize(num_buffers, nullptr); + auto start_ts = getCurrentTimeInNano(); + for (int i = 0; i < num_buffers; ++i) { + auto location = device_prefix + ":" + std::to_string(start_idx + i); + if (!xport_type.empty()) { + MemoryOptions options; + options.type = getTransportType(xport_type); + options.location = std::move(location); + CHECK_FAIL(engine_->allocateLocalMemory( + &pinned_buffer_list_[i], total_buffer_size, options)); + } else { + CHECK_FAIL(engine_->allocateLocalMemory( + &pinned_buffer_list_[i], total_buffer_size, location)); + } + } + + // Register + auto allocated_ts = getCurrentTimeInNano(); + std::vector buffers_size(num_buffers, total_buffer_size); + CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_, buffers_size)); + auto registered_ts = getCurrentTimeInNano(); + + LOG(INFO) << "Allocated " << total_buffer_size * num_buffers << " bytes " + << seg_type << " buffers in " << (allocated_ts - start_ts) / 1e6 + << " ms, registered in " << (registered_ts - allocated_ts) / 1e6 + << " ms"; return 0; } @@ -254,8 +261,8 @@ int TENTBenchRunner::runInitiatorTasks( current_task_[id] = func; pending_ = (int)threads_.size(); cv_task_.notify_all(); - cv_done_.wait(lk, [&] { return g_tent_running && pending_ == 0; }); - return 0; + cv_done_.wait(lk, [&] { return !g_tent_running || pending_ == 0; }); + return g_tent_running ? 0 : -1; } double TENTBenchRunner::runSingleTransfer(uint64_t local_addr, @@ -297,4 +304,4 @@ double TENTBenchRunner::runSingleTransfer(uint64_t local_addr, } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp index 27f91ddc..2d7ca5e5 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/mooncake-transfer-engine/benchmark/utils.cpp @@ -35,12 +35,15 @@ DEFINE_int32(start_num_threads, 1, "Start number of concurrent worker threads."); DEFINE_int32(max_num_threads, 1, "Maximum number of concurrent worker threads."); -DEFINE_int32(local_gpu_id, 0, "Local GPU ID to be used"); -DEFINE_int32(target_gpu_id, 0, "Target GPU ID to be used"); +DEFINE_int32(local_gpu_id, 0, "Local GPU ID to be used, -1 for all GPUs"); +DEFINE_int32(target_gpu_id, 0, "Target GPU ID to be used, -1 for all GPUs"); DEFINE_string(metadata_type, "p2p", "Type of metadata service: p2p|etcd|redis|http"); DEFINE_string(metadata_url_list, "", "List of metadata service URLs, comma-separated."); +DEFINE_int32( + rpc_server_port, 0, + "RPC server port used for p2p metadata service (0 = auto-select)."); DEFINE_string(xport_type, "", "Transport type: rdma|shm|mnnvl|gds|iouring"); DEFINE_string(backend, "tent", "Transport backend: classic|tent"); DEFINE_bool(notifi, false, @@ -65,6 +68,7 @@ int XferBenchConfig::start_num_threads = 0; std::string XferBenchConfig::metadata_type; std::string XferBenchConfig::metadata_url_list; +int XferBenchConfig::rpc_server_port = 0; std::string XferBenchConfig::xport_type; std::string XferBenchConfig::backend; bool XferBenchConfig::notifi = false; @@ -90,6 +94,7 @@ void XferBenchConfig::loadFromFlags() { metadata_type = FLAGS_metadata_type; metadata_url_list = FLAGS_metadata_url_list; + rpc_server_port = FLAGS_rpc_server_port; xport_type = FLAGS_xport_type; backend = FLAGS_backend; @@ -157,4 +162,4 @@ void printStats(size_t block_size, size_t batch_size, XferBenchStats& stats, } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/utils.h b/mooncake-transfer-engine/benchmark/utils.h index bcb9657d..ae48d43b 100644 --- a/mooncake-transfer-engine/benchmark/utils.h +++ b/mooncake-transfer-engine/benchmark/utils.h @@ -65,6 +65,7 @@ struct XferBenchConfig { static std::string metadata_type; static std::string metadata_url_list; + static int rpc_server_port; static std::string xport_type; static std::string backend; static bool notifi; @@ -191,4 +192,4 @@ enum OpCode { READ, WRITE }; } // namespace tent } // namespace mooncake -#endif // XFER_UTILS_H \ No newline at end of file +#endif // XFER_UTILS_H diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 9581152c..f3d4becd 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -16,6 +16,7 @@ #define TENT_YLT_RPC_H #include +#include #include #include #include @@ -82,4 +83,4 @@ class CoroRpcAgent { } // namespace tent } // namespace mooncake -#endif // TENT_YLT_RPC_H \ No newline at end of file +#endif // TENT_YLT_RPC_H