add support in nixlbench for O_DIRECT

benchmark/kvbench/commands/args.py

@@ -17,7 +17,9 @@
 
 
 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 +148,41 @@ 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)",
     )

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,54 @@ 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 +198,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 +234,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 +278,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,

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

@@ -203,11 +203,24 @@ static void iovListToNixlXferDlist(const std::vector<xferBenchIOV> &iov_list,
 
 std::optional<xferBenchIOV> 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<int> createFileFds(std::string name, bool is_gds) {
 std::optional<xferBenchIOV> xferBenchNixlWorker::initBasicDescFile(size_t buffer_size, int fd, int mem_dev_id) {
     auto ret = std::optional<xferBenchIOV>(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<xferBenchIOV> xferBenchNixlWorker::initBasicDescFile(size_t buffer
     }
     free(buf);
 
-    gds_running_ptr += (buffer_size * mem_dev_id);
-
     return ret;
 }
 
@@ -420,7 +452,14 @@ std::vector<std::vector<xferBenchIOV>> 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, &params), "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;