From bb7a13aa838e0af7dc77ca9bf99874c17025a71c Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 8 Dec 2025 18:41:51 +0200 Subject: [PATCH 01/67] UCP/RNDV: Throttle rndv fragment requests (both pipeline and standalone mtype requests) --- src/ucp/core/ucp_context.c | 10 +++++ src/ucp/core/ucp_context.h | 4 ++ src/ucp/core/ucp_request.h | 5 ++- src/ucp/core/ucp_worker.c | 4 ++ src/ucp/core/ucp_worker.h | 6 +++ src/ucp/rndv/rndv_get.c | 11 +++++ src/ucp/rndv/rndv_mtype.inl | 80 +++++++++++++++++++++++++++++++++++++ src/ucp/rndv/rndv_ppln.c | 3 +- src/ucp/rndv/rndv_put.c | 13 ++++-- src/ucp/rndv/rndv_rtr.c | 11 +++++ 10 files changed, 142 insertions(+), 5 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 4543f0b6bda..bb341266cd9 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,6 +404,16 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, + {"RNDV_PIPELINE_WORKER_FC_ENABLE", "n", + "Enable worker-level flow control to limit total concurrent pipeline fragments\n" + "across all requests, preventing memory exhaustion", + ucs_offsetof(ucp_context_config_t, rndv_ppln_worker_fc_enable), UCS_CONFIG_TYPE_BOOL}, + + {"RNDV_PIPELINE_WORKER_MAX_FRAGS", "5000", + "Maximum number of concurrent pipeline fragments per worker\n" + "(only applies when RNDV_PIPELINE_WORKER_FC_ENABLE=y)", + ucs_offsetof(ucp_context_config_t, rndv_ppln_worker_max_frags), UCS_CONFIG_TYPE_ULUNITS}, + {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" "the flush operation in a bounded time even if there are new requests on\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 912a09feb8d..b6cbec20913 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -132,6 +132,10 @@ typedef struct ucp_context_config { int rndv_errh_ppln_enable; /** Force-enable the RMA rendezvous put/get protocols */ int rma_ppln_enable; + /** Enable flow control for rndv pipeline fragments at worker level */ + int rndv_ppln_worker_fc_enable; + /** Maximum number of concurrent pipeline fragments per worker */ + size_t rndv_ppln_worker_max_frags; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 26079db4ada..80583ddb977 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -326,7 +326,10 @@ struct ucp_request { /* Used by rndv/send/ppln and rndv/recv/ppln */ struct { /* Size to send in ack message */ - ssize_t ack_data_size; + ssize_t ack_data_size; + /* Element in worker-level pending queue + * for throttled ppln requests */ + ucs_queue_elem_t queue_elem; } ppln; /* Used by rndv/rkey_ptr */ diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 9f3591bb593..3132b390315 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2639,6 +2639,10 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_closures = 0; worker->counters.ep_failures = 0; + /* Initialize RNDV pipeline flow control */ + worker->rndv_ppln_fc.active_frags = 0; + ucs_queue_head_init(&worker->rndv_ppln_fc.pending_q); + /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & UCS_MASK(UCP_WORKER_INTERNAL_FLAGS_SHIFT); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 3f917ea63a1..e472bd77c41 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -396,6 +396,12 @@ typedef struct ucp_worker { uint64_t ep_failures; } counters; + struct { + /* Worker-level ppln fragment flow control */ + size_t active_frags; /* Current active fragments */ + ucs_queue_head_t pending_q; /* Queue of throttled ppln requests */ + } rndv_ppln_fc; + struct { /* Usage tracker handle */ ucs_usage_tracker_h handle; diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 6636ef717b4..d22e11e85de 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -274,6 +274,8 @@ ucp_proto_rndv_get_mtype_unpack_completion(uct_completion_t *uct_comp) send.state.uct_comp); ucs_mpool_put_inline(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_fc_decrement(req); + if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 1, 0); } else { @@ -304,6 +306,12 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + /* Check throttling. If no resource at the moment, queue the request + * in the throttle pending queue and return UCS_OK. */ + if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { @@ -311,6 +319,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) return UCS_OK; } + ucp_proto_rndv_mtype_fc_increment(req); ucp_proto_rndv_get_common_request_init(req); ucp_proto_completion_init(&req->send.state.uct_comp, ucp_proto_rndv_get_mtype_fetch_completion); @@ -376,6 +385,8 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) req->send.rndv.mdesc = NULL; req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; + ucp_proto_rndv_mtype_fc_decrement(req); + if ((req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_FETCH) && (req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_ATS)) { ucp_proto_fatal_invalid_stage(req, "reset"); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 4b11edb3fd8..939e58a7691 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -168,6 +168,86 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_copy( return status; } +/* Reschedule callback for throttled mtype requests */ +static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +{ + ucp_request_t *req = arg; + ucp_request_send(req); + return 1; +} + +/** + * Check if request should be throttled due to flow control limit. + * If throttled, the request is queued to pending_q. + * + * @return UCS_OK if not throttled (caller should continue), + * UCS_ERR_NO_RESOURCE if throttled and queued (caller should return UCS_OK). + */ +static UCS_F_ALWAYS_INLINE ucs_status_t +ucp_proto_rndv_mtype_fc_check(ucp_request_t *req) +{ + ucp_worker_h worker = req->send.ep->worker; + ucp_context_h context = worker->context; + + if (!context->config.ext.rndv_ppln_worker_fc_enable) { + return UCS_OK; + } + + if (worker->rndv_ppln_fc.active_frags >= + context->config.ext.rndv_ppln_worker_max_frags) { + ucs_trace_req("mtype_fc: throttle limit reached active_frags=%zu max=%zu", + worker->rndv_ppln_fc.active_frags, + context->config.ext.rndv_ppln_worker_max_frags); + ucs_queue_push(&worker->rndv_ppln_fc.pending_q, + &req->send.rndv.ppln.queue_elem); + return UCS_ERR_NO_RESOURCE; + } + + return UCS_OK; +} + +/** + * Increment active_frags counter after successful mtype allocation. + */ +static UCS_F_ALWAYS_INLINE void +ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) +{ + ucp_worker_h worker = req->send.ep->worker; + + if (worker->context->config.ext.rndv_ppln_worker_fc_enable) { + worker->rndv_ppln_fc.active_frags++; + } +} + +/** + * Decrement active_frags counter and reschedule pending request if any. + */ +static UCS_F_ALWAYS_INLINE void +ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) +{ + ucp_worker_h worker = req->send.ep->worker; + ucp_context_h context = worker->context; + + if (!context->config.ext.rndv_ppln_worker_fc_enable) { + return; + } + + ucs_assert(worker->rndv_ppln_fc.active_frags > 0); + worker->rndv_ppln_fc.active_frags--; + + if (!ucs_queue_is_empty(&worker->rndv_ppln_fc.pending_q)) { + ucp_request_t *pending_req; + ucs_queue_elem_t *elem; + + elem = ucs_queue_pull(&worker->rndv_ppln_fc.pending_q); + pending_req = ucs_container_of(elem, ucp_request_t, + send.rndv.ppln.queue_elem); + ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, + ucp_proto_rndv_mtype_fc_reschedule_cb, + pending_req); + } +} + static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mdesc_mtype_copy(ucp_request_t *req, uct_ep_put_zcopy_func_t copy_func, diff --git a/src/ucp/rndv/rndv_ppln.c b/src/ucp/rndv/rndv_ppln.c index 9b5c7fc0a6c..049ce9ca1f5 100644 --- a/src/ucp/rndv/rndv_ppln.c +++ b/src/ucp/rndv/rndv_ppln.c @@ -13,6 +13,7 @@ #include #include #include +#include enum { @@ -261,7 +262,7 @@ void ucp_proto_rndv_ppln_recv_frag_complete(ucp_request_t *freq, int send_ack, static ucs_status_t ucp_proto_rndv_ppln_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_ppln_priv_t *rpriv; ucp_datatype_iter_t next_iter; diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 0cab5a465d1..677fabec8b3 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,14 +520,18 @@ static UCS_F_INLINE_OPTIMIZED ucs_status_t ucp_proto_rndv_put_mtype_send_func( static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, - ucp_request_t, - send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; ucs_status_t status; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); + /* Check throttling. If no resource at the moment, queue the request + * in the throttle pending queue and return UCS_OK. */ + if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev); if (status != UCS_OK) { @@ -535,6 +539,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) return UCS_OK; } + ucp_proto_rndv_mtype_fc_increment(req); ucp_proto_rndv_put_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; ucp_proto_rndv_mdesc_mtype_copy(req, uct_ep_get_zcopy, @@ -566,6 +571,7 @@ static void ucp_proto_rndv_put_mtype_completion(uct_completion_t *uct_comp) ucp_trace_req(req, "rndv_put_mtype_completion"); ucs_mpool_put(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_fc_decrement(req); ucp_proto_rndv_put_common_complete(req); } @@ -576,6 +582,7 @@ static void ucp_proto_rndv_put_mtype_frag_completion(uct_completion_t *uct_comp) ucp_trace_req(req, "rndv_put_mtype_frag_completion"); ucs_mpool_put(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_fc_decrement(req); ucp_proto_rndv_ppln_send_frag_complete(req, 1); } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 7494a47c415..2352f4a84ca 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -362,6 +362,9 @@ ucp_proto_rndv_rtr_mtype_complete(ucp_request_t *req, int abort) if (!abort || (req->send.rndv.mdesc != NULL)) { ucs_mpool_put_inline(req->send.rndv.mdesc); } + + ucp_proto_rndv_mtype_fc_decrement(req); + if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 0, abort); } else { @@ -392,6 +395,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_reset(ucp_request_t *req) if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { ucs_mpool_put_inline(req->send.rndv.mdesc); req->send.rndv.mdesc = NULL; + ucp_proto_rndv_mtype_fc_decrement(req); } return ucp_proto_request_zcopy_id_reset(req); @@ -428,6 +432,12 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + /* Check throttling. If no resource at the moment, queue the request + * in the throttle pending queue and return UCS_OK. */ + if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { @@ -435,6 +445,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) return UCS_OK; } + ucp_proto_rndv_mtype_fc_increment(req); ucp_proto_rtr_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; } From 8d97668f1b56ca69ee6ed8c16107baec79d8e69b Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 15 Dec 2025 16:59:18 +0200 Subject: [PATCH 02/67] UCP/RNDV: change 'ppln' prefix to 'mtype' --- src/ucp/core/ucp_context.c | 14 +++++++------- src/ucp/core/ucp_context.h | 8 ++++---- src/ucp/core/ucp_worker.c | 4 ++-- src/ucp/core/ucp_worker.h | 6 +++--- src/ucp/rndv/rndv_mtype.inl | 26 +++++++++++++------------- 5 files changed, 29 insertions(+), 29 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index bb341266cd9..ee9c5e57fa0 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,15 +404,15 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_PIPELINE_WORKER_FC_ENABLE", "n", - "Enable worker-level flow control to limit total concurrent pipeline fragments\n" + {"RNDV_MTYPE_WORKER_FC_ENABLE", "n", + "Enable worker-level flow control to limit total concurrent mtype fragments\n" "across all requests, preventing memory exhaustion", - ucs_offsetof(ucp_context_config_t, rndv_ppln_worker_fc_enable), UCS_CONFIG_TYPE_BOOL}, + ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_fc_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_PIPELINE_WORKER_MAX_FRAGS", "5000", - "Maximum number of concurrent pipeline fragments per worker\n" - "(only applies when RNDV_PIPELINE_WORKER_FC_ENABLE=y)", - ucs_offsetof(ucp_context_config_t, rndv_ppln_worker_max_frags), UCS_CONFIG_TYPE_ULUNITS}, + {"RNDV_MTYPE_WORKER_MAX_FRAGS", "1024", + "Maximum number of concurrent mtype fragments per worker\n" + "(only applies when RNDV_MTYPE_WORKER_FC_ENABLE=y)", + ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_frags), UCS_CONFIG_TYPE_ULUNITS}, {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index b6cbec20913..302163d8e82 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -132,10 +132,10 @@ typedef struct ucp_context_config { int rndv_errh_ppln_enable; /** Force-enable the RMA rendezvous put/get protocols */ int rma_ppln_enable; - /** Enable flow control for rndv pipeline fragments at worker level */ - int rndv_ppln_worker_fc_enable; - /** Maximum number of concurrent pipeline fragments per worker */ - size_t rndv_ppln_worker_max_frags; + /** Enable flow control for rndv mtype fragments at worker level */ + int rndv_mtype_worker_fc_enable; + /** Maximum number of concurrent rndv mtype fragments per worker */ + size_t rndv_mtype_worker_max_frags; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 3132b390315..6729f0dd9ad 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2640,8 +2640,8 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_failures = 0; /* Initialize RNDV pipeline flow control */ - worker->rndv_ppln_fc.active_frags = 0; - ucs_queue_head_init(&worker->rndv_ppln_fc.pending_q); + worker->rndv_mtype_fc.active_frags = 0; + ucs_queue_head_init(&worker->rndv_mtype_fc.pending_q); /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index e472bd77c41..ac16cfd382f 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -397,10 +397,10 @@ typedef struct ucp_worker { } counters; struct { - /* Worker-level ppln fragment flow control */ + /* Worker-level mtype fragment flow control */ size_t active_frags; /* Current active fragments */ - ucs_queue_head_t pending_q; /* Queue of throttled ppln requests */ - } rndv_ppln_fc; + ucs_queue_head_t pending_q; /* Queue of throttled mtype requests */ + } rndv_mtype_fc; struct { /* Usage tracker handle */ diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 939e58a7691..6b56c444509 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -189,16 +189,16 @@ ucp_proto_rndv_mtype_fc_check(ucp_request_t *req) ucp_worker_h worker = req->send.ep->worker; ucp_context_h context = worker->context; - if (!context->config.ext.rndv_ppln_worker_fc_enable) { + if (!context->config.ext.rndv_mtype_worker_fc_enable) { return UCS_OK; } - if (worker->rndv_ppln_fc.active_frags >= - context->config.ext.rndv_ppln_worker_max_frags) { + if (worker->rndv_mtype_fc.active_frags >= + context->config.ext.rndv_mtype_worker_max_frags) { ucs_trace_req("mtype_fc: throttle limit reached active_frags=%zu max=%zu", - worker->rndv_ppln_fc.active_frags, - context->config.ext.rndv_ppln_worker_max_frags); - ucs_queue_push(&worker->rndv_ppln_fc.pending_q, + worker->rndv_mtype_fc.active_frags, + context->config.ext.rndv_mtype_worker_max_frags); + ucs_queue_push(&worker->rndv_mtype_fc.pending_q, &req->send.rndv.ppln.queue_elem); return UCS_ERR_NO_RESOURCE; } @@ -214,8 +214,8 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; - if (worker->context->config.ext.rndv_ppln_worker_fc_enable) { - worker->rndv_ppln_fc.active_frags++; + if (worker->context->config.ext.rndv_mtype_worker_fc_enable) { + worker->rndv_mtype_fc.active_frags++; } } @@ -228,18 +228,18 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) ucp_worker_h worker = req->send.ep->worker; ucp_context_h context = worker->context; - if (!context->config.ext.rndv_ppln_worker_fc_enable) { + if (!context->config.ext.rndv_mtype_worker_fc_enable) { return; } - ucs_assert(worker->rndv_ppln_fc.active_frags > 0); - worker->rndv_ppln_fc.active_frags--; + ucs_assert(worker->rndv_mtype_fc.active_frags > 0); + worker->rndv_mtype_fc.active_frags--; - if (!ucs_queue_is_empty(&worker->rndv_ppln_fc.pending_q)) { + if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q)) { ucp_request_t *pending_req; ucs_queue_elem_t *elem; - elem = ucs_queue_pull(&worker->rndv_ppln_fc.pending_q); + elem = ucs_queue_pull(&worker->rndv_mtype_fc.pending_q); pending_req = ucs_container_of(elem, ucp_request_t, send.rndv.ppln.queue_elem); ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, From 30c5741f46a1bf50652a6ea40a1de0348537b9a4 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 23 Dec 2025 17:52:32 +0200 Subject: [PATCH 03/67] UCP/RNDV: prioritize PUT ops over RTR --- src/ucp/core/ucp_worker.c | 8 ++-- src/ucp/core/ucp_worker.h | 5 ++- src/ucp/rndv/rndv_mtype.inl | 76 +++++++++++++++++++++++++++++-------- 3 files changed, 70 insertions(+), 19 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 6729f0dd9ad..b5af3a461b4 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2639,9 +2639,11 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_closures = 0; worker->counters.ep_failures = 0; - /* Initialize RNDV pipeline flow control */ - worker->rndv_mtype_fc.active_frags = 0; - ucs_queue_head_init(&worker->rndv_mtype_fc.pending_q); + /* Initialize RNDV mtype flow control with separate priority queues */ + worker->rndv_mtype_fc.active_frags = 0; + ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index ac16cfd382f..391ac3e2559 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -399,7 +399,10 @@ typedef struct ucp_worker { struct { /* Worker-level mtype fragment flow control */ size_t active_frags; /* Current active fragments */ - ucs_queue_head_t pending_q; /* Queue of throttled mtype requests */ + /* Separate pending queues by priority (PUT > GET > RTR) */ + ucs_queue_head_t put_pending_q; /* Throttled PUT (RNDV_SEND) requests */ + ucs_queue_head_t get_pending_q; /* Throttled GET requests */ + ucs_queue_head_t rtr_pending_q; /* Throttled RTR requests */ } rndv_mtype_fc; struct { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 6b56c444509..faf5fc0a93d 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -10,6 +10,7 @@ #include "rndv.h" #include +#include static ucp_ep_h ucp_proto_rndv_mtype_ep(ucp_worker_t *worker, @@ -176,9 +177,25 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) return 1; } +/** + * Check if this is an RTR protocol (vs GET protocol). + * Both are UCP_OP_ID_RNDV_RECV, but RTR holds memory longer (waiting for PUT). + */ +static UCS_F_ALWAYS_INLINE int +ucp_proto_rndv_mtype_is_rtr(ucp_request_t *req) +{ + const char *proto_name = req->send.proto_config->proto->name; + return (strstr(proto_name, "rtr") != NULL); +} + /** * Check if request should be throttled due to flow control limit. - * If throttled, the request is queued to pending_q. + * If throttled, the request is queued to the appropriate priority queue. + * + * Priority and quota allocation: + * - PUT (RNDV_SEND): Full quota, highest dequeue priority (releases both sides' memory) + * - GET (RNDV_RECV): 75% quota, medium priority (releases receiver memory) + * - RTR (RNDV_RECV): 50% quota, lowest priority (holds memory until remote PUT) * * @return UCS_OK if not throttled (caller should continue), * UCS_ERR_NO_RESOURCE if throttled and queued (caller should return UCS_OK). @@ -186,20 +203,38 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_fc_check(ucp_request_t *req) { - ucp_worker_h worker = req->send.ep->worker; - ucp_context_h context = worker->context; + ucp_worker_h worker = req->send.ep->worker; + ucp_context_h context = worker->context; + size_t max_frags = context->config.ext.rndv_mtype_worker_max_frags; + ucp_operation_id_t op_id; + ucs_queue_head_t *pending_q; + size_t effective_max; if (!context->config.ext.rndv_mtype_worker_fc_enable) { return UCS_OK; } - if (worker->rndv_mtype_fc.active_frags >= - context->config.ext.rndv_mtype_worker_max_frags) { - ucs_trace_req("mtype_fc: throttle limit reached active_frags=%zu max=%zu", - worker->rndv_mtype_fc.active_frags, - context->config.ext.rndv_mtype_worker_max_frags); - ucs_queue_push(&worker->rndv_mtype_fc.pending_q, - &req->send.rndv.ppln.queue_elem); + op_id = ucp_proto_select_op_id(&req->send.proto_config->select_param); + + if (op_id == UCP_OP_ID_RNDV_SEND) { + /* PUT - can use full quota (highest priority, releases both sides) */ + effective_max = max_frags; + pending_q = &worker->rndv_mtype_fc.put_pending_q; + } else if (ucp_proto_rndv_mtype_is_rtr(req)) { + /* RTR - reserve 50% for PUT/GET (holds memory waiting for remote PUT) */ + effective_max = max_frags / 2; + pending_q = &worker->rndv_mtype_fc.rtr_pending_q; + } else { + /* GET - reserve 25% for PUT (releases receiver memory on completion) */ + effective_max = (max_frags * 3) / 4; + pending_q = &worker->rndv_mtype_fc.get_pending_q; + } + + if (worker->rndv_mtype_fc.active_frags >= effective_max) { + ucs_trace_req("mtype_fc: throttle op=%d active=%zu max=%zu effective=%zu", + op_id, worker->rndv_mtype_fc.active_frags, + max_frags, effective_max); + ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); return UCS_ERR_NO_RESOURCE; } @@ -220,13 +255,19 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) } /** - * Decrement active_frags counter and reschedule pending request if any. + * Decrement active_frags counter and reschedule pending request. + * Dequeue priority: PUT > GET > RTR (by memory release impact) + * - PUT completion releases memory on both sender and receiver + * - GET completion releases memory on receiver + * - RTR completion doesn't release memory (waiting for PUT) */ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; ucp_context_h context = worker->context; + ucs_queue_elem_t *elem = NULL; + ucp_request_t *pending_req; if (!context->config.ext.rndv_mtype_worker_fc_enable) { return; @@ -235,11 +276,16 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) ucs_assert(worker->rndv_mtype_fc.active_frags > 0); worker->rndv_mtype_fc.active_frags--; - if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q)) { - ucp_request_t *pending_req; - ucs_queue_elem_t *elem; + /* Dequeue with priority: PUT > GET > RTR */ + if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) { + elem = ucs_queue_pull(&worker->rndv_mtype_fc.put_pending_q); + } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.get_pending_q)) { + elem = ucs_queue_pull(&worker->rndv_mtype_fc.get_pending_q); + } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.rtr_pending_q)) { + elem = ucs_queue_pull(&worker->rndv_mtype_fc.rtr_pending_q); + } - elem = ucs_queue_pull(&worker->rndv_mtype_fc.pending_q); + if (elem != NULL) { pending_req = ucs_container_of(elem, ucp_request_t, send.rndv.ppln.queue_elem); ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, From 7716784e330e328bd3de699c96a50bf144998c83 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 10:02:48 +0200 Subject: [PATCH 04/67] UCP/RNDV: unite RTR, PUT and GET in a generic ucp_proto_rndv_mtype_fc_check --- src/ucp/rndv/rndv_get.c | 14 +++++++++---- src/ucp/rndv/rndv_mtype.inl | 39 +++++++++---------------------------- src/ucp/rndv/rndv_put.c | 14 +++++++++---- src/ucp/rndv/rndv_rtr.c | 11 ++++++++--- 4 files changed, 37 insertions(+), 41 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index d22e11e85de..63d4dc56e95 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -298,17 +298,23 @@ ucp_proto_rndv_get_mtype_fetch_completion(uct_completion_t *uct_comp) static ucs_status_t ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_context_h context = req->send.ep->worker->context; const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; + size_t max_frags; + ucs_queue_head_t *pending_q; /* coverity[tainted_data_downcast] */ rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - /* Check throttling. If no resource at the moment, queue the request - * in the throttle pending queue and return UCS_OK. */ - if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + /* Check throttling limit. If no resource at the moment, queue the request + * in GET pending queue and return UCS_OK. */ + max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 4; + pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; + if (ucp_proto_rndv_mtype_fc_check( + req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index faf5fc0a93d..c517c60565d 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -179,7 +179,6 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) /** * Check if this is an RTR protocol (vs GET protocol). - * Both are UCP_OP_ID_RNDV_RECV, but RTR holds memory longer (waiting for PUT). */ static UCS_F_ALWAYS_INLINE int ucp_proto_rndv_mtype_is_rtr(ucp_request_t *req) @@ -201,39 +200,19 @@ ucp_proto_rndv_mtype_is_rtr(ucp_request_t *req) * UCS_ERR_NO_RESOURCE if throttled and queued (caller should return UCS_OK). */ static UCS_F_ALWAYS_INLINE ucs_status_t -ucp_proto_rndv_mtype_fc_check(ucp_request_t *req) +ucp_proto_rndv_mtype_fc_check(ucp_request_t *req, size_t max_frags, + ucs_queue_head_t *pending_q) { - ucp_worker_h worker = req->send.ep->worker; - ucp_context_h context = worker->context; - size_t max_frags = context->config.ext.rndv_mtype_worker_max_frags; - ucp_operation_id_t op_id; - ucs_queue_head_t *pending_q; - size_t effective_max; + ucp_worker_h worker = req->send.ep->worker; + ucp_context_h context = worker->context; if (!context->config.ext.rndv_mtype_worker_fc_enable) { return UCS_OK; } - op_id = ucp_proto_select_op_id(&req->send.proto_config->select_param); - - if (op_id == UCP_OP_ID_RNDV_SEND) { - /* PUT - can use full quota (highest priority, releases both sides) */ - effective_max = max_frags; - pending_q = &worker->rndv_mtype_fc.put_pending_q; - } else if (ucp_proto_rndv_mtype_is_rtr(req)) { - /* RTR - reserve 50% for PUT/GET (holds memory waiting for remote PUT) */ - effective_max = max_frags / 2; - pending_q = &worker->rndv_mtype_fc.rtr_pending_q; - } else { - /* GET - reserve 25% for PUT (releases receiver memory on completion) */ - effective_max = (max_frags * 3) / 4; - pending_q = &worker->rndv_mtype_fc.get_pending_q; - } - - if (worker->rndv_mtype_fc.active_frags >= effective_max) { - ucs_trace_req("mtype_fc: throttle op=%d active=%zu max=%zu effective=%zu", - op_id, worker->rndv_mtype_fc.active_frags, - max_frags, effective_max); + if (worker->rndv_mtype_fc.active_frags >= max_frags) { + ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", + worker->rndv_mtype_fc.active_frags, max_frags); ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); return UCS_ERR_NO_RESOURCE; } @@ -264,8 +243,8 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) { - ucp_worker_h worker = req->send.ep->worker; - ucp_context_h context = worker->context; + ucp_worker_h worker = req->send.ep->worker; + ucp_context_h context = worker->context; ucs_queue_elem_t *elem = NULL; ucp_request_t *pending_req; diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 677fabec8b3..ab267122426 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,15 +520,21 @@ static UCS_F_INLINE_OPTIMIZED ucs_status_t ucp_proto_rndv_put_mtype_send_func( static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_context_h context = req->send.ep->worker->context; const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; ucs_status_t status; + size_t max_frags; + ucs_queue_head_t *pending_q; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); - /* Check throttling. If no resource at the moment, queue the request - * in the throttle pending queue and return UCS_OK. */ - if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + /* Check throttling limit. If no resource at the moment, queue the request + * in PUT pending queue and return UCS_OK. */ + max_frags = context->config.ext.rndv_mtype_worker_max_frags; + pending_q = &req->send.ep->worker->rndv_mtype_fc.put_pending_q; + if (ucp_proto_rndv_mtype_fc_check( + req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 2352f4a84ca..86918b1002f 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -429,12 +429,17 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); const ucp_proto_rndv_rtr_mtype_priv_t *rpriv = req->send.proto_config->priv; + size_t max_frags; + ucs_queue_head_t *pending_q; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - /* Check throttling. If no resource at the moment, queue the request - * in the throttle pending queue and return UCS_OK. */ - if (ucp_proto_rndv_mtype_fc_check(req) == UCS_ERR_NO_RESOURCE) { + /* Check throttling limit. If no resource at the moment, queue the + * request in RTR pending queue and return UCS_OK. */ + max_frags = context->config.ext.rndv_mtype_worker_max_frags / 2; + pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; + if (ucp_proto_rndv_mtype_fc_check( + req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } From 3eddbacef33525c83b5520241e9b79a843317694 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 10:21:47 +0200 Subject: [PATCH 05/67] UCP/RNDV: some fixes --- src/ucp/rndv/rndv_mtype.inl | 25 ++++++++++++++++--------- src/ucp/rndv/rndv_put.c | 3 ++- src/ucp/rndv/rndv_rtr.c | 6 ++++-- 3 files changed, 22 insertions(+), 12 deletions(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index c517c60565d..f00987e8fc1 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -191,13 +191,15 @@ ucp_proto_rndv_mtype_is_rtr(ucp_request_t *req) * Check if request should be throttled due to flow control limit. * If throttled, the request is queued to the appropriate priority queue. * - * Priority and quota allocation: - * - PUT (RNDV_SEND): Full quota, highest dequeue priority (releases both sides' memory) - * - GET (RNDV_RECV): 75% quota, medium priority (releases receiver memory) - * - RTR (RNDV_RECV): 50% quota, lowest priority (holds memory until remote PUT) + * Priority and quota allocation (based on memory release and system impact): + * - PUT: Full quota, remote already sent RTR and is waiting. Completing PUT + * unblocks remote. + * - GET: 75% quota, self-contained operation, no remote allocation triggered. + * - RTR: 50% quota, scheduling RTR causes additional allocation on remote + * (RNDV PUT). * * @return UCS_OK if not throttled (caller should continue), - * UCS_ERR_NO_RESOURCE if throttled and queued (caller should return UCS_OK). + * UCS_ERR_NO_RESOURCE if throttled and queued. */ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_fc_check(ucp_request_t *req, size_t max_frags, @@ -235,10 +237,15 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) /** * Decrement active_frags counter and reschedule pending request. - * Dequeue priority: PUT > GET > RTR (by memory release impact) - * - PUT completion releases memory on both sender and receiver - * - GET completion releases memory on receiver - * - RTR completion doesn't release memory (waiting for PUT) + * Dequeue priority: PUT > GET > RTR + * + * Priority rationale: + * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote + * immediately. + * GET - Self-contained fetch operation. Completes without causing remote + * allocations. + * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total + * memory pressure. */ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index ab267122426..80b2d32af09 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -522,7 +522,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); ucp_context_h context = req->send.ep->worker->context; - const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; + const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; size_t max_frags; ucs_queue_head_t *pending_q; @@ -538,6 +538,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) return UCS_OK; } + rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev); if (status != UCS_OK) { diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 86918b1002f..804567e6e74 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -427,8 +427,9 @@ ucp_proto_rndv_rtr_mtype_data_received(ucp_request_t *req, int in_buffer) static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { - ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); - const ucp_proto_rndv_rtr_mtype_priv_t *rpriv = req->send.proto_config->priv; + ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + ucp_context_h context = req->send.ep->worker->context; + const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; size_t max_frags; ucs_queue_head_t *pending_q; ucs_status_t status; @@ -443,6 +444,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) return UCS_OK; } + rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { From a9685bbce57c9c610bca44d252061d62351a6d96 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 12:03:52 +0200 Subject: [PATCH 06/67] UCP/RNDV: clang --- src/ucp/rndv/rndv_get.c | 4 ++-- src/ucp/rndv/rndv_ppln.c | 2 +- src/ucp/rndv/rndv_put.c | 6 +++--- src/ucp/rndv/rndv_rtr.c | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 63d4dc56e95..da29369a682 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -313,8 +313,8 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) * in GET pending queue and return UCS_OK. */ max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 4; pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; - if (ucp_proto_rndv_mtype_fc_check( - req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { + if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_ppln.c b/src/ucp/rndv/rndv_ppln.c index 049ce9ca1f5..bf8b9f64ebb 100644 --- a/src/ucp/rndv/rndv_ppln.c +++ b/src/ucp/rndv/rndv_ppln.c @@ -262,7 +262,7 @@ void ucp_proto_rndv_ppln_recv_frag_complete(ucp_request_t *freq, int send_ack, static ucs_status_t ucp_proto_rndv_ppln_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_ppln_priv_t *rpriv; ucp_datatype_iter_t next_iter; diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 80b2d32af09..a5d6d9d6d1d 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -533,12 +533,12 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) * in PUT pending queue and return UCS_OK. */ max_frags = context->config.ext.rndv_mtype_worker_max_frags; pending_q = &req->send.ep->worker->rndv_mtype_fc.put_pending_q; - if (ucp_proto_rndv_mtype_fc_check( - req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { + if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + UCS_ERR_NO_RESOURCE) { return UCS_OK; } - rpriv = req->send.proto_config->priv; + rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev); if (status != UCS_OK) { diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 804567e6e74..4f5107b4526 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -439,12 +439,12 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) * request in RTR pending queue and return UCS_OK. */ max_frags = context->config.ext.rndv_mtype_worker_max_frags / 2; pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; - if (ucp_proto_rndv_mtype_fc_check( - req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { + if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + UCS_ERR_NO_RESOURCE) { return UCS_OK; } - rpriv = req->send.proto_config->priv; + rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { From 31eefc11cd6d93e34a50914eae90d1f49611a605 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 13:50:01 +0200 Subject: [PATCH 07/67] UCP/RNDV: some fixes --- src/ucp/rndv/rndv_get.c | 4 ++-- src/ucp/rndv/rndv_mtype.inl | 26 ++++++-------------------- 2 files changed, 8 insertions(+), 22 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index da29369a682..4d6d3f0339b 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -309,8 +309,8 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - /* Check throttling limit. If no resource at the moment, queue the request - * in GET pending queue and return UCS_OK. */ + /* Check throttling limit. If no resource at the moment, queue the + * request in GET pending queue and return UCS_OK. */ max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 4; pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index f00987e8fc1..305a34b1696 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -177,29 +177,15 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) return 1; } -/** - * Check if this is an RTR protocol (vs GET protocol). - */ -static UCS_F_ALWAYS_INLINE int -ucp_proto_rndv_mtype_is_rtr(ucp_request_t *req) -{ - const char *proto_name = req->send.proto_config->proto->name; - return (strstr(proto_name, "rtr") != NULL); -} - /** * Check if request should be throttled due to flow control limit. * If throttled, the request is queued to the appropriate priority queue. * - * Priority and quota allocation (based on memory release and system impact): - * - PUT: Full quota, remote already sent RTR and is waiting. Completing PUT - * unblocks remote. - * - GET: 75% quota, self-contained operation, no remote allocation triggered. - * - RTR: 50% quota, scheduling RTR causes additional allocation on remote - * (RNDV PUT). + * @param req The request to check. + * @param max_frags The maximum number of fragments allowed. + * @param pending_q The queue to add the request to if it is throttled. * - * @return UCS_OK if not throttled (caller should continue), - * UCS_ERR_NO_RESOURCE if throttled and queued. + * @return UCS_OK if not throttled, UCS_ERR_NO_RESOURCE if throttled and queued. */ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_fc_check(ucp_request_t *req, size_t max_frags, @@ -241,9 +227,9 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) * * Priority rationale: * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote - * immediately. + * as well. * GET - Self-contained fetch operation. Completes without causing remote - * allocations. + * allocations, but scheduling it doesn't unblock another buffer. * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total * memory pressure. */ From 8275c3660cb27f4cc3700faa59b5336e3d6a35d1 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 13:53:09 +0200 Subject: [PATCH 08/67] UCP/RNDV: change fragments distibution across different protos --- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_rtr.c | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 4d6d3f0339b..84ba27aa018 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -311,7 +311,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 4; + max_frags = context->config.ext.rndv_mtype_worker_max_frags * 4 / 5; pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 4f5107b4526..d844f4d7ca6 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -437,7 +437,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags / 2; + max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 5; pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { From 2701d7fcdb055e47ccddbc799d4f32b455d642dd Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 17:26:09 +0200 Subject: [PATCH 09/67] UCP/RNDV: some changes --- src/ucp/core/ucp_worker.c | 2 +- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_mtype.inl | 1 - src/ucp/rndv/rndv_rtr.c | 2 +- 4 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index b5af3a461b4..1330bf49cf0 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2639,7 +2639,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_closures = 0; worker->counters.ep_failures = 0; - /* Initialize RNDV mtype flow control with separate priority queues */ + /* Initialize RNDV mtype flow control */ worker->rndv_mtype_fc.active_frags = 0; ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 84ba27aa018..7ebd364312a 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -311,7 +311,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags * 4 / 5; + max_frags = context->config.ext.rndv_mtype_worker_max_frags / 5 * 4; pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 305a34b1696..bf783ddabbd 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -10,7 +10,6 @@ #include "rndv.h" #include -#include static ucp_ep_h ucp_proto_rndv_mtype_ep(ucp_worker_t *worker, diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index d844f4d7ca6..99b47e704e6 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -437,7 +437,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags * 3 / 5; + max_frags = context->config.ext.rndv_mtype_worker_max_frags / 5 * 3; pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { From 7766da1633073e40278bba55007dd69f04cca1e4 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 31 Dec 2025 17:26:41 +0200 Subject: [PATCH 10/67] UCP/RNDV: Added gtest --- src/ucp/core/ucp_worker.c | 26 ++++--- src/ucp/core/ucp_worker.h | 2 + src/ucp/rndv/rndv_mtype.inl | 4 + test/gtest/ucp/test_ucp_am.cc | 142 ++++++++++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 12 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 1330bf49cf0..6ee58f05254 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -88,18 +88,20 @@ static ucs_stats_class_t ucp_worker_stats_class = { .num_counters = UCP_WORKER_STAT_LAST, .class_id = UCS_STATS_CLASS_ID_INVALID, .counter_names = { - [UCP_WORKER_STAT_TAG_RX_EAGER_MSG] = "tag_rx_eager_msg", - [UCP_WORKER_STAT_TAG_RX_EAGER_SYNC_MSG] = "tag_rx_sync_msg", - [UCP_WORKER_STAT_TAG_RX_EAGER_CHUNK_EXP] = "tag_rx_eager_chunk_exp", - [UCP_WORKER_STAT_TAG_RX_EAGER_CHUNK_UNEXP] = "tag_rx_eager_chunk_unexp", - [UCP_WORKER_STAT_RNDV_RX_EXP] = "rndv_rx_exp", - [UCP_WORKER_STAT_RNDV_RX_UNEXP] = "rndv_rx_unexp", - [UCP_WORKER_STAT_RNDV_PUT_ZCOPY] = "rndv_put_zcopy", - [UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY] = "rndv_put_mtype_zcopy", - [UCP_WORKER_STAT_RNDV_GET_ZCOPY] = "rndv_get_zcopy", - [UCP_WORKER_STAT_RNDV_RTR] = "rndv_rtr", - [UCP_WORKER_STAT_RNDV_RTR_MTYPE] = "rndv_rtr_mtype", - [UCP_WORKER_STAT_RNDV_RKEY_PTR] = "rndv_rkey_ptr" + [UCP_WORKER_STAT_TAG_RX_EAGER_MSG] = "tag_rx_eager_msg", + [UCP_WORKER_STAT_TAG_RX_EAGER_SYNC_MSG] = "tag_rx_sync_msg", + [UCP_WORKER_STAT_TAG_RX_EAGER_CHUNK_EXP] = "tag_rx_eager_chunk_exp", + [UCP_WORKER_STAT_TAG_RX_EAGER_CHUNK_UNEXP] = "tag_rx_eager_chunk_unexp", + [UCP_WORKER_STAT_RNDV_RX_EXP] = "rndv_rx_exp", + [UCP_WORKER_STAT_RNDV_RX_UNEXP] = "rndv_rx_unexp", + [UCP_WORKER_STAT_RNDV_PUT_ZCOPY] = "rndv_put_zcopy", + [UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY] = "rndv_put_mtype_zcopy", + [UCP_WORKER_STAT_RNDV_GET_ZCOPY] = "rndv_get_zcopy", + [UCP_WORKER_STAT_RNDV_RTR] = "rndv_rtr", + [UCP_WORKER_STAT_RNDV_RTR_MTYPE] = "rndv_rtr_mtype", + [UCP_WORKER_STAT_RNDV_RKEY_PTR] = "rndv_rkey_ptr", + [UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED] = "rndv_mtype_fc_throttled", + [UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED] = "rndv_mtype_fc_incremented" } }; #endif diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 391ac3e2559..50e8cc1dc71 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -154,6 +154,8 @@ enum { UCP_WORKER_STAT_RNDV_RTR, UCP_WORKER_STAT_RNDV_RTR_MTYPE, UCP_WORKER_STAT_RNDV_RKEY_PTR, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, UCP_WORKER_STAT_LAST }; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index bf783ddabbd..b2eb5710329 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -200,6 +200,8 @@ ucp_proto_rndv_mtype_fc_check(ucp_request_t *req, size_t max_frags, if (worker->rndv_mtype_fc.active_frags >= max_frags) { ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", worker->rndv_mtype_fc.active_frags, max_frags); + UCS_STATS_UPDATE_COUNTER(worker->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); return UCS_ERR_NO_RESOURCE; } @@ -217,6 +219,8 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) if (worker->context->config.ext.rndv_mtype_worker_fc_enable) { worker->rndv_mtype_fc.active_frags++; + UCS_STATS_UPDATE_COUNTER(worker->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, 1); } } diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index af41696f2e2..15b63d0cf88 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2346,6 +2346,148 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_ppln, cuda_managed_buff, UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_ppln); + +class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { +protected: + void check_stats_ge(entity &e, uint64_t cntr, uint64_t min_value) + { + auto stats_node = e.worker()->stats; + auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); + + EXPECT_GE(value, min_value) << "counter " + << stats_node->cls->counter_names[cntr] + << " expected >= " << min_value + << " but got " << value; + } + + void check_active_frags_zero(entity &e) + { + EXPECT_EQ(0u, e.worker()->rndv_mtype_fc.active_frags) + << "active_frags should be 0 after completion"; + } + + void check_pending_queues_empty(entity &e) + { + ucp_worker_h worker = e.worker(); + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) + << "put_pending_q should be empty"; + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.get_pending_q)) + << "get_pending_q should be empty"; + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.rtr_pending_q)) + << "rtr_pending_q should be empty"; + } + + void send_message(size_t num_frags) + { + set_mem_type(UCS_MEMORY_TYPE_CUDA_MANAGED); + test_am_send_recv(get_rndv_frag_size(UCS_MEMORY_TYPE_CUDA) * num_frags); + } + + void verify_clean_fc_state() + { + check_active_frags_zero(sender()); + check_active_frags_zero(receiver()); + check_pending_queues_empty(sender()); + check_pending_queues_empty(receiver()); + } +}; + +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, + "RNDV_MTYPE_WORKER_FC_ENABLE=y", + "RNDV_MTYPE_WORKER_MAX_FRAGS=128", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + if (!sender().is_rndv_put_ppln_supported()) { + UCS_TEST_SKIP_R("RNDV is not supported"); + } + + send_message(8); + + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + + /* Throttling should NOT have occurred */ + auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + EXPECT_EQ(0u, sender_throttled) << "sender should not be throttled"; + EXPECT_EQ(0u, receiver_throttled) << "receiver should not be throttled"; + + /* FC should have been active (incremented) even though no throttling */ + auto sender_incremented = UCS_STATS_GET_COUNTER( + sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); + auto receiver_incremented = UCS_STATS_GET_COUNTER( + receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); + EXPECT_GT(sender_incremented + receiver_incremented, 0u) + << "FC should be active and tracking fragments"; + + verify_clean_fc_state(); +} + +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, + "RNDV_MTYPE_WORKER_FC_ENABLE=y", + "RNDV_MTYPE_WORKER_MAX_FRAGS=5", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + if (!sender().is_rndv_put_ppln_supported()) { + UCS_TEST_SKIP_R("RNDV is not supported"); + } + + send_message(20); + + /* Verify mtype protocols were used */ + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + + /* Throttling SHOULD have occurred (at least once on sender or receiver) */ + auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + EXPECT_GT(sender_throttled + receiver_throttled, 0u) + << "throttling should have occurred with MAX_FRAGS=5"; + + verify_clean_fc_state(); +} + +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, + "RNDV_MTYPE_WORKER_FC_ENABLE=n", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + if (!sender().is_rndv_put_ppln_supported()) { + UCS_TEST_SKIP_R("RNDV is not supported"); + } + + send_message(8); + + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + + /* No throttling should have occurred (FC is disabled) */ + auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + EXPECT_EQ(0u, sender_throttled) << "FC disabled - no throttling expected"; + EXPECT_EQ(0u, receiver_throttled) << "FC disabled - no throttling expected"; + + /* With FC disabled, increment should NOT be called */ + auto sender_incremented = UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); + auto receiver_incremented = UCS_STATS_GET_COUNTER(receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); + EXPECT_EQ(0u, sender_incremented) << "FC disabled - no increment expected"; + EXPECT_EQ(0u, receiver_incremented) << "FC disabled - no increment expected"; + + verify_clean_fc_state(); +} + + +UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_mtype_fc); + #endif /* Test class for AM PSN protocol */ From 4e515c3a8b6c79f423b81c1857047c9756ddc937 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Sun, 4 Jan 2026 14:29:30 +0200 Subject: [PATCH 11/67] UCP/RNDV: pre-compute max frags --- src/ucp/core/ucp_context.c | 9 +++++---- src/ucp/core/ucp_context.h | 4 ++-- src/ucp/rndv/proto_rndv.c | 3 +++ src/ucp/rndv/proto_rndv.h | 3 +++ src/ucp/rndv/rndv_get.c | 7 ++++--- src/ucp/rndv/rndv_mtype.inl | 19 +++++++++++++++++++ src/ucp/rndv/rndv_put.c | 9 ++++----- src/ucp/rndv/rndv_rtr.c | 15 ++++++++++----- 8 files changed, 50 insertions(+), 19 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index ee9c5e57fa0..b38a226cd9c 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -409,10 +409,11 @@ static ucs_config_field_t ucp_context_config_table[] = { "across all requests, preventing memory exhaustion", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_fc_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_MAX_FRAGS", "1024", - "Maximum number of concurrent mtype fragments per worker\n" - "(only applies when RNDV_MTYPE_WORKER_FC_ENABLE=y)", - ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_frags), UCS_CONFIG_TYPE_ULUNITS}, + {"RNDV_MTYPE_WORKER_MAX_MEM", "4g", + "Maximum memory for concurrent mtype fragments per worker.\n" + "This value is translated to a fragment count based on RNDV_FRAG_SIZE\n" + "for each memory type (only applies when RNDV_MTYPE_WORKER_FC_ENABLE=y)", + ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 302163d8e82..61f2e7aefbb 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -134,8 +134,8 @@ typedef struct ucp_context_config { int rma_ppln_enable; /** Enable flow control for rndv mtype fragments at worker level */ int rndv_mtype_worker_fc_enable; - /** Maximum number of concurrent rndv mtype fragments per worker */ - size_t rndv_mtype_worker_max_frags; + /** Maximum memory for concurrent rndv mtype fragments per worker (bytes) */ + size_t rndv_mtype_worker_max_mem; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 33e5515b612..8d4c0b2d31b 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -9,6 +9,7 @@ #endif #include "proto_rndv.inl" +#include "rndv_mtype.inl" #include #include @@ -732,6 +733,8 @@ ucp_proto_rndv_bulk_init(const ucp_proto_multi_init_params_t *init_params, rpriv->frag_mem_type = init_params->super.reg_mem_info.type; rpriv->frag_sys_dev = init_params->super.reg_mem_info.sys_dev; + rpriv->fc_max_frags = ucp_proto_rndv_mtype_fc_max_frags( + context, rpriv->frag_mem_type); if (rpriv->super.lane == UCP_NULL_LANE) { /* Add perf without ACK in case of pipeline */ diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 3ee14ac1464..fd3ce5f30de 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -93,6 +93,9 @@ typedef struct { ucs_memory_type_t frag_mem_type; ucs_sys_device_t frag_sys_dev; + /* max fragments for flow control */ + size_t fc_max_frags; + /* Multi-lane common part. Must be the last field, see @ref ucp_proto_multi_priv_t */ ucp_proto_multi_priv_t mpriv; diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 7ebd364312a..5ae58c9956d 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -298,8 +298,7 @@ ucp_proto_rndv_get_mtype_fetch_completion(uct_completion_t *uct_comp) static ucs_status_t ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); - ucp_context_h context = req->send.ep->worker->context; + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; size_t max_frags; @@ -309,9 +308,11 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + /* GET priority: 80% of total fragments */ + max_frags = rpriv->fc_max_frags / 5 * 4; + /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags / 5 * 4; pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index b2eb5710329..869b3dbdfdb 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -176,6 +176,25 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) return 1; } +/** + * Compute maximum number of fragments allowed based on configured max memory + * and fragment size for the given memory type. + * + * @param context The UCP context. + * @param frag_mem_type Memory type used for fragments. + * + * @return Maximum number of fragments that fit within the configured memory limit. + */ +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, + ucs_memory_type_t frag_mem_type) +{ + size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; + size_t frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; + + return (frag_size > 0) ? (max_mem / frag_size) : SIZE_MAX; +} + /** * Check if request should be throttled due to flow control limit. * If throttled, the request is queued to the appropriate priority queue. diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index a5d6d9d6d1d..e9f2559bc79 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,8 +520,7 @@ static UCS_F_INLINE_OPTIMIZED ucs_status_t ucp_proto_rndv_put_mtype_send_func( static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); - ucp_context_h context = req->send.ep->worker->context; + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; size_t max_frags; @@ -529,16 +528,16 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); + rpriv = req->send.proto_config->priv; + /* Check throttling limit. If no resource at the moment, queue the request * in PUT pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags; + max_frags = rpriv->bulk.fc_max_frags; pending_q = &req->send.ep->worker->rndv_mtype_fc.put_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } - - rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev); if (status != UCS_OK) { diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 99b47e704e6..cf6a628d8e5 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -39,6 +39,8 @@ typedef struct { ucp_proto_rndv_rtr_priv_t super; ucs_memory_type_t frag_mem_type; ucs_sys_device_t frag_sys_dev; + /* max fragments for flow control */ + size_t fc_max_frags; } ucp_proto_rndv_rtr_mtype_priv_t; @@ -427,24 +429,25 @@ ucp_proto_rndv_rtr_mtype_data_received(ucp_request_t *req, int in_buffer) static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { - ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); - ucp_context_h context = req->send.ep->worker->context; + ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; size_t max_frags; ucs_queue_head_t *pending_q; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + rpriv = req->send.proto_config->priv; + + /* RTR priority: 60% of total fragments */ + max_frags = rpriv->fc_max_frags / 5 * 3; + /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ - max_frags = context->config.ext.rndv_mtype_worker_max_frags / 5 * 3; pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } - - rpriv = req->send.proto_config->priv; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { @@ -549,6 +552,8 @@ ucp_proto_rndv_rtr_mtype_probe(const ucp_proto_init_params_t *init_params) rpriv.frag_mem_type = frag_mem_type; rpriv.frag_sys_dev = params.super.reg_mem_info.sys_dev; params.flags = ucp_proto_rndv_ctrl_init_flags(¶ms); + rpriv.fc_max_frags = ucp_proto_rndv_mtype_fc_max_frags( + context, frag_mem_type); ucp_proto_rndv_ctrl_probe(¶ms, &rpriv, sizeof(rpriv)); out_unpack_perf_destroy: From 50b208fc1a80f405bc0ff440687b8129a684101e Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Sun, 4 Jan 2026 16:37:34 +0200 Subject: [PATCH 12/67] UCP/RNDV: align max frags to mpool chunk size --- src/ucp/rndv/rndv_mtype.inl | 29 ++++++++++++++++++++++++----- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 869b3dbdfdb..46852ce73be 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -178,21 +178,40 @@ static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) /** * Compute maximum number of fragments allowed based on configured max memory - * and fragment size for the given memory type. + * and fragment size for the given memory type. The result is rounded down to + * the allocation chunk granularity (rndv_num_frags). * * @param context The UCP context. * @param frag_mem_type Memory type used for fragments. * - * @return Maximum number of fragments that fit within the configured memory limit. + * @return Maximum number of fragments that fit within the configured memory + * limit, aligned to allocation chunk size. */ static UCS_F_ALWAYS_INLINE size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, ucs_memory_type_t frag_mem_type) { - size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; - size_t frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; + size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; + size_t frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; + size_t frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; + size_t max_frags; + + /* frag_size must be > 0 for mtype protocols - validated during proto init */ + ucs_assert(frag_size > 0); + + /* Compute max fragments and round down to allocation chunk granularity */ + max_frags = max_mem / frag_size; + max_frags = (max_frags / frags_in_chunk) * frags_in_chunk; + + if (max_frags == 0) { + ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " + "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " + "frags", max_mem, ucs_memory_type_names[frag_mem_type], + frag_size, frags_in_chunk, frags_in_chunk); + return frags_in_chunk; + } - return (frag_size > 0) ? (max_mem / frag_size) : SIZE_MAX; + return max_frags; } /** From ac696c76d4187b53b1e3f3725b2476be750e1761 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 5 Jan 2026 10:54:19 +0200 Subject: [PATCH 13/67] UCP/RNDV: adjust gtest to MAX_MEM --- src/ucp/rndv/rndv_mtype.inl | 1 - test/gtest/ucp/test_ucp_am.cc | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 46852ce73be..9d145ceff02 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -196,7 +196,6 @@ ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, size_t frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; size_t max_frags; - /* frag_size must be > 0 for mtype protocols - validated during proto init */ ucs_assert(frag_size > 0); /* Compute max fragments and round down to allocation chunk granularity */ diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 15b63d0cf88..3bd5f2ad6a9 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2394,7 +2394,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, "RNDV_MTYPE_WORKER_FC_ENABLE=y", - "RNDV_MTYPE_WORKER_MAX_FRAGS=128", + "RNDV_MTYPE_WORKER_MAX_MEM=1g", "RNDV_FRAG_MEM_TYPE=cuda") { if (!sender().is_rndv_put_ppln_supported()) { @@ -2429,14 +2429,14 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, "RNDV_MTYPE_WORKER_FC_ENABLE=y", - "RNDV_MTYPE_WORKER_MAX_FRAGS=5", + "RNDV_MTYPE_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") { if (!sender().is_rndv_put_ppln_supported()) { UCS_TEST_SKIP_R("RNDV is not supported"); } - send_message(20); + send_message(200); /* Verify mtype protocols were used */ check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); @@ -2448,7 +2448,7 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); EXPECT_GT(sender_throttled + receiver_throttled, 0u) - << "throttling should have occurred with MAX_FRAGS=5"; + << "throttling should have occurred with MAX_MEM=600mb"; verify_clean_fc_state(); } From e97bc99c6b2dd59b710c0ad3f20bfbadefc2e917 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 7 Jan 2026 15:37:03 +0200 Subject: [PATCH 14/67] UCP/RNDV: pre-define stats-counter-array size --- src/ucs/stats/libstats.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ucs/stats/libstats.h b/src/ucs/stats/libstats.h index 3b72c5fb7eb..63e5df1a533 100644 --- a/src/ucs/stats/libstats.h +++ b/src/ucs/stats/libstats.h @@ -56,7 +56,9 @@ struct ucs_stats_class { const char *name; unsigned num_counters; unsigned class_id; - const char* counter_names[]; + const char* counter_names[14]; /* Need to maintain the number of + counters once it gets bigger + than 14 in one of the classes */ }; /* From 45ef50455add3ef5bc7963d65c8d97e341774b21 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Sun, 8 Mar 2026 18:38:34 +0200 Subject: [PATCH 15/67] UCP/RNDV: code review fixes --- src/ucp/core/ucp_context.c | 9 +-- src/ucp/core/ucp_context.h | 9 ++- src/ucp/core/ucp_worker.c | 13 ++-- src/ucp/core/ucp_worker.h | 1 + src/ucp/rndv/proto_rndv.c | 43 ++++++++++++ src/ucp/rndv/proto_rndv.h | 44 ++++++++++++ src/ucp/rndv/rndv_get.c | 5 +- src/ucp/rndv/rndv_mtype.inl | 83 ++++++---------------- src/ucp/rndv/rndv_put.c | 6 +- src/ucp/rndv/rndv_rtr.c | 5 +- src/ucs/stats/libstats.h | 4 +- test/gtest/ucp/test_ucp_am.cc | 128 ++++++++++++++-------------------- 12 files changed, 189 insertions(+), 161 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index b38a226cd9c..916f07bfd17 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,15 +404,10 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_FC_ENABLE", "n", - "Enable worker-level flow control to limit total concurrent mtype fragments\n" - "across all requests, preventing memory exhaustion", - ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_fc_enable), UCS_CONFIG_TYPE_BOOL}, - - {"RNDV_MTYPE_WORKER_MAX_MEM", "4g", + {"RNDV_MTYPE_WORKER_MAX_MEM", "inf", "Maximum memory for concurrent mtype fragments per worker.\n" "This value is translated to a fragment count based on RNDV_FRAG_SIZE\n" - "for each memory type (only applies when RNDV_MTYPE_WORKER_FC_ENABLE=y)", + "for each memory type. Set to \"inf\" to disable flow control.", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 61f2e7aefbb..2ddabd30755 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -132,8 +132,6 @@ typedef struct ucp_context_config { int rndv_errh_ppln_enable; /** Force-enable the RMA rendezvous put/get protocols */ int rma_ppln_enable; - /** Enable flow control for rndv mtype fragments at worker level */ - int rndv_mtype_worker_fc_enable; /** Maximum memory for concurrent rndv mtype fragments per worker (bytes) */ size_t rndv_mtype_worker_max_mem; /** Threshold for using tag matching offload capabilities. Smaller buffers @@ -807,6 +805,13 @@ ucp_context_print_transport_tables_enabled(ucp_context_h context) return value == UCS_CONFIG_ON; } +static UCS_F_ALWAYS_INLINE int +ucp_context_rndv_mtype_fc_enabled(ucp_context_h context) +{ + size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; + return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); +} + void ucp_context_memaccess_tl_bitmap(ucp_context_h context, uint64_t mem_type_bitmap, uint64_t md_reg_flags, diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 6ee58f05254..7f47ae7d0fe 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2641,11 +2641,14 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_closures = 0; worker->counters.ep_failures = 0; - /* Initialize RNDV mtype flow control */ - worker->rndv_mtype_fc.active_frags = 0; - ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); + /* Initialize RNDV mtype flow control if configured */ + worker->rndv_mtype_fc.enabled = ucp_context_rndv_mtype_fc_enabled(context); + if (worker->rndv_mtype_fc.enabled) { + worker->rndv_mtype_fc.active_frags = 0; + ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); + } /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 50e8cc1dc71..37352387382 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -400,6 +400,7 @@ typedef struct ucp_worker { struct { /* Worker-level mtype fragment flow control */ + int enabled; /* Whether FC is enabled (cached) */ size_t active_frags; /* Current active fragments */ /* Separate pending queues by priority (PUT > GET > RTR) */ ucs_queue_head_t put_pending_q; /* Throttled PUT (RNDV_SEND) requests */ diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 8d4c0b2d31b..b7395c104c9 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -680,6 +680,49 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } +/** + * Compute maximum number of fragments allowed based on configured max memory + * and fragment size for the given memory type. The result is rounded down to + * the allocation chunk granularity (rndv_num_frags). + * + * @param context The UCP context. + * @param frag_mem_type Memory type used for fragments. + * + * @return Maximum number of fragments that fit within the configured memory + * limit, aligned to allocation chunk size. + */ +size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, + ucs_memory_type_t frag_mem_type) +{ + size_t max_mem; + size_t frag_size; + size_t frags_in_chunk; + size_t max_frags; + + if (!ucp_context_rndv_mtype_fc_enabled(context)) { + return SIZE_MAX; + } + + max_mem = context->config.ext.rndv_mtype_worker_max_mem; + frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; + frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; + ucs_assert(frag_size > 0); + + /* Compute max fragments and round down to allocation chunk granularity */ + max_frags = max_mem / frag_size; + max_frags = (max_frags / frags_in_chunk) * frags_in_chunk; + + if (max_frags == 0) { + ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " + "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " + "frags", max_mem, ucs_memory_type_names[frag_mem_type], + frag_size, frags_in_chunk, frags_in_chunk); + return frags_in_chunk; + } + + return max_frags; +} + ucs_status_t ucp_proto_rndv_ack_init(const ucp_proto_common_init_params_t *init_params, const char *name, double overhead, diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index fd3ce5f30de..203c152f07f 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -38,6 +38,35 @@ ucp_proto_rndv_rts_tag_op_flags(ucp_rndv_rts_opcode_t opcode) UCP_PROTO_SELECT_OP_FLAG_TAG_RNDV : 0; } +/* + * Staging-buffer flow-control throttle limit per operation type. + * + * To prevent deadlock under memory pressure, each operation type is capped at a + * different share of the total fragment budget (fc_max_frags). Operations that + * free more resources on completion receive a larger cap, so they can always + * make progress and unblock others: + * + * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. + * GET (level 1, ~90%) – frees local staging buffer only. + * RTR (level 2, ~80%) – triggers a remote PUT allocation, adding pressure. + * + * The step between tiers is 10% of the budget (1/10). The exact fractions are + * not performance-sensitive; only the strict ordering PUT > GET > RTR matters. + * The same ordering is used when dequeueing pending requests, see + * ucp_proto_rndv_mtype_fc_decrement(). + */ +#define UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, _level) \ + ((_fc_max) - (_level) * ((_fc_max) / 10)) + +#define UCP_PROTO_RNDV_MTYPE_FC_PUT_LIMIT(_fc_max) \ + UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 0) + +#define UCP_PROTO_RNDV_MTYPE_FC_GET_LIMIT(_fc_max) \ + UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 1) + +#define UCP_PROTO_RNDV_MTYPE_FC_RTR_LIMIT(_fc_max) \ + UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 2) + /** * Rendezvous protocol which sends a control message to the remote peer, and not @@ -196,6 +225,21 @@ void ucp_proto_rndv_rts_abort(ucp_request_t *req, ucs_status_t status); ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); +/** + * Compute maximum number of fragments allowed based on configured max memory + * and fragment size for the given memory type. The result is rounded down to + * the allocation chunk granularity (rndv_num_frags). + * + * @param context The UCP context. + * @param frag_mem_type Memory type used for fragments. + * + * @return Maximum number of fragments that fit within the configured memory + * limit, aligned to allocation chunk size. + */ +size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, + ucs_memory_type_t frag_mem_type); + + ucs_status_t ucp_proto_rndv_ack_init(const ucp_proto_common_init_params_t *init_params, const char *name, double overhead, diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 5ae58c9956d..ce5653c95a7 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -308,13 +308,12 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - /* GET priority: 80% of total fragments */ - max_frags = rpriv->fc_max_frags / 5 * 4; + max_frags = UCP_PROTO_RNDV_MTYPE_FC_GET_LIMIT(rpriv->fc_max_frags); /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; - if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 9d145ceff02..eec4dd47c37 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -169,50 +169,14 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_copy( } /* Reschedule callback for throttled mtype requests */ -static unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +static UCS_F_ALWAYS_INLINE +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) { ucp_request_t *req = arg; ucp_request_send(req); return 1; } -/** - * Compute maximum number of fragments allowed based on configured max memory - * and fragment size for the given memory type. The result is rounded down to - * the allocation chunk granularity (rndv_num_frags). - * - * @param context The UCP context. - * @param frag_mem_type Memory type used for fragments. - * - * @return Maximum number of fragments that fit within the configured memory - * limit, aligned to allocation chunk size. - */ -static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, - ucs_memory_type_t frag_mem_type) -{ - size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; - size_t frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; - size_t frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; - size_t max_frags; - - ucs_assert(frag_size > 0); - - /* Compute max fragments and round down to allocation chunk granularity */ - max_frags = max_mem / frag_size; - max_frags = (max_frags / frags_in_chunk) * frags_in_chunk; - - if (max_frags == 0) { - ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " - "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " - "frags", max_mem, ucs_memory_type_names[frag_mem_type], - frag_size, frags_in_chunk, frags_in_chunk); - return frags_in_chunk; - } - - return max_frags; -} - /** * Check if request should be throttled due to flow control limit. * If throttled, the request is queued to the appropriate priority queue. @@ -224,13 +188,12 @@ ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, * @return UCS_OK if not throttled, UCS_ERR_NO_RESOURCE if throttled and queued. */ static UCS_F_ALWAYS_INLINE ucs_status_t -ucp_proto_rndv_mtype_fc_check(ucp_request_t *req, size_t max_frags, - ucs_queue_head_t *pending_q) +ucp_proto_rndv_mtype_fc_throttle(ucp_request_t *req, const size_t max_frags, + ucs_queue_head_t *pending_q) { - ucp_worker_h worker = req->send.ep->worker; - ucp_context_h context = worker->context; + ucp_worker_h worker = req->send.ep->worker; - if (!context->config.ext.rndv_mtype_worker_fc_enable) { + if (!worker->rndv_mtype_fc.enabled) { return UCS_OK; } @@ -254,16 +217,15 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; - if (worker->context->config.ext.rndv_mtype_worker_fc_enable) { - worker->rndv_mtype_fc.active_frags++; - UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, 1); - } + worker->rndv_mtype_fc.active_frags++; + UCS_STATS_UPDATE_COUNTER(worker->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, 1); } /** * Decrement active_frags counter and reschedule pending request. - * Dequeue priority: PUT > GET > RTR + * Dequeue priority: PUT > GET > RTR (same ordering as the throttle limits + * defined by UCP_PROTO_RNDV_MTYPE_FC_LIMIT). * * Priority rationale: * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote @@ -277,17 +239,16 @@ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; - ucp_context_h context = worker->context; ucs_queue_elem_t *elem = NULL; ucp_request_t *pending_req; - if (!context->config.ext.rndv_mtype_worker_fc_enable) { - return; - } - ucs_assert(worker->rndv_mtype_fc.active_frags > 0); worker->rndv_mtype_fc.active_frags--; + if (!worker->rndv_mtype_fc.enabled) { + return; + } + /* Dequeue with priority: PUT > GET > RTR */ if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) { elem = ucs_queue_pull(&worker->rndv_mtype_fc.put_pending_q); @@ -297,13 +258,15 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) elem = ucs_queue_pull(&worker->rndv_mtype_fc.rtr_pending_q); } - if (elem != NULL) { - pending_req = ucs_container_of(elem, ucp_request_t, - send.rndv.ppln.queue_elem); - ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, - ucp_proto_rndv_mtype_fc_reschedule_cb, - pending_req); + if (elem == NULL) { + return; } + + pending_req = ucs_container_of(elem, ucp_request_t, + send.rndv.ppln.queue_elem); + ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, + ucp_proto_rndv_mtype_fc_reschedule_cb, + pending_req); } static UCS_F_ALWAYS_INLINE ucs_status_t diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index e9f2559bc79..1fb5884033e 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -530,11 +530,9 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; - /* Check throttling limit. If no resource at the moment, queue the request - * in PUT pending queue and return UCS_OK. */ - max_frags = rpriv->bulk.fc_max_frags; + max_frags = UCP_PROTO_RNDV_MTYPE_FC_PUT_LIMIT(rpriv->bulk.fc_max_frags); pending_q = &req->send.ep->worker->rndv_mtype_fc.put_pending_q; - if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index cf6a628d8e5..c0f14272111 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -438,13 +438,12 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { rpriv = req->send.proto_config->priv; - /* RTR priority: 60% of total fragments */ - max_frags = rpriv->fc_max_frags / 5 * 3; + max_frags = UCP_PROTO_RNDV_MTYPE_FC_RTR_LIMIT(rpriv->fc_max_frags); /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; - if (ucp_proto_rndv_mtype_fc_check(req, max_frags, pending_q) == + if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucs/stats/libstats.h b/src/ucs/stats/libstats.h index 63e5df1a533..3b72c5fb7eb 100644 --- a/src/ucs/stats/libstats.h +++ b/src/ucs/stats/libstats.h @@ -56,9 +56,7 @@ struct ucs_stats_class { const char *name; unsigned num_counters; unsigned class_id; - const char* counter_names[14]; /* Need to maintain the number of - counters once it gets bigger - than 14 in one of the classes */ + const char* counter_names[]; }; /* diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 3bd5f2ad6a9..7218faee5ba 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2349,7 +2349,14 @@ UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_ppln); class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { protected: - void check_stats_ge(entity &e, uint64_t cntr, uint64_t min_value) + struct fc_counters { + uint64_t sender_throttled; + uint64_t receiver_throttled; + uint64_t sender_incremented; + uint64_t receiver_incremented; + }; + + void check_stats_ge(const entity &e, uint64_t cntr, uint64_t min_value) { auto stats_node = e.worker()->stats; auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); @@ -2360,13 +2367,13 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { << " but got " << value; } - void check_active_frags_zero(entity &e) + void check_active_frags_zero(const entity &e) { EXPECT_EQ(0u, e.worker()->rndv_mtype_fc.active_frags) << "active_frags should be 0 after completion"; } - void check_pending_queues_empty(entity &e) + void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) @@ -2385,104 +2392,77 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { void verify_clean_fc_state() { - check_active_frags_zero(sender()); - check_active_frags_zero(receiver()); - check_pending_queues_empty(sender()); - check_pending_queues_empty(receiver()); - } -}; - -UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, - "RNDV_MTYPE_WORKER_FC_ENABLE=y", - "RNDV_MTYPE_WORKER_MAX_MEM=1g", - "RNDV_FRAG_MEM_TYPE=cuda") -{ - if (!sender().is_rndv_put_ppln_supported()) { - UCS_TEST_SKIP_R("RNDV is not supported"); + for (auto *ep : { &sender(), &receiver() }) + { + check_active_frags_zero(*ep); + check_pending_queues_empty(*ep); + } } - send_message(8); + void run_fc_test(size_t num_frags, fc_counters &fc) + { + if (!sender().is_rndv_put_ppln_supported()) { + UCS_TEST_SKIP_R("RNDV is not supported"); + } - check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); - check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + send_message(num_frags); - /* Throttling should NOT have occurred */ - auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - EXPECT_EQ(0u, sender_throttled) << "sender should not be throttled"; - EXPECT_EQ(0u, receiver_throttled) << "receiver should not be throttled"; + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); - /* FC should have been active (incremented) even though no throttling */ - auto sender_incremented = UCS_STATS_GET_COUNTER( + fc.sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + fc.receiver_throttled = UCS_STATS_GET_COUNTER( + receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + fc.sender_incremented = UCS_STATS_GET_COUNTER( sender().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); - auto receiver_incremented = UCS_STATS_GET_COUNTER( + fc.receiver_incremented = UCS_STATS_GET_COUNTER( receiver().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); - EXPECT_GT(sender_incremented + receiver_incremented, 0u) - << "FC should be active and tracking fragments"; + } +}; + +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, + "RNDV_MTYPE_WORKER_MAX_MEM=1g", "RNDV_FRAG_MEM_TYPE=cuda") +{ + fc_counters fc; + + run_fc_test(8, fc); + + EXPECT_EQ(0u, fc.sender_throttled) << "sender should not be throttled"; + EXPECT_EQ(0u, fc.receiver_throttled) << "receiver should not be throttled"; + EXPECT_GT(fc.sender_incremented + fc.receiver_incremented, 0u) + << "FC should be active and tracking fragments"; verify_clean_fc_state(); } UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, - "RNDV_MTYPE_WORKER_FC_ENABLE=y", - "RNDV_MTYPE_WORKER_MAX_MEM=600mb", - "RNDV_FRAG_MEM_TYPE=cuda") + "RNDV_MTYPE_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") { - if (!sender().is_rndv_put_ppln_supported()) { - UCS_TEST_SKIP_R("RNDV is not supported"); - } - - send_message(200); + fc_counters fc; - /* Verify mtype protocols were used */ - check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); - check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + run_fc_test(200, fc); - /* Throttling SHOULD have occurred (at least once on sender or receiver) */ - auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - EXPECT_GT(sender_throttled + receiver_throttled, 0u) + EXPECT_GT(fc.sender_throttled + fc.receiver_throttled, 0u) << "throttling should have occurred with MAX_MEM=600mb"; verify_clean_fc_state(); } UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, - "RNDV_MTYPE_WORKER_FC_ENABLE=n", "RNDV_FRAG_MEM_TYPE=cuda") { - if (!sender().is_rndv_put_ppln_supported()) { - UCS_TEST_SKIP_R("RNDV is not supported"); - } - - send_message(8); + fc_counters fc; - check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); - check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + run_fc_test(8, fc); - /* No throttling should have occurred (FC is disabled) */ - auto sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - auto receiver_throttled = UCS_STATS_GET_COUNTER(receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - EXPECT_EQ(0u, sender_throttled) << "FC disabled - no throttling expected"; - EXPECT_EQ(0u, receiver_throttled) << "FC disabled - no throttling expected"; - - /* With FC disabled, increment should NOT be called */ - auto sender_incremented = UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); - auto receiver_incremented = UCS_STATS_GET_COUNTER(receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); - EXPECT_EQ(0u, sender_incremented) << "FC disabled - no increment expected"; - EXPECT_EQ(0u, receiver_incremented) << "FC disabled - no increment expected"; - - verify_clean_fc_state(); + EXPECT_EQ(0u, fc.sender_throttled) << "FC disabled - no throttling expected"; + EXPECT_EQ(0u, fc.receiver_throttled) << "FC disabled - no throttling expected"; + EXPECT_EQ(0u, fc.sender_incremented) << "FC disabled - no increment expected"; + EXPECT_EQ(0u, fc.receiver_incremented) << "FC disabled - no increment expected"; } From 9c3bce903ed409ffe36b072cc46de3a4e01b7a7f Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 16:22:39 +0200 Subject: [PATCH 16/67] UCP/RNDV: cr fixes --- src/ucp/rndv/rndv_mtype.inl | 4 ---- src/ucp/rndv/rndv_rtr.c | 1 + 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index eec4dd47c37..7bef21b7549 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -193,10 +193,6 @@ ucp_proto_rndv_mtype_fc_throttle(ucp_request_t *req, const size_t max_frags, { ucp_worker_h worker = req->send.ep->worker; - if (!worker->rndv_mtype_fc.enabled) { - return UCS_OK; - } - if (worker->rndv_mtype_fc.active_frags >= max_frags) { ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", worker->rndv_mtype_fc.active_frags, max_frags); diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index c0f14272111..edbe73e2d40 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -447,6 +447,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) UCS_ERR_NO_RESOURCE) { return UCS_OK; } + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev); if (status != UCS_OK) { From 5c67c594034ec01dc908a21ed71147a99adf9955 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 16:23:51 +0200 Subject: [PATCH 17/67] UCP/RNDV: STAT_RNDV_MTYPE_FC_THROTTLED -> STAT_RNDV_MTYPE_FC_OK --- src/ucp/core/ucp_worker.c | 4 ++-- src/ucp/core/ucp_worker.h | 2 +- src/ucp/rndv/rndv_mtype.inl | 2 +- test/gtest/ucp/test_ucp_am.cc | 22 +++++++++++----------- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 7f47ae7d0fe..25c8b87bd78 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -100,8 +100,8 @@ static ucs_stats_class_t ucp_worker_stats_class = { [UCP_WORKER_STAT_RNDV_RTR] = "rndv_rtr", [UCP_WORKER_STAT_RNDV_RTR_MTYPE] = "rndv_rtr_mtype", [UCP_WORKER_STAT_RNDV_RKEY_PTR] = "rndv_rkey_ptr", - [UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED] = "rndv_mtype_fc_throttled", - [UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED] = "rndv_mtype_fc_incremented" + [UCP_WORKER_STAT_RNDV_MTYPE_FC_OK] = "rndv_mtype_fc_ok", + [UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED] = "rndv_mtype_fc_throttled" } }; #endif diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 37352387382..ff8869d6b0e 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -154,8 +154,8 @@ enum { UCP_WORKER_STAT_RNDV_RTR, UCP_WORKER_STAT_RNDV_RTR_MTYPE, UCP_WORKER_STAT_RNDV_RKEY_PTR, + UCP_WORKER_STAT_RNDV_MTYPE_FC_OK, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, UCP_WORKER_STAT_LAST }; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 7bef21b7549..d69aafe69c7 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -215,7 +215,7 @@ ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) worker->rndv_mtype_fc.active_frags++; UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED, 1); + UCP_WORKER_STAT_RNDV_MTYPE_FC_OK, 1); } /** diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 7218faee5ba..48bc9bcf907 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2350,10 +2350,10 @@ UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_ppln); class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { protected: struct fc_counters { + uint64_t sender_ok; + uint64_t receiver_ok; uint64_t sender_throttled; uint64_t receiver_throttled; - uint64_t sender_incremented; - uint64_t receiver_incremented; }; void check_stats_ge(const entity &e, uint64_t cntr, uint64_t min_value) @@ -2410,17 +2410,17 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + fc.sender_ok = UCS_STATS_GET_COUNTER( + sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_OK); + fc.receiver_ok = UCS_STATS_GET_COUNTER( + receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_OK); fc.sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); fc.receiver_throttled = UCS_STATS_GET_COUNTER( receiver().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - fc.sender_incremented = UCS_STATS_GET_COUNTER( - sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); - fc.receiver_incremented = UCS_STATS_GET_COUNTER( - receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_INCREMENTED); } }; @@ -2433,7 +2433,7 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, EXPECT_EQ(0u, fc.sender_throttled) << "sender should not be throttled"; EXPECT_EQ(0u, fc.receiver_throttled) << "receiver should not be throttled"; - EXPECT_GT(fc.sender_incremented + fc.receiver_incremented, 0u) + EXPECT_GT(fc.sender_ok + fc.receiver_ok, 0u) << "FC should be active and tracking fragments"; verify_clean_fc_state(); @@ -2461,8 +2461,8 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, EXPECT_EQ(0u, fc.sender_throttled) << "FC disabled - no throttling expected"; EXPECT_EQ(0u, fc.receiver_throttled) << "FC disabled - no throttling expected"; - EXPECT_EQ(0u, fc.sender_incremented) << "FC disabled - no increment expected"; - EXPECT_EQ(0u, fc.receiver_incremented) << "FC disabled - no increment expected"; + EXPECT_EQ(0u, fc.sender_ok) << "FC disabled - no increment expected"; + EXPECT_EQ(0u, fc.receiver_ok) << "FC disabled - no increment expected"; } From 123982177cf2e35bb076cf917b538911e7cab0ee Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 16:46:34 +0200 Subject: [PATCH 18/67] UCP/RNDV: remove UCP_WORKER_STAT_RNDV_MTYPE_FC_OK stats + remove fc_increment func --- src/ucp/core/ucp_worker.c | 1 - src/ucp/core/ucp_worker.h | 1 - src/ucp/rndv/rndv_get.c | 7 ++++--- src/ucp/rndv/rndv_mtype.inl | 13 ------------- src/ucp/rndv/rndv_put.c | 7 ++++--- src/ucp/rndv/rndv_rtr.c | 7 ++++--- test/gtest/ucp/test_ucp_am.cc | 25 ------------------------- 7 files changed, 12 insertions(+), 49 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 25c8b87bd78..dd4bfce4b42 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -100,7 +100,6 @@ static ucs_stats_class_t ucp_worker_stats_class = { [UCP_WORKER_STAT_RNDV_RTR] = "rndv_rtr", [UCP_WORKER_STAT_RNDV_RTR_MTYPE] = "rndv_rtr_mtype", [UCP_WORKER_STAT_RNDV_RKEY_PTR] = "rndv_rkey_ptr", - [UCP_WORKER_STAT_RNDV_MTYPE_FC_OK] = "rndv_mtype_fc_ok", [UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED] = "rndv_mtype_fc_throttled" } }; diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index ff8869d6b0e..8b28abb004c 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -154,7 +154,6 @@ enum { UCP_WORKER_STAT_RNDV_RTR, UCP_WORKER_STAT_RNDV_RTR_MTYPE, UCP_WORKER_STAT_RNDV_RKEY_PTR, - UCP_WORKER_STAT_RNDV_MTYPE_FC_OK, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, UCP_WORKER_STAT_LAST diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index ce5653c95a7..0652d26a66e 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -298,7 +298,8 @@ ucp_proto_rndv_get_mtype_fetch_completion(uct_completion_t *uct_comp) static ucs_status_t ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; size_t max_frags; @@ -312,7 +313,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ - pending_q = &req->send.ep->worker->rndv_mtype_fc.get_pending_q; + pending_q = &worker->rndv_mtype_fc.get_pending_q; if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; @@ -325,7 +326,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) return UCS_OK; } - ucp_proto_rndv_mtype_fc_increment(req); + worker->rndv_mtype_fc.active_frags++; ucp_proto_rndv_get_common_request_init(req); ucp_proto_completion_init(&req->send.state.uct_comp, ucp_proto_rndv_get_mtype_fetch_completion); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index d69aafe69c7..b7e49a16f80 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -205,19 +205,6 @@ ucp_proto_rndv_mtype_fc_throttle(ucp_request_t *req, const size_t max_frags, return UCS_OK; } -/** - * Increment active_frags counter after successful mtype allocation. - */ -static UCS_F_ALWAYS_INLINE void -ucp_proto_rndv_mtype_fc_increment(ucp_request_t *req) -{ - ucp_worker_h worker = req->send.ep->worker; - - worker->rndv_mtype_fc.active_frags++; - UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_OK, 1); -} - /** * Decrement active_frags counter and reschedule pending request. * Dequeue priority: PUT > GET > RTR (same ordering as the throttle limits diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 1fb5884033e..3299741f3cd 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,7 +520,8 @@ static UCS_F_INLINE_OPTIMIZED ucs_status_t ucp_proto_rndv_put_mtype_send_func( static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct);\ + ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; size_t max_frags; @@ -531,7 +532,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; max_frags = UCP_PROTO_RNDV_MTYPE_FC_PUT_LIMIT(rpriv->bulk.fc_max_frags); - pending_q = &req->send.ep->worker->rndv_mtype_fc.put_pending_q; + pending_q = &worker->rndv_mtype_fc.put_pending_q; if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; @@ -543,7 +544,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) return UCS_OK; } - ucp_proto_rndv_mtype_fc_increment(req); + worker->rndv_mtype_fc.active_frags++; ucp_proto_rndv_put_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; ucp_proto_rndv_mdesc_mtype_copy(req, uct_ep_get_zcopy, diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index edbe73e2d40..9e72f68c918 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -429,7 +429,8 @@ ucp_proto_rndv_rtr_mtype_data_received(ucp_request_t *req, int in_buffer) static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { - ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; size_t max_frags; ucs_queue_head_t *pending_q; @@ -442,7 +443,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ - pending_q = &req->send.ep->worker->rndv_mtype_fc.rtr_pending_q; + pending_q = &worker->rndv_mtype_fc.rtr_pending_q; if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { return UCS_OK; @@ -455,7 +456,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) return UCS_OK; } - ucp_proto_rndv_mtype_fc_increment(req); + worker->rndv_mtype_fc.active_frags++; ucp_proto_rtr_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; } diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 48bc9bcf907..7f92bd03cd4 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2350,8 +2350,6 @@ UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_ppln); class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { protected: struct fc_counters { - uint64_t sender_ok; - uint64_t receiver_ok; uint64_t sender_throttled; uint64_t receiver_throttled; }; @@ -2410,12 +2408,6 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); - fc.sender_ok = UCS_STATS_GET_COUNTER( - sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_OK); - fc.receiver_ok = UCS_STATS_GET_COUNTER( - receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_OK); fc.sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); fc.receiver_throttled = UCS_STATS_GET_COUNTER( @@ -2424,21 +2416,6 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { } }; -UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_under_cap, - "RNDV_MTYPE_WORKER_MAX_MEM=1g", "RNDV_FRAG_MEM_TYPE=cuda") -{ - fc_counters fc; - - run_fc_test(8, fc); - - EXPECT_EQ(0u, fc.sender_throttled) << "sender should not be throttled"; - EXPECT_EQ(0u, fc.receiver_throttled) << "receiver should not be throttled"; - EXPECT_GT(fc.sender_ok + fc.receiver_ok, 0u) - << "FC should be active and tracking fragments"; - - verify_clean_fc_state(); -} - UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, "RNDV_MTYPE_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") { @@ -2461,8 +2438,6 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, EXPECT_EQ(0u, fc.sender_throttled) << "FC disabled - no throttling expected"; EXPECT_EQ(0u, fc.receiver_throttled) << "FC disabled - no throttling expected"; - EXPECT_EQ(0u, fc.sender_ok) << "FC disabled - no increment expected"; - EXPECT_EQ(0u, fc.receiver_ok) << "FC disabled - no increment expected"; } From a757e006e2c603107b918822f79f8090d8af8300 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 17:57:24 +0200 Subject: [PATCH 19/67] UCP/RNDV: more fixes --- src/ucp/rndv/proto_rndv.h | 31 -------------------------- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_mtype.inl | 43 ++++++++++++++++++++++++++++++++++++- src/ucp/rndv/rndv_ppln.c | 1 - src/ucp/rndv/rndv_put.c | 2 +- src/ucp/rndv/rndv_rtr.c | 2 +- 6 files changed, 45 insertions(+), 36 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 203c152f07f..35ffdc57a6a 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -37,37 +37,6 @@ ucp_proto_rndv_rts_tag_op_flags(ucp_rndv_rts_opcode_t opcode) return (opcode == UCP_RNDV_RTS_TAG_OK) ? UCP_PROTO_SELECT_OP_FLAG_TAG_RNDV : 0; } - -/* - * Staging-buffer flow-control throttle limit per operation type. - * - * To prevent deadlock under memory pressure, each operation type is capped at a - * different share of the total fragment budget (fc_max_frags). Operations that - * free more resources on completion receive a larger cap, so they can always - * make progress and unblock others: - * - * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. - * GET (level 1, ~90%) – frees local staging buffer only. - * RTR (level 2, ~80%) – triggers a remote PUT allocation, adding pressure. - * - * The step between tiers is 10% of the budget (1/10). The exact fractions are - * not performance-sensitive; only the strict ordering PUT > GET > RTR matters. - * The same ordering is used when dequeueing pending requests, see - * ucp_proto_rndv_mtype_fc_decrement(). - */ -#define UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, _level) \ - ((_fc_max) - (_level) * ((_fc_max) / 10)) - -#define UCP_PROTO_RNDV_MTYPE_FC_PUT_LIMIT(_fc_max) \ - UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 0) - -#define UCP_PROTO_RNDV_MTYPE_FC_GET_LIMIT(_fc_max) \ - UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 1) - -#define UCP_PROTO_RNDV_MTYPE_FC_RTR_LIMIT(_fc_max) \ - UCP_PROTO_RNDV_MTYPE_FC_LIMIT(_fc_max, 2) - - /** * Rendezvous protocol which sends a control message to the remote peer, and not * actually transferring bulk data. The remote peer is expected to perform the diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 0652d26a66e..76adbe9f5b6 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -309,7 +309,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - max_frags = UCP_PROTO_RNDV_MTYPE_FC_GET_LIMIT(rpriv->fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_get_limit(rpriv->fc_max_frags); /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index b7e49a16f80..0aaf18b448a 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -177,6 +177,47 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) return 1; } +/* + * Staging-buffer flow-control throttle limit per operation type. + * + * To prevent deadlock under memory pressure, each operation type is capped at a + * different share of the total fragment budget (fc_max_frags). Operations that + * free more resources on completion receive a larger cap, so they can always + * make progress and unblock others: + * + * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. + * GET (level 1, ~90%) – frees local staging buffer only. + * RTR (level 2, ~80%) – triggers a remote PUT allocation, adding pressure. + * + * The step between tiers is 10% of the budget (1/10). The exact fractions are + * not performance-sensitive; only the strict ordering PUT > GET > RTR matters. + * The same ordering is used when dequeueing pending requests, see + * ucp_proto_rndv_mtype_fc_decrement(). + */ +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_rndv_mtype_fc_limit(size_t fc_max, unsigned level) +{ + return fc_max - level * (fc_max / 10); +} + +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_rndv_mtype_fc_put_limit(size_t fc_max) +{ + return ucp_proto_rndv_mtype_fc_limit(fc_max, 0); +} + +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_rndv_mtype_fc_get_limit(size_t fc_max) +{ + return ucp_proto_rndv_mtype_fc_limit(fc_max, 1); +} + +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max) +{ + return ucp_proto_rndv_mtype_fc_limit(fc_max, 2); +} + /** * Check if request should be throttled due to flow control limit. * If throttled, the request is queued to the appropriate priority queue. @@ -208,7 +249,7 @@ ucp_proto_rndv_mtype_fc_throttle(ucp_request_t *req, const size_t max_frags, /** * Decrement active_frags counter and reschedule pending request. * Dequeue priority: PUT > GET > RTR (same ordering as the throttle limits - * defined by UCP_PROTO_RNDV_MTYPE_FC_LIMIT). + * defined by ucp_proto_rndv_mtype_fc_limit()). * * Priority rationale: * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote diff --git a/src/ucp/rndv/rndv_ppln.c b/src/ucp/rndv/rndv_ppln.c index bf8b9f64ebb..9b5c7fc0a6c 100644 --- a/src/ucp/rndv/rndv_ppln.c +++ b/src/ucp/rndv/rndv_ppln.c @@ -13,7 +13,6 @@ #include #include #include -#include enum { diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 3299741f3cd..66920b91395 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -531,7 +531,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; - max_frags = UCP_PROTO_RNDV_MTYPE_FC_PUT_LIMIT(rpriv->bulk.fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags); pending_q = &worker->rndv_mtype_fc.put_pending_q; if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == UCS_ERR_NO_RESOURCE) { diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 9e72f68c918..b74817d33a4 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -439,7 +439,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { rpriv = req->send.proto_config->priv; - max_frags = UCP_PROTO_RNDV_MTYPE_FC_RTR_LIMIT(rpriv->fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags); /* Check throttling limit. If no resource at the moment, queue the * request in RTR pending queue and return UCS_OK. */ From 8c447d675f12457bd9723d62f3e5a6176ceb3656 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 18:01:07 +0200 Subject: [PATCH 20/67] UCP/RNDV: put throttle in request_init --- src/ucp/rndv/rndv_get.c | 8 +++---- src/ucp/rndv/rndv_mtype.inl | 43 ++++++++++++------------------------- src/ucp/rndv/rndv_put.c | 9 ++++---- src/ucp/rndv/rndv_rtr.c | 14 +++++------- 4 files changed, 28 insertions(+), 46 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 76adbe9f5b6..48f8485d6e6 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -314,13 +314,13 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ pending_q = &worker->rndv_mtype_fc.get_pending_q; - if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == - UCS_ERR_NO_RESOURCE) { + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, + rpriv->frag_sys_dev, + max_frags, pending_q); + if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } - status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, - rpriv->frag_sys_dev); if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 0aaf18b448a..87e8778d690 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -47,10 +47,23 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params, static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, - ucs_sys_device_t frag_sys_dev) + ucs_sys_device_t frag_sys_dev, + const size_t max_frags, + ucs_queue_head_t *pending_q) { ucp_worker_h worker = req->send.ep->worker; + /* Check throttling limit. If no resource at the moment, queue the + * request in RTR pending queue and return NO_RESOURCE. */ + if (worker->rndv_mtype_fc.active_frags >= max_frags) { + ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", + worker->rndv_mtype_fc.active_frags, max_frags); + UCS_STATS_UPDATE_COUNTER(worker->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); + return UCS_ERR_NO_RESOURCE; + } + req->send.rndv.mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev); if (req->send.rndv.mdesc == NULL) { @@ -218,34 +231,6 @@ ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max) return ucp_proto_rndv_mtype_fc_limit(fc_max, 2); } -/** - * Check if request should be throttled due to flow control limit. - * If throttled, the request is queued to the appropriate priority queue. - * - * @param req The request to check. - * @param max_frags The maximum number of fragments allowed. - * @param pending_q The queue to add the request to if it is throttled. - * - * @return UCS_OK if not throttled, UCS_ERR_NO_RESOURCE if throttled and queued. - */ -static UCS_F_ALWAYS_INLINE ucs_status_t -ucp_proto_rndv_mtype_fc_throttle(ucp_request_t *req, const size_t max_frags, - ucs_queue_head_t *pending_q) -{ - ucp_worker_h worker = req->send.ep->worker; - - if (worker->rndv_mtype_fc.active_frags >= max_frags) { - ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", - worker->rndv_mtype_fc.active_frags, max_frags); - UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); - ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); - return UCS_ERR_NO_RESOURCE; - } - - return UCS_OK; -} - /** * Decrement active_frags counter and reschedule pending request. * Dequeue priority: PUT > GET > RTR (same ordering as the throttle limits diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 66920b91395..5955412422c 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -533,12 +533,13 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags); pending_q = &worker->rndv_mtype_fc.put_pending_q; - if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == - UCS_ERR_NO_RESOURCE) { + status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, + rpriv->bulk.frag_sys_dev, + max_frags, pending_q); + if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } - status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, - rpriv->bulk.frag_sys_dev); + if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index b74817d33a4..4ab4b8c76d7 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -363,10 +363,9 @@ ucp_proto_rndv_rtr_mtype_complete(ucp_request_t *req, int abort) { if (!abort || (req->send.rndv.mdesc != NULL)) { ucs_mpool_put_inline(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_fc_decrement(req); } - ucp_proto_rndv_mtype_fc_decrement(req); - if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 0, abort); } else { @@ -440,17 +439,14 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) rpriv = req->send.proto_config->priv; max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags); - - /* Check throttling limit. If no resource at the moment, queue the - * request in RTR pending queue and return UCS_OK. */ pending_q = &worker->rndv_mtype_fc.rtr_pending_q; - if (ucp_proto_rndv_mtype_fc_throttle(req, max_frags, pending_q) == - UCS_ERR_NO_RESOURCE) { + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, + rpriv->frag_sys_dev, + max_frags, pending_q); + if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } - status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, - rpriv->frag_sys_dev); if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; From 8f5f4ca6208b98305cfb121b5e47d32049c72d7c Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 18:39:26 +0200 Subject: [PATCH 21/67] UCP/RNDV: put decrement in mdesc_release --- src/ucp/rndv/rndv_get.c | 7 ++----- src/ucp/rndv/rndv_mtype.inl | 12 ++++++++++++ src/ucp/rndv/rndv_put.c | 6 ++---- src/ucp/rndv/rndv_rtr.c | 6 ++---- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 48f8485d6e6..70b0c864b68 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -273,8 +273,7 @@ ucp_proto_rndv_get_mtype_unpack_completion(uct_completion_t *uct_comp) ucp_request_t *req = ucs_container_of(uct_comp, ucp_request_t, send.state.uct_comp); - ucs_mpool_put_inline(req->send.rndv.mdesc); - ucp_proto_rndv_mtype_fc_decrement(req); + ucp_proto_rndv_mtype_mdesc_release(req); if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 1, 0); @@ -388,12 +387,10 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) return UCS_OK; } - ucs_mpool_put_inline(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_mdesc_release(req); req->send.rndv.mdesc = NULL; req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; - ucp_proto_rndv_mtype_fc_decrement(req); - if ((req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_FETCH) && (req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_ATS)) { ucp_proto_fatal_invalid_stage(req, "reset"); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 87e8778d690..0bd25db5663 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -278,6 +278,18 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) pending_req); } +/** + * Release the staging buffer and decrement the FC active fragments counter. + * This pairs with ucp_proto_rndv_mtype_request_init() which allocates the + * mdesc and increments the counter. + */ +static UCS_F_ALWAYS_INLINE void +ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) +{ + ucs_mpool_put_inline(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_fc_decrement(req); +} + static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mdesc_mtype_copy(ucp_request_t *req, uct_ep_put_zcopy_func_t copy_func, diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 5955412422c..9e8477660a2 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -576,8 +576,7 @@ static void ucp_proto_rndv_put_mtype_completion(uct_completion_t *uct_comp) send.state.uct_comp); ucp_trace_req(req, "rndv_put_mtype_completion"); - ucs_mpool_put(req->send.rndv.mdesc); - ucp_proto_rndv_mtype_fc_decrement(req); + ucp_proto_rndv_mtype_mdesc_release(req); ucp_proto_rndv_put_common_complete(req); } @@ -587,8 +586,7 @@ static void ucp_proto_rndv_put_mtype_frag_completion(uct_completion_t *uct_comp) send.state.uct_comp); ucp_trace_req(req, "rndv_put_mtype_frag_completion"); - ucs_mpool_put(req->send.rndv.mdesc); - ucp_proto_rndv_mtype_fc_decrement(req); + ucp_proto_rndv_mtype_mdesc_release(req); ucp_proto_rndv_ppln_send_frag_complete(req, 1); } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 4ab4b8c76d7..476c7c98ffc 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -362,8 +362,7 @@ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_rtr_mtype_complete(ucp_request_t *req, int abort) { if (!abort || (req->send.rndv.mdesc != NULL)) { - ucs_mpool_put_inline(req->send.rndv.mdesc); - ucp_proto_rndv_mtype_fc_decrement(req); + ucp_proto_rndv_mtype_mdesc_release(req); } if (ucp_proto_rndv_request_is_ppln_frag(req)) { @@ -394,9 +393,8 @@ ucp_proto_rndv_rtr_mtype_abort(ucp_request_t *req, ucs_status_t status) static ucs_status_t ucp_proto_rndv_rtr_mtype_reset(ucp_request_t *req) { if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { - ucs_mpool_put_inline(req->send.rndv.mdesc); + ucp_proto_rndv_mtype_mdesc_release(req); req->send.rndv.mdesc = NULL; - ucp_proto_rndv_mtype_fc_decrement(req); } return ucp_proto_request_zcopy_id_reset(req); From 7627bcbc6ba3053eeb855ab80284a35a4eecdcf9 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 18:58:54 +0200 Subject: [PATCH 22/67] UCP/RNDV: some fixes (enable by default, remove enable field) --- src/ucp/core/ucp_context.c | 2 +- src/ucp/core/ucp_worker.c | 10 +++------- src/ucp/rndv/rndv_mtype.inl | 4 ---- 3 files changed, 4 insertions(+), 12 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 916f07bfd17..6529e9d3d67 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,7 +404,7 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_MAX_MEM", "inf", + {"RNDV_MTYPE_WORKER_MAX_MEM", "8g", "Maximum memory for concurrent mtype fragments per worker.\n" "This value is translated to a fragment count based on RNDV_FRAG_SIZE\n" "for each memory type. Set to \"inf\" to disable flow control.", diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index dd4bfce4b42..d86971911c1 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2641,13 +2641,9 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_failures = 0; /* Initialize RNDV mtype flow control if configured */ - worker->rndv_mtype_fc.enabled = ucp_context_rndv_mtype_fc_enabled(context); - if (worker->rndv_mtype_fc.enabled) { - worker->rndv_mtype_fc.active_frags = 0; - ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); - } + ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 0bd25db5663..e3c5b5dfd95 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -254,10 +254,6 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) ucs_assert(worker->rndv_mtype_fc.active_frags > 0); worker->rndv_mtype_fc.active_frags--; - if (!worker->rndv_mtype_fc.enabled) { - return; - } - /* Dequeue with priority: PUT > GET > RTR */ if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) { elem = ucs_queue_pull(&worker->rndv_mtype_fc.put_pending_q); From 415a289e4b9f9240fbd8e2f3dcdf4fe7c877f661 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 16 Mar 2026 18:59:23 +0200 Subject: [PATCH 23/67] UCP/RNDV: make ratios configurable --- src/ucp/core/ucp_context.c | 8 ++++++++ src/ucp/core/ucp_context.h | 2 ++ src/ucp/core/ucp_worker.c | 1 + src/ucp/core/ucp_worker.h | 1 + src/ucp/rndv/rndv_get.c | 3 ++- src/ucp/rndv/rndv_mtype.inl | 28 +++++++++++++++------------- src/ucp/rndv/rndv_put.c | 3 ++- src/ucp/rndv/rndv_rtr.c | 3 ++- 8 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 6529e9d3d67..ee9d1bdc9c0 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -410,6 +410,14 @@ static ucs_config_field_t ucp_context_config_table[] = { "for each memory type. Set to \"inf\" to disable flow control.", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, + {"RNDV_MTYPE_FC_TIER_STEP", "10", + "Percentage step between flow-control throttle tiers for mtype fragments.\n" + "Each operation type (PUT, GET, RTR) is capped at a different share of the\n" + "total fragment budget. This value controls the gap between tiers.\n" + "For example, 10 means PUT=100%, GET=90%, RTR=80% of the budget.\n" + "Valid values are 1-49.", + ucs_offsetof(ucp_context_config_t, rndv_mtype_fc_tier_step), UCS_CONFIG_TYPE_UINT}, + {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" "the flush operation in a bounded time even if there are new requests on\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 2ddabd30755..d7df3f47d67 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -134,6 +134,8 @@ typedef struct ucp_context_config { int rma_ppln_enable; /** Maximum memory for concurrent rndv mtype fragments per worker (bytes) */ size_t rndv_mtype_worker_max_mem; + /** Percentage step between FC throttle tiers (PUT/GET/RTR) */ + unsigned rndv_mtype_fc_tier_step; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index d86971911c1..429ea7cb637 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2641,6 +2641,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_failures = 0; /* Initialize RNDV mtype flow control if configured */ + worker->rndv_mtype_fc.tier_step = context->config.ext.rndv_mtype_fc_tier_step; ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 8b28abb004c..de28fbfc0b0 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -400,6 +400,7 @@ typedef struct ucp_worker { struct { /* Worker-level mtype fragment flow control */ int enabled; /* Whether FC is enabled (cached) */ + unsigned tier_step; /* Pct step between throttle tiers (cached) */ size_t active_frags; /* Current active fragments */ /* Separate pending queues by priority (PUT > GET > RTR) */ ucs_queue_head_t put_pending_q; /* Throttled PUT (RNDV_SEND) requests */ diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 70b0c864b68..10a632aca6b 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -308,7 +308,8 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - max_frags = ucp_proto_rndv_mtype_fc_get_limit(rpriv->fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_get_limit(rpriv->fc_max_frags, + worker->rndv_mtype_fc.tier_step); /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index e3c5b5dfd95..97145aed6c0 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -198,37 +198,39 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) * free more resources on completion receive a larger cap, so they can always * make progress and unblock others: * - * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. - * GET (level 1, ~90%) – frees local staging buffer only. - * RTR (level 2, ~80%) – triggers a remote PUT allocation, adding pressure. + * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. + * GET (level 1, 100%-tier_step) – frees local staging buffer only. + * RTR (level 2, 100%-2*tier_step) – triggers a remote PUT allocation, adding pressure. * - * The step between tiers is 10% of the budget (1/10). The exact fractions are - * not performance-sensitive; only the strict ordering PUT > GET > RTR matters. + * The step between tiers is configurable via UCX_RNDV_MTYPE_FC_TIER_STEP + * (default 10%). The exact fractions are not performance-sensitive; only the + * strict ordering PUT > GET > RTR matters. * The same ordering is used when dequeueing pending requests, see * ucp_proto_rndv_mtype_fc_decrement(). */ static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_limit(size_t fc_max, unsigned level) +ucp_proto_rndv_mtype_fc_limit(size_t fc_max, unsigned level, + unsigned tier_step) { - return fc_max - level * (fc_max / 10); + return fc_max - level * (fc_max * tier_step / 100); } static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_put_limit(size_t fc_max) +ucp_proto_rndv_mtype_fc_put_limit(size_t fc_max, unsigned tier_step) { - return ucp_proto_rndv_mtype_fc_limit(fc_max, 0); + return ucp_proto_rndv_mtype_fc_limit(fc_max, 0, tier_step); } static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_get_limit(size_t fc_max) +ucp_proto_rndv_mtype_fc_get_limit(size_t fc_max, unsigned tier_step) { - return ucp_proto_rndv_mtype_fc_limit(fc_max, 1); + return ucp_proto_rndv_mtype_fc_limit(fc_max, 1, tier_step); } static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max) +ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max, unsigned tier_step) { - return ucp_proto_rndv_mtype_fc_limit(fc_max, 2); + return ucp_proto_rndv_mtype_fc_limit(fc_max, 2, tier_step); } /** diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 9e8477660a2..a50ce28a6df 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -531,7 +531,8 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) rpriv = req->send.proto_config->priv; - max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags, + worker->rndv_mtype_fc.tier_step); pending_q = &worker->rndv_mtype_fc.put_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev, diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 476c7c98ffc..6e366bc52b1 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -436,7 +436,8 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { rpriv = req->send.proto_config->priv; - max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags); + max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags, + worker->rndv_mtype_fc.tier_step); pending_q = &worker->rndv_mtype_fc.rtr_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev, From a84c68466eebc244892ba1a0768a7283e7a0e2d2 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 17 Mar 2026 13:29:32 +0200 Subject: [PATCH 24/67] UCP/RNDV: two pending queues instead of 3 - high prio and low prio --- src/ucp/core/ucp_worker.c | 5 ++--- src/ucp/core/ucp_worker.h | 7 +++---- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_mtype.inl | 12 +++++------- src/ucp/rndv/rndv_put.c | 2 +- src/ucp/rndv/rndv_rtr.c | 2 +- test/gtest/ucp/test_ucp_am.cc | 10 ++++------ 7 files changed, 17 insertions(+), 23 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 429ea7cb637..8686c88bbca 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2642,9 +2642,8 @@ ucs_status_t ucp_worker_create(ucp_context_h context, /* Initialize RNDV mtype flow control if configured */ worker->rndv_mtype_fc.tier_step = context->config.ext.rndv_mtype_fc_tier_step; - ucs_queue_head_init(&worker->rndv_mtype_fc.put_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.get_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.rtr_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.hi_pending_q); + ucs_queue_head_init(&worker->rndv_mtype_fc.lo_pending_q); /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index de28fbfc0b0..6a8897167e7 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -402,10 +402,9 @@ typedef struct ucp_worker { int enabled; /* Whether FC is enabled (cached) */ unsigned tier_step; /* Pct step between throttle tiers (cached) */ size_t active_frags; /* Current active fragments */ - /* Separate pending queues by priority (PUT > GET > RTR) */ - ucs_queue_head_t put_pending_q; /* Throttled PUT (RNDV_SEND) requests */ - ucs_queue_head_t get_pending_q; /* Throttled GET requests */ - ucs_queue_head_t rtr_pending_q; /* Throttled RTR requests */ + /* Separate pending queues by priority (GET/PUT > RTR) */ + ucs_queue_head_t hi_pending_q; /* Throttled GET/PUT requests (high prio) */ + ucs_queue_head_t lo_pending_q; /* Throttled RTR requests (low prio) */ } rndv_mtype_fc; struct { diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 10a632aca6b..b6393f5652e 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -313,7 +313,7 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) /* Check throttling limit. If no resource at the moment, queue the * request in GET pending queue and return UCS_OK. */ - pending_q = &worker->rndv_mtype_fc.get_pending_q; + pending_q = &worker->rndv_mtype_fc.hi_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev, max_frags, pending_q); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 97145aed6c0..8ad54f31f87 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -256,13 +256,11 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) ucs_assert(worker->rndv_mtype_fc.active_frags > 0); worker->rndv_mtype_fc.active_frags--; - /* Dequeue with priority: PUT > GET > RTR */ - if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) { - elem = ucs_queue_pull(&worker->rndv_mtype_fc.put_pending_q); - } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.get_pending_q)) { - elem = ucs_queue_pull(&worker->rndv_mtype_fc.get_pending_q); - } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.rtr_pending_q)) { - elem = ucs_queue_pull(&worker->rndv_mtype_fc.rtr_pending_q); + /* Dequeue with priority: GET/PUT (hi) > RTR (lo) */ + if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.hi_pending_q)) { + elem = ucs_queue_pull(&worker->rndv_mtype_fc.hi_pending_q); + } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.lo_pending_q)) { + elem = ucs_queue_pull(&worker->rndv_mtype_fc.lo_pending_q); } if (elem == NULL) { diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index a50ce28a6df..32e102e3bf4 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -533,7 +533,7 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags, worker->rndv_mtype_fc.tier_step); - pending_q = &worker->rndv_mtype_fc.put_pending_q; + pending_q = &worker->rndv_mtype_fc.hi_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev, max_frags, pending_q); diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 6e366bc52b1..fd6c0e9e0bf 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -438,7 +438,7 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags, worker->rndv_mtype_fc.tier_step); - pending_q = &worker->rndv_mtype_fc.rtr_pending_q; + pending_q = &worker->rndv_mtype_fc.lo_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev, max_frags, pending_q); diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 7f92bd03cd4..7d05b46df74 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2374,12 +2374,10 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.put_pending_q)) - << "put_pending_q should be empty"; - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.get_pending_q)) - << "get_pending_q should be empty"; - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.rtr_pending_q)) - << "rtr_pending_q should be empty"; + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.hi_pending_q)) + << "hi_pending_q should be empty"; + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.lo_pending_q)) + << "lo_pending_q should be empty"; } void send_message(size_t num_frags) From aed4dc6fb7b99bae2bc1d0b0fc71b262371ea45b Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 17 Mar 2026 14:27:28 +0200 Subject: [PATCH 25/67] UCP/RNDV: queue indexing --- src/ucp/core/ucp_worker.c | 7 +++++-- src/ucp/core/ucp_worker.h | 22 +++++++++++++++++++--- src/ucp/rndv/rndv_get.c | 7 +++---- src/ucp/rndv/rndv_mtype.inl | 31 ++++++++++++++++++++----------- src/ucp/rndv/rndv_put.c | 5 ++--- src/ucp/rndv/rndv_rtr.c | 5 ++--- test/gtest/ucp/test_ucp_am.cc | 8 ++++---- 7 files changed, 55 insertions(+), 30 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 8686c88bbca..530f06852fb 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2610,6 +2610,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context, unsigned name_length; ucp_worker_h worker; ucs_status_t status; + int q_index; worker = ucs_calloc(1, sizeof(*worker), "ucp worker"); if (worker == NULL) { @@ -2642,8 +2643,10 @@ ucs_status_t ucp_worker_create(ucp_context_h context, /* Initialize RNDV mtype flow control if configured */ worker->rndv_mtype_fc.tier_step = context->config.ext.rndv_mtype_fc_tier_step; - ucs_queue_head_init(&worker->rndv_mtype_fc.hi_pending_q); - ucs_queue_head_init(&worker->rndv_mtype_fc.lo_pending_q); + worker->rndv_mtype_fc.best_q = UCP_WORKER_RNDV_FC_OP_LAST; + for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { + ucs_queue_head_init(&worker->rndv_mtype_fc.pending_q[q_index]); + } /* Copy user flags, and mask-out unsupported flags for compatibility */ worker->flags = UCP_PARAM_VALUE(WORKER, params, flags, FLAGS, 0) & diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 6a8897167e7..7c137a4c9b6 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -297,6 +297,19 @@ UCS_PTR_MAP_TYPE(request, 0); UCS_ARRAY_DECLARE_TYPE(ucp_rkey_config_arr_t, unsigned, ucp_rkey_config_t); +/** + * Rendezvous mtype flow-control operation types. + * Used as indices into the per-worker pending queue array. + * PUT and GET share the high-priority queue; RTR uses the low-priority queue. + */ +enum { + UCP_WORKER_RNDV_FC_OP_PUT = 0, /* GET/PUT (hi prio) */ + UCP_WORKER_RNDV_FC_OP_GET = UCP_WORKER_RNDV_FC_OP_PUT, /* Same as PUT */ + UCP_WORKER_RNDV_FC_OP_RTR = 1, /* RTR (lo prio) */ + UCP_WORKER_RNDV_FC_OP_LAST = 2 +}; + + /** * UCP worker (thread context). */ @@ -402,9 +415,12 @@ typedef struct ucp_worker { int enabled; /* Whether FC is enabled (cached) */ unsigned tier_step; /* Pct step between throttle tiers (cached) */ size_t active_frags; /* Current active fragments */ - /* Separate pending queues by priority (GET/PUT > RTR) */ - ucs_queue_head_t hi_pending_q; /* Throttled GET/PUT requests (high prio) */ - ucs_queue_head_t lo_pending_q; /* Throttled RTR requests (low prio) */ + /* Pending queues indexed by ucp_worker_rndv_fc_op_t, ordered by + * descending priority (GET/PUT = 0, RTR = 1) */ + ucs_queue_head_t pending_q[UCP_WORKER_RNDV_FC_OP_LAST]; + /* Index of highest-priority non-empty queue + * (UCP_WORKER_RNDV_FC_OP_LAST if all queues are empty) */ + int best_q; } rndv_mtype_fc; struct { diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index b6393f5652e..61138660afb 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -302,7 +302,6 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; size_t max_frags; - ucs_queue_head_t *pending_q; /* coverity[tainted_data_downcast] */ rpriv = req->send.proto_config->priv; @@ -312,11 +311,11 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) worker->rndv_mtype_fc.tier_step); /* Check throttling limit. If no resource at the moment, queue the - * request in GET pending queue and return UCS_OK. */ - pending_q = &worker->rndv_mtype_fc.hi_pending_q; + * request in GET/PUT pending queue and return UCS_OK. */ status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev, - max_frags, pending_q); + max_frags, + UCP_WORKER_RNDV_FC_OP_GET); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 8ad54f31f87..cf794a54fa9 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -49,18 +49,23 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, ucs_sys_device_t frag_sys_dev, const size_t max_frags, - ucs_queue_head_t *pending_q) + int fc_op) { ucp_worker_h worker = req->send.ep->worker; /* Check throttling limit. If no resource at the moment, queue the - * request in RTR pending queue and return NO_RESOURCE. */ + * request in the appropriate pending queue and return NO_RESOURCE. */ if (worker->rndv_mtype_fc.active_frags >= max_frags) { ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", worker->rndv_mtype_fc.active_frags, max_frags); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); - ucs_queue_push(pending_q, &req->send.rndv.ppln.queue_elem); + ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], + &req->send.rndv.ppln.queue_elem); + if (fc_op < worker->rndv_mtype_fc.best_q) { + worker->rndv_mtype_fc.best_q = fc_op; + } + return UCS_ERR_NO_RESOURCE; } @@ -250,21 +255,25 @@ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; - ucs_queue_elem_t *elem = NULL; + ucs_queue_elem_t *elem; ucp_request_t *pending_req; ucs_assert(worker->rndv_mtype_fc.active_frags > 0); worker->rndv_mtype_fc.active_frags--; - /* Dequeue with priority: GET/PUT (hi) > RTR (lo) */ - if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.hi_pending_q)) { - elem = ucs_queue_pull(&worker->rndv_mtype_fc.hi_pending_q); - } else if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.lo_pending_q)) { - elem = ucs_queue_pull(&worker->rndv_mtype_fc.lo_pending_q); + /* Dequeue from highest-priority non-empty queue */ + if (worker->rndv_mtype_fc.best_q >= UCP_WORKER_RNDV_FC_OP_LAST) { + return; } - if (elem == NULL) { - return; + elem = ucs_queue_pull( + &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q]); + + /* Advance best_q past empty queues */ + while ((worker->rndv_mtype_fc.best_q < UCP_WORKER_RNDV_FC_OP_LAST) && + ucs_queue_is_empty( + &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q])) { + worker->rndv_mtype_fc.best_q++; } pending_req = ucs_container_of(elem, ucp_request_t, diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 32e102e3bf4..46b8d9088da 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -525,7 +525,6 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; size_t max_frags; - ucs_queue_head_t *pending_q; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); @@ -533,10 +532,10 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags, worker->rndv_mtype_fc.tier_step); - pending_q = &worker->rndv_mtype_fc.hi_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev, - max_frags, pending_q); + max_frags, + UCP_WORKER_RNDV_FC_OP_PUT); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index fd6c0e9e0bf..5f7e1f069bb 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -430,7 +430,6 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; size_t max_frags; - ucs_queue_head_t *pending_q; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { @@ -438,10 +437,10 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags, worker->rndv_mtype_fc.tier_step); - pending_q = &worker->rndv_mtype_fc.lo_pending_q; status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, rpriv->frag_sys_dev, - max_frags, pending_q); + max_frags, + UCP_WORKER_RNDV_FC_OP_RTR); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 7d05b46df74..b302bbdebfe 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2374,10 +2374,10 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.hi_pending_q)) - << "hi_pending_q should be empty"; - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.lo_pending_q)) - << "lo_pending_q should be empty"; + for (int i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; i++) { + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) + << "pending_q[" << i << "] should be empty"; + } } void send_message(size_t num_frags) From ea99ea199eb5083eb3533fc715c8f1ad689ff7a6 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 17 Mar 2026 15:03:00 +0200 Subject: [PATCH 26/67] UCP/RNDV: limit mem using mpool::max_elems - removing differentiation between ops --- src/ucp/core/ucp_context.c | 8 ---- src/ucp/core/ucp_context.h | 2 - src/ucp/core/ucp_worker.c | 5 +-- src/ucp/core/ucp_worker.h | 13 +++++- src/ucp/rndv/proto_rndv.c | 24 ++++++----- src/ucp/rndv/proto_rndv.h | 21 +++++----- src/ucp/rndv/rndv.c | 2 + src/ucp/rndv/rndv_get.c | 17 +++----- src/ucp/rndv/rndv_mtype.inl | 79 ++++++----------------------------- src/ucp/rndv/rndv_put.c | 8 +--- src/ucp/rndv/rndv_rtr.c | 18 +++----- test/gtest/ucp/test_ucp_am.cc | 7 ---- 12 files changed, 64 insertions(+), 140 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index ee9d1bdc9c0..6529e9d3d67 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -410,14 +410,6 @@ static ucs_config_field_t ucp_context_config_table[] = { "for each memory type. Set to \"inf\" to disable flow control.", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, - {"RNDV_MTYPE_FC_TIER_STEP", "10", - "Percentage step between flow-control throttle tiers for mtype fragments.\n" - "Each operation type (PUT, GET, RTR) is capped at a different share of the\n" - "total fragment budget. This value controls the gap between tiers.\n" - "For example, 10 means PUT=100%, GET=90%, RTR=80% of the budget.\n" - "Valid values are 1-49.", - ucs_offsetof(ucp_context_config_t, rndv_mtype_fc_tier_step), UCS_CONFIG_TYPE_UINT}, - {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" "the flush operation in a bounded time even if there are new requests on\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index d7df3f47d67..2ddabd30755 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -134,8 +134,6 @@ typedef struct ucp_context_config { int rma_ppln_enable; /** Maximum memory for concurrent rndv mtype fragments per worker (bytes) */ size_t rndv_mtype_worker_max_mem; - /** Percentage step between FC throttle tiers (PUT/GET/RTR) */ - unsigned rndv_mtype_fc_tier_step; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 530f06852fb..b35354dfe09 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2641,9 +2641,8 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_closures = 0; worker->counters.ep_failures = 0; - /* Initialize RNDV mtype flow control if configured */ - worker->rndv_mtype_fc.tier_step = context->config.ext.rndv_mtype_fc_tier_step; - worker->rndv_mtype_fc.best_q = UCP_WORKER_RNDV_FC_OP_LAST; + /* Initialize RNDV mtype flow control pending queues */ + worker->rndv_mtype_fc.best_q = UCP_WORKER_RNDV_FC_OP_LAST; for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { ucs_queue_head_init(&worker->rndv_mtype_fc.pending_q[q_index]); } diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 7c137a4c9b6..02406ccd9b9 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -411,8 +411,17 @@ typedef struct ucp_worker { } counters; struct { - /* Worker-level mtype fragment flow control */ - int enabled; /* Whether FC is enabled (cached) */ + /* Worker-level mtype fragment flow control. + * + * Two levels of limiting cooperate: + * 1. Soft (tiered) limit – the active_frags counter tracks how many + * fragments are in-flight. Each operation type (PUT/GET/RTR) is + * capped at a different fraction of the budget so that higher- + * priority operations always have headroom, preventing deadlock. + * 2. Hard limit – the per-mpool quota (ucs_mpool_data::quota, set via + * max_elems when the mpool is created) caps total memory. If the + * mpool is exhausted the request is throttled the same way. + */ unsigned tier_step; /* Pct step between throttle tiers (cached) */ size_t active_frags; /* Current active fragments */ /* Pending queues indexed by ucp_worker_rndv_fc_op_t, ordered by diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index b7395c104c9..d4feeb5afbd 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -681,18 +681,22 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) } /** - * Compute maximum number of fragments allowed based on configured max memory - * and fragment size for the given memory type. The result is rounded down to - * the allocation chunk granularity (rndv_num_frags). + * Compute maximum number of elements allowed for the rndv fragment mpool + * based on configured max memory and fragment size for the given memory type. + * The result is rounded down to the allocation chunk granularity + * (rndv_num_frags) and returned as an unsigned value suitable for + * ucs_mpool_params_t::max_elems. + * + * When flow control is disabled (UCX_RNDV_MTYPE_WORKER_MAX_MEM="inf"), + * returns UINT_MAX so that the mpool is effectively unlimited. * * @param context The UCP context. * @param frag_mem_type Memory type used for fragments. * - * @return Maximum number of fragments that fit within the configured memory - * limit, aligned to allocation chunk size. + * @return Maximum number of mpool elements, aligned to allocation chunk size. */ -size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, - ucs_memory_type_t frag_mem_type) +unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, + ucs_memory_type_t frag_mem_type) { size_t max_mem; size_t frag_size; @@ -700,7 +704,7 @@ size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, size_t max_frags; if (!ucp_context_rndv_mtype_fc_enabled(context)) { - return SIZE_MAX; + return UINT_MAX; } max_mem = context->config.ext.rndv_mtype_worker_max_mem; @@ -720,7 +724,7 @@ size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, return frags_in_chunk; } - return max_frags; + return ucs_min(max_frags, UINT_MAX); } ucs_status_t @@ -776,8 +780,6 @@ ucp_proto_rndv_bulk_init(const ucp_proto_multi_init_params_t *init_params, rpriv->frag_mem_type = init_params->super.reg_mem_info.type; rpriv->frag_sys_dev = init_params->super.reg_mem_info.sys_dev; - rpriv->fc_max_frags = ucp_proto_rndv_mtype_fc_max_frags( - context, rpriv->frag_mem_type); if (rpriv->super.lane == UCP_NULL_LANE) { /* Add perf without ACK in case of pipeline */ diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 35ffdc57a6a..7e98f9e709c 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -91,9 +91,6 @@ typedef struct { ucs_memory_type_t frag_mem_type; ucs_sys_device_t frag_sys_dev; - /* max fragments for flow control */ - size_t fc_max_frags; - /* Multi-lane common part. Must be the last field, see @ref ucp_proto_multi_priv_t */ ucp_proto_multi_priv_t mpriv; @@ -195,18 +192,22 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); /** - * Compute maximum number of fragments allowed based on configured max memory - * and fragment size for the given memory type. The result is rounded down to - * the allocation chunk granularity (rndv_num_frags). + * Compute maximum number of elements allowed for the rndv fragment mpool + * based on configured max memory and fragment size for the given memory type. + * The result is rounded down to the allocation chunk granularity + * (rndv_num_frags) and returned as an unsigned value suitable for + * ucs_mpool_params_t::max_elems. + * + * When flow control is disabled (UCX_RNDV_MTYPE_WORKER_MAX_MEM="inf"), + * returns UINT_MAX so that the mpool is effectively unlimited. * * @param context The UCP context. * @param frag_mem_type Memory type used for fragments. * - * @return Maximum number of fragments that fit within the configured memory - * limit, aligned to allocation chunk size. + * @return Maximum number of mpool elements, aligned to allocation chunk size. */ -size_t ucp_proto_rndv_mtype_fc_max_frags(ucp_context_h context, - ucs_memory_type_t frag_mem_type); +unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, + ucs_memory_type_t frag_mem_type); ucs_status_t diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 31042738a40..2250fe33595 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1173,6 +1173,8 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.elem_size = sizeof(ucp_mem_desc_t); mp_params.alignment = 1; mp_params.elems_per_chunk = num_frags; + mp_params.max_elems = ucp_proto_rndv_mtype_fc_max_elems( + worker->context, key.mem_type); mp_params.ops = &ucp_frag_mpool_ops; mp_params.name = "ucp_rndv_frags"; status = ucs_mpool_init(&mp_params, mpool); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 61138660afb..e076c0f4b1a 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -298,24 +298,18 @@ static ucs_status_t ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) { ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); - ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; - size_t max_frags; /* coverity[tainted_data_downcast] */ rpriv = req->send.proto_config->priv; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { - max_frags = ucp_proto_rndv_mtype_fc_get_limit(rpriv->fc_max_frags, - worker->rndv_mtype_fc.tier_step); - - /* Check throttling limit. If no resource at the moment, queue the - * request in GET/PUT pending queue and return UCS_OK. */ - status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, - rpriv->frag_sys_dev, - max_frags, - UCP_WORKER_RNDV_FC_OP_GET); + /* Try to allocate a staging buffer. If the mpool quota is exhausted + * the request is queued and will be rescheduled later. */ + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, + rpriv->frag_sys_dev, + UCP_WORKER_RNDV_FC_OP_GET); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } @@ -325,7 +319,6 @@ ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) return UCS_OK; } - worker->rndv_mtype_fc.active_frags++; ucp_proto_rndv_get_common_request_init(req); ucp_proto_completion_init(&req->send.state.uct_comp, ucp_proto_rndv_get_mtype_fetch_completion); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index cf794a54fa9..0823baca099 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -48,16 +48,16 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, ucs_sys_device_t frag_sys_dev, - const size_t max_frags, int fc_op) { ucp_worker_h worker = req->send.ep->worker; - /* Check throttling limit. If no resource at the moment, queue the - * request in the appropriate pending queue and return NO_RESOURCE. */ - if (worker->rndv_mtype_fc.active_frags >= max_frags) { - ucs_trace_req("mtype_fc: fragments throttle limit reached (%zu/%zu)", - worker->rndv_mtype_fc.active_frags, max_frags); + req->send.rndv.mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, + frag_sys_dev); + if (req->send.rndv.mdesc == NULL) { + /* Mpool quota exhausted - throttle by queuing the request in the + * appropriate pending queue ordered by priority. */ + ucs_trace_req("mtype_fc: frag mpool exhausted, throttling request"); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], @@ -69,12 +69,6 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, return UCS_ERR_NO_RESOURCE; } - req->send.rndv.mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, - frag_sys_dev); - if (req->send.rndv.mdesc == NULL) { - return UCS_ERR_NO_MEMORY; - } - return UCS_OK; } @@ -195,53 +189,9 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) return 1; } -/* - * Staging-buffer flow-control throttle limit per operation type. - * - * To prevent deadlock under memory pressure, each operation type is capped at a - * different share of the total fragment budget (fc_max_frags). Operations that - * free more resources on completion receive a larger cap, so they can always - * make progress and unblock others: - * - * PUT (level 0, 100%) – frees local staging buffer AND remote RTR buffer. - * GET (level 1, 100%-tier_step) – frees local staging buffer only. - * RTR (level 2, 100%-2*tier_step) – triggers a remote PUT allocation, adding pressure. - * - * The step between tiers is configurable via UCX_RNDV_MTYPE_FC_TIER_STEP - * (default 10%). The exact fractions are not performance-sensitive; only the - * strict ordering PUT > GET > RTR matters. - * The same ordering is used when dequeueing pending requests, see - * ucp_proto_rndv_mtype_fc_decrement(). - */ -static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_limit(size_t fc_max, unsigned level, - unsigned tier_step) -{ - return fc_max - level * (fc_max * tier_step / 100); -} - -static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_put_limit(size_t fc_max, unsigned tier_step) -{ - return ucp_proto_rndv_mtype_fc_limit(fc_max, 0, tier_step); -} - -static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_get_limit(size_t fc_max, unsigned tier_step) -{ - return ucp_proto_rndv_mtype_fc_limit(fc_max, 1, tier_step); -} - -static UCS_F_ALWAYS_INLINE size_t -ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max, unsigned tier_step) -{ - return ucp_proto_rndv_mtype_fc_limit(fc_max, 2, tier_step); -} - /** - * Decrement active_frags counter and reschedule pending request. - * Dequeue priority: PUT > GET > RTR (same ordering as the throttle limits - * defined by ucp_proto_rndv_mtype_fc_limit()). + * Reschedule a pending throttled request after a fragment is released back to + * the mpool. Dequeue priority: PUT > GET > RTR. * * Priority rationale: * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote @@ -252,15 +202,12 @@ ucp_proto_rndv_mtype_fc_rtr_limit(size_t fc_max, unsigned tier_step) * memory pressure. */ static UCS_F_ALWAYS_INLINE void -ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) +ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; ucs_queue_elem_t *elem; ucp_request_t *pending_req; - ucs_assert(worker->rndv_mtype_fc.active_frags > 0); - worker->rndv_mtype_fc.active_frags--; - /* Dequeue from highest-priority non-empty queue */ if (worker->rndv_mtype_fc.best_q >= UCP_WORKER_RNDV_FC_OP_LAST) { return; @@ -284,15 +231,15 @@ ucp_proto_rndv_mtype_fc_decrement(ucp_request_t *req) } /** - * Release the staging buffer and decrement the FC active fragments counter. - * This pairs with ucp_proto_rndv_mtype_request_init() which allocates the - * mdesc and increments the counter. + * Release the staging buffer back to the mpool and reschedule any pending + * throttled request. This pairs with ucp_proto_rndv_mtype_request_init() + * which allocates the mdesc from the mpool. */ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) { ucs_mpool_put_inline(req->send.rndv.mdesc); - ucp_proto_rndv_mtype_fc_decrement(req); + ucp_proto_rndv_mtype_fc_reschedule_pending(req); } static UCS_F_ALWAYS_INLINE ucs_status_t diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 46b8d9088da..a6981c85237 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -521,20 +521,17 @@ static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct);\ - ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; - size_t max_frags; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); rpriv = req->send.proto_config->priv; - max_frags = ucp_proto_rndv_mtype_fc_put_limit(rpriv->bulk.fc_max_frags, - worker->rndv_mtype_fc.tier_step); + /* Try to allocate a staging buffer. If the mpool quota is exhausted + * the request is queued and will be rescheduled later. */ status = ucp_proto_rndv_mtype_request_init(req, rpriv->bulk.frag_mem_type, rpriv->bulk.frag_sys_dev, - max_frags, UCP_WORKER_RNDV_FC_OP_PUT); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; @@ -545,7 +542,6 @@ ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) return UCS_OK; } - worker->rndv_mtype_fc.active_frags++; ucp_proto_rndv_put_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; ucp_proto_rndv_mdesc_mtype_copy(req, uct_ep_get_zcopy, diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 5f7e1f069bb..4045cc2b2bd 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -39,8 +39,6 @@ typedef struct { ucp_proto_rndv_rtr_priv_t super; ucs_memory_type_t frag_mem_type; ucs_sys_device_t frag_sys_dev; - /* max fragments for flow control */ - size_t fc_max_frags; } ucp_proto_rndv_rtr_mtype_priv_t; @@ -427,20 +425,17 @@ ucp_proto_rndv_rtr_mtype_data_received(ucp_request_t *req, int in_buffer) static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); - ucp_worker_h worker = req->send.ep->worker; const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; - size_t max_frags; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { rpriv = req->send.proto_config->priv; - max_frags = ucp_proto_rndv_mtype_fc_rtr_limit(rpriv->fc_max_frags, - worker->rndv_mtype_fc.tier_step); - status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, - rpriv->frag_sys_dev, - max_frags, - UCP_WORKER_RNDV_FC_OP_RTR); + /* Try to allocate a staging buffer. If the mpool quota is exhausted + * the request is queued and will be rescheduled later. */ + status = ucp_proto_rndv_mtype_request_init(req, rpriv->frag_mem_type, + rpriv->frag_sys_dev, + UCP_WORKER_RNDV_FC_OP_RTR); if (status == UCS_ERR_NO_RESOURCE) { return UCS_OK; } @@ -450,7 +445,6 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) return UCS_OK; } - worker->rndv_mtype_fc.active_frags++; ucp_proto_rtr_common_request_init(req); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; } @@ -547,8 +541,6 @@ ucp_proto_rndv_rtr_mtype_probe(const ucp_proto_init_params_t *init_params) rpriv.frag_mem_type = frag_mem_type; rpriv.frag_sys_dev = params.super.reg_mem_info.sys_dev; params.flags = ucp_proto_rndv_ctrl_init_flags(¶ms); - rpriv.fc_max_frags = ucp_proto_rndv_mtype_fc_max_frags( - context, frag_mem_type); ucp_proto_rndv_ctrl_probe(¶ms, &rpriv, sizeof(rpriv)); out_unpack_perf_destroy: diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index b302bbdebfe..0249106c187 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2365,12 +2365,6 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { << " but got " << value; } - void check_active_frags_zero(const entity &e) - { - EXPECT_EQ(0u, e.worker()->rndv_mtype_fc.active_frags) - << "active_frags should be 0 after completion"; - } - void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); @@ -2390,7 +2384,6 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { { for (auto *ep : { &sender(), &receiver() }) { - check_active_frags_zero(*ep); check_pending_queues_empty(*ep); } } From 5aacd1b123dbac85506cb8a74148728d1280a016 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 17 Mar 2026 16:55:54 +0200 Subject: [PATCH 27/67] UCP/RNDV: clang --- src/ucp/rndv/proto_rndv.c | 5 +++-- src/ucp/rndv/rndv.c | 2 +- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_put.c | 2 +- src/ucp/rndv/rndv_rtr.c | 2 +- test/gtest/ucp/test_ucp_am.cc | 31 ++++++++++++++++--------------- 6 files changed, 23 insertions(+), 21 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index d4feeb5afbd..e4de223d988 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -719,8 +719,9 @@ unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, if (max_frags == 0) { ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " - "frags", max_mem, ucs_memory_type_names[frag_mem_type], - frag_size, frags_in_chunk, frags_in_chunk); + "frags", + max_mem, ucs_memory_type_names[frag_mem_type], frag_size, + frags_in_chunk, frags_in_chunk); return frags_in_chunk; } diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 2250fe33595..c2bbce48765 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1174,7 +1174,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.alignment = 1; mp_params.elems_per_chunk = num_frags; mp_params.max_elems = ucp_proto_rndv_mtype_fc_max_elems( - worker->context, key.mem_type); + worker->context, key.mem_type); mp_params.ops = &ucp_frag_mpool_ops; mp_params.name = "ucp_rndv_frags"; status = ucs_mpool_init(&mp_params, mpool); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index e076c0f4b1a..4df09b703c0 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -297,7 +297,7 @@ ucp_proto_rndv_get_mtype_fetch_completion(uct_completion_t *uct_comp) static ucs_status_t ucp_proto_rndv_get_mtype_fetch_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); const ucp_proto_rndv_bulk_priv_t *rpriv; ucs_status_t status; diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index a6981c85237..de5cfe6b21f 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,7 +520,7 @@ static UCS_F_INLINE_OPTIMIZED ucs_status_t ucp_proto_rndv_put_mtype_send_func( static ucs_status_t ucp_proto_rndv_put_mtype_copy_progress(uct_pending_req_t *uct_req) { - ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct);\ + ucp_request_t *req = ucs_container_of(uct_req, ucp_request_t, send.uct); const ucp_proto_rndv_put_priv_t *rpriv; ucs_status_t status; diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 4045cc2b2bd..4dd411778c8 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -424,7 +424,7 @@ ucp_proto_rndv_rtr_mtype_data_received(ucp_request_t *req, int in_buffer) static ucs_status_t ucp_proto_rndv_rtr_mtype_progress(uct_pending_req_t *self) { - ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; ucs_status_t status; diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 0249106c187..c2992ef1ad3 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2359,10 +2359,9 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { auto stats_node = e.worker()->stats; auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); - EXPECT_GE(value, min_value) << "counter " - << stats_node->cls->counter_names[cntr] - << " expected >= " << min_value - << " but got " << value; + EXPECT_GE(value, min_value) + << "counter " << stats_node->cls->counter_names[cntr] + << " expected >= " << min_value << " but got " << value; } void check_pending_queues_empty(const entity &e) @@ -2370,7 +2369,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { ucp_worker_h worker = e.worker(); for (int i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; i++) { EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) - << "pending_q[" << i << "] should be empty"; + << "pending_q[" << i << "] should be empty"; } } @@ -2382,8 +2381,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { void verify_clean_fc_state() { - for (auto *ep : { &sender(), &receiver() }) - { + for (auto *ep : {&sender(), &receiver()}) { check_pending_queues_empty(*ep); } } @@ -2399,11 +2397,12 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); - fc.sender_throttled = UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - fc.receiver_throttled = UCS_STATS_GET_COUNTER( - receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + fc.sender_throttled = + UCS_STATS_GET_COUNTER(sender().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); + fc.receiver_throttled = + UCS_STATS_GET_COUNTER(receiver().worker()->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); } }; @@ -2415,7 +2414,7 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, run_fc_test(200, fc); EXPECT_GT(fc.sender_throttled + fc.receiver_throttled, 0u) - << "throttling should have occurred with MAX_MEM=600mb"; + << "throttling should have occurred with MAX_MEM=600mb"; verify_clean_fc_state(); } @@ -2427,8 +2426,10 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, run_fc_test(8, fc); - EXPECT_EQ(0u, fc.sender_throttled) << "FC disabled - no throttling expected"; - EXPECT_EQ(0u, fc.receiver_throttled) << "FC disabled - no throttling expected"; + EXPECT_EQ(0u, fc.sender_throttled) + << "FC disabled - no throttling expected"; + EXPECT_EQ(0u, fc.receiver_throttled) + << "FC disabled - no throttling expected"; } From dce0ce4678f7db210d91232089ed57937124aac2 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 17 Mar 2026 17:13:24 +0200 Subject: [PATCH 28/67] UCP/RNDV: small fixes --- src/ucp/core/ucp_context.h | 2 +- src/ucp/core/ucp_worker.h | 13 ------------- src/ucp/rndv/proto_rndv.c | 17 +---------------- src/ucp/rndv/rndv_mtype.inl | 16 +++++++--------- 4 files changed, 9 insertions(+), 39 deletions(-) diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 2ddabd30755..bf41003d971 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -806,7 +806,7 @@ ucp_context_print_transport_tables_enabled(ucp_context_h context) } static UCS_F_ALWAYS_INLINE int -ucp_context_rndv_mtype_fc_enabled(ucp_context_h context) +ucp_context_rndv_mtype_mem_limit_enabled(ucp_context_h context) { size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 02406ccd9b9..3ad4482ec8d 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -411,19 +411,6 @@ typedef struct ucp_worker { } counters; struct { - /* Worker-level mtype fragment flow control. - * - * Two levels of limiting cooperate: - * 1. Soft (tiered) limit – the active_frags counter tracks how many - * fragments are in-flight. Each operation type (PUT/GET/RTR) is - * capped at a different fraction of the budget so that higher- - * priority operations always have headroom, preventing deadlock. - * 2. Hard limit – the per-mpool quota (ucs_mpool_data::quota, set via - * max_elems when the mpool is created) caps total memory. If the - * mpool is exhausted the request is throttled the same way. - */ - unsigned tier_step; /* Pct step between throttle tiers (cached) */ - size_t active_frags; /* Current active fragments */ /* Pending queues indexed by ucp_worker_rndv_fc_op_t, ordered by * descending priority (GET/PUT = 0, RTR = 1) */ ucs_queue_head_t pending_q[UCP_WORKER_RNDV_FC_OP_LAST]; diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index e4de223d988..d99e625948c 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -680,21 +680,6 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } -/** - * Compute maximum number of elements allowed for the rndv fragment mpool - * based on configured max memory and fragment size for the given memory type. - * The result is rounded down to the allocation chunk granularity - * (rndv_num_frags) and returned as an unsigned value suitable for - * ucs_mpool_params_t::max_elems. - * - * When flow control is disabled (UCX_RNDV_MTYPE_WORKER_MAX_MEM="inf"), - * returns UINT_MAX so that the mpool is effectively unlimited. - * - * @param context The UCP context. - * @param frag_mem_type Memory type used for fragments. - * - * @return Maximum number of mpool elements, aligned to allocation chunk size. - */ unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type) { @@ -703,7 +688,7 @@ unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, size_t frags_in_chunk; size_t max_frags; - if (!ucp_context_rndv_mtype_fc_enabled(context)) { + if (!ucp_context_rndv_mtype_mem_limit_enabled(context)) { return UINT_MAX; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 0823baca099..7952831f63e 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -191,15 +191,13 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) /** * Reschedule a pending throttled request after a fragment is released back to - * the mpool. Dequeue priority: PUT > GET > RTR. + * the mpool. Dequeue priority: PUT/GET > RTR. * * Priority rationale: - * PUT - Remote is blocked waiting for our data. Scheduling PUT unblocks remote - * as well. - * GET - Self-contained fetch operation. Completes without causing remote - * allocations, but scheduling it doesn't unblock another buffer. - * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total - * memory pressure. + * PUT/GET - Completing a PUT or GET frees a staging buffer, reducing memory + * pressure. PUT also unblocks the remote side. + * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total + * memory pressure. */ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) @@ -218,8 +216,8 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) /* Advance best_q past empty queues */ while ((worker->rndv_mtype_fc.best_q < UCP_WORKER_RNDV_FC_OP_LAST) && - ucs_queue_is_empty( - &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q])) { + ucs_queue_is_empty( + &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q])) { worker->rndv_mtype_fc.best_q++; } From 0f877beb2521155fe53731eabce39d8108211948 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 11 Aug 2026 17:12:41 +0300 Subject: [PATCH 29/67] UCP/RNDV: disable throttling by default (RNDV_MTYPE_WORKER_MAX_MEM=inf) --- src/ucp/core/ucp_context.c | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 6529e9d3d67..bb75751a556 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,10 +404,13 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_MAX_MEM", "8g", - "Maximum memory for concurrent mtype fragments per worker.\n" - "This value is translated to a fragment count based on RNDV_FRAG_SIZE\n" - "for each memory type. Set to \"inf\" to disable flow control.", + {"RNDV_MTYPE_WORKER_MAX_MEM", "inf", + "Maximum amount of memory a worker may use for rendezvous staging\n" + "fragments, or \"inf\" for no limit. The limit is enforced separately for\n" + "each fragment memory type and device, and is translated to a fragment\n" + "count using RNDV_FRAG_SIZE and RNDV_FRAG_ALLOC_COUNT. While a limit is\n" + "reached, requests that need a staging fragment are queued until\n" + "fragments are released.", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", From b548a0b35638d995c1a38500c2de177cb73b41f8 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 12 Aug 2026 11:03:20 +0300 Subject: [PATCH 30/67] UCP/RNDV: code review fixes --- src/ucp/core/ucp_context.h | 8 ------ src/ucp/core/ucp_worker.c | 2 +- src/ucp/core/ucp_worker.h | 2 +- src/ucp/rndv/proto_rndv.c | 8 +++++- src/ucp/rndv/rndv_mtype.inl | 4 +-- test/gtest/ucp/test_ucp_am.cc | 53 ++++++++++++++++++----------------- 6 files changed, 39 insertions(+), 38 deletions(-) diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index bf41003d971..509e6cb910c 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -804,14 +804,6 @@ ucp_context_print_transport_tables_enabled(ucp_context_h context) return value == UCS_CONFIG_ON; } - -static UCS_F_ALWAYS_INLINE int -ucp_context_rndv_mtype_mem_limit_enabled(ucp_context_h context) -{ - size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; - return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); -} - void ucp_context_memaccess_tl_bitmap(ucp_context_h context, uint64_t mem_type_bitmap, uint64_t md_reg_flags, diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index b35354dfe09..c1c78d7483b 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2610,7 +2610,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context, unsigned name_length; ucp_worker_h worker; ucs_status_t status; - int q_index; + unsigned q_index; worker = ucs_calloc(1, sizeof(*worker), "ucp worker"); if (worker == NULL) { diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 3ad4482ec8d..dbdea76c9dd 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -416,7 +416,7 @@ typedef struct ucp_worker { ucs_queue_head_t pending_q[UCP_WORKER_RNDV_FC_OP_LAST]; /* Index of highest-priority non-empty queue * (UCP_WORKER_RNDV_FC_OP_LAST if all queues are empty) */ - int best_q; + unsigned best_q; } rndv_mtype_fc; struct { diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index d99e625948c..f7d4a6dc709 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -9,7 +9,6 @@ #endif #include "proto_rndv.inl" -#include "rndv_mtype.inl" #include #include @@ -680,6 +679,13 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } +static int +ucp_context_rndv_mtype_mem_limit_enabled(ucp_context_h context) +{ + const size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; + return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); +} + unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type) { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 7952831f63e..72df1fa5f02 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -48,7 +48,7 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, ucs_sys_device_t frag_sys_dev, - int fc_op) + unsigned fc_op) { ucp_worker_h worker = req->send.ep->worker; @@ -202,7 +202,7 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) { - ucp_worker_h worker = req->send.ep->worker; + ucp_worker_h worker = req->send.ep->worker; ucs_queue_elem_t *elem; ucp_request_t *pending_req; diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index c2992ef1ad3..492d5d38aaa 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2354,31 +2354,6 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { uint64_t receiver_throttled; }; - void check_stats_ge(const entity &e, uint64_t cntr, uint64_t min_value) - { - auto stats_node = e.worker()->stats; - auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); - - EXPECT_GE(value, min_value) - << "counter " << stats_node->cls->counter_names[cntr] - << " expected >= " << min_value << " but got " << value; - } - - void check_pending_queues_empty(const entity &e) - { - ucp_worker_h worker = e.worker(); - for (int i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; i++) { - EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) - << "pending_q[" << i << "] should be empty"; - } - } - - void send_message(size_t num_frags) - { - set_mem_type(UCS_MEMORY_TYPE_CUDA_MANAGED); - test_am_send_recv(get_rndv_frag_size(UCS_MEMORY_TYPE_CUDA) * num_frags); - } - void verify_clean_fc_state() { for (auto *ep : {&sender(), &receiver()}) { @@ -2404,6 +2379,34 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { UCS_STATS_GET_COUNTER(receiver().worker()->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); } + +private: + static void check_stats_ge(const entity &e, uint64_t cntr, + uint64_t min_value) + { + const auto stats_node = e.worker()->stats; + const auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); + + EXPECT_GE(value, min_value) + << "counter " << stats_node->cls->counter_names[cntr] + << " expected >= " << min_value << " but got " << value; + } + + static void check_pending_queues_empty(const entity &e) + { + const ucp_worker_h worker = e.worker(); + + for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; i++) { + EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) + << "pending_q[" << i << "] should be empty"; + } + } + + void send_message(size_t num_frags) + { + set_mem_type(UCS_MEMORY_TYPE_CUDA_MANAGED); + test_am_send_recv(get_rndv_frag_size(UCS_MEMORY_TYPE_CUDA) * num_frags); + } }; UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, From b8021655f9b0005005f62bb1e810767ce179a9f5 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 12 Aug 2026 14:14:16 +0300 Subject: [PATCH 31/67] UCP/RNDV: code review fixes 2 --- src/ucp/core/ucp_context.c | 12 +++++------ src/ucp/core/ucp_worker.c | 1 - src/ucp/core/ucp_worker.h | 3 --- src/ucp/rndv/proto_rndv.c | 18 ++++------------ src/ucp/rndv/proto_rndv.h | 17 ++------------- src/ucp/rndv/rndv_mtype.inl | 42 ++++++++++++++++++------------------- 6 files changed, 33 insertions(+), 60 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index bb75751a556..111e2d398c3 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,13 +404,13 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_MAX_MEM", "inf", + {"RNDV_MTYPE_WORKER_MAX_MEM", "auto", "Maximum amount of memory a worker may use for rendezvous staging\n" - "fragments, or \"inf\" for no limit. The limit is enforced separately for\n" - "each fragment memory type and device, and is translated to a fragment\n" - "count using RNDV_FRAG_SIZE and RNDV_FRAG_ALLOC_COUNT. While a limit is\n" - "reached, requests that need a staging fragment are queued until\n" - "fragments are released.", + "fragments, or \"inf\" for no limit. The same limit is applied\n" + "independently to each fragment memory type, and per device when\n" + "available. It is converted to a fragment count using RNDV_FRAG_SIZE and\n" + "RNDV_FRAG_ALLOC_COUNT. When the limit is reached, further fragment\n" + "requests are queued until fragments are released", ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index c1c78d7483b..fa6d78f5e08 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2642,7 +2642,6 @@ ucs_status_t ucp_worker_create(ucp_context_h context, worker->counters.ep_failures = 0; /* Initialize RNDV mtype flow control pending queues */ - worker->rndv_mtype_fc.best_q = UCP_WORKER_RNDV_FC_OP_LAST; for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { ucs_queue_head_init(&worker->rndv_mtype_fc.pending_q[q_index]); } diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index dbdea76c9dd..8a39a7f2bc3 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -414,9 +414,6 @@ typedef struct ucp_worker { /* Pending queues indexed by ucp_worker_rndv_fc_op_t, ordered by * descending priority (GET/PUT = 0, RTR = 1) */ ucs_queue_head_t pending_q[UCP_WORKER_RNDV_FC_OP_LAST]; - /* Index of highest-priority non-empty queue - * (UCP_WORKER_RNDV_FC_OP_LAST if all queues are empty) */ - unsigned best_q; } rndv_mtype_fc; struct { diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index f7d4a6dc709..b5c2cd6e06c 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -679,35 +679,25 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } -static int -ucp_context_rndv_mtype_mem_limit_enabled(ucp_context_h context) -{ - const size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; - return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); -} - unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type) { - size_t max_mem; + const size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; size_t frag_size; size_t frags_in_chunk; size_t max_frags; - if (!ucp_context_rndv_mtype_mem_limit_enabled(context)) { + if ((max_mem == UCS_MEMUNITS_INF) || (max_mem == UCS_MEMUNITS_AUTO)) { return UINT_MAX; } - max_mem = context->config.ext.rndv_mtype_worker_max_mem; frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; ucs_assert(frag_size > 0); - /* Compute max fragments and round down to allocation chunk granularity */ max_frags = max_mem / frag_size; - max_frags = (max_frags / frags_in_chunk) * frags_in_chunk; - - if (max_frags == 0) { + if (max_frags < frags_in_chunk) { + /* mpool requires max_elems >= elems_per_chunk */ ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " "frags", diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 7e98f9e709c..5ea74580a0a 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -191,21 +191,8 @@ void ucp_proto_rndv_rts_abort(ucp_request_t *req, ucs_status_t status); ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); -/** - * Compute maximum number of elements allowed for the rndv fragment mpool - * based on configured max memory and fragment size for the given memory type. - * The result is rounded down to the allocation chunk granularity - * (rndv_num_frags) and returned as an unsigned value suitable for - * ucs_mpool_params_t::max_elems. - * - * When flow control is disabled (UCX_RNDV_MTYPE_WORKER_MAX_MEM="inf"), - * returns UINT_MAX so that the mpool is effectively unlimited. - * - * @param context The UCP context. - * @param frag_mem_type Memory type used for fragments. - * - * @return Maximum number of mpool elements, aligned to allocation chunk size. - */ +/* Convert UCX_RNDV_MTYPE_WORKER_MAX_MEM to mpool max_elems (frag count), or + * UINT_MAX if unlimited. */ unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 72df1fa5f02..ffd51c9f450 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -10,6 +10,7 @@ #include "rndv.h" #include +#include static ucp_ep_h ucp_proto_rndv_mtype_ep(ucp_worker_t *worker, @@ -57,14 +58,16 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, if (req->send.rndv.mdesc == NULL) { /* Mpool quota exhausted - throttle by queuing the request in the * appropriate pending queue ordered by priority. */ - ucs_trace_req("mtype_fc: frag mpool exhausted, throttling request"); + ucp_trace_req(req, + "mtype_fc: mpool exhausted, queue %s mem_type %s " + "sys_dev %s", + (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", + ucs_memory_type_names[frag_mem_type], + ucs_topo_sys_device_get_name(frag_sys_dev)); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.ppln.queue_elem); - if (fc_op < worker->rndv_mtype_fc.best_q) { - worker->rndv_mtype_fc.best_q = fc_op; - } return UCS_ERR_NO_RESOURCE; } @@ -205,27 +208,24 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) ucp_worker_h worker = req->send.ep->worker; ucs_queue_elem_t *elem; ucp_request_t *pending_req; + unsigned q_index; - /* Dequeue from highest-priority non-empty queue */ - if (worker->rndv_mtype_fc.best_q >= UCP_WORKER_RNDV_FC_OP_LAST) { - return; - } + /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ + for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { + if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { + continue; + } - elem = ucs_queue_pull( - &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q]); + elem = ucs_queue_pull(&worker->rndv_mtype_fc.pending_q[q_index]); + ucs_assert(elem != NULL); - /* Advance best_q past empty queues */ - while ((worker->rndv_mtype_fc.best_q < UCP_WORKER_RNDV_FC_OP_LAST) && - ucs_queue_is_empty( - &worker->rndv_mtype_fc.pending_q[worker->rndv_mtype_fc.best_q])) { - worker->rndv_mtype_fc.best_q++; + pending_req = ucs_container_of(elem, ucp_request_t, + send.rndv.ppln.queue_elem); + ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, + ucp_proto_rndv_mtype_fc_reschedule_cb, + pending_req); + return; } - - pending_req = ucs_container_of(elem, ucp_request_t, - send.rndv.ppln.queue_elem); - ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, - ucp_proto_rndv_mtype_fc_reschedule_cb, - pending_req); } /** From bc5cd27d430c6f45c6570b225ae9ec5dde5dcb08 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 12 Aug 2026 15:31:50 +0300 Subject: [PATCH 32/67] UCP/RNDV: rename RNDV_MTYPE_WORKER_MAX_MEM to RNDV_FRAG_WORKER_MAX_MEM --- src/ucp/core/ucp_context.c | 4 ++-- src/ucp/core/ucp_context.h | 4 ++-- src/ucp/rndv/proto_rndv.c | 8 ++++---- src/ucp/rndv/proto_rndv.h | 6 +++--- src/ucp/rndv/rndv.c | 2 +- test/gtest/ucp/test_ucp_am.cc | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 111e2d398c3..d28903e642d 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -404,14 +404,14 @@ static ucs_config_field_t ucp_context_config_table[] = { "Force-enable the RMA rendezvous put/get protocols.", ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, - {"RNDV_MTYPE_WORKER_MAX_MEM", "auto", + {"RNDV_FRAG_WORKER_MAX_MEM", "auto", "Maximum amount of memory a worker may use for rendezvous staging\n" "fragments, or \"inf\" for no limit. The same limit is applied\n" "independently to each fragment memory type, and per device when\n" "available. It is converted to a fragment count using RNDV_FRAG_SIZE and\n" "RNDV_FRAG_ALLOC_COUNT. When the limit is reached, further fragment\n" "requests are queued until fragments are released", - ucs_offsetof(ucp_context_config_t, rndv_mtype_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, + ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 509e6cb910c..047f76e88bd 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -132,8 +132,8 @@ typedef struct ucp_context_config { int rndv_errh_ppln_enable; /** Force-enable the RMA rendezvous put/get protocols */ int rma_ppln_enable; - /** Maximum memory for concurrent rndv mtype fragments per worker (bytes) */ - size_t rndv_mtype_worker_max_mem; + /** Maximum memory for concurrent rndv fragments per worker (bytes) */ + size_t rndv_frag_worker_max_mem; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index b5c2cd6e06c..5870034da5c 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -679,10 +679,10 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } -unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, - ucs_memory_type_t frag_mem_type) +unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, + ucs_memory_type_t frag_mem_type) { - const size_t max_mem = context->config.ext.rndv_mtype_worker_max_mem; + const size_t max_mem = context->config.ext.rndv_frag_worker_max_mem; size_t frag_size; size_t frags_in_chunk; size_t max_frags; @@ -698,7 +698,7 @@ unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, max_frags = max_mem / frag_size; if (max_frags < frags_in_chunk) { /* mpool requires max_elems >= elems_per_chunk */ - ucs_warn("RNDV_MTYPE_WORKER_MAX_MEM (%zu) is too low for %s " + ucs_warn("RNDV_FRAG_WORKER_MAX_MEM (%zu) is too low for %s " "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " "frags", max_mem, ucs_memory_type_names[frag_mem_type], frag_size, diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 5ea74580a0a..7b3cb08acda 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -191,10 +191,10 @@ void ucp_proto_rndv_rts_abort(ucp_request_t *req, ucs_status_t status); ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); -/* Convert UCX_RNDV_MTYPE_WORKER_MAX_MEM to mpool max_elems (frag count), or +/* Convert UCX_RNDV_FRAG_WORKER_MAX_MEM to mpool max_elems (frag count), or * UINT_MAX if unlimited. */ -unsigned ucp_proto_rndv_mtype_fc_max_elems(ucp_context_h context, - ucs_memory_type_t frag_mem_type); +unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, + ucs_memory_type_t frag_mem_type); ucs_status_t diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index c2bbce48765..9db77a938b4 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1173,7 +1173,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.elem_size = sizeof(ucp_mem_desc_t); mp_params.alignment = 1; mp_params.elems_per_chunk = num_frags; - mp_params.max_elems = ucp_proto_rndv_mtype_fc_max_elems( + mp_params.max_elems = ucp_proto_rndv_frag_max_elems( worker->context, key.mem_type); mp_params.ops = &ucp_frag_mpool_ops; mp_params.name = "ucp_rndv_frags"; diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 492d5d38aaa..7776581097f 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2410,7 +2410,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { }; UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, - "RNDV_MTYPE_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") + "RNDV_FRAG_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") { fc_counters fc; From 398d835a079f4751f661ea8073a2c684194dbdd5 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 12 Aug 2026 18:25:46 +0300 Subject: [PATCH 33/67] UCP/RNDV: throttle only in case of NO_RESOURCE --- src/ucp/rndv/rndv.c | 52 ++++++++++++++------- src/ucp/rndv/rndv.h | 4 +- src/ucp/rndv/rndv_mtype.inl | 37 +++++++-------- test/gtest/ucp/test_ucp_mmap.cc | 82 ++++++++++++++++++++++++++++++++- 4 files changed, 136 insertions(+), 39 deletions(-) diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 9db77a938b4..aed57ebb879 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1136,9 +1136,9 @@ ucs_mpool_ops_t ucp_frag_mpool_ops = { .obj_cleanup = (ucs_mpool_obj_cleanup_func_t)ucs_empty_function }; -ucp_mem_desc_t * +ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, - ucs_sys_device_t sys_dev) + ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p) { ucp_rndv_mpool_priv_t *mpriv; ucp_worker_mpool_key_t key; @@ -1160,7 +1160,8 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, khiter = kh_put(ucp_worker_mpool_hash, &worker->mpool_hash, key, &khret); if (khret == UCS_KH_PUT_FAILED) { - return NULL; + status = UCS_ERR_NO_MEMORY; + goto err; } ucs_assert_always(khret != UCS_KH_PUT_KEY_PRESENT); @@ -1179,16 +1180,29 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.name = "ucp_rndv_frags"; status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { - return NULL; + kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); + goto err; } - mpriv = ucs_mpool_priv(mpool); - mpriv->worker = worker; - mpriv->mem_type = key.mem_type; - mpriv->sys_dev = sys_dev; + mpriv = ucs_mpool_priv(mpool); + mpriv->worker = worker; + mpriv->mem_type = key.mem_type; + mpriv->sys_dev = sys_dev; out_mp_get: - return ucp_worker_mpool_get(mpool); + *mdesc_p = ucp_worker_mpool_get(mpool); + if (*mdesc_p != NULL) { + return UCS_OK; + } + + /* Quota exhaustion can be retried later. Any other NULL means that + * allocation or setup failed. */ + return ucs_mpool_is_empty(mpool) ? UCS_ERR_NO_RESOURCE : + UCS_ERR_NO_MEMORY; + +err: + *mdesc_p = NULL; + return status; } static void ucp_rndv_send_frag_get_mem_type(ucp_request_t *sreq, size_t length, @@ -1203,6 +1217,7 @@ static void ucp_rndv_send_frag_get_mem_type(ucp_request_t *sreq, size_t length, ucs_memory_type_t frag_mem_type = ucp_rndv_frag_mem_type(worker->context); ucp_request_t *freq; ucp_mem_desc_t *mdesc; + ucs_status_t status; /* GET fragment to stage buffer */ @@ -1211,10 +1226,11 @@ static void ucp_rndv_send_frag_get_mem_type(ucp_request_t *sreq, size_t length, ucs_fatal("failed to allocate fragment receive request"); } - mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN); - if (ucs_unlikely(mdesc == NULL)) { - ucs_fatal("failed to allocate fragment memory desc"); + status = ucp_rndv_mpool_get(worker, frag_mem_type, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc); + if (ucs_unlikely(status != UCS_OK)) { + ucs_fatal("failed to allocate fragment memory desc: %s", + ucs_status_string(status)); } freq->send.ep = sreq->send.ep; @@ -1363,6 +1379,7 @@ static void ucp_rndv_send_frag_rtr(ucp_worker_h worker, ucp_request_t *rndv_req, ucp_request_t *freq; ucp_request_t *frndv_req; ucp_md_map_t md_map; + ucs_status_t status; ucp_trace_req(rreq, "using rndv pipeline protocol rndv_req %p", rndv_req); @@ -1388,10 +1405,11 @@ static void ucp_rndv_send_frag_rtr(ucp_worker_h worker, ucp_request_t *rndv_req, } /* allocate fragment recv buffer desc*/ - mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN); - if (mdesc == NULL) { - ucs_fatal("failed to allocate fragment memory buffer"); + status = ucp_rndv_mpool_get(worker, frag_mem_type, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc); + if (status != UCS_OK) { + ucs_fatal("failed to allocate fragment memory buffer: %s", + ucs_status_string(status)); } /* Pass datatype to avoid Coverity warning */ diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index 7156e32d14f..45acfbfabb9 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -114,9 +114,9 @@ ucs_status_t ucp_proto_progress_rndv_rtr(uct_pending_req_t *self); ucs_status_t ucp_rndv_reg_send_buffer(ucp_request_t *sreq, const ucp_request_param_t *param); -ucp_mem_desc_t * +ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, - ucs_sys_device_t sys_dev); + ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p); void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index ffd51c9f450..dcab9fc8919 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -52,27 +52,28 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, unsigned fc_op) { ucp_worker_h worker = req->send.ep->worker; + ucs_status_t status; - req->send.rndv.mdesc = ucp_rndv_mpool_get(worker, frag_mem_type, - frag_sys_dev); - if (req->send.rndv.mdesc == NULL) { - /* Mpool quota exhausted - throttle by queuing the request in the - * appropriate pending queue ordered by priority. */ - ucp_trace_req(req, - "mtype_fc: mpool exhausted, queue %s mem_type %s " - "sys_dev %s", - (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", - ucs_memory_type_names[frag_mem_type], - ucs_topo_sys_device_get_name(frag_sys_dev)); - UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); - ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], - &req->send.rndv.ppln.queue_elem); - - return UCS_ERR_NO_RESOURCE; + status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, + &req->send.rndv.mdesc); + if (status != UCS_ERR_NO_RESOURCE) { + return status; } - return UCS_OK; + /* Mpool quota exhausted - throttle by queuing the request in the + * appropriate pending queue ordered by priority. */ + ucp_trace_req(req, + "mtype_fc: mpool exhausted, queue %s mem_type %s " + "sys_dev %s", + (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", + ucs_memory_type_names[frag_mem_type], + ucs_topo_sys_device_get_name(frag_sys_dev)); + UCS_STATS_UPDATE_COUNTER(worker->stats, + UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], + &req->send.rndv.ppln.queue_elem); + + return UCS_ERR_NO_RESOURCE; } static UCS_F_ALWAYS_INLINE uct_mem_h diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 29d78753232..2ae350ff118 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1013,17 +1013,95 @@ UCS_TEST_P(test_ucp_mmap, gva, "GVA_ENABLE=y") } } +static ucs_status_t +rndv_mpool_chunk_alloc_fail(ucs_mpool_t*, size_t*, void**) +{ + return UCS_ERR_NO_MEMORY; +} + UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) { ucp_worker_h worker = sender().worker(); for (auto mem_type : mem_buffer::supported_mem_types()) { - ucp_mem_desc_t *mdesc = ucp_rndv_mpool_get(worker, mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN); + ucp_mem_desc_t *mdesc; + + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, mem_type, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc)); + ASSERT_NE(nullptr, mdesc); EXPECT_EQ(mdesc->memh, mdesc->memh->parent); ucs_mpool_put(mdesc); } } +UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_exhausted, + "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", + "RNDV_FRAG_WORKER_MAX_MEM=8K") +{ + ucp_mem_desc_t *mdesc1; + ucp_mem_desc_t *mdesc2; + ucp_mem_desc_t *mdesc3; + ucp_worker_h worker = sender().worker(); + + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); + EXPECT_EQ(UCS_ERR_NO_RESOURCE, + ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3)); + EXPECT_EQ(nullptr, mdesc3); + + ucs_mpool_put(mdesc2); + ucs_mpool_put(mdesc1); +} + +UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, + "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", + "RNDV_FRAG_WORKER_MAX_MEM=16K") +{ + ucp_worker_mpool_key_t key; + ucp_mem_desc_t *mdesc1; + ucp_mem_desc_t *mdesc2; + ucp_mem_desc_t *mdesc3; + ucp_worker_h worker = sender().worker(); + ucs_mpool_ops_t fail_ops; + const ucs_mpool_ops_t *orig_ops; + ucs_mpool_t *mpool; + ucs_status_t status; + khiter_t khiter; + + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); + + key.mem_type = UCS_MEMORY_TYPE_HOST; + key.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); + ASSERT_NE(kh_end(&worker->mpool_hash), khiter); + mpool = &kh_val(&worker->mpool_hash, khiter); + + /* The first chunk is exhausted, but the pool still has quota. Make the + * next chunk allocation fail and verify it is not reported as quota + * exhaustion. */ + orig_ops = mpool->data->ops; + fail_ops = *orig_ops; + fail_ops.chunk_alloc = rndv_mpool_chunk_alloc_fail; + mpool->data->ops = &fail_ops; + { + scoped_log_handler slh(hide_errors_logger); + status = ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3); + } + mpool->data->ops = orig_ops; + + EXPECT_EQ(UCS_ERR_NO_MEMORY, status); + EXPECT_EQ(nullptr, mdesc3); + + ucs_mpool_put(mdesc2); + ucs_mpool_put(mdesc1); +} + UCS_TEST_P(test_ucp_mmap, no_sys_dev_with_zero_md_map) { const std::vector &mem_types = From e683f95a7649bd5a55b00bdb327b257030a877f1 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 12 Aug 2026 18:31:00 +0300 Subject: [PATCH 34/67] UCP/RNDV: more fixes --- src/ucp/rndv/proto_rndv.c | 2 +- src/ucp/rndv/rndv_mtype.inl | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 5870034da5c..8667736375d 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -693,7 +693,7 @@ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; - ucs_assert(frag_size > 0); + ucs_assert_always(frag_size > 0); max_frags = max_mem / frag_size; if (max_frags < frags_in_chunk) { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index dcab9fc8919..91c0714b2f0 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -207,7 +207,6 @@ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) { ucp_worker_h worker = req->send.ep->worker; - ucs_queue_elem_t *elem; ucp_request_t *pending_req; unsigned q_index; @@ -217,11 +216,11 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) continue; } - elem = ucs_queue_pull(&worker->rndv_mtype_fc.pending_q[q_index]); - ucs_assert(elem != NULL); - - pending_req = ucs_container_of(elem, ucp_request_t, - send.rndv.ppln.queue_elem); + pending_req = ucs_queue_pull_elem_non_empty( + &worker->rndv_mtype_fc.pending_q[q_index], ucp_request_t, + send.rndv.ppln.queue_elem); + ucp_trace_req(pending_req, "mtype_fc: dequeue %s", + (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, ucp_proto_rndv_mtype_fc_reschedule_cb, pending_req); From c974c97ce9f640ea2d9224e4e5253eebea2812ac Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 19 Aug 2026 16:45:41 +0300 Subject: [PATCH 35/67] UCP/RNDV: add abort path gtests --- test/gtest/ucp/test_ucp_peer_failure.cc | 60 ++++++++++++++++++++++--- 1 file changed, 55 insertions(+), 5 deletions(-) diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index f814d792185..4d43f5d405c 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -745,7 +745,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { enum class rndv_mode { rndv_get, rndv_put, - put_ppln + put_ppln, + rtr_mtype_fc_pending }; void init() @@ -837,6 +838,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { auto *req = ucs_container_of(uct_req, ucp_request_t, send.uct); auto stage = req->send.proto_stage; auto *proto = req->send.proto_config->proto; + auto *worker = req->send.ep->worker; + ucs_status_t status; ucs::mock mock; if (proto->name == self->m_proto_name) { @@ -844,14 +847,42 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { mock_rndv_ops(req->send.ep, mock); } - if (stage == self->m_proto_xfer_stage) { + if ((stage == self->m_proto_xfer_stage) && + !self->m_close_on_fc_pending) { self->close_peer(); } } /* Call original proto progress */ - return self->m_progress_mock.orig_func(&proto->progress[stage], - uct_req); + status = self->m_progress_mock.orig_func(&proto->progress[stage], + uct_req); + + if ((proto->name == self->m_proto_name) && + (stage == self->m_proto_xfer_stage) && + self->m_close_on_fc_pending && + self->has_fc_pending(worker)) { + self->m_fc_was_pending = true; + self->close_peer(); + } + + return status; + } + + static bool has_fc_pending(ucp_worker_h worker) + { + for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; ++i) { + if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) { + return true; + } + } + + return false; + } + + void verify_fc_queues_empty() + { + EXPECT_FALSE(has_fc_pending(sender().worker())); + EXPECT_FALSE(has_fc_pending(receiver().worker())); } virtual void cleanup() @@ -863,7 +894,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { void define_test_settings(rndv_mode mode, bool replace_ops) { - m_replace_ops = replace_ops; + m_replace_ops = replace_ops; + m_close_on_fc_pending = false; switch (mode) { case rndv_mode::rndv_get: @@ -881,6 +913,12 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { m_proto_xfer_stage = UCP_PROTO_RNDV_PUT_MTYPE_STAGE_SEND; m_peer_to_close = &self->receiver(); break; + case rndv_mode::rtr_mtype_fc_pending: + m_proto_name = "rndv/rtr/mtype"; + m_proto_xfer_stage = 0; + m_peer_to_close = &self->sender(); + m_close_on_fc_pending = true; + break; default: EXPECT_TRUE(false) << "Wrong RNDV mode"; } @@ -920,6 +958,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { ucs::mock m_progress_mock; bool m_is_peer_closed{false}; + bool m_close_on_fc_pending{false}; + bool m_fc_was_pending{false}; std::string m_proto_name{}; /* Protocol stage during which data transfer happens */ uint8_t m_proto_xfer_stage{}; @@ -1019,4 +1059,14 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, pipeline, rndv_progress_failure_test(rndv_mode::put_ppln, true); } +UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_fc_pending, + "RNDV_FRAG_SIZE=host:8K,cuda:8K", + "RNDV_FRAG_WORKER_MAX_MEM=8K") +{ + rndv_progress_failure_test(rndv_mode::rtr_mtype_fc_pending, false); + + EXPECT_TRUE(m_fc_was_pending); + verify_fc_queues_empty(); +} + UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); From 89505d2a8dbb846bf352e3a072f365ddb3d97d40 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 19 Aug 2026 16:46:08 +0300 Subject: [PATCH 36/67] UCP/RNDV: add abort path support --- src/ucp/core/ucp_request.h | 4 ++- src/ucp/rndv/rndv.c | 11 ++++++++ src/ucp/rndv/rndv.h | 2 ++ src/ucp/rndv/rndv_get.c | 9 ++++++- src/ucp/rndv/rndv_mtype.inl | 51 ++++++++++++++++++++++++++++++++----- src/ucp/rndv/rndv_put.c | 9 ++++++- src/ucp/rndv/rndv_rtr.c | 1 + 7 files changed, 78 insertions(+), 9 deletions(-) diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 80583ddb977..dd99bb844b3 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -69,7 +69,9 @@ enum { UCP_REQUEST_FLAG_RNDV_SEND_INTERNAL = UCS_BIT(26), UCP_REQUEST_FLAG_RNDV_GET_REQ = UCS_BIT(27), UCP_REQUEST_FLAG_RNDV_FLUSH = UCS_BIT(28), - UCP_REQUEST_FLAG_RNDV_START_FLUSH = UCS_BIT(29) + UCP_REQUEST_FLAG_RNDV_START_FLUSH = UCS_BIT(29), + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED = UCS_BIT(30), + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED = UCS_BIT(31) }; diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index aed57ebb879..c2072c2e3e1 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -22,6 +22,17 @@ #include +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +{ + ucp_request_t *req = arg; + + UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(req->send.ep->worker); + ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED); + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; + ucp_request_send(req); + return 1; +} + static UCS_F_ALWAYS_INLINE int ucp_rndv_frag_mem_type(ucp_context_t *context) { diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index 45acfbfabb9..b62bca378c1 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -118,6 +118,8 @@ ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p); +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg); + void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, const void *rkey_buf); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 4df09b703c0..eda8196eb37 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -392,6 +392,13 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) return UCS_OK; } +static void +ucp_proto_rndv_get_mtype_abort(ucp_request_t *req, ucs_status_t status) +{ + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_GET); + ucp_proto_abort_fatal_not_implemented(req, status); +} + ucp_proto_t ucp_rndv_get_mtype_proto = { .name = "rndv/get/mtype", .desc = NULL, @@ -403,6 +410,6 @@ ucp_proto_t ucp_rndv_get_mtype_proto = { [UCP_PROTO_RNDV_GET_STAGE_FETCH] = ucp_proto_rndv_get_mtype_fetch_progress, [UCP_PROTO_RNDV_GET_STAGE_ATS] = ucp_proto_rndv_ats_progress }, - .abort = ucp_proto_abort_fatal_not_implemented, + .abort = ucp_proto_rndv_get_mtype_abort, .reset = ucp_proto_rndv_get_mtype_reset }; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 91c0714b2f0..a6fd0bcff8b 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -54,6 +54,7 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucp_worker_h worker = req->send.ep->worker; ucs_status_t status; + UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, &req->send.rndv.mdesc); if (status != UCS_ERR_NO_RESOURCE) { @@ -70,6 +71,10 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_topo_sys_device_get_name(frag_sys_dev)); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + ucs_assert(!(req->flags & + (UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED))); + req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.ppln.queue_elem); @@ -184,13 +189,39 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_copy( return status; } -/* Reschedule callback for throttled mtype requests */ -static UCS_F_ALWAYS_INLINE -unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +static UCS_F_ALWAYS_INLINE int +ucp_proto_rndv_mtype_fc_reschedule_pred(const ucs_callbackq_elem_t *elem, + void *arg) { - ucp_request_t *req = arg; - ucp_request_send(req); - return 1; + return (elem->cb == ucp_proto_rndv_mtype_fc_reschedule_cb) && + (elem->arg == arg); +} + +static UCS_F_ALWAYS_INLINE void +ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) +{ + ucp_worker_h worker = req->send.ep->worker; + + UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); + ucs_assert(fc_op < UCP_WORKER_RNDV_FC_OP_LAST); + ucs_assert(!ucs_test_all_flags(req->flags, + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED)); + + if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED) { + ucp_trace_req(req, "mtype_fc: remove aborted request from queue"); + ucs_queue_remove(&worker->rndv_mtype_fc.pending_q[fc_op], + &req->send.rndv.ppln.queue_elem); + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; + } + + if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED) { + ucp_trace_req(req, "mtype_fc: remove aborted reschedule callback"); + ucs_callbackq_remove_oneshot( + &worker->uct->progress_q, req, + ucp_proto_rndv_mtype_fc_reschedule_pred, req); + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; + } } /** @@ -210,6 +241,8 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) ucp_request_t *pending_req; unsigned q_index; + UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); + /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { @@ -219,6 +252,12 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) pending_req = ucs_queue_pull_elem_non_empty( &worker->rndv_mtype_fc.pending_q[q_index], ucp_request_t, send.rndv.ppln.queue_elem); + ucs_assert(pending_req->flags & + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED); + ucs_assert(!(pending_req->flags & + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED)); + pending_req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; + pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; ucp_trace_req(pending_req, "mtype_fc: dequeue %s", (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index de5cfe6b21f..efa1097065e 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -667,6 +667,13 @@ ucp_proto_rndv_put_mtype_query(const ucp_proto_query_params_t *params, put_desc); } +static void +ucp_proto_rndv_put_mtype_abort(ucp_request_t *req, ucs_status_t status) +{ + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_PUT); + ucp_proto_rndv_stub_abort(req, status); +} + ucp_proto_t ucp_rndv_put_mtype_proto = { .name = "rndv/put/mtype", .desc = NULL, @@ -681,6 +688,6 @@ ucp_proto_t ucp_rndv_put_mtype_proto = { [UCP_PROTO_RNDV_PUT_STAGE_ATP] = ucp_proto_rndv_put_common_atp_progress, [UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP] = ucp_proto_rndv_put_common_fenced_atp_progress, }, - .abort = ucp_proto_rndv_stub_abort, + .abort = ucp_proto_rndv_put_mtype_abort, .reset = (ucp_request_reset_func_t)ucp_proto_reset_fatal_not_implemented }; diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 4dd411778c8..465f9df8105 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -375,6 +375,7 @@ ucp_proto_rndv_rtr_mtype_abort(ucp_request_t *req, ucs_status_t status) { ucp_request_t *super_req = ucp_request_get_super(req); + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_RTR); super_req->status = status; /* When pipeline is used, there is a top-level 'recv' request that we also From d52d74c3ebdee8ea0eb4fc44e7528b4760967386 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 20 Aug 2026 17:56:56 +0300 Subject: [PATCH 37/67] UCP/RNDV: abort gtest fix --- src/ucp/core/ucp_request.h | 4 +- src/ucp/rndv/rndv.c | 4 +- src/ucp/rndv/rndv_mtype.inl | 12 +- test/gtest/ucp/test_ucp_peer_failure.cc | 158 +++++++++++++++++++----- 4 files changed, 138 insertions(+), 40 deletions(-) diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index dd99bb844b3..33f5460f24e 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -70,8 +70,8 @@ enum { UCP_REQUEST_FLAG_RNDV_GET_REQ = UCS_BIT(27), UCP_REQUEST_FLAG_RNDV_FLUSH = UCS_BIT(28), UCP_REQUEST_FLAG_RNDV_START_FLUSH = UCS_BIT(29), - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED = UCS_BIT(30), - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED = UCS_BIT(31) + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED = UCS_BIT(30), + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED = UCS_BIT(31) }; diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index c2072c2e3e1..3785f892c0f 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -27,8 +27,8 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) ucp_request_t *req = arg; UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(req->send.ep->worker); - ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED); - req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; + ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; ucp_request_send(req); return 1; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index a6fd0bcff8b..c95d7bbde3e 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -73,7 +73,7 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); ucs_assert(!(req->flags & (UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED))); + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED))); req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.ppln.queue_elem); @@ -206,7 +206,7 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) ucs_assert(fc_op < UCP_WORKER_RNDV_FC_OP_LAST); ucs_assert(!ucs_test_all_flags(req->flags, UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED)); + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED)); if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED) { ucp_trace_req(req, "mtype_fc: remove aborted request from queue"); @@ -215,12 +215,12 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; } - if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED) { + if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED) { ucp_trace_req(req, "mtype_fc: remove aborted reschedule callback"); ucs_callbackq_remove_oneshot( &worker->uct->progress_q, req, ucp_proto_rndv_mtype_fc_reschedule_pred, req); - req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; } } @@ -255,9 +255,9 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) ucs_assert(pending_req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED); ucs_assert(!(pending_req->flags & - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED)); + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED)); pending_req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; - pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_SCHEDULED; + pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; ucp_trace_req(pending_req, "mtype_fc: dequeue %s", (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 4d43f5d405c..76d99d0457b 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -13,7 +13,9 @@ extern "C" { #include /* for testing EP RNDV configuration */ #include /* for debug */ #include /* for testing memory consumption */ +#include #include +#include } #include @@ -745,8 +747,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { enum class rndv_mode { rndv_get, rndv_put, - put_ppln, - rtr_mtype_fc_pending + put_ppln }; void init() @@ -816,6 +817,32 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } } + static unsigned abort_fc_pending_cb(void *arg) + { + auto *test = static_cast(arg); + ucp_worker_h worker = test->receiver().worker(); + std::vector pending_reqs; + ucp_request_t *req; + + /* Abort may remove the current request, so snapshot the queue first. */ + for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; ++i) { + auto *queue = &worker->rndv_mtype_fc.pending_q[i]; + + ucs_queue_for_each(req, queue, send.rndv.ppln.queue_elem) { + pending_reqs.push_back(req); + } + } + + for (auto *pending_req : pending_reqs) { + ucp_proto_request_abort(pending_req, UCS_ERR_CONNECTION_RESET); + } + + /* Complete the top-level send after aborting receiver fragments. */ + test->sender().close_all_eps(*test, 0, UCP_EP_CLOSE_FLAG_FORCE); + test->m_is_peer_closed = true; + return 1; + } + void close_peer() { if (m_is_peer_closed) { @@ -835,9 +862,9 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { static ucs_status_t progress_wrapper(uct_pending_req_t *uct_req) { - auto *req = ucs_container_of(uct_req, ucp_request_t, send.uct); - auto stage = req->send.proto_stage; - auto *proto = req->send.proto_config->proto; + auto *req = ucs_container_of(uct_req, ucp_request_t, send.uct); + auto stage = req->send.proto_stage; + auto *proto = req->send.proto_config->proto; auto *worker = req->send.ep->worker; ucs_status_t status; ucs::mock mock; @@ -847,8 +874,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { mock_rndv_ops(req->send.ep, mock); } - if ((stage == self->m_proto_xfer_stage) && - !self->m_close_on_fc_pending) { + if (stage == self->m_proto_xfer_stage) { self->close_peer(); } } @@ -857,12 +883,13 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { status = self->m_progress_mock.orig_func(&proto->progress[stage], uct_req); - if ((proto->name == self->m_proto_name) && - (stage == self->m_proto_xfer_stage) && - self->m_close_on_fc_pending && + if (self->m_abort_fc_pending && + (worker == self->receiver().worker()) && self->has_fc_pending(worker)) { - self->m_fc_was_pending = true; - self->close_peer(); + self->m_abort_fc_pending = false; + /* Abort after the current protocol-progress call returns. */ + ucs_callbackq_add_oneshot(&worker->uct->progress_q, self, + abort_fc_pending_cb, self); } return status; @@ -879,10 +906,53 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { return false; } - void verify_fc_queues_empty() + static std::vector + hold_cuda_fragments(ucp_worker_h worker) { - EXPECT_FALSE(has_fc_pending(sender().worker())); - EXPECT_FALSE(has_fc_pending(receiver().worker())); + std::vector held; + ucp_mem_desc_t *mdesc; + khiter_t iter; + + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (iter = kh_begin(&worker->mpool_hash); + iter != kh_end(&worker->mpool_hash); ++iter) { + if (!kh_exist(&worker->mpool_hash, iter) || + (kh_key(&worker->mpool_hash, iter).mem_type != + UCS_MEMORY_TYPE_CUDA)) { + continue; + } + + while ((mdesc = static_cast( + ucs_mpool_get_inline( + &kh_val(&worker->mpool_hash, iter)))) != + nullptr) { + held.push_back(mdesc); + } + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); + + return held; + } + + static void + release_cuda_fragments(ucp_worker_h worker, + std::vector &held) + { + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (auto *mdesc : held) { + /* + * Return the descriptor without invoking FC rescheduling. Aborted + * requests must already have been removed from pending_q. + */ + ucs_mpool_put_inline(mdesc); + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); + + held.clear(); } virtual void cleanup() @@ -894,8 +964,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { void define_test_settings(rndv_mode mode, bool replace_ops) { - m_replace_ops = replace_ops; - m_close_on_fc_pending = false; + m_replace_ops = replace_ops; switch (mode) { case rndv_mode::rndv_get: @@ -913,12 +982,6 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { m_proto_xfer_stage = UCP_PROTO_RNDV_PUT_MTYPE_STAGE_SEND; m_peer_to_close = &self->receiver(); break; - case rndv_mode::rtr_mtype_fc_pending: - m_proto_name = "rndv/rtr/mtype"; - m_proto_xfer_stage = 0; - m_peer_to_close = &self->sender(); - m_close_on_fc_pending = true; - break; default: EXPECT_TRUE(false) << "Wrong RNDV mode"; } @@ -956,10 +1019,46 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { ASSERT_TRUE(m_is_peer_closed); } + void rndv_fc_pending_abort_test() + { + ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); + std::vector held_mdescs; + std::pair result; + + if (sender_config->key.rma_bw_lanes[0] == UCP_NULL_LANE) { + UCS_TEST_SKIP_R("transport has no rma_bw lanes"); + } + + init_buffers(16 * UCS_KBYTE); + + /* Complete one transfer so the CUDA fragment mpool is created. */ + smoke_test(true); + held_mdescs = hold_cuda_fragments(receiver().worker()); + ASSERT_FALSE(held_mdescs.empty()); + + m_abort_fc_pending = true; + setup_progress_mock(sender().worker(), m_progress_mock); + setup_progress_mock(receiver().worker(), m_progress_mock); + + { + scoped_log_handler err_wrapper(wrap_errors_logger); + scoped_log_handler warn_wrapper(wrap_warns_logger); + result = smoke_test(true); + } + + release_cuda_fragments(receiver().worker(), held_mdescs); + m_progress_mock.cleanup(); + + EXPECT_TRUE(UCS_STATUS_IS_ERR(result.first)); + EXPECT_TRUE(UCS_STATUS_IS_ERR(result.second)); + EXPECT_TRUE(m_is_peer_closed); + EXPECT_FALSE(has_fc_pending(sender().worker())); + EXPECT_FALSE(has_fc_pending(receiver().worker())); + } + ucs::mock m_progress_mock; bool m_is_peer_closed{false}; - bool m_close_on_fc_pending{false}; - bool m_fc_was_pending{false}; + bool m_abort_fc_pending{false}; std::string m_proto_name{}; /* Protocol stage during which data transfer happens */ uint8_t m_proto_xfer_stage{}; @@ -1061,12 +1160,11 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, pipeline, UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_fc_pending, "RNDV_FRAG_SIZE=host:8K,cuda:8K", - "RNDV_FRAG_WORKER_MAX_MEM=8K") + "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=8K", + "RNDV_FRAG_MEM_TYPE=cuda") { - rndv_progress_failure_test(rndv_mode::rtr_mtype_fc_pending, false); - - EXPECT_TRUE(m_fc_was_pending); - verify_fc_queues_empty(); + rndv_fc_pending_abort_test(); } UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); From 4c746e907235b8cac4fe3ed03f3881828b6e9490 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 26 Aug 2026 15:49:51 +0300 Subject: [PATCH 38/67] UCP/RNDV: code review fixes --- src/ucp/rndv/proto_rndv.h | 2 ++ src/ucp/rndv/rndv.c | 12 ++++++++---- test/gtest/ucp/test_ucp_mmap.cc | 2 -- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 7b3cb08acda..803ff915e2e 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -37,6 +37,8 @@ ucp_proto_rndv_rts_tag_op_flags(ucp_rndv_rts_opcode_t opcode) return (opcode == UCP_RNDV_RTS_TAG_OK) ? UCP_PROTO_SELECT_OP_FLAG_TAG_RNDV : 0; } + + /** * Rendezvous protocol which sends a control message to the remote peer, and not * actually transferring bulk data. The remote peer is expected to perform the diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 3785f892c0f..7cd4bf96444 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -21,7 +21,6 @@ #include #include - unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) { ucp_request_t *req = arg; @@ -1152,6 +1151,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p) { ucp_rndv_mpool_priv_t *mpriv; + ucp_mem_desc_t *mdesc; ucp_worker_mpool_key_t key; ucs_status_t status; unsigned num_frags; @@ -1192,6 +1192,10 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); + if (status == UCS_ERR_NO_RESOURCE) { + status = UCS_ERR_NO_MEMORY; + } + goto err; } @@ -1201,8 +1205,9 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mpriv->sys_dev = sys_dev; out_mp_get: - *mdesc_p = ucp_worker_mpool_get(mpool); - if (*mdesc_p != NULL) { + mdesc = ucp_worker_mpool_get(mpool); + if (ucs_likely(mdesc != NULL)) { + *mdesc_p = mdesc; return UCS_OK; } @@ -1212,7 +1217,6 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, UCS_ERR_NO_MEMORY; err: - *mdesc_p = NULL; return status; } diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 2ae350ff118..19c9c5a63c9 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1049,7 +1049,6 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_exhausted, EXPECT_EQ(UCS_ERR_NO_RESOURCE, ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3)); - EXPECT_EQ(nullptr, mdesc3); ucs_mpool_put(mdesc2); ucs_mpool_put(mdesc1); @@ -1096,7 +1095,6 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, mpool->data->ops = orig_ops; EXPECT_EQ(UCS_ERR_NO_MEMORY, status); - EXPECT_EQ(nullptr, mdesc3); ucs_mpool_put(mdesc2); ucs_mpool_put(mdesc1); From 9b6186051f44d3ebe0763cd3feb9e4545623eea3 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 27 Aug 2026 19:20:05 +0300 Subject: [PATCH 39/67] UCP/RNDV: more code review fixes + fixed mtype initialization --- src/ucp/rndv/rndv_get.c | 10 +--------- src/ucp/rndv/rndv_mtype.inl | 8 +++++--- src/ucp/rndv/rndv_put.c | 17 +++++++++++++++-- src/ucp/rndv/rndv_rtr.c | 3 +-- 4 files changed, 22 insertions(+), 16 deletions(-) diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index eda8196eb37..bb625d9b741 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -381,7 +381,6 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) } ucp_proto_rndv_mtype_mdesc_release(req); - req->send.rndv.mdesc = NULL; req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; if ((req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_FETCH) && @@ -392,13 +391,6 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) return UCS_OK; } -static void -ucp_proto_rndv_get_mtype_abort(ucp_request_t *req, ucs_status_t status) -{ - ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_GET); - ucp_proto_abort_fatal_not_implemented(req, status); -} - ucp_proto_t ucp_rndv_get_mtype_proto = { .name = "rndv/get/mtype", .desc = NULL, @@ -410,6 +402,6 @@ ucp_proto_t ucp_rndv_get_mtype_proto = { [UCP_PROTO_RNDV_GET_STAGE_FETCH] = ucp_proto_rndv_get_mtype_fetch_progress, [UCP_PROTO_RNDV_GET_STAGE_ATS] = ucp_proto_rndv_ats_progress }, - .abort = ucp_proto_rndv_get_mtype_abort, + .abort = ucp_proto_abort_fatal_not_implemented, .reset = ucp_proto_rndv_get_mtype_reset }; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index c95d7bbde3e..0f41550d9d0 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -54,9 +54,10 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucp_worker_h worker = req->send.ep->worker; ucs_status_t status; - UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); - status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, - &req->send.rndv.mdesc); + req->send.rndv.mdesc = NULL; + status = ucp_rndv_mpool_get(worker, frag_mem_type, + frag_sys_dev, + &req->send.rndv.mdesc); if (status != UCS_ERR_NO_RESOURCE) { return status; } @@ -276,6 +277,7 @@ static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) { ucs_mpool_put_inline(req->send.rndv.mdesc); + req->send.rndv.mdesc = NULL; ucp_proto_rndv_mtype_fc_reschedule_pending(req); } diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index efa1097065e..a9dd5c636a5 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -572,7 +572,10 @@ static void ucp_proto_rndv_put_mtype_completion(uct_completion_t *uct_comp) send.state.uct_comp); ucp_trace_req(req, "rndv_put_mtype_completion"); - ucp_proto_rndv_mtype_mdesc_release(req); + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucp_proto_rndv_mtype_mdesc_release(req); + } + ucp_proto_rndv_put_common_complete(req); } @@ -582,7 +585,10 @@ static void ucp_proto_rndv_put_mtype_frag_completion(uct_completion_t *uct_comp) send.state.uct_comp); ucp_trace_req(req, "rndv_put_mtype_frag_completion"); - ucp_proto_rndv_mtype_mdesc_release(req); + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucp_proto_rndv_mtype_mdesc_release(req); + } + ucp_proto_rndv_ppln_send_frag_complete(req, 1); } @@ -670,7 +676,14 @@ ucp_proto_rndv_put_mtype_query(const ucp_proto_query_params_t *params, static void ucp_proto_rndv_put_mtype_abort(ucp_request_t *req, ucs_status_t status) { + const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_PUT); + if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + ucp_proto_completion_init(&req->send.state.uct_comp, + rpriv->put_comp_cb); + } + ucp_proto_rndv_stub_abort(req, status); } diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 465f9df8105..b159c7f75f6 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -359,7 +359,7 @@ static size_t ucp_proto_rndv_rtr_mtype_pack(void *dest, void *arg) static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_rtr_mtype_complete(ucp_request_t *req, int abort) { - if (!abort || (req->send.rndv.mdesc != NULL)) { + if (!abort || (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { ucp_proto_rndv_mtype_mdesc_release(req); } @@ -393,7 +393,6 @@ static ucs_status_t ucp_proto_rndv_rtr_mtype_reset(ucp_request_t *req) { if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { ucp_proto_rndv_mtype_mdesc_release(req); - req->send.rndv.mdesc = NULL; } return ucp_proto_request_zcopy_id_reset(req); From 9f30c96d05cc765ca1aee358374faa7bf15e2d87 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Sun, 30 Aug 2026 11:55:31 +0300 Subject: [PATCH 40/67] UCP/RNDV: small fix --- src/ucp/core/ucp_worker.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 8a39a7f2bc3..407558e9a7d 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -411,8 +411,8 @@ typedef struct ucp_worker { } counters; struct { - /* Pending queues indexed by ucp_worker_rndv_fc_op_t, ordered by - * descending priority (GET/PUT = 0, RTR = 1) */ + /* Pending queues indexed by UCP_WORKER_RNDV_FC_OP_*, + * ordered by descending priority (GET/PUT = 0, RTR = 1) */ ucs_queue_head_t pending_q[UCP_WORKER_RNDV_FC_OP_LAST]; } rndv_mtype_fc; From 231bdebaf3199806d22838a9e31f9d1be9b2c8d4 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 2 Sep 2026 17:28:25 +0300 Subject: [PATCH 41/67] UCP/RNDV: some code review fixes --- test/gtest/ucp/test_ucp_mmap.cc | 13 ++++++------- test/gtest/ucp/test_ucp_peer_failure.cc | 9 ++++----- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 19c9c5a63c9..08c1e8f48b9 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1058,27 +1058,26 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=16K") { + ucp_worker_h worker = sender().worker(); ucp_worker_mpool_key_t key; ucp_mem_desc_t *mdesc1; ucp_mem_desc_t *mdesc2; ucp_mem_desc_t *mdesc3; - ucp_worker_h worker = sender().worker(); ucs_mpool_ops_t fail_ops; const ucs_mpool_ops_t *orig_ops; - ucs_mpool_t *mpool; ucs_status_t status; - khiter_t khiter; ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); - key.mem_type = UCS_MEMORY_TYPE_HOST; - key.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; - khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); + key.mem_type = UCS_MEMORY_TYPE_HOST; + key.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + const khiter_t khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, + key); ASSERT_NE(kh_end(&worker->mpool_hash), khiter); - mpool = &kh_val(&worker->mpool_hash, khiter); + ucs_mpool_t *mpool = &kh_val(&worker->mpool_hash, khiter); /* The first chunk is exhausted, but the pool still has quota. Make the * next chunk allocation fail and verify it is not reported as quota diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 76d99d0457b..48dcfa65f92 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -866,7 +866,6 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { auto stage = req->send.proto_stage; auto *proto = req->send.proto_config->proto; auto *worker = req->send.ep->worker; - ucs_status_t status; ucs::mock mock; if (proto->name == self->m_proto_name) { @@ -880,8 +879,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } /* Call original proto progress */ - status = self->m_progress_mock.orig_func(&proto->progress[stage], - uct_req); + const ucs_status_t status = self->m_progress_mock.orig_func( + &proto->progress[stage], uct_req); if (self->m_abort_fc_pending && (worker == self->receiver().worker()) && @@ -1022,7 +1021,6 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { void rndv_fc_pending_abort_test() { ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); - std::vector held_mdescs; std::pair result; if (sender_config->key.rma_bw_lanes[0] == UCP_NULL_LANE) { @@ -1033,7 +1031,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { /* Complete one transfer so the CUDA fragment mpool is created. */ smoke_test(true); - held_mdescs = hold_cuda_fragments(receiver().worker()); + std::vector held_mdescs = hold_cuda_fragments( + receiver().worker()); ASSERT_FALSE(held_mdescs.empty()); m_abort_fc_pending = true; From b58d48dfa10062308c94f3cc27b42034cf5fdc31 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 7 Sep 2026 10:24:19 +0300 Subject: [PATCH 42/67] UCP/RNDV: code review --- src/ucp/core/ucp_request.h | 8 ++++++-- src/ucp/rndv/rndv.c | 1 - src/ucp/rndv/rndv_mtype.inl | 8 ++------ 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 33f5460f24e..01ff027c73a 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -329,8 +329,12 @@ struct ucp_request { struct { /* Size to send in ack message */ ssize_t ack_data_size; - /* Element in worker-level pending queue - * for throttled ppln requests */ + /* Element in worker-level pending queue for + * throttled ppln, put/mtype, get/mtype and + * rtr/mtype requests. Sibling union members + * (e.g. put.*) must stay unused until + * UCP_REQUEST_FLAG_PROTO_INITIALIZED is + * set. */ ucs_queue_elem_t queue_elem; } ppln; diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 7cd4bf96444..a65a5010d70 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -25,7 +25,6 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) { ucp_request_t *req = arg; - UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(req->send.ep->worker); ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; ucp_request_send(req); diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 0f41550d9d0..4f65636bf61 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -10,7 +10,6 @@ #include "rndv.h" #include -#include static ucp_ep_h ucp_proto_rndv_mtype_ep(ucp_worker_t *worker, @@ -66,10 +65,10 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, * appropriate pending queue ordered by priority. */ ucp_trace_req(req, "mtype_fc: mpool exhausted, queue %s mem_type %s " - "sys_dev %s", + "sys_dev %u", (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", ucs_memory_type_names[frag_mem_type], - ucs_topo_sys_device_get_name(frag_sys_dev)); + frag_sys_dev); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); ucs_assert(!(req->flags & @@ -203,7 +202,6 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) { ucp_worker_h worker = req->send.ep->worker; - UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); ucs_assert(fc_op < UCP_WORKER_RNDV_FC_OP_LAST); ucs_assert(!ucs_test_all_flags(req->flags, UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | @@ -242,8 +240,6 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) ucp_request_t *pending_req; unsigned q_index; - UCP_WORKER_THREAD_CS_CHECK_IS_BLOCKED_CONDITIONAL(worker); - /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { From 586d57d255a4722ae2954be89bb0c7ee6e3a5bb8 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 8 Sep 2026 17:09:30 +0300 Subject: [PATCH 43/67] UCP/RNDV: Apply fragment mpool limit only to proto v2 --- src/ucp/core/ucp_context.c | 13 +++++++------ src/ucp/rndv/rndv.c | 6 ++++-- test/gtest/ucp/test_ucp_mmap.cc | 24 +++++++++++++++++++++++- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index d28903e642d..53356acb586 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -405,12 +405,13 @@ static ucs_config_field_t ucp_context_config_table[] = { ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, {"RNDV_FRAG_WORKER_MAX_MEM", "auto", - "Maximum amount of memory a worker may use for rendezvous staging\n" - "fragments, or \"inf\" for no limit. The same limit is applied\n" - "independently to each fragment memory type, and per device when\n" - "available. It is converted to a fragment count using RNDV_FRAG_SIZE and\n" - "RNDV_FRAG_ALLOC_COUNT. When the limit is reached, further fragment\n" - "requests are queued until fragments are released", + "Maximum amount of memory a worker may use for proto v2 rendezvous\n" + "staging fragments, or \"inf\" for no limit. This setting has no effect\n" + "when PROTO_ENABLE=n. The same limit is applied independently to each\n" + "fragment memory type, and per device when available. It is converted to\n" + "a fragment count using RNDV_FRAG_SIZE and RNDV_FRAG_ALLOC_COUNT. When\n" + "the limit is reached, further fragment requests are queued until\n" + "fragments are released", ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index a65a5010d70..ec4bc79994a 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1184,8 +1184,10 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.elem_size = sizeof(ucp_mem_desc_t); mp_params.alignment = 1; mp_params.elems_per_chunk = num_frags; - mp_params.max_elems = ucp_proto_rndv_frag_max_elems( - worker->context, key.mem_type); + mp_params.max_elems = worker->context->config.ext.proto_enable ? + ucp_proto_rndv_frag_max_elems( + worker->context, key.mem_type) : + UINT_MAX; mp_params.ops = &ucp_frag_mpool_ops; mp_params.name = "ucp_rndv_frags"; status = ucs_mpool_init(&mp_params, mpool); diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 08c1e8f48b9..59db33021a9 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1034,7 +1034,8 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) } UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_exhausted, - "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", + "PROTO_ENABLE=y", "RNDV_FRAG_SIZE=host:4K", + "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=8K") { ucp_mem_desc_t *mdesc1; @@ -1054,6 +1055,27 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_exhausted, ucs_mpool_put(mdesc1); } +UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_disabled_with_proto_v1, + "PROTO_ENABLE=n", "RNDV_FRAG_SIZE=host:4K", + "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=8K") +{ + ucp_mem_desc_t *mdesc1; + ucp_mem_desc_t *mdesc2; + ucp_mem_desc_t *mdesc3; + ucp_worker_h worker = sender().worker(); + + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); + ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3)); + + ucs_mpool_put(mdesc3); + ucs_mpool_put(mdesc2); + ucs_mpool_put(mdesc1); +} + UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=16K") From c356ccc990a559fe6ff863a413a502edf5f1e8d9 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 15 Sep 2026 18:35:53 +0300 Subject: [PATCH 44/67] UCP/RNDV: Abort throttled mtype requests on EP purge --- src/ucp/core/ucp_ep.c | 7 ++- src/ucp/core/ucp_ep.h | 2 + src/ucp/core/ucp_request.h | 22 ++++--- src/ucp/rndv/proto_rndv.c | 50 +++++++++++++++ src/ucp/rndv/rndv.c | 10 --- src/ucp/rndv/rndv.h | 11 ++++ src/ucp/rndv/rndv_get.c | 31 +++++++++- src/ucp/rndv/rndv_mtype.inl | 45 +++++++------- src/ucp/rndv/rndv_rtr.c | 2 + test/gtest/common/test_obj_size.cc | 2 +- test/gtest/ucp/test_ucp_peer_failure.cc | 82 +++++++++++-------------- 11 files changed, 175 insertions(+), 89 deletions(-) diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index a822b85d580..4370c1d9b18 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -280,6 +280,7 @@ static ucp_ep_h ucp_ep_allocate(ucp_worker_h worker, const char *peer_name) memset(&ep->ext->ep_match, 0, sizeof(ep->ext->ep_match)); ucs_hlist_head_init(&ep->ext->proto_reqs); + ucs_hlist_head_init(&ep->ext->rndv_mtype_fc_reqs); for (lane = 0; lane < UCP_MAX_FAST_PATH_LANES; ++lane) { ucp_ep_set_lane(ep, lane, NULL); @@ -524,7 +525,8 @@ static int ucp_ep_remove_filter(const ucs_callbackq_elem_t *elem, void *arg) ucp_listener_accept_cb_remove_filter(elem, arg) || ucp_ep_local_disconnect_progress_remove_filter(elem, arg) || ucp_ep_set_failed_remove_filter(elem, arg) || - ucp_ep_wireup_eps_progress_filter(elem, arg)) { + ucp_ep_wireup_eps_progress_filter(elem, arg) || + ucp_proto_rndv_mtype_fc_reschedule_filter(elem, arg)) { return 1; } @@ -542,6 +544,7 @@ void ucp_ep_destroy_base(ucp_ep_h ep) ucp_ep_refcount_assert(ep, discard, ==, 0); ucp_ep_refcount_assert(ep, probe, ==, 0); ucs_assert(ucs_hlist_is_empty(&ep->ext->proto_reqs)); + ucs_assert(ucs_hlist_is_empty(&ep->ext->rndv_mtype_fc_reqs)); if (!(ep->flags & UCP_EP_FLAG_INTERNAL)) { ucs_assert(worker->num_all_eps > 0); @@ -4506,6 +4509,8 @@ void ucp_ep_reqs_purge(ucp_ep_h ucp_ep, ucs_status_t status) ucp_ep_flush_state_t *flush_state; ucp_request_t *req; + ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep, status); + while (!ucs_hlist_is_empty(proto_reqs)) { req = ucs_hlist_head_elem(proto_reqs, ucp_request_t, send.list); if (ucp_ep->worker->context->config.ext.proto_enable) { diff --git a/src/ucp/core/ucp_ep.h b/src/ucp/core/ucp_ep.h index 6df732b7a72..32fe613973a 100644 --- a/src/ucp/core/ucp_ep.h +++ b/src/ucp/core/ucp_ep.h @@ -539,6 +539,8 @@ typedef struct ucp_ep_ext { used by 2-stage ppln rndv proto */ /* List of requests which are waiting for remote completion */ ucs_hlist_head_t proto_reqs; + /* List of rendezvous mtype requests waiting for a fragment */ + ucs_hlist_head_t rndv_mtype_fc_reqs; #if UCS_ENABLE_ASSERT ucs_time_t ka_last_round; /* Time of last KA round done */ #endif diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 01ff027c73a..8803e7e9653 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -74,6 +74,10 @@ enum { UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED = UCS_BIT(31) }; +#define UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK \ + (UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | \ + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED) + /** * Protocols enumerator to work with send request state @@ -328,16 +332,18 @@ struct ucp_request { /* Used by rndv/send/ppln and rndv/recv/ppln */ struct { /* Size to send in ack message */ - ssize_t ack_data_size; - /* Element in worker-level pending queue for - * throttled ppln, put/mtype, get/mtype and - * rtr/mtype requests. Sibling union members - * (e.g. put.*) must stay unused until - * UCP_REQUEST_FLAG_PROTO_INITIALIZED is - * set. */ - ucs_queue_elem_t queue_elem; + ssize_t ack_data_size; } ppln; + /* Used by throttled rndv mtype requests before + * UCP_REQUEST_FLAG_PROTO_INITIALIZED is set. */ + struct { + /* Element in worker-level pending queue */ + ucs_queue_elem_t queue_elem; + /* Element in per-EP list */ + ucs_hlist_link_t ep_list; + } fc; + /* Used by rndv/rkey_ptr */ struct { /* Element in queue for segmented RKEY ptr */ diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 8667736375d..68757655d9f 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -16,6 +16,56 @@ #include +void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req) +{ + ucp_ep_h ep = req->send.ep; + + ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK); + ucs_hlist_del(&ep->ext->rndv_mtype_fc_reqs, + &req->send.rndv.fc.ep_list); + req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK; +} + +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +{ + ucp_request_t *req = arg; + + ucs_assert((req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK) == + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); + ucp_proto_rndv_mtype_fc_leave(req); + ucp_request_send(req); + return 1; +} + +int ucp_proto_rndv_mtype_fc_reschedule_filter( + const ucs_callbackq_elem_t *elem, void *arg) +{ + if (elem->cb != ucp_proto_rndv_mtype_fc_reschedule_cb) { + return 0; + } + + ucs_assertv(0, "ep %p still has mtype FC reschedule callback for req %p", + arg, elem->arg); + ucs_error("ep %p still has mtype FC reschedule callback for req %p", arg, + elem->arg); + return 1; +} + +void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) +{ + ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; + ucp_request_t *req; + + while (!ucs_hlist_is_empty(fc_reqs)) { + req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, + send.rndv.fc.ep_list); + ucp_proto_request_abort(req, status); + ucs_assert(ucs_hlist_is_empty(fc_reqs) || + (ucs_hlist_head_elem(fc_reqs, ucp_request_t, + send.rndv.fc.ep_list) != req)); + } +} + static void ucp_proto_rndv_ctrl_get_md_map(const ucp_proto_rndv_ctrl_init_params_t *params, ucp_md_map_t *md_map, diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index ec4bc79994a..53a0b870381 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -21,16 +21,6 @@ #include #include -unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) -{ - ucp_request_t *req = arg; - - ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); - req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; - ucp_request_send(req); - return 1; -} - static UCS_F_ALWAYS_INLINE int ucp_rndv_frag_mem_type(ucp_context_t *context) { diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index b62bca378c1..a4f2e90547a 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -120,6 +120,17 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg); +int ucp_proto_rndv_mtype_fc_reschedule_filter( + const ucs_callbackq_elem_t *elem, void *arg); + +void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req); + +/** + * Abort every request of an endpoint that is waiting for an rndv mtype + * fragment. + */ +void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status); + void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, const void *rkey_buf); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index bb625d9b741..0dbc26ed941 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -376,6 +376,8 @@ ucp_proto_rndv_get_mtype_query(const ucp_proto_query_params_t *params, static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) { + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_GET); + if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { return UCS_OK; } @@ -391,6 +393,33 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) return UCS_OK; } +static void +ucp_proto_rndv_get_mtype_abort(ucp_request_t *req, ucs_status_t status) +{ + ucp_request_t *super_req; + + if (!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK)) { + ucp_proto_abort_fatal_not_implemented(req, status); + return; + } + + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_GET); + if (ucp_proto_rndv_request_is_ppln_frag(req)) { + super_req = ucp_request_get_super(req); + super_req->status = status; + + /* The pipeline has a top-level receive request that also needs to be + * completed with the error. */ + ucp_request_get_super(super_req)->status = status; + ucp_proto_rndv_ppln_recv_frag_complete(req, 0, 1); + } else { + if (req->send.rndv.rkey != NULL) { + ucp_proto_rndv_rkey_destroy(req); + } + ucp_proto_rndv_recv_complete_status(req, status); + } +} + ucp_proto_t ucp_rndv_get_mtype_proto = { .name = "rndv/get/mtype", .desc = NULL, @@ -402,6 +431,6 @@ ucp_proto_t ucp_rndv_get_mtype_proto = { [UCP_PROTO_RNDV_GET_STAGE_FETCH] = ucp_proto_rndv_get_mtype_fetch_progress, [UCP_PROTO_RNDV_GET_STAGE_ATS] = ucp_proto_rndv_ats_progress }, - .abort = ucp_proto_abort_fatal_not_implemented, + .abort = ucp_proto_rndv_get_mtype_abort, .reset = ucp_proto_rndv_get_mtype_reset }; diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 4f65636bf61..996ede03339 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -50,7 +50,8 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_sys_device_t frag_sys_dev, unsigned fc_op) { - ucp_worker_h worker = req->send.ep->worker; + ucp_ep_h ep = req->send.ep; + ucp_worker_h worker = ep->worker; ucs_status_t status; req->send.rndv.mdesc = NULL; @@ -71,12 +72,12 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, frag_sys_dev); UCS_STATS_UPDATE_COUNTER(worker->stats, UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); - ucs_assert(!(req->flags & - (UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED))); + ucs_assert(!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK)); req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], - &req->send.rndv.ppln.queue_elem); + &req->send.rndv.fc.queue_elem); + ucs_hlist_add_tail(&ep->ext->rndv_mtype_fc_reqs, + &req->send.rndv.fc.ep_list); return UCS_ERR_NO_RESOURCE; } @@ -200,27 +201,26 @@ ucp_proto_rndv_mtype_fc_reschedule_pred(const ucs_callbackq_elem_t *elem, static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) { - ucp_worker_h worker = req->send.ep->worker; + ucp_ep_h ep = req->send.ep; ucs_assert(fc_op < UCP_WORKER_RNDV_FC_OP_LAST); ucs_assert(!ucs_test_all_flags(req->flags, - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED | - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED)); + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK)); if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED) { ucp_trace_req(req, "mtype_fc: remove aborted request from queue"); - ucs_queue_remove(&worker->rndv_mtype_fc.pending_q[fc_op], - &req->send.rndv.ppln.queue_elem); - req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; - } - - if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED) { + ucs_queue_remove(&ep->worker->rndv_mtype_fc.pending_q[fc_op], + &req->send.rndv.fc.queue_elem); + } else if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED) { ucp_trace_req(req, "mtype_fc: remove aborted reschedule callback"); - ucs_callbackq_remove_oneshot( - &worker->uct->progress_q, req, - ucp_proto_rndv_mtype_fc_reschedule_pred, req); - req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; + ucs_callbackq_remove_oneshot(&ep->worker->uct->progress_q, ep, + ucp_proto_rndv_mtype_fc_reschedule_pred, + req); + } else { + return; } + + ucp_proto_rndv_mtype_fc_leave(req); } /** @@ -248,16 +248,15 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) pending_req = ucs_queue_pull_elem_non_empty( &worker->rndv_mtype_fc.pending_q[q_index], ucp_request_t, - send.rndv.ppln.queue_elem); - ucs_assert(pending_req->flags & + send.rndv.fc.queue_elem); + ucs_assert((pending_req->flags & + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK) == UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED); - ucs_assert(!(pending_req->flags & - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED)); pending_req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; ucp_trace_req(pending_req, "mtype_fc: dequeue %s", (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); - ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req, + ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req->send.ep, ucp_proto_rndv_mtype_fc_reschedule_cb, pending_req); return; diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index b159c7f75f6..53af8f418e4 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -391,6 +391,8 @@ ucp_proto_rndv_rtr_mtype_abort(ucp_request_t *req, ucs_status_t status) static ucs_status_t ucp_proto_rndv_rtr_mtype_reset(ucp_request_t *req) { + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_RTR); + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { ucp_proto_rndv_mtype_mdesc_release(req); } diff --git a/test/gtest/common/test_obj_size.cc b/test/gtest/common/test_obj_size.cc index 076cf38ab1c..d1f070f591d 100644 --- a/test/gtest/common/test_obj_size.cc +++ b/test/gtest/common/test_obj_size.cc @@ -50,7 +50,7 @@ UCS_TEST_F(test_obj_size, size) { UCS_TEST_SKIP_R("Assert enabled"); #else EXPECTED_SIZE(ucp_ep_t, 64); - EXPECTED_SIZE(ucp_ep_ext_t, 216); + EXPECTED_SIZE(ucp_ep_ext_t, 224); #if ENABLE_PARAMS_CHECK EXPECTED_SIZE(ucp_rkey_t, 24 + sizeof(ucp_ep_h)); #else diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 48dcfa65f92..6a328445e90 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -817,28 +817,13 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } } - static unsigned abort_fc_pending_cb(void *arg) + static unsigned close_fc_pending_ep_cb(void *arg) { auto *test = static_cast(arg); - ucp_worker_h worker = test->receiver().worker(); - std::vector pending_reqs; - ucp_request_t *req; - /* Abort may remove the current request, so snapshot the queue first. */ - for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; ++i) { - auto *queue = &worker->rndv_mtype_fc.pending_q[i]; - - ucs_queue_for_each(req, queue, send.rndv.ppln.queue_elem) { - pending_reqs.push_back(req); - } - } - - for (auto *pending_req : pending_reqs) { - ucp_proto_request_abort(pending_req, UCS_ERR_CONNECTION_RESET); - } - - /* Complete the top-level send after aborting receiver fragments. */ - test->sender().close_all_eps(*test, 0, UCP_EP_CLOSE_FLAG_FORCE); + /* Endpoint purge must abort requests waiting for a staging fragment. */ + test->m_fc_pending_entity->close_all_eps(*test, 0, + UCP_EP_CLOSE_FLAG_FORCE); test->m_is_peer_closed = true; return 1; } @@ -868,7 +853,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { auto *worker = req->send.ep->worker; ucs::mock mock; - if (proto->name == self->m_proto_name) { + if (!self->m_proto_name.empty() && + (proto->name == self->m_proto_name)) { if (self->m_replace_ops) { mock_rndv_ops(req->send.ep, mock); } @@ -882,27 +868,21 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { const ucs_status_t status = self->m_progress_mock.orig_func( &proto->progress[stage], uct_req); - if (self->m_abort_fc_pending && - (worker == self->receiver().worker()) && - self->has_fc_pending(worker)) { - self->m_abort_fc_pending = false; - /* Abort after the current protocol-progress call returns. */ + if (self->m_close_fc_pending_ep && + (worker == self->m_fc_pending_entity->worker()) && + self->has_fc_pending(worker, self->m_fc_op)) { + self->m_close_fc_pending_ep = false; + /* Close after the current protocol-progress call returns. */ ucs_callbackq_add_oneshot(&worker->uct->progress_q, self, - abort_fc_pending_cb, self); + close_fc_pending_ep_cb, self); } return status; } - static bool has_fc_pending(ucp_worker_h worker) + static bool has_fc_pending(ucp_worker_h worker, unsigned fc_op) { - for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; ++i) { - if (!ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) { - return true; - } - } - - return false; + return !ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[fc_op]); } static std::vector @@ -943,8 +923,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { for (auto *mdesc : held) { /* - * Return the descriptor without invoking FC rescheduling. Aborted - * requests must already have been removed from pending_q. + * Return the descriptor without invoking FC rescheduling. Endpoint + * purge must already have removed requests from pending_q. */ ucs_mpool_put_inline(mdesc); } @@ -1018,7 +998,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { ASSERT_TRUE(m_is_peer_closed); } - void rndv_fc_pending_abort_test() + void rndv_fc_pending_abort_test(entity &fc_entity, unsigned fc_op) { ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); std::pair result; @@ -1031,11 +1011,13 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { /* Complete one transfer so the CUDA fragment mpool is created. */ smoke_test(true); - std::vector held_mdescs = hold_cuda_fragments( - receiver().worker()); + std::vector held_mdescs = + hold_cuda_fragments(fc_entity.worker()); ASSERT_FALSE(held_mdescs.empty()); - m_abort_fc_pending = true; + m_fc_pending_entity = &fc_entity; + m_fc_op = fc_op; + m_close_fc_pending_ep = true; setup_progress_mock(sender().worker(), m_progress_mock); setup_progress_mock(receiver().worker(), m_progress_mock); @@ -1045,23 +1027,24 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { result = smoke_test(true); } - release_cuda_fragments(receiver().worker(), held_mdescs); + release_cuda_fragments(fc_entity.worker(), held_mdescs); m_progress_mock.cleanup(); EXPECT_TRUE(UCS_STATUS_IS_ERR(result.first)); EXPECT_TRUE(UCS_STATUS_IS_ERR(result.second)); EXPECT_TRUE(m_is_peer_closed); - EXPECT_FALSE(has_fc_pending(sender().worker())); - EXPECT_FALSE(has_fc_pending(receiver().worker())); + EXPECT_FALSE(has_fc_pending(fc_entity.worker(), fc_op)); } ucs::mock m_progress_mock; bool m_is_peer_closed{false}; - bool m_abort_fc_pending{false}; + bool m_close_fc_pending_ep{false}; std::string m_proto_name{}; /* Protocol stage during which data transfer happens */ uint8_t m_proto_xfer_stage{}; entity *m_peer_to_close{nullptr}; + entity *m_fc_pending_entity{nullptr}; + unsigned m_fc_op{UCP_WORKER_RNDV_FC_OP_LAST}; /* Even if we close peer EP with the force flag, the next proto progress call probably would return UCS_OK. This option enables emulation of certain progress call failure. */ @@ -1163,7 +1146,16 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_fc_pending, "RNDV_FRAG_WORKER_MAX_MEM=8K", "RNDV_FRAG_MEM_TYPE=cuda") { - rndv_fc_pending_abort_test(); + rndv_fc_pending_abort_test(receiver(), UCP_WORKER_RNDV_FC_OP_RTR); +} + +UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_fc_pending, + "RNDV_FRAG_SIZE=host:8K,cuda:8K", + "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=8K", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + rndv_fc_pending_abort_test(sender(), UCP_WORKER_RNDV_FC_OP_PUT); } UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); From af47dde6eb5c983702b521757f969af3f40dbbd0 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 16 Sep 2026 10:35:50 +0300 Subject: [PATCH 45/67] UCP/RNDV: Pass mtype FC wakeup to next waiter on abort --- src/ucp/rndv/proto_rndv.c | 4 +- src/ucp/rndv/rndv_mtype.inl | 111 ++++++++++++++++++++++-------------- 2 files changed, 72 insertions(+), 43 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 68757655d9f..ac7bc802c11 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -32,7 +32,9 @@ unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) ucs_assert((req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK) == UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); - ucp_proto_rndv_mtype_fc_leave(req); + /* Keep the request in RESCHED state, so that if it is aborted before + * ucp_proto_rndv_mtype_request_init() retries the allocation, the wakeup + * is passed to the next waiter rather than dropped. */ ucp_request_send(req); return 1; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 996ede03339..46aa26e1087 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -44,6 +44,45 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params, return UCS_OK; } +/** + * Reschedule a pending throttled request after a fragment is released back to + * the mpool. Dequeue priority: PUT/GET > RTR. + * + * Priority rationale: + * PUT/GET - Completing a PUT or GET frees a staging buffer, reducing memory + * pressure. PUT also unblocks the remote side. + * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total + * memory pressure. + */ +static UCS_F_ALWAYS_INLINE void +ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_worker_h worker) +{ + ucp_request_t *pending_req; + unsigned q_index; + + /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ + for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { + if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { + continue; + } + + pending_req = ucs_queue_pull_elem_non_empty( + &worker->rndv_mtype_fc.pending_q[q_index], ucp_request_t, + send.rndv.fc.queue_elem); + ucs_assert((pending_req->flags & + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK) == + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED); + pending_req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; + pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; + ucp_trace_req(pending_req, "mtype_fc: dequeue %s", + (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); + ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req->send.ep, + ucp_proto_rndv_mtype_fc_reschedule_cb, + pending_req); + return; + } +} + static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, @@ -52,13 +91,26 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, { ucp_ep_h ep = req->send.ep; ucp_worker_h worker = ep->worker; + /* A rescheduled request owns the wakeup of a fragment which was released + * back to the mpool, and must pass it on if it does not consume one. */ + int owns_wakeup = req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; ucs_status_t status; + ucs_assert(!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED)); + if (owns_wakeup) { + ucp_proto_rndv_mtype_fc_leave(req); + } + req->send.rndv.mdesc = NULL; status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, &req->send.rndv.mdesc); if (status != UCS_ERR_NO_RESOURCE) { + if ((status != UCS_OK) && owns_wakeup) { + /* The caller aborts the request, so the wakeup would be lost. */ + ucp_proto_rndv_mtype_fc_reschedule_pending(worker); + } + return status; } @@ -201,7 +253,9 @@ ucp_proto_rndv_mtype_fc_reschedule_pred(const ucs_callbackq_elem_t *elem, static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) { - ucp_ep_h ep = req->send.ep; + ucp_ep_h ep = req->send.ep; + ucp_worker_h worker = ep->worker; + int owns_wakeup; ucs_assert(fc_op < UCP_WORKER_RNDV_FC_OP_LAST); ucs_assert(!ucs_test_all_flags(req->flags, @@ -209,57 +263,28 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED) { ucp_trace_req(req, "mtype_fc: remove aborted request from queue"); - ucs_queue_remove(&ep->worker->rndv_mtype_fc.pending_q[fc_op], + ucs_queue_remove(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.fc.queue_elem); + owns_wakeup = 0; } else if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED) { ucp_trace_req(req, "mtype_fc: remove aborted reschedule callback"); - ucs_callbackq_remove_oneshot(&ep->worker->uct->progress_q, ep, + /* No-op if the callback was already dispatched and the retry is in + * progress. */ + ucs_callbackq_remove_oneshot(&worker->uct->progress_q, ep, ucp_proto_rndv_mtype_fc_reschedule_pred, req); + owns_wakeup = 1; } else { return; } ucp_proto_rndv_mtype_fc_leave(req); -} -/** - * Reschedule a pending throttled request after a fragment is released back to - * the mpool. Dequeue priority: PUT/GET > RTR. - * - * Priority rationale: - * PUT/GET - Completing a PUT or GET frees a staging buffer, reducing memory - * pressure. PUT also unblocks the remote side. - * RTR - Scheduling RTR triggers a remote PUT allocation, increasing total - * memory pressure. - */ -static UCS_F_ALWAYS_INLINE void -ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) -{ - ucp_worker_h worker = req->send.ep->worker; - ucp_request_t *pending_req; - unsigned q_index; - - /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ - for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { - if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { - continue; - } - - pending_req = ucs_queue_pull_elem_non_empty( - &worker->rndv_mtype_fc.pending_q[q_index], ucp_request_t, - send.rndv.fc.queue_elem); - ucs_assert((pending_req->flags & - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK) == - UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED); - pending_req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; - pending_req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; - ucp_trace_req(pending_req, "mtype_fc: dequeue %s", - (q_index == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get"); - ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req->send.ep, - ucp_proto_rndv_mtype_fc_reschedule_cb, - pending_req); - return; + if (owns_wakeup) { + /* A fragment was already released on behalf of this request, so pass + * the wakeup to the next waiter rather than dropping it, otherwise the + * remaining waiters could stall while the fragment is free. */ + ucp_proto_rndv_mtype_fc_reschedule_pending(worker); } } @@ -271,9 +296,11 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_request_t *req) static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) { + ucp_worker_h worker = req->send.ep->worker; + ucs_mpool_put_inline(req->send.rndv.mdesc); req->send.rndv.mdesc = NULL; - ucp_proto_rndv_mtype_fc_reschedule_pending(req); + ucp_proto_rndv_mtype_fc_reschedule_pending(worker); } static UCS_F_ALWAYS_INLINE ucs_status_t From 3176441bd2744e82a789dd27ba0893bd2c06a59a Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 16 Sep 2026 10:42:23 +0300 Subject: [PATCH 46/67] UCP/RNDV: small fixes --- src/ucp/core/ucp_context.c | 17 +++++++++-------- src/ucp/rndv/proto_rndv.h | 2 +- test/gtest/ucp/test_ucp_am.cc | 2 +- test/gtest/ucp/test_ucp_peer_failure.cc | 5 ++--- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 53356acb586..ac03e40e575 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -405,14 +405,15 @@ static ucs_config_field_t ucp_context_config_table[] = { ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL}, {"RNDV_FRAG_WORKER_MAX_MEM", "auto", - "Maximum amount of memory a worker may use for proto v2 rendezvous\n" - "staging fragments, or \"inf\" for no limit. This setting has no effect\n" - "when PROTO_ENABLE=n. The same limit is applied independently to each\n" - "fragment memory type, and per device when available. It is converted to\n" - "a fragment count using RNDV_FRAG_SIZE and RNDV_FRAG_ALLOC_COUNT. When\n" - "the limit is reached, further fragment requests are queued until\n" - "fragments are released", - ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, + "Maximum amount of memory a worker may use for rendezvous staging\n" + "fragments. Both \"auto\" and \"inf\" mean no limit. This setting has no\n" + "effect when PROTO_ENABLE=n. The same limit is applied independently to\n" + "each fragment memory type, and per device when available. It is\n" + "converted to a fragment count using RNDV_FRAG_SIZE and\n" + "RNDV_FRAG_ALLOC_COUNT. When the limit is reached, further fragment\n" + "requests are queued until fragments are released", + ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), + UCS_CONFIG_TYPE_MEMUNITS}, {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 803ff915e2e..7811c385769 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -196,7 +196,7 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); /* Convert UCX_RNDV_FRAG_WORKER_MAX_MEM to mpool max_elems (frag count), or * UINT_MAX if unlimited. */ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, - ucs_memory_type_t frag_mem_type); + ucs_memory_type_t frag_mem_type); ucs_status_t diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 7776581097f..9726129e425 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2394,7 +2394,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { static void check_pending_queues_empty(const entity &e) { - const ucp_worker_h worker = e.worker(); + ucp_worker_h worker = e.worker(); for (unsigned i = 0; i < UCP_WORKER_RNDV_FC_OP_LAST; i++) { EXPECT_TRUE(ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[i])) diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 6a328445e90..bce9f41cd49 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -890,11 +890,10 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { { std::vector held; ucp_mem_desc_t *mdesc; - khiter_t iter; UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - for (iter = kh_begin(&worker->mpool_hash); + for (khiter_t iter = kh_begin(&worker->mpool_hash); iter != kh_end(&worker->mpool_hash); ++iter) { if (!kh_exist(&worker->mpool_hash, iter) || (kh_key(&worker->mpool_hash, iter).mem_type != @@ -1000,7 +999,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { void rndv_fc_pending_abort_test(entity &fc_entity, unsigned fc_op) { - ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); + const ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); std::pair result; if (sender_config->key.rma_bw_lanes[0] == UCP_NULL_LANE) { From 7ab57303e856b72ddcde6b1c6416c00ea868afdd Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 16 Sep 2026 17:52:39 +0300 Subject: [PATCH 47/67] UCP/RNDV: Skip proto-v1 test cases --- test/gtest/ucp/test_ucp_mmap.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 59db33021a9..a2758890f93 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1033,10 +1033,10 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) } } -UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_exhausted, - "PROTO_ENABLE=y", "RNDV_FRAG_SIZE=host:4K", - "RNDV_FRAG_ALLOC_COUNT=host:2", - "RNDV_FRAG_WORKER_MAX_MEM=8K") +UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_quota_exhausted, + (get_variant_value() == VARIANT_PROTO_DISABLE), + "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", + "RNDV_FRAG_WORKER_MAX_MEM=8K") { ucp_mem_desc_t *mdesc1; ucp_mem_desc_t *mdesc2; From e1be1df4feea4f317dff8208200f0d35143c80ca Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 16 Sep 2026 19:29:02 +0300 Subject: [PATCH 48/67] UCP/RNDV: Replay throttled mtype requests on EP reconfiguration --- src/ucp/rndv/proto_rndv.c | 26 +++++++++++++++++++++++++ src/ucp/rndv/rndv.h | 13 ++++++++++--- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_put.c | 14 ++++++++++++- src/ucp/wireup/wireup.c | 4 ++++ test/gtest/ucp/test_ucp_peer_failure.cc | 6 +++--- 6 files changed, 57 insertions(+), 8 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 8fc61598dfa..cc0355bf38d 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -68,6 +68,32 @@ void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) } } +void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, + ucs_queue_head_t *replay_queue) +{ + ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; + ucp_request_t *req; + ucs_status_t status; + + /* The replay does not always restart (and reset) the request, so clear + * its flow-control state here. This may wake up another request of this + * ep, which is then extracted on a later iteration. */ + while (!ucs_hlist_is_empty(fc_reqs)) { + req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, + send.rndv.fc.ep_list); + ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); + ucp_trace_req(req, "mtype_fc: extract for replay"); + + status = req->send.proto_config->proto->reset(req); + ucs_assertv_always(status == UCS_OK, "req %p, failed to reset: %s", + req, ucs_status_string(status)); + ucs_assert(ucs_hlist_is_empty(fc_reqs) || + (ucs_hlist_head_elem(fc_reqs, ucp_request_t, + send.rndv.fc.ep_list) != req)); + ucs_queue_push(replay_queue, (ucs_queue_elem_t*)&req->send.uct.priv); + } +} + static int ucp_proto_rndv_ctrl_skip_inter_node_md( const ucp_proto_common_init_params_t *params, diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index a4f2e90547a..438500f0a8d 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -10,6 +10,7 @@ #include #include #include +#include typedef enum { @@ -125,11 +126,17 @@ int ucp_proto_rndv_mtype_fc_reschedule_filter( void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req); +/* Abort every request of an endpoint that is waiting for a rndv mtype + * fragment */ +void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status); + /** - * Abort every request of an endpoint that is waiting for an rndv mtype - * fragment. + * Move every request of an endpoint that is waiting for a rndv mtype + * fragment to @a replay_queue, so it is replayed after the endpoint is + * reconfigured. */ -void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status); +void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, + ucs_queue_head_t *replay_queue); void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 0dbc26ed941..b90c9e8a9d5 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -383,7 +383,7 @@ static ucs_status_t ucp_proto_rndv_get_mtype_reset(ucp_request_t *req) } ucp_proto_rndv_mtype_mdesc_release(req); - req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; + req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; if ((req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_FETCH) && (req->send.proto_stage != UCP_PROTO_RNDV_GET_STAGE_ATS)) { diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index a9dd5c636a5..452acc869a7 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -687,6 +687,18 @@ ucp_proto_rndv_put_mtype_abort(ucp_request_t *req, ucs_status_t status) ucp_proto_rndv_stub_abort(req, status); } +static ucs_status_t ucp_proto_rndv_put_mtype_reset(ucp_request_t *req) +{ + ucp_proto_rndv_mtype_fc_cancel(req, UCP_WORKER_RNDV_FC_OP_PUT); + + /* Only a request which is still waiting for a fragment can be reset */ + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucp_proto_reset_fatal_not_implemented(req); + } + + return UCS_OK; +} + ucp_proto_t ucp_rndv_put_mtype_proto = { .name = "rndv/put/mtype", .desc = NULL, @@ -702,5 +714,5 @@ ucp_proto_t ucp_rndv_put_mtype_proto = { [UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP] = ucp_proto_rndv_put_common_fenced_atp_progress, }, .abort = ucp_proto_rndv_put_mtype_abort, - .reset = (ucp_request_reset_func_t)ucp_proto_reset_fatal_not_implemented + .reset = ucp_proto_rndv_put_mtype_reset }; diff --git a/src/ucp/wireup/wireup.c b/src/ucp/wireup/wireup.c index f5d527953a8..a7d08b94a91 100644 --- a/src/ucp/wireup/wireup.c +++ b/src/ucp/wireup/wireup.c @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -2240,6 +2241,9 @@ ucp_wireup_gather_pending_requests(ucp_ep_h ep, } } + /* Requests throttled while waiting for a rndv mtype fragment */ + ucp_proto_rndv_mtype_fc_ep_extract(ep, replay_pending_queue); + if (ep->cfg_index == UCP_WORKER_CFG_INDEX_NULL) { return; } diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index bddfcdc712b..81715ab590b 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -1014,9 +1014,9 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { hold_cuda_fragments(fc_entity.worker()); ASSERT_FALSE(held_mdescs.empty()); - m_fc_pending_entity = &fc_entity; - m_fc_op = fc_op; - m_close_fc_pending_ep = true; + m_fc_pending_entity = &fc_entity; + m_fc_op = fc_op; + m_close_fc_pending_ep = true; setup_progress_mock(sender().worker(), m_progress_mock); setup_progress_mock(receiver().worker(), m_progress_mock); From b3ea1eed1abcb828f3cc1c5b74517bd24bc2c36b Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 17 Sep 2026 15:09:11 +0300 Subject: [PATCH 49/67] UCP/RNDV: Keep request size unchanged by moving mtype FC EP link --- src/ucp/core/ucp_request.h | 7 +++++-- src/ucp/rndv/proto_rndv.c | 10 +++++----- src/ucp/rndv/rndv_mtype.inl | 2 +- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 01e0944771b..c24d8940400 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -203,6 +203,11 @@ struct ucp_request { * Used by rkey_ptr to track copied data size */ ssize_t completed_size; + + /* Element in per-EP list of rndv mtype requests throttled + * while waiting for a fragment. Used only before + * UCP_REQUEST_FLAG_PROTO_INITIALIZED is set. */ + ucs_hlist_link_t rndv_fc_ep_list; }; } state; @@ -345,8 +350,6 @@ struct ucp_request { struct { /* Element in worker-level pending queue */ ucs_queue_elem_t queue_elem; - /* Element in per-EP list */ - ucs_hlist_link_t ep_list; } fc; /* Used by rndv/rkey_ptr */ diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index cc0355bf38d..ddc4f355f77 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -22,7 +22,7 @@ void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req) ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK); ucs_hlist_del(&ep->ext->rndv_mtype_fc_reqs, - &req->send.rndv.fc.ep_list); + &req->send.state.rndv_fc_ep_list); req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK; } @@ -60,11 +60,11 @@ void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) while (!ucs_hlist_is_empty(fc_reqs)) { req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.rndv.fc.ep_list); + send.state.rndv_fc_ep_list); ucp_proto_request_abort(req, status); ucs_assert(ucs_hlist_is_empty(fc_reqs) || (ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.rndv.fc.ep_list) != req)); + send.state.rndv_fc_ep_list) != req)); } } @@ -80,7 +80,7 @@ void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, * ep, which is then extracted on a later iteration. */ while (!ucs_hlist_is_empty(fc_reqs)) { req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.rndv.fc.ep_list); + send.state.rndv_fc_ep_list); ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); ucp_trace_req(req, "mtype_fc: extract for replay"); @@ -89,7 +89,7 @@ void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, req, ucs_status_string(status)); ucs_assert(ucs_hlist_is_empty(fc_reqs) || (ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.rndv.fc.ep_list) != req)); + send.state.rndv_fc_ep_list) != req)); ucs_queue_push(replay_queue, (ucs_queue_elem_t*)&req->send.uct.priv); } } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 46aa26e1087..26c19876d1d 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -129,7 +129,7 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.fc.queue_elem); ucs_hlist_add_tail(&ep->ext->rndv_mtype_fc_reqs, - &req->send.rndv.fc.ep_list); + &req->send.state.rndv_fc_ep_list); return UCS_ERR_NO_RESOURCE; } From 65e91292674f7a0ca61737e6a4c802b12368165e Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 17 Sep 2026 16:33:29 +0300 Subject: [PATCH 50/67] UCP/RNDV: code review fixes - Fix a datatype iterator leak when a throttled rndv/get/mtype request is aborted. - Move the mtype flow-control declarations to the proto-v2 header. - Apply logging and style fixes. - Shrink the AM throttling test to a few MB of device memory. - Skip the peer-failure FC test when no CUDA fragments are available. - Share the stats helpers between the test classes. --- src/ucp/core/ucp_ep.c | 1 + src/ucp/rndv/proto_rndv.c | 4 +- src/ucp/rndv/proto_rndv.h | 21 +++++++++ src/ucp/rndv/rndv.h | 20 --------- src/ucp/rndv/rndv_get.c | 3 ++ src/ucp/rndv/rndv_mtype.inl | 3 +- src/ucp/wireup/wireup.c | 2 +- test/gtest/ucp/test_ucp_am.cc | 60 +++++++++++++------------ test/gtest/ucp/test_ucp_peer_failure.cc | 6 ++- 9 files changed, 66 insertions(+), 54 deletions(-) diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index d786100bbbc..a75178b109e 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index ddc4f355f77..4089861eb71 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -805,8 +805,8 @@ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, max_frags = max_mem / frag_size; if (max_frags < frags_in_chunk) { /* mpool requires max_elems >= elems_per_chunk */ - ucs_warn("RNDV_FRAG_WORKER_MAX_MEM (%zu) is too low for %s " - "(frag_size=%zu, frags_per_alloc=%zu), using minimum %zu " + ucs_warn("RNDV_FRAG_WORKER_MAX_MEM %zu is too low for %s " + "(frag size %zu, frags per alloc %zu), using minimum %zu " "frags", max_mem, ucs_memory_type_names[frag_mem_type], frag_size, frags_in_chunk, frags_in_chunk); diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 7811c385769..d44f2fa23b3 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -9,6 +9,8 @@ #include "rndv.h" +#include + #include @@ -198,6 +200,25 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type); +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg); + +int ucp_proto_rndv_mtype_fc_reschedule_filter( + const ucs_callbackq_elem_t *elem, void *arg); + +void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req); + +/* Abort every request of an endpoint that is waiting for a rndv mtype + * fragment */ +void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status); + +/** + * Move every request of an endpoint that is waiting for a rndv mtype + * fragment to @a replay_queue, so it is replayed after the endpoint is + * reconfigured. + */ +void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, + ucs_queue_head_t *replay_queue); + ucs_status_t ucp_proto_rndv_ack_init(const ucp_proto_common_init_params_t *init_params, diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index 438500f0a8d..45acfbfabb9 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -10,7 +10,6 @@ #include #include #include -#include typedef enum { @@ -119,25 +118,6 @@ ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p); -unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg); - -int ucp_proto_rndv_mtype_fc_reschedule_filter( - const ucs_callbackq_elem_t *elem, void *arg); - -void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req); - -/* Abort every request of an endpoint that is waiting for a rndv mtype - * fragment */ -void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status); - -/** - * Move every request of an endpoint that is waiting for a rndv mtype - * fragment to @a replay_queue, so it is replayed after the endpoint is - * reconfigured. - */ -void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, - ucs_queue_head_t *replay_queue); - void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, const void *rkey_buf); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index b90c9e8a9d5..583b72b19f6 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -416,6 +416,9 @@ ucp_proto_rndv_get_mtype_abort(ucp_request_t *req, ucs_status_t status) if (req->send.rndv.rkey != NULL) { ucp_proto_rndv_rkey_destroy(req); } + + ucp_datatype_iter_cleanup(&req->send.state.dt_iter, 1, + UCP_DT_MASK_ALL); ucp_proto_rndv_recv_complete_status(req, status); } } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 26c19876d1d..bc0cccf8efb 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -93,7 +93,8 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucp_worker_h worker = ep->worker; /* A rescheduled request owns the wakeup of a fragment which was released * back to the mpool, and must pass it on if it does not consume one. */ - int owns_wakeup = req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED; + int owns_wakeup = !!(req->flags & + UCP_REQUEST_FLAG_RNDV_MTYPE_FC_RESCHED); ucs_status_t status; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED)); diff --git a/src/ucp/wireup/wireup.c b/src/ucp/wireup/wireup.c index a7d08b94a91..135d4de49a9 100644 --- a/src/ucp/wireup/wireup.c +++ b/src/ucp/wireup/wireup.c @@ -22,7 +22,7 @@ #include #include #include -#include +#include #include #include diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 9726129e425..8aa436c28fa 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2272,6 +2272,28 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { m_mem_type = mem_type; } + static uint64_t get_stats(const entity &e, uint64_t cntr) + { + return UCS_STATS_GET_COUNTER(e.worker()->stats, cntr); + } + + static const char *stats_name(const entity &e, uint64_t cntr) + { + return e.worker()->stats->cls->counter_names[cntr]; + } + + void check_stats(const entity &e, uint64_t cntr, uint64_t exp_value) + { + EXPECT_EQ(exp_value, get_stats(e, cntr)) + << "counter is " << stats_name(e, cntr); + } + + void check_stats_ge(const entity &e, uint64_t cntr, uint64_t min_value) + { + EXPECT_GE(get_stats(e, cntr), min_value) + << "counter is " << stats_name(e, cntr); + } + private: ucs_memory_type_t tx_memtype() const override { @@ -2283,15 +2305,6 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { return m_mem_type; } - void check_stats(entity &e, uint64_t cntr, uint64_t exp_value) - { - auto stats_node = e.worker()->stats; - auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); - - EXPECT_EQ(exp_value, value) << "counter is " - << stats_node->cls->counter_names[cntr]; - } - ucp_err_handling_mode_t get_err_mode() const { return static_cast(get_variant_value(1)); @@ -2372,26 +2385,13 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); - fc.sender_throttled = - UCS_STATS_GET_COUNTER(sender().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - fc.receiver_throttled = - UCS_STATS_GET_COUNTER(receiver().worker()->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED); - } + const uint64_t cntr = UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED; -private: - static void check_stats_ge(const entity &e, uint64_t cntr, - uint64_t min_value) - { - const auto stats_node = e.worker()->stats; - const auto value = UCS_STATS_GET_COUNTER(stats_node, cntr); - - EXPECT_GE(value, min_value) - << "counter " << stats_node->cls->counter_names[cntr] - << " expected >= " << min_value << " but got " << value; + fc.sender_throttled = get_stats(sender(), cntr); + fc.receiver_throttled = get_stats(receiver(), cntr); } +private: static void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); @@ -2410,14 +2410,16 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { }; UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, - "RNDV_FRAG_WORKER_MAX_MEM=600mb", "RNDV_FRAG_MEM_TYPE=cuda") + "RNDV_FRAG_SIZE=cuda:256K", "RNDV_FRAG_ALLOC_COUNT=cuda:4", + "RNDV_FRAG_WORKER_MAX_MEM=1M", "RNDV_FRAG_MEM_TYPE=cuda") { fc_counters fc; - run_fc_test(200, fc); + /* 16 fragments against a 4-fragment quota */ + run_fc_test(16, fc); EXPECT_GT(fc.sender_throttled + fc.receiver_throttled, 0u) - << "throttling should have occurred with MAX_MEM=600mb"; + << "throttling should have occurred with MAX_MEM=1M"; verify_clean_fc_state(); } diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 81715ab590b..164404722a4 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -901,6 +901,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { continue; } + /* Terminates only because the mpool quota is finite, which + * requires proto v2 */ while ((mdesc = static_cast( ucs_mpool_get_inline( &kh_val(&worker->mpool_hash, iter)))) != @@ -1012,7 +1014,9 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { smoke_test(true); std::vector held_mdescs = hold_cuda_fragments(fc_entity.worker()); - ASSERT_FALSE(held_mdescs.empty()); + if (held_mdescs.empty()) { + UCS_TEST_SKIP_R("no CUDA fragments were allocated"); + } m_fc_pending_entity = &fc_entity; m_fc_op = fc_op; From 3096ac96c81a351215911eba2781a9ff680700bb Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 17 Sep 2026 18:31:25 +0300 Subject: [PATCH 51/67] UCP/RNDV: small fixes --- src/ucp/rndv/proto_rndv.c | 2 -- src/ucp/rndv/rndv.c | 4 ---- 2 files changed, 6 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 4089861eb71..7c64c65d2ff 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -46,8 +46,6 @@ int ucp_proto_rndv_mtype_fc_reschedule_filter( return 0; } - ucs_assertv(0, "ep %p still has mtype FC reschedule callback for req %p", - arg, elem->arg); ucs_error("ep %p still has mtype FC reschedule callback for req %p", arg, elem->arg); return 1; diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 53a0b870381..855a7aebb4b 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1183,10 +1183,6 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); - if (status == UCS_ERR_NO_RESOURCE) { - status = UCS_ERR_NO_MEMORY; - } - goto err; } From ac56848ca90084307b850485e536ff96f9f242f6 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 22 Sep 2026 13:39:07 +0300 Subject: [PATCH 52/67] UCP/RNDV: CI failure - allow aborting RTS request with completed fragments --- src/ucp/rndv/proto_rndv.c | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 7c64c65d2ff..a3a3f9201eb 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -763,16 +763,16 @@ void ucp_proto_rndv_rts_query(const ucp_proto_query_params_t *params, void ucp_proto_rndv_rts_abort(ucp_request_t *req, ucs_status_t status) { + int invalidating; + ucp_am_release_user_header(req); ucp_request_rndv_flush_complete(req); - if (ucp_request_memh_invalidate(req, status)) { - ucp_proto_rndv_rts_reset(req); - return; + invalidating = ucp_request_memh_invalidate(req, status); + ucp_proto_request_zcopy_id_reset(req); + if (!invalidating) { + ucp_request_complete_send(req, status); } - - ucp_proto_rndv_rts_reset(req); - ucp_request_complete_send(req, status); } ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) From c1efc3b0e3874cdbe36332573fc3cbecc9eef391 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 22 Sep 2026 17:11:44 +0300 Subject: [PATCH 53/67] UCP/RNDV: small fixes --- src/ucp/rndv/AGENTS.md | 4 ++++ src/ucp/rndv/proto_rndv.c | 2 ++ src/ucp/rndv/rndv_mtype.inl | 5 +++-- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/ucp/rndv/AGENTS.md b/src/ucp/rndv/AGENTS.md index 74a431a5ff9..56fad2ad225 100644 --- a/src/ucp/rndv/AGENTS.md +++ b/src/ucp/rndv/AGENTS.md @@ -59,6 +59,10 @@ Defined in `rndv.h`: pool growth. - For GPU paths, always go through `rndv_mtype.inl` helpers and the `memtype_cache`. Direct CUDA driver-API calls in this dir are wrong. +- Acquire mtype staging fragments with `ucp_proto_rndv_mtype_request_init()` + and return them with `ucp_proto_rndv_mtype_mdesc_release()`. A plain + `ucs_mpool_put()` on `send.rndv.mdesc` skips the flow-control wakeup and + can leave throttled requests queued forever. - Pipelining (`ppln`) is a meta-protocol: it owns a parent request that spawns child requests selected with `UCP_PROTO_SELECT_OP_FLAG_PPLN_FRAG`. Never inline-progress a pipeline child — schedule it through the diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index a3a3f9201eb..5dbf1ea43ee 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -768,6 +768,8 @@ void ucp_proto_rndv_rts_abort(ucp_request_t *req, ucs_status_t status) ucp_am_release_user_header(req); ucp_request_rndv_flush_complete(req); + /* Fragments of this request may have already completed (e.g. aborted + * during EP purge), so ucp_proto_rndv_rts_reset() cannot be used here */ invalidating = ucp_request_memh_invalidate(req, status); ucp_proto_request_zcopy_id_reset(req); if (!invalidating) { diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index bc0cccf8efb..5b58e28c1f1 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -123,9 +123,10 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", ucs_memory_type_names[frag_mem_type], frag_sys_dev); - UCS_STATS_UPDATE_COUNTER(worker->stats, - UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + UCP_WORKER_STAT_RNDV(worker, MTYPE_FC_THROTTLED, 1); ucs_assert(!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK)); + /* The EP list link aliases send.state fields used once initialized */ + ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED; ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.fc.queue_elem); From b61d637324e24048a845caa3f547ea5ffe116ec4 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 23 Sep 2026 10:35:28 +0300 Subject: [PATCH 54/67] UCP/RNDV: Fix mtype FC wakeup, GET abort gate and fragment limit rounding --- src/ucp/core/ucp_context.c | 6 +++--- src/ucp/rndv/proto_rndv.c | 17 ++++++----------- src/ucp/rndv/rndv.c | 13 ++++++++----- src/ucp/rndv/rndv_get.c | 2 +- src/ucp/rndv/rndv_mtype.inl | 2 ++ 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index d5c2db96c36..3f687baedce 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -405,9 +405,9 @@ static ucs_config_field_t ucp_context_config_table[] = { "fragments. Both \"auto\" and \"inf\" mean no limit. This setting has no\n" "effect when PROTO_ENABLE=n. The same limit is applied independently to\n" "each fragment memory type, and per device when available. It is\n" - "converted to a fragment count using RNDV_FRAG_SIZE and\n" - "RNDV_FRAG_ALLOC_COUNT. When the limit is reached, further fragment\n" - "requests are queued until fragments are released", + "converted to a fragment count using RNDV_FRAG_SIZE. When the limit is\n" + "reached, further fragment requests are queued until fragments are\n" + "released", ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 5dbf1ea43ee..549b82f0eaf 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -791,26 +791,21 @@ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, { const size_t max_mem = context->config.ext.rndv_frag_worker_max_mem; size_t frag_size; - size_t frags_in_chunk; size_t max_frags; if ((max_mem == UCS_MEMUNITS_INF) || (max_mem == UCS_MEMUNITS_AUTO)) { return UINT_MAX; } - frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; - frags_in_chunk = context->config.ext.rndv_num_frags[frag_mem_type]; + frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; ucs_assert_always(frag_size > 0); max_frags = max_mem / frag_size; - if (max_frags < frags_in_chunk) { - /* mpool requires max_elems >= elems_per_chunk */ - ucs_warn("RNDV_FRAG_WORKER_MAX_MEM %zu is too low for %s " - "(frag size %zu, frags per alloc %zu), using minimum %zu " - "frags", - max_mem, ucs_memory_type_names[frag_mem_type], frag_size, - frags_in_chunk, frags_in_chunk); - return frags_in_chunk; + if (max_frags == 0) { + ucs_warn("RNDV_FRAG_WORKER_MAX_MEM %zu is smaller than one %s " + "fragment (%zu), using one fragment", + max_mem, ucs_memory_type_names[frag_mem_type], frag_size); + return 1; } return ucs_min(max_frags, UINT_MAX); diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 855a7aebb4b..8b4b66cecc1 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1144,6 +1144,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, ucp_worker_mpool_key_t key; ucs_status_t status; unsigned num_frags; + unsigned max_elems; ucs_mpool_t *mpool; khiter_t khiter; int khret; @@ -1168,16 +1169,18 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mpool = &kh_value(&worker->mpool_hash, khiter); num_frags = worker->context->config.ext.rndv_num_frags[key.mem_type]; + max_elems = worker->context->config.ext.proto_enable ? + ucp_proto_rndv_frag_max_elems(worker->context, + key.mem_type) : + UINT_MAX; ucs_mpool_params_reset(&mp_params); mp_params.priv_size = sizeof(ucp_rndv_mpool_priv_t); mp_params.elem_size = sizeof(ucp_mem_desc_t); mp_params.alignment = 1; - mp_params.elems_per_chunk = num_frags; - mp_params.max_elems = worker->context->config.ext.proto_enable ? - ucp_proto_rndv_frag_max_elems( - worker->context, key.mem_type) : - UINT_MAX; + /* A chunk may not exceed the limit */ + mp_params.elems_per_chunk = ucs_min(num_frags, max_elems); + mp_params.max_elems = max_elems; mp_params.ops = &ucp_frag_mpool_ops; mp_params.name = "ucp_rndv_frags"; status = ucs_mpool_init(&mp_params, mpool); diff --git a/src/ucp/rndv/rndv_get.c b/src/ucp/rndv/rndv_get.c index 583b72b19f6..452e719f008 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -398,7 +398,7 @@ ucp_proto_rndv_get_mtype_abort(ucp_request_t *req, ucs_status_t status) { ucp_request_t *super_req; - if (!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_STATE_MASK)) { + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { ucp_proto_abort_fatal_not_implemented(req, status); return; } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 5b58e28c1f1..efe67cc6b08 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -79,6 +79,8 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_worker_h worker) ucs_callbackq_add_oneshot(&worker->uct->progress_q, pending_req->send.ep, ucp_proto_rndv_mtype_fc_reschedule_cb, pending_req); + /* The retry is not tied to any event, so wake up a blocked worker */ + ucp_worker_signal_internal(worker); return; } } From c984f1134c2e9ec4d0a0f3f6bd02f1a1990875ac Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 23 Sep 2026 10:47:25 +0300 Subject: [PATCH 55/67] UCP/RNDV: Dequeue throttled requests of an endpoint in one pass --- src/ucp/rndv/proto_rndv.c | 74 +++++++++++++++++++++++++++++++------ src/ucp/rndv/rndv_mtype.inl | 2 + 2 files changed, 65 insertions(+), 11 deletions(-) diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 549b82f0eaf..1552a108af6 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -51,11 +51,47 @@ int ucp_proto_rndv_mtype_fc_reschedule_filter( return 1; } +/* Move all queued requests of the endpoint out of the worker pending queues + * and out of the flow-control state into @a reqs, in a single pass. Requests + * which were already woken up are not queued and stay on the endpoint list. */ +static void +ucp_proto_rndv_mtype_fc_ep_dequeue(ucp_ep_h ep, ucs_queue_head_t *reqs) +{ + ucp_worker_h worker = ep->worker; + ucs_queue_head_t *pending_q; + ucs_queue_iter_t iter; + ucp_request_t *req; + unsigned q_index; + + for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; ++q_index) { + pending_q = &worker->rndv_mtype_fc.pending_q[q_index]; + ucs_queue_for_each_safe(req, iter, pending_q, + send.rndv.fc.queue_elem) { + if (req->send.ep != ep) { + continue; + } + + ucs_queue_del_iter(pending_q, iter); + ucp_proto_rndv_mtype_fc_leave(req); + ucs_queue_push(reqs, &req->send.rndv.fc.queue_elem); + } + } +} + void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) { ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; + ucs_queue_head_t reqs; ucp_request_t *req; + ucs_queue_head_init(&reqs); + ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs); + ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { + ucp_proto_request_abort(req, status); + } + + /* Aborting a woken-up request may wake up another request of this ep, + * which is then aborted on a later iteration */ while (!ucs_hlist_is_empty(fc_reqs)) { req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, send.state.rndv_fc_ep_list); @@ -66,29 +102,45 @@ void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) } } +/* Reset the request, since the replay does not always restart it, and add it + * to the replay queue */ +static void +ucp_proto_rndv_mtype_fc_extract_one(ucp_request_t *req, + ucs_queue_head_t *replay_queue) +{ + ucs_status_t status; + + ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); + ucp_trace_req(req, "mtype_fc: extract for replay"); + + status = req->send.proto_config->proto->reset(req); + ucs_assertv_always(status == UCS_OK, "req %p, failed to reset: %s", req, + ucs_status_string(status)); + ucs_queue_push(replay_queue, (ucs_queue_elem_t*)&req->send.uct.priv); +} + void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, ucs_queue_head_t *replay_queue) { ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; + ucs_queue_head_t reqs; ucp_request_t *req; - ucs_status_t status; - /* The replay does not always restart (and reset) the request, so clear - * its flow-control state here. This may wake up another request of this - * ep, which is then extracted on a later iteration. */ + ucs_queue_head_init(&reqs); + ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs); + ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { + ucp_proto_rndv_mtype_fc_extract_one(req, replay_queue); + } + + /* Resetting a woken-up request may wake up another request of this ep, + * which is then extracted on a later iteration */ while (!ucs_hlist_is_empty(fc_reqs)) { req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, send.state.rndv_fc_ep_list); - ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); - ucp_trace_req(req, "mtype_fc: extract for replay"); - - status = req->send.proto_config->proto->reset(req); - ucs_assertv_always(status == UCS_OK, "req %p, failed to reset: %s", - req, ucs_status_string(status)); + ucp_proto_rndv_mtype_fc_extract_one(req, replay_queue); ucs_assert(ucs_hlist_is_empty(fc_reqs) || (ucs_hlist_head_elem(fc_reqs, ucp_request_t, send.state.rndv_fc_ep_list) != req)); - ucs_queue_push(replay_queue, (ucs_queue_elem_t*)&req->send.uct.priv); } } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index efe67cc6b08..ce34e90dcee 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -267,6 +267,8 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) if (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC_QUEUED) { ucp_trace_req(req, "mtype_fc: remove aborted request from queue"); + /* O(n) for a single request; endpoint purge and reconfiguration + * dequeue all requests of the endpoint in one pass instead */ ucs_queue_remove(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.fc.queue_elem); owns_wakeup = 0; From 5a8bdb4b8d5201885eb9ba6c535dc26bed1f109e Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 6 Oct 2026 12:22:03 +0300 Subject: [PATCH 56/67] UCP/RNDV: Keep part of the fragment quota in a pool for PUT/GET only --- src/ucp/core/ucp_context.c | 21 +++++++- src/ucp/core/ucp_context.h | 2 + src/ucp/core/ucp_mm.h | 1 + src/ucp/core/ucp_worker.h | 14 ++++++ src/ucp/core/ucp_worker.inl | 6 ++- src/ucp/rndv/proto_rndv.c | 15 ++++++ src/ucp/rndv/proto_rndv.h | 5 ++ src/ucp/rndv/rndv.c | 45 +++++++++++++---- src/ucp/rndv/rndv.h | 8 +++- src/ucp/rndv/rndv_mtype.inl | 66 +++++++++++++++++++------ test/gtest/ucp/test_ucp_mmap.cc | 85 +++++++++++++++++++++++++-------- 11 files changed, 220 insertions(+), 48 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 17d20a1c56d..08a43103690 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -425,10 +425,21 @@ static ucs_config_field_t ucp_context_config_table[] = { "each fragment memory type, and per device when available. It is\n" "converted to a fragment count using RNDV_FRAG_SIZE. When the limit is\n" "reached, further fragment requests are queued until fragments are\n" - "released", + "released. See also RNDV_FRAG_RTR_RATIO", ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, + {"RNDV_FRAG_RTR_RATIO", "0.8", + "Fraction of RNDV_FRAG_WORKER_MAX_MEM which fragments posted to receive\n" + "data (RTR) may use at once, in the range (0, 1]. The rest is kept for\n" + "the put/get operations which move the data, so that they can always\n" + "make progress: an RTR fragment is released only when the peer's data\n" + "arrives, and the peer needs a fragment of its own to send it. A value\n" + "of 1 keeps nothing, which is safe only for a worker that never moves\n" + "rendezvous data through staging fragments", + ucs_offsetof(ucp_context_config_t, rndv_frag_rtr_ratio), + UCS_CONFIG_TYPE_POS_DOUBLE}, + {"FLUSH_WORKER_EPS", "y", "Enable flushing the worker by flushing its endpoints. Allows completing\n" "the flush operation in a bounded time even if there are new requests on\n" @@ -2701,6 +2712,14 @@ static ucs_status_t ucp_fill_config(ucp_context_h context, goto err_free_alloc_methods; } + if (UCS_CONFIG_DBL_IS_AUTO(context->config.ext.rndv_frag_rtr_ratio) || + !(context->config.ext.rndv_frag_rtr_ratio <= 1.0)) { + ucs_error("UCX_RNDV_FRAG_RTR_RATIO value %g is not in the range " + "(0, 1]", context->config.ext.rndv_frag_rtr_ratio); + status = UCS_ERR_INVALID_PARAM; + goto err_free_alloc_methods; + } + if (!ucp_dynamic_tl_switch_config_valid(&context->config.ext)) { status = UCS_ERR_INVALID_PARAM; goto err_free_alloc_methods; diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index f355c78aed6..76ca6d175d1 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -132,6 +132,8 @@ typedef struct ucp_context_config { int rndv_errh_ppln_enable; /** Maximum memory for concurrent rndv fragments per worker (bytes) */ size_t rndv_frag_worker_max_mem; + /** Fraction of rndv_frag_worker_max_mem which RTR fragments may use */ + double rndv_frag_rtr_ratio; /** Threshold for using tag matching offload capabilities. Smaller buffers * will not be posted to the transport. */ size_t tm_thresh; diff --git a/src/ucp/core/ucp_mm.h b/src/ucp/core/ucp_mm.h index a135d104080..9b7c02a1d32 100644 --- a/src/ucp/core/ucp_mm.h +++ b/src/ucp/core/ucp_mm.h @@ -117,6 +117,7 @@ typedef struct ucp_rndv_mpool_priv { ucp_worker_h worker; ucs_memory_type_t mem_type; ucs_sys_device_t sys_dev; + uint8_t pool; /* UCP_WORKER_RNDV_FRAG_POOL_* */ } ucp_rndv_mpool_priv_t; diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 99c6ac73b0b..f80b5edeefa 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -224,10 +224,24 @@ KHASH_TYPE(ucp_worker_discard_uct_ep_hash, uct_ep_h, ucp_request_t*); typedef khash_t(ucp_worker_discard_uct_ep_hash) ucp_worker_discard_uct_ep_hash_t; +/** + * Rendezvous fragment pools of a memory type and device. RTR requests use + * only the shared pool, while PUT/GET requests use the reserved pool first. + * This keeps fragments for PUT/GET: an RTR fragment is released only when + * the peer's data arrives, and the peer needs a fragment of its own to send + * it, so RTR fragments on both peers could otherwise wait for each other. + */ +typedef enum { + UCP_WORKER_RNDV_FRAG_POOL_SHARED, + UCP_WORKER_RNDV_FRAG_POOL_RESERVED +} ucp_worker_rndv_frag_pool_t; + + typedef struct ucp_worker_mpool_key { ucs_memory_type_t mem_type; /* memory type of the buffer pool */ ucs_sys_device_t sys_dev; /* identifier for the device, UINT_MAX for default device */ + uint8_t pool; /* UCP_WORKER_RNDV_FRAG_POOL_* */ } ucp_worker_mpool_key_t; diff --git a/src/ucp/core/ucp_worker.inl b/src/ucp/core/ucp_worker.inl index 4f7dce3762c..75b6e353652 100644 --- a/src/ucp/core/ucp_worker.inl +++ b/src/ucp/core/ucp_worker.inl @@ -48,7 +48,8 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_worker_rkey_config_get( static UCS_F_ALWAYS_INLINE khint_t ucp_worker_mpool_hash_func(ucp_worker_mpool_key_t mpool_key) { - return (khint_t)mpool_key.mem_type ^ (mpool_key.sys_dev << 8); + return (khint_t)mpool_key.mem_type ^ (mpool_key.sys_dev << 8) ^ + (mpool_key.pool << 16); } static UCS_F_ALWAYS_INLINE int @@ -56,7 +57,8 @@ ucp_worker_mpool_key_is_equal(ucp_worker_mpool_key_t mpool_key1, ucp_worker_mpool_key_t mpool_key2) { return (mpool_key1.sys_dev == mpool_key2.sys_dev) && - (mpool_key1.mem_type == mpool_key2.mem_type); + (mpool_key1.mem_type == mpool_key2.mem_type) && + (mpool_key1.pool == mpool_key2.pool); } KHASH_IMPL(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t, diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 93e58de916c..238dceedee0 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -863,6 +863,21 @@ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, return ucs_min(max_frags, UINT_MAX); } +unsigned ucp_proto_rndv_frag_shared_elems(ucp_context_h context, + unsigned max_elems) +{ + const double ratio = context->config.ext.rndv_frag_rtr_ratio; + unsigned shared; + + /* A reserve needs at least two fragments: one to keep and one to share */ + if ((max_elems == UINT_MAX) || (ratio >= 1.0) || (max_elems < 2)) { + return max_elems; + } + + shared = ucs_min((unsigned)(max_elems * ratio), max_elems - 1); + return ucs_max(shared, 1); +} + ucs_status_t ucp_proto_rndv_ack_init(const ucp_proto_common_init_params_t *init_params, const char *name, double overhead, diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index d44f2fa23b3..28e3c7f6c81 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -200,6 +200,11 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req); unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, ucs_memory_type_t frag_mem_type); +/* Number of fragments out of @a max_elems in the shared pool, which is + * UCX_RNDV_FRAG_RTR_RATIO of them; the rest go to the reserved pool. */ +unsigned ucp_proto_rndv_frag_shared_elems(ucp_context_h context, + unsigned max_elems); + unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg); int ucp_proto_rndv_mtype_fc_reschedule_filter( diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 8b4b66cecc1..127e1a15018 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1135,16 +1135,24 @@ ucs_mpool_ops_t ucp_frag_mpool_ops = { .obj_cleanup = (ucs_mpool_obj_cleanup_func_t)ucs_empty_function }; +static const char *ucp_rndv_frag_mpool_names[] = { + [UCP_WORKER_RNDV_FRAG_POOL_SHARED] = "ucp_rndv_frags", + [UCP_WORKER_RNDV_FRAG_POOL_RESERVED] = "ucp_rndv_frags_reserved" +}; + ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, - ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p) + ucs_sys_device_t sys_dev, unsigned pool, + ucp_mem_desc_t **mdesc_p) { + ucp_context_h context = worker->context; ucp_rndv_mpool_priv_t *mpriv; ucp_mem_desc_t *mdesc; ucp_worker_mpool_key_t key; ucs_status_t status; unsigned num_frags; unsigned max_elems; + unsigned shared_elems; ucs_mpool_t *mpool; khiter_t khiter; int khret; @@ -1152,6 +1160,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, key.sys_dev = sys_dev; key.mem_type = mem_type; + key.pool = pool; khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); if (ucs_likely(khiter != kh_end(&worker->mpool_hash))) { @@ -1159,6 +1168,24 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, goto out_mp_get; } + /* The shared pool is for all operations, the reserved pool is for PUT/GET + * only. Without a limit, or without a reserve, there is no reserved + * pool. */ + max_elems = context->config.ext.proto_enable ? + ucp_proto_rndv_frag_max_elems(context, + key.mem_type) : + UINT_MAX; + shared_elems = ucp_proto_rndv_frag_shared_elems(context, max_elems); + if (pool == UCP_WORKER_RNDV_FRAG_POOL_SHARED) { + max_elems = shared_elems; + } else { + ucs_assert(pool == UCP_WORKER_RNDV_FRAG_POOL_RESERVED); + max_elems -= shared_elems; + if (max_elems == 0) { + return UCS_ERR_NO_RESOURCE; + } + } + khiter = kh_put(ucp_worker_mpool_hash, &worker->mpool_hash, key, &khret); if (khret == UCS_KH_PUT_FAILED) { status = UCS_ERR_NO_MEMORY; @@ -1168,11 +1195,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, ucs_assert_always(khret != UCS_KH_PUT_KEY_PRESENT); mpool = &kh_value(&worker->mpool_hash, khiter); - num_frags = worker->context->config.ext.rndv_num_frags[key.mem_type]; - max_elems = worker->context->config.ext.proto_enable ? - ucp_proto_rndv_frag_max_elems(worker->context, - key.mem_type) : - UINT_MAX; + num_frags = context->config.ext.rndv_num_frags[key.mem_type]; ucs_mpool_params_reset(&mp_params); mp_params.priv_size = sizeof(ucp_rndv_mpool_priv_t); @@ -1182,7 +1205,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.elems_per_chunk = ucs_min(num_frags, max_elems); mp_params.max_elems = max_elems; mp_params.ops = &ucp_frag_mpool_ops; - mp_params.name = "ucp_rndv_frags"; + mp_params.name = ucp_rndv_frag_mpool_names[pool]; status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); @@ -1193,6 +1216,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mpriv->worker = worker; mpriv->mem_type = key.mem_type; mpriv->sys_dev = sys_dev; + mpriv->pool = pool; out_mp_get: mdesc = ucp_worker_mpool_get(mpool); @@ -1232,7 +1256,8 @@ static void ucp_rndv_send_frag_get_mem_type(ucp_request_t *sreq, size_t length, } status = ucp_rndv_mpool_get(worker, frag_mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc); + UCS_SYS_DEVICE_ID_UNKNOWN, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc); if (ucs_unlikely(status != UCS_OK)) { ucs_fatal("failed to allocate fragment memory desc: %s", ucs_status_string(status)); @@ -1411,7 +1436,9 @@ static void ucp_rndv_send_frag_rtr(ucp_worker_h worker, ucp_request_t *rndv_req, /* allocate fragment recv buffer desc*/ status = ucp_rndv_mpool_get(worker, frag_mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc); + UCS_SYS_DEVICE_ID_UNKNOWN, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, + &mdesc); if (status != UCS_OK) { ucs_fatal("failed to allocate fragment memory buffer: %s", ucs_status_string(status)); diff --git a/src/ucp/rndv/rndv.h b/src/ucp/rndv/rndv.h index 45acfbfabb9..2d32141c4ef 100644 --- a/src/ucp/rndv/rndv.h +++ b/src/ucp/rndv/rndv.h @@ -114,9 +114,15 @@ ucs_status_t ucp_proto_progress_rndv_rtr(uct_pending_req_t *self); ucs_status_t ucp_rndv_reg_send_buffer(ucp_request_t *sreq, const ucp_request_param_t *param); +/** + * Get a staging fragment from the given UCP_WORKER_RNDV_FRAG_POOL_* pool of + * the memory type and device. Returns UCS_ERR_NO_RESOURCE if the pool quota + * is exhausted, or if the pool has no quota at all. + */ ucs_status_t ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, - ucs_sys_device_t sys_dev, ucp_mem_desc_t **mdesc_p); + ucs_sys_device_t sys_dev, unsigned pool, + ucp_mem_desc_t **mdesc_p); void ucp_rndv_receive(ucp_worker_h worker, ucp_request_t *rreq, const ucp_rndv_rts_hdr_t *rndv_rts_hdr, diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 460d293fd04..5a580d5a179 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -46,7 +46,8 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params, /** * Reschedule a pending throttled request after a fragment is released back to - * the mpool. Dequeue priority: PUT/GET > RTR. + * the mpool, from the first @a num_queues queues. Dequeue priority: + * PUT/GET > RTR. * * Priority rationale: * PUT/GET - Completing a PUT or GET frees a staging buffer, reducing memory @@ -55,13 +56,16 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params, * memory pressure. */ static UCS_F_ALWAYS_INLINE void -ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_worker_h worker) +ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_worker_h worker, + unsigned num_queues) { ucp_request_t *pending_req; unsigned q_index; + ucs_assert(num_queues <= UCP_WORKER_RNDV_FC_OP_LAST); + /* Dequeue from highest-priority non-empty queue (PUT/GET before RTR) */ - for (q_index = 0; q_index < UCP_WORKER_RNDV_FC_OP_LAST; q_index++) { + for (q_index = 0; q_index < num_queues; q_index++) { if (ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[q_index])) { continue; } @@ -84,6 +88,16 @@ ucp_proto_rndv_mtype_fc_reschedule_pending(ucp_worker_h worker) } } +/* Whether a reserved fragment pool can exist, see ucp_rndv_mpool_get() */ +static UCS_F_ALWAYS_INLINE int +ucp_proto_rndv_mtype_frag_has_reserve(ucp_context_h context) +{ + const size_t max_mem = context->config.ext.rndv_frag_worker_max_mem; + + return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO) && + (context->config.ext.rndv_frag_rtr_ratio < 1.0); +} + static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucs_memory_type_t frag_mem_type, @@ -106,22 +120,37 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, } req->send.rndv.mdesc = NULL; - status = ucp_rndv_mpool_get(worker, frag_mem_type, - frag_sys_dev, - &req->send.rndv.mdesc); + status = UCS_ERR_NO_RESOURCE; + + /* PUT/GET take from the reserved pool first, so that the shared pool + * stays available for RTR as long as possible */ + if ((fc_op != UCP_WORKER_RNDV_FC_OP_RTR) && + ucp_proto_rndv_mtype_frag_has_reserve(worker->context)) { + status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, + UCP_WORKER_RNDV_FRAG_POOL_RESERVED, + &req->send.rndv.mdesc); + } + + if (status == UCS_ERR_NO_RESOURCE) { + status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, + &req->send.rndv.mdesc); + } + if (status != UCS_ERR_NO_RESOURCE) { if ((status != UCS_OK) && owns_wakeup) { /* The caller aborts the request, so the wakeup would be lost. */ - ucp_proto_rndv_mtype_fc_reschedule_pending(worker); + ucp_proto_rndv_mtype_fc_reschedule_pending( + worker, UCP_WORKER_RNDV_FC_OP_LAST); } return status; } - /* Mpool quota exhausted - throttle by queuing the request in the - * appropriate pending queue ordered by priority. */ + /* Quota exhausted - throttle by queuing the request in the appropriate + * pending queue ordered by priority. */ ucp_trace_req(req, - "mtype_fc: mpool exhausted, queue %s mem_type %s " + "mtype_fc: quota exhausted, queue %s mem_type %s " "sys_dev %u", (fc_op == UCP_WORKER_RNDV_FC_OP_RTR) ? "rtr" : "put/get", ucs_memory_type_names[frag_mem_type], @@ -294,7 +323,8 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) /* A fragment was already released on behalf of this request, so pass * the wakeup to the next waiter rather than dropping it, otherwise the * remaining waiters could stall while the fragment is free. */ - ucp_proto_rndv_mtype_fc_reschedule_pending(worker); + ucp_proto_rndv_mtype_fc_reschedule_pending(worker, + UCP_WORKER_RNDV_FC_OP_LAST); } } @@ -306,11 +336,19 @@ ucp_proto_rndv_mtype_fc_cancel(ucp_request_t *req, unsigned fc_op) static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) { - ucp_worker_h worker = req->send.ep->worker; + ucp_worker_h worker = req->send.ep->worker; + ucp_mem_desc_t *mdesc = req->send.rndv.mdesc; + ucp_rndv_mpool_priv_t *mpriv = ucs_mpool_priv(ucs_mpool_obj_owner(mdesc)); + unsigned num_queues; + + /* Only PUT/GET requests can use a fragment of the reserved pool */ + num_queues = (mpriv->pool == UCP_WORKER_RNDV_FRAG_POOL_RESERVED) ? + (UCP_WORKER_RNDV_FC_OP_PUT + 1) : + UCP_WORKER_RNDV_FC_OP_LAST; - ucs_mpool_put_inline(req->send.rndv.mdesc); + ucs_mpool_put_inline(mdesc); req->send.rndv.mdesc = NULL; - ucp_proto_rndv_mtype_fc_reschedule_pending(worker); + ucp_proto_rndv_mtype_fc_reschedule_pending(worker, num_queues); } static UCS_F_ALWAYS_INLINE ucs_status_t diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index a2758890f93..ea0795db51c 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1019,14 +1019,23 @@ rndv_mpool_chunk_alloc_fail(ucs_mpool_t*, size_t*, void**) return UCS_ERR_NO_MEMORY; } +static ucs_status_t rndv_mpool_get(ucp_worker_h worker, + ucs_memory_type_t mem_type, unsigned pool, + ucp_mem_desc_t **mdesc_p) +{ + return ucp_rndv_mpool_get(worker, mem_type, UCS_SYS_DEVICE_ID_UNKNOWN, + pool, mdesc_p); +} + UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) { ucp_worker_h worker = sender().worker(); for (auto mem_type : mem_buffer::supported_mem_types()) { ucp_mem_desc_t *mdesc; - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc)); + ASSERT_UCS_OK(rndv_mpool_get(worker, mem_type, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, + &mdesc)); ASSERT_NE(nullptr, mdesc); EXPECT_EQ(mdesc->memh, mdesc->memh->parent); ucs_mpool_put(mdesc); @@ -1036,25 +1045,58 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_quota_exhausted, (get_variant_value() == VARIANT_PROTO_DISABLE), "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", - "RNDV_FRAG_WORKER_MAX_MEM=8K") + "RNDV_FRAG_WORKER_MAX_MEM=8K", "RNDV_FRAG_RTR_RATIO=1") { ucp_mem_desc_t *mdesc1; ucp_mem_desc_t *mdesc2; ucp_mem_desc_t *mdesc3; ucp_worker_h worker = sender().worker(); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc1)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc2)); EXPECT_EQ(UCS_ERR_NO_RESOURCE, - ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3)); + rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc3)); ucs_mpool_put(mdesc2); ucs_mpool_put(mdesc1); } +/* 4 fragments: 3 in the shared pool for any operation, 1 kept in the + * reserved pool for PUT/GET */ +UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_rtr_ratio, + (get_variant_value() == VARIANT_PROTO_DISABLE), + "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:4", + "RNDV_FRAG_WORKER_MAX_MEM=16K", "RNDV_FRAG_RTR_RATIO=0.8") +{ + ucp_worker_h worker = sender().worker(); + std::vector shared(3); + ucp_mem_desc_t *reserved; + ucp_mem_desc_t *mdesc; + + for (auto &m : shared) { + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &m)); + } + + EXPECT_EQ(UCS_ERR_NO_RESOURCE, + rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_RESERVED, + &reserved)); + EXPECT_EQ(UCS_ERR_NO_RESOURCE, + rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_RESERVED, &mdesc)); + + ucs_mpool_put(reserved); + for (auto m : shared) { + ucs_mpool_put(m); + } +} + UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_disabled_with_proto_v1, "PROTO_ENABLE=n", "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=8K") @@ -1064,12 +1106,12 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_disabled_with_proto_v1, ucp_mem_desc_t *mdesc3; ucp_worker_h worker = sender().worker(); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc1)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc2)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc3)); ucs_mpool_put(mdesc3); ucs_mpool_put(mdesc2); @@ -1089,13 +1131,14 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, const ucs_mpool_ops_t *orig_ops; ucs_status_t status; - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc1)); - ASSERT_UCS_OK(ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc2)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc1)); + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc2)); key.mem_type = UCS_MEMORY_TYPE_HOST; key.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + key.pool = UCP_WORKER_RNDV_FRAG_POOL_SHARED; const khiter_t khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); ASSERT_NE(kh_end(&worker->mpool_hash), khiter); @@ -1110,8 +1153,8 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, mpool->data->ops = &fail_ops; { scoped_log_handler slh(hide_errors_logger); - status = ucp_rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCS_SYS_DEVICE_ID_UNKNOWN, &mdesc3); + status = rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc3); } mpool->data->ops = orig_ops; From 342492a9eccf04509776ae8becb279994950cc22 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 6 Oct 2026 14:25:28 +0300 Subject: [PATCH 57/67] UCP/RNDV: Assert fragment mpool creation does not report quota exhaustion --- src/ucp/rndv/rndv.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 127e1a15018..0b319669421 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1208,6 +1208,8 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, mp_params.name = ucp_rndv_frag_mpool_names[pool]; status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { + /* Callers retry on UCS_ERR_NO_RESOURCE, meaning quota exhaustion */ + ucs_assert(status != UCS_ERR_NO_RESOURCE); kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); goto err; } From 29be1062934c980769f68fd8a222332c96a86d1d Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 7 Oct 2026 11:27:51 +0300 Subject: [PATCH 58/67] UCP/RNDV: small fix --- src/ucp/rndv/rndv_mtype.inl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 5a580d5a179..13a03b52592 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -108,9 +108,9 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, ucp_worker_h worker = ep->worker; /* A rescheduled request owns the wakeup of a fragment which was released * back to the mpool, and must pass it on if it does not consume one. */ - int owns_wakeup = (req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC) && - (req->send.rndv.fc.state == - UCP_REQUEST_RNDV_MTYPE_FC_RESCHED); + int owns_wakeup = !!((req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC) && + (req->send.rndv.fc.state == + UCP_REQUEST_RNDV_MTYPE_FC_RESCHED)); ucs_status_t status; /* A queued request is retried only after being rescheduled */ From ca2f87c7c190afb8de9b953d94a0b826e5def3e5 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 7 Oct 2026 16:54:40 +0300 Subject: [PATCH 59/67] UCP/RNDV: Find throttled requests of an endpoint without a per-EP list --- src/ucp/core/ucp_ep.c | 2 - src/ucp/core/ucp_ep.h | 2 - src/ucp/core/ucp_request.h | 5 -- src/ucp/rndv/proto_rndv.c | 75 ++++++++++++++++-------------- src/ucp/rndv/rndv_mtype.inl | 4 -- test/gtest/common/test_obj_size.cc | 2 +- 6 files changed, 40 insertions(+), 50 deletions(-) diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index 194bf9e2f23..125f7a2a928 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -286,7 +286,6 @@ static ucp_ep_h ucp_ep_allocate(ucp_worker_h worker, const char *peer_name) memset(&ep->ext->ep_match, 0, sizeof(ep->ext->ep_match)); ucs_hlist_head_init(&ep->ext->proto_reqs); - ucs_hlist_head_init(&ep->ext->rndv_mtype_fc_reqs); for (lane = 0; lane < UCP_MAX_FAST_PATH_LANES; ++lane) { ucp_ep_set_lane(ep, lane, NULL); @@ -550,7 +549,6 @@ void ucp_ep_destroy_base(ucp_ep_h ep) ucp_ep_refcount_assert(ep, discard, ==, 0); ucp_ep_refcount_assert(ep, probe, ==, 0); ucs_assert(ucs_hlist_is_empty(&ep->ext->proto_reqs)); - ucs_assert(ucs_hlist_is_empty(&ep->ext->rndv_mtype_fc_reqs)); if (!(ep->flags & UCP_EP_FLAG_INTERNAL)) { ucs_assert(worker->num_all_eps > 0); diff --git a/src/ucp/core/ucp_ep.h b/src/ucp/core/ucp_ep.h index 3b97c05cfb6..24a801a13c0 100644 --- a/src/ucp/core/ucp_ep.h +++ b/src/ucp/core/ucp_ep.h @@ -578,8 +578,6 @@ typedef struct ucp_ep_ext { used by 2-stage ppln rndv proto */ /* List of requests which are waiting for remote completion */ ucs_hlist_head_t proto_reqs; - /* List of rendezvous mtype requests waiting for a fragment */ - ucs_hlist_head_t rndv_mtype_fc_reqs; #if UCS_ENABLE_ASSERT ucs_time_t ka_last_round; /* Time of last KA round done */ #endif diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 710639b6a71..f5a8877164e 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -211,11 +211,6 @@ struct ucp_request { * Used by rkey_ptr to track copied data size */ ssize_t completed_size; - - /* Element in per-EP list of rndv mtype requests throttled - * while waiting for a fragment. Used only before - * UCP_REQUEST_FLAG_PROTO_INITIALIZED is set. */ - ucs_hlist_link_t rndv_fc_ep_list; }; } state; diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index 238dceedee0..93c9ade651c 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -18,11 +18,7 @@ void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req) { - ucp_ep_h ep = req->send.ep; - ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC); - ucs_hlist_del(&ep->ext->rndv_mtype_fc_reqs, - &req->send.state.rndv_fc_ep_list); req->flags &= ~UCP_REQUEST_FLAG_RNDV_MTYPE_FC; } @@ -51,10 +47,30 @@ int ucp_proto_rndv_mtype_fc_reschedule_filter( return 1; } -/* Move all queued requests of the endpoint out of the worker pending queues - * and out of the flow-control state into @a reqs, in a single pass. Requests - * which were already woken up are not queued and stay on the endpoint list. */ -static void +static int +ucp_proto_rndv_mtype_fc_resched_collect(const ucs_callbackq_elem_t *elem, + void *arg) +{ + ucs_queue_head_t *reqs = arg; + ucp_request_t *req; + + if (elem->cb != ucp_proto_rndv_mtype_fc_reschedule_cb) { + return 0; + } + + req = elem->arg; + ucs_assert(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC); + ucs_assert(req->send.rndv.fc.state == UCP_REQUEST_RNDV_MTYPE_FC_RESCHED); + ucs_queue_push(reqs, &req->send.rndv.fc.queue_elem); + return 1; +} + +/* Move all throttled requests of the endpoint into @a reqs. Queued requests + * are removed from the worker pending queues and leave the flow-control state. + * Woken-up requests have their reschedule callback removed, but stay in the + * flow-control state, so that aborting or resetting them passes their wakeup + * to the next waiter. Returns nonzero if any request was found. */ +static int ucp_proto_rndv_mtype_fc_ep_dequeue(ucp_ep_h ep, ucs_queue_head_t *reqs) { ucp_worker_h worker = ep->worker; @@ -76,29 +92,25 @@ ucp_proto_rndv_mtype_fc_ep_dequeue(ucp_ep_h ep, ucs_queue_head_t *reqs) ucs_queue_push(reqs, &req->send.rndv.fc.queue_elem); } } + + ucs_callbackq_remove_oneshot(&worker->uct->progress_q, ep, + ucp_proto_rndv_mtype_fc_resched_collect, + reqs); + return !ucs_queue_is_empty(reqs); } void ucp_proto_rndv_mtype_fc_ep_purge(ucp_ep_h ep, ucs_status_t status) { - ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; ucs_queue_head_t reqs; ucp_request_t *req; - ucs_queue_head_init(&reqs); - ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs); - ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { - ucp_proto_request_abort(req, status); - } - /* Aborting a woken-up request may wake up another request of this ep, * which is then aborted on a later iteration */ - while (!ucs_hlist_is_empty(fc_reqs)) { - req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.state.rndv_fc_ep_list); - ucp_proto_request_abort(req, status); - ucs_assert(ucs_hlist_is_empty(fc_reqs) || - (ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.state.rndv_fc_ep_list) != req)); + ucs_queue_head_init(&reqs); + while (ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs)) { + ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { + ucp_proto_request_abort(req, status); + } } } @@ -122,25 +134,16 @@ ucp_proto_rndv_mtype_fc_extract_one(ucp_request_t *req, void ucp_proto_rndv_mtype_fc_ep_extract(ucp_ep_h ep, ucs_queue_head_t *replay_queue) { - ucs_hlist_head_t *fc_reqs = &ep->ext->rndv_mtype_fc_reqs; ucs_queue_head_t reqs; ucp_request_t *req; - ucs_queue_head_init(&reqs); - ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs); - ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { - ucp_proto_rndv_mtype_fc_extract_one(req, replay_queue); - } - /* Resetting a woken-up request may wake up another request of this ep, * which is then extracted on a later iteration */ - while (!ucs_hlist_is_empty(fc_reqs)) { - req = ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.state.rndv_fc_ep_list); - ucp_proto_rndv_mtype_fc_extract_one(req, replay_queue); - ucs_assert(ucs_hlist_is_empty(fc_reqs) || - (ucs_hlist_head_elem(fc_reqs, ucp_request_t, - send.state.rndv_fc_ep_list) != req)); + ucs_queue_head_init(&reqs); + while (ucp_proto_rndv_mtype_fc_ep_dequeue(ep, &reqs)) { + ucs_queue_for_each_extract(req, &reqs, send.rndv.fc.queue_elem, 1) { + ucp_proto_rndv_mtype_fc_extract_one(req, replay_queue); + } } } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index 13a03b52592..ff2c2b1aedc 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -157,14 +157,10 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, frag_sys_dev); UCP_WORKER_STAT_RNDV(worker, MTYPE_FC_THROTTLED, 1); ucs_assert(!(req->flags & UCP_REQUEST_FLAG_RNDV_MTYPE_FC)); - /* The EP list link aliases send.state fields used once initialized */ - ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); req->flags |= UCP_REQUEST_FLAG_RNDV_MTYPE_FC; req->send.rndv.fc.state = UCP_REQUEST_RNDV_MTYPE_FC_QUEUED; ucs_queue_push(&worker->rndv_mtype_fc.pending_q[fc_op], &req->send.rndv.fc.queue_elem); - ucs_hlist_add_tail(&ep->ext->rndv_mtype_fc_reqs, - &req->send.state.rndv_fc_ep_list); return UCS_ERR_NO_RESOURCE; } diff --git a/test/gtest/common/test_obj_size.cc b/test/gtest/common/test_obj_size.cc index 5bc0016ddee..8ddbe8bf3b0 100644 --- a/test/gtest/common/test_obj_size.cc +++ b/test/gtest/common/test_obj_size.cc @@ -50,7 +50,7 @@ UCS_TEST_F(test_obj_size, size) { UCS_TEST_SKIP_R("Assert enabled"); #else EXPECTED_SIZE(ucp_ep_t, 64); - EXPECTED_SIZE(ucp_ep_ext_t, 224); + EXPECTED_SIZE(ucp_ep_ext_t, 216); #if ENABLE_PARAMS_CHECK EXPECTED_SIZE(ucp_rkey_t, 24 + sizeof(ucp_ep_h)); #else From 99d72b29f41d6a036bb0020b6b3496a72279e482 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Thu, 8 Oct 2026 11:52:07 +0300 Subject: [PATCH 60/67] UCP/RNDV: Add rndv mtype flow-control tests for request reset and GET abort --- test/gtest/ucp/test_ucp_am.cc | 17 ++++++++ test/gtest/ucp/test_ucp_peer_failure.cc | 58 ++++++++++++++++++++++++- 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 8aa436c28fa..ad0d2bd359e 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2424,6 +2424,23 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_cap_reached, verify_clean_fc_state(); } +/* Every replayed request is restarted, so the reset of an initialized + * rndv/rtr/mtype request releases its fragment and wakes up a waiter */ +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_request_reset, + "PROTO_REQUEST_RESET=y", "RNDV_FRAG_SIZE=cuda:256K", + "RNDV_FRAG_ALLOC_COUNT=cuda:4", "RNDV_FRAG_WORKER_MAX_MEM=1M", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + fc_counters fc; + + run_fc_test(16, fc); + + EXPECT_GT(fc.sender_throttled + fc.receiver_throttled, 0u) + << "throttling should have occurred with MAX_MEM=1M"; + + verify_clean_fc_state(); +} + UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, "RNDV_FRAG_MEM_TYPE=cuda") { diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 93816c3f8e0..915ad218b78 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -83,6 +83,16 @@ class test_ucp_peer_failure : public ucp_test { void init_buffers(size_t msg_size); virtual ucs_memory_type_t memtype() const; + virtual ucs_memory_type_t tx_memtype() const + { + return memtype(); + } + + virtual ucs_memory_type_t rx_memtype() const + { + return memtype(); + } + void *m_sreq, *m_rreq; std::queue m_am_rndv_descs; size_t m_am_rx_count; @@ -484,8 +494,8 @@ void test_ucp_peer_failure::cleanup_rndv_descs() { } void test_ucp_peer_failure::init_buffers(size_t msg_size) { - m_sbuf.reset(new mem_buffer(msg_size, memtype(), TX_SEED)); - m_rbuf.reset(new mem_buffer(msg_size, memtype(), RX_SEED)); + m_sbuf.reset(new mem_buffer(msg_size, tx_memtype(), TX_SEED)); + m_rbuf.reset(new mem_buffer(msg_size, rx_memtype(), RX_SEED)); } ucs_memory_type_t test_ucp_peer_failure::memtype() const @@ -1166,3 +1176,47 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_fc_pending, } UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); + + +/* The sender buffer is in host memory, so the RTS carries a remote key and + * the receiver can read the data. The receiver buffer is in managed memory, + * which the network lanes cannot register, so the receiver stages the data + * through a fragment using rndv/get/mtype. The message fits in a single + * fragment, so the pipeline protocol is not used. */ +class test_ucp_peer_failure_get_mtype_abort : + public test_ucp_peer_failure_rndv_abort { +public: + static void get_test_variants(variant_vec_t &variants) + { + if (!mem_buffer::is_gpu_supported()) { + return; + } + + test_ucp_peer_failure_rndv_abort::get_test_variants(variants); + } + + ucs_memory_type_t rx_memtype() const override + { + return UCS_MEMORY_TYPE_CUDA_MANAGED; + } + + void init() override + { + /* rndv/get/mtype is enabled by the get_ppln scheme only */ + modify_config("RNDV_SCHEME", "get_ppln"); + modify_config("RNDV_PIPELINE_SHM_ENABLE", "n"); + modify_config("RNDV_PIPELINE_ERROR_HANDLING", "y"); + test_ucp_peer_failure_rndv_abort::init(); + } +}; + +UCS_TEST_P(test_ucp_peer_failure_get_mtype_abort, get_mtype_fc_pending, + "RNDV_FRAG_SIZE=host:32K,cuda:32K", + "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=32K", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + rndv_fc_pending_abort_test(receiver(), UCP_WORKER_RNDV_FC_OP_GET); +} + +UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_get_mtype_abort); From dfe0c1d699e056825296cec17b1a3eeadd0a7d57 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Fri, 9 Oct 2026 13:45:20 +0300 Subject: [PATCH 61/67] UCP/RNDV: Add bidirectional rndv mtype flow-control test --- test/gtest/ucp/test_ucp_am.cc | 129 ++++++++++++++++++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 2d47fbdf26d..edb9d665c06 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -4,6 +4,7 @@ * Copyright (C) Los Alamos National Security, LLC. 2018. ALL RIGHTS RESERVED. * */ +#include #include #include #include @@ -2442,7 +2443,122 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { fc.receiver_throttled = get_stats(receiver(), cntr); } + /* Both workers send to each other at the same time, so each of them holds + * RTR fragments for the peer's messages while needing fragments for its + * own PUTs. Without fragments kept for PUT/GET, the pools of both workers + * could fill up with RTR fragments and wait for each other forever. */ + void run_bidirectional_fc_test(size_t num_frags, size_t num_msgs) + { + const ucs_memory_type_t mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED; + const size_t size = get_rndv_frag_size(UCS_MEMORY_TYPE_CUDA) * + num_frags; + bidirectional_side sides[2] = {{&sender(), mem_type, {}, 0}, + {&receiver(), mem_type, {}, 0}}; + std::vector sbufs, sreqs; + ucp_request_param_t param; + + if (!sender().is_rndv_put_ppln_supported()) { + UCS_TEST_SKIP_R("RNDV is not supported"); + } + + for (auto &side : sides) { + set_am_data_handler(*side.e, TEST_AM_NBX_ID, + am_bidirectional_data_cb, &side); + } + + param.op_attr_mask = 0; + for (size_t i = 0; i < num_msgs; ++i) { + for (auto &side : sides) { + void *sbuf = mem_buffer::allocate(size, mem_type); + mem_buffer::pattern_fill(sbuf, size, SEED, mem_type); + sbufs.push_back(sbuf); + sreqs.push_back(ucp_am_send_nbx(side.e->ep(), TEST_AM_NBX_ID, + NULL, 0, sbuf, size, ¶m)); + } + } + + /* A deadlock shows up as a timeout rather than a hanging test */ + const ucs_time_t deadline = ucs::get_deadline(60.0); + while (ucs_get_time() < deadline) { + progress(); + if ((sides[0].recv_count == num_msgs) && + (sides[1].recv_count == num_msgs) && + std::all_of(sreqs.begin(), sreqs.end(), [](void *req) { + return ucp_request_check_status(req) != UCS_INPROGRESS; + })) { + break; + } + } + + EXPECT_EQ(num_msgs, sides[0].recv_count) << "sender did not receive"; + EXPECT_EQ(num_msgs, sides[1].recv_count) << "receiver did not receive"; + for (void *req : sreqs) { + EXPECT_NE(UCS_INPROGRESS, ucp_request_check_status(req)); + if (UCS_PTR_IS_PTR(req)) { + ucp_request_release(req); + } + } + + for (auto &side : sides) { + for (void *rbuf : side.rx_bufs) { + mem_buffer::pattern_check(rbuf, size, SEED, mem_type); + mem_buffer::release(rbuf, mem_type); + } + } + + for (void *sbuf : sbufs) { + mem_buffer::release(sbuf, mem_type); + } + } + private: + struct bidirectional_side { + entity *e; + ucs_memory_type_t mem_type; + std::vector rx_bufs; + volatile size_t recv_count; + }; + + static void am_bidirectional_recv_cb(void *request, ucs_status_t status, + size_t length, void *user_data) + { + EXPECT_UCS_OK(status); + ++static_cast(user_data)->recv_count; + } + + static ucs_status_t am_bidirectional_data_cb( + void *arg, const void *header, + size_t header_length, void *data, + size_t length, + const ucp_am_recv_param_t *rx_param) + { + bidirectional_side *side = static_cast(arg); + ucp_request_param_t param; + ucs_status_ptr_t sp; + void *rbuf; + + EXPECT_TRUE(rx_param->recv_attr & UCP_AM_RECV_ATTR_FLAG_RNDV); + + rbuf = mem_buffer::allocate(length, side->mem_type); + side->rx_bufs.push_back(rbuf); + + param.op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | + UCP_OP_ATTR_FIELD_USER_DATA; + param.cb.recv_am = am_bidirectional_recv_cb; + param.user_data = side; + + sp = ucp_am_recv_data_nbx(side->e->worker(), data, rbuf, length, + ¶m); + if (UCS_PTR_IS_PTR(sp)) { + ucp_request_release(sp); + return UCS_INPROGRESS; + } + + EXPECT_EQ(NULL, sp); + ++side->recv_count; + return UCS_OK; + } + static void check_pending_queues_empty(const entity &e) { ucp_worker_h worker = e.worker(); @@ -2492,6 +2608,19 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_request_reset, verify_clean_fc_state(); } +UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_enabled_bidirectional, + "RNDV_FRAG_SIZE=cuda:256K", "RNDV_FRAG_ALLOC_COUNT=cuda:4", + "RNDV_FRAG_WORKER_MAX_MEM=1M", "RNDV_FRAG_MEM_TYPE=cuda") +{ + /* 4 messages of 16 fragments in each direction against a 4-fragment + * quota per worker */ + run_bidirectional_fc_test(16, 4); + + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED, 1); + verify_clean_fc_state(); +} + UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, "RNDV_FRAG_MEM_TYPE=cuda") { From 1c38b2ade4f01e6d0f095cadae635e89eae79b5e Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Fri, 9 Oct 2026 13:59:40 +0300 Subject: [PATCH 62/67] UCP/RNDV: Abort a throttled rndv mtype request on fragment allocation failure --- test/gtest/ucp/test_ucp_am.cc | 2 +- test/gtest/ucp/test_ucp_peer_failure.cc | 162 ++++++++++++++++++++++++ 2 files changed, 163 insertions(+), 1 deletion(-) diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index edb9d665c06..13e2215f983 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2520,7 +2520,7 @@ class test_ucp_am_nbx_rndv_mtype_fc : public test_ucp_am_nbx_rndv_ppln { }; static void am_bidirectional_recv_cb(void *request, ucs_status_t status, - size_t length, void *user_data) + size_t length, void *user_data) { EXPECT_UCS_OK(status); ++static_cast(user_data)->recv_count; diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index da912a2134e..b5ad8229980 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -889,9 +889,79 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { close_fc_pending_ep_cb, self); } + if (self->m_close_on_alloc_failure && self->m_alloc_failed && + (worker == self->m_fc_pending_entity->worker())) { + self->m_close_on_alloc_failure = false; + /* The request was aborted locally, the peer is still waiting for + * it. Close the eps so the peer fails as well. */ + ucs_callbackq_add_oneshot(&worker->uct->progress_q, self, + close_fc_pending_ep_cb, self); + } + return status; } + static ucs_status_t fail_chunk_alloc(ucs_mpool_t *mp, size_t *size_p, + void **chunk_p) + { + self->m_alloc_failed = true; + return UCS_ERR_NO_MEMORY; + } + + /* Make every CUDA fragment mpool of the worker fail to grow */ + void fail_cuda_fragment_alloc(ucp_worker_h worker) + { + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (khiter_t iter = kh_begin(&worker->mpool_hash); + iter != kh_end(&worker->mpool_hash); ++iter) { + if (!kh_exist(&worker->mpool_hash, iter) || + (kh_key(&worker->mpool_hash, iter).mem_type != + UCS_MEMORY_TYPE_CUDA)) { + continue; + } + + ucs_mpool_t *mpool = &kh_val(&worker->mpool_hash, iter); + m_orig_mpool_ops.emplace_back(mpool, mpool->data->ops); + m_fail_mpool_ops = *mpool->data->ops; + m_fail_mpool_ops.chunk_alloc = fail_chunk_alloc; + mpool->data->ops = &m_fail_mpool_ops; + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); + } + + void restore_cuda_fragment_alloc(ucp_worker_h worker) + { + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (auto &mpool_ops : m_orig_mpool_ops) { + mpool_ops.first->data->ops = mpool_ops.second; + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); + + m_orig_mpool_ops.clear(); + } + + /* Whether the next fragment allocation must grow a failing pool: with no + * free fragment, a pool without quota would throttle the request instead, + * and a pool created later would not fail */ + bool cuda_fragment_alloc_fails(ucp_worker_h worker) const + { + bool fails = !m_orig_mpool_ops.empty(); + + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (auto &mpool_ops : m_orig_mpool_ops) { + fails = fails && !ucs_mpool_is_empty(mpool_ops.first); + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); + + return fails; + } + static bool has_fc_pending(ucp_worker_h worker, unsigned fc_op) { return !ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[fc_op]); @@ -1051,9 +1121,78 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { EXPECT_FALSE(has_fc_pending(fc_entity.worker(), fc_op)); } + /* A fragment allocation which fails for a reason other than the quota + * aborts the request with the allocation error instead of queuing it */ + void rndv_fc_alloc_failure_abort_test(entity &fc_entity, size_t msg_size) + { + const ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); + std::pair result; + std::vector held_mdescs; + + if (sender_config->key.rma_bw_lanes[0] == UCP_NULL_LANE) { + UCS_TEST_SKIP_R("transport has no rma_bw lanes"); + } + + init_buffers(msg_size); + + /* Complete one transfer so the CUDA fragment mpools are created. Then + * make them fail to grow before holding their fragments: holding + * grows a pool up to its quota otherwise, while the next request must + * find no free fragment and quota left, so that it has to grow. */ + smoke_test(true); + fail_cuda_fragment_alloc(fc_entity.worker()); + { + /* Failing to grow a pool is logged as an error */ + scoped_log_handler slh(hide_errors_logger); + held_mdescs = hold_cuda_fragments(fc_entity.worker()); + } + + if (held_mdescs.empty()) { + restore_cuda_fragment_alloc(fc_entity.worker()); + UCS_TEST_SKIP_R("no CUDA fragments were allocated"); + } + + if (!cuda_fragment_alloc_fails(fc_entity.worker())) { + restore_cuda_fragment_alloc(fc_entity.worker()); + release_cuda_fragments(fc_entity.worker(), held_mdescs); + FAIL() << "a CUDA fragment pool has no quota left to grow"; + } + + m_alloc_failed = false; + m_fc_pending_entity = &fc_entity; + m_close_on_alloc_failure = true; + setup_progress_mock(sender().worker(), m_progress_mock); + setup_progress_mock(receiver().worker(), m_progress_mock); + + { + scoped_log_handler err_wrapper(wrap_errors_logger); + scoped_log_handler warn_wrapper(wrap_warns_logger); + result = smoke_test(true); + } + + m_progress_mock.cleanup(); + restore_cuda_fragment_alloc(fc_entity.worker()); + release_cuda_fragments(fc_entity.worker(), held_mdescs); + + EXPECT_TRUE(m_alloc_failed); + EXPECT_TRUE(m_is_peer_closed); + EXPECT_EQ(UCS_ERR_NO_MEMORY, (&fc_entity == &sender()) ? + result.first : result.second); + EXPECT_TRUE(UCS_STATUS_IS_ERR(result.first)); + EXPECT_TRUE(UCS_STATUS_IS_ERR(result.second)); + for (unsigned fc_op = 0; fc_op < UCP_WORKER_RNDV_FC_OP_LAST; ++fc_op) { + EXPECT_FALSE(has_fc_pending(fc_entity.worker(), fc_op)); + } + } + ucs::mock m_progress_mock; bool m_is_peer_closed{false}; bool m_close_fc_pending_ep{false}; + bool m_close_on_alloc_failure{false}; + volatile bool m_alloc_failed{false}; + ucs_mpool_ops_t m_fail_mpool_ops{}; + std::vector> + m_orig_mpool_ops; std::string m_proto_name{}; /* Protocol stage during which data transfer happens */ uint8_t m_proto_xfer_stage{}; @@ -1176,6 +1315,29 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_fc_pending, rndv_fc_pending_abort_test(sender(), UCP_WORKER_RNDV_FC_OP_PUT); } +UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_alloc_failure, + "RNDV_FRAG_SIZE=host:8K,cuda:8K", + "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=32K", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + rndv_fc_alloc_failure_abort_test(receiver(), 16 * UCS_KBYTE); +} + +/* A single fragment, since the abort of a pipelined send does not report the + * error to the send request (see ucp_proto_rndv_stub_abort()). No reserved + * pool, which a single fragment would exhaust while the shared pool would be + * created only after the allocation failure is set up. */ +UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_alloc_failure, + "RNDV_FRAG_SIZE=host:8K,cuda:8K", + "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=32K", + "RNDV_FRAG_RTR_RATIO=1", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + rndv_fc_alloc_failure_abort_test(sender(), 8 * UCS_KBYTE); +} + UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); From 9f8df930281122a9da8914da0f4a1e8fbb5170b6 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Wed, 14 Oct 2026 12:49:50 +0300 Subject: [PATCH 63/67] UCP/RNDV: Add rndv/get/mtype flow-control tests --- test/gtest/Makefile.am | 1 + test/gtest/ucp/rndv_frag_mpools.cc | 121 ++++++++++++++ test/gtest/ucp/rndv_frag_mpools.h | 78 +++++++++ test/gtest/ucp/test_ucp_am.cc | 77 ++++++++- test/gtest/ucp/test_ucp_peer_failure.cc | 211 +++--------------------- test/gtest/ucp/ucp_test.cc | 17 +- test/gtest/ucp/ucp_test.h | 5 + 7 files changed, 309 insertions(+), 201 deletions(-) create mode 100644 test/gtest/ucp/rndv_frag_mpools.cc create mode 100644 test/gtest/ucp/rndv_frag_mpools.h diff --git a/test/gtest/Makefile.am b/test/gtest/Makefile.am index fa8db098f87..700ac735777 100644 --- a/test/gtest/Makefile.am +++ b/test/gtest/Makefile.am @@ -181,6 +181,7 @@ gtest_SOURCES = \ ucp/test_ucp_sockaddr.cc \ ucp/ucp_test.cc \ ucp/ucp_datatype.cc \ + ucp/rndv_frag_mpools.cc \ \ ucs/test_algorithm.cc \ ucs/test_arbiter.cc \ diff --git a/test/gtest/ucp/rndv_frag_mpools.cc b/test/gtest/ucp/rndv_frag_mpools.cc new file mode 100644 index 00000000000..e88ec249451 --- /dev/null +++ b/test/gtest/ucp/rndv_frag_mpools.cc @@ -0,0 +1,121 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#include "rndv_frag_mpools.h" + +extern "C" { +#include +#include +} + + +unsigned rndv_frag_mpools::s_num_alloc_failures = 0; + +rndv_frag_mpools::rndv_frag_mpools(ucp_worker_h worker, + ucs_memory_type_t mem_type) : + m_worker(worker) +{ + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + + for (khiter_t iter = kh_begin(&worker->mpool_hash); + iter != kh_end(&worker->mpool_hash); ++iter) { + if (kh_exist(&worker->mpool_hash, iter) && + (kh_key(&worker->mpool_hash, iter).mem_type == mem_type)) { + m_mpools.push_back(&kh_val(&worker->mpool_hash, iter)); + } + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); +} + +rndv_frag_mpools::~rndv_frag_mpools() +{ + restore_alloc(); + release(); +} + +void rndv_frag_mpools::hold() +{ + ucp_mem_desc_t *mdesc; + + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); + + for (ucs_mpool_t *mpool : m_mpools) { + /* Terminates only because the mpool quota is finite, which requires + * proto v2 */ + while ((mdesc = static_cast( + ucs_mpool_get_inline(mpool))) != nullptr) { + m_held.push_back(mdesc); + } + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); +} + +void rndv_frag_mpools::release() +{ + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); + + for (ucp_mem_desc_t *mdesc : m_held) { + /* Return the fragment without the flow-control wakeup: no request + * may be left waiting for it when the test completes */ + ucs_mpool_put_inline(mdesc); + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); + + m_held.clear(); +} + +ucs_status_t rndv_frag_mpools::fail_chunk_alloc(ucs_mpool_t *mp, size_t *size_p, + void **chunk_p) +{ + ++s_num_alloc_failures; + return UCS_ERR_NO_MEMORY; +} + +void rndv_frag_mpools::fail_alloc() +{ + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); + + for (ucs_mpool_t *mpool : m_mpools) { + m_orig_ops.push_back(mpool->data->ops); + m_fail_ops = *mpool->data->ops; + m_fail_ops.chunk_alloc = fail_chunk_alloc; + mpool->data->ops = &m_fail_ops; + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); +} + +void rndv_frag_mpools::restore_alloc() +{ + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); + + for (size_t i = 0; i < m_orig_ops.size(); ++i) { + m_mpools[i]->data->ops = m_orig_ops[i]; + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); + + m_orig_ops.clear(); +} + +bool rndv_frag_mpools::can_grow() const +{ + bool can_grow = !m_mpools.empty(); + + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); + + for (ucs_mpool_t *mpool : m_mpools) { + /* With no free fragment, empty means the quota is exhausted */ + can_grow = can_grow && !ucs_mpool_is_empty(mpool); + } + + UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); + + return can_grow; +} diff --git a/test/gtest/ucp/rndv_frag_mpools.h b/test/gtest/ucp/rndv_frag_mpools.h new file mode 100644 index 00000000000..fa7a2140c92 --- /dev/null +++ b/test/gtest/ucp/rndv_frag_mpools.h @@ -0,0 +1,78 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifndef UCP_TEST_RNDV_FRAG_MPOOLS_H_ +#define UCP_TEST_RNDV_FRAG_MPOOLS_H_ + +#include + +extern "C" { +#include +#include +} + +#include + + +/** + * Controls the rendezvous fragment memory pools of one memory type on a + * worker. Requires the pools to exist, i.e. a transfer which used fragments of + * that memory type must have completed on the worker. + */ +class rndv_frag_mpools { +public: + rndv_frag_mpools(ucp_worker_h worker, ucs_memory_type_t mem_type); + ~rndv_frag_mpools(); + + /* Take every free fragment. Grows the pools up to their quota, unless + * growing was made to fail. */ + void hold(); + + /* Return the held fragments without waking up throttled requests */ + void release(); + + /* Make the pools fail to grow with UCS_ERR_NO_MEMORY */ + void fail_alloc(); + void restore_alloc(); + + /* Whether every pool can still grow, so that fail_alloc() takes effect + * for the next allocation which finds no free fragment */ + bool can_grow() const; + + size_t num_pools() const + { + return m_mpools.size(); + } + + size_t num_held() const + { + return m_held.size(); + } + + /* Number of failed growth attempts since the last reset */ + static unsigned num_alloc_failures() + { + return s_num_alloc_failures; + } + + static void reset_alloc_failures() + { + s_num_alloc_failures = 0; + } + +private: + static ucs_status_t fail_chunk_alloc(ucs_mpool_t *mp, size_t *size_p, + void **chunk_p); + + static unsigned s_num_alloc_failures; + ucp_worker_h m_worker; + std::vector m_mpools; + std::vector m_held; + std::vector m_orig_ops; + ucs_mpool_ops_t m_fail_ops; +}; + +#endif diff --git a/test/gtest/ucp/test_ucp_am.cc b/test/gtest/ucp/test_ucp_am.cc index 13e2215f983..262a725198f 100644 --- a/test/gtest/ucp/test_ucp_am.cc +++ b/test/gtest/ucp/test_ucp_am.cc @@ -2346,6 +2346,11 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { << "counter is " << stats_name(e, cntr); } + ucp_err_handling_mode_t get_err_mode() const + { + return static_cast(get_variant_value(1)); + } + private: ucs_memory_type_t tx_memtype() const override { @@ -2357,11 +2362,6 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { return m_mem_type; } - ucp_err_handling_mode_t get_err_mode() const - { - return static_cast(get_variant_value(1)); - } - ucs_memory_type_t m_mem_type; }; @@ -2637,6 +2637,73 @@ UCS_TEST_P(test_ucp_am_nbx_rndv_mtype_fc, fc_disabled, UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_mtype_fc); + +/* The sender buffer is in host memory, so the RTS carries a remote key and + * the receiver can read the data. The receiver buffer is in managed memory, + * which the network lanes cannot register, so the receiver stages the data + * through a fragment using rndv/get/mtype, which is enabled by the get_ppln + * scheme only and does not support error handling. */ +class test_ucp_am_nbx_rndv_get_mtype_fc : + public test_ucp_am_nbx_rndv_mtype_fc { +public: + void init() override + { + if (get_err_mode() != UCP_ERR_HANDLING_MODE_NONE) { + UCS_TEST_SKIP_R("rndv/get/mtype does not support error handling"); + } + + test_ucp_am_nbx_rndv_mtype_fc::init(); + } + +protected: + ucs_memory_type_t tx_memtype() const override + { + return UCS_MEMORY_TYPE_HOST; + } + + ucs_memory_type_t rx_memtype() const override + { + return UCS_MEMORY_TYPE_CUDA_MANAGED; + } + + /* Only the receiver uses fragments: the sender buffer is read directly */ + void throttle_test(size_t num_frags) + { + const uint64_t cntr = UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED; + + /* The receiver reads the sender buffer, which is in host memory */ + if (!receiver().is_rndv_zcopy_supported(UCT_IFACE_FLAG_GET_ZCOPY, + UCS_MEMORY_TYPE_HOST)) { + UCS_TEST_SKIP_R("rndv/get/mtype is not supported"); + } + + test_am_send_recv(get_rndv_frag_size(UCS_MEMORY_TYPE_CUDA) * + num_frags); + + if (get_stats(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE) > 0) { + UCS_TEST_SKIP_R("rndv/get/mtype was not selected"); + } + + EXPECT_EQ(0u, get_stats(sender(), + UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY)); + EXPECT_EQ(0u, get_stats(sender(), cntr)); + EXPECT_GT(get_stats(receiver(), cntr), 0u) + << "the receiver should have been throttled"; + verify_clean_fc_state(); + } +}; + +UCS_TEST_P(test_ucp_am_nbx_rndv_get_mtype_fc, fc_enabled_cap_reached, + "RNDV_SCHEME=get_ppln", "RNDV_FRAG_SIZE=cuda:256K", + "RNDV_FRAG_ALLOC_COUNT=cuda:4", "RNDV_FRAG_WORKER_MAX_MEM=1M", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + /* 16 fragments against a 4-fragment quota */ + throttle_test(16); +} + +UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_am_nbx_rndv_get_mtype_fc); + #endif /* Test class for AM PSN protocol */ diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index b5ad8229980..ea318c2bacb 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -8,6 +8,7 @@ #include "ucp_test.h" #include "ucp_datatype.h" +#include "rndv_frag_mpools.h" extern "C" { #include /* for testing EP RNDV configuration */ @@ -83,16 +84,6 @@ class test_ucp_peer_failure : public ucp_test { void init_buffers(size_t msg_size); virtual ucs_memory_type_t memtype() const; - virtual ucs_memory_type_t tx_memtype() const - { - return memtype(); - } - - virtual ucs_memory_type_t rx_memtype() const - { - return memtype(); - } - void *m_sreq, *m_rreq; std::queue m_am_rndv_descs; size_t m_am_rx_count; @@ -494,8 +485,8 @@ void test_ucp_peer_failure::cleanup_rndv_descs() { } void test_ucp_peer_failure::init_buffers(size_t msg_size) { - m_sbuf.reset(new mem_buffer(msg_size, tx_memtype(), TX_SEED)); - m_rbuf.reset(new mem_buffer(msg_size, rx_memtype(), RX_SEED)); + m_sbuf.reset(new mem_buffer(msg_size, memtype(), TX_SEED)); + m_rbuf.reset(new mem_buffer(msg_size, memtype(), RX_SEED)); } ucs_memory_type_t test_ucp_peer_failure::memtype() const @@ -889,7 +880,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { close_fc_pending_ep_cb, self); } - if (self->m_close_on_alloc_failure && self->m_alloc_failed && + if (self->m_close_on_alloc_failure && + (rndv_frag_mpools::num_alloc_failures() > 0) && (worker == self->m_fc_pending_entity->worker())) { self->m_close_on_alloc_failure = false; /* The request was aborted locally, the peer is still waiting for @@ -901,122 +893,11 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { return status; } - static ucs_status_t fail_chunk_alloc(ucs_mpool_t *mp, size_t *size_p, - void **chunk_p) - { - self->m_alloc_failed = true; - return UCS_ERR_NO_MEMORY; - } - - /* Make every CUDA fragment mpool of the worker fail to grow */ - void fail_cuda_fragment_alloc(ucp_worker_h worker) - { - UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - - for (khiter_t iter = kh_begin(&worker->mpool_hash); - iter != kh_end(&worker->mpool_hash); ++iter) { - if (!kh_exist(&worker->mpool_hash, iter) || - (kh_key(&worker->mpool_hash, iter).mem_type != - UCS_MEMORY_TYPE_CUDA)) { - continue; - } - - ucs_mpool_t *mpool = &kh_val(&worker->mpool_hash, iter); - m_orig_mpool_ops.emplace_back(mpool, mpool->data->ops); - m_fail_mpool_ops = *mpool->data->ops; - m_fail_mpool_ops.chunk_alloc = fail_chunk_alloc; - mpool->data->ops = &m_fail_mpool_ops; - } - - UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); - } - - void restore_cuda_fragment_alloc(ucp_worker_h worker) - { - UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - - for (auto &mpool_ops : m_orig_mpool_ops) { - mpool_ops.first->data->ops = mpool_ops.second; - } - - UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); - - m_orig_mpool_ops.clear(); - } - - /* Whether the next fragment allocation must grow a failing pool: with no - * free fragment, a pool without quota would throttle the request instead, - * and a pool created later would not fail */ - bool cuda_fragment_alloc_fails(ucp_worker_h worker) const - { - bool fails = !m_orig_mpool_ops.empty(); - - UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - - for (auto &mpool_ops : m_orig_mpool_ops) { - fails = fails && !ucs_mpool_is_empty(mpool_ops.first); - } - - UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); - - return fails; - } - static bool has_fc_pending(ucp_worker_h worker, unsigned fc_op) { return !ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[fc_op]); } - static std::vector - hold_cuda_fragments(ucp_worker_h worker) - { - std::vector held; - ucp_mem_desc_t *mdesc; - - UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - - for (khiter_t iter = kh_begin(&worker->mpool_hash); - iter != kh_end(&worker->mpool_hash); ++iter) { - if (!kh_exist(&worker->mpool_hash, iter) || - (kh_key(&worker->mpool_hash, iter).mem_type != - UCS_MEMORY_TYPE_CUDA)) { - continue; - } - - /* Terminates only because the mpool quota is finite, which - * requires proto v2 */ - while ((mdesc = static_cast( - ucs_mpool_get_inline( - &kh_val(&worker->mpool_hash, iter)))) != - nullptr) { - held.push_back(mdesc); - } - } - - UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); - - return held; - } - - static void - release_cuda_fragments(ucp_worker_h worker, - std::vector &held) - { - UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); - - for (auto *mdesc : held) { - /* - * Return the descriptor without invoking FC rescheduling. Endpoint - * purge must already have removed requests from pending_q. - */ - ucs_mpool_put_inline(mdesc); - } - - UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker); - - held.clear(); - } - virtual void cleanup() { test_ucp_peer_failure::cleanup(); @@ -1094,9 +975,10 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { /* Complete one transfer so the CUDA fragment mpool is created. */ smoke_test(true); - std::vector held_mdescs = - hold_cuda_fragments(fc_entity.worker()); - if (held_mdescs.empty()) { + rndv_frag_mpools frag_mpools(fc_entity.worker(), + UCS_MEMORY_TYPE_CUDA); + frag_mpools.hold(); + if (frag_mpools.num_held() == 0) { UCS_TEST_SKIP_R("no CUDA fragments were allocated"); } @@ -1112,7 +994,7 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { result = smoke_test(true); } - release_cuda_fragments(fc_entity.worker(), held_mdescs); + frag_mpools.release(); m_progress_mock.cleanup(); EXPECT_TRUE(UCS_STATUS_IS_ERR(result.first)); @@ -1127,7 +1009,6 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { { const ucp_ep_config_t *sender_config = ucp_ep_config(sender().ep()); std::pair result; - std::vector held_mdescs; if (sender_config->key.rma_bw_lanes[0] == UCP_NULL_LANE) { UCS_TEST_SKIP_R("transport has no rma_bw lanes"); @@ -1140,25 +1021,23 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { * grows a pool up to its quota otherwise, while the next request must * find no free fragment and quota left, so that it has to grow. */ smoke_test(true); - fail_cuda_fragment_alloc(fc_entity.worker()); + rndv_frag_mpools frag_mpools(fc_entity.worker(), + UCS_MEMORY_TYPE_CUDA); + frag_mpools.fail_alloc(); { /* Failing to grow a pool is logged as an error */ scoped_log_handler slh(hide_errors_logger); - held_mdescs = hold_cuda_fragments(fc_entity.worker()); + frag_mpools.hold(); } - if (held_mdescs.empty()) { - restore_cuda_fragment_alloc(fc_entity.worker()); + if (frag_mpools.num_held() == 0) { UCS_TEST_SKIP_R("no CUDA fragments were allocated"); } - if (!cuda_fragment_alloc_fails(fc_entity.worker())) { - restore_cuda_fragment_alloc(fc_entity.worker()); - release_cuda_fragments(fc_entity.worker(), held_mdescs); - FAIL() << "a CUDA fragment pool has no quota left to grow"; - } + ASSERT_TRUE(frag_mpools.can_grow()) + << "a CUDA fragment pool has no quota left to grow"; - m_alloc_failed = false; + rndv_frag_mpools::reset_alloc_failures(); m_fc_pending_entity = &fc_entity; m_close_on_alloc_failure = true; setup_progress_mock(sender().worker(), m_progress_mock); @@ -1171,10 +1050,10 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } m_progress_mock.cleanup(); - restore_cuda_fragment_alloc(fc_entity.worker()); - release_cuda_fragments(fc_entity.worker(), held_mdescs); + frag_mpools.restore_alloc(); + frag_mpools.release(); - EXPECT_TRUE(m_alloc_failed); + EXPECT_GT(rndv_frag_mpools::num_alloc_failures(), 0u); EXPECT_TRUE(m_is_peer_closed); EXPECT_EQ(UCS_ERR_NO_MEMORY, (&fc_entity == &sender()) ? result.first : result.second); @@ -1189,10 +1068,6 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { bool m_is_peer_closed{false}; bool m_close_fc_pending_ep{false}; bool m_close_on_alloc_failure{false}; - volatile bool m_alloc_failed{false}; - ucs_mpool_ops_t m_fail_mpool_ops{}; - std::vector> - m_orig_mpool_ops; std::string m_proto_name{}; /* Protocol stage during which data transfer happens */ uint8_t m_proto_xfer_stage{}; @@ -1339,47 +1214,3 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_alloc_failure, } UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_rndv_put_ppln_abort); - - -/* The sender buffer is in host memory, so the RTS carries a remote key and - * the receiver can read the data. The receiver buffer is in managed memory, - * which the network lanes cannot register, so the receiver stages the data - * through a fragment using rndv/get/mtype. The message fits in a single - * fragment, so the pipeline protocol is not used. */ -class test_ucp_peer_failure_get_mtype_abort : - public test_ucp_peer_failure_rndv_abort { -public: - static void get_test_variants(variant_vec_t &variants) - { - if (!mem_buffer::is_gpu_supported()) { - return; - } - - test_ucp_peer_failure_rndv_abort::get_test_variants(variants); - } - - ucs_memory_type_t rx_memtype() const override - { - return UCS_MEMORY_TYPE_CUDA_MANAGED; - } - - void init() override - { - /* rndv/get/mtype is enabled by the get_ppln scheme only */ - modify_config("RNDV_SCHEME", "get_ppln"); - modify_config("RNDV_PIPELINE_SHM_ENABLE", "n"); - modify_config("RNDV_PIPELINE_ERROR_HANDLING", "y"); - test_ucp_peer_failure_rndv_abort::init(); - } -}; - -UCS_TEST_P(test_ucp_peer_failure_get_mtype_abort, get_mtype_fc_pending, - "RNDV_FRAG_SIZE=host:32K,cuda:32K", - "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", - "RNDV_FRAG_WORKER_MAX_MEM=32K", - "RNDV_FRAG_MEM_TYPE=cuda") -{ - rndv_fc_pending_abort_test(receiver(), UCP_WORKER_RNDV_FC_OP_GET); -} - -UCP_INSTANTIATE_TEST_CASE_GPU_AWARE(test_ucp_peer_failure_get_mtype_abort); diff --git a/test/gtest/ucp/ucp_test.cc b/test/gtest/ucp/ucp_test.cc index 4958cde6763..65fdd99fd39 100644 --- a/test/gtest/ucp/ucp_test.cc +++ b/test/gtest/ucp/ucp_test.cc @@ -1210,13 +1210,10 @@ bool ucp_test_base::entity::has_lane_with_caps(uint64_t caps, return false; } -bool ucp_test_base::entity::is_rndv_put_ppln_supported() const +bool ucp_test_base::entity::is_rndv_zcopy_supported( + uint64_t iface_cap, ucs_memory_type_t mem_type) const { const auto config = ucp_ep_config(ep()); - ucs_memory_type_t mem_type; - - mem_type = (ucs_memory_type_t)ucs_ffs64( - ucph()->config.ext.rndv_frag_mem_types); for (auto i = 0; i < config->key.num_lanes; ++i) { const auto lane = config->key.rma_bw_lanes[i]; @@ -1227,7 +1224,7 @@ bool ucp_test_base::entity::is_rndv_put_ppln_supported() const const auto iface_attr = ucp_worker_iface_get_attr(worker(), ucp_ep_get_rsc_index(ep(), lane)); - if (iface_attr->cap.flags & UCT_IFACE_FLAG_PUT_ZCOPY) { + if (iface_attr->cap.flags & iface_cap) { return true; } } @@ -1236,6 +1233,14 @@ bool ucp_test_base::entity::is_rndv_put_ppln_supported() const return false; } +bool ucp_test_base::entity::is_rndv_put_ppln_supported() const +{ + const auto mem_type = (ucs_memory_type_t)ucs_ffs64( + ucph()->config.ext.rndv_frag_mem_types); + + return is_rndv_zcopy_supported(UCT_IFACE_FLAG_PUT_ZCOPY, mem_type); +} + bool ucp_test_base::entity::is_rndv_supported() const { const auto config = ucp_ep_config(ep()); diff --git a/test/gtest/ucp/ucp_test.h b/test/gtest/ucp/ucp_test.h index 53bc2c830e9..700045f9d71 100644 --- a/test/gtest/ucp/ucp_test.h +++ b/test/gtest/ucp/ucp_test.h @@ -157,6 +157,11 @@ class ucp_test_base : public ucs::test_base { bool has_lane_with_caps(uint64_t caps, uint64_t v2_caps = 0) const; + /* Whether a rma_bw lane which can register the memory type has + * the given UCT_IFACE_FLAG_*_ZCOPY capability */ + bool is_rndv_zcopy_supported(uint64_t iface_cap, + ucs_memory_type_t mem_type) const; + bool is_rndv_put_ppln_supported() const; bool is_rndv_supported() const; From 45a77f13e3db32f36069ffbfe1225d52df73c136 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Sun, 18 Oct 2026 17:10:32 +0300 Subject: [PATCH 64/67] UCP/RNDV: Pass the wakeup on when a woken PUT/GET takes a reserved fragment --- src/ucp/rndv/rndv_mtype.inl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index cc6c1cf5120..ad2f7e2fda6 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -127,6 +127,12 @@ ucp_proto_rndv_mtype_request_init(ucp_request_t *req, status = ucp_rndv_mpool_get(worker, frag_mem_type, frag_sys_dev, UCP_WORKER_RNDV_FRAG_POOL_RESERVED, &req->send.rndv.mdesc); + if ((status == UCS_OK) && owns_wakeup) { + /* The wakeup may have been for a shared fragment, which is then + * still free, so pass it on rather than leave it unused. */ + ucp_proto_rndv_mtype_fc_reschedule_pending( + worker, UCP_WORKER_RNDV_FC_OP_LAST); + } } if (status == UCS_ERR_NO_RESOURCE) { @@ -334,7 +340,10 @@ ucp_proto_rndv_mtype_mdesc_release(ucp_request_t *req) ucp_rndv_mpool_priv_t *mpriv = ucs_mpool_priv(ucs_mpool_obj_owner(mdesc)); unsigned num_queues; - /* Only PUT/GET requests can use a fragment of the reserved pool */ + /* Only PUT/GET requests can use a fragment of the reserved pool. A woken + * up PUT/GET which takes a reserved fragment passes its wakeup on, see + * ucp_proto_rndv_mtype_request_init(), so a shared fragment freed in the + * meantime is not left without a waiter. */ num_queues = (mpriv->pool == UCP_WORKER_RNDV_FRAG_POOL_RESERVED) ? (UCP_WORKER_RNDV_FC_OP_PUT + 1) : UCP_WORKER_RNDV_FC_OP_LAST; From 21d5bb874e3b3cca3a95155a00a6d71a3d670bfd Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 19 Oct 2026 10:58:12 +0300 Subject: [PATCH 65/67] UCP/RNDV: Always keep a fragment for PUT/GET --- src/ucp/core/ucp_context.c | 14 ++++++------- src/ucp/rndv/proto_rndv.c | 16 ++++++++------ src/ucp/rndv/rndv_mtype.inl | 4 ++-- test/gtest/ucp/rndv_frag_mpools.cc | 28 +++++++++++++++++++++++-- test/gtest/ucp/rndv_frag_mpools.h | 10 +++++---- test/gtest/ucp/test_ucp_mmap.cc | 23 ++++++++++++-------- test/gtest/ucp/test_ucp_peer_failure.cc | 11 +++++----- 7 files changed, 70 insertions(+), 36 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index becd741190f..cfe60954cf6 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -425,18 +425,18 @@ static ucs_config_field_t ucp_context_config_table[] = { "each fragment memory type, and per device when available. It is\n" "converted to a fragment count using RNDV_FRAG_SIZE. When the limit is\n" "reached, further fragment requests are queued until fragments are\n" - "released. See also RNDV_FRAG_RTR_RATIO", + "released. The limit is at least two fragments. See also\n" + "RNDV_FRAG_RTR_RATIO", ucs_offsetof(ucp_context_config_t, rndv_frag_worker_max_mem), UCS_CONFIG_TYPE_MEMUNITS}, {"RNDV_FRAG_RTR_RATIO", "0.8", "Fraction of RNDV_FRAG_WORKER_MAX_MEM which fragments posted to receive\n" - "data (RTR) may use at once, in the range (0, 1]. The rest is kept for\n" + "data (RTR) may use at once, in the range (0, 1). The rest is kept for\n" "the put/get operations which move the data, so that they can always\n" "make progress: an RTR fragment is released only when the peer's data\n" - "arrives, and the peer needs a fragment of its own to send it. A value\n" - "of 1 keeps nothing, which is safe only for a worker that never moves\n" - "rendezvous data through staging fragments", + "arrives, and the peer needs a fragment of its own to send it, so\n" + "without a reserve both peers could wait for each other forever", ucs_offsetof(ucp_context_config_t, rndv_frag_rtr_ratio), UCS_CONFIG_TYPE_POS_DOUBLE}, @@ -2784,9 +2784,9 @@ static ucs_status_t ucp_fill_config(ucp_context_h context, } if (UCS_CONFIG_DBL_IS_AUTO(context->config.ext.rndv_frag_rtr_ratio) || - !(context->config.ext.rndv_frag_rtr_ratio <= 1.0)) { + !(context->config.ext.rndv_frag_rtr_ratio < 1.0)) { ucs_error("UCX_RNDV_FRAG_RTR_RATIO value %g is not in the range " - "(0, 1]", context->config.ext.rndv_frag_rtr_ratio); + "(0, 1)", context->config.ext.rndv_frag_rtr_ratio); status = UCS_ERR_INVALID_PARAM; goto err_free_alloc_methods; } diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index ed1961adb32..875f4e72a0f 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -853,12 +853,13 @@ unsigned ucp_proto_rndv_frag_max_elems(ucp_context_h context, frag_size = context->config.ext.rndv_frag_size[frag_mem_type]; ucs_assert_always(frag_size > 0); + /* One fragment is kept for PUT/GET and at least one is shared */ max_frags = max_mem / frag_size; - if (max_frags == 0) { - ucs_warn("RNDV_FRAG_WORKER_MAX_MEM %zu is smaller than one %s " - "fragment (%zu), using one fragment", + if (max_frags < 2) { + ucs_warn("RNDV_FRAG_WORKER_MAX_MEM %zu is smaller than two %s " + "fragments (%zu each), using two fragments", max_mem, ucs_memory_type_names[frag_mem_type], frag_size); - return 1; + return 2; } return ucs_min(max_frags, UINT_MAX); @@ -870,11 +871,14 @@ unsigned ucp_proto_rndv_frag_shared_elems(ucp_context_h context, const double ratio = context->config.ext.rndv_frag_rtr_ratio; unsigned shared; - /* A reserve needs at least two fragments: one to keep and one to share */ - if ((max_elems == UINT_MAX) || (ratio >= 1.0) || (max_elems < 2)) { + if (max_elems == UINT_MAX) { return max_elems; } + /* Keep at least one fragment for PUT/GET and share at least one, see + * ucp_proto_rndv_frag_max_elems() */ + ucs_assert(ratio < 1.0); + ucs_assert(max_elems >= 2); shared = ucs_min((unsigned)(max_elems * ratio), max_elems - 1); return ucs_max(shared, 1); } diff --git a/src/ucp/rndv/rndv_mtype.inl b/src/ucp/rndv/rndv_mtype.inl index ad2f7e2fda6..0c9c062544d 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -93,8 +93,8 @@ ucp_proto_rndv_mtype_frag_has_reserve(ucp_context_h context) { const size_t max_mem = context->config.ext.rndv_frag_worker_max_mem; - return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO) && - (context->config.ext.rndv_frag_rtr_ratio < 1.0); + /* A reserved pool exists whenever the quota is limited */ + return (max_mem != UCS_MEMUNITS_INF) && (max_mem != UCS_MEMUNITS_AUTO); } static UCS_F_ALWAYS_INLINE ucs_status_t diff --git a/test/gtest/ucp/rndv_frag_mpools.cc b/test/gtest/ucp/rndv_frag_mpools.cc index e88ec249451..880192b8960 100644 --- a/test/gtest/ucp/rndv_frag_mpools.cc +++ b/test/gtest/ucp/rndv_frag_mpools.cc @@ -8,6 +8,7 @@ extern "C" { #include +#include #include } @@ -18,8 +19,31 @@ rndv_frag_mpools::rndv_frag_mpools(ucp_worker_h worker, ucs_memory_type_t mem_type) : m_worker(worker) { + std::vector sys_devs; + ucp_mem_desc_t *mdesc; + UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(worker); + for (khiter_t iter = kh_begin(&worker->mpool_hash); + iter != kh_end(&worker->mpool_hash); ++iter) { + if (kh_exist(&worker->mpool_hash, iter) && + (kh_key(&worker->mpool_hash, iter).mem_type == mem_type)) { + sys_devs.push_back(kh_key(&worker->mpool_hash, iter).sys_dev); + } + } + + /* A pool is created by the first allocation from it, so the other pool + * of the same device may not exist yet */ + for (ucs_sys_device_t sys_dev : sys_devs) { + for (unsigned pool : {UCP_WORKER_RNDV_FRAG_POOL_SHARED, + UCP_WORKER_RNDV_FRAG_POOL_RESERVED}) { + if (ucp_rndv_mpool_get(worker, mem_type, sys_dev, pool, &mdesc) == + UCS_OK) { + ucs_mpool_put_inline(mdesc); + } + } + } + for (khiter_t iter = kh_begin(&worker->mpool_hash); iter != kh_end(&worker->mpool_hash); ++iter) { if (kh_exist(&worker->mpool_hash, iter) && @@ -106,13 +130,13 @@ void rndv_frag_mpools::restore_alloc() bool rndv_frag_mpools::can_grow() const { - bool can_grow = !m_mpools.empty(); + bool can_grow = false; UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(m_worker); for (ucs_mpool_t *mpool : m_mpools) { /* With no free fragment, empty means the quota is exhausted */ - can_grow = can_grow && !ucs_mpool_is_empty(mpool); + can_grow = can_grow || !ucs_mpool_is_empty(mpool); } UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(m_worker); diff --git a/test/gtest/ucp/rndv_frag_mpools.h b/test/gtest/ucp/rndv_frag_mpools.h index fa7a2140c92..acdd79176af 100644 --- a/test/gtest/ucp/rndv_frag_mpools.h +++ b/test/gtest/ucp/rndv_frag_mpools.h @@ -19,8 +19,9 @@ extern "C" { /** * Controls the rendezvous fragment memory pools of one memory type on a - * worker. Requires the pools to exist, i.e. a transfer which used fragments of - * that memory type must have completed on the worker. + * worker. Requires a pool to exist, i.e. a transfer which used fragments of + * that memory type must have completed on the worker; the other pool of the + * same device is created if needed. */ class rndv_frag_mpools { public: @@ -38,8 +39,9 @@ class rndv_frag_mpools { void fail_alloc(); void restore_alloc(); - /* Whether every pool can still grow, so that fail_alloc() takes effect - * for the next allocation which finds no free fragment */ + /* Whether some pool can still grow, so that fail_alloc() takes effect + * for the next allocation which finds no free fragment: a request which + * finds the quota of a pool exhausted tries the next one */ bool can_grow() const; size_t num_pools() const diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index ea0795db51c..e8922aa2604 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -1045,23 +1045,28 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_mdesc_no_rcache) UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_quota_exhausted, (get_variant_value() == VARIANT_PROTO_DISABLE), "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", - "RNDV_FRAG_WORKER_MAX_MEM=8K", "RNDV_FRAG_RTR_RATIO=1") + "RNDV_FRAG_WORKER_MAX_MEM=8K") { - ucp_mem_desc_t *mdesc1; - ucp_mem_desc_t *mdesc2; - ucp_mem_desc_t *mdesc3; ucp_worker_h worker = sender().worker(); + ucp_mem_desc_t *shared; + ucp_mem_desc_t *reserved; + ucp_mem_desc_t *mdesc; + /* Two fragments: one in each pool */ ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc1)); + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &shared)); ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc2)); + UCP_WORKER_RNDV_FRAG_POOL_RESERVED, + &reserved)); + EXPECT_EQ(UCS_ERR_NO_RESOURCE, + rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc)); EXPECT_EQ(UCS_ERR_NO_RESOURCE, rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, - UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc3)); + UCP_WORKER_RNDV_FRAG_POOL_RESERVED, &mdesc)); - ucs_mpool_put(mdesc2); - ucs_mpool_put(mdesc1); + ucs_mpool_put(reserved); + ucs_mpool_put(shared); } /* 4 fragments: 3 in the shared pool for any operation, 1 kept in the diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index ea318c2bacb..499729502a1 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -1175,7 +1175,7 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, pipeline, UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_fc_pending, "RNDV_FRAG_SIZE=host:8K,cuda:8K", "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", - "RNDV_FRAG_WORKER_MAX_MEM=8K", + "RNDV_FRAG_WORKER_MAX_MEM=16K", "RNDV_FRAG_MEM_TYPE=cuda") { rndv_fc_pending_abort_test(receiver(), UCP_WORKER_RNDV_FC_OP_RTR); @@ -1184,7 +1184,7 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_fc_pending, UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_fc_pending, "RNDV_FRAG_SIZE=host:8K,cuda:8K", "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", - "RNDV_FRAG_WORKER_MAX_MEM=8K", + "RNDV_FRAG_WORKER_MAX_MEM=16K", "RNDV_FRAG_MEM_TYPE=cuda") { rndv_fc_pending_abort_test(sender(), UCP_WORKER_RNDV_FC_OP_PUT); @@ -1200,14 +1200,13 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype_alloc_failure, } /* A single fragment, since the abort of a pipelined send does not report the - * error to the send request (see ucp_proto_rndv_stub_abort()). No reserved - * pool, which a single fragment would exhaust while the shared pool would be - * created only after the allocation failure is set up. */ + * error to the send request (see ucp_proto_rndv_stub_abort()). The one + * fragment of the reserved pool is held, so the request falls through to the + * shared pool and fails to grow it. */ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, put_mtype_alloc_failure, "RNDV_FRAG_SIZE=host:8K,cuda:8K", "RNDV_FRAG_ALLOC_COUNT=host:1,cuda:1", "RNDV_FRAG_WORKER_MAX_MEM=32K", - "RNDV_FRAG_RTR_RATIO=1", "RNDV_FRAG_MEM_TYPE=cuda") { rndv_fc_alloc_failure_abort_test(sender(), 8 * UCS_KBYTE); From 58b352340961587325f75730b53a155839c8e1e7 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Mon, 19 Oct 2026 11:13:25 +0300 Subject: [PATCH 66/67] UCP/RNDV: Allocate rndv fragment mpools outside of the worker hash table --- src/ucp/core/ucp_worker.c | 3 +- src/ucp/core/ucp_worker.h | 4 ++- src/ucp/core/ucp_worker.inl | 2 +- src/ucp/rndv/rndv.c | 16 ++++++--- test/gtest/ucp/rndv_frag_mpools.cc | 2 +- test/gtest/ucp/test_ucp_mmap.cc | 55 +++++++++++++++++++++++++++++- 6 files changed, 73 insertions(+), 9 deletions(-) diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 1a062526ac4..9bee150d05a 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -2209,7 +2209,8 @@ static void ucp_worker_destroy_mpools(ucp_worker_h worker) if (!kh_exist(&worker->mpool_hash, iter)) { continue; } - ucs_mpool_cleanup(&kh_val(&worker->mpool_hash, iter), 1); + ucs_mpool_cleanup(kh_val(&worker->mpool_hash, iter), 1); + ucs_free(kh_val(&worker->mpool_hash, iter)); } kh_destroy_inplace(ucp_worker_mpool_hash, &worker->mpool_hash); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index f9086113581..b493dd72045 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -246,7 +246,9 @@ typedef struct ucp_worker_mpool_key { /* Hash map to find mpool by mpool key */ -KHASH_TYPE(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t); +/* Pools are allocated separately, since fragments point to their pool and the + * hash table moves its values when it grows */ +KHASH_TYPE(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t*); typedef khash_t(ucp_worker_mpool_hash) ucp_worker_mpool_hash_t; diff --git a/src/ucp/core/ucp_worker.inl b/src/ucp/core/ucp_worker.inl index 75b6e353652..a8d1177d04f 100644 --- a/src/ucp/core/ucp_worker.inl +++ b/src/ucp/core/ucp_worker.inl @@ -61,7 +61,7 @@ ucp_worker_mpool_key_is_equal(ucp_worker_mpool_key_t mpool_key1, (mpool_key1.pool == mpool_key2.pool); } -KHASH_IMPL(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t, +KHASH_IMPL(ucp_worker_mpool_hash, ucp_worker_mpool_key_t, ucs_mpool_t*, 1, ucp_worker_mpool_hash_func, ucp_worker_mpool_key_is_equal); static UCS_F_ALWAYS_INLINE khint_t diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 0b319669421..cd3d0cd5b25 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -1164,7 +1164,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); if (ucs_likely(khiter != kh_end(&worker->mpool_hash))) { - mpool = &kh_val(&worker->mpool_hash, khiter); + mpool = kh_val(&worker->mpool_hash, khiter); goto out_mp_get; } @@ -1186,15 +1186,21 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, } } + mpool = ucs_malloc(sizeof(*mpool), "ucp_rndv_frag_mpool"); + if (mpool == NULL) { + status = UCS_ERR_NO_MEMORY; + goto err; + } + khiter = kh_put(ucp_worker_mpool_hash, &worker->mpool_hash, key, &khret); if (khret == UCS_KH_PUT_FAILED) { status = UCS_ERR_NO_MEMORY; - goto err; + goto err_free_mpool; } ucs_assert_always(khret != UCS_KH_PUT_KEY_PRESENT); - mpool = &kh_value(&worker->mpool_hash, khiter); + kh_value(&worker->mpool_hash, khiter) = mpool; num_frags = context->config.ext.rndv_num_frags[key.mem_type]; ucs_mpool_params_reset(&mp_params); @@ -1211,7 +1217,7 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, /* Callers retry on UCS_ERR_NO_RESOURCE, meaning quota exhaustion */ ucs_assert(status != UCS_ERR_NO_RESOURCE); kh_del(ucp_worker_mpool_hash, &worker->mpool_hash, khiter); - goto err; + goto err_free_mpool; } mpriv = ucs_mpool_priv(mpool); @@ -1232,6 +1238,8 @@ ucp_rndv_mpool_get(ucp_worker_h worker, ucs_memory_type_t mem_type, return ucs_mpool_is_empty(mpool) ? UCS_ERR_NO_RESOURCE : UCS_ERR_NO_MEMORY; +err_free_mpool: + ucs_free(mpool); err: return status; } diff --git a/test/gtest/ucp/rndv_frag_mpools.cc b/test/gtest/ucp/rndv_frag_mpools.cc index 880192b8960..e99bc0a7942 100644 --- a/test/gtest/ucp/rndv_frag_mpools.cc +++ b/test/gtest/ucp/rndv_frag_mpools.cc @@ -48,7 +48,7 @@ rndv_frag_mpools::rndv_frag_mpools(ucp_worker_h worker, iter != kh_end(&worker->mpool_hash); ++iter) { if (kh_exist(&worker->mpool_hash, iter) && (kh_key(&worker->mpool_hash, iter).mem_type == mem_type)) { - m_mpools.push_back(&kh_val(&worker->mpool_hash, iter)); + m_mpools.push_back(kh_val(&worker->mpool_hash, iter)); } } diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index e8922aa2604..f89d55404b1 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -13,6 +13,7 @@ extern "C" { #include #include #include +#include #include #include #include @@ -1102,6 +1103,58 @@ UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_rtr_ratio, } } +/* Fragments point to their pool, so the pool must stay in place when the hash + * table of pools grows */ +UCS_TEST_SKIP_COND_P(test_ucp_mmap, rndv_mpool_hash_growth, + (get_variant_value() == VARIANT_PROTO_DISABLE), + "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:1", + "RNDV_FRAG_WORKER_MAX_MEM=8K") +{ + ucp_worker_h worker = sender().worker(); + ucp_worker_mpool_hash_t *mpools = &worker->mpool_hash; + std::vector placeholders; + ucp_worker_mpool_key_t key; + ucp_mem_desc_t *mdesc; + khint_t num_buckets; + khiter_t khiter; + int khret; + + /* A fragment points to its pool */ + ASSERT_UCS_OK(rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc)); + const ucs_mpool_t *owner = ucs_mpool_obj_owner(mdesc); + + /* Grow the hash table with placeholder entries for other devices */ + key.mem_type = UCS_MEMORY_TYPE_HOST; + key.pool = UCP_WORKER_RNDV_FRAG_POOL_SHARED; + num_buckets = kh_n_buckets(mpools); + for (key.sys_dev = 0; kh_n_buckets(mpools) == num_buckets; ++key.sys_dev) { + khiter = kh_put(ucp_worker_mpool_hash, mpools, key, &khret); + ASSERT_NE(UCS_KH_PUT_FAILED, khret); + if (khret == UCS_KH_PUT_KEY_PRESENT) { + continue; + } + + kh_val(mpools, khiter) = NULL; + placeholders.push_back(key); + } + + /* The pool is still where its fragment points to */ + key.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN; + khiter = kh_get(ucp_worker_mpool_hash, mpools, key); + ASSERT_NE(kh_end(mpools), khiter); + EXPECT_EQ(owner, kh_val(mpools, khiter)) + << "the pool moved while its fragment was allocated"; + ucs_mpool_put(mdesc); + + /* The worker cleans up every pool in the hash table */ + for (const auto &placeholder : placeholders) { + khiter = kh_get(ucp_worker_mpool_hash, mpools, placeholder); + ASSERT_NE(kh_end(mpools), khiter); + kh_del(ucp_worker_mpool_hash, mpools, khiter); + } +} + UCS_TEST_P(test_ucp_mmap, rndv_mpool_quota_disabled_with_proto_v1, "PROTO_ENABLE=n", "RNDV_FRAG_SIZE=host:4K", "RNDV_FRAG_ALLOC_COUNT=host:2", "RNDV_FRAG_WORKER_MAX_MEM=8K") @@ -1147,7 +1200,7 @@ UCS_TEST_P(test_ucp_mmap, rndv_mpool_allocation_failure, const khiter_t khiter = kh_get(ucp_worker_mpool_hash, &worker->mpool_hash, key); ASSERT_NE(kh_end(&worker->mpool_hash), khiter); - ucs_mpool_t *mpool = &kh_val(&worker->mpool_hash, khiter); + ucs_mpool_t *mpool = kh_val(&worker->mpool_hash, khiter); /* The first chunk is exhausted, but the pool still has quota. Make the * next chunk allocation fail and verify it is not reported as quota From ebc7ce7b6c6d2156468b34ce93fffbd5d9443809 Mon Sep 17 00:00:00 2001 From: Alma Mastbaum Date: Tue, 20 Oct 2026 16:13:36 +0300 Subject: [PATCH 67/67] UCP/RNDV: Initialize the fragment mpool failure ops in the constructor --- test/gtest/ucp/rndv_frag_mpools.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/gtest/ucp/rndv_frag_mpools.cc b/test/gtest/ucp/rndv_frag_mpools.cc index e99bc0a7942..19c740af047 100644 --- a/test/gtest/ucp/rndv_frag_mpools.cc +++ b/test/gtest/ucp/rndv_frag_mpools.cc @@ -17,7 +17,7 @@ unsigned rndv_frag_mpools::s_num_alloc_failures = 0; rndv_frag_mpools::rndv_frag_mpools(ucp_worker_h worker, ucs_memory_type_t mem_type) : - m_worker(worker) + m_worker(worker), m_fail_ops() { std::vector sys_devs; ucp_mem_desc_t *mdesc;