From 7e94ed47738c83b4939e3bf79feabf535664153f Mon Sep 17 00:00:00 2001 From: Michal Shalev Date: Wed, 23 Sep 2026 19:52:57 -0700 Subject: [PATCH 1/4] UCT/CUDA_IPC: Avoid UUID query for non-local rkeys --- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 27 ++++++------ test/gtest/uct/cuda/test_cuda_ipc_md.cc | 56 +++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 13 deletions(-) diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c index 69a25dc0e83..fa685bc5237 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -848,20 +848,21 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_map_memhandle, ucs_status_t status; CUuuid uuid; - status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetUuid(&uuid, cu_dev)); - if (status != UCS_OK) { - return status; - } + if (uct_cuda_ipc_is_rkey_local(key->pid, ext_key->pid_ns)) { + status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetUuid(&uuid, cu_dev)); + if (status != UCS_OK) { + return status; + } - if (uct_cuda_ipc_is_rkey_local(key->pid, ext_key->pid_ns) && - (memcmp(uuid.bytes, key->uuid.bytes, sizeof(uuid.bytes)) == 0)) { - /* TODO: added for test purpose to enable cuda_ipc tests in gtest - * mapped addrr is set to be same as d_bptr avoiding any calls to - * uct_cuda_ipc_open_memhandle which would fail with invalid argument - * error - */ - *mapped_addr = (CUdeviceptr*)key->d_bptr; - return UCS_OK; + if (memcmp(uuid.bytes, key->uuid.bytes, sizeof(uuid.bytes)) == 0) { + /* TODO: added for test purpose to enable cuda_ipc tests in gtest + * mapped addrr is set to be same as d_bptr avoiding any calls to + * uct_cuda_ipc_open_memhandle which would fail with invalid argument + * error + */ + *mapped_addr = (CUdeviceptr*)key->d_bptr; + return UCS_OK; + } } ucs_rw_spinlock_write_lock(&uct_cuda_ipc_remote_cache.lock); diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index 5414ca7a59c..f3af6dacf08 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -176,6 +176,34 @@ class test_cuda_ipc_md : public test_md { EXPECT_UCS_OK(uct_md_mem_dereg_v2(md(), &dereg_params)); } + void reg_and_unpack(CUdeviceptr ptr, size_t size, uct_mem_h *memh, + uct_rkey_bundle_t *rkey_bundle) + { + uct_md_mem_reg_params_t reg_params = {}; + uct_rkey_unpack_params_t unpack_params = {}; + uct_cuda_ipc_extended_rkey_t rkey; + + reg_params.field_mask = UCT_MD_MEM_REG_FIELD_MEM_TYPE; + reg_params.mem_type = UCS_MEMORY_TYPE_CUDA; + ASSERT_UCS_OK(uct_md_mem_reg_v2(md(), (void*)ptr, size, ®_params, + memh)); + ASSERT_UCS_OK(md()->ops->mkey_pack(md(), *memh, (void*)ptr, size, NULL, + &rkey)); + ASSERT_UCS_OK(uct_rkey_unpack_v2(md()->component, &rkey, + &unpack_params, rkey_bundle)); + } + + void release_and_dereg(uct_mem_h memh, uct_rkey_bundle_t *rkey_bundle) + { + uct_md_mem_dereg_params_t dereg_params; + + EXPECT_UCS_OK(uct_rkey_release(md()->component, rkey_bundle)); + + dereg_params.field_mask = UCT_MD_MEM_DEREG_FIELD_MEMH; + dereg_params.memh = memh; + EXPECT_UCS_OK(uct_md_mem_dereg_v2(md(), &dereg_params)); + } + #if HAVE_DECL_SYS_PIDFD_GETFD ucs_status_t pack_posix_fd_key(cuda_posix_fd_mem_buffer &buf, size_t size, uct_mem_h *memh, uct_cuda_ipc_rkey_t *rkey) @@ -418,6 +446,34 @@ UCS_TEST_P(test_cuda_ipc_md, posix_fd_same_node_ipc) #endif } +UCS_TEST_P(test_cuda_ipc_md, map_memhandle_same_process) +{ + constexpr size_t size = 4096; + uct_rkey_bundle_t rkey_bundle = {}; + uct_mem_h memh; + CUdeviceptr ptr; + CUdevice cu_dev; + CUuuid uuid; + void *mapped_addr; + + ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc(&ptr, size)); + ASSERT_NO_FATAL_FAILURE(reg_and_unpack(ptr, size, &memh, &rkey_bundle)); + + auto unpacked = (uct_cuda_ipc_unpacked_rkey_t*)rkey_bundle.rkey; + ASSERT_EQ(CUDA_SUCCESS, cuCtxGetDevice(&cu_dev)); + ASSERT_EQ(CUDA_SUCCESS, cuDeviceGetUuid(&uuid, cu_dev)); + EXPECT_EQ(0, memcmp(unpacked->super.super.uuid.bytes, uuid.bytes, + sizeof(uuid.bytes))); + + ASSERT_UCS_OK(uct_cuda_ipc_map_memhandle(&unpacked->super, cu_dev, + &mapped_addr, + UCS_LOG_LEVEL_ERROR)); + EXPECT_EQ((void*)unpacked->super.super.d_bptr, mapped_addr); + + release_and_dereg(memh, &rkey_bundle); + EXPECT_EQ(CUDA_SUCCESS, cuMemFree(ptr)); +} + _UCT_MD_INSTANTIATE_TEST_CASE(test_cuda_ipc_md, cuda_ipc); From 94789e7ebc019e72afc76109338ed04328af4bcc Mon Sep 17 00:00:00 2001 From: Michal Shalev Date: Tue, 3 Nov 2026 20:35:12 -0800 Subject: [PATCH 2/4] TEST/CUDA_IPC: Use uct_md_mkey_pack_v2 in map_memhandle test --- test/gtest/uct/cuda/test_cuda_ipc_md.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index f3af6dacf08..caff0d799c0 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -180,6 +180,7 @@ class test_cuda_ipc_md : public test_md { uct_rkey_bundle_t *rkey_bundle) { uct_md_mem_reg_params_t reg_params = {}; + uct_md_mkey_pack_params_t pack_params = {}; uct_rkey_unpack_params_t unpack_params = {}; uct_cuda_ipc_extended_rkey_t rkey; @@ -187,8 +188,8 @@ class test_cuda_ipc_md : public test_md { reg_params.mem_type = UCS_MEMORY_TYPE_CUDA; ASSERT_UCS_OK(uct_md_mem_reg_v2(md(), (void*)ptr, size, ®_params, memh)); - ASSERT_UCS_OK(md()->ops->mkey_pack(md(), *memh, (void*)ptr, size, NULL, - &rkey)); + ASSERT_UCS_OK(uct_md_mkey_pack_v2(md(), *memh, (void*)ptr, size, + &pack_params, &rkey)); ASSERT_UCS_OK(uct_rkey_unpack_v2(md()->component, &rkey, &unpack_params, rkey_bundle)); } From 9e81befb11f7904c6ebdb38d54cb186875fa1b3c Mon Sep 17 00:00:00 2001 From: Michal Shalev Date: Wed, 4 Nov 2026 02:44:02 -0800 Subject: [PATCH 3/4] UCT/CUDA_IPC: Re-wrap comment to 80 columns --- src/uct/cuda/cuda_ipc/cuda_ipc_cache.c | 4 ++-- 1 file changed, 2 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 fa685bc5237..d3b3d45bf56 100644 --- a/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c +++ b/src/uct/cuda/cuda_ipc/cuda_ipc_cache.c @@ -857,8 +857,8 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_map_memhandle, if (memcmp(uuid.bytes, key->uuid.bytes, sizeof(uuid.bytes)) == 0) { /* TODO: added for test purpose to enable cuda_ipc tests in gtest * mapped addrr is set to be same as d_bptr avoiding any calls to - * uct_cuda_ipc_open_memhandle which would fail with invalid argument - * error + * uct_cuda_ipc_open_memhandle which would fail with invalid + * argument error */ *mapped_addr = (CUdeviceptr*)key->d_bptr; return UCS_OK; From 4bd76b175eb3ece2791e48166c1fd8c3f23c5f5a Mon Sep 17 00:00:00 2001 From: Michal Shalev Date: Wed, 4 Nov 2026 23:59:19 -0800 Subject: [PATCH 4/4] TEST/CUDA_IPC: Drop map_memhandle_same_process test --- test/gtest/uct/cuda/test_cuda_ipc_md.cc | 57 ------------------------- 1 file changed, 57 deletions(-) diff --git a/test/gtest/uct/cuda/test_cuda_ipc_md.cc b/test/gtest/uct/cuda/test_cuda_ipc_md.cc index caff0d799c0..5414ca7a59c 100644 --- a/test/gtest/uct/cuda/test_cuda_ipc_md.cc +++ b/test/gtest/uct/cuda/test_cuda_ipc_md.cc @@ -176,35 +176,6 @@ class test_cuda_ipc_md : public test_md { EXPECT_UCS_OK(uct_md_mem_dereg_v2(md(), &dereg_params)); } - void reg_and_unpack(CUdeviceptr ptr, size_t size, uct_mem_h *memh, - uct_rkey_bundle_t *rkey_bundle) - { - uct_md_mem_reg_params_t reg_params = {}; - uct_md_mkey_pack_params_t pack_params = {}; - uct_rkey_unpack_params_t unpack_params = {}; - uct_cuda_ipc_extended_rkey_t rkey; - - reg_params.field_mask = UCT_MD_MEM_REG_FIELD_MEM_TYPE; - reg_params.mem_type = UCS_MEMORY_TYPE_CUDA; - ASSERT_UCS_OK(uct_md_mem_reg_v2(md(), (void*)ptr, size, ®_params, - memh)); - ASSERT_UCS_OK(uct_md_mkey_pack_v2(md(), *memh, (void*)ptr, size, - &pack_params, &rkey)); - ASSERT_UCS_OK(uct_rkey_unpack_v2(md()->component, &rkey, - &unpack_params, rkey_bundle)); - } - - void release_and_dereg(uct_mem_h memh, uct_rkey_bundle_t *rkey_bundle) - { - uct_md_mem_dereg_params_t dereg_params; - - EXPECT_UCS_OK(uct_rkey_release(md()->component, rkey_bundle)); - - dereg_params.field_mask = UCT_MD_MEM_DEREG_FIELD_MEMH; - dereg_params.memh = memh; - EXPECT_UCS_OK(uct_md_mem_dereg_v2(md(), &dereg_params)); - } - #if HAVE_DECL_SYS_PIDFD_GETFD ucs_status_t pack_posix_fd_key(cuda_posix_fd_mem_buffer &buf, size_t size, uct_mem_h *memh, uct_cuda_ipc_rkey_t *rkey) @@ -447,34 +418,6 @@ UCS_TEST_P(test_cuda_ipc_md, posix_fd_same_node_ipc) #endif } -UCS_TEST_P(test_cuda_ipc_md, map_memhandle_same_process) -{ - constexpr size_t size = 4096; - uct_rkey_bundle_t rkey_bundle = {}; - uct_mem_h memh; - CUdeviceptr ptr; - CUdevice cu_dev; - CUuuid uuid; - void *mapped_addr; - - ASSERT_EQ(CUDA_SUCCESS, cuMemAlloc(&ptr, size)); - ASSERT_NO_FATAL_FAILURE(reg_and_unpack(ptr, size, &memh, &rkey_bundle)); - - auto unpacked = (uct_cuda_ipc_unpacked_rkey_t*)rkey_bundle.rkey; - ASSERT_EQ(CUDA_SUCCESS, cuCtxGetDevice(&cu_dev)); - ASSERT_EQ(CUDA_SUCCESS, cuDeviceGetUuid(&uuid, cu_dev)); - EXPECT_EQ(0, memcmp(unpacked->super.super.uuid.bytes, uuid.bytes, - sizeof(uuid.bytes))); - - ASSERT_UCS_OK(uct_cuda_ipc_map_memhandle(&unpacked->super, cu_dev, - &mapped_addr, - UCS_LOG_LEVEL_ERROR)); - EXPECT_EQ((void*)unpacked->super.super.d_bptr, mapped_addr); - - release_and_dereg(memh, &rkey_bundle); - EXPECT_EQ(CUDA_SUCCESS, cuMemFree(ptr)); -} - _UCT_MD_INSTANTIATE_TEST_CASE(test_cuda_ipc_md, cuda_ipc);