Skip to content

[5/N] PLUGINS/UCX: Simplify notification #2341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
iyastreb merged 7 commits into from
Oct 6, 2026
Merged

[5/N] PLUGINS/UCX: Simplify notification #2341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d0b9d02
[5/N] PLUGINS/UCX: Simplify notification
iyastreb Oct 5, 2026
f422d7b
CR fixes from #2293
iyastreb Oct 5, 2026
5ac0350
Merge branch 'main' into iyastreb/tp-5-notification
iyastreb Oct 5, 2026
2755a80
CR fix
iyastreb Oct 5, 2026
ef46d22
Merge branch 'main' into iyastreb/tp-5-notification
iyastreb Oct 5, 2026
d7ea9dc
Removed optional from notif string
iyastreb Oct 6, 2026
5baffe2
Merge branch 'main' into iyastreb/tp-5-notification
iyastreb Oct 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 26 additions & 57 deletions src/plugins/ucx/ucx_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,11 +44,14 @@ sglEnabledFromConfig() {
return false;
#endif
}
} // namespace

// A transfer to a single endpoint posts at most three requests:
// one data request, one flush request, and one notification request.
constexpr size_t single_ep_request_count = 3;
[[nodiscard]] ucp_err_handling_mode_t
errHandlingModeFromParams(const nixl_b_params_t *custom_params) {
const auto opt = nixl::getBackendParamOptional<std::string>(
custom_params, std::string(nixl_ucx_err_handling_param_name));
return opt ? ucx_err_mode_from_string(*opt) : UCP_ERR_HANDLING_MODE_PEER;
}
} // namespace

/****************************************
* Constructor/Destructor
Expand All @@ -71,6 +74,7 @@ nixlUcxEngine::create(const nixlBackendInitParams &init_params) {

nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params, size_t num_dedicated_workers)
: nixlBackendEngine(&init_params),
errHandlingMode_(errHandlingModeFromParams(init_params.customParams)),
sharedWorkerIndex_(1),
sglEnabled_(sglEnabledFromConfig()) {
std::vector<std::string> devs; /* Empty vector */
Expand All @@ -89,14 +93,6 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params, size_t nu
const size_t num_device_channels =
nixl::getBackendParamDefaulted(custom_params, "ucx_num_device_channels", 4u);


if (const auto opt = nixl::getBackendParamOptional<std::string>(
custom_params, std::string(nixl_ucx_err_handling_param_name))) {
errHandlingMode_ = ucx_err_mode_from_string(*opt);
} else {
errHandlingMode_ = UCP_ERR_HANDLING_MODE_PEER;
}

const auto engine_config =
nixl::getBackendParamDefaulted(custom_params, "engine_config", std::string());

Expand Down Expand Up @@ -456,18 +452,17 @@ nixlUcxEngine::sendXferSgl(nixlBackendReqH *handle) const {
const ucx_connection_ptr_t &conn = sgl.conn();

auto &ep = conn->getEp(int_handle->getWorkerId());

int_handle->reserve(single_ep_request_count);
int_handle->init(conn, *ep);

nixlUcxReq req;
const nixl_status_t post_ret = sgl.post(*ep, req);
if (int_handle->append(post_ret, req, conn) != NIXL_SUCCESS) {
if (int_handle->append(post_ret, req) != NIXL_SUCCESS) {
return post_ret;
}

nixlUcxReq flush_req;
const nixl_status_t flush_ret = ep->flushEp(flush_req);
if (int_handle->append(flush_ret, flush_req, conn) != NIXL_SUCCESS) {
if (int_handle->append(flush_ret, flush_req) != NIXL_SUCCESS) {
return flush_ret;
}

Expand Down Expand Up @@ -499,11 +494,10 @@ nixlUcxEngine::sendXferRange(const nixl_xfer_op_t &operation,
}
#endif

int_handle->reserve(single_ep_request_count);

const ucx_connection_ptr_t &conn =
static_cast<nixlUcxPublicMetadata *>(remote[start_idx].metadataP)->conn;
auto &ep = conn->getEp(worker_id);
int_handle->init(conn, *ep);

nixl_status_t status = NIXL_SUCCESS;
nixlUcxReq pending_req = nullptr;
Expand Down Expand Up @@ -542,7 +536,7 @@ nixlUcxEngine::sendXferRange(const nixl_xfer_op_t &operation,
status = NIXL_IN_PROG;
}

if (int_handle->append(status, pending_req, conn) != NIXL_SUCCESS) {
if (int_handle->append(status, pending_req) != NIXL_SUCCESS) {
return status;
}

Expand All @@ -552,7 +546,7 @@ nixlUcxEngine::sendXferRange(const nixl_xfer_op_t &operation,
*/
nixlUcxReq flush_req;
const nixl_status_t flush_ret = ep->flushEp(flush_req);
if (int_handle->append(flush_ret, flush_req, conn) != NIXL_SUCCESS) {
if (int_handle->append(flush_ret, flush_req) != NIXL_SUCCESS) {
return flush_ret;
}

Expand Down Expand Up @@ -588,16 +582,14 @@ nixlUcxEngine::postXfer(const nixl_xfer_op_t &operation,
if (opt_args && opt_args->hasNotif) {
if (ret == NIXL_SUCCESS) {
nixlUcxReq req;
const auto rmd = static_cast<nixlUcxPublicMetadata *>(remote[0].metadataP);
const nixlUcxEp &ep = *rmd->conn->getEp(int_handle->getWorkerId());
ret = notifSendPriv(remote_agent, opt_args->notifMsg, ep, &req);
if (int_handle->append(ret, req, rmd->conn) != NIXL_SUCCESS) {
ret = notifSendPriv(remote_agent, opt_args->notifMsg, int_handle->getEp(), &req);
if (int_handle->append(ret, req) != NIXL_SUCCESS) {
return ret;
}

ret = int_handle->status();
} else if (ret == NIXL_IN_PROG) {
int_handle->notif.emplace(remote_agent, buildNotif(opt_args->notifMsg));
int_handle->notif = buildNotif(opt_args->notifMsg);
Comment thread
ColinNV marked this conversation as resolved.
}
}

Expand All @@ -609,27 +601,20 @@ nixl_status_t nixlUcxEngine::checkXfer (nixlBackendReqH* handle) const
const auto int_handle = static_cast<nixlUcxBackendReqH *>(handle);
const nixl_status_t handle_status = int_handle->status();

if ((handle_status == NIXL_IN_PROG) || !int_handle->notif) {
if ((handle_status == NIXL_IN_PROG) || int_handle->notif.empty()) {
return handle_status;
}

nixlUcxBackendReqH::Notif notif(std::move(int_handle->notif).value());
int_handle->notif.reset();

if (handle_status != NIXL_SUCCESS) [[unlikely]] {
int_handle->notif.clear();
return handle_status;
}

const ucx_connection_ptr_t conn = getConnection(notif.agent);
if (!conn) [[unlikely]] {
return NIXL_ERR_NOT_FOUND;
}

nixlUcxReq req;
const nixlUcxEp &ep = *conn->getEp(int_handle->getWorkerId());
const nixl_status_t status = sendNotif(std::move(notif.msg), ep, &req);
const nixl_status_t status = sendNotif(std::move(int_handle->notif), int_handle->getEp(), &req);
int_handle->notif.clear();

if (int_handle->append(status, req, conn) != NIXL_SUCCESS) {
if (int_handle->append(status, req) != NIXL_SUCCESS) {
return status;
}

Expand Down Expand Up @@ -667,35 +652,19 @@ nixlUcxEngine::progressLoop() {
* Notifications
*****************************************/

std::unique_ptr<std::string>
std::string
nixlUcxEngine::buildNotif(const std::string &msg) const {
nixlSerDes ser_des;

ser_des.addStr("name", localAgent);
ser_des.addStr("msg", msg);
// TODO: replace with mpool for performance
return std::make_unique<std::string>(ser_des.exportStr());
return ser_des.exportStr();
}

nixl_status_t
nixlUcxEngine::sendNotif(std::unique_ptr<std::string> &&msg, const nixlUcxEp &ep, nixlUcxReq *req) {
std::string *buffer = msg.release();
auto cleanup = [buffer, req](void *completed_request, void *ptr) {
delete buffer;
if ((req == nullptr) && (completed_request != nullptr)) {
/* Caller is not interested in the request, free it */
ucp_request_free(completed_request);
}
};

return ep.sendAm(nixl::ucx::am_cb_op_t::NOTIF_STR,
nullptr,
0,
buffer->data(),
buffer->size(),
UCP_AM_SEND_FLAG_EAGER,
req,
std::move(cleanup));
nixlUcxEngine::sendNotif(std::string &&msg, const nixlUcxEp &ep, nixlUcxReq *req) {
return ep.sendAm(nixl::ucx::am_cb_op_t::NOTIF_STR, std::move(msg), UCP_AM_SEND_FLAG_EAGER, req);
}

nixl_status_t
Expand Down
6 changes: 3 additions & 3 deletions src/plugins/ucx/ucx_backend.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,11 +268,11 @@ class nixlUcxEngine : public nixlBackendEngine {
size_t length,
const ucp_am_recv_param_t *param);

[[nodiscard]] std::unique_ptr<std::string>
[[nodiscard]] std::string
buildNotif(const std::string &msg) const;

[[nodiscard]] static nixl_status_t
sendNotif(std::unique_ptr<std::string> &&msg, const nixlUcxEp &ep, nixlUcxReq *req);
sendNotif(std::string &&msg, const nixlUcxEp &ep, nixlUcxReq *req);

nixl_status_t
notifSendPriv(const std::string &remote_agent,
Expand Down Expand Up @@ -304,7 +304,7 @@ class nixlUcxEngine : public nixlBackendEngine {
std::unique_ptr<nixlUcxContext> uc;
std::vector<std::unique_ptr<nixlUcxWorker>> workers_;
size_t numSharedWorkers_;
ucp_err_handling_mode_t errHandlingMode_;
const ucp_err_handling_mode_t errHandlingMode_;
std::string workerAddr;
mutable std::atomic<size_t> sharedWorkerIndex_;
const bool sglEnabled_;
Expand Down
58 changes: 33 additions & 25 deletions src/plugins/ucx/ucx_backend_req.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,9 @@

class nixlUcxBackendReqH : public nixlBackendReqH {
public:
// Notification to be sent after completion of all requests
struct Notif {
const std::string agent;
std::unique_ptr<std::string> msg;

Notif(const std::string &remote_agent, std::unique_ptr<std::string> msg)
: agent(remote_agent),
msg(std::move(msg)) {}
};

std::optional<Notif> notif;
// Notification to be sent over the bound connection after completion of all requests.
// Empty if there is no pending notification.
std::string notif;

#ifdef HAVE_UCX_SGL_API
std::optional<nixl::ucx::sglXfer> sgl;
Expand All @@ -52,13 +44,21 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
}

void
reserve(size_t size) {
requests_.reserve(size);
NIXL_ASSERT(conn_ == nullptr);
init(const ucx_connection_ptr_t &conn, const nixlUcxEp &ep) {
NIXL_ASSERT(requests_.empty());
requests_.reserve(max_requests);
conn_ = conn;
ep_ = &ep;
}

[[nodiscard]] const nixlUcxEp &
getEp() const {
NIXL_ASSERT(ep_ != nullptr);
return *ep_;
}

[[nodiscard]] nixl_status_t
append(nixl_status_t status, nixlUcxReq req, const ucx_connection_ptr_t &conn) {
append(nixl_status_t status, nixlUcxReq req) {
if (status == NIXL_IN_PROG) [[likely]] {
requests_.push_back(req);
} else if (status != NIXL_SUCCESS) {
Expand All @@ -67,8 +67,6 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
return status;
}

NIXL_ASSERT(conn_ == nullptr || conn_ == conn);
conn_ = conn;
return NIXL_SUCCESS;
}

Expand All @@ -89,15 +87,13 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
}
worker_->reqRelease(req);
}
requests_.clear();
conn_.reset();
reset();
Comment thread
iyastreb marked this conversation as resolved.
}

[[nodiscard]] virtual nixl_status_t
status() {
if (requests_.empty()) {
/* No pending transmissions */
conn_.reset();
return NIXL_SUCCESS;
}

Expand Down Expand Up @@ -133,9 +129,6 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
}

requests_.resize(incomplete_reqs);
if (requests_.empty()) {
conn_.reset();
}
return out_ret;
}

Expand All @@ -157,14 +150,29 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
}

private:
// A post to a single endpoint issues at most three requests:
// one data request, one flush request, and one notification request.
static constexpr size_t max_requests = 3;

void
reset() noexcept {
requests_.clear();
conn_.reset();
ep_ = nullptr;
}

[[nodiscard]] nixl_status_t
checkConnection(const nixl_status_t status = NIXL_SUCCESS) const {
NIXL_ASSERT(conn_ != nullptr);
const nixl_status_t conn_status = conn_->getEp(getWorkerId())->checkTxState();
NIXL_ASSERT(ep_ != nullptr);
const nixl_status_t conn_status = ep_->checkTxState();
return (conn_status != NIXL_SUCCESS) ? conn_status : status;
}

// Keeps the connection (which owns the endpoint) alive for the lifetime
// of the request handle.
ucx_connection_ptr_t conn_;
// Resolved endpoint over which data and notifications are sent.
const nixlUcxEp *ep_ = nullptr;
Comment thread
iyastreb marked this conversation as resolved.
std::vector<nixlUcxReq> requests_;
nixlUcxWorker *worker_ = nullptr;
};
Expand Down
14 changes: 8 additions & 6 deletions src/plugins/ucx/ucx_thread_pool_engine.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,8 +271,6 @@ class nixlUcxCompositeBackendReqH : public nixlUcxBackendReqH {
std::shared_ptr<nixlUcxBackendSharedState> sharedState_;
};

} // namespace

class nixlUcxDedicatedThread final : public nixlUcxThread {
public:
nixlUcxDedicatedThread(nixlUcxEngine *engine, nixlUcxDedicatedWorker &worker);
Expand Down Expand Up @@ -348,8 +346,6 @@ class nixlUcxDedicatedThread final : public nixlUcxThread {
std::jthread thread_;
};

namespace {

/**
* @brief UCX worker that owns its dedicated thread.
*/
Expand All @@ -376,14 +372,14 @@ class nixlUcxDedicatedWorker : public nixlUcxWorker {
std::optional<nixlUcxDedicatedThread> thread_;
};

} // namespace

nixlUcxDedicatedThread::nixlUcxDedicatedThread(nixlUcxEngine *engine,
nixlUcxDedicatedWorker &worker)
: nixlUcxThread(engine, {&worker}),
worker_(worker),
thread_(startThread()) {}

} // namespace

nixlUcxThreadPoolEngine::nixlUcxThreadPoolEngine(const nixlBackendInitParams &init_params,
size_t num_threads)
: nixlUcxThreadEngine(init_params, num_threads),
Expand Down Expand Up @@ -454,6 +450,12 @@ nixlUcxThreadPoolEngine::sendXferRange(const nixl_xfer_op_t &operation,

const auto comp_handle = static_cast<nixlUcxCompositeBackendReqH *>(int_handle);
comp_handle->startXfer();

// Notifications of the composite request are sent over its shared worker
const ucx_connection_ptr_t &conn =
static_cast<nixlUcxPublicMetadata *>(remote[start_idx].metadataP)->conn;
comp_handle->init(conn, *conn->getEp(comp_handle->getWorkerId()));

const size_t batch_size = local.descCount();
const size_t num_chunks = comp_handle->getNumChunks();
NIXL_TRACE << "sending " << *comp_handle;
Expand Down
Loading
Loading