NIXL/UCX threadpool

benchmark/kvbench/README.md

@@ -137,6 +137,7 @@ These arguments are used by both `plan` and `profile` commands:
 | `--num_initiator_dev` | Number of devices in initiator processes (default: 1) |
 | `--num_target_dev` | Number of devices in target processes (default: 1) |
 | `--enable_pt` | Enable progress thread |
+| `--progress_threads` |  Number of progress threads (default: 0) |
 | `--device_list` | Comma-separated device names (default: all) |
 | `--runtime_type` | Type of runtime to use [ETCD] (default: ETCD) |
 | `--etcd-endpoints` | ETCD server URL for coordination (default: http://localhost:2379) |

benchmark/kvbench/commands/args.py

@@ -144,6 +144,9 @@ def nixl_bench_args(func):
     func = click.option("--enable_pt", is_flag=True, help="Enable progress thread")(
         func
     )
+    func = click.option(
+        "--progress_threads", type=int, help="Number of progress threads (default: 0)"
+    )(func)
     func = click.option(
         "--device_list", type=str, help="Comma-separated device names (default: all)"
     )(func)

benchmark/kvbench/commands/nixlbench.py

@@ -34,6 +34,7 @@ def __init__(
         check_consistency=False,
         device_list="all",
         enable_pt=False,
+        progress_threads=0,
         etcd_endpoints="http://localhost:2379",
         storage_enable_direct=False,
         filepath="",
@@ -72,6 +73,7 @@ def __init__(
             check_consistency (bool, optional): Whether to check consistency. Defaults to False.
             device_list (str, optional): List of devices to use. Defaults to "all".
             enable_pt (bool, optional): Whether to enable peer-to-peer transfer. Defaults to False.
+            progress_threads (int, optional): Number of progress threads (default: 0).
             etcd_endpoints (str, optional): ETCD endpoints for runtime. Defaults to "http://localhost:2379".
             storage_enable_direct (bool, optional): Whether to enable direct I/O for storage operations. Defaults to False.
             filepath (str, optional): Path for GDS and POSIX operations. Defaults to "".
@@ -105,6 +107,7 @@ def __init__(
         self.check_consistency = check_consistency
         self.device_list = device_list
         self.enable_pt = enable_pt
+        self.progress_threads = progress_threads
         self.etcd_endpoints = etcd_endpoints
         self.storage_enable_direct = storage_enable_direct
         self.filepath = filepath
@@ -251,6 +254,7 @@ def _params(self):
             "check_consistency": self.check_consistency,
             "device_list": self.device_list,
             "enable_pt": self.enable_pt,
+            "progress_threads": self.progress_threads,
             "etcd_endpoints": self.etcd_endpoints,
             "storage_enable_direct": self.storage_enable_direct,
             "filepath": self.filepath,
@@ -294,6 +298,7 @@ def defaults():
             "check_consistency": False,
             "device_list": "all",
             "enable_pt": False,
+            "progress_threads": 0,
             "etcd_endpoints": "http://localhost:2379",
             "storage_enable_direct": False,
             "filepath": "",

benchmark/nixlbench/README.md

@@ -94,6 +94,7 @@ meson configure /path/to/build/dir
 --num_initiator_dev NUM    # Number of devices in initiator processes (default: 1)
 --num_target_dev NUM       # Number of devices in target processes (default: 1)
 --enable_pt                # Enable progress thread
+--progress_threads         # Number of progress threads (default: 0)
 --device_list LIST         # Comma-separated device names (default: all)
 --runtime_type NAME        # Type of runtime to use [ETCD] (default: ETCD)
 --etcd-endpoints URL       # ETCD server URL for coordination (default: http://localhost:2379)

benchmark/nixlbench/src/utils/utils.cpp

@@ -79,6 +79,7 @@ DEFINE_int32 (
 DEFINE_int32(num_initiator_dev, 1, "Number of device in initiator process");
 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)");
+DEFINE_uint64(progress_threads, 0, "Number of progress threads (default: 0)");
 DEFINE_bool(enable_vmm, false, "Enable VMM memory allocation when DRAM is requested");
 
 // Storage backend(GDS, GDS_MT, POSIX, HF3FS, OBJ) options
@@ -141,6 +142,7 @@ int xferBenchConfig::large_blk_iter_ftr = 16;
 int xferBenchConfig::warmup_iter = 0;
 int xferBenchConfig::num_threads = 0;
 bool xferBenchConfig::enable_pt = false;
+size_t xferBenchConfig::progress_threads = 0;
 bool xferBenchConfig::enable_vmm = false;
 std::string xferBenchConfig::device_list = "";
 std::string xferBenchConfig::etcd_endpoints = "";
@@ -175,6 +177,7 @@ xferBenchConfig::loadFromFlags() {
     if (worker_type == XFERBENCH_WORKER_NIXL) {
         backend = FLAGS_backend;
         enable_pt = FLAGS_enable_pt;
+        progress_threads = FLAGS_progress_threads;
         device_list = FLAGS_device_list;
         enable_vmm = FLAGS_enable_vmm;
 
@@ -371,6 +374,7 @@ xferBenchConfig::printConfig() {
         printOption("Backend (--backend=[UCX,UCX_MO,GDS,GDS_MT,POSIX,Mooncake,HF3FS,OBJ])",
                     backend);
         printOption ("Enable pt (--enable_pt=[0,1])", std::to_string (enable_pt));
+        printOption("Progress threads (--progress_threads=N)", std::to_string(progress_threads));
         printOption ("Device list (--device_list=dev1,dev2,...)", device_list);
         printOption ("Enable VMM (--enable_vmm=[0,1])", std::to_string (enable_vmm));
 

benchmark/nixlbench/src/utils/utils.h

@@ -143,6 +143,7 @@ class xferBenchConfig {
         static int warmup_iter;
         static int num_threads;
         static bool enable_pt;
+        static size_t progress_threads;
         static std::string device_list;
         static std::string etcd_endpoints;
         static std::string benchmark_group;

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

@@ -111,6 +111,8 @@ xferBenchNixlWorker::xferBenchNixlWorker(int *argc, char ***argv, std::vector<st
 
     if (0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_UCX) ||
         0 == xferBenchConfig::backend.compare(XFERBENCH_BACKEND_UCX_MO)) {
+        backend_params["num_threads"] = std::to_string(xferBenchConfig::progress_threads);
+
         // No need to set device_list if all is specified
         // fallback to backend preference
         if (devices[0] != "all" && devices.size() >= 1) {

src/api/python/_api.py

@@ -51,13 +51,15 @@ def __init__(
         enable_prog_thread: bool = True,
         enable_listen_thread: bool = False,
         listen_port: int = 0,
+        num_threads: int = 0,
         backends: list[str] = ["UCX"],
     ):
         # TODO: add backend init parameters
         self.backends = backends
         self.enable_pthread = enable_prog_thread
         self.enable_listen = enable_listen_thread
         self.port = listen_port
+        self.num_threads = num_threads
 
 
 """
@@ -119,22 +121,24 @@ def __init__(
             self.plugin_b_options[plugin] = backend_options
             self.plugin_mem_types[plugin] = mem_types
 
-        # TODO: populate init from default parameters, or define a set of params in python
-        init: dict[str, str] = {}
-
         if instantiate_all:
-            for plugin in self.plugin_list:
-                self.create_backend(plugin, init)
-        else:
-            for bknd in nixl_conf.backends:
-                # TODO: populate init from nixl_conf when added
-                if bknd not in self.plugin_list:
-                    logger.warning(
-                        "Skipping backend registration %s due to the missing plugin.",
-                        bknd,
-                    )
-                else:
-                    self.create_backend(bknd, init)
+            nixl_conf.backends = self.plugin_list
+
+        for bknd in nixl_conf.backends:
+            if bknd not in self.plugin_list:
+                logger.warning(
+                    "Skipping backend registration %s due to the missing plugin.",
+                    bknd,
+                )
+            else:
+                # TODO: improve population of init from nixl_conf
+                init: dict[str, str] = {}
+                if nixl_conf.num_threads > 0:
+                    if bknd == "UCX" or bknd == "OBJ":
+                        init["num_threads"] = str(nixl_conf.num_threads)
+                    elif bknd == "GDS_MT":
+                        init["thread_count"] = str(nixl_conf.num_threads)
+                self.create_backend(bknd, init)
 
         self.nixl_mems = {
             "DRAM": nixlBind.DRAM_SEG,

src/core/meson.build

@@ -22,7 +22,7 @@ if etcd_dep.found()
 endif
 
 if 'UCX' in static_plugins
-    nixl_lib_deps += [ ucx_backend_interface, cuda_dep ]
+    nixl_lib_deps += [ ucx_backend_interface, asio_dep, cuda_dep ]
 endif
 
 if 'UCX_MO' in static_plugins

src/plugins/ucx/meson.build

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 ucx_utils_dep = declare_dependency(link_with: ucx_utils_lib, include_directories: utils_inc_dirs )
+asio_dep = [dependency('asio', required: true)]
 
 compile_flags = []
 if cuda_dep.found()
@@ -23,15 +24,15 @@ endif
 if 'UCX' in static_plugins
     ucx_backend_lib = static_library('UCX',
                'ucx_backend.cpp', 'ucx_backend.h', 'ucx_plugin.cpp',
-               dependencies: [nixl_infra, ucx_utils_dep, serdes_interface, cuda_dep, ucx_dep, thread_dep, nixl_common_dep],
+               dependencies: [nixl_infra, ucx_utils_dep, serdes_interface, cuda_dep, ucx_dep, thread_dep, nixl_common_dep, asio_dep],
                include_directories: nixl_inc_dirs,
                install: false,
                cpp_args : compile_flags,
                name_prefix: 'libplugin_')  # Custom prefix for plugin libraries
 else
     ucx_backend_lib = shared_library('UCX',
                'ucx_backend.cpp', 'ucx_backend.h', 'ucx_plugin.cpp',
-               dependencies: [nixl_infra, ucx_utils_dep, serdes_interface, cuda_dep, ucx_dep, thread_dep, nixl_common_dep],
+               dependencies: [nixl_infra, ucx_utils_dep, serdes_interface, cuda_dep, ucx_dep, thread_dep, nixl_common_dep, asio_dep],
                include_directories: nixl_inc_dirs,
                install: true,
                cpp_args : compile_flags + ['-fPIC'],