From d75d8e1df70030861d53829ddb848a3c983e8df7 Mon Sep 17 00:00:00 2001 From: Wilson Liu Date: Thu, 13 Aug 2026 16:43:51 -0500 Subject: [PATCH 1/2] UCP/RNDV: keep the rkey when a put completion restarts on failover Admitting the put scheme to failover routed its completion through ucp_proto_request_zcopy_complete_cb, which restarts the request when the status is an error and the endpoint is still recovering. But ucp_proto_rndv_put_common_complete released the remote key before calling it, and ucp_proto_rndv_rkey_destroy ends with rkey = NULL, so the restarted write dereferenced a NULL rkey in ucp_rkey_get_tl_rkey and segfaulted. In a debug build it trips the assert at the top of ucp_proto_rndv_rkey_destroy instead. The get scheme already had this right - it keeps the rkey across a restart because the remote buffer stays registered until ATS - but the condition was open-coded in each place, so the put path could be written without it. Hoist the predicate into ucp_proto_request_is_failover_restart() and have all three sites ask it: the completion hook that performs the restart, and the two schemes that must not release resources the restarted send still needs. The RTR paths already guard their rkey release with a NULL check and complete through ucp_proto_rndv_recv_complete rather than the restart hook, so they were never exposed. --- src/ucp/proto/proto_common.inl | 17 ++++++++++++++--- src/ucp/rndv/rndv_get.c | 3 +-- src/ucp/rndv/rndv_put.c | 16 +++++++++++++--- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/src/ucp/proto/proto_common.inl b/src/ucp/proto/proto_common.inl index 44a1a46efb4..6576ac5f04c 100644 --- a/src/ucp/proto/proto_common.inl +++ b/src/ucp/proto/proto_common.inl @@ -93,6 +93,19 @@ ucp_proto_request_zcopy_clean(ucp_request_t *req, unsigned dt_mask) req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; } +/* Whether a completion with this status will restart the request instead of + * completing it. Callers that release per-request resources before completing + * must ask first: anything the restarted send still needs (a remote rkey, for + * one) has to outlive the completion. */ +static UCS_F_ALWAYS_INLINE int +ucp_proto_request_is_failover_restart(const ucp_request_t *req, + ucs_status_t status) +{ + return ucs_unlikely(status != UCS_OK) && + ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) && + !(req->send.ep->flags & UCP_EP_FLAG_FAILED); +} + static UCS_F_ALWAYS_INLINE void ucp_proto_request_zcopy_complete_cb(ucp_request_t *req, ucs_status_t status, ucp_request_callback_t complete_cb) @@ -105,9 +118,7 @@ ucp_proto_request_zcopy_complete_cb(ucp_request_t *req, ucs_status_t status, UCP_EP_STAT_TAG_OP(req->send.ep, EAGER) } - if (ucs_unlikely(status != UCS_OK) && - ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) && - !(req->send.ep->flags & UCP_EP_FLAG_FAILED)) { + if (ucp_proto_request_is_failover_restart(req, status)) { ucp_proto_request_restart(req); } else { if (complete_cb != NULL) { diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index e3eece612d2..c8e65dd47ca 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -115,8 +115,7 @@ ucp_proto_rndv_get_zcopy_fetch_completion(uct_completion_t *uct_comp) ucp_datatype_iter_mem_dereg(&req->send.state.dt_iter, UCS_BIT(UCP_DATATYPE_CONTIG)); if (ucs_unlikely(uct_comp->status != UCS_OK)) { - if (ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) && - !(req->send.ep->flags & UCP_EP_FLAG_FAILED)) { + if (ucp_proto_request_is_failover_restart(req, uct_comp->status)) { /* A lane failed but the endpoint is recovering: restart the * fetch on the surviving lanes instead of failing the receive. * Keep the rkey - the remote buffer is still registered until diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index ad37cee1186..3ac06ba785c 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -43,12 +43,22 @@ ucp_proto_rndv_put_common_complete(ucp_request_t *req) { const ucp_proto_rndv_put_priv_t UCS_V_UNUSED *rpriv = req->send.proto_config->priv; + ucs_status_t status = req->send.state.uct_comp.status; + ucp_trace_req(req, "rndv_put_common_complete"); UCS_STATS_UPDATE_COUNTER(req->send.ep->worker->stats, rpriv->stat_counter, +1); - ucp_proto_rndv_rkey_destroy(req); - ucp_proto_rndv_request_zcopy_complete(req, - req->send.state.uct_comp.status); + + /* Keep the rkey when the completion below is going to restart this write + * on the surviving lanes rather than complete it: the remote buffer stays + * registered until ATP, and the restarted send reads the rkey again. The + * get scheme keeps it for the same reason. Destroying it here first left + * the restart dereferencing a NULL rkey. */ + if (!ucp_proto_request_is_failover_restart(req, status)) { + ucp_proto_rndv_rkey_destroy(req); + } + + ucp_proto_rndv_request_zcopy_complete(req, status); } static void ucp_proto_rndv_put_zcopy_completion(uct_completion_t *uct_comp) From 0b04ed7100747e79254bcd85ea3791c075e5691e Mon Sep 17 00:00:00 2001 From: Wilson Liu Date: Thu, 13 Aug 2026 17:03:55 -0500 Subject: [PATCH 2/2] UCP/PROTO: fail a restart the protocol cannot resume instead of crashing ucp_proto_request_restart() re-enters protocol selection, so the request must still own everything the newly selected protocol will use. reset() can rewind a datatype iterator but cannot resurrect a released remote key, and nothing said so - which is how the put scheme came to release its rkey before a completion that restarts. Give the protocol a way to say it cannot resume: a reset status other than UCS_OK or UCS_ERR_CANCELED now aborts the request rather than tripping ucs_assertv_always. The assert fired in every build, so this replaces a crash with a failed transfer that the caller already knows how to retry. Both rndv zcopy protocols answer it by refusing to restart without their remote key. This is the backstop, not the fix - the rkey is now kept across a restart, so the refusal should never fire. It exists because the invariant is subtle enough to be missed again: a restart REQUIRES some cleanup to have happened (the get scheme deregisters its datatype iterator on purpose, and reset re-registers) while forbidding other cleanup (the peer's rkey, which only the RTS could provide). Release only what the restart re-acquires is not something the next protocol admitted to failover can infer, and asserts are no help in the release builds where this surfaced as a segfault. --- src/ucp/proto/proto_common.c | 18 +++++++++++++++--- src/ucp/rndv/rndv_get.c | 7 +++++++ src/ucp/rndv/rndv_put.c | 8 ++++++++ 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/src/ucp/proto/proto_common.c b/src/ucp/proto/proto_common.c index ba221ecc9e8..eb626f20715 100644 --- a/src/ucp/proto/proto_common.c +++ b/src/ucp/proto/proto_common.c @@ -876,10 +876,22 @@ void ucp_proto_request_restart(ucp_request_t *req) proto_config->proto->name, req->send.proto_stage); status = proto_config->proto->reset(req); + if (status == UCS_ERR_CANCELED) { + /* The request was already completed by cancellation. */ + return; + } + if (status != UCS_OK) { - ucs_assertv_always(status == UCS_ERR_CANCELED, - "req %p, failed to reset: status %s", req, - ucs_status_string(status)); + /* The protocol cannot resume from this state - a rendezvous protocol + * whose remote key has already been released, for instance. Restart + * re-enters protocol selection, so the request must still own + * everything the selected protocol will use, and reset is the only + * place that can tell. Fail the request rather than re-selecting on + * top of resources it no longer holds: the caller sees a failed + * transfer instead of the worker dereferencing a released one. */ + ucs_error("req %p: proto %s cannot restart (%s); aborting the request", + req, proto_config->proto->name, ucs_status_string(status)); + ucp_proto_request_abort(req, status); return; } diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index c8e65dd47ca..279906b2f99 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -227,6 +227,13 @@ static ucs_status_t ucp_rndv_get_zcopy_proto_reset(ucp_request_t *req) return UCS_OK; } + /* Same contract as the put scheme: the fetch resumes against the peer's + * remote key, so a restart is only possible while the request still owns + * it. */ + if (req->send.rndv.rkey == NULL) { + return UCS_ERR_CONNECTION_RESET; + } + req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; switch (req->send.proto_stage) { diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 3ac06ba785c..ab9b157b0a5 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -465,6 +465,14 @@ static ucs_status_t ucp_proto_rndv_put_zcopy_reset(ucp_request_t *req) { const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; + /* The write resumes against the peer's remote key, which came from the + * RTS and cannot be re-acquired here. Refuse the restart if a completion + * path released it, so the request fails instead of re-selecting a + * protocol that would dereference it. */ + if (req->send.rndv.rkey == NULL) { + return UCS_ERR_CONNECTION_RESET; + } + if (req->send.rndv.put.atp_count == rpriv->atp_num_lanes) { /* Sent all ATPs so the iterator should be at the end */ ucs_assertv_always(ucp_datatype_iter_is_end(&req->send.state.dt_iter),