diff --git a/.gitlab/test_ep.sh b/.gitlab/test_ep.sh index b5fafa059e..56378cce07 100755 --- a/.gitlab/test_ep.sh +++ b/.gitlab/test_ep.sh @@ -64,7 +64,7 @@ run_elastic_test() { unset UCX_NET_DEVICES # let UCX auto-select GPU-capable transport # Force NVLink-only transports. if [[ "$extra_flags" != *--disable-ll-nvlink* ]]; then - export UCX_TLS=cuda_copy,cuda_ipc,sm,self + export UCX_TLS=^rc_gda fi PYTHONPATH="${NIXL_BUILD_DIR}/${EP_SRC_DIR}:${EP_SRC_DIR}/tests:${EP_SRC_DIR}/tests/elastic${PYTHONPATH:+:$PYTHONPATH}" \ timeout 300 python3 ${EP_SRC_DIR}/tests/elastic/elastic.py \ diff --git a/examples/device/ep/csrc/nixl_ep.cpp b/examples/device/ep/csrc/nixl_ep.cpp index de5db4a0d2..5594c0aed3 100644 --- a/examples/device/ep/csrc/nixl_ep.cpp +++ b/examples/device/ep/csrc/nixl_ep.cpp @@ -1407,7 +1407,8 @@ void Buffer::_nixl_agent_init() { const char* num_channels_env = std::getenv("NIXL_EP_NUM_CHANNELS"); init_params["ucx_num_device_channels"] = num_channels_env ? num_channels_env : "4"; - init_params["ucx_error_handling_mode"] = "none"; + init_params["ucx_error_handling_mode"] = "peer"; + init_params["ucx_ep_close_force"] = "yes"; init_params["num_workers"] = std::to_string(1); nixlBackendH* ucx_backend = nullptr; diff --git a/src/plugins/ucx/ucx_backend.cpp b/src/plugins/ucx/ucx_backend.cpp index 8f8835c32a..39089e1288 100644 --- a/src/plugins/ucx/ucx_backend.cpp +++ b/src/plugins/ucx/ucx_backend.cpp @@ -32,6 +32,15 @@ #include "absl/strings/str_split.h" #include +namespace { +[[nodiscard]] uint32_t +epCloseFlags(const nixl_b_params_t *custom_params) { + return nixl::getBackendParamDefaulted(custom_params, "ucx_ep_close_force", false) ? + UCP_EP_CLOSE_FLAG_FORCE : + 0; +} +} // namespace + /**************************************** * Backend request management *****************************************/ @@ -823,6 +832,8 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params) err_handling_mode = ucx_err_mode_from_string(*opt); } + const uint32_t ep_close_flags = epCloseFlags(custom_params); + const auto engine_config = nixl::getBackendParamDefaulted(custom_params, "engine_config", std::string()); @@ -836,7 +847,7 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params) uc->warnAboutHardwareSupportMismatch(); for (size_t i = 0; i < num_workers; i++) { - uws.emplace_back(std::make_unique(*uc, err_handling_mode)); + uws.emplace_back(std::make_unique(*uc, err_handling_mode, ep_close_flags)); } auto &uw = uws.front(); diff --git a/src/plugins/ucx/ucx_utils.cpp b/src/plugins/ucx/ucx_utils.cpp index a0ad10bd56..a2ed24551b 100644 --- a/src/plugins/ucx/ucx_utils.cpp +++ b/src/plugins/ucx/ucx_utils.cpp @@ -121,9 +121,10 @@ nixlUcxEp::setState(nixl::ucx::ep_state_t new_state) { } nixl_status_t -nixlUcxEp::closeImpl(ucp_ep_close_flags_t flags) { +nixlUcxEp::closeImpl() { ucs_status_ptr_t request = nullptr; - ucp_request_param_t req_param = {.op_attr_mask = UCP_OP_ATTR_FIELD_FLAGS, .flags = flags}; + const ucp_request_param_t req_param = {.op_attr_mask = UCP_OP_ATTR_FIELD_FLAGS, + .flags = closeFlags_}; switch (state) { case nixl::ucx::ep_state_t::UNINITIALIZED: @@ -156,7 +157,11 @@ nixlUcxEp::closeImpl(ucp_ep_close_flags_t flags) { std::terminate(); } -nixlUcxEp::nixlUcxEp(ucp_worker_h worker, void *addr, ucp_err_handling_mode_t err_handling_mode) { +nixlUcxEp::nixlUcxEp(ucp_worker_h worker, + void *addr, + ucp_err_handling_mode_t err_handling_mode, + uint32_t close_flags) + : closeFlags_{close_flags} { ucp_ep_params_t ep_params; nixl_status_t status; @@ -185,7 +190,7 @@ nixlUcxEp::~nixlUcxEp() { nixl_status_t nixlUcxEp::disconnect_nb() { - nixl_status_t status = closeImpl(ucp_ep_close_flags_t(0)); + const nixl_status_t status = closeImpl(); // At step of disconnect we can ignore the remote disconnect error. return (status == NIXL_ERR_REMOTE_DISCONNECT) ? NIXL_SUCCESS : status; @@ -517,9 +522,12 @@ nixlUcxWorker::createUcpWorker(const nixlUcxContext &ctx) { return worker; } -nixlUcxWorker::nixlUcxWorker(const nixlUcxContext &ctx, ucp_err_handling_mode_t err_handling_mode) +nixlUcxWorker::nixlUcxWorker(const nixlUcxContext &ctx, + ucp_err_handling_mode_t err_handling_mode, + uint32_t ep_close_flags) : worker(createUcpWorker(ctx), &ucp_worker_destroy), - err_handling_mode_(err_handling_mode) {} + err_handling_mode_(err_handling_mode), + epCloseFlags_(ep_close_flags) {} std::string nixlUcxWorker::epAddr() { @@ -540,7 +548,7 @@ nixlUcxWorker::epAddr() { std::unique_ptr nixlUcxWorker::connect(void *addr, std::size_t size) { try { - return std::make_unique(worker.get(), addr, err_handling_mode_); + return std::make_unique(worker.get(), addr, err_handling_mode_, epCloseFlags_); } catch (const std::exception &e) { NIXL_ERROR << "UCX endpoint create failed: " << e.what(); diff --git a/src/plugins/ucx/ucx_utils.h b/src/plugins/ucx/ucx_utils.h index f22fd7b327..6c928461bd 100644 --- a/src/plugins/ucx/ucx_utils.h +++ b/src/plugins/ucx/ucx_utils.h @@ -45,11 +45,12 @@ class nixlUcxEp { private: ucp_ep_h eph{nullptr}; nixl::ucx::ep_state_t state = nixl::ucx::ep_state_t::UNINITIALIZED; + const uint32_t closeFlags_; void setState(nixl::ucx::ep_state_t new_state); nixl_status_t - closeImpl(ucp_ep_close_flags_t flags); + closeImpl(); /* Connection */ nixl_status_t @@ -67,7 +68,10 @@ class nixlUcxEp { return nixl::ucx::toNixlStatus(state); } - nixlUcxEp(ucp_worker_h worker, void *addr, ucp_err_handling_mode_t err_handling_mode); + nixlUcxEp(ucp_worker_h worker, + void *addr, + ucp_err_handling_mode_t err_handling_mode, + uint32_t close_flags); ~nixlUcxEp(); nixlUcxEp(const nixlUcxEp &) = delete; nixlUcxEp & @@ -187,7 +191,8 @@ class nixlUcxWorker { public: explicit nixlUcxWorker( const nixlUcxContext &, - ucp_err_handling_mode_t ucp_err_handling_mode = UCP_ERR_HANDLING_MODE_NONE); + ucp_err_handling_mode_t ucp_err_handling_mode = UCP_ERR_HANDLING_MODE_NONE, + uint32_t ep_close_flags = 0); nixlUcxWorker(nixlUcxWorker &&) = delete; nixlUcxWorker(const nixlUcxWorker &) = delete; @@ -241,7 +246,8 @@ class nixlUcxWorker { createUcpWorker(const nixlUcxContext &); const std::unique_ptr worker; - ucp_err_handling_mode_t err_handling_mode_; + const ucp_err_handling_mode_t err_handling_mode_; + const uint32_t epCloseFlags_; }; [[nodiscard]] nixl_b_params_t