From c3979417dcc3f4ff6cab9851ce42102c56979922 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sat, 15 Mar 2025 23:13:13 -0700 Subject: [PATCH 01/32] pybind: exception throwing Create NIXL exception classes inside the python bindings that are mapped to python exceptions we can throw. Signed-off-by: Timothy Stamler --- src/bindings/python/nixl_bindings.cpp | 123 +++++++++++++++++++++++--- 1 file changed, 113 insertions(+), 10 deletions(-) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index fb661b21f3..d40e5c36a9 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -26,6 +26,84 @@ namespace py = pybind11; +class nixlNotPostedError : public std::runtime_error { + public: + nixlNotPostedError(const char* what) : runtime_error(what) {} +}; + +class nixlInvalidParamError : public std::runtime_error { + public: + nixlInvalidParamError(const char* what) : runtime_error(what) {} +}; + +class nixlBackendError : public std::runtime_error { + public: + nixlBackendError(const char* what) : runtime_error(what) {} +}; + + +class nixlNotFoundError : public std::runtime_error { + public: + nixlNotFoundError(const char* what) : runtime_error(what) {} +}; + + +class nixlMismatchError : public std::runtime_error { + public: + nixlMismatchError(const char* what) : runtime_error(what) {} +}; + + +class nixlNotAllowedError : public std::runtime_error { + public: + nixlNotAllowedError(const char* what) : runtime_error(what) {} +}; + + +class nixlRepostActiveError : public std::runtime_error { + public: + nixlRepostActiveError(const char* what) : runtime_error(what) {} +}; + + +class nixlUnknownError : public std::runtime_error { + public: + nixlUnknownError(const char* what) : runtime_error(what) {} +}; + +void throw_nixl_exception(const nixl_status_t &status) { + switch (status) { + case NIXL_IN_PROG: return; //not an error + case NIXL_SUCCESS: return; //not an error + case NIXL_ERR_NOT_POSTED: + throw nixlNotPostedError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_INVALID_PARAM: + throw nixlInvalidParamError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_BACKEND: + throw nixlBackendError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_NOT_FOUND: + throw nixlNotFoundError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_MISMATCH: + throw nixlMismatchError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_NOT_ALLOWED: + throw nixlNotAllowedError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_REPOST_ACTIVE: + throw nixlRepostActiveError(nixlEnumStrings::statusStr(status).c_str()); + break; + case NIXL_ERR_UNKNOWN: + throw nixlUnknownError(nixlEnumStrings::statusStr(status).c_str()); + break; + default: + throw std::runtime_error("BAD_STATUS"); + } +} + PYBIND11_MODULE(_bindings, m) { //TODO: each nixl class and/or function can be documented in place @@ -58,6 +136,15 @@ PYBIND11_MODULE(_bindings, m) { .value("NIXL_ERR_UNKNOWN", NIXL_ERR_UNKNOWN) .export_values(); + py::register_exception(m, "nixlNotPostedError"); + py::register_exception(m, "nixlInvalidParamError"); + py::register_exception(m, "nixlBackendError"); + py::register_exception(m, "nixlNotFoundError"); + py::register_exception(m, "nixlMismatchError"); + py::register_exception(m, "nixlNotAllowedError"); + py::register_exception(m, "nixlRepostActiveError"); + py::register_exception(m, "nixlUnknownError"); + py::class_(m, "nixlXferDList") .def(py::init(), py::arg("type"), py::arg("unifiedAddr")=true, py::arg("sorted")=false, py::arg("init_size")=0) .def(py::init([](nixl_mem_t mem, std::vector descs, bool unifiedAddr, bool sorted) { @@ -179,7 +266,7 @@ PYBIND11_MODULE(_bindings, m) { nixl_b_params_t params; nixl_mem_list_t mems; nixl_status_t ret = agent.getPluginParams(type, mems, params); - if(ret < 0); //throw exception + if(ret < 0) throw_nixl_exception(ret); // TODO merge the mems return params; }) @@ -187,14 +274,17 @@ PYBIND11_MODULE(_bindings, m) { nixl_b_params_t params; nixl_mem_list_t mems; nixl_status_t ret = agent.getBackendParams((nixlBackendH*) backend, mems, params); - if(ret < 0); //throw exception + if(ret < 0) throw_nixl_exception(ret); // TODO merge the mems return params; }) .def("createBackend", [](nixlAgent &agent, const nixl_backend_t &type, const nixl_b_params_t &initParams) -> uintptr_t { nixlBackendH* backend; nixl_status_t ret = agent.createBackend(type, initParams, backend); - if(ret < 0) return (uintptr_t) nullptr; //throw exception + if(ret < 0){ + throw_nixl_exception(ret); + return (uintptr_t) nullptr; + } return (uintptr_t) backend; }) .def("registerMem", [](nixlAgent &agent, nixl_reg_dlist_t descs, uintptr_t backend) -> nixl_status_t { @@ -225,8 +315,12 @@ PYBIND11_MODULE(_bindings, m) { extra_params.hasNotif = true; } nixl_status_t ret = agent.createXferReq(operation, local_descs, remote_descs, remote_agent, handle, &extra_params); - if (ret != NIXL_SUCCESS) return (uintptr_t) nullptr; - else return (uintptr_t) handle; + + if(ret < 0){ + throw_nixl_exception(ret); + return (uintptr_t) nullptr; + } + return (uintptr_t) handle; }, py::arg("local_descs"), py::arg("remote_descs"), py::arg("remote_agent"), py::arg("notif_msg"), py::arg("operation"), @@ -234,7 +328,10 @@ PYBIND11_MODULE(_bindings, m) { .def("queryXferBackend", [](nixlAgent &agent, uintptr_t reqh) -> uintptr_t { nixlBackendH* handle; nixl_status_t ret = agent.queryXferBackend((nixlXferReqH*) reqh, handle); - if(ret < 0) return (uintptr_t) nullptr; + if(ret < 0){ + throw_nixl_exception(ret); + return (uintptr_t) nullptr; + } return (uintptr_t) handle; }) .def("prepXferDlist", [](nixlAgent &agent, @@ -245,8 +342,11 @@ PYBIND11_MODULE(_bindings, m) { nixl_opt_args_t extra_params; extra_params.backends.push_back((nixlBackendH*) backend); nixl_status_t ret = agent.prepXferDlist(descs, remote_agent, handle, &extra_params); - if (ret != NIXL_SUCCESS) return (uintptr_t) nullptr; - else return (uintptr_t) handle; + if(ret < 0){ + throw_nixl_exception(ret); + return (uintptr_t) nullptr; + } + return (uintptr_t) handle; }) .def("makeXferReq", [](nixlAgent &agent, uintptr_t local_side, @@ -265,8 +365,11 @@ PYBIND11_MODULE(_bindings, m) { (nixlDlistH*) local_side, local_indices, (nixlDlistH*) remote_side, remote_indices, handle, &extra_params); - if (ret != NIXL_SUCCESS) return (uintptr_t) nullptr; - else return (uintptr_t) handle; + if(ret < 0){ + throw_nixl_exception(ret); + return (uintptr_t) nullptr; + } + return (uintptr_t) handle; }) .def("releaseXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { return agent.releaseXferReq((nixlXferReqH*) reqh); From fa0c921045d72618ec88d0cfb296b61bce4efbef Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 01:01:41 -0700 Subject: [PATCH 02/32] python: initial backend params API Signed-off-by: Timothy Stamler --- src/api/python/_api.py | 36 +++++++++++++++++++++++++-- src/bindings/python/nixl_bindings.cpp | 12 ++++----- test/python/nixl_wrapper_test.py | 12 +++++++++ 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 77beb6d205..63e8226fd6 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -25,13 +25,27 @@ def __init__(self, agent_name, nixl_config): # Read available backends and device info from nixl_config # For now setting the multithreading to enabled. devices = nixlBind.nixlAgentConfig(True) - init = {} + # init = {} self.name = agent_name self.notifs = {} self.backends = {} self.agent = nixlBind.nixlAgent(agent_name, devices) - self.backends["UCX"] = self.agent.createBackend("UCX", init) + + self.plugin_list = nixlBind.getAvailPlugins() + + self.backend_option_map = {} + self.mem_type_map = {} + + for plugin in self.plugin_list: + (backend_options, mem_types) = self.agent.getPluginParams(plugin) + self.backend_option_map[plugin] = backend_options + self.mem_type_map[plugin] = mem_types + + # TODO: make explicit call later + # self.backends["UCX"] = self.agent.createBackend("UCX", init) + if len(self.plugin_list) == 0: + print("No plugins available, cannot start transfers!") self.nixl_mems = { "DRAM": nixlBind.DRAM_SEG, @@ -46,6 +60,24 @@ def __init__(self, agent_name, nixl_config): print("Initializied NIXL agent:", agent_name) + def get_plugin_list(self): + return self.plugin_list + + def get_backend_mem_types(self, backend): + return self.mem_types[backend] + + def get_backend_params(self, backend): + return self.backend_options_map[backend] + + def create_backend(self, backend, initParams=None): + self.backends[backend] = self.agent.createBackend(backend, initParams) + + (backend_options, mem_types) = self.agent.getBackendParams( + self.backends[backend] + ) + self.backend_option_map[backend] = backend_options + self.mem_type_map[backend] = mem_types + def get_xfer_descs( self, descs, mem_type=None, is_unified_addr=True, is_sorted=False ): diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index d40e5c36a9..0695bff40a 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -262,26 +262,24 @@ PYBIND11_MODULE(_bindings, m) { py::class_(m, "nixlAgent") .def(py::init()) .def("getAvailPlugins", &nixlAgent::getAvailPlugins) - .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> nixl_b_params_t { + .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> std::pair { nixl_b_params_t params; nixl_mem_list_t mems; nixl_status_t ret = agent.getPluginParams(type, mems, params); if(ret < 0) throw_nixl_exception(ret); - // TODO merge the mems - return params; + return std::make_pair(params, mems); }) - .def("getBackendParams", [](nixlAgent &agent, uintptr_t backend) -> nixl_b_params_t { + .def("getBackendParams", [](nixlAgent &agent, uintptr_t backend) -> std::pair { nixl_b_params_t params; nixl_mem_list_t mems; nixl_status_t ret = agent.getBackendParams((nixlBackendH*) backend, mems, params); if(ret < 0) throw_nixl_exception(ret); - // TODO merge the mems - return params; + return std::make_pair(params, mems); }) .def("createBackend", [](nixlAgent &agent, const nixl_backend_t &type, const nixl_b_params_t &initParams) -> uintptr_t { nixlBackendH* backend; nixl_status_t ret = agent.createBackend(type, initParams, backend); - if(ret < 0){ + if(ret < 0) { throw_nixl_exception(ret); return (uintptr_t) nullptr; } diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index 1a64da9de5..e93ada1126 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -24,6 +24,18 @@ buf_size = 256 # Allocate memory and register with NIXL nixl_agent1 = nixl_agent("target", None) + + plugin_list = nixl_agent1.get_plugin_list("UCX") + assert "UCX" in plugin_list + print(plugin_list) + + print(nixl_agent1.get_backend_mem_types("UCX")) + print(nixl_agent1.get_backend_params("UCX")) + + nixl_agent1.create_backend("UCX") + + print(nixl_agent1.get_backend_params("UCX")) + addr1 = nixl_utils.malloc_passthru(buf_size * 2) addr2 = addr1 + buf_size From 3c306f9c5be93d7c7ddccda8192e453b77932886 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 01:20:52 -0700 Subject: [PATCH 03/32] Python API: Added initial support for GDS backend to the wrapper * Also added appened method to the pybinds and fixed formatting issues. Signed-off-by: Timothy Stamler Signed-off-by: Moein Khazraee --- src/api/python/_api.py | 64 +++++++++++++++++++++++++++++++++--------- 1 file changed, 51 insertions(+), 13 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 63e8226fd6..3d1b3ba216 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -19,9 +19,12 @@ import nixl._bindings as nixlBind +class nixl_config: + def __init(self, backends=["UCX", "GDS"]): + self.backends = backends class nixl_agent: - def __init__(self, agent_name, nixl_config): + def __init__(self, agent_name, nixl_config=None): # Read available backends and device info from nixl_config # For now setting the multithreading to enabled. devices = nixlBind.nixlAgentConfig(True) @@ -47,6 +50,13 @@ def __init__(self, agent_name, nixl_config): if len(self.plugin_list) == 0: print("No plugins available, cannot start transfers!") + if nixl_config: + for x in nixl_config.backends: + self.backends[x] = self.agent.createBackend(x, init) + else: # Defaulting to UCX and GDS for now + self.backends["UCX"] = self.agent.createBackend("UCX", init) + self.backends["GDS"] = self.agent.createBackend("GDS", init) + self.nixl_mems = { "DRAM": nixlBind.DRAM_SEG, "VRAM": nixlBind.VRAM_SEG, @@ -170,19 +180,31 @@ def deserialize_descs(self, serialized_descs): # The returned descriptor object can be used for call to deregister def register_memory( - self, reg_list, mem_type=None, is_unified_addr=True, is_sorted=False + self, reg_list, mem_type=None, is_unified_addr=True, is_sorted=False, backend=None ): - # based on backend type and mem_type, figure what registrations are meaningful reg_descs = self.get_reg_descs(reg_list, mem_type, is_unified_addr, is_sorted) - ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) + # based on backend type and mem_type, figure what registrations are meaningful + if backend: + ret = self.agent.registerMem(reg_descs, self.backends[backend]) + else: + if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) + else: + ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) if ret != 0: return None return reg_descs - def deregister_memory(self, dereg_descs): + def deregister_memory(self, dereg_descs, backend=None): # based on backend type and mem_type, figure what deregistrations are needed - self.agent.deregisterMem(dereg_descs, self.backends["UCX"]) + if backend: + self.agent.deregisterMem(dereg_descs, self.backends[backend]) + else: + if (dereg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + self.agent.deregisterMem(dereg_descs, self.backends["GDS"]) + else: + self.agent.deregisterMem(dereg_descs, self.backends["UCX"]) # No return # Optional proactive make connection @@ -200,14 +222,24 @@ def remove_remote_agent(self, agent): self.agent.invalidateRemoteMD(agent) def initialize_xfer( - self, local_descs, remote_descs, remote_agent, notif_msg, operation + self, local_descs, remote_descs, remote_agent, notif_msg, operation, xfer_backend=None ): op = self.nixl_ops[operation] if op: - handle = self.agent.createXferReq( - local_descs, remote_descs, remote_agent, notif_msg, op - ) - return handle # In case of error it will be 0 + if xfer_backend: + handle = self.agent.createXferReq( + local_descs, + remote_descs, + remote_agent, + notif_msg, + op, + xfer_backend, + ) + else: + handle = self.agent.createXferReq( + local_descs, remote_descs, remote_agent, notif_msg, op + ) + return handle # In case of error it will be None else: return None @@ -229,8 +261,14 @@ def prep_xfer_side( backend = self.agent.getXferBackend(example_xfer) handle = self.agent.prepXferDlist(descs, remote_agent, backend) else: - # Or use same logic that we used in register_memory - handle = self.agent.prepXferDlist(descs, remote_agent, self.backends["UCX"]) + if (descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + handle = self.agent.prepXferDlist( + descs, remote_agent, self.backends["GDS"] + ) + else: + handle = self.agent.prepXferDlist( + descs, remote_agent, self.backends["UCX"] + ) if handle == 0: return None From 2435abbda5b3f96931fb95a073afe15679b34abe Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 02:02:11 -0700 Subject: [PATCH 04/32] Added more functioanlity to the nixl_config and separated methods to get backend info Signed-off-by: Moein Khazraee --- src/api/python/_api.py | 81 +++++++++++++++++------ src/bindings/python/nixl_bindings.cpp | 4 +- test/python/blocking_send_recv_example.py | 2 +- test/python/nixl_wrapper_test.py | 8 +-- 4 files changed, 68 insertions(+), 27 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 3d1b3ba216..db1361ff9b 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -20,42 +20,57 @@ import nixl._bindings as nixlBind class nixl_config: - def __init(self, backends=["UCX", "GDS"]): + def __init(self, enable_prog_thread = True, backends=["UCX", "GDS"]): + # TODO: add backend init parameters self.backends = backends + self.enable_pthread = enable_prog_thread class nixl_agent: def __init__(self, agent_name, nixl_config=None): - # Read available backends and device info from nixl_config - # For now setting the multithreading to enabled. - devices = nixlBind.nixlAgentConfig(True) - # init = {} + # Set agent config and instantiate an agent + if nixl_config: + agent_config = nixlBind.nixlAgentConfig(nixl_config.enable_pthread) + else: + agent_config = nixlBind.nixlAgentConfig(True) + self.agent = nixlBind.nixlAgent(agent_name, agent_config) self.name = agent_name self.notifs = {} self.backends = {} - self.agent = nixlBind.nixlAgent(agent_name, devices) + self.backend_mems = {} + self.backend_options = {} self.plugin_list = nixlBind.getAvailPlugins() + if len(self.plugin_list) == 0: + print("No plugins available, cannot start transfers!") - self.backend_option_map = {} - self.mem_type_map = {} - + self.plugin_b_options = {} + self.plugin_mem_types = {} for plugin in self.plugin_list: (backend_options, mem_types) = self.agent.getPluginParams(plugin) - self.backend_option_map[plugin] = backend_options - self.mem_type_map[plugin] = mem_types + self.plugin_b_options[plugin] = backend_options + self.plugin_mem_types[plugin] = mem_types - # TODO: make explicit call later # self.backends["UCX"] = self.agent.createBackend("UCX", init) - if len(self.plugin_list) == 0: - print("No plugins available, cannot start transfers!") + init = {} if nixl_config: for x in nixl_config.backends: - self.backends[x] = self.agent.createBackend(x, init) + # TODO: populate init from nixl_config when added + if x not in self.plugin_list: + print("Skiping backend registration", x, "due to the missing plugin.") + else: + self.backends[x] = self.agent.createBackend(x, init) else: # Defaulting to UCX and GDS for now - self.backends["UCX"] = self.agent.createBackend("UCX", init) - self.backends["GDS"] = self.agent.createBackend("GDS", init) + if "UCX" in self.plugin_list: + self.backends["UCX"] = self.agent.createBackend("UCX", init) + if "GDS" in self.plugin_list: + self.backends["GDS"] = self.agent.createBackend("GDS", init) + + for backend in self.backends: + (backend_options, mem_types) = self.agent.getBackendParams(backend) + self.backend_mems[backend] = mem_types + self.backend_options[backend] = backend_options self.nixl_mems = { "DRAM": nixlBind.DRAM_SEG, @@ -73,11 +88,33 @@ def __init__(self, agent_name, nixl_config=None): def get_plugin_list(self): return self.plugin_list + def get_plugin_mem_types(self, backend): + if backend in self.plugin_mem_types: + return self.plugin_mem_types[backend] + else: + print ("Plugin", backend, "is not available to get its supported mem types.") + return None + + def get_plugin_params(self, backend): + if backend in self.plugin_b_options: + return self.plugin_b_options[backend] + else: + print ("Plugin", backend, "is not available to get its parameters.") + return None + def get_backend_mem_types(self, backend): - return self.mem_types[backend] + if backend in self.backend_mems: + return self.backend_mems[backend] + else: + print ("Backend", backend, "not instantiated to get its supported mem types.") + return None def get_backend_params(self, backend): - return self.backend_options_map[backend] + if backend in self.backend_options: + return self.backend_options[backend] + else: + print ("Backend", backend, "not instantiated to get its parameters.") + return None def create_backend(self, backend, initParams=None): self.backends[backend] = self.agent.createBackend(backend, initParams) @@ -282,6 +319,7 @@ def make_prepped_xfer( remote_indices, notif_msg, operation, + skip_desc_merge = False ): op = self.nixl_ops[operation] if op: @@ -292,6 +330,7 @@ def make_prepped_xfer( remote_indices, notif_msg, op, + skip_desc_merge, ) if handle == 0: return None @@ -300,7 +339,7 @@ def make_prepped_xfer( else: return None - def delete_xfer_side(self, handle): + def delete_dlist_handle(self, handle): # frees the handle too self.agent.releasedDlistH(handle) @@ -335,7 +374,7 @@ def check_remote_xfer_done(self, remote_agent_name, lookup_msg): self.notifs[remote_agent_name].remove(message) return message - def abort_xfer(self, handle): + def release_xfer_handle(self, handle): # frees the handle too self.agent.releaseXferReq(handle) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 0695bff40a..cee020dfd6 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -352,13 +352,15 @@ PYBIND11_MODULE(_bindings, m) { uintptr_t remote_side, const std::vector &remote_indices, const std::string ¬if_msg, - const nixl_xfer_op_t &operation) -> uintptr_t { + const nixl_xfer_op_t &operation, + bool skip_desc_merge) -> uintptr_t { nixlXferReqH* handle; nixl_opt_args_t extra_params; if (notif_msg.size()>0) { extra_params.notifMsg = notif_msg; extra_params.hasNotif = true; } + extra_params.skipDescMerge = skip_desc_merge; nixl_status_t ret = agent.makeXferReq(operation, (nixlDlistH*) local_side, local_indices, (nixlDlistH*) remote_side, remote_indices, diff --git a/test/python/blocking_send_recv_example.py b/test/python/blocking_send_recv_example.py index b28ecd806b..b77e0e3328 100755 --- a/test/python/blocking_send_recv_example.py +++ b/test/python/blocking_send_recv_example.py @@ -105,7 +105,7 @@ def parse_args(): if args.mode != "target": agent.remove_remote_agent(peer_name) - agent.abort_xfer(xfer_handle) + agent.release_xfer(xfer_handle) agent.deregister_memory(reg_descs) diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index e93ada1126..b7fff24d14 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -146,10 +146,10 @@ target_done = True print("Target done") - nixl_agent2.abort_xfer(xfer_handle_1) - nixl_agent2.abort_xfer(xfer_handle_2) - nixl_agent2.delete_xfer_side(local_prep_handle) - nixl_agent2.delete_xfer_side(remote_prep_handle) + nixl_agent2.release_xfer(xfer_handle_1) + nixl_agent2.release_xfer(xfer_handle_2) + nixl_agent2.delete_dlist_handle(local_prep_handle) + nixl_agent2.delete_dlist_handle(remote_prep_handle) nixl_agent2.remove_remote_agent("target") nixl_agent1.deregister_memory(agent1_reg_descs) nixl_agent2.deregister_memory(agent2_reg_descs) From 1c86516060464191729ee2e05c99f49b8987eba9 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 02:55:42 -0700 Subject: [PATCH 05/32] More changes to python API --- src/api/python/_api.py | 43 +++++++++++++++++++++++++----------------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index db1361ff9b..128158dea9 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -51,8 +51,6 @@ def __init__(self, agent_name, nixl_config=None): self.plugin_b_options[plugin] = backend_options self.plugin_mem_types[plugin] = mem_types - # self.backends["UCX"] = self.agent.createBackend("UCX", init) - init = {} if nixl_config: for x in nixl_config.backends: @@ -61,11 +59,10 @@ def __init__(self, agent_name, nixl_config=None): print("Skiping backend registration", x, "due to the missing plugin.") else: self.backends[x] = self.agent.createBackend(x, init) - else: # Defaulting to UCX and GDS for now - if "UCX" in self.plugin_list: - self.backends["UCX"] = self.agent.createBackend("UCX", init) - if "GDS" in self.plugin_list: - self.backends["GDS"] = self.agent.createBackend("GDS", init) + else: + # TODO: populate init from default parameters, or define a set of params in python + for plugin in self.plugin_list: + self.backends[plugin] = self.agent.createBackend(plugin, init) for backend in self.backends: (backend_options, mem_types) = self.agent.getBackendParams(backend) @@ -225,10 +222,15 @@ def register_memory( if backend: ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: + # TODO: rely on underlying capability to register with all when supported if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) - else: + else if (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) + else if (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): + ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) + else if (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): + ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) if ret != 0: return None return reg_descs @@ -238,11 +240,18 @@ def deregister_memory(self, dereg_descs, backend=None): if backend: self.agent.deregisterMem(dereg_descs, self.backends[backend]) else: - if (dereg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): - self.agent.deregisterMem(dereg_descs, self.backends["GDS"]) - else: - self.agent.deregisterMem(dereg_descs, self.backends["UCX"]) - # No return + # TODO: rely on underlying capability to register with all when supported + if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) + else if (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): + ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) + else if (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): + ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) + else if (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): + ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) + if ret != 0: + return None + return reg_descs # Optional proactive make connection def make_connection(self, remote_agent): @@ -258,7 +267,7 @@ def add_remote_agent(self, metadata): def remove_remote_agent(self, agent): self.agent.invalidateRemoteMD(agent) - def initialize_xfer( + def create_xfer( self, local_descs, remote_descs, remote_agent, notif_msg, operation, xfer_backend=None ): op = self.nixl_ops[operation] @@ -281,7 +290,7 @@ def initialize_xfer( return None # "" remote agent means local. example xfer can be used to know the backend - def prep_xfer_side( + def prep_xfer_dlist( self, remote_agent, xfer_list, @@ -311,7 +320,7 @@ def prep_xfer_side( return handle - def make_prepped_xfer( + def make_resolved_xfer( self, local_xfer_side, local_indices, @@ -339,7 +348,7 @@ def make_prepped_xfer( else: return None - def delete_dlist_handle(self, handle): + def release_dlist_handle(self, handle): # frees the handle too self.agent.releasedDlistH(handle) From 89ef6697ed556ac70733ee8e19087ad8204eecb8 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 08:12:35 -0700 Subject: [PATCH 06/32] cleanup and all exceptions --- src/api/python/_api.py | 2 +- src/bindings/python/nixl_bindings.cpp | 103 +++++++++++++------------- 2 files changed, 53 insertions(+), 52 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 128158dea9..9296a08f0b 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -40,7 +40,7 @@ def __init__(self, agent_name, nixl_config=None): self.backend_mems = {} self.backend_options = {} - self.plugin_list = nixlBind.getAvailPlugins() + self.plugin_list = self.getAvailPlugins() if len(self.plugin_list) == 0: print("No plugins available, cannot start transfers!") diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index cee020dfd6..8bf9dfdc87 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -265,38 +265,41 @@ PYBIND11_MODULE(_bindings, m) { .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> std::pair { nixl_b_params_t params; nixl_mem_list_t mems; - nixl_status_t ret = agent.getPluginParams(type, mems, params); - if(ret < 0) throw_nixl_exception(ret); + throw_nixl_exception(agent.getPluginParams(type, mems, params)); return std::make_pair(params, mems); }) .def("getBackendParams", [](nixlAgent &agent, uintptr_t backend) -> std::pair { nixl_b_params_t params; nixl_mem_list_t mems; - nixl_status_t ret = agent.getBackendParams((nixlBackendH*) backend, mems, params); - if(ret < 0) throw_nixl_exception(ret); + throw_nixl_exception(agent.getBackendParams((nixlBackendH*) backend, mems, params)); return std::make_pair(params, mems); }) .def("createBackend", [](nixlAgent &agent, const nixl_backend_t &type, const nixl_b_params_t &initParams) -> uintptr_t { - nixlBackendH* backend; - nixl_status_t ret = agent.createBackend(type, initParams, backend); - if(ret < 0) { - throw_nixl_exception(ret); - return (uintptr_t) nullptr; - } + nixlBackendH* backend = nullptr; + throw_nixl_exception(agent.createBackend(type, initParams, backend)); return (uintptr_t) backend; }) .def("registerMem", [](nixlAgent &agent, nixl_reg_dlist_t descs, uintptr_t backend) -> nixl_status_t { nixl_opt_args_t extra_params; + nixl_status_t ret; extra_params.backends.push_back((nixlBackendH*) backend); - return agent.registerMem(descs, &extra_params); + + ret = agent.registerMem(descs, &extra_params); + return ret; }) .def("deregisterMem", [](nixlAgent &agent, nixl_reg_dlist_t descs, uintptr_t backend) -> nixl_status_t { nixl_opt_args_t extra_params; + nixl_status_t ret; extra_params.backends.push_back((nixlBackendH*) backend); - return agent.deregisterMem(descs, &extra_params); + + ret = agent.deregisterMem(descs, &extra_params); + return ret; + }) + .def("makeConnection", [](nixlAgent &agent, const std::string &remote_agent) { + nixl_status_t ret = agent.makeConnection(remote_agent); + throw_nixl_exception(ret); + return ret; }) - .def("makeConnection", &nixlAgent::makeConnection) - //note: slight API change, python cannot receive values by passing refs, so handle must be returned .def("createXferReq", [](nixlAgent &agent, const nixl_xfer_dlist_t &local_descs, const nixl_xfer_dlist_t &remote_descs, @@ -304,7 +307,7 @@ PYBIND11_MODULE(_bindings, m) { const std::string ¬if_msg, const nixl_xfer_op_t &operation, uintptr_t backend) -> uintptr_t { - nixlXferReqH* handle; + nixlXferReqH* handle = nullptr; nixl_opt_args_t extra_params; if (backend!=0) extra_params.backends.push_back((nixlBackendH*) backend); @@ -314,36 +317,26 @@ PYBIND11_MODULE(_bindings, m) { } nixl_status_t ret = agent.createXferReq(operation, local_descs, remote_descs, remote_agent, handle, &extra_params); - if(ret < 0){ - throw_nixl_exception(ret); - return (uintptr_t) nullptr; - } + throw_nixl_exception(ret); return (uintptr_t) handle; }, py::arg("local_descs"), py::arg("remote_descs"), py::arg("remote_agent"), py::arg("notif_msg"), py::arg("operation"), py::arg("backend") = ((uintptr_t) nullptr)) .def("queryXferBackend", [](nixlAgent &agent, uintptr_t reqh) -> uintptr_t { - nixlBackendH* handle; - nixl_status_t ret = agent.queryXferBackend((nixlXferReqH*) reqh, handle); - if(ret < 0){ - throw_nixl_exception(ret); - return (uintptr_t) nullptr; - } + nixlBackendH* handle = nullptr; + throw_nixl_exception(agent.queryXferBackend((nixlXferReqH*) reqh, handle)); return (uintptr_t) handle; }) .def("prepXferDlist", [](nixlAgent &agent, const nixl_xfer_dlist_t &descs, const std::string &remote_agent, uintptr_t backend) -> uintptr_t { - nixlDlistH* handle; + nixlDlistH* handle = nullptr; nixl_opt_args_t extra_params; extra_params.backends.push_back((nixlBackendH*) backend); - nixl_status_t ret = agent.prepXferDlist(descs, remote_agent, handle, &extra_params); - if(ret < 0){ - throw_nixl_exception(ret); - return (uintptr_t) nullptr; - } + throw_nixl_exception(agent.prepXferDlist(descs, remote_agent, handle, &extra_params)); + return (uintptr_t) handle; }) .def("makeXferReq", [](nixlAgent &agent, @@ -354,38 +347,44 @@ PYBIND11_MODULE(_bindings, m) { const std::string ¬if_msg, const nixl_xfer_op_t &operation, bool skip_desc_merge) -> uintptr_t { - nixlXferReqH* handle; + nixlXferReqH* handle = nullptr; nixl_opt_args_t extra_params; if (notif_msg.size()>0) { extra_params.notifMsg = notif_msg; extra_params.hasNotif = true; } extra_params.skipDescMerge = skip_desc_merge; - nixl_status_t ret = agent.makeXferReq(operation, - (nixlDlistH*) local_side, local_indices, - (nixlDlistH*) remote_side, remote_indices, - handle, &extra_params); - if(ret < 0){ - throw_nixl_exception(ret); - return (uintptr_t) nullptr; - } + throw_nixl_exception(agent.makeXferReq(operation, + (nixlDlistH*) local_side, local_indices, + (nixlDlistH*) remote_side, remote_indices, + handle, &extra_params)); + return (uintptr_t) handle; }) .def("releaseXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { - return agent.releaseXferReq((nixlXferReqH*) reqh); + nixl_status_t ret = agent.releaseXferReq((nixlXferReqH*) reqh); + throw_nixl_exception(ret); + return ret; }) .def("releasedDlistH", [](nixlAgent &agent, uintptr_t handle) -> nixl_status_t { - return agent.releasedDlistH((nixlDlistH*) handle); + nixl_status_t ret = agent.releasedDlistH((nixlDlistH*) handle); + throw_nixl_exception(ret); + return ret; }) .def("postXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { - return agent.postXferReq((nixlXferReqH*) reqh); + nixl_status_t ret = agent.postXferReq((nixlXferReqH*) reqh); + throw_nixl_exception(ret); + return ret; }) .def("getXferStatus", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { - return agent.getXferStatus((nixlXferReqH*) reqh); + nixl_status_t ret = agent.getXferStatus((nixlXferReqH*) reqh); + throw_nixl_exception(ret); + return ret; }) .def("getNotifs", [](nixlAgent &agent, nixl_notifs_t notif_map) -> nixl_notifs_t { nixl_status_t ret = agent.getNotifs(notif_map); + throw_nixl_exception(ret); if (ret != NIXL_SUCCESS || notif_map.size() == 0) return notif_map; nixl_notifs_t ret_map; @@ -404,21 +403,23 @@ PYBIND11_MODULE(_bindings, m) { const std::string &msg, uintptr_t backend) { nixl_opt_args_t extra_params; + nixl_status_t ret; extra_params.backends.push_back((nixlBackendH*) backend); - return agent.genNotif(remote_agent, msg, &extra_params); + ret = agent.genNotif(remote_agent, msg, &extra_params); + + throw_nixl_exception(ret); + return ret; }) .def("getLocalMD", [](nixlAgent &agent) -> py::bytes { //python can only interpret text strings - std::string ret_str; - nixl_status_t ret = agent.getLocalMD(ret_str); - if(ret != NIXL_SUCCESS) return ""; + std::string ret_str(""); + throw_nixl_exception(agent.getLocalMD(ret_str)); return py::bytes(ret_str); }) .def("loadRemoteMD", [](nixlAgent &agent, const std::string &remote_metadata) -> py::bytes { //python can only interpret text strings - std::string remote_name; - nixl_status_t ret = agent.loadRemoteMD(remote_metadata, remote_name); - if(ret != NIXL_SUCCESS) return ""; + std::string remote_name(""); + throw_nixl_exception(agent.loadRemoteMD(remote_metadata, remote_name)); return py::bytes(remote_name); }) .def("invalidateRemoteMD", &nixlAgent::invalidateRemoteMD); From 57658a53ea47b7a38a52dfce38e2c1cc111be120 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 08:20:35 -0700 Subject: [PATCH 07/32] query backend and README --- README.md | 7 ++----- src/bindings/python/nixl_bindings.cpp | 5 +++++ 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 7d874a6e0b..33b9c9e655 100644 --- a/README.md +++ b/README.md @@ -63,11 +63,8 @@ The pybind11 bindings for the public facing NIXL API are available in src/bindin The Python module library (.so) will be built by default. To make it easy to import you can include either the build or install directory to your PYTHONPATH, see: -`$ export PYTHONPATH=$PYTHONPATH:/src/pybind` +`$ export PYTHONPATH=$PYTHONPATH:/src/pybind:/src/api/python` -or - -` $export PYTHONPATH=$PYTHONPATH:/lib64/python3.9/site-packages/` The preferred way is to build it through meson-python, which will just let it be installed with pip. This can be done from the root nixl directory: @@ -77,4 +74,4 @@ The preferred way is to build it through meson-python, which will just let it be * [C++ examples](https://github.com/ai-dynamo/nixl/tree/main/test/nixl) -* [Python examples](https://github.com/ai-dynamo/nixl/tree/main/test/python) \ No newline at end of file +* [Python examples](https://github.com/ai-dynamo/nixl/tree/main/test/python) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 8bf9dfdc87..e7dd8074f4 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -381,6 +381,11 @@ PYBIND11_MODULE(_bindings, m) { throw_nixl_exception(ret); return ret; }) + .def("queryXferBackend", [](nixlAgent &agent, uintptr_t reqh) -> uintptr_t { + nixlBackendH* backend = nullptr; + throw_nixl_exception(agent.queryXferBackend((nixlXferReqH*) reqh, backend)); + return (uintptr_t) backend; + }) .def("getNotifs", [](nixlAgent &agent, nixl_notifs_t notif_map) -> nixl_notifs_t { nixl_status_t ret = agent.getNotifs(notif_map); From 490d3ef7be13343852dccf287a7c2316bfece354 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 08:40:24 -0700 Subject: [PATCH 08/32] python api changes --- src/api/python/_api.py | 100 +++++++++++++------------- src/bindings/python/nixl_bindings.cpp | 42 ++++++----- 2 files changed, 77 insertions(+), 65 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 9296a08f0b..88a5a37d35 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -206,6 +206,7 @@ def get_reg_descs( return new_descs + #TODO: these not necessarily agent specific, maybe separate somehow? def get_serialized_descs(self, descs): return pickle.dumps(descs) @@ -257,56 +258,21 @@ def deregister_memory(self, dereg_descs, backend=None): def make_connection(self, remote_agent): self.agent.makeConnection(remote_agent) - def get_agent_metadata(self): - return self.agent.getLocalMD() - - def add_remote_agent(self, metadata): - agent_name = self.agent.loadRemoteMD(metadata) - return agent_name - - def remove_remote_agent(self, agent): - self.agent.invalidateRemoteMD(agent) - - def create_xfer( - self, local_descs, remote_descs, remote_agent, notif_msg, operation, xfer_backend=None - ): - op = self.nixl_ops[operation] - if op: - if xfer_backend: - handle = self.agent.createXferReq( - local_descs, - remote_descs, - remote_agent, - notif_msg, - op, - xfer_backend, - ) - else: - handle = self.agent.createXferReq( - local_descs, remote_descs, remote_agent, notif_msg, op - ) - return handle # In case of error it will be None - else: - return None - # "" remote agent means local. example xfer can be used to know the backend def prep_xfer_dlist( self, - remote_agent, xfer_list, + remote_agent, mem_type=None, is_unified_addr=True, is_sorted=False, - xfer_backend=None, - example_xfer=None, + xfer_backend=None ): descs = self.get_xfer_descs(xfer_list, mem_type, is_unified_addr, is_sorted) if xfer_backend: handle = self.agent.prepXferDlist(descs, remote_agent, xfer_backend) - elif example_xfer: - backend = self.agent.getXferBackend(example_xfer) - handle = self.agent.prepXferDlist(descs, remote_agent, backend) else: + #TODO: need better way to select backend if not specified if (descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] @@ -322,12 +288,12 @@ def prep_xfer_dlist( def make_resolved_xfer( self, + operation, local_xfer_side, local_indices, remote_xfer_side, remote_indices, - notif_msg, - operation, + notif_msg = "", skip_desc_merge = False ): op = self.nixl_ops[operation] @@ -348,12 +314,36 @@ def make_resolved_xfer( else: return None - def release_dlist_handle(self, handle): - # frees the handle too - self.agent.releasedDlistH(handle) + def create_xfer( + self, + operation, + local_descs, + remote_descs, + remote_agent, + notif_msg = "", + xfer_backend = None + ): + op = self.nixl_ops[operation] + if op: + if xfer_backend: + handle = self.agent.createXferReq( + local_descs, + remote_descs, + remote_agent, + notif_msg, + op, + xfer_backend, + ) + else: + handle = self.agent.createXferReq( + local_descs, remote_descs, remote_agent, notif_msg, op + ) + return handle # In case of error it will be None + else: + return None - def transfer(self, handle): - status = self.agent.postXferReq(handle) + def transfer(self, handle, notif_msg = ""): + status = self.agent.postXferReq(handle, notif_msg) if status == nixlBind.NIXL_SUCCESS: return "DONE" elif status == nixlBind.NIXL_IN_PROG: @@ -387,15 +377,29 @@ def release_xfer_handle(self, handle): # frees the handle too self.agent.releaseXferReq(handle) + def release_dlist_handle(self, handle): + # frees the handle too + self.agent.releasedDlistH(handle) + # Extra notification APIs def send_notif(self, remote_agent_name, notif_msg): self.agent.genNotif(remote_agent_name, notif_msg) + # Returns new notifs, without touching self.notifs + def get_new_notifs(self): + return self.agent.getNotifs({}) + # Adds new notifs to self.notifs and returns it def update_notifs(self): self.notifs = self.agent.getNotifs(self.notifs) return self.notifs - # Returns new notifs, without touching self.notifs - def get_new_notifs(self): - return self.agent.getNotifs({}) + def get_agent_metadata(self): + return self.agent.getLocalMD() + + def add_remote_agent(self, metadata): + agent_name = self.agent.loadRemoteMD(metadata) + return agent_name + + def remove_remote_agent(self, agent): + self.agent.invalidateRemoteMD(agent) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index e7dd8074f4..3e77ce7f58 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -301,11 +301,11 @@ PYBIND11_MODULE(_bindings, m) { return ret; }) .def("createXferReq", [](nixlAgent &agent, + const nixl_xfer_op_t &operation, const nixl_xfer_dlist_t &local_descs, const nixl_xfer_dlist_t &remote_descs, const std::string &remote_agent, const std::string ¬if_msg, - const nixl_xfer_op_t &operation, uintptr_t backend) -> uintptr_t { nixlXferReqH* handle = nullptr; nixl_opt_args_t extra_params; @@ -319,9 +319,9 @@ PYBIND11_MODULE(_bindings, m) { throw_nixl_exception(ret); return (uintptr_t) handle; - }, py::arg("local_descs"), + }, py::arg("operation"), py::arg("local_descs"), py::arg("remote_descs"), py::arg("remote_agent"), - py::arg("notif_msg"), py::arg("operation"), + py::arg("notif_msg") = std::string(""), py::arg("backend") = ((uintptr_t) nullptr)) .def("queryXferBackend", [](nixlAgent &agent, uintptr_t reqh) -> uintptr_t { nixlBackendH* handle = nullptr; @@ -340,12 +340,12 @@ PYBIND11_MODULE(_bindings, m) { return (uintptr_t) handle; }) .def("makeXferReq", [](nixlAgent &agent, + const nixl_xfer_op_t &operation, uintptr_t local_side, const std::vector &local_indices, uintptr_t remote_side, const std::vector &remote_indices, const std::string ¬if_msg, - const nixl_xfer_op_t &operation, bool skip_desc_merge) -> uintptr_t { nixlXferReqH* handle = nullptr; nixl_opt_args_t extra_params; @@ -360,19 +360,17 @@ PYBIND11_MODULE(_bindings, m) { handle, &extra_params)); return (uintptr_t) handle; - }) - .def("releaseXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { - nixl_status_t ret = agent.releaseXferReq((nixlXferReqH*) reqh); - throw_nixl_exception(ret); - return ret; - }) - .def("releasedDlistH", [](nixlAgent &agent, uintptr_t handle) -> nixl_status_t { - nixl_status_t ret = agent.releasedDlistH((nixlDlistH*) handle); - throw_nixl_exception(ret); - return ret; - }) - .def("postXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { - nixl_status_t ret = agent.postXferReq((nixlXferReqH*) reqh); + }, py::arg("operation"), py::arg("local_side"), + py::arg("local_indices"), py::arg("remote_side"), + py::arg("remote_indices"), py::arg("notif_msg") = std::string(""), + py::arg("skip_desc_merg") = false) + .def("postXferReq", [](nixlAgent &agent, uintptr_t reqh, std::string notif_msg) -> nixl_status_t { + nixl_opt_args_t extra_params; + if (notif_msg.size()>0) { + extra_params.notifMsg = notif_msg; + extra_params.hasNotif = true; + } + nixl_status_t ret = agent.postXferReq((nixlXferReqH*) reqh, &extra_params); throw_nixl_exception(ret); return ret; }) @@ -386,6 +384,16 @@ PYBIND11_MODULE(_bindings, m) { throw_nixl_exception(agent.queryXferBackend((nixlXferReqH*) reqh, backend)); return (uintptr_t) backend; }) + .def("releaseXferReq", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { + nixl_status_t ret = agent.releaseXferReq((nixlXferReqH*) reqh); + throw_nixl_exception(ret); + return ret; + }) + .def("releasedDlistH", [](nixlAgent &agent, uintptr_t handle) -> nixl_status_t { + nixl_status_t ret = agent.releasedDlistH((nixlDlistH*) handle); + throw_nixl_exception(ret); + return ret; + }) .def("getNotifs", [](nixlAgent &agent, nixl_notifs_t notif_map) -> nixl_notifs_t { nixl_status_t ret = agent.getNotifs(notif_map); From 51bb461f7a04d016e36db4339922da668a931d88 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 08:42:27 -0700 Subject: [PATCH 09/32] flake --- src/api/python/_api.py | 53 +++++++++++++++++++++++++----------------- 1 file changed, 32 insertions(+), 21 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 88a5a37d35..3cced1ced2 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -19,12 +19,14 @@ import nixl._bindings as nixlBind + class nixl_config: - def __init(self, enable_prog_thread = True, backends=["UCX", "GDS"]): + def __init(self, enable_prog_thread=True, backends=["UCX", "GDS"]): # TODO: add backend init parameters self.backends = backends self.enable_pthread = enable_prog_thread + class nixl_agent: def __init__(self, agent_name, nixl_config=None): # Set agent config and instantiate an agent @@ -56,7 +58,9 @@ def __init__(self, agent_name, nixl_config=None): for x in nixl_config.backends: # TODO: populate init from nixl_config when added if x not in self.plugin_list: - print("Skiping backend registration", x, "due to the missing plugin.") + print( + "Skipping backend registration", x, "due to the missing plugin." + ) else: self.backends[x] = self.agent.createBackend(x, init) else: @@ -89,28 +93,30 @@ def get_plugin_mem_types(self, backend): if backend in self.plugin_mem_types: return self.plugin_mem_types[backend] else: - print ("Plugin", backend, "is not available to get its supported mem types.") + print("Plugin", backend, "is not available to get its supported mem types.") return None def get_plugin_params(self, backend): if backend in self.plugin_b_options: return self.plugin_b_options[backend] else: - print ("Plugin", backend, "is not available to get its parameters.") + print("Plugin", backend, "is not available to get its parameters.") return None def get_backend_mem_types(self, backend): if backend in self.backend_mems: return self.backend_mems[backend] else: - print ("Backend", backend, "not instantiated to get its supported mem types.") + print( + "Backend", backend, "not instantiated to get its supported mem types." + ) return None def get_backend_params(self, backend): if backend in self.backend_options: return self.backend_options[backend] else: - print ("Backend", backend, "not instantiated to get its parameters.") + print("Backend", backend, "not instantiated to get its parameters.") return None def create_backend(self, backend, initParams=None): @@ -206,7 +212,7 @@ def get_reg_descs( return new_descs - #TODO: these not necessarily agent specific, maybe separate somehow? + # TODO: these not necessarily agent specific, maybe separate somehow? def get_serialized_descs(self, descs): return pickle.dumps(descs) @@ -215,7 +221,12 @@ def deserialize_descs(self, serialized_descs): # The returned descriptor object can be used for call to deregister def register_memory( - self, reg_list, mem_type=None, is_unified_addr=True, is_sorted=False, backend=None + self, + reg_list, + mem_type=None, + is_unified_addr=True, + is_sorted=False, + backend=None, ): reg_descs = self.get_reg_descs(reg_list, mem_type, is_unified_addr, is_sorted) @@ -226,11 +237,11 @@ def register_memory( # TODO: rely on underlying capability to register with all when supported if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) - else if (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - else if (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - else if (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): + elif (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) if ret != 0: return None @@ -244,11 +255,11 @@ def deregister_memory(self, dereg_descs, backend=None): # TODO: rely on underlying capability to register with all when supported if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) - else if (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) - else if (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) - else if (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): + elif (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) if ret != 0: return None @@ -266,13 +277,13 @@ def prep_xfer_dlist( mem_type=None, is_unified_addr=True, is_sorted=False, - xfer_backend=None + xfer_backend=None, ): descs = self.get_xfer_descs(xfer_list, mem_type, is_unified_addr, is_sorted) if xfer_backend: handle = self.agent.prepXferDlist(descs, remote_agent, xfer_backend) else: - #TODO: need better way to select backend if not specified + # TODO: need better way to select backend if not specified if (descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] @@ -293,8 +304,8 @@ def make_resolved_xfer( local_indices, remote_xfer_side, remote_indices, - notif_msg = "", - skip_desc_merge = False + notif_msg="", + skip_desc_merge=False, ): op = self.nixl_ops[operation] if op: @@ -320,8 +331,8 @@ def create_xfer( local_descs, remote_descs, remote_agent, - notif_msg = "", - xfer_backend = None + notif_msg="", + xfer_backend=None, ): op = self.nixl_ops[operation] if op: @@ -342,7 +353,7 @@ def create_xfer( else: return None - def transfer(self, handle, notif_msg = ""): + def transfer(self, handle, notif_msg=""): status = self.agent.postXferReq(handle, notif_msg) if status == nixlBind.NIXL_SUCCESS: return "DONE" From 6cfde2e4051f09b685ce216e017202542d7cb5d1 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 08:48:04 -0700 Subject: [PATCH 10/32] more flake --- src/api/python/_api.py | 39 +++++++++++++++++++++++---------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 3cced1ced2..a738aba9df 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -235,35 +235,42 @@ def register_memory( ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported - if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + if (reg_descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) - elif (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backend): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) if ret != 0: return None return reg_descs - def deregister_memory(self, dereg_descs, backend=None): + def deregister_memory(self, dereg_list, backend=None): # based on backend type and mem_type, figure what deregistrations are needed if backend: - self.agent.deregisterMem(dereg_descs, self.backends[backend]) + self.agent.deregisterMem(dereg_list, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported - if (reg_descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): - ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) - elif (reg_descs.getType() == nixl.DRAM_SEG) and ("UCX" in self.backend): - ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixl.VRAM_SEG) and ("UCX" in self.backend): - ret = self.agent.deregisterMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixl.VRAM_SEG) and ("GDS" in self.backend): - ret = self.agent.deregisterMem(reg_descs, self.backends["GDS"]) + if (dereg_list.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): + ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) + elif (dereg_list.getType() == nixlBind.DRAM_SEG) and ( + "UCX" in self.backend + ): + ret = self.agent.deregisterMem(dereg_list, self.backends["UCX"]) + elif (dereg_list.getType() == nixlBind.VRAM_SEG) and ( + "UCX" in self.backend + ): + ret = self.agent.deregisterMem(dereg_list, self.backends["UCX"]) + elif (dereg_list.getType() == nixlBind.VRAM_SEG) and ( + "GDS" in self.backend + ): + ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) if ret != 0: return None - return reg_descs + # is this the best ret value? + return dereg_list # Optional proactive make connection def make_connection(self, remote_agent): @@ -284,7 +291,7 @@ def prep_xfer_dlist( handle = self.agent.prepXferDlist(descs, remote_agent, xfer_backend) else: # TODO: need better way to select backend if not specified - if (descs.getType() == nixl.FILE_SEG) and ("GDS" in self.backend): + if (descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] ) From 04d7b89167d1befdd1cf5f7acfa59f535408f279 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 09:44:17 -0700 Subject: [PATCH 11/32] queryBackend API --- src/api/python/_api.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index a738aba9df..532dbac5fa 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -369,6 +369,13 @@ def transfer(self, handle, notif_msg=""): else: return "ERR" + def query_xfer_backend(self, handle): + b_handle = self.agent.queryXferBackend(handle) + # this works because there should not be multiple matching handles in the Dict + return next( + backendS for backendS, backendH in self.backends if backendH == b_handle + ) + def check_xfer_state(self, handle): status = self.agent.getXferStatus(handle) if status == nixlBind.NIXL_SUCCESS: From 0491565fd9b0948f45bffd465cff24cb83eaf507 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 10:48:38 -0700 Subject: [PATCH 12/32] Bug fix and minor improvements to _api --- src/api/python/_api.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 532dbac5fa..5de2fc3344 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -28,7 +28,7 @@ def __init(self, enable_prog_thread=True, backends=["UCX", "GDS"]): class nixl_agent: - def __init__(self, agent_name, nixl_config=None): + def __init__(self, agent_name, nixl_config=None, instantiate_all=False): # Set agent config and instantiate an agent if nixl_config: agent_config = nixlBind.nixlAgentConfig(nixl_config.enable_pthread) @@ -42,7 +42,7 @@ def __init__(self, agent_name, nixl_config=None): self.backend_mems = {} self.backend_options = {} - self.plugin_list = self.getAvailPlugins() + self.plugin_list = self.agent.getAvailPlugins() if len(self.plugin_list) == 0: print("No plugins available, cannot start transfers!") @@ -63,7 +63,7 @@ def __init__(self, agent_name, nixl_config=None): ) else: self.backends[x] = self.agent.createBackend(x, init) - else: + elif instantiate_all: # TODO: populate init from default parameters, or define a set of params in python for plugin in self.plugin_list: self.backends[plugin] = self.agent.createBackend(plugin, init) @@ -212,7 +212,7 @@ def get_reg_descs( return new_descs - # TODO: these not necessarily agent specific, maybe separate somehow? + # Since we create descriptor lists in agent, their SerDes methods are in the agent too def get_serialized_descs(self, descs): return pickle.dumps(descs) @@ -290,15 +290,23 @@ def prep_xfer_dlist( if xfer_backend: handle = self.agent.prepXferDlist(descs, remote_agent, xfer_backend) else: - # TODO: need better way to select backend if not specified + # TODO: rely on underlying capability to register with all when supported if (descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] ) - else: + elif (descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backend): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["UCX"] ) + elif (descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backend): + handle = self.agent.prepXferDlist( + descs, remote_agent, self.backends["UCX"] + ) + elif (descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backend): + handle = self.agent.prepXferDlist( + descs, remote_agent, self.backends["GDS"] + ) if handle == 0: return None From 4e6d821a5341251eff1b8404d99345c973f2b079 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 11:18:36 -0700 Subject: [PATCH 13/32] fix to getAvailPlugins in the bindings and wrapper test --- src/bindings/python/nixl_bindings.cpp | 6 +++++- test/python/nixl_wrapper_test.py | 9 +++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 3e77ce7f58..733b62bb39 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -261,7 +261,11 @@ PYBIND11_MODULE(_bindings, m) { py::class_(m, "nixlAgent") .def(py::init()) - .def("getAvailPlugins", &nixlAgent::getAvailPlugins) + .def("getAvailPlugins", [](nixlAgent &agent) -> std::vector { + std::vector backends; + throw_nixl_exception(agent.getAvailPlugins(backends)); + return backends; + }) .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> std::pair { nixl_b_params_t params; nixl_mem_list_t mems; diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index b7fff24d14..93bf8cbd52 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -23,17 +23,18 @@ if __name__ == "__main__": buf_size = 256 # Allocate memory and register with NIXL - nixl_agent1 = nixl_agent("target", None) + nixl_agent1 = nixl_agent("target", None, False) - plugin_list = nixl_agent1.get_plugin_list("UCX") + plugin_list = nixl_agent1.get_plugin_list() assert "UCX" in plugin_list print(plugin_list) - print(nixl_agent1.get_backend_mem_types("UCX")) - print(nixl_agent1.get_backend_params("UCX")) + print(nixl_agent1.get_plugin_mem_types("UCX")) + print(nixl_agent1.get_plugin_params("UCX")) nixl_agent1.create_backend("UCX") + print(nixl_agent1.get_backend_mem_types("UCX")) print(nixl_agent1.get_backend_params("UCX")) addr1 = nixl_utils.malloc_passthru(buf_size * 2) From a6be20b69f72c0418c3e2dd43fdc19c5e46a9d93 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 11:48:10 -0700 Subject: [PATCH 14/32] Reverted python API changes to be backwards compatible. --- src/api/python/_api.py | 20 ++++++++++---------- test/python/blocking_send_recv_example.py | 2 +- test/python/nixl_wrapper_test.py | 10 +++++----- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 5de2fc3344..e91566fd68 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -277,10 +277,10 @@ def make_connection(self, remote_agent): self.agent.makeConnection(remote_agent) # "" remote agent means local. example xfer can be used to know the backend - def prep_xfer_dlist( + def prep_xfer_side( self, - xfer_list, remote_agent, + xfer_list, mem_type=None, is_unified_addr=True, is_sorted=False, @@ -312,14 +312,14 @@ def prep_xfer_dlist( return handle - def make_resolved_xfer( + def make_prepped_xfer( self, - operation, local_xfer_side, local_indices, remote_xfer_side, remote_indices, - notif_msg="", + notif_msg, + operation, skip_desc_merge=False, ): op = self.nixl_ops[operation] @@ -340,13 +340,13 @@ def make_resolved_xfer( else: return None - def create_xfer( + def initialize_xfer( self, - operation, local_descs, remote_descs, remote_agent, - notif_msg="", + notif_msg, + operation, xfer_backend=None, ): op = self.nixl_ops[operation] @@ -406,11 +406,11 @@ def check_remote_xfer_done(self, remote_agent_name, lookup_msg): self.notifs[remote_agent_name].remove(message) return message - def release_xfer_handle(self, handle): + def abort_xfer(self, handle): # frees the handle too self.agent.releaseXferReq(handle) - def release_dlist_handle(self, handle): + def delete_xfer_side(self, handle): # frees the handle too self.agent.releasedDlistH(handle) diff --git a/test/python/blocking_send_recv_example.py b/test/python/blocking_send_recv_example.py index b77e0e3328..b28ecd806b 100755 --- a/test/python/blocking_send_recv_example.py +++ b/test/python/blocking_send_recv_example.py @@ -105,7 +105,7 @@ def parse_args(): if args.mode != "target": agent.remove_remote_agent(peer_name) - agent.release_xfer(xfer_handle) + agent.abort_xfer(xfer_handle) agent.deregister_memory(reg_descs) diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index 93bf8cbd52..caed1cee05 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -23,7 +23,7 @@ if __name__ == "__main__": buf_size = 256 # Allocate memory and register with NIXL - nixl_agent1 = nixl_agent("target", None, False) + nixl_agent1 = nixl_agent("target") plugin_list = nixl_agent1.get_plugin_list() assert "UCX" in plugin_list @@ -147,10 +147,10 @@ target_done = True print("Target done") - nixl_agent2.release_xfer(xfer_handle_1) - nixl_agent2.release_xfer(xfer_handle_2) - nixl_agent2.delete_dlist_handle(local_prep_handle) - nixl_agent2.delete_dlist_handle(remote_prep_handle) + nixl_agent2.abort_xfer(xfer_handle_1) + nixl_agent2.abort_xfer(xfer_handle_2) + nixl_agent2.delete_xfer_side(local_prep_handle) + nixl_agent2.delete_xfer_side(remote_prep_handle) nixl_agent2.remove_remote_agent("target") nixl_agent1.deregister_memory(agent1_reg_descs) nixl_agent2.deregister_memory(agent2_reg_descs) From 85ced6cd660295c35f89481360f1f84296a15fce Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 12:35:23 -0700 Subject: [PATCH 15/32] Fixes. --- src/api/python/_api.py | 200 +++++++++++++++++++++-------------------- 1 file changed, 101 insertions(+), 99 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index e91566fd68..6a6a0eac57 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -55,14 +55,14 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): init = {} if nixl_config: - for x in nixl_config.backends: + for bknd in nixl_config.backends: # TODO: populate init from nixl_config when added - if x not in self.plugin_list: + if bknd not in self.plugin_list: print( - "Skipping backend registration", x, "due to the missing plugin." + "Skipping backend registration", bknd, "due to the missing plugin." ) else: - self.backends[x] = self.agent.createBackend(x, init) + self.backends[bknd] = self.agent.createBackend(bknd, init) elif instantiate_all: # TODO: populate init from default parameters, or define a set of params in python for plugin in self.plugin_list: @@ -70,7 +70,7 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): for backend in self.backends: (backend_options, mem_types) = self.agent.getBackendParams(backend) - self.backend_mems[backend] = mem_types + self.backend_mems[backend] = mem_types self.backend_options[backend] = backend_options self.nixl_mems = { @@ -119,105 +119,14 @@ def get_backend_params(self, backend): print("Backend", backend, "not instantiated to get its parameters.") return None - def create_backend(self, backend, initParams=None): + def create_backend(self, backend, initParams={}): self.backends[backend] = self.agent.createBackend(backend, initParams) (backend_options, mem_types) = self.agent.getBackendParams( self.backends[backend] ) - self.backend_option_map[backend] = backend_options - self.mem_type_map[backend] = mem_types - - def get_xfer_descs( - self, descs, mem_type=None, is_unified_addr=True, is_sorted=False - ): - # can add check for DLPack input - - if isinstance(descs, nixlBind.nixlXferDList): - return descs - elif isinstance(descs[0], tuple): - if mem_type is not None and len(descs[0]) == 3: - new_descs = nixlBind.nixlXferDList( - self.nixl_mems[mem_type], descs, is_unified_addr, is_sorted - ) - elif mem_type is None: - print("Please specify a mem type if not using Tensors") - new_descs = None - else: - print("3-tuple list needed for transfer") - new_descs = None - elif isinstance(descs[0], torch.Tensor): # List[torch.Tensor]: - tensor_type = descs[0].device - dlist = [(0, 0, 0)] * len(descs) - - for i in range(len(descs)): - if descs[i].device != tensor_type: - return None - base_addr = descs[i].data_ptr() - region_len = descs[i].numel() * descs[i].element_size() - gpu_id = descs[i].get_device() - if gpu_id == -1: # DRAM - gpu_id = 0 - dlist[i] = (base_addr, region_len, gpu_id) - new_descs = nixlBind.nixlXferDList( - self.nixl_mems[str(tensor_type)], dlist, is_unified_addr, is_sorted - ) - elif isinstance(descs, nixlBind.nixlRegDList): - print("RegList type detected for transfer, please use XferList") - new_descs = None - else: - new_descs = None - - return new_descs - - def get_reg_descs( - self, descs, mem_type=None, is_unified_addr=True, is_sorted=False - ): - # can add check for DLPack input - - if isinstance(descs, nixlBind.nixlRegDList): - return descs - elif isinstance(descs[0], tuple): - if mem_type is not None and len(descs[0]) == 4: - new_descs = nixlBind.nixlRegDList( - self.nixl_mems[mem_type], descs, is_unified_addr, is_sorted - ) - elif mem_type is None: - print("Please specify a mem type if not using Tensors") - new_descs = None - else: - print("4-tuple list needed for registration") - new_descs = None - elif isinstance(descs[0], torch.Tensor): # List[torch.Tensor]: - tensor_type = descs[0].device - dlist = [(0, 0, 0, "")] * len(descs) - - for i in range(len(descs)): - if descs[i].device != tensor_type: - return None - base_addr = descs[i].data_ptr() - region_len = descs[i].numel() * descs[i].element_size() - gpu_id = descs[i].get_device() - if gpu_id == -1: # DRAM - gpu_id = 0 - dlist[i] = (base_addr, region_len, gpu_id, "") - new_descs = nixlBind.nixlRegDList( - self.nixl_mems[str(tensor_type)], dlist, is_unified_addr, is_sorted - ) - elif isinstance(descs, nixlBind.nixlXferDList): - print("XferList type detected for registration, please use RegList") - new_descs = None - else: - new_descs = None - - return new_descs - - # Since we create descriptor lists in agent, their SerDes methods are in the agent too - def get_serialized_descs(self, descs): - return pickle.dumps(descs) - - def deserialize_descs(self, serialized_descs): - return pickle.loads(serialized_descs) + self.backend_mems[backend] = mem_types + self.backend_options[backend] = backend_options # The returned descriptor object can be used for call to deregister def register_memory( @@ -436,3 +345,96 @@ def add_remote_agent(self, metadata): def remove_remote_agent(self, agent): self.agent.invalidateRemoteMD(agent) + + ### Methods to create and serialize/deserialize descriptors, provided through Agent + def get_xfer_descs( + self, descs, mem_type=None, is_unified_addr=True, is_sorted=False + ): + # can add check for DLPack input + + if isinstance(descs, nixlBind.nixlXferDList): + return descs + elif isinstance(descs[0], tuple): + if mem_type is not None and len(descs[0]) == 3: + new_descs = nixlBind.nixlXferDList( + self.nixl_mems[mem_type], descs, is_unified_addr, is_sorted + ) + elif mem_type is None: + print("Please specify a mem type if not using Tensors") + new_descs = None + else: + print("3-tuple list needed for transfer") + new_descs = None + elif isinstance(descs[0], torch.Tensor): # List[torch.Tensor]: + tensor_type = descs[0].device + dlist = [(0, 0, 0)] * len(descs) + + for i in range(len(descs)): + if descs[i].device != tensor_type: + return None + base_addr = descs[i].data_ptr() + region_len = descs[i].numel() * descs[i].element_size() + gpu_id = descs[i].get_device() + if gpu_id == -1: # DRAM + gpu_id = 0 + dlist[i] = (base_addr, region_len, gpu_id) + new_descs = nixlBind.nixlXferDList( + self.nixl_mems[str(tensor_type)], dlist, is_unified_addr, is_sorted + ) + elif isinstance(descs, nixlBind.nixlRegDList): + print("RegList type detected for transfer, please use XferList") + new_descs = None + else: + new_descs = None + + return new_descs + + def get_reg_descs( + self, descs, mem_type=None, is_unified_addr=True, is_sorted=False + ): + # can add check for DLPack input + + if isinstance(descs, nixlBind.nixlRegDList): + return descs + elif isinstance(descs[0], tuple): + if mem_type is not None and len(descs[0]) == 4: + new_descs = nixlBind.nixlRegDList( + self.nixl_mems[mem_type], descs, is_unified_addr, is_sorted + ) + elif mem_type is None: + print("Please specify a mem type if not using Tensors") + new_descs = None + else: + print("4-tuple list needed for registration") + new_descs = None + elif isinstance(descs[0], torch.Tensor): # List[torch.Tensor]: + tensor_type = descs[0].device + dlist = [(0, 0, 0, "")] * len(descs) + + for i in range(len(descs)): + if descs[i].device != tensor_type: + return None + base_addr = descs[i].data_ptr() + region_len = descs[i].numel() * descs[i].element_size() + gpu_id = descs[i].get_device() + if gpu_id == -1: # DRAM + gpu_id = 0 + dlist[i] = (base_addr, region_len, gpu_id, "") + new_descs = nixlBind.nixlRegDList( + self.nixl_mems[str(tensor_type)], dlist, is_unified_addr, is_sorted + ) + elif isinstance(descs, nixlBind.nixlXferDList): + print("XferList type detected for registration, please use RegList") + new_descs = None + else: + new_descs = None + + return new_descs + + def get_serialized_descs(self, descs): + return pickle.dumps(descs) + + def deserialize_descs(self, serialized_descs): + return pickle.loads(serialized_descs) + + From 632be9345341c3da82a066a1322c63c627856b9a Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 12:40:33 -0700 Subject: [PATCH 16/32] more fixes. --- src/api/python/_api.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 6a6a0eac57..271f9a3784 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -144,13 +144,13 @@ def register_memory( ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported - if (reg_descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): + if (reg_descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backends): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) - elif (reg_descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backends): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backend): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backends): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backend): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backends): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) if ret != 0: return None @@ -162,18 +162,20 @@ def deregister_memory(self, dereg_list, backend=None): self.agent.deregisterMem(dereg_list, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported - if (dereg_list.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): + if (dereg_list.getType() == nixlBind.FILE_SEG) and ( + "GDS" in self.backends + ): ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) elif (dereg_list.getType() == nixlBind.DRAM_SEG) and ( - "UCX" in self.backend + "UCX" in self.backends ): ret = self.agent.deregisterMem(dereg_list, self.backends["UCX"]) elif (dereg_list.getType() == nixlBind.VRAM_SEG) and ( - "UCX" in self.backend + "UCX" in self.backends ): ret = self.agent.deregisterMem(dereg_list, self.backends["UCX"]) elif (dereg_list.getType() == nixlBind.VRAM_SEG) and ( - "GDS" in self.backend + "GDS" in self.backends ): ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) if ret != 0: @@ -200,19 +202,19 @@ def prep_xfer_side( handle = self.agent.prepXferDlist(descs, remote_agent, xfer_backend) else: # TODO: rely on underlying capability to register with all when supported - if (descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backend): + if (descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backends): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] ) - elif (descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backend): + elif (descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backends): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["UCX"] ) - elif (descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backend): + elif (descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backends): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["UCX"] ) - elif (descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backend): + elif (descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backends): handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] ) @@ -346,7 +348,8 @@ def add_remote_agent(self, metadata): def remove_remote_agent(self, agent): self.agent.invalidateRemoteMD(agent) - ### Methods to create and serialize/deserialize descriptors, provided through Agent + # 4 methods to create and serialize/deserialize descriptors, provided through Agent + def get_xfer_descs( self, descs, mem_type=None, is_unified_addr=True, is_sorted=False ): From 77da6a09fd0d13a798365606b016f60b8476f910 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 12:43:09 -0700 Subject: [PATCH 17/32] black fixes --- src/api/python/_api.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 271f9a3784..71c92b345b 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -59,7 +59,9 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): # TODO: populate init from nixl_config when added if bknd not in self.plugin_list: print( - "Skipping backend registration", bknd, "due to the missing plugin." + "Skipping backend registration", + bknd, + "due to the missing plugin.", ) else: self.backends[bknd] = self.agent.createBackend(bknd, init) @@ -70,7 +72,7 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): for backend in self.backends: (backend_options, mem_types) = self.agent.getBackendParams(backend) - self.backend_mems[backend] = mem_types + self.backend_mems[backend] = mem_types self.backend_options[backend] = backend_options self.nixl_mems = { @@ -125,7 +127,7 @@ def create_backend(self, backend, initParams={}): (backend_options, mem_types) = self.agent.getBackendParams( self.backends[backend] ) - self.backend_mems[backend] = mem_types + self.backend_mems[backend] = mem_types self.backend_options[backend] = backend_options # The returned descriptor object can be used for call to deregister @@ -146,11 +148,17 @@ def register_memory( # TODO: rely on underlying capability to register with all when supported if (reg_descs.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backends): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) - elif (reg_descs.getType() == nixlBind.DRAM_SEG) and ("UCX" in self.backends): + elif (reg_descs.getType() == nixlBind.DRAM_SEG) and ( + "UCX" in self.backends + ): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("UCX" in self.backends): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ( + "UCX" in self.backends + ): ret = self.agent.registerMem(reg_descs, self.backends["UCX"]) - elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ("GDS" in self.backends): + elif (reg_descs.getType() == nixlBind.VRAM_SEG) and ( + "GDS" in self.backends + ): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) if ret != 0: return None @@ -162,9 +170,7 @@ def deregister_memory(self, dereg_list, backend=None): self.agent.deregisterMem(dereg_list, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported - if (dereg_list.getType() == nixlBind.FILE_SEG) and ( - "GDS" in self.backends - ): + if (dereg_list.getType() == nixlBind.FILE_SEG) and ("GDS" in self.backends): ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) elif (dereg_list.getType() == nixlBind.DRAM_SEG) and ( "UCX" in self.backends @@ -439,5 +445,3 @@ def get_serialized_descs(self, descs): def deserialize_descs(self, serialized_descs): return pickle.loads(serialized_descs) - - From d82fbdc96085ad36357f11855f341f762bfa00e1 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 13:15:13 -0700 Subject: [PATCH 18/32] Improved get mems outputs to be string, fix for nixl_wrapper --- src/api/python/_api.py | 9 +++++++++ src/bindings/python/nixl_bindings.cpp | 14 ++++++++++---- test/python/nixl_wrapper_test.py | 2 +- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 71c92b345b..3d01cf4344 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -160,6 +160,9 @@ def register_memory( "GDS" in self.backends ): ret = self.agent.registerMem(reg_descs, self.backends["GDS"]) + else: + return None + if ret != 0: return None return reg_descs @@ -184,6 +187,9 @@ def deregister_memory(self, dereg_list, backend=None): "GDS" in self.backends ): ret = self.agent.deregisterMem(dereg_list, self.backends["GDS"]) + else: + return None + if ret != 0: return None # is this the best ret value? @@ -224,6 +230,9 @@ def prep_xfer_side( handle = self.agent.prepXferDlist( descs, remote_agent, self.backends["GDS"] ) + else: + return None + if handle == 0: return None diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 733b62bb39..9f094404f4 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -266,17 +266,23 @@ PYBIND11_MODULE(_bindings, m) { throw_nixl_exception(agent.getAvailPlugins(backends)); return backends; }) - .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> std::pair { + .def("getPluginParams", [](nixlAgent &agent, const nixl_backend_t type) -> std::pair> { nixl_b_params_t params; nixl_mem_list_t mems; + std::vector mems_vec; throw_nixl_exception(agent.getPluginParams(type, mems, params)); - return std::make_pair(params, mems); + for (const auto& elm: mems) + mems_vec.push_back(nixlEnumStrings::memTypeStr(elm)); + return std::make_pair(params, mems_vec); }) - .def("getBackendParams", [](nixlAgent &agent, uintptr_t backend) -> std::pair { + .def("getBackendParams", [](nixlAgent &agent, uintptr_t backend) -> std::pair> { nixl_b_params_t params; nixl_mem_list_t mems; + std::vector mems_vec; throw_nixl_exception(agent.getBackendParams((nixlBackendH*) backend, mems, params)); - return std::make_pair(params, mems); + for (const auto& elm: mems) + mems_vec.push_back(nixlEnumStrings::memTypeStr(elm)); + return std::make_pair(params, mems_vec); }) .def("createBackend", [](nixlAgent &agent, const nixl_backend_t &type, const nixl_b_params_t &initParams) -> uintptr_t { nixlBackendH* backend = nullptr; diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index caed1cee05..146cb45779 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -63,7 +63,7 @@ agent2_reg_descs = nixl_agent1.get_reg_descs(agent2_strings, "DRAM", True) agent2_xfer_descs = nixl_agent1.get_xfer_descs(agent2_addrs, "DRAM", True) - agent2_descs = nixl_agent2.register_memory(agent2_reg_descs, True) + agent2_descs = nixl_agent2.register_memory(agent2_reg_descs, is_sorted=True) assert agent2_descs is not None # Exchange metadata From 8e028669b96080433ed9c336a25709cccf7754e8 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 14:12:02 -0700 Subject: [PATCH 19/32] fixes to _api and tester --- src/api/python/_api.py | 6 +++--- test/python/nixl_wrapper_test.py | 11 ++++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 3d01cf4344..12f60b5c84 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -251,12 +251,12 @@ def make_prepped_xfer( op = self.nixl_ops[operation] if op: handle = self.agent.makeXferReq( + op, local_xfer_side, local_indices, remote_xfer_side, remote_indices, notif_msg, - op, skip_desc_merge, ) if handle == 0: @@ -279,16 +279,16 @@ def initialize_xfer( if op: if xfer_backend: handle = self.agent.createXferReq( + op, local_descs, remote_descs, remote_agent, notif_msg, - op, xfer_backend, ) else: handle = self.agent.createXferReq( - local_descs, remote_descs, remote_agent, notif_msg, op + op, local_descs, remote_descs, remote_agent, notif_msg ) return handle # In case of error it will be None else: diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index 146cb45779..f3e5be075e 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -26,7 +26,7 @@ nixl_agent1 = nixl_agent("target") plugin_list = nixl_agent1.get_plugin_list() - assert "UCX" in plugin_list + # assert "UCX" in plugin_list # not supported in static linking mode yet print(plugin_list) print(nixl_agent1.get_plugin_mem_types("UCX")) @@ -43,8 +43,8 @@ agent1_addrs = [(addr1, buf_size, 0), (addr2, buf_size, 0)] agent1_strings = [(addr1, buf_size, 0, "a"), (addr2, buf_size, 0, "b")] - agent1_reg_descs = nixl_agent1.get_reg_descs(agent1_strings, "DRAM", True) - agent1_xfer_descs = nixl_agent1.get_xfer_descs(agent1_addrs, "DRAM", True) + agent1_reg_descs = nixl_agent1.get_reg_descs(agent1_strings, "DRAM", is_sorted=True) + agent1_xfer_descs = nixl_agent1.get_xfer_descs(agent1_addrs, "DRAM", is_sorted=True) # Just for tensor test tensors = [torch.zeros(10, dtype=torch.float32) for _ in range(2)] @@ -54,14 +54,15 @@ assert nixl_agent1.register_memory(agent1_reg_descs) is not None nixl_agent2 = nixl_agent("initiator", None) + nixl_agent2.create_backend("UCX") addr3 = nixl_utils.malloc_passthru(buf_size * 2) addr4 = addr3 + buf_size agent2_addrs = [(addr3, buf_size, 0), (addr4, buf_size, 0)] agent2_strings = [(addr3, buf_size, 0, "a"), (addr4, buf_size, 0, "b")] - agent2_reg_descs = nixl_agent1.get_reg_descs(agent2_strings, "DRAM", True) - agent2_xfer_descs = nixl_agent1.get_xfer_descs(agent2_addrs, "DRAM", True) + agent2_reg_descs = nixl_agent2.get_reg_descs(agent2_strings, "DRAM", is_sorted=True) + agent2_xfer_descs = nixl_agent2.get_xfer_descs(agent2_addrs, "DRAM", is_sorted=True) agent2_descs = nixl_agent2.register_memory(agent2_reg_descs, is_sorted=True) assert agent2_descs is not None From a5d9cc1aef04c98b17094882d206391aa2135849 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 14:22:28 -0700 Subject: [PATCH 20/32] reordering. --- src/api/python/_api.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 12f60b5c84..77934faa7b 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -319,19 +319,6 @@ def check_xfer_state(self, handle): else: return "ERR" - # Only removes the specific notification from self.notifs - def check_remote_xfer_done(self, remote_agent_name, lookup_msg): - self.notifs = self.agent.getNotifs(self.notifs) # Adds new notifs - message = None - if remote_agent_name in self.notifs: - for msg in self.notifs[remote_agent_name]: - if lookup_msg in msg: - message = msg - break - if message: - self.notifs[remote_agent_name].remove(message) - return message - def abort_xfer(self, handle): # frees the handle too self.agent.releaseXferReq(handle) @@ -340,10 +327,6 @@ def delete_xfer_side(self, handle): # frees the handle too self.agent.releasedDlistH(handle) - # Extra notification APIs - def send_notif(self, remote_agent_name, notif_msg): - self.agent.genNotif(remote_agent_name, notif_msg) - # Returns new notifs, without touching self.notifs def get_new_notifs(self): return self.agent.getNotifs({}) @@ -353,6 +336,23 @@ def update_notifs(self): self.notifs = self.agent.getNotifs(self.notifs) return self.notifs + # Only removes the specific notification from self.notifs + def check_remote_xfer_done(self, remote_agent_name, lookup_msg): + self.notifs = self.agent.getNotifs(self.notifs) # Adds new notifs + message = None + if remote_agent_name in self.notifs: + for msg in self.notifs[remote_agent_name]: + if lookup_msg in msg: + message = msg + break + if message: + self.notifs[remote_agent_name].remove(message) + return message + + # Extra notification APIs + def send_notif(self, remote_agent_name, notif_msg): + self.agent.genNotif(remote_agent_name, notif_msg) + def get_agent_metadata(self): return self.agent.getLocalMD() From e62c6d693b1f8d7847bc3c5fd16797c69483f98e Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 15:06:50 -0700 Subject: [PATCH 21/32] bug fix --- src/bindings/python/nixl_bindings.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 9f094404f4..707f0f6ccb 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -376,11 +376,14 @@ PYBIND11_MODULE(_bindings, m) { py::arg("skip_desc_merg") = false) .def("postXferReq", [](nixlAgent &agent, uintptr_t reqh, std::string notif_msg) -> nixl_status_t { nixl_opt_args_t extra_params; + nixl_status_t ret; if (notif_msg.size()>0) { extra_params.notifMsg = notif_msg; extra_params.hasNotif = true; + ret = agent.postXferReq((nixlXferReqH*) reqh, &extra_params); + } else { + ret = agent.postXferReq((nixlXferReqH*) reqh); } - nixl_status_t ret = agent.postXferReq((nixlXferReqH*) reqh, &extra_params); throw_nixl_exception(ret); return ret; }) From 2c1113f44576ef657a6265cf673ce81f6532c15a Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 15:25:13 -0700 Subject: [PATCH 22/32] README --- README.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/README.md b/README.md index 33b9c9e655..6d6c47b982 100644 --- a/README.md +++ b/README.md @@ -61,11 +61,6 @@ $ ninja-install ### pybind11 Python Interface The pybind11 bindings for the public facing NIXL API are available in src/bindings/python. These bindings implement the headers in the src/api/cpp directory. -The Python module library (.so) will be built by default. To make it easy to import you can include either the build or install directory to your PYTHONPATH, see: - -`$ export PYTHONPATH=$PYTHONPATH:/src/pybind:/src/api/python` - - The preferred way is to build it through meson-python, which will just let it be installed with pip. This can be done from the root nixl directory: ` $pip install .` From facf92e1fc1fe6b254cdc68070c92c76c17847f3 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 15:55:03 -0700 Subject: [PATCH 23/32] Added NIXL_ERR_NOT_SUPPORTED --- src/bindings/python/nixl_bindings.cpp | 5 +++++ src/core/nixl_agent.cpp | 1 + 2 files changed, 6 insertions(+) diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 707f0f6ccb..52e63253a9 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -99,6 +99,9 @@ void throw_nixl_exception(const nixl_status_t &status) { case NIXL_ERR_UNKNOWN: throw nixlUnknownError(nixlEnumStrings::statusStr(status).c_str()); break; + case NIXL_ERR_NOT_SUPPORTED: + throw nixlUnknownError(nixlEnumStrings::statusStr(status).c_str()); + break; default: throw std::runtime_error("BAD_STATUS"); } @@ -134,6 +137,7 @@ PYBIND11_MODULE(_bindings, m) { .value("NIXL_ERR_NOT_ALLOWED", NIXL_ERR_NOT_ALLOWED) .value("NIXL_ERR_REPOST_ACTIVE", NIXL_ERR_REPOST_ACTIVE) .value("NIXL_ERR_UNKNOWN", NIXL_ERR_UNKNOWN) + .value("NIXL_ERR_NOT_SUPPORTED", NIXL_ERR_NOT_SUPPORTED) .export_values(); py::register_exception(m, "nixlNotPostedError"); @@ -144,6 +148,7 @@ PYBIND11_MODULE(_bindings, m) { py::register_exception(m, "nixlNotAllowedError"); py::register_exception(m, "nixlRepostActiveError"); py::register_exception(m, "nixlUnknownError"); + py::register_exception(m, "nixlNotSupportedError"); py::class_(m, "nixlXferDList") .def(py::init(), py::arg("type"), py::arg("unifiedAddr")=true, py::arg("sorted")=false, py::arg("init_size")=0) diff --git a/src/core/nixl_agent.cpp b/src/core/nixl_agent.cpp index 51f309a1a6..cbdc0e3724 100644 --- a/src/core/nixl_agent.cpp +++ b/src/core/nixl_agent.cpp @@ -52,6 +52,7 @@ std::string nixlEnumStrings::statusStr (const nixl_status_t &status) { case NIXL_ERR_NOT_ALLOWED: return "NIXL_ERR_NOT_ALLOWED"; case NIXL_ERR_REPOST_ACTIVE: return "NIXL_ERR_REPOST_ACTIVE"; case NIXL_ERR_UNKNOWN: return "NIXL_ERR_UNKNOWN"; + case NIXL_ERR_NOT_SUPPORTED: return "NIXL_ERR_NOT_SUPPORTED"; default: return "BAD_STATUS"; } } From f945fa22b554b4a80d6a5d3fac07eef12479cafc Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 16:00:22 -0700 Subject: [PATCH 24/32] exception test --- src/api/python/_api.py | 2 ++ src/bindings/python/nixl_bindings.cpp | 13 ++++++-- test/python/nixl_exception_test.py | 43 +++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 3 deletions(-) create mode 100755 test/python/nixl_exception_test.py diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 77934faa7b..46723ed0da 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -143,6 +143,8 @@ def register_memory( # based on backend type and mem_type, figure what registrations are meaningful if backend: + print("trying backend") + print(self.backends[backend]) ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 52e63253a9..6ef5205e53 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -65,6 +65,10 @@ class nixlRepostActiveError : public std::runtime_error { nixlRepostActiveError(const char* what) : runtime_error(what) {} }; +class nixlNotSupportedError : public std::runtime_error { + public: + nixlNotSupportedError(const char* what) : runtime_error(what) {} +}; class nixlUnknownError : public std::runtime_error { public: @@ -100,7 +104,7 @@ void throw_nixl_exception(const nixl_status_t &status) { throw nixlUnknownError(nixlEnumStrings::statusStr(status).c_str()); break; case NIXL_ERR_NOT_SUPPORTED: - throw nixlUnknownError(nixlEnumStrings::statusStr(status).c_str()); + throw nixlNotSupportedError(nixlEnumStrings::statusStr(status).c_str()); break; default: throw std::runtime_error("BAD_STATUS"); @@ -148,7 +152,7 @@ PYBIND11_MODULE(_bindings, m) { py::register_exception(m, "nixlNotAllowedError"); py::register_exception(m, "nixlRepostActiveError"); py::register_exception(m, "nixlUnknownError"); - py::register_exception(m, "nixlNotSupportedError"); + py::register_exception(m, "nixlNotSupportedError"); py::class_(m, "nixlXferDList") .def(py::init(), py::arg("type"), py::arg("unifiedAddr")=true, py::arg("sorted")=false, py::arg("init_size")=0) @@ -297,9 +301,11 @@ PYBIND11_MODULE(_bindings, m) { .def("registerMem", [](nixlAgent &agent, nixl_reg_dlist_t descs, uintptr_t backend) -> nixl_status_t { nixl_opt_args_t extra_params; nixl_status_t ret; - extra_params.backends.push_back((nixlBackendH*) backend); + if(backend != 0) + extra_params.backends.push_back((nixlBackendH*) backend); ret = agent.registerMem(descs, &extra_params); + throw_nixl_exception(ret); return ret; }) .def("deregisterMem", [](nixlAgent &agent, nixl_reg_dlist_t descs, uintptr_t backend) -> nixl_status_t { @@ -308,6 +314,7 @@ PYBIND11_MODULE(_bindings, m) { extra_params.backends.push_back((nixlBackendH*) backend); ret = agent.deregisterMem(descs, &extra_params); + throw_nixl_exception(ret); return ret; }) .def("makeConnection", [](nixlAgent &agent, const std::string &remote_agent) { diff --git a/test/python/nixl_exception_test.py b/test/python/nixl_exception_test.py new file mode 100755 index 0000000000..05b9ff3d33 --- /dev/null +++ b/test/python/nixl_exception_test.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import nixl._utils as nixl_utils +from nixl._api import nixl_agent +import nixl._bindings as nixl_bindings + +if __name__ == "__main__": + + nixl_agent1 = nixl_agent("bad agent", None) + + buf_size = 256 + addr1 = nixl_utils.malloc_passthru(buf_size * 2) + addr2 = addr1 + buf_size + + agent1_addrs = [(addr1, buf_size, 0), (addr2, buf_size, 0)] + agent1_strings = [(addr1, buf_size, 0, "a"), (addr2, buf_size, 0, "b")] + + agent1_reg_descs = nixl_agent1.get_reg_descs(agent1_strings, "DRAM", True) + + agent1_xfer_descs = nixl_agent1.get_xfer_descs(agent1_addrs, "DRAM", True) + + try: + #null backend not supported + nixl_agent1.backends["UVX"] = 0 + nixl_agent1.register_memory(agent1_reg_descs, backend="UVX") + except nixl_bindings.nixlNotSupportedError as e: + print("Caught you!") + print(e) From ba299fc26a22cfa16ce10832ef8e32bbd1ead600 Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 16:13:09 -0700 Subject: [PATCH 25/32] exception done --- test/python/nixl_exception_test.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/test/python/nixl_exception_test.py b/test/python/nixl_exception_test.py index 05b9ff3d33..293445627b 100755 --- a/test/python/nixl_exception_test.py +++ b/test/python/nixl_exception_test.py @@ -17,7 +17,6 @@ import nixl._utils as nixl_utils from nixl._api import nixl_agent -import nixl._bindings as nixl_bindings if __name__ == "__main__": @@ -35,9 +34,9 @@ agent1_xfer_descs = nixl_agent1.get_xfer_descs(agent1_addrs, "DRAM", True) try: - #null backend not supported + # null backend not supported nixl_agent1.backends["UVX"] = 0 nixl_agent1.register_memory(agent1_reg_descs, backend="UVX") - except nixl_bindings.nixlNotSupportedError as e: + except Exception as e: print("Caught you!") print(e) From bac359e150a8e60611ee09bc9b901c6c39779a4d Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 16:17:31 -0700 Subject: [PATCH 26/32] thank you black --- test/python/nixl_exception_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/python/nixl_exception_test.py b/test/python/nixl_exception_test.py index 293445627b..fdbcb91347 100755 --- a/test/python/nixl_exception_test.py +++ b/test/python/nixl_exception_test.py @@ -19,7 +19,6 @@ from nixl._api import nixl_agent if __name__ == "__main__": - nixl_agent1 = nixl_agent("bad agent", None) buf_size = 256 From 1227931d73643da86bafc074385f5cc70884464e Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 16:22:44 -0700 Subject: [PATCH 27/32] fix nixl_bindings_test --- pyproject.toml | 2 +- src/bindings/python/nixl_bindings.cpp | 2 +- test/python/nixl_bindings_test.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a940c613b7..fb1c7a3e21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ readme = 'README.md' license = 'LICENSE' requires-python = '>=3.9' authors = [ - {name = 'NIXL Developers', email = 'tstamler@nvidia.com'}, + {name = 'Timothy Stamler', email = 'tstamler@nvidia.com'}, ] [tool.meson-python.args] diff --git a/src/bindings/python/nixl_bindings.cpp b/src/bindings/python/nixl_bindings.cpp index 6ef5205e53..8687e78773 100644 --- a/src/bindings/python/nixl_bindings.cpp +++ b/src/bindings/python/nixl_bindings.cpp @@ -398,7 +398,7 @@ PYBIND11_MODULE(_bindings, m) { } throw_nixl_exception(ret); return ret; - }) + }, py::arg("reqh"), py::arg("notif_msg") = std::string("")) .def("getXferStatus", [](nixlAgent &agent, uintptr_t reqh) -> nixl_status_t { nixl_status_t ret = agent.getXferStatus((nixlXferReqH*) reqh); throw_nixl_exception(ret); diff --git a/test/python/nixl_bindings_test.py b/test/python/nixl_bindings_test.py index f51a64ac96..f74bfc8e77 100755 --- a/test/python/nixl_bindings_test.py +++ b/test/python/nixl_bindings_test.py @@ -115,7 +115,7 @@ def test_agent(): print(src_list) print(dst_list) - handle = agent1.createXferReq(src_list, dst_list, name2, noti_str, nixl.NIXL_WRITE) + handle = agent1.createXferReq(nixl.NIXL_WRITE, src_list, dst_list, name2, noti_str) assert handle != 0 print(handle) From c45f74694c56ae2e18947916cbd3f12d6ca8872d Mon Sep 17 00:00:00 2001 From: Timothy Stamler Date: Sun, 16 Mar 2025 16:29:21 -0700 Subject: [PATCH 28/32] toml --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index fb1c7a3e21..1f8407c46a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,9 @@ license = 'LICENSE' requires-python = '>=3.9' authors = [ {name = 'Timothy Stamler', email = 'tstamler@nvidia.com'}, + {name = 'Moein Khazraee', email = 'moein@nvidia.com'}, + {name = 'Vishwanath Venkatesan', email = 'vvenkatesan@nvidia.com'}, + {name = 'Omri Kahalon', email = 'omrik@nvidia.com'}, ] [tool.meson-python.args] From 2286139b0f072510a9f06e68020e9078a3ba7bd6 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 17:26:25 -0700 Subject: [PATCH 29/32] more corner case handling in _api, and renamed the wrapper test. --- src/api/python/_api.py | 8 +- test/python/nixl_api_test.py | 161 +++++++++++++++++++++++++++++++ test/python/nixl_wrapper_test.py | 3 +- 3 files changed, 169 insertions(+), 3 deletions(-) create mode 100755 test/python/nixl_api_test.py diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 46723ed0da..b05c0727a2 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -28,7 +28,7 @@ def __init(self, enable_prog_thread=True, backends=["UCX", "GDS"]): class nixl_agent: - def __init__(self, agent_name, nixl_config=None, instantiate_all=False): + def __init__(self, agent_name, nixl_config=None, instantiate_all=True): # Set agent config and instantiate an agent if nixl_config: agent_config = nixlBind.nixlAgentConfig(nixl_config.enable_pthread) @@ -45,6 +45,10 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): self.plugin_list = self.agent.getAvailPlugins() if len(self.plugin_list) == 0: print("No plugins available, cannot start transfers!") + raise RuntimeError("No plugins available for NIXL, cannot start transfers!") + else: + print("Found plugins during Agent initialization:") + print(self.plugin_list) self.plugin_b_options = {} self.plugin_mem_types = {} @@ -69,6 +73,8 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=False): # TODO: populate init from default parameters, or define a set of params in python for plugin in self.plugin_list: self.backends[plugin] = self.agent.createBackend(plugin, init) + else: + print("No backends were created during agent creation.") for backend in self.backends: (backend_options, mem_types) = self.agent.getBackendParams(backend) diff --git a/test/python/nixl_api_test.py b/test/python/nixl_api_test.py new file mode 100755 index 0000000000..2062a7aca1 --- /dev/null +++ b/test/python/nixl_api_test.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +import nixl._utils as nixl_utils +from nixl._api import nixl_agent + +if __name__ == "__main__": + buf_size = 256 + # Allocate memory and register with NIXL + nixl_agent1 = nixl_agent("target") + + plugin_list = nixl_agent1.get_plugin_list() + assert "UCX" in plugin_list + + print(nixl_agent1.get_plugin_mem_types("UCX")) + print(nixl_agent1.get_plugin_params("UCX")) + + nixl_agent1.create_backend("UCX") + + print(nixl_agent1.get_backend_mem_types("UCX")) + print(nixl_agent1.get_backend_params("UCX")) + + addr1 = nixl_utils.malloc_passthru(buf_size * 2) + addr2 = addr1 + buf_size + + agent1_addrs = [(addr1, buf_size, 0), (addr2, buf_size, 0)] + agent1_strings = [(addr1, buf_size, 0, "a"), (addr2, buf_size, 0, "b")] + + agent1_reg_descs = nixl_agent1.get_reg_descs(agent1_strings, "DRAM", is_sorted=True) + agent1_xfer_descs = nixl_agent1.get_xfer_descs(agent1_addrs, "DRAM", is_sorted=True) + + # Just for tensor test + tensors = [torch.zeros(10, dtype=torch.float32) for _ in range(2)] + agent1_tensor_reg_descs = nixl_agent1.get_reg_descs(tensors) + agent1_tensor_xfer_descs = nixl_agent1.get_xfer_descs(tensors) + + assert nixl_agent1.register_memory(agent1_reg_descs) is not None + + nixl_agent2 = nixl_agent("initiator", None) + nixl_agent2.create_backend("UCX") + addr3 = nixl_utils.malloc_passthru(buf_size * 2) + addr4 = addr3 + buf_size + + agent2_addrs = [(addr3, buf_size, 0), (addr4, buf_size, 0)] + agent2_strings = [(addr3, buf_size, 0, "a"), (addr4, buf_size, 0, "b")] + + agent2_reg_descs = nixl_agent2.get_reg_descs(agent2_strings, "DRAM", is_sorted=True) + agent2_xfer_descs = nixl_agent2.get_xfer_descs(agent2_addrs, "DRAM", is_sorted=True) + + agent2_descs = nixl_agent2.register_memory(agent2_reg_descs, is_sorted=True) + assert agent2_descs is not None + + # Exchange metadata + meta = nixl_agent1.get_agent_metadata() + remote_name = nixl_agent2.add_remote_agent(meta) + print("Loaded name from metadata:", remote_name) + + serdes = nixl_agent1.get_serialized_descs(agent1_reg_descs) + src_descs_recvd = nixl_agent2.deserialize_descs(serdes) + assert src_descs_recvd == agent1_reg_descs + + # initialize transfer mode + xfer_handle_1 = nixl_agent2.initialize_xfer( + agent2_xfer_descs, agent1_xfer_descs, remote_name, "UUID1", "READ" + ) + if not xfer_handle_1: + print("Creating transfer failed.") + exit() + + state = nixl_agent2.transfer(xfer_handle_1) + assert state != "ERR" + + target_done = False + init_done = False + + while (not init_done) or (not target_done): + if not init_done: + state = nixl_agent2.check_xfer_state(xfer_handle_1) + if state == "ERR": + print("Transfer got to Error state.") + exit() + elif state == "DONE": + init_done = True + print("Initiator done") + + if not target_done: + if nixl_agent1.check_remote_xfer_done("initiator", "UUID1"): + target_done = True + print("Target done") + + # prep transfer mode + local_prep_handle = nixl_agent2.prep_xfer_side( + "", [(addr3, buf_size, 0), (addr4, buf_size, 0)], "DRAM", True + ) + remote_prep_handle = nixl_agent2.prep_xfer_side( + remote_name, agent1_xfer_descs, "DRAM" + ) + + assert local_prep_handle != 0 + assert remote_prep_handle != 0 + + xfer_handle_2 = nixl_agent2.make_prepped_xfer( + local_prep_handle, [0, 1], remote_prep_handle, [1, 0], "UUID2", "WRITE" + ) + if not local_prep_handle or not remote_prep_handle: + print("Preparing transfer side handles failed.") + exit() + + if not xfer_handle_2: + print("Make prepped transfer failed.") + exit() + + state = nixl_agent2.transfer(xfer_handle_2) + assert state != "ERR" + + target_done = False + init_done = False + + while (not init_done) or (not target_done): + if not init_done: + state = nixl_agent2.check_xfer_state(xfer_handle_2) + if state == "ERR": + print("Transfer got to Error state.") + exit() + elif state == "DONE": + init_done = True + print("Initiator done") + + if not target_done: + if nixl_agent1.check_remote_xfer_done("initiator", "UUID2"): + target_done = True + print("Target done") + + nixl_agent2.abort_xfer(xfer_handle_1) + nixl_agent2.abort_xfer(xfer_handle_2) + nixl_agent2.delete_xfer_side(local_prep_handle) + nixl_agent2.delete_xfer_side(remote_prep_handle) + nixl_agent2.remove_remote_agent("target") + nixl_agent1.deregister_memory(agent1_reg_descs) + nixl_agent2.deregister_memory(agent2_reg_descs) + + nixl_utils.free_passthru(addr1) + nixl_utils.free_passthru(addr3) + + print("Test Complete.") diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index f3e5be075e..2062a7aca1 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -26,8 +26,7 @@ nixl_agent1 = nixl_agent("target") plugin_list = nixl_agent1.get_plugin_list() - # assert "UCX" in plugin_list # not supported in static linking mode yet - print(plugin_list) + assert "UCX" in plugin_list print(nixl_agent1.get_plugin_mem_types("UCX")) print(nixl_agent1.get_plugin_params("UCX")) From 0c733376aa6e7554152b23c5b79923ff842524fa Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 18:12:41 -0700 Subject: [PATCH 30/32] Made default behavior of agent to load UCX only --- src/api/python/_api.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index b05c0727a2..00f929c3a2 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -21,17 +21,17 @@ class nixl_config: - def __init(self, enable_prog_thread=True, backends=["UCX", "GDS"]): + def __init(self, enable_prog_thread=True, backends=["UCX"]): # TODO: add backend init parameters self.backends = backends self.enable_pthread = enable_prog_thread class nixl_agent: - def __init__(self, agent_name, nixl_config=None, instantiate_all=True): + def __init__(self, agent_name, nixl_conf=None, instantiate_all=False): # Set agent config and instantiate an agent - if nixl_config: - agent_config = nixlBind.nixlAgentConfig(nixl_config.enable_pthread) + if nixl_conf: + agent_config = nixlBind.nixlAgentConfig(nixl_conf.enable_pthread) else: agent_config = nixlBind.nixlAgentConfig(True) self.agent = nixlBind.nixlAgent(agent_name, agent_config) @@ -46,9 +46,6 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=True): if len(self.plugin_list) == 0: print("No plugins available, cannot start transfers!") raise RuntimeError("No plugins available for NIXL, cannot start transfers!") - else: - print("Found plugins during Agent initialization:") - print(self.plugin_list) self.plugin_b_options = {} self.plugin_mem_types = {} @@ -58,9 +55,21 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=True): self.plugin_mem_types[plugin] = mem_types init = {} - if nixl_config: - for bknd in nixl_config.backends: - # TODO: populate init from nixl_config when added + if nixl_conf and instantiate_all: + instantiate_all = False + print( + "Overruling instantiate_all based on the provided config in agent creation." + ) + + if instantiate_all: + # TODO: populate init from default parameters, or define a set of params in python + for plugin in self.plugin_list: + self.backends[plugin] = self.agent.createBackend(plugin, init) + else: + if not nixl_conf: + nixl_conf = nixl_config() + for bknd in nixl_conf.backends: + # TODO: populate init from nixl_conf when added if bknd not in self.plugin_list: print( "Skipping backend registration", @@ -69,12 +78,6 @@ def __init__(self, agent_name, nixl_config=None, instantiate_all=True): ) else: self.backends[bknd] = self.agent.createBackend(bknd, init) - elif instantiate_all: - # TODO: populate init from default parameters, or define a set of params in python - for plugin in self.plugin_list: - self.backends[plugin] = self.agent.createBackend(plugin, init) - else: - print("No backends were created during agent creation.") for backend in self.backends: (backend_options, mem_types) = self.agent.getBackendParams(backend) From c34f38b3194c701c828bda5f8f2cc8e5fc623673 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 19:08:30 -0700 Subject: [PATCH 31/32] More fixes --- src/api/python/_api.py | 16 ++++++++++------ test/python/nixl_api_test.py | 6 +++--- test/python/nixl_wrapper_test.py | 6 +++--- 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 00f929c3a2..216ff2b6e9 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -20,8 +20,8 @@ import nixl._bindings as nixlBind -class nixl_config: - def __init(self, enable_prog_thread=True, backends=["UCX"]): +class nixl_agent_config: + def __init__(self, enable_prog_thread=True, backends=["UCX"]): # TODO: add backend init parameters self.backends = backends self.enable_pthread = enable_prog_thread @@ -67,7 +67,7 @@ def __init__(self, agent_name, nixl_conf=None, instantiate_all=False): self.backends[plugin] = self.agent.createBackend(plugin, init) else: if not nixl_conf: - nixl_conf = nixl_config() + nixl_conf = nixl_agent_config() for bknd in nixl_conf.backends: # TODO: populate init from nixl_conf when added if bknd not in self.plugin_list: @@ -80,7 +80,9 @@ def __init__(self, agent_name, nixl_conf=None, instantiate_all=False): self.backends[bknd] = self.agent.createBackend(bknd, init) for backend in self.backends: - (backend_options, mem_types) = self.agent.getBackendParams(backend) + (backend_options, mem_types) = self.agent.getBackendParams( + self.backends[backend] + ) self.backend_mems[backend] = mem_types self.backend_options[backend] = backend_options @@ -152,8 +154,7 @@ def register_memory( # based on backend type and mem_type, figure what registrations are meaningful if backend: - print("trying backend") - print(self.backends[backend]) + print("Registering backend", self.backends[backend]) ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported @@ -301,6 +302,9 @@ def initialize_xfer( handle = self.agent.createXferReq( op, local_descs, remote_descs, remote_agent, notif_msg ) + + if handle == 0: + return None return handle # In case of error it will be None else: return None diff --git a/test/python/nixl_api_test.py b/test/python/nixl_api_test.py index 2062a7aca1..b9e96b7179 100755 --- a/test/python/nixl_api_test.py +++ b/test/python/nixl_api_test.py @@ -28,13 +28,14 @@ plugin_list = nixl_agent1.get_plugin_list() assert "UCX" in plugin_list + print("Plugin parameters") print(nixl_agent1.get_plugin_mem_types("UCX")) print(nixl_agent1.get_plugin_params("UCX")) - nixl_agent1.create_backend("UCX") - + print("\nLoaded backend parameters") print(nixl_agent1.get_backend_mem_types("UCX")) print(nixl_agent1.get_backend_params("UCX")) + print() addr1 = nixl_utils.malloc_passthru(buf_size * 2) addr2 = addr1 + buf_size @@ -53,7 +54,6 @@ assert nixl_agent1.register_memory(agent1_reg_descs) is not None nixl_agent2 = nixl_agent("initiator", None) - nixl_agent2.create_backend("UCX") addr3 = nixl_utils.malloc_passthru(buf_size * 2) addr4 = addr3 + buf_size diff --git a/test/python/nixl_wrapper_test.py b/test/python/nixl_wrapper_test.py index 2062a7aca1..b9e96b7179 100755 --- a/test/python/nixl_wrapper_test.py +++ b/test/python/nixl_wrapper_test.py @@ -28,13 +28,14 @@ plugin_list = nixl_agent1.get_plugin_list() assert "UCX" in plugin_list + print("Plugin parameters") print(nixl_agent1.get_plugin_mem_types("UCX")) print(nixl_agent1.get_plugin_params("UCX")) - nixl_agent1.create_backend("UCX") - + print("\nLoaded backend parameters") print(nixl_agent1.get_backend_mem_types("UCX")) print(nixl_agent1.get_backend_params("UCX")) + print() addr1 = nixl_utils.malloc_passthru(buf_size * 2) addr2 = addr1 + buf_size @@ -53,7 +54,6 @@ assert nixl_agent1.register_memory(agent1_reg_descs) is not None nixl_agent2 = nixl_agent("initiator", None) - nixl_agent2.create_backend("UCX") addr3 = nixl_utils.malloc_passthru(buf_size * 2) addr4 = addr3 + buf_size From 27ea328586e2f86aad0ede56a12f5efead2ed211 Mon Sep 17 00:00:00 2001 From: Moein Khazraee Date: Sun, 16 Mar 2025 19:26:33 -0700 Subject: [PATCH 32/32] Minor implementation improvement --- src/api/python/_api.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/src/api/python/_api.py b/src/api/python/_api.py index 216ff2b6e9..6158329a91 100644 --- a/src/api/python/_api.py +++ b/src/api/python/_api.py @@ -29,11 +29,16 @@ def __init__(self, enable_prog_thread=True, backends=["UCX"]): class nixl_agent: def __init__(self, agent_name, nixl_conf=None, instantiate_all=False): + if nixl_conf and instantiate_all: + instantiate_all = False + print( + "Ignoring instantiate_all based on the provided config in agent creation." + ) + if not nixl_conf: + nixl_conf = nixl_agent_config() # Using defaults set in nixl_agent_config + # Set agent config and instantiate an agent - if nixl_conf: - agent_config = nixlBind.nixlAgentConfig(nixl_conf.enable_pthread) - else: - agent_config = nixlBind.nixlAgentConfig(True) + agent_config = nixlBind.nixlAgentConfig(nixl_conf.enable_pthread) self.agent = nixlBind.nixlAgent(agent_name, agent_config) self.name = agent_name @@ -55,19 +60,12 @@ def __init__(self, agent_name, nixl_conf=None, instantiate_all=False): self.plugin_mem_types[plugin] = mem_types init = {} - if nixl_conf and instantiate_all: - instantiate_all = False - print( - "Overruling instantiate_all based on the provided config in agent creation." - ) if instantiate_all: # TODO: populate init from default parameters, or define a set of params in python for plugin in self.plugin_list: self.backends[plugin] = self.agent.createBackend(plugin, init) else: - if not nixl_conf: - nixl_conf = nixl_agent_config() for bknd in nixl_conf.backends: # TODO: populate init from nixl_conf when added if bknd not in self.plugin_list: @@ -154,7 +152,6 @@ def register_memory( # based on backend type and mem_type, figure what registrations are meaningful if backend: - print("Registering backend", self.backends[backend]) ret = self.agent.registerMem(reg_descs, self.backends[backend]) else: # TODO: rely on underlying capability to register with all when supported