From 27973c7859b2d3b007db4db81498d94347e54819 Mon Sep 17 00:00:00 2001 From: Patrick Riel Date: Thu, 26 Jun 2025 05:21:34 +0000 Subject: [PATCH 1/2] add support in nixlbench for O_DIRECT --- benchmark/kvbench/commands/args.py | 35 +++++++- benchmark/kvbench/commands/nixlbench.py | 89 +++++++++++++++---- .../nixlbench/src/worker/nixl/nixl_worker.cpp | 71 +++++++++++---- 3 files changed, 157 insertions(+), 38 deletions(-) diff --git a/benchmark/kvbench/commands/args.py b/benchmark/kvbench/commands/args.py index ba75ff3768..1ea1ea1197 100644 --- a/benchmark/kvbench/commands/args.py +++ b/benchmark/kvbench/commands/args.py @@ -17,7 +17,7 @@ def add_common_args(subparser: argparse.ArgumentParser): - subparser.add_argument("--model", type=str, help="Model name (e.g., 'llama3.1-8b')") + subparser.add_argument("--model", type=str, help="Path to a model architecture YAML file") subparser.add_argument( "--model_config", type=str, help="Path to a single model config YAML file" ) @@ -146,15 +146,42 @@ def add_nixl_bench_args(subparser: argparse.ArgumentParser): ) subparser.add_argument( "--storage_enable_direct", - type=bool, + action="store_true", help="Enable direct I/O for storage operations (only used with POSIX backend)", - default=False, ) subparser.add_argument( "--gds_filepath", type=str, help="(File path for GDS operations" - ) + ), subparser.add_argument( "--enable_vmm", action="store_true", help="Enable VMM memory allocation when DRAM is requested", + ), + subparser.add_argument( + "--gds_batch_pool_size", + type=int, + default=32, + help="Batch pool size for GDS operations (default: 32, only used with GDS backend)" + ) + subparser.add_argument( + "--gds_batch_limit", + type=int, + default=128, + help="Batch limit for GDS operations (default: 128, only used with GDS backend)" + ) + subparser.add_argument( + "--posix_filepath", + type=str, + help="POSIX filepath (default: '')", + ) + subparser.add_argument( + "--posix_api_type", + type=str, + help="POSIX API type [AIO, O_DIRECT] (default: AIO)", + ) + subparser.add_argument( + "--num_files", + type=int, + help="Number of files (default: 1)", ) + \ No newline at end of file diff --git a/benchmark/kvbench/commands/nixlbench.py b/benchmark/kvbench/commands/nixlbench.py index 36107a924b..68d95f656a 100644 --- a/benchmark/kvbench/commands/nixlbench.py +++ b/benchmark/kvbench/commands/nixlbench.py @@ -37,16 +37,21 @@ def __init__( etcd_endpoints="http://localhost:2379", storage_enable_direct=False, gds_filepath="", + gds_batch_pool_size=32, + gds_batch_limit=128, initiator_seg_type="DRAM", enable_vmm=False, max_batch_size=None, max_block_size=None, mode="SG", + num_files=1, num_initiator_dev=1, num_iter=1000, num_target_dev=1, num_threads=1, op_type="WRITE", + posix_api_type="AIO", + posix_filepath="", runtime_type="ETCD", scheme="pairwise", start_batch_size=None, @@ -70,15 +75,20 @@ def __init__( storage_enable_direct (bool, optional): Whether to enable direct I/O for storage operations. Defaults to False. gds_filepath (str, optional): Path for GDS file. Defaults to "". enable_vmm (bool, optional): Whether to use VMM memory allocation. Defaults to False. + gds_batch_pool_size (int, optional): Batch pool size for GDS operations. Defaults to 32. + gds_batch_limit (int, optional): Batch limit for GDS operations. Defaults to 128. initiator_seg_type (str, optional): Type of initiator segment. Defaults to "DRAM". max_batch_size (int, optional): Maximum batch size for testing. Defaults to model_config value. max_block_size (int, optional): Maximum block size for testing. Defaults to tp_size * isl. mode (str, optional): Benchmarking mode. Defaults to "SG". + num_files (int, optional): Number of files. Defaults to 1. num_initiator_dev (int, optional): Number of initiator devices. Defaults to 1. num_iter (int, optional): Number of iterations. Defaults to 1000. num_target_dev (int, optional): Number of target devices. Defaults to 1. num_threads (int, optional): Number of threads. Defaults to 1. op_type (str, optional): Operation type. Defaults to "WRITE". + posix_api_type (str, optional): POSIX API type. Defaults to "AIO". + posix_filepath (str, optional): POSIX filepath. Defaults to "". runtime_type (str, optional): Runtime type. Defaults to "ETCD". scheme (str, optional): Communication scheme. Defaults to "pairwise". start_batch_size (int, optional): Starting batch size. Defaults to 1. @@ -98,15 +108,20 @@ def __init__( self.storage_enable_direct = storage_enable_direct self.gds_filepath = gds_filepath self.enable_vmm = enable_vmm + self.gds_batch_pool_size = gds_batch_pool_size + self.gds_batch_limit = gds_batch_limit self.initiator_seg_type = initiator_seg_type self.max_batch_size = max_batch_size self.max_block_size = max_block_size self.mode = mode + self.num_files = num_files self.num_initiator_dev = num_initiator_dev self.num_iter = num_iter self.num_target_dev = num_target_dev self.num_threads = num_threads self.op_type = op_type + self.posix_api_type = posix_api_type + self.posix_filepath = posix_filepath self.runtime_type = runtime_type self.scheme = scheme self.start_batch_size = start_batch_size @@ -121,25 +136,55 @@ def set_io_size(self, io_size: int): self.start_block_size = io_size self.max_block_size = io_size + + def _configure_gds(self, source: str, destination: str): + if source == "file": + # this is a READ from GDS to GPU + self.op_type = "READ" + self.target_seg_type = "VRAM" + elif source == "gpu": + # this is a WRITE from GPU to GDS + self.op_type = "WRITE" + self.target_seg_type = "VRAM" + else: + raise ValueError(f"Invalid source for GDS: {source}") + + def _configure_posix(self, source: str, destination: str): + if source == "file": + self.op_type = "READ" + self.target_seg_type = "DRAM" + elif source == "memory": + self.op_type = "WRITE" + self.initiator_seg_type = "DRAM" + else: + raise ValueError(f"Invalid source for POSIX: {source}") + def configure_segment_type(self, backend: str, source: str, destination: str): - if backend == "GDS": - if source == "file": - # this is a READ from GDS to GPU - self.op_type = "READ" - self.target_seg_type = "VRAM" - elif source == "gpu": - # this is a WRITE from GPU to GDS - self.op_type = "WRITE" - self.target_seg_type = "VRAM" - - elif source == "memory": - # this is a WRITE from memory to GDS - self.op_type = "WRITE" - self.initiator_seg_type = "DRAM" - self.target_seg_type = "DRAM" + if backend.lower() == "gds": + self._configure_gds(source, destination) + elif backend.lower() == "posix": + self._configure_posix(source, destination) else: raise ValueError(f"Invalid backend: {backend}") + # if backend == "GDS" or backend == "POSIX": + # if source == "file": + # # this is a READ from GDS to GPU + # self.op_type = "READ" + # self.target_seg_type = "VRAM" + # elif source == "gpu": + # # this is a WRITE from GPU to GDS + # self.op_type = "WRITE" + # self.target_seg_type = "VRAM" + + # elif source == "memory": + # # this is a WRITE from memory to GDS + # self.op_type = "WRITE" + # self.initiator_seg_type = "DRAM" + # self.target_seg_type = "DRAM" + # else: + # raise ValueError(f"Invalid backend: {backend}") + def configure_scheme(self, scheme: str = "pairwise", direction: str = "isl"): """ Configure the scheme based on the model configuration. @@ -154,6 +199,10 @@ def configure_scheme(self, scheme: str = "pairwise", direction: str = "isl"): self.num_target_dev = 1 def set_batch_size(self, batch_size: int): + """ + Set the batch size for benchmarking. + """ + self.start_batch_size = batch_size self.max_batch_size = batch_size @@ -186,15 +235,20 @@ def _params(self): "storage_enable_direct": self.storage_enable_direct, "gds_filepath": self.gds_filepath, "enable_vmm": self.enable_vmm, + "gds_batch_pool_size": self.gds_batch_pool_size, + "gds_batch_limit": self.gds_batch_limit, "initiator_seg_type": self.initiator_seg_type, "max_batch_size": self.max_batch_size, "max_block_size": self.max_block_size, "mode": self.mode, + "num_files": self.num_files, "num_initiator_dev": self.num_initiator_dev, "num_iter": self.num_iter, "num_target_dev": self.num_target_dev, "num_threads": self.num_threads, "op_type": self.op_type, + "posix_api_type": self.posix_api_type, + "posix_filepath": self.posix_filepath, "runtime_type": self.runtime_type, "scheme": self.scheme, "start_batch_size": self.start_batch_size, @@ -225,15 +279,20 @@ def defaults(): "storage_enable_direct": False, "gds_filepath": "", "enable_vmm": False, + "gds_batch_pool_size": 32, + "gds_batch_limit": 128, "initiator_seg_type": "DRAM", "max_batch_size": 1, # ios per gpu "max_block_size": 67108864, # io size "mode": "SG", + "num_files": 1, "num_initiator_dev": 1, "num_iter": 1000, "num_target_dev": 1, "num_threads": 1, "op_type": "WRITE", + "posix_api_type": "AIO", + "posix_filepath": "", "runtime_type": "ETCD", "scheme": "pairwise", "start_batch_size": 1, diff --git a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp index 3d1f1aeb47..e11306c408 100644 --- a/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp +++ b/benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp @@ -203,11 +203,24 @@ static void iovListToNixlXferDlist(const std::vector &iov_list, std::optional xferBenchNixlWorker::initBasicDescDram(size_t buffer_size, int mem_dev_id) { void *addr; - - addr = calloc(1, buffer_size); - if (!addr) { - std::cerr << "Failed to allocate " << buffer_size << " bytes of DRAM memory" << std::endl; - return std::nullopt; + if (xferBenchConfig::storage_enable_direct) { + long page_size = sysconf(_SC_PAGESIZE); + if (page_size == 0) { + std::cerr << "Error: Invalid page size returned by sysconf" << std::endl; + return std::nullopt; + } + int rc = posix_memalign(&addr, page_size, buffer_size); + if (rc != 0 || !addr) { + std::cerr << "Failed to allocate " << buffer_size << " bytes of page-aligned DRAM memory" << std::endl; + return std::nullopt; + } + memset(addr, 0, buffer_size); + } else { + addr = calloc(1, buffer_size); + if (!addr) { + std::cerr << "Failed to allocate " << buffer_size << " bytes of DRAM memory" << std::endl; + return std::nullopt; + } } if (isInitiator()) { @@ -362,14 +375,35 @@ static std::vector createFileFds(std::string name, bool is_gds) { std::optional xferBenchNixlWorker::initBasicDescFile(size_t buffer_size, int fd, int mem_dev_id) { auto ret = std::optional(std::in_place, (uintptr_t)gds_running_ptr, buffer_size, fd); // Fill up with data - void *buf = (void *)malloc(buffer_size); - if (!buf) { - std::cerr << "Failed to allocate " << buffer_size - << " bytes of memory" << std::endl; - return std::nullopt; + void *buf; + long page_size = sysconf(_SC_PAGESIZE); + if (page_size == 0) { + std::cerr << "Error: Invalid page size returned by sysconf" << std::endl; + exit(EXIT_FAILURE); + } + if (xferBenchConfig::storage_enable_direct) { + int rc = posix_memalign(&buf, (int)page_size, buffer_size); + if (rc != 0) { + std::cerr << "Error: " << strerror(rc) << std::endl; + std::cerr << "Failed to allocate " << buffer_size + << " bytes of memory" << std::endl; + return std::nullopt; + } + } else { + buf = (void *)malloc(buffer_size); + if (!buf) { + std::cerr << "Failed to allocate " << buffer_size + << " bytes of memory" << std::endl; + return std::nullopt; + } } // File is always initialized with XFERBENCH_TARGET_BUFFER_ELEMENT memset(buf, XFERBENCH_TARGET_BUFFER_ELEMENT, buffer_size); + if (xferBenchConfig::storage_enable_direct) { + gds_running_ptr = ((gds_running_ptr + page_size - 1) / page_size) * page_size; + } else { + gds_running_ptr += (buffer_size * mem_dev_id); + } int rc = pwrite(fd, buf, buffer_size, gds_running_ptr); if (rc < 0) { std::cerr << "Failed to write to file: " << fd @@ -378,8 +412,6 @@ std::optional xferBenchNixlWorker::initBasicDescFile(size_t buffer } free(buf); - gds_running_ptr += (buffer_size * mem_dev_id); - return ret; } @@ -420,7 +452,14 @@ std::vector> xferBenchNixlWorker::allocateMemory(int n num_devices = xferBenchConfig::num_target_dev; } buffer_size = xferBenchConfig::total_buffer_size / (num_devices * num_lists); - + if (xferBenchConfig::storage_enable_direct) { + long page_size = sysconf(_SC_PAGESIZE); + if (page_size == 0) { + std::cerr << "Error: Invalid page size returned by sysconf" << std::endl; + exit(EXIT_FAILURE); + } + buffer_size = ((buffer_size + page_size - 1) / page_size) * page_size; + } opt_args.backends.push_back(backend_engine); if (XFERBENCH_BACKEND_GDS == xferBenchConfig::backend || @@ -696,7 +735,6 @@ static int execTransfer(nixlAgent *agent, CHECK_NIXL_ERROR(agent->createXferReq(op, local_desc, remote_desc, target, req, ¶ms), "createTransferReq failed"); - for (int i = 0; i < num_iter && !error; i++) { rc = agent->postXferReq(req); if (NIXL_ERR_BACKEND == rc) { @@ -714,12 +752,7 @@ static int execTransfer(nixlAgent *agent, } while (NIXL_SUCCESS != rc); } } - agent->releaseXferReq(req); - if (error) { - std::cout << "NIXL releaseXferReq failed" << std::endl; - ret = -1; - } } return ret; From 7a73d9097281867268c29390474178e1b832d29e Mon Sep 17 00:00:00 2001 From: Patrick Riel Date: Thu, 26 Jun 2025 15:42:59 +0000 Subject: [PATCH 2/2] lint fix --- benchmark/kvbench/commands/args.py | 9 +++++---- benchmark/kvbench/commands/nixlbench.py | 1 - 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/benchmark/kvbench/commands/args.py b/benchmark/kvbench/commands/args.py index 1ea1ea1197..ff214eecd8 100644 --- a/benchmark/kvbench/commands/args.py +++ b/benchmark/kvbench/commands/args.py @@ -17,7 +17,9 @@ def add_common_args(subparser: argparse.ArgumentParser): - subparser.add_argument("--model", type=str, help="Path to a model architecture YAML file") + subparser.add_argument( + "--model", type=str, help="Path to a model architecture YAML file" + ) subparser.add_argument( "--model_config", type=str, help="Path to a single model config YAML file" ) @@ -161,13 +163,13 @@ def add_nixl_bench_args(subparser: argparse.ArgumentParser): "--gds_batch_pool_size", type=int, default=32, - help="Batch pool size for GDS operations (default: 32, only used with GDS backend)" + help="Batch pool size for GDS operations (default: 32, only used with GDS backend)", ) subparser.add_argument( "--gds_batch_limit", type=int, default=128, - help="Batch limit for GDS operations (default: 128, only used with GDS backend)" + help="Batch limit for GDS operations (default: 128, only used with GDS backend)", ) subparser.add_argument( "--posix_filepath", @@ -184,4 +186,3 @@ def add_nixl_bench_args(subparser: argparse.ArgumentParser): type=int, help="Number of files (default: 1)", ) - \ No newline at end of file diff --git a/benchmark/kvbench/commands/nixlbench.py b/benchmark/kvbench/commands/nixlbench.py index 68d95f656a..5f17ca8b08 100644 --- a/benchmark/kvbench/commands/nixlbench.py +++ b/benchmark/kvbench/commands/nixlbench.py @@ -136,7 +136,6 @@ def set_io_size(self, io_size: int): self.start_block_size = io_size self.max_block_size = io_size - def _configure_gds(self, source: str, destination: str): if source == "file": # this is a READ from GDS to GPU