UCP/RNDV: Throttle rndv fragment requests

src/ucp/core/ucp_context.c

@@ -418,6 +418,28 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "even if invalidation workflow isn't supported",
    ucs_offsetof(ucp_context_config_t, rndv_errh_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. 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. When the limit is\n"
+   "reached, further fragment requests are queued until fragments are\n"
+   "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"
+   "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, 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},
+
   {"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"
@@ -2785,6 +2807,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;

src/ucp/core/ucp_context.h

@@ -130,6 +130,10 @@ typedef struct ucp_context_config {
     int                                    rndv_shm_cuda_staging_force;
     /** Enable error handling for rndv pipeline protocol */
     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;
@@ -807,7 +811,6 @@ ucp_context_print_transport_tables_enabled(ucp_context_h context)
 
     return value == UCS_CONFIG_ON;
 }
-
 void ucp_context_memaccess_tl_bitmap(ucp_context_h context,
                                      uint64_t mem_type_bitmap,
                                      uint64_t md_reg_flags,

src/ucp/core/ucp_ep.c

@@ -26,6 +26,7 @@
 #include <ucp/proto/proto_common.inl>
 #include <ucp/proto/proto_debug.h>
 #include <ucp/rndv/rndv.h>
+#include <ucp/rndv/proto_rndv.h>
 #include <ucp/stream/stream.h>
 #include <ucp/core/ucp_listener.h>
 #include <ucp/rma/rma.inl>
@@ -532,6 +533,7 @@ static int ucp_ep_remove_filter(const ucs_callbackq_elem_t *elem, void *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_proto_rndv_mtype_fc_reschedule_filter(elem, arg) ||
         ucp_ep_tf_purge_remove_filter(elem, arg)) {
         return 1;
     }
@@ -5476,6 +5478,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) {

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;
 
 

src/ucp/core/ucp_request.h

@@ -73,6 +73,22 @@ enum {
 };
 
 
+/**
+ * Fragment flow control state of a rndv request. Every rndv request must be
+ * initialized with ucp_proto_rndv_req_fc_init() when it is set up, since
+ * ucp_proto_rndv_mtype_fc_cancel() reads the state for any aborted or reset
+ * request.
+ */
+enum {
+    /* Not throttled */
+    UCP_REQUEST_RNDV_MTYPE_FC_NONE,
+    /* Waiting in the worker-level pending queue */
+    UCP_REQUEST_RNDV_MTYPE_FC_QUEUED,
+    /* Dequeued, a retry is scheduled from the progress loop */
+    UCP_REQUEST_RNDV_MTYPE_FC_RESCHED
+};
+
+
 /**
  * Protocols enumerator to work with send request state
  */
@@ -322,6 +338,9 @@ struct ucp_request {
                             /* Data start offset of this request */
                             size_t offset;
 
+                            /* UCP_REQUEST_RNDV_MTYPE_FC_* state */
+                            uint8_t fc_state;
+
                             union {
                                 /* Used by rndv/put and rndv/put/frag */
                                 struct {
@@ -343,6 +362,13 @@ struct ucp_request {
                                     ssize_t ack_data_size;
                                 } ppln;
 
+                                /* Used by rndv mtype requests while fc_state
+                                 * is not UCP_REQUEST_RNDV_MTYPE_FC_NONE */
+                                struct {
+                                    /* Element in worker-level pending queue */
+                                    ucs_queue_elem_t queue_elem;
+                                } fc;
+
                                 /* Used by rndv/rkey_ptr */
                                 struct {
                                     /* Element in queue for segmented RKEY ptr */

src/ucp/core/ucp_worker.c

@@ -88,18 +88,19 @@ 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"
     }
 };
 #endif
@@ -2217,7 +2218,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);
@@ -2712,6 +2714,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context,
     unsigned name_length;
     ucp_worker_h worker;
     ucs_status_t status;
+    unsigned q_index;
 
     worker = ucs_calloc(1, sizeof(*worker), "ucp worker");
     if (worker == NULL) {
@@ -2743,6 +2746,11 @@ 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 pending queues */
+    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) &
                     UCS_MASK(UCP_WORKER_INTERNAL_FLAGS_SHIFT);

src/ucp/core/ucp_worker.h

@@ -159,6 +159,7 @@ 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_LAST
 };
@@ -223,15 +224,31 @@ 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;
 
 
 /* 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;
 
 
@@ -302,6 +319,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).
  */
@@ -403,6 +433,12 @@ typedef struct ucp_worker {
         uint64_t                     ep_failures;
     } counters;
 
+    struct {
+        /* 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;
+
     struct {
         /* Usage tracker handle */
         ucs_usage_tracker_h          handle;

src/ucp/core/ucp_worker.inl

@@ -48,18 +48,20 @@ 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
 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,
+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

src/ucp/rma/rma_rndv.c

@@ -510,6 +510,7 @@ static ucs_status_t ucp_proto_get_rndv_init(ucp_request_t *get_req,
     rndv_req->send.rndv.remote_address     = address;
     rndv_req->send.rndv.rkey               = get_req->send.rma.rkey;
     rndv_req->send.rndv.offset             = 0;
+    ucp_proto_rndv_req_fc_init(rndv_req);
 
     UCS_PROFILE_CALL_VOID(ucp_datatype_iter_move,
                           &rndv_req->send.state.dt_iter,

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