diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 70f97660c4f..38b685894c5 100644 --- a/src/ucp/core/ucp_context.c +++ b/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; diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 7d19b9bf163..6ec11c89170 100644 --- a/src/ucp/core/ucp_context.h +++ b/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, diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index f185398e6a7..1884c3a774c 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -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) { 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_request.h b/src/ucp/core/ucp_request.h index f65774a9b43..7f101c94ca4 100644 --- a/src/ucp/core/ucp_request.h +++ b/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 */ diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index b674833d888..9123c09a814 100644 --- a/src/ucp/core/ucp_worker.c +++ b/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); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 32ae9d1cac6..b493dd72045 100644 --- a/src/ucp/core/ucp_worker.h +++ b/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; diff --git a/src/ucp/core/ucp_worker.inl b/src/ucp/core/ucp_worker.inl index 4f7dce3762c..a8d1177d04f 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,10 +57,11 @@ 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 diff --git a/src/ucp/rma/rma_rndv.c b/src/ucp/rma/rma_rndv.c index 2c9d42f8e2e..40943b12c8e 100644 --- a/src/ucp/rma/rma_rndv.c +++ b/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, 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 f790386ce51..f3e4f6dedd2 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -16,6 +16,135 @@ #include +void ucp_proto_rndv_mtype_fc_leave(ucp_request_t *req) +{ + ucs_assert(req->send.rndv.fc_state != UCP_REQUEST_RNDV_MTYPE_FC_NONE); + req->send.rndv.fc_state = UCP_REQUEST_RNDV_MTYPE_FC_NONE; +} + +unsigned ucp_proto_rndv_mtype_fc_reschedule_cb(void *arg) +{ + ucp_request_t *req = arg; + + ucs_assert(req->send.rndv.fc_state == UCP_REQUEST_RNDV_MTYPE_FC_RESCHED); + /* 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; +} + +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_error("ep %p still has mtype FC reschedule callback for req %p", arg, + elem->arg); + return 1; +} + +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->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; + 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); + } + } + + 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_queue_head_t reqs; + ucp_request_t *req; + + /* Aborting a woken-up request may wake up another request of this ep, + * which is then aborted on a later iteration */ + 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); + } + } +} + +/* 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_queue_head_t reqs; + ucp_request_t *req; + + /* Resetting a woken-up request may wake up another request of this ep, + * which is then extracted on a later iteration */ + 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); + } + } +} + static int ucp_proto_rndv_ctrl_skip_inter_node_md( const ucp_proto_common_init_params_t *params, @@ -687,16 +816,18 @@ 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; + /* 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) { + 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) @@ -708,6 +839,50 @@ ucs_status_t ucp_proto_rndv_rts_reset(ucp_request_t *req) return ucp_proto_request_zcopy_id_reset(req); } +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_frag_worker_max_mem; + size_t frag_size; + 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]; + 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 < 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 2; + } + + 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; + + 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); +} + ucs_status_t ucp_proto_rndv_ack_init(const ucp_proto_common_init_params_t *init_params, const char *name, double overhead, @@ -885,6 +1060,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_proto_rndv_send_reply, /* Caching rkey_buffer pointer for later unpacking of shm keys in * rkey_ptr mtype ppln protocol. */ req->send.rndv.rkey_buffer = rkey_buffer; + ucp_proto_rndv_req_fc_init(req); ucp_trace_req(req, "%s rva 0x%" PRIx64 " length %zd rreq_id 0x%" PRIx64 diff --git a/src/ucp/rndv/proto_rndv.h b/src/ucp/rndv/proto_rndv.h index 3ee14ac1464..28e3c7f6c81 100644 --- a/src/ucp/rndv/proto_rndv.h +++ b/src/ucp/rndv/proto_rndv.h @@ -9,6 +9,8 @@ #include "rndv.h" +#include + #include @@ -193,6 +195,36 @@ 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_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); + +/* 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( + 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, const char *name, double overhead, diff --git a/src/ucp/rndv/proto_rndv.inl b/src/ucp/rndv/proto_rndv.inl index d4ae0f720d3..a4d072a8b76 100644 --- a/src/ucp/rndv/proto_rndv.inl +++ b/src/ucp/rndv/proto_rndv.inl @@ -17,6 +17,13 @@ #include +/* Must be called when a rndv request is set up, see + * UCP_REQUEST_RNDV_MTYPE_FC_* */ +static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_req_fc_init(ucp_request_t *req) +{ + req->send.rndv.fc_state = UCP_REQUEST_RNDV_MTYPE_FC_NONE; +} + static UCS_F_ALWAYS_INLINE int ucp_proto_rndv_shm_pipeline_force_enabled(ucp_context_h context) { diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 80d47e02d4c..cd74a5a0925 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -21,7 +21,6 @@ #include #include - static UCS_F_ALWAYS_INLINE int ucp_rndv_frag_mem_type(ucp_context_t *context) { @@ -1137,14 +1136,24 @@ ucs_mpool_ops_t ucp_frag_mpool_ops = { .obj_cleanup = (ucs_mpool_obj_cleanup_func_t)ucs_empty_function }; -ucp_mem_desc_t * +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) + 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,42 +1161,88 @@ 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))) { - mpool = &kh_val(&worker->mpool_hash, khiter); + mpool = kh_val(&worker->mpool_hash, khiter); 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; + } + } + + 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) { - return NULL; + status = UCS_ERR_NO_MEMORY; + goto err_free_mpool; } 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]; + kh_value(&worker->mpool_hash, khiter) = mpool; + 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); mp_params.elem_size = sizeof(ucp_mem_desc_t); mp_params.alignment = 1; - mp_params.elems_per_chunk = num_frags; + /* 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"; + mp_params.name = ucp_rndv_frag_mpool_names[pool]; status = ucs_mpool_init(&mp_params, mpool); if (status != UCS_OK) { - return NULL; + /* 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_free_mpool; } - 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; + mpriv->pool = pool; out_mp_get: - return ucp_worker_mpool_get(mpool); + mdesc = ucp_worker_mpool_get(mpool); + if (ucs_likely(mdesc != NULL)) { + *mdesc_p = mdesc; + 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_free_mpool: + ucs_free(mpool); +err: + return status; } static void ucp_rndv_send_frag_get_mem_type(ucp_request_t *sreq, size_t length, @@ -1202,6 +1257,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 */ @@ -1210,10 +1266,12 @@ 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, + 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)); } freq->send.ep = sreq->send.ep; @@ -1362,6 +1420,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); @@ -1387,10 +1446,13 @@ 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, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, + &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..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); -ucp_mem_desc_t * +/** + * 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); + 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_get.c b/src/ucp/rndv/rndv_get.c index 6636ef717b4..452e719f008 100644 --- a/src/ucp/rndv/rndv_get.c +++ b/src/ucp/rndv/rndv_get.c @@ -273,7 +273,8 @@ 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_mdesc_release(req); + if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 1, 0); } else { @@ -304,8 +305,15 @@ 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)) { + /* 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); + rpriv->frag_sys_dev, + UCP_WORKER_RNDV_FC_OP_GET); + if (status == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; @@ -368,13 +376,14 @@ 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; } - ucs_mpool_put_inline(req->send.rndv.mdesc); - req->send.rndv.mdesc = NULL; - req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED; + ucp_proto_rndv_mtype_mdesc_release(req); + 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)) { @@ -384,6 +393,36 @@ 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_PROTO_INITIALIZED) { + 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_datatype_iter_cleanup(&req->send.state.dt_iter, 1, + UCP_DT_MASK_ALL); + ucp_proto_rndv_recv_complete_status(req, status); + } +} + ucp_proto_t ucp_rndv_get_mtype_proto = { .name = "rndv/get/mtype", .desc = NULL, @@ -395,6 +434,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 4b11edb3fd8..0c9c062544d 100644 --- a/src/ucp/rndv/rndv_mtype.inl +++ b/src/ucp/rndv/rndv_mtype.inl @@ -44,20 +44,128 @@ 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, 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 + * 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, + 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 < num_queues; 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->send.rndv.fc_state == + UCP_REQUEST_RNDV_MTYPE_FC_QUEUED); + pending_req->send.rndv.fc_state = UCP_REQUEST_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); + /* The retry is not tied to any event, so wake up a blocked worker */ + ucp_worker_signal_internal(worker); + return; + } +} + +/* 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; + + /* 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 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, + unsigned fc_op) { - ucp_worker_h worker = req->send.ep->worker; + 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->send.rndv.fc_state == + UCP_REQUEST_RNDV_MTYPE_FC_RESCHED); + 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) { - return UCS_ERR_NO_MEMORY; + /* A queued request is retried only after being rescheduled */ + ucs_assert(req->send.rndv.fc_state != UCP_REQUEST_RNDV_MTYPE_FC_QUEUED); + if (owns_wakeup) { + ucp_proto_rndv_mtype_fc_leave(req); } - return UCS_OK; + req->send.rndv.mdesc = NULL; + 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_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) { + 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_WORKER_RNDV_FC_OP_LAST); + } + + return status; + } + + /* Quota exhausted - throttle by queuing the request in the appropriate + * pending queue ordered by priority. */ + ucp_trace_req(req, + "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], + frag_sys_dev); + UCP_WORKER_STAT_RNDV(worker, MTYPE_FC_THROTTLED, 1); + ucs_assert(req->send.rndv.fc_state == UCP_REQUEST_RNDV_MTYPE_FC_NONE); + 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); + + return UCS_ERR_NO_RESOURCE; } static UCS_F_ALWAYS_INLINE uct_mem_h @@ -168,6 +276,83 @@ static UCS_F_ALWAYS_INLINE ucs_status_t ucp_proto_rndv_mtype_copy( return status; } +static UCS_F_ALWAYS_INLINE int +ucp_proto_rndv_mtype_fc_reschedule_pred(const ucs_callbackq_elem_t *elem, + void *arg) +{ + 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_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); + + if (req->send.rndv.fc_state == UCP_REQUEST_RNDV_MTYPE_FC_NONE) { + return; + } + + if (req->send.rndv.fc_state == UCP_REQUEST_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; + } else { + ucs_assert(req->send.rndv.fc_state == + UCP_REQUEST_RNDV_MTYPE_FC_RESCHED); + ucp_trace_req(req, "mtype_fc: remove aborted reschedule callback"); + /* 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; + } + + ucp_proto_rndv_mtype_fc_leave(req); + + 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, + UCP_WORKER_RNDV_FC_OP_LAST); + } +} + +/** + * 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) +{ + 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. 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; + + ucs_mpool_put_inline(mdesc); + req->send.rndv.mdesc = NULL; + ucp_proto_rndv_mtype_fc_reschedule_pending(worker, num_queues); +} + 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..0437887f8c4 100644 --- a/src/ucp/rndv/rndv_ppln.c +++ b/src/ucp/rndv/rndv_ppln.c @@ -304,6 +304,7 @@ static ucs_status_t ucp_proto_rndv_ppln_progress(uct_pending_req_t *uct_req) freq->send.rndv.rkey = req->send.rndv.rkey; freq->send.rndv.offset = req->send.rndv.offset + req->send.state.dt_iter.offset; + ucp_proto_rndv_req_fc_init(freq); ucp_proto_request_set_proto(freq, &rpriv->frag_proto_cfg, freq->send.state.dt_iter.length); diff --git a/src/ucp/rndv/rndv_put.c b/src/ucp/rndv/rndv_put.c index 0cab5a465d1..452acc869a7 100644 --- a/src/ucp/rndv/rndv_put.c +++ b/src/ucp/rndv/rndv_put.c @@ -520,16 +520,23 @@ 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); - const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv; + 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; ucs_assert(!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)); + rpriv = req->send.proto_config->priv; + + /* 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); + rpriv->bulk.frag_sys_dev, + UCP_WORKER_RNDV_FC_OP_PUT); + if (status == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; @@ -565,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"); - ucs_mpool_put(req->send.rndv.mdesc); + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucp_proto_rndv_mtype_mdesc_release(req); + } + ucp_proto_rndv_put_common_complete(req); } @@ -575,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"); - ucs_mpool_put(req->send.rndv.mdesc); + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucp_proto_rndv_mtype_mdesc_release(req); + } + ucp_proto_rndv_ppln_send_frag_complete(req, 1); } @@ -660,6 +673,32 @@ 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) +{ + 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); +} + +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, @@ -674,6 +713,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, - .reset = (ucp_request_reset_func_t)ucp_proto_reset_fatal_not_implemented + .abort = ucp_proto_rndv_put_mtype_abort, + .reset = ucp_proto_rndv_put_mtype_reset }; diff --git a/src/ucp/rndv/rndv_rtr.c b/src/ucp/rndv/rndv_rtr.c index 7494a47c415..53af8f418e4 100644 --- a/src/ucp/rndv/rndv_rtr.c +++ b/src/ucp/rndv/rndv_rtr.c @@ -359,9 +359,10 @@ 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)) { - ucs_mpool_put_inline(req->send.rndv.mdesc); + if (!abort || (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + ucp_proto_rndv_mtype_mdesc_release(req); } + if (ucp_proto_rndv_request_is_ppln_frag(req)) { ucp_proto_rndv_ppln_recv_frag_complete(req, 0, abort); } else { @@ -374,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 @@ -389,9 +391,10 @@ 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) { - ucs_mpool_put_inline(req->send.rndv.mdesc); - req->send.rndv.mdesc = NULL; + ucp_proto_rndv_mtype_mdesc_release(req); } return ucp_proto_request_zcopy_id_reset(req); @@ -424,12 +427,21 @@ 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; + const ucp_proto_rndv_rtr_mtype_priv_t *rpriv; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + rpriv = req->send.proto_config->priv; + + /* 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); + rpriv->frag_sys_dev, + UCP_WORKER_RNDV_FC_OP_RTR); + if (status == UCS_ERR_NO_RESOURCE) { + return UCS_OK; + } + if (status != UCS_OK) { ucp_proto_request_abort(req, status); return UCS_OK; diff --git a/src/ucp/wireup/wireup.c b/src/ucp/wireup/wireup.c index 0e7bcc8ab3a..f532ac938a0 100644 --- a/src/ucp/wireup/wireup.c +++ b/src/ucp/wireup/wireup.c @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -2538,6 +2539,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/Makefile.am b/test/gtest/Makefile.am index 31100454790..d31a8629427 100644 --- a/test/gtest/Makefile.am +++ b/test/gtest/Makefile.am @@ -182,6 +182,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..19c740af047 --- /dev/null +++ b/test/gtest/ucp/rndv_frag_mpools.cc @@ -0,0 +1,145 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#include "rndv_frag_mpools.h" + +extern "C" { +#include +#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), m_fail_ops() +{ + 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) && + (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 = 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); + } + + 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..acdd79176af --- /dev/null +++ b/test/gtest/ucp/rndv_frag_mpools.h @@ -0,0 +1,80 @@ +/** + * 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 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: + 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 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 + { + 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 7c1022e4559..262a725198f 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 @@ -2323,24 +2324,26 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { m_mem_type = mem_type; } -private: - ucs_memory_type_t tx_memtype() const override + static uint64_t get_stats(const entity &e, uint64_t cntr) { - return m_mem_type; + return UCS_STATS_GET_COUNTER(e.worker()->stats, cntr); } - ucs_memory_type_t rx_memtype() const override + static const char *stats_name(const entity &e, uint64_t cntr) { - return m_mem_type; + return e.worker()->stats->cls->counter_names[cntr]; } - void check_stats(entity &e, uint64_t cntr, uint64_t exp_value) + void check_stats(const 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, get_stats(e, cntr)) + << "counter is " << stats_name(e, cntr); + } - EXPECT_EQ(exp_value, value) << "counter is " - << stats_node->cls->counter_names[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); } ucp_err_handling_mode_t get_err_mode() const @@ -2348,6 +2351,17 @@ class test_ucp_am_nbx_rndv_ppln : public test_ucp_am_nbx_rndv { return static_cast(get_variant_value(1)); } +private: + ucs_memory_type_t tx_memtype() const override + { + return m_mem_type; + } + + ucs_memory_type_t rx_memtype() const override + { + return m_mem_type; + } + ucs_memory_type_t m_mem_type; }; @@ -2397,6 +2411,299 @@ 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: + struct fc_counters { + uint64_t sender_throttled; + uint64_t receiver_throttled; + }; + + void verify_clean_fc_state() + { + for (auto *ep : {&sender(), &receiver()}) { + check_pending_queues_empty(*ep); + } + } + + 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"); + } + + send_message(num_frags); + + check_stats_ge(sender(), UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, 1); + check_stats_ge(receiver(), UCP_WORKER_STAT_RNDV_RTR_MTYPE, 1); + + const uint64_t cntr = UCP_WORKER_STAT_RNDV_MTYPE_FC_THROTTLED; + + fc.sender_throttled = get_stats(sender(), cntr); + 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(); + + 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, + "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; + + /* 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=1M"; + + 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_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") +{ + fc_counters fc; + + 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"; +} + + +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_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 479f1b7402d..eae86e27805 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 @@ -1013,17 +1014,214 @@ 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; +} + +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 = ucp_rndv_mpool_get(worker, mem_type, - UCS_SYS_DEVICE_ID_UNKNOWN); + ucp_mem_desc_t *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); } } +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_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, &shared)); + 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_SHARED, &mdesc)); + 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); + ucs_mpool_put(shared); +} + +/* 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); + } +} + +/* 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") +{ + ucp_mem_desc_t *mdesc1; + ucp_mem_desc_t *mdesc2; + ucp_mem_desc_t *mdesc3; + ucp_worker_h worker = sender().worker(); + + 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); + 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_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; + ucs_mpool_ops_t fail_ops; + const ucs_mpool_ops_t *orig_ops; + ucs_status_t status; + + 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); + 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 + * 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 = rndv_mpool_get(worker, UCS_MEMORY_TYPE_HOST, + UCP_WORKER_RNDV_FRAG_POOL_SHARED, &mdesc3); + } + mpool->data->ops = orig_ops; + + EXPECT_EQ(UCS_ERR_NO_MEMORY, status); + + 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 = diff --git a/test/gtest/ucp/test_ucp_peer_failure.cc b/test/gtest/ucp/test_ucp_peer_failure.cc index 9b0b658dac3..499729502a1 100644 --- a/test/gtest/ucp/test_ucp_peer_failure.cc +++ b/test/gtest/ucp/test_ucp_peer_failure.cc @@ -8,12 +8,15 @@ #include "ucp_test.h" #include "ucp_datatype.h" +#include "rndv_frag_mpools.h" extern "C" { #include /* for testing EP RNDV configuration */ #include /* for debug */ #include /* for testing memory consumption */ +#include #include +#include } #include @@ -817,6 +820,17 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } } + static unsigned close_fc_pending_ep_cb(void *arg) + { + auto *test = static_cast(arg); + + /* 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; + } + void close_peer() { if (m_is_peer_closed) { @@ -836,12 +850,14 @@ 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::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); } @@ -852,8 +868,34 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { } /* Call original proto progress */ - return 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_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, + close_fc_pending_ep_cb, self); + } + + 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 + * 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 bool has_fc_pending(ucp_worker_h worker, unsigned fc_op) + { + return !ucs_queue_is_empty(&worker->rndv_mtype_fc.pending_q[fc_op]); } virtual void cleanup() @@ -920,12 +962,118 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure { ASSERT_TRUE(m_is_peer_closed); } + void rndv_fc_pending_abort_test(entity &fc_entity, unsigned fc_op) + { + 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) { + 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); + 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"); + } + + 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); + + { + scoped_log_handler err_wrapper(wrap_errors_logger); + scoped_log_handler warn_wrapper(wrap_warns_logger); + result = smoke_test(true); + } + + frag_mpools.release(); + 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(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; + + 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); + 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); + frag_mpools.hold(); + } + + if (frag_mpools.num_held() == 0) { + UCS_TEST_SKIP_R("no CUDA fragments were allocated"); + } + + ASSERT_TRUE(frag_mpools.can_grow()) + << "a CUDA fragment pool has no quota left to grow"; + + 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); + 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(); + frag_mpools.restore_alloc(); + frag_mpools.release(); + + 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); + 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}; 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. */ @@ -1024,4 +1172,44 @@ 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_ALLOC_COUNT=host:1,cuda:1", + "RNDV_FRAG_WORKER_MAX_MEM=16K", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + 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=16K", + "RNDV_FRAG_MEM_TYPE=cuda") +{ + 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()). 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_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); diff --git a/test/gtest/ucp/ucp_test.cc b/test/gtest/ucp/ucp_test.cc index ad8c70ce154..a9a5039b6b5 100644 --- a/test/gtest/ucp/ucp_test.cc +++ b/test/gtest/ucp/ucp_test.cc @@ -1220,13 +1220,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]; @@ -1237,7 +1234,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; } } @@ -1246,6 +1243,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 80032511894..86e084ec0b0 100644 --- a/test/gtest/ucp/ucp_test.h +++ b/test/gtest/ucp/ucp_test.h @@ -162,6 +162,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;