From a1fb5905d2ebb43ff10b5a2f4ed63b03c9b6f627 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Mon, 19 Oct 2026 13:18:45 +0300 Subject: [PATCH 1/8] UCT/API/V2: Added flag to free unused mappings at rkey release. --- src/uct/api/v2/uct_v2.h | 23 +++++++++- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 33 +++++++++++++++ src/uct/cuda/cuda_ipc/cuda_ipc_cache.h | 23 ++++++++++ src/uct/cuda/cuda_ipc/cuda_ipc_md.c | 33 +++++++++++---- test/gtest/uct/cuda/test_cuda_ipc_md.cc | 56 +++++++++++++++++++++++++ 5 files changed, 158 insertions(+), 10 deletions(-) diff --git a/src/uct/api/v2/uct_v2.h b/src/uct/api/v2/uct_v2.h index a7115711a7a..b1a6db5f9d2 100644 --- a/src/uct/api/v2/uct_v2.h +++ b/src/uct/api/v2/uct_v2.h @@ -764,10 +764,25 @@ typedef struct uct_rkey_compare_params { * @brief Rkey unpack parameters field mask. */ typedef enum { - UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0) /**< sys_device field */ + UCT_RKEY_UNPACK_FIELD_SYS_DEVICE = UCS_BIT(0), /**< sys_device field */ + UCT_RKEY_UNPACK_FIELD_FLAGS = UCS_BIT(1) /**< flags field */ } uct_rkey_unpack_field_mask_t; +/** + * @ingroup UCT_MD + * @brief Rkey unpack flags. + */ +typedef enum { + /** + * When the remote key is released (see @ref uct_rkey_release), free the + * local mapping of the remote memory which was created for this remote key + * and is not used by any other remote key or operation. + */ + UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING = UCS_BIT(0) +} uct_rkey_unpack_flags_t; + + /** * @ingroup UCT_MD * @brief Parameters for unpacking remote key using @ref uct_rkey_unpack_v2. @@ -785,6 +800,12 @@ typedef struct uct_rkey_unpack_params { * (default behavior). */ ucs_sys_device_t sys_device; + + /** + * Flags to unpack the rkey with, using bits from + * @ref uct_rkey_unpack_flags_t. + */ + uint64_t flags; } uct_rkey_unpack_params_t; diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index bafd8048132..9d8a4941970 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -661,6 +661,39 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, pthread_rwlock_unlock(&cache->lock); } +void uct_cuda_ipc_cache_release_unused(uct_cuda_ipc_cache_t *cache, + uintptr_t d_bptr) +{ + ucs_pgt_region_t *pgt_region; + uct_cuda_ipc_cache_region_t *region; + + pthread_rwlock_wrlock(&cache->lock); + pgt_region = ucs_pgtable_lookup(&cache->pgtable, d_bptr); + if (pgt_region != NULL) { + region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); + if (region->refcount == 0) { + uct_cuda_ipc_cache_region_destroy(cache, region); + } + } + + pthread_rwlock_unlock(&cache->lock); +} + +void uct_cuda_ipc_cache_release_unused_region(pid_t pid, ucs_sys_ns_t pid_ns, + uintptr_t d_bptr, CUdevice cu_dev) +{ + const uct_cuda_ipc_cache_hash_key_t key = {pid, pid_ns, cu_dev}; + uct_cuda_ipc_cache_t *cache; + + ucs_rw_spinlock_read_lock(&uct_cuda_ipc_remote_cache.lock); + cache = uct_cuda_ipc_remote_cache_get(key); + if (cache != NULL) { + uct_cuda_ipc_cache_release_unused(cache, d_bptr); + } + + ucs_rw_spinlock_read_unlock(&uct_cuda_ipc_remote_cache.lock); +} + void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, uintptr_t d_bptr, const void *mapped_addr, CUdevice cu_dev, int cache_enabled) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h index c1d553f55b1..6e15196b767 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h @@ -75,6 +75,29 @@ void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, CUdevice cu_dev, int cache_enabled); +/** + * @brief Close the region of a buffer in the given cache if it is not in use + * + * @param cache Cache to look for the region in + * @param d_bptr Remote base address of the buffer + */ +void uct_cuda_ipc_cache_release_unused(uct_cuda_ipc_cache_t *cache, + uintptr_t d_bptr); + + +/** + * @brief Close the cached region of a buffer if it is not in use + * + * @param pid Remote process ID + * @param pid_ns Remote process ID namespace + * @param d_bptr Remote base address of the buffer + * @param cu_dev CUDA device handle the region was mapped on + */ +void uct_cuda_ipc_cache_release_unused_region(pid_t pid, ucs_sys_ns_t pid_ns, + uintptr_t d_bptr, + CUdevice cu_dev); + + /** * @brief Set global cache limits for newly created CUDA IPC remote caches * diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index ba7e1e649aa..57b7ef9dd95 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -39,6 +39,7 @@ typedef struct { const void *mapped_addr; CUdevice cu_dev; + int free_unused_mapping; } uct_cuda_ipc_rkey_handle_t; typedef struct { @@ -474,7 +475,8 @@ uct_cuda_ipc_mkey_pack(uct_md_h md, uct_mem_h tl_memh, void *address, memh->dev_num)); } -static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle() +static uct_cuda_ipc_rkey_handle_t * +uct_cuda_ipc_create_rkey_handle(uint64_t flags) { uct_cuda_ipc_rkey_handle_t *rkey_handle; @@ -484,8 +486,10 @@ static uct_cuda_ipc_rkey_handle_t *uct_cuda_ipc_create_rkey_handle() return NULL; } - rkey_handle->mapped_addr = NULL; - rkey_handle->cu_dev = CU_DEVICE_INVALID; + rkey_handle->mapped_addr = NULL; + rkey_handle->cu_dev = CU_DEVICE_INVALID; + rkey_handle->free_unused_mapping = !!( + flags & UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING); return rkey_handle; } @@ -562,6 +566,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, uct_cuda_ipc_unpacked_rkey_t *unpacked; uct_cuda_ipc_rkey_handle_t *rkey_handle; ucs_sys_device_t sys_dev; + uint64_t flags; CUdevice cuda_device; CUdevice avail_cuda_device; ucs_status_t status; @@ -569,6 +574,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, sys_dev = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, sys_device, SYS_DEVICE, UCS_SYS_DEVICE_ID_UNKNOWN); + flags = UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, flags, FLAGS, 0); unpacked = ucs_malloc(sizeof(*unpacked), "uct_cuda_ipc_unpacked_rkey_t"); if (NULL == unpacked) { @@ -580,7 +586,7 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack, unpacked->super.super = *packed; unpacked->super.super.pid &= ~UCT_CUDA_IPC_RKEY_FLAG_PID_NS; - rkey_handle = uct_cuda_ipc_create_rkey_handle(); + rkey_handle = uct_cuda_ipc_create_rkey_handle(flags); if (rkey_handle == NULL) { status = UCS_ERR_NO_MEMORY; goto err_free_key; @@ -633,8 +639,7 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, uct_cuda_ipc_extended_rkey_t *extended_rkey; uct_cuda_ipc_rkey_t *rkey; - if ((uct_rkey == 0) || (rkey_handle == NULL) || - (rkey_handle->mapped_addr == NULL)) { + if ((uct_rkey == 0) || (rkey_handle == NULL)) { return; } @@ -642,9 +647,19 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, extended_rkey = &unpacked_rkey->super; rkey = &extended_rkey->super; - uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, rkey->d_bptr, - rkey_handle->mapped_addr, rkey_handle->cu_dev, - uct_cuda_ipc_component.enable_remote_cache); + if (rkey_handle->mapped_addr != NULL) { + uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, + rkey->d_bptr, rkey_handle->mapped_addr, + rkey_handle->cu_dev, + uct_cuda_ipc_component.enable_remote_cache); + } + + if (rkey_handle->free_unused_mapping) { + uct_cuda_ipc_cache_release_unused_region(rkey->pid, + extended_rkey->pid_ns, + rkey->d_bptr, + rkey_handle->cu_dev); + } } static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component, diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index 0cfbf5f8d84..0accf90df89 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -716,6 +716,62 @@ UCS_TEST_F(test_cuda_ipc_cache_lru, stale_destroy_while_in_use) { EXPECT_TRUE(ucs_list_is_empty(&m_cache->lru_list)); } +UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_keeps_in_use) { + create_cache(ULONG_MAX, SIZE_MAX); + + insert_region(0); + uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); + + /* The region is in use, so it must stay in the cache */ + EXPECT_EQ(1UL, m_cache->num_regions); + EXPECT_EQ(REGION_SIZE, m_cache->total_size); + EXPECT_TRUE(pgtable_has(0)); +} + +UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_keeps_reacquired) { + create_cache(ULONG_MAX, SIZE_MAX); + + uct_cuda_ipc_cache_region_t *region = insert_region(0); + release_region(region); + reacquire_region(region); + uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); + + EXPECT_EQ(1UL, m_cache->num_regions); + EXPECT_TRUE(pgtable_has(0)); +} + +UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_no_region) { + create_cache(ULONG_MAX, SIZE_MAX); + + insert_region(0); + /* No region at this address, the cache must not change */ + uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR + 4 * REGION_SIZE); + + EXPECT_EQ(1UL, m_cache->num_regions); + EXPECT_TRUE(pgtable_has(0)); +} + +UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_closes_idle) { + create_cache(ULONG_MAX, SIZE_MAX); + + uct_cuda_ipc_cache_region_t *region0 = insert_region(0); + insert_region(1); + release_region(region0); + + { + /* The test regions have no real memory handle behind them, so closing + * them makes the CUDA driver fail and print a warning */ + scoped_log_handler slh(hide_warns_logger); + uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); + } + + /* Only the idle region is closed, the in-use one stays */ + EXPECT_EQ(1UL, m_cache->num_regions); + EXPECT_EQ(REGION_SIZE, m_cache->total_size); + EXPECT_FALSE(pgtable_has(0)); + EXPECT_TRUE(pgtable_has(1)); +} + class test_cuda_copy_md : public test_md { }; From 63867df3e95223cb9f77e1cb7808cdc9aad84355 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Mon, 19 Oct 2026 13:18:45 +0300 Subject: [PATCH 2/8] UCP/CORE: Free unused mapping when a user remote key is released. --- src/ucp/core/ucp_ep.c | 2 +- src/ucp/core/ucp_rkey.c | 31 +++++++++++++++++++++++++------ src/ucp/core/ucp_rkey.h | 14 +++++++++----- src/ucp/core/ucp_rkey.inl | 5 +++-- src/ucp/rndv/proto_rndv.c | 2 +- src/ucp/rndv/rndv.c | 18 ++++++++++-------- test/gtest/ucp/test_ucp_mmap.cc | 4 ++-- 7 files changed, 51 insertions(+), 25 deletions(-) diff --git a/src/ucp/core/ucp_ep.c b/src/ucp/core/ucp_ep.c index 6629640c002..7f96bd0a8f2 100644 --- a/src/ucp/core/ucp_ep.c +++ b/src/ucp/core/ucp_ep.c @@ -371,7 +371,7 @@ ucp_ep_peer_mem_get(ucp_context_h context, ucp_ep_h ep, uint64_t address, data->size = size; ucp_ep_rkey_unpack_internal(ep, rkey_buf, 0, UCS_BIT(rkey_ptr_md_index), 0, - UCS_SYS_DEVICE_ID_UNKNOWN, &data->rkey); + UCS_SYS_DEVICE_ID_UNKNOWN, 0, &data->rkey); rkey_index = ucs_bitmap2idx(data->rkey->md_map, rkey_ptr_md_index); status = uct_rkey_ptr(data->rkey->tl_rkey[rkey_index].cmpt, &data->rkey->tl_rkey[rkey_index].rkey, address, diff --git a/src/ucp/core/ucp_rkey.c b/src/ucp/core/ucp_rkey.c index e3dea958b66..54ee070d004 100644 --- a/src/ucp/core/ucp_rkey.c +++ b/src/ucp/core/ucp_rkey.c @@ -914,10 +914,11 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rkey_proto_resolve, UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, (ep, buffer, length, unpack_md_map, skip_md_map, sys_dev, - rkey_p), + uct_flags, rkey_p), ucp_ep_h ep, const void *buffer, size_t length, ucp_md_map_t unpack_md_map, ucp_md_map_t skip_md_map, - ucs_sys_device_t sys_dev, ucp_rkey_h *rkey_p) + ucs_sys_device_t sys_dev, uint64_t uct_flags, + ucp_rkey_h *rkey_p) { ucp_worker_h worker = ep->worker; const ucp_ep_config_t *ep_config = ucp_ep_config(ep); @@ -974,8 +975,10 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, rkey->ep = ep; #endif - unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE; + unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE | + UCT_RKEY_UNPACK_FIELD_FLAGS; unpack_params.sys_device = sys_dev; + unpack_params.flags = uct_flags; /* Go over remote MD indices and unpack rkey of each UCT MD */ rkey_index = 0; /* Index of the rkey in the array */ @@ -1056,18 +1059,34 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal, return status; } -ucs_status_t ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer, - ucp_rkey_h *rkey_p) +static ucs_status_t +ucp_ep_rkey_unpack_flags(ucp_ep_h ep, const void *rkey_buffer, + uint64_t uct_flags, ucp_rkey_h *rkey_p) { ucs_status_t status; UCP_WORKER_THREAD_CS_ENTER_CONDITIONAL(ep->worker); - status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, rkey_p); + status = ucp_ep_rkey_unpack_reachable(ep, rkey_buffer, 0, uct_flags, + rkey_p); UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(ep->worker); return status; } +ucs_status_t +ucp_ep_rkey_unpack(ucp_ep_h ep, const void *rkey_buffer, ucp_rkey_h *rkey_p) +{ + return ucp_ep_rkey_unpack_flags(ep, rkey_buffer, + UCT_RKEY_UNPACK_FLAG_FREE_UNUSED_MAPPING, + rkey_p); +} + +ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer, + ucp_rkey_h *rkey_p) +{ + return ucp_ep_rkey_unpack_flags(ep, rkey_buffer, 0, rkey_p); +} + void ucp_rkey_dump_packed(const void *buffer, size_t length, ucs_string_buffer_t *strb) { diff --git a/src/ucp/core/ucp_rkey.h b/src/ucp/core/ucp_rkey.h index 4b1402a8028..5c54b5c944f 100644 --- a/src/ucp/core/ucp_rkey.h +++ b/src/ucp/core/ucp_rkey.h @@ -231,11 +231,15 @@ ucp_memh_exported_unpack(ucp_context_h context, const void *export_mkey_buffer, int ucp_memh_buffer_is_dummy(const void *exported_memh_buffer); -ucs_status_t ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer, - size_t length, - ucp_md_map_t unpack_md_map, - ucp_md_map_t skip_md_map, - ucs_sys_device_t sys_device, +ucs_status_t +ucp_ep_rkey_unpack_internal(ucp_ep_h ep, const void *buffer, size_t length, + ucp_md_map_t unpack_md_map, + ucp_md_map_t skip_md_map, + ucs_sys_device_t sys_device, uint64_t uct_flags, + ucp_rkey_h *rkey_p); + +/* Unpack a remote key passing zero flags to uct_rkey_unpack */ +ucs_status_t ucp_ep_rkey_unpack_wo_flags(ucp_ep_h ep, const void *rkey_buffer, ucp_rkey_h *rkey_p); diff --git a/src/ucp/core/ucp_rkey.inl b/src/ucp/core/ucp_rkey.inl index 0552a193460..7b1d8c68ada 100644 --- a/src/ucp/core/ucp_rkey.inl +++ b/src/ucp/core/ucp_rkey.inl @@ -67,13 +67,14 @@ ucp_rkey_get_tl_rkey(ucp_rkey_h rkey, ucp_md_index_t rkey_index) static UCS_F_ALWAYS_INLINE ucs_status_t ucp_ep_rkey_unpack_reachable(ucp_ep_h ep, const void *buffer, size_t length, - ucp_rkey_h *rkey_p) + uint64_t uct_flags, ucp_rkey_h *rkey_p) { ucp_ep_config_t *config = &ucs_array_elem(&ep->worker->ep_config, ep->cfg_index); return ucp_ep_rkey_unpack_internal(ep, buffer, length, config->key.reachable_md_map, 0, - UCS_SYS_DEVICE_ID_UNKNOWN, rkey_p); + UCS_SYS_DEVICE_ID_UNKNOWN, uct_flags, + rkey_p); } static UCS_F_ALWAYS_INLINE int diff --git a/src/ucp/rndv/proto_rndv.c b/src/ucp/rndv/proto_rndv.c index ea5485e1ea2..f790386ce51 100644 --- a/src/ucp/rndv/proto_rndv.c +++ b/src/ucp/rndv/proto_rndv.c @@ -857,7 +857,7 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_proto_rndv_send_reply, */ status = ucp_ep_rkey_unpack_internal( ep, rkey_buffer, rkey_length, ep_config->key.reachable_md_map, - ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, &rkey); + ep_config->rndv.proto_rndv_rkey_skip_mds, sys_dev, 0, &rkey); if (status != UCS_OK) { goto err; } diff --git a/src/ucp/rndv/rndv.c b/src/ucp/rndv/rndv.c index 31042738a40..80d47e02d4c 100644 --- a/src/ucp/rndv/rndv.c +++ b/src/ucp/rndv/rndv.c @@ -952,7 +952,8 @@ static ucs_status_t ucp_rndv_req_send_rma_get(ucp_request_t *rndv_req, ucp_request_set_super(rndv_req, rreq); - status = ucp_ep_rkey_unpack(ep, rkey_buf, &rndv_req->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf, + &rndv_req->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -1124,7 +1125,7 @@ ucp_rndv_send_frag_update_get_rkey(ucp_worker_h worker, ucp_request_t *freq, &rkey_buffer, &rkey_size); ucs_assert_always(status == UCS_OK); - status = ucp_ep_rkey_unpack(mem_type_ep, rkey_buffer, rkey_p); + status = ucp_ep_rkey_unpack_wo_flags(mem_type_ep, rkey_buffer, rkey_p); ucs_assert_always(status == UCS_OK); ucp_rkey_buffer_release(rkey_buffer); } @@ -1303,8 +1304,8 @@ ucp_rndv_recv_start_get_pipeline(ucp_worker_h worker, ucp_request_t *rndv_req, * Step 3: Send ATS for RNDV request */ - status = ucp_ep_rkey_unpack(rndv_req->send.ep, rkey_buffer, - &rndv_req->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(rndv_req->send.ep, rkey_buffer, + &rndv_req->send.rndv.rkey); if (ucs_unlikely(status != UCS_OK)) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(rndv_req->send.ep), ucs_status_string(status)); @@ -1500,7 +1501,7 @@ static void ucp_rndv_do_rkey_ptr(ucp_request_t *rndv_req, ucp_request_t *rreq, ucp_trace_req(rndv_req, "start rkey_ptr rndv rreq %p", rreq); - status = ucp_ep_rkey_unpack(ep, rkey_buf, &rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rkey_buf, &rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -2123,7 +2124,8 @@ static ucs_status_t ucp_rndv_send_start_put_pipeline(ucp_request_t *sreq, ucp_trace_req(sreq, "using put rndv pipeline protocol: va %p size %zu", (void*)rndv_rtr_hdr->address, rndv_rtr_hdr->size); - status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1, &sreq->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1, + &sreq->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); @@ -2325,8 +2327,8 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler, * and we have to use PUT_ZCOPY RNDV scheme instead */ } - status = ucp_ep_rkey_unpack(ep, rndv_rtr_hdr + 1, - &sreq->send.rndv.rkey); + status = ucp_ep_rkey_unpack_wo_flags(ep, rndv_rtr_hdr + 1, + &sreq->send.rndv.rkey); if (status != UCS_OK) { ucs_fatal("failed to unpack rendezvous remote key received from %s: %s", ucp_ep_peer_name(ep), ucs_status_string(status)); diff --git a/test/gtest/ucp/test_ucp_mmap.cc b/test/gtest/ucp/test_ucp_mmap.cc index 15ecc612383..b124bdc43ba 100644 --- a/test/gtest/ucp/test_ucp_mmap.cc +++ b/test/gtest/ucp/test_ucp_mmap.cc @@ -245,7 +245,7 @@ ucp_rkey_h test_ucp_mmap::mem_chunk::unpack(ucp_ep_h ep, ucp_md_map_t md_map) // Different MD map means different config index on proto v2 ASSERT_UCS_OK(ucp_ep_rkey_unpack_internal(ep, rkey_buffer, rkey_size, md_map, 0, - UCS_SYS_DEVICE_ID_UNKNOWN, + UCS_SYS_DEVICE_ID_UNKNOWN, 0, &rkey)); } @@ -497,7 +497,7 @@ void test_ucp_mmap::test_rkey_proto(ucp_mem_h memh) /* Unpack remote key buffer */ ucp_rkey_h rkey; status = ucp_ep_rkey_unpack_reachable(receiver().ep(), &rkey_buffer[0], - rkey_size, &rkey); + rkey_size, 0, &rkey); ASSERT_UCS_OK(status); /* Check rkey configuration and lanes distances */ From 5ee5af2724281a3c63d1ee14fcaac4bf592376ee Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Fri, 16 Oct 2026 19:25:13 +0300 Subject: [PATCH 3/8] TEST/MPI: Use remote mem list in test_ucp_rkey_destroy. --- contrib/test_jenkins.sh | 10 +++-- test/mpi/test_rma_cuda.c | 2 +- test/mpi/test_ucp.c | 8 ++-- test/mpi/test_ucp.h | 2 +- test/mpi/test_ucp_rkey_destroy.c | 77 +++++++++++++++++++++++++++----- 5 files changed, 79 insertions(+), 20 deletions(-) diff --git a/contrib/test_jenkins.sh b/contrib/test_jenkins.sh index d1b7b4ca845..1b3d82fbefa 100755 --- a/contrib/test_jenkins.sh +++ b/contrib/test_jenkins.sh @@ -1010,10 +1010,12 @@ run_mpi_tests() { # The test checks that after destroying a UCP remote key on the # remote process, the local memory can be released. However, the # internal caches may hold references to the memory, e.g. via - # opened CUDA IPC memory handle. All the caches are disabled - # to be able to release the memory after the remote key is - # destroyed. - ${MPIRUN_COMMON} -np 2 -x UCX_MEMTYPE_CACHE=n \ + # opened CUDA IPC memory handle. All unrelated caches are + # disabled to be able to release the memory after the remote key + # is destroyed. + ${MPIRUN_COMMON} -np 2 \ + -x UCX_TLS=^gdr_copy \ + -x UCX_RCACHE_ENABLE=n \ ./test/mpi/test_ucp_rkey_destroy fi diff --git a/test/mpi/test_rma_cuda.c b/test/mpi/test_rma_cuda.c index 36609a19821..3e9319f48a7 100644 --- a/test/mpi/test_rma_cuda.c +++ b/test/mpi/test_rma_cuda.c @@ -167,7 +167,7 @@ int main(int argc, char **argv) return -1; } - xfer_args.ucp = create_ucp(); + xfer_args.ucp = create_ucp(UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP); for (operation_idx = 0; operation_idx < ucs_static_array_size(operations); ++operation_idx) { diff --git a/test/mpi/test_ucp.c b/test/mpi/test_ucp.c index 22657b3ad5e..207c87de625 100644 --- a/test/mpi/test_ucp.c +++ b/test/mpi/test_ucp.c @@ -11,14 +11,14 @@ #include -static ucp_context_h create_ucp_context() +static ucp_context_h create_ucp_context(uint64_t features) { ucp_params_t params; ucp_config_t *config; ucp_context_h context; params.field_mask = UCP_PARAM_FIELD_FEATURES | UCP_PARAM_FIELD_NAME; - params.features = UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP; + params.features = features; params.name = "test_context"; UCX_CHECK(ucp_config_read(NULL, NULL, &config)); UCX_CHECK(ucp_init(¶ms, config, &context)); @@ -124,10 +124,10 @@ create_ucp_mem(void *buffer, size_t size, ucp_context_h context) return mem; } -ucp_t create_ucp() +ucp_t create_ucp(uint64_t features) { ucp_t ucp; - ucp.context = create_ucp_context(); + ucp.context = create_ucp_context(features); ucp.worker = create_ucp_worker(ucp.context); ucp.ep = create_ucp_ep(ucp.worker); return ucp; diff --git a/test/mpi/test_ucp.h b/test/mpi/test_ucp.h index bc679ab43b2..3f6a615163e 100644 --- a/test/mpi/test_ucp.h +++ b/test/mpi/test_ucp.h @@ -22,7 +22,7 @@ typedef struct { } rkey_t; -ucp_t create_ucp(); +ucp_t create_ucp(uint64_t features); void destroy_ucp(ucp_t); diff --git a/test/mpi/test_ucp_rkey_destroy.c b/test/mpi/test_ucp_rkey_destroy.c index 822f1a18f9a..8ce5bef1a89 100644 --- a/test/mpi/test_ucp_rkey_destroy.c +++ b/test/mpi/test_ucp_rkey_destroy.c @@ -13,6 +13,8 @@ #include "test_ucp.h" #include "test_ucx_check_def.h" +#include "ucp/api/device/ucp_host.h" + #include #include @@ -23,7 +25,28 @@ #define MPI_COMM_SIZE 2 #define SIZE (1024 * 1024 * 1024) -static int rank0(ucp_context_h ucp_context) +static void mpi_barrier(void) +{ + MPI_Request request; + + MPI_Ibarrier(MPI_COMM_WORLD, &request); + MPI_Wait(&request, MPI_STATUS_IGNORE); +} + +/* Progress the worker while waiting, to serve the peer's wireup messages */ +static void mpi_barrier_progress(ucp_worker_h worker) +{ + int done = 0; + MPI_Request request; + + MPI_Ibarrier(MPI_COMM_WORLD, &request); + while (!done) { + ucp_worker_progress(worker); + MPI_Test(&request, &done, MPI_STATUS_IGNORE); + } +} + +static int rank0(ucp_t ucp) { CUdeviceptr ptr; ucp_mem_h ucp_mem; @@ -32,11 +55,11 @@ static int rank0(ucp_context_h ucp_context) CUDA_CHECK(cuMemGetInfo(&free_bytes_before, &total_bytes)); CUDA_CHECK(cuMemAlloc(&ptr, SIZE)); - ucp_mem = send_rkey(1, (void*)ptr, SIZE, ucp_context); + ucp_mem = send_rkey(1, (void*)ptr, SIZE, ucp.context); - MPI_Barrier(MPI_COMM_WORLD); + mpi_barrier_progress(ucp.worker); - UCX_CHECK(ucp_mem_unmap(ucp_context, ucp_mem)); + UCX_CHECK(ucp_mem_unmap(ucp.context, ucp_mem)); CUDA_CHECK(cuMemFree(ptr)); CUDA_CHECK(cuCtxSynchronize()); @@ -47,17 +70,51 @@ static int rank0(ucp_context_h ucp_context) return (unreleased_memory == 0) ? EXIT_SUCCESS : EXIT_FAILURE; } -static void rank1(ucp_ep_h ucp_ep) +static void create_remote_mem_list(ucp_t ucp, const rkey_t *rkey, + ucp_device_remote_mem_list_h *mem_list_h) { + ucp_device_mem_list_elem_t elem; + ucp_device_mem_list_params_t params; + ucs_status_t status; + + elem.field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_EP | + UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR | + UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY; + elem.ep = ucp.ep; + elem.remote_addr = rkey->remote_address; + elem.rkey = rkey->rkey; + + params.field_mask = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS | + UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE | + UCP_DEVICE_MEM_LIST_PARAMS_FIELD_NUM_ELEMENTS; + params.elements = &elem; + params.element_size = sizeof(elem); + params.num_elements = 1; + + do { + ucp_worker_progress(ucp.worker); + status = ucp_device_remote_mem_list_create(¶ms, mem_list_h); + } while (status == UCS_ERR_NOT_CONNECTED); + + UCX_CHECK(status); +} + +static void rank1(ucp_t ucp) +{ + ucp_device_remote_mem_list_h mem_list_h; rkey_t rkey; void *ptr; - rkey = recv_rkey(0, ucp_ep); + rkey = recv_rkey(0, ucp.ep); UCX_CHECK(ucp_rkey_ptr(rkey.rkey, rkey.remote_address, &ptr)); + + create_remote_mem_list(ucp, &rkey, &mem_list_h); + ucp_device_mem_list_release(mem_list_h); + ucp_rkey_destroy(rkey.rkey); - MPI_Barrier(MPI_COMM_WORLD); + mpi_barrier(); } int main(int argc, char **argv) @@ -88,12 +145,12 @@ int main(int argc, char **argv) CUDA_CHECK(cuDevicePrimaryCtxRetain(&cu_ctx, cu_dev)); CUDA_CHECK(cuCtxSetCurrent(cu_ctx)); - ucp = create_ucp(); + ucp = create_ucp(UCP_FEATURE_RMA | UCP_FEATURE_WAKEUP | UCP_FEATURE_DEVICE); if (rank == 0) { - exit_status = rank0(ucp.context); + exit_status = rank0(ucp); } else { - rank1(ucp.ep); + rank1(ucp); } MPI_Bcast(&exit_status, 1, MPI_INT, 0, MPI_COMM_WORLD); From 5905bee2772d98fe5adcd71d3fea986d941cf460 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Tue, 27 Oct 2026 17:01:52 +0200 Subject: [PATCH 4/8] CONTRIB: Reverted merge changes. --- contrib/test_jenkins.sh | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/contrib/test_jenkins.sh b/contrib/test_jenkins.sh index fe21b2f2cd4..9f0428efd6d 100755 --- a/contrib/test_jenkins.sh +++ b/contrib/test_jenkins.sh @@ -1000,22 +1000,8 @@ run_mpi_tests() { ./test/mpi/test_mpi_cuda fi - if [ "X$have_cuda" != "Xno" ] && [ -x ./test/mpi/test_ucp_rkey_destroy ] - then - echo "==== Running UCP RKEY destroy tests ====" - # The test checks that after destroying a UCP remote key on the - # remote process, the local memory can be released. However, the - # internal caches may hold references to the memory, e.g. via - # opened CUDA IPC memory handle. All unrelated caches are - # disabled to be able to release the memory after the remote key - # is destroyed. - ${MPIRUN_COMMON} -np 2 \ - -x UCX_TLS=^gdr_copy \ - -x UCX_RCACHE_ENABLE=n \ - ./test/mpi/test_ucp_rkey_destroy - fi - - # test_rma_cuda is not run: it fails on the CI GPU hosts + # test_rma_cuda and test_ucp_rkey_destroy are not run: they fail + # on the CI GPU hosts # Restore LD_LIBRARY_PATH so subsequent tests will not take UCX libs # from installation directory From 81170eb3e84f0a11796111543d7a46adb3554973 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Tue, 27 Oct 2026 17:55:46 +0200 Subject: [PATCH 5/8] UCT/CUDA/CUDA_IPC: Code refactoring. --- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 56 +++++++----------------- src/uct/cuda/cuda_ipc/cuda_ipc_cache.h | 26 +---------- src/uct/cuda/cuda_ipc/cuda_ipc_iface.c | 2 +- src/uct/cuda/cuda_ipc/cuda_ipc_iface.h | 2 +- src/uct/cuda/cuda_ipc/cuda_ipc_md.c | 26 +++++------ test/gtest/uct/cuda/test_cuda_ipc_md.cc | 58 +------------------------ 6 files changed, 33 insertions(+), 137 deletions(-) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index 9d8a4941970..03971cd89f1 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -639,7 +639,8 @@ uct_cuda_ipc_remote_cache_put(uct_cuda_ipc_cache_hash_key_t key, static void uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, - const void *mapped_addr, int cache_enabled) + const void *mapped_addr, int enable_cache, + int enable_assert) { ucs_pgt_region_t *pgt_region; uct_cuda_ipc_cache_region_t *region; @@ -647,56 +648,31 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, /* use write lock because cache maybe modified */ pthread_rwlock_wrlock(&cache->lock); pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr); - ucs_assert(pgt_region != NULL); + if (enable_assert) { + ucs_assert(pgt_region != NULL); + } + region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); + if (enable_assert) { + ucs_assert(region->refcount >= 1); + } - ucs_assert(region->refcount >= 1); region->refcount--; + if (!region->refcount && !enable_cache) { + if (enable_assert) { + ucs_assert(region->mapped_addr == mapped_addr); + } - if (!region->refcount && !cache_enabled) { - ucs_assert(region->mapped_addr == mapped_addr); uct_cuda_ipc_cache_region_destroy(cache, region); } pthread_rwlock_unlock(&cache->lock); } -void uct_cuda_ipc_cache_release_unused(uct_cuda_ipc_cache_t *cache, - uintptr_t d_bptr) -{ - ucs_pgt_region_t *pgt_region; - uct_cuda_ipc_cache_region_t *region; - - pthread_rwlock_wrlock(&cache->lock); - pgt_region = ucs_pgtable_lookup(&cache->pgtable, d_bptr); - if (pgt_region != NULL) { - region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); - if (region->refcount == 0) { - uct_cuda_ipc_cache_region_destroy(cache, region); - } - } - - pthread_rwlock_unlock(&cache->lock); -} - -void uct_cuda_ipc_cache_release_unused_region(pid_t pid, ucs_sys_ns_t pid_ns, - uintptr_t d_bptr, CUdevice cu_dev) -{ - const uct_cuda_ipc_cache_hash_key_t key = {pid, pid_ns, cu_dev}; - uct_cuda_ipc_cache_t *cache; - - ucs_rw_spinlock_read_lock(&uct_cuda_ipc_remote_cache.lock); - cache = uct_cuda_ipc_remote_cache_get(key); - if (cache != NULL) { - uct_cuda_ipc_cache_release_unused(cache, d_bptr); - } - - ucs_rw_spinlock_read_unlock(&uct_cuda_ipc_remote_cache.lock); -} - void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, uintptr_t d_bptr, const void *mapped_addr, - CUdevice cu_dev, int cache_enabled) + CUdevice cu_dev, int enable_cache, + int enable_assert) { const uct_cuda_ipc_cache_hash_key_t key = {pid, pid_ns, cu_dev}; uct_cuda_ipc_cache_t *cache; @@ -713,7 +689,7 @@ void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, cache = uct_cuda_ipc_remote_cache_get(key); if (cache != NULL) { uct_cuda_ipc_cache_destroy_region(cache, d_bptr, mapped_addr, - cache_enabled); + enable_cache, enable_assert); } else { ucs_debug("no remote cache found for key: %d:%u:%d", pid, pid_ns, cu_dev); diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h index 6e15196b767..2c888823429 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h @@ -72,30 +72,8 @@ ucs_status_t uct_cuda_ipc_map_memhandle(uct_cuda_ipc_extended_rkey_t *key, void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, uintptr_t d_bptr, const void *mapped_addr, - CUdevice cu_dev, int cache_enabled); - - -/** - * @brief Close the region of a buffer in the given cache if it is not in use - * - * @param cache Cache to look for the region in - * @param d_bptr Remote base address of the buffer - */ -void uct_cuda_ipc_cache_release_unused(uct_cuda_ipc_cache_t *cache, - uintptr_t d_bptr); - - -/** - * @brief Close the cached region of a buffer if it is not in use - * - * @param pid Remote process ID - * @param pid_ns Remote process ID namespace - * @param d_bptr Remote base address of the buffer - * @param cu_dev CUDA device handle the region was mapped on - */ -void uct_cuda_ipc_cache_release_unused_region(pid_t pid, ucs_sys_ns_t pid_ns, - uintptr_t d_bptr, - CUdevice cu_dev); + CUdevice cu_dev, int enable_cache, + int enabel_assert); /** diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c index f955cc110c0..22ced4ba1ff 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c @@ -400,7 +400,7 @@ static void uct_cuda_ipc_complete_event(uct_iface_h tl_iface, cuda_ipc_event->d_bptr, cuda_ipc_event->mapped_addr, cuda_ipc_event->cuda_device, - uct_cuda_ipc_component.enable_remote_cache); + uct_cuda_ipc_component.enable_remote_cache, 1); } static uct_iface_ops_t uct_cuda_ipc_iface_ops = { diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h index 9c8fd16c218..12fcedc0eb4 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h @@ -68,7 +68,7 @@ uct_cuda_ipc_sgl_mapping_destroy(uct_cuda_ipc_sgl_mapping_t *mapping, mapping->entries[i].pid_ns, mapping->entries[i].d_bptr, mapping->entries[i].mapped_addr, - cuda_device, enable_cache); + cuda_device, enable_cache, 1); } ucs_free(mapping); diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index 57b7ef9dd95..3d554758683 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -643,23 +643,21 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, return; } + if ((rkey_handle->mapped_addr == NULL) && + !rkey_handle->free_unused_mapping) { + return; + } + unpacked_rkey = (uct_cuda_ipc_unpacked_rkey_t*)uct_rkey; extended_rkey = &unpacked_rkey->super; rkey = &extended_rkey->super; - if (rkey_handle->mapped_addr != NULL) { - uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, - rkey->d_bptr, rkey_handle->mapped_addr, - rkey_handle->cu_dev, - uct_cuda_ipc_component.enable_remote_cache); - } - - if (rkey_handle->free_unused_mapping) { - uct_cuda_ipc_cache_release_unused_region(rkey->pid, - extended_rkey->pid_ns, - rkey->d_bptr, - rkey_handle->cu_dev); - } + uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, + rkey->d_bptr, rkey_handle->mapped_addr, + rkey_handle->cu_dev, + uct_cuda_ipc_component.enable_remote_cache && + !rkey_handle->free_unused_mapping, + rkey_handle->mapped_addr != NULL); } static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component, @@ -829,7 +827,7 @@ static void uct_cuda_ipc_md_mem_elem_release(uct_md_h md, void *release_handle) uct_cuda_ipc_unmap_memhandle(handle->pid, handle->pid_ns, handle->d_bptr, handle->mapped_addr, handle->cu_dev, - uct_cuda_ipc_component.enable_remote_cache); + uct_cuda_ipc_component.enable_remote_cache, 1); ucs_free(handle); } diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index 0accf90df89..4455d175e93 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -372,7 +372,7 @@ UCS_TEST_P(test_cuda_ipc_md, posix_fd_same_node_ipc) uct_cuda_ipc_unmap_memhandle( unpacked->super.super.pid, unpacked->super.pid_ns, - unpacked->super.super.d_bptr, mapped_addr, dev, 0); + unpacked->super.super.d_bptr, mapped_addr, dev, 0, 1); EXPECT_UCS_OK(uct_rkey_release(component, &rkey_bundle)); EXPECT_EQ(cuCtxDestroy(ctx), CUDA_SUCCESS); } catch (...) { @@ -716,62 +716,6 @@ UCS_TEST_F(test_cuda_ipc_cache_lru, stale_destroy_while_in_use) { EXPECT_TRUE(ucs_list_is_empty(&m_cache->lru_list)); } -UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_keeps_in_use) { - create_cache(ULONG_MAX, SIZE_MAX); - - insert_region(0); - uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); - - /* The region is in use, so it must stay in the cache */ - EXPECT_EQ(1UL, m_cache->num_regions); - EXPECT_EQ(REGION_SIZE, m_cache->total_size); - EXPECT_TRUE(pgtable_has(0)); -} - -UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_keeps_reacquired) { - create_cache(ULONG_MAX, SIZE_MAX); - - uct_cuda_ipc_cache_region_t *region = insert_region(0); - release_region(region); - reacquire_region(region); - uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); - - EXPECT_EQ(1UL, m_cache->num_regions); - EXPECT_TRUE(pgtable_has(0)); -} - -UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_no_region) { - create_cache(ULONG_MAX, SIZE_MAX); - - insert_region(0); - /* No region at this address, the cache must not change */ - uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR + 4 * REGION_SIZE); - - EXPECT_EQ(1UL, m_cache->num_regions); - EXPECT_TRUE(pgtable_has(0)); -} - -UCS_TEST_F(test_cuda_ipc_cache_lru, release_unused_closes_idle) { - create_cache(ULONG_MAX, SIZE_MAX); - - uct_cuda_ipc_cache_region_t *region0 = insert_region(0); - insert_region(1); - release_region(region0); - - { - /* The test regions have no real memory handle behind them, so closing - * them makes the CUDA driver fail and print a warning */ - scoped_log_handler slh(hide_warns_logger); - uct_cuda_ipc_cache_release_unused(m_cache, BASE_ADDR); - } - - /* Only the idle region is closed, the in-use one stays */ - EXPECT_EQ(1UL, m_cache->num_regions); - EXPECT_EQ(REGION_SIZE, m_cache->total_size); - EXPECT_FALSE(pgtable_has(0)); - EXPECT_TRUE(pgtable_has(1)); -} - class test_cuda_copy_md : public test_md { }; From 9303c0d2464819a9e176ced7cffd907cde7763f7 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Tue, 27 Oct 2026 18:24:01 +0200 Subject: [PATCH 6/8] UCT/CUDA/CUDA_IPC: Removed extra argument. --- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 18 ++++++++---------- src/uct/cuda/cuda_ipc/cuda_ipc_cache.h | 3 +-- src/uct/cuda/cuda_ipc/cuda_ipc_iface.c | 2 +- src/uct/cuda/cuda_ipc/cuda_ipc_iface.h | 2 +- src/uct/cuda/cuda_ipc/cuda_ipc_md.c | 5 ++--- test/gtest/uct/cuda/test_cuda_ipc_md.cc | 2 +- 6 files changed, 14 insertions(+), 18 deletions(-) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index 03971cd89f1..6cbaf52ece9 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -639,8 +639,7 @@ uct_cuda_ipc_remote_cache_put(uct_cuda_ipc_cache_hash_key_t key, static void uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, - const void *mapped_addr, int enable_cache, - int enable_assert) + const void *mapped_addr, int cache_enabled) { ucs_pgt_region_t *pgt_region; uct_cuda_ipc_cache_region_t *region; @@ -648,18 +647,18 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, /* use write lock because cache maybe modified */ pthread_rwlock_wrlock(&cache->lock); pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr); - if (enable_assert) { + if (mapped_addr != NULL) { ucs_assert(pgt_region != NULL); } region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); - if (enable_assert) { + if (mapped_addr != NULL) { ucs_assert(region->refcount >= 1); + region->refcount--; } - region->refcount--; - if (!region->refcount && !enable_cache) { - if (enable_assert) { + if (!region->refcount && !cache_enabled) { + if (mapped_addr != NULL) { ucs_assert(region->mapped_addr == mapped_addr); } @@ -671,8 +670,7 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, uintptr_t d_bptr, const void *mapped_addr, - CUdevice cu_dev, int enable_cache, - int enable_assert) + CUdevice cu_dev, int cache_enabled) { const uct_cuda_ipc_cache_hash_key_t key = {pid, pid_ns, cu_dev}; uct_cuda_ipc_cache_t *cache; @@ -689,7 +687,7 @@ void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, cache = uct_cuda_ipc_remote_cache_get(key); if (cache != NULL) { uct_cuda_ipc_cache_destroy_region(cache, d_bptr, mapped_addr, - enable_cache, enable_assert); + cache_enabled); } else { ucs_debug("no remote cache found for key: %d:%u:%d", pid, pid_ns, cu_dev); diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h index 2c888823429..c1d553f55b1 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.h @@ -72,8 +72,7 @@ ucs_status_t uct_cuda_ipc_map_memhandle(uct_cuda_ipc_extended_rkey_t *key, void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns, uintptr_t d_bptr, const void *mapped_addr, - CUdevice cu_dev, int enable_cache, - int enabel_assert); + CUdevice cu_dev, int cache_enabled); /** diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c index 22ced4ba1ff..f955cc110c0 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c @@ -400,7 +400,7 @@ static void uct_cuda_ipc_complete_event(uct_iface_h tl_iface, cuda_ipc_event->d_bptr, cuda_ipc_event->mapped_addr, cuda_ipc_event->cuda_device, - uct_cuda_ipc_component.enable_remote_cache, 1); + uct_cuda_ipc_component.enable_remote_cache); } static uct_iface_ops_t uct_cuda_ipc_iface_ops = { diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h index 12fcedc0eb4..9c8fd16c218 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.h @@ -68,7 +68,7 @@ uct_cuda_ipc_sgl_mapping_destroy(uct_cuda_ipc_sgl_mapping_t *mapping, mapping->entries[i].pid_ns, mapping->entries[i].d_bptr, mapping->entries[i].mapped_addr, - cuda_device, enable_cache, 1); + cuda_device, enable_cache); } ucs_free(mapping); diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index 3d554758683..d3c071578dc 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -656,8 +656,7 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, rkey->d_bptr, rkey_handle->mapped_addr, rkey_handle->cu_dev, uct_cuda_ipc_component.enable_remote_cache && - !rkey_handle->free_unused_mapping, - rkey_handle->mapped_addr != NULL); + !rkey_handle->free_unused_mapping); } static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component, @@ -827,7 +826,7 @@ static void uct_cuda_ipc_md_mem_elem_release(uct_md_h md, void *release_handle) uct_cuda_ipc_unmap_memhandle(handle->pid, handle->pid_ns, handle->d_bptr, handle->mapped_addr, handle->cu_dev, - uct_cuda_ipc_component.enable_remote_cache, 1); + uct_cuda_ipc_component.enable_remote_cache); ucs_free(handle); } diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index 4455d175e93..0cfbf5f8d84 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -372,7 +372,7 @@ UCS_TEST_P(test_cuda_ipc_md, posix_fd_same_node_ipc) uct_cuda_ipc_unmap_memhandle( unpacked->super.super.pid, unpacked->super.pid_ns, - unpacked->super.super.d_bptr, mapped_addr, dev, 0, 1); + unpacked->super.super.d_bptr, mapped_addr, dev, 0); EXPECT_UCS_OK(uct_rkey_release(component, &rkey_bundle)); EXPECT_EQ(cuCtxDestroy(ctx), CUDA_SUCCESS); } catch (...) { From 834f54f7b318a34b430499d4012181a38d6557f6 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Wed, 28 Oct 2026 12:22:45 +0200 Subject: [PATCH 7/8] UCT/CUDA/CUDA_IPC: Fixed NULL dereference. --- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index 6cbaf52ece9..69a25dc0e83 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -647,8 +647,10 @@ uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr, /* use write lock because cache maybe modified */ pthread_rwlock_wrlock(&cache->lock); pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr); - if (mapped_addr != NULL) { - ucs_assert(pgt_region != NULL); + if (pgt_region == NULL) { + ucs_assert(mapped_addr == NULL); + pthread_rwlock_unlock(&cache->lock); + return; } region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t); From cf4689ff2f46fd10207e6fe4c2941ef4a058fd36 Mon Sep 17 00:00:00 2001 From: Raul Akhmetshin Date: Wed, 28 Oct 2026 12:28:59 +0200 Subject: [PATCH 8/8] UCT/CUDA/CUDA_IPC: Fixed code format issues. --- src/uct/cuda/cuda_ipc/cuda_ipc_md.c | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c index d3c071578dc..f34f3ed62ef 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_md.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_md.c @@ -652,11 +652,10 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey, extended_rkey = &unpacked_rkey->super; rkey = &extended_rkey->super; - uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, - rkey->d_bptr, rkey_handle->mapped_addr, - rkey_handle->cu_dev, + uct_cuda_ipc_unmap_memhandle(rkey->pid, extended_rkey->pid_ns, rkey->d_bptr, + rkey_handle->mapped_addr, rkey_handle->cu_dev, uct_cuda_ipc_component.enable_remote_cache && - !rkey_handle->free_unused_mapping); + !rkey_handle->free_unused_mapping); } static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component,