NIXLBench: add support for specifying GDS batch and pool sizes (#394)

* add support for specifying gds batch and pool sizes
* update readme

Author: cheese-head

benchmark/kvbench/README.md

@@ -33,15 +33,15 @@ A utility for generating NIXL Bench commands that test KVCache transfer across v
 
 ### Docker
 ```bash
-git clone https://github.com/cheese-head/nixl.git
+git clone https://github.com/ai-dynamo/nixl.git
 export NIXL_SRC=/path/to/nixl/
 cd nixl/benchmark/nixlbench/contrib
 ./build.sh --nixl $NIXL_SRC
 ```
 
 ### Python
 ```bash
-git clone https://github.com/cheese-head/nixl.git
+git clone https://github.com/ai-dynamo/nixl.git
 cd nixl/benchmark/kvbench
 python3 -m venv venv
 source venv/bin/activate

benchmark/nixlbench/src/utils/utils.cpp

@@ -66,6 +66,8 @@ DEFINE_int32(num_target_dev, 1, "Number of device in target process");
 DEFINE_bool(enable_pt, false, "Enable Progress Thread (only used with nixl worker)");
 // GDS options - only used when backend is GDS
 DEFINE_string(gds_filepath, "", "File path for GDS operations (only used with GDS backend)");
+DEFINE_int32(gds_batch_pool_size, 32, "Batch pool size for GDS operations (default: 32, only used with GDS backend)");
+DEFINE_int32(gds_batch_limit, 128, "Batch limit for GDS operations (default: 128, only used with GDS backend)");
 
 // TODO: We should take rank wise device list as input to extend support
 // <rank>:<device_list>, ...
@@ -102,7 +104,8 @@ bool xferBenchConfig::enable_pt = false;
 std::string xferBenchConfig::device_list = "";
 std::string xferBenchConfig::etcd_endpoints = "";
 std::string xferBenchConfig::gds_filepath = "";
-
+int xferBenchConfig::gds_batch_pool_size = 0;
+int xferBenchConfig::gds_batch_limit = 0;
 std::vector<std::string> devices = { };
 int xferBenchConfig::num_files = 0;
 std::string xferBenchConfig::posix_api_type = "";
@@ -122,6 +125,8 @@ int xferBenchConfig::loadFromFlags() {
         // Load GDS-specific configurations if backend is GDS
         if (backend == XFERBENCH_BACKEND_GDS) {
             gds_filepath = FLAGS_gds_filepath;
+            gds_batch_pool_size = FLAGS_gds_batch_pool_size;
+            gds_batch_limit = FLAGS_gds_batch_limit;
             num_files = FLAGS_num_files;
             storage_enable_direct = FLAGS_storage_enable_direct;
         }
@@ -254,6 +259,10 @@ void xferBenchConfig::printConfig() {
         if (backend == XFERBENCH_BACKEND_GDS) {
             std::cout << std::left << std::setw(60) << "GDS filepath (--gds_filepath=path)" << ": "
                       << gds_filepath << std::endl;
+            std::cout << std::left << std::setw(60) << "GDS batch pool size (--gds_batch_pool_size=N)" << ": "
+                      << gds_batch_pool_size << std::endl;
+            std::cout << std::left << std::setw(60) << "GDS batch limit (--gds_batch_limit=N)" << ": "
+                      << gds_batch_limit << std::endl;
             std::cout << std::left << std::setw(60) << "GDS enable direct (--gds_enable_direct=[0,1])" << ": "
                       << storage_enable_direct << std::endl;
             std::cout << std::left << std::setw(60) << "Number of files (--num_files=N)" << ": "

benchmark/nixlbench/src/utils/utils.h

@@ -128,6 +128,8 @@ class xferBenchConfig {
         static std::string posix_api_type;
         static std::string posix_filepath;
         static bool storage_enable_direct;
+        static int gds_batch_pool_size;
+        static int gds_batch_limit;
 
         static int loadFromFlags();
         static void printConfig();

benchmark/nixlbench/src/worker/nixl/nixl_worker.cpp

@@ -140,6 +140,10 @@ xferBenchNixlWorker::xferBenchNixlWorker(int *argc, char ***argv, std::vector<st
     } else if (0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_GDS)) {
         // Using default param values for GDS backend
         std::cout << "GDS backend" << std::endl;
+        backend_params["batch_pool_size"] = std::to_string(xferBenchConfig::gds_batch_pool_size);
+        backend_params["batch_limit"] = std::to_string(xferBenchConfig::gds_batch_limit);
+        std::cout << "GDS batch pool size: " << xferBenchConfig::gds_batch_pool_size << std::endl;
+        std::cout << "GDS batch limit: " << xferBenchConfig::gds_batch_limit << std::endl;
     } else if (0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_POSIX)) {
         // Set API type parameter for POSIX backend
         if (xferBenchConfig::posix_api_type == XFERBENCH_POSIX_API_AIO) {
@@ -404,9 +408,10 @@ std::vector<std::vector<xferBenchIOV>> xferBenchNixlWorker::allocateMemory(int n
 
     if (XFERBENCH_BACKEND_GDS == xferBenchConfig::backend ||
         XFERBENCH_BACKEND_POSIX == xferBenchConfig::backend) {
-        remote_fds = createFileFds(getName(), true);
+        bool is_gds = XFERBENCH_BACKEND_GDS == xferBenchConfig::backend;
+        remote_fds = createFileFds(getName(), is_gds);
         if (remote_fds.empty()) {
-            std::cerr << "Failed to create GDS file" << std::endl;
+            std::cerr << "Failed to create " << ((is_gds) ? "GDS" : "POSIX") << " file" << std::endl;
             exit(EXIT_FAILURE);
         }
         for (int list_idx = 0; list_idx < num_lists; list_idx++) {