[TENT] Improve tebench: GPU selection, graceful interruption, and build fixes (#1537)
* build: add csignal header to rpc.h - Include csignal for signal handling in RPC module Signed-off-by: staryxchen <staryxchen@tencent.com> * build(benchmark): set RPATH for tebench executable - Configure tebench target to use relative RPATH for runtime library resolution - Ensure libasio.so can be found at runtime from the ../lib director Signed-off-by: staryxchen <staryxchen@tencent.com> * feat(benchmark): support graceful interruption during execution - Update runners to return error when stopped - Add early exit logic in main loop Signed-off-by: staryxchen <staryxchen@tencent.com> * feat(benchmark): support specifying single GPU ID or all GPUs - Update flags to accept -1 for all GPUs. - Adjust buffer allocation and address calculation for specific GPU selection. Signed-off-by: staryxchen <staryxchen@tencent.com> * refactor(benchmark): unify buffer allocation logic for DRAM and VRAM - Extract common logic for device prefix and buffer count - Merge allocation and registration loops to reduce code duplication - Improve variable const-correctness Signed-off-by: staryxchen <staryxchen@tencent.com> * feat(benchmark): add rpc server port configuration - Add`rpc_server_port`field to`XferBenchConfig` - Load port value from command line flags Signed-off-by: staryxchen <staryxchen@tencent.com> * Update mooncake-transfer-engine/benchmark/tent_backend.cpp Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * refactor(benchmark): simplify gpu id validation and offset calculation - Replace ternary operators with`std::max`for offset calculation - Update assertions to check for non-negative GPU IDs Signed-off-by: staryxchen <staryxchen@tencent.com> * retrigger CI Signed-off-by: staryxchen <staryxchen@tencent.com> --------- Signed-off-by: staryxchen <staryxchen@tencent.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Stary
GitHub
parent
dc432620cf
commit
1eec5a5b6b
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ std::shared_ptr<Config> 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<size_t> 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<size_t> 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
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -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
|
||||
} // namespace mooncake
|
||||
|
|
|
|||
|
|
@ -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
|
||||
#endif // XFER_UTILS_H
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@
|
|||
#define TENT_YLT_RPC_H
|
||||
|
||||
#include <atomic>
|
||||
#include <csignal>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
|
|
@ -82,4 +83,4 @@ class CoroRpcAgent {
|
|||
|
||||
} // namespace tent
|
||||
} // namespace mooncake
|
||||
#endif // TENT_YLT_RPC_H
|
||||
#endif // TENT_YLT_RPC_H
|
||||
|
|
|
|||
Loading…
Reference in New Issue