UCP/API: Endpoint parameter controlling remote key resources.

contrib/test_jenkins.sh

@@ -972,10 +972,14 @@ 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. The test creates its endpoint
+				# with UCP_EP_PARAMS_FLAGS_RKEY_BOUND_LIFETIME so the CUDA IPC
+				# handle is released when the remote key is destroyed; other
+				# caches unrelated to that are disabled here so they don't
+				# interfere with the memory-release check.
+				${MPIRUN_COMMON} -np 2 \
+								-x UCX_TLS=^gdr_copy \
+								-x UCX_RCACHE_ENABLE=n \
 								./test/mpi/test_ucp_rkey_destroy
 			fi
 

src/ucp/api/ucp.h

@@ -297,14 +297,20 @@ enum ucp_ep_params_flags_field {
                                                            send to a particular
                                                            remote endpoint, for
                                                            example stream */
-    UCP_EP_PARAMS_FLAGS_SEND_CLIENT_ID = UCS_BIT(2)   /**< Send client id
+    UCP_EP_PARAMS_FLAGS_SEND_CLIENT_ID = UCS_BIT(2), /**< Send client id
                                                            when connecting to remote
                                                            socket address as part of the
                                                            connection request payload.
                                                            On the remote side value
                                                            can be obtained from
                                                            @ref ucp_conn_request_h using
                                                            @ref ucp_conn_request_query */
+    /**
+     * Resources used by remote keys unpacked (see @ref ucp_ep_rkey_unpack) on
+     * this endpoint have their lifetime bound to the remote key, and are
+     * released once it is destroyed (see @ref ucp_rkey_destroy).
+     */
+    UCP_EP_PARAMS_FLAGS_RKEY_BOUND_LIFETIME = UCS_BIT(3)
 };
 
 

src/ucp/core/ucp_device.c

@@ -250,8 +250,8 @@ static ucs_status_t ucp_device_local_mem_list_element_pack(
         return UCS_ERR_INVALID_PARAM;
     }
 
-    status = uct_md_mem_elem_pack(ucp_md->md, uct_memh, UCT_INVALID_RKEY,
-                                  mem_element, &release_handle);
+    status = uct_md_mem_elem_pack(ucp_md->md, uct_memh, NULL, mem_element,
+                                  &release_handle);
     if (status != UCS_OK) {
         ucs_error("failed to pack local mem element for memh=%p", memh);
         return status;
@@ -534,7 +534,7 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack(
     ucp_ep_config_t *ep_config = ucp_ep_config(ep);
     uct_md_h md;
     uint8_t rkey_index;
-    uct_rkey_t uct_rkey;
+    uct_rkey_bundle_t *uct_rkey;
     uct_ep_h uct_ep;
     uct_device_ep_h device_ep;
     void *release_handle;
@@ -556,8 +556,8 @@ static ucs_status_t ucp_device_remote_mem_list_element_pack(
 
     rkey_index = ucs_bitmap2idx(rkey->md_map,
                                 ep_config->key.lanes[lane].dst_md_index);
-    uct_rkey   = ucp_rkey_get_tl_rkey(rkey, rkey_index);
-    ucs_assert(uct_rkey != UCT_INVALID_RKEY);
+    uct_rkey   = &rkey->tl_rkey[rkey_index].rkey;
+    ucs_assert(uct_rkey->rkey != UCT_INVALID_RKEY);
 
     md              = ucp_ep_md(ep, lane);
     mem_element->ep = device_ep;

src/ucp/core/ucp_ep.c

@@ -274,6 +274,7 @@ static ucp_ep_h ucp_ep_allocate(ucp_worker_h worker, const char *peer_name)
     ep->ext->fence_seq                    = 0;
     ep->ext->uct_eps                      = NULL;
     ep->ext->flush_sys_dev_map            = 0;
+    ep->ext->flags                        = 0;
 
     UCS_STATIC_ASSERT(sizeof(ep->ext->ep_match) >=
                       sizeof(ep->ext->flush_state));
@@ -1284,6 +1285,10 @@ ucs_status_t ucp_ep_create(ucp_worker_h worker, const ucp_ep_params_t *params,
         }
 #endif
 
+        if (flags & UCP_EP_PARAMS_FLAGS_RKEY_BOUND_LIFETIME) {
+            ep->ext->flags |= UCP_EP_EXT_FLAG_RKEY_BOUND_LIFETIME;
+        }
+
         ucp_ep_params_check_err_handling(ep, params);
         ucp_ep_update_flags(ep, UCP_EP_FLAG_USED, 0);
         *ep_p = ep;

src/ucp/core/ucp_ep.h

@@ -517,6 +517,11 @@ typedef struct ucp_ep_recovery_arg {
 } ucp_ep_recovery_arg_t;
 
 
+enum ucp_ep_ext_flags {
+    UCP_EP_EXT_FLAG_RKEY_BOUND_LIFETIME = UCS_BIT(0)
+};
+
+
 /**
  * Endpoint extension
  */
@@ -581,6 +586,9 @@ typedef struct ucp_ep_ext {
      * Map of system devices that require a flush operation
      */
     ucp_sys_dev_map_t             flush_sys_dev_map;
+
+    /* Flags from ucp_ep_ext_flags */
+    uint8_t flags;
 } ucp_ep_ext_t;
 
 

src/ucp/core/ucp_rkey.c

@@ -976,6 +976,10 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_ep_rkey_unpack_internal,
 
     unpack_params.field_mask = UCT_RKEY_UNPACK_FIELD_SYS_DEVICE;
     unpack_params.sys_device = sys_dev;
+    if (ep->ext->flags & UCP_EP_EXT_FLAG_RKEY_BOUND_LIFETIME) {
+        unpack_params.field_mask |= UCT_RKEY_UNPACK_FIELD_FLAGS;
+        unpack_params.flags       = UCT_RKEY_UNPACK_FLAG_BOUND_LIFETIME;
+    }
 
     /* Go over remote MD indices and unpack rkey of each UCT MD */
     rkey_index = 0; /* Index of the rkey in the array */

src/uct/api/v2/uct_v2.h

@@ -765,10 +765,26 @@ 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)
 } uct_rkey_unpack_field_mask_t;
 
 
+/**
+ * @ingroup UCT_MD
+ * @brief Rkey unpack flags.
+ */
+typedef enum {
+    /**
+     * Indicate that resources required for accessing the unpacked remote key
+     * are created at unpack time, and their lifetime is bound to the remote
+     * key: they must be released once the key is released (see
+     * @ref uct_rkey_release).
+     */
+    UCT_RKEY_UNPACK_FLAG_BOUND_LIFETIME = UCS_BIT(0)
+} uct_rkey_unpack_flags_t;
+
+
 /**
  * @ingroup UCT_MD
  * @brief Parameters for unpacking remote key using @ref uct_rkey_unpack_v2.
@@ -786,6 +802,11 @@ typedef struct uct_rkey_unpack_params {
      * (default behavior).
      */
     ucs_sys_device_t     sys_device;
+
+    /**
+     * Flags to unpack rkey with, using bits from @ref uct_rkey_unpack_flags_t.
+     */
+    uint64_t             flags;
 } uct_rkey_unpack_params_t;
 
 
@@ -1592,14 +1613,17 @@ ucs_status_t uct_rkey_unpack_v2(uct_component_h component,
  *
  * @param [in]  md                Memory domain.
  * @param [in]  memh              Memory handle to pack (can be NULL).
- * @param [in]  rkey              Remote key to pack (can be UCT_INVALID_RKEY).
+ * @param [in]  rkey_ob           Remote key bundle as returned by
+ *                                the @ref uct_rkey_unpack_v2 function
+ *                                (can be NULL).
  * @param [out] mem_elem          Filled with the packed memh and rkey.
  * @param [out] release_handle_p  Handle for releasing resources allocated
  *                                during packing.
  *
  * @return UCS_OK on success or error code in case of failure.
  */
-ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
+ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh,
+                                  const uct_rkey_bundle_t *rkey_ob,
                                   uct_device_mem_elem_t *mem_elem,
                                   void **release_handle_p);
 

src/uct/base/uct_md.c

@@ -751,11 +751,15 @@ ucs_status_t uct_md_dummy_mem_dereg(uct_md_h uct_md,
     return UCS_OK;
 }
 
-ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
+ucs_status_t uct_md_mem_elem_pack(uct_md_h md, uct_mem_h memh,
+                                  const uct_rkey_bundle_t *rkey_ob,
                                   uct_device_mem_elem_t *mem_elem,
                                   void **release_handle_p)
 {
-    return md->ops->mem_elem_pack(md, memh, rkey, mem_elem, release_handle_p);
+    uct_rkey_t rkey           = rkey_ob ? rkey_ob->rkey : UCT_INVALID_RKEY;
+    void *rkey_release_handle = rkey_ob ? rkey_ob->handle : NULL;
+    return md->ops->mem_elem_pack(md, memh, rkey, rkey_release_handle, mem_elem,
+                                  release_handle_p);
 }
 
 void uct_md_mem_elem_release(uct_md_h md, void *release_handle)

src/uct/base/uct_md.h

@@ -127,7 +127,7 @@ typedef ucs_status_t (*uct_md_detect_memory_type_func_t)(uct_md_h md,
                                                          ucs_memory_type_t *mem_type_p);
 
 typedef ucs_status_t (*uct_md_mem_elem_pack_func_t)(
-        uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
+        uct_md_h md, uct_mem_h memh, uct_rkey_t rkey, void *rkey_release_handle,
         uct_device_mem_elem_t *mem_elem, void **release_handle_p);
 
 typedef void (*uct_md_mem_elem_release_func_t)(uct_md_h md,

src/uct/cuda/cuda_ipc/cuda_ipc_md.c

@@ -39,6 +39,7 @@
 typedef struct {
     const void *mapped_addr;
     CUdevice   cu_dev;
+    int        cache_enabled;
 } uct_cuda_ipc_rkey_handle_t;
 
 typedef struct {
@@ -47,6 +48,7 @@ typedef struct {
     uintptr_t    d_bptr;
     void         *mapped_addr;
     CUdevice     cu_dev;
+    int          cache_enabled;
 } uct_cuda_ipc_mem_elem_release_handle_t;
 
 static ucs_config_field_t uct_cuda_ipc_md_config_table[] = {
@@ -493,7 +495,7 @@ static ucs_status_t
 uct_cuda_ipc_is_peer_accessible(uct_cuda_ipc_component_t *component,
                                 uct_cuda_ipc_unpacked_rkey_t *rkey,
                                 uct_cuda_ipc_rkey_handle_t *rkey_handle,
-                                CUdevice cu_dev)
+                                CUdevice cu_dev, int bound_lifetime)
 {
     ucs_status_t status;
     void *d_mapped;
@@ -515,8 +517,8 @@ uct_cuda_ipc_is_peer_accessible(uct_cuda_ipc_component_t *component,
      * stream sequentialization */
     rkey->stream_id = cache->dev_num;
     accessible      = &cache->accessible[cu_dev];
-    if (ucs_unlikely(*accessible == UCS_TRY)) { /* unchecked, add to cache */
-
+    if (ucs_unlikely(*accessible == UCS_TRY) || /* unchecked, add to cache */
+        (bound_lifetime && (*accessible == UCS_YES))) {
         /* Check if peer is reachable by trying to open memory handle. This is
          * necessary when the device is not visible through CUDA_VISIBLE_DEVICES
          * and checking peer accessibility through CUDA driver API is not
@@ -539,8 +541,12 @@ uct_cuda_ipc_is_peer_accessible(uct_cuda_ipc_component_t *component,
             rkey_handle->mapped_addr = d_mapped;
         }
 
-        *accessible = ((status == UCS_OK) || (status == UCS_ERR_ALREADY_EXISTS))
-                      ? UCS_YES : UCS_NO;
+        if (*accessible == UCS_TRY) {
+            *accessible = ((status == UCS_OK) ||
+                           (status == UCS_ERR_ALREADY_EXISTS)) ?
+                                  UCS_YES :
+                                  UCS_NO;
+        }
     }
 
     status = (*accessible == UCS_YES) ? UCS_OK : UCS_ERR_UNREACHABLE;
@@ -550,6 +556,13 @@ uct_cuda_ipc_is_peer_accessible(uct_cuda_ipc_component_t *component,
     return status;
 }
 
+static int
+uct_cuda_ipc_rkey_unpack_bound_lifetime(const uct_rkey_unpack_params_t *params)
+{
+    return !!(UCS_PARAM_VALUE(UCT_RKEY_UNPACK_FIELD, params, flags, FLAGS, 0) &
+              UCT_RKEY_UNPACK_FLAG_BOUND_LIFETIME);
+}
+
 UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack,
                  (component, rkey_buffer, params, rkey_p, handle_p),
                  uct_component_t *component, const void *rkey_buffer,
@@ -562,6 +575,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;
+    int bound_lifetime;
     CUdevice cuda_device;
     CUdevice avail_cuda_device;
     ucs_status_t status;
@@ -570,6 +584,8 @@ 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);
 
+    bound_lifetime = uct_cuda_ipc_rkey_unpack_bound_lifetime(params);
+
     unpacked = ucs_malloc(sizeof(*unpacked), "uct_cuda_ipc_unpacked_rkey_t");
     if (NULL == unpacked) {
         ucs_error("failed to allocate memory for uct_cuda_ipc_unpacked_rkey_t");
@@ -605,14 +621,17 @@ UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_rkey_unpack,
     /* Check if peer is accessible, and if that check required mapping a memory
      * handle, save it in the rkey handle */
     status = uct_cuda_ipc_is_peer_accessible(com, unpacked, rkey_handle,
-                                             avail_cuda_device);
+                                             avail_cuda_device, bound_lifetime);
     if (cuda_device != avail_cuda_device) {
         uct_cuda_ctx_primary_pop_and_release(avail_cuda_device);
     }
     if (status != UCS_OK) {
         goto err_free_rkey_handle;
     }
 
+    rkey_handle->cache_enabled =
+            bound_lifetime ? 0 : uct_cuda_ipc_component.enable_remote_cache;
+
     *handle_p = rkey_handle;
     *rkey_p   = (uintptr_t) unpacked;
     return UCS_OK;
@@ -644,7 +663,7 @@ static void uct_cuda_ipc_rkey_release_unmap_memhandle(uct_rkey_t uct_rkey,
 
     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->cache_enabled);
 }
 
 static ucs_status_t uct_cuda_ipc_rkey_release(uct_component_t *component,
@@ -767,17 +786,23 @@ static void uct_cuda_ipc_md_close(uct_md_h md)
 
 static ucs_status_t
 uct_cuda_ipc_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
+                              void *rkey_release_handle,
                               uct_device_mem_elem_t *mem_elem,
                               void **release_handle_p)
 {
     uct_cuda_ipc_unpacked_rkey_t *key = (uct_cuda_ipc_unpacked_rkey_t*)rkey;
     uct_cuda_ipc_md_device_mem_element_t *cuda_ipc_md_mem_element =
             (uct_cuda_ipc_md_device_mem_element_t*)mem_elem;
+    uct_cuda_ipc_rkey_handle_t *cuda_ipc_rkey_handle = rkey_release_handle;
     uct_cuda_ipc_mem_elem_release_handle_t *release_handle;
     ucs_status_t status;
     CUdevice cuda_device;
     void *mapped_addr;
 
+    if ((rkey == UCT_INVALID_RKEY) || (rkey_release_handle == NULL)) {
+        return UCS_ERR_INVALID_PARAM;
+    }
+
     if (UCT_CUDADRV_FUNC_LOG_DEBUG(cuCtxGetDevice(&cuda_device)) != UCS_OK) {
         return UCS_ERR_UNREACHABLE;
     }
@@ -798,12 +823,13 @@ uct_cuda_ipc_md_mem_elem_pack(uct_md_h md, uct_mem_h memh, uct_rkey_t rkey,
     cuda_ipc_md_mem_element->mapped_offset =
             UCS_PTR_BYTE_DIFF(key->super.super.d_bptr, mapped_addr);
 
-    release_handle->pid         = key->super.super.pid;
-    release_handle->pid_ns      = key->super.pid_ns;
-    release_handle->d_bptr      = key->super.super.d_bptr;
-    release_handle->mapped_addr = mapped_addr;
-    release_handle->cu_dev      = cuda_device;
-    *release_handle_p           = release_handle;
+    release_handle->pid           = key->super.super.pid;
+    release_handle->pid_ns        = key->super.pid_ns;
+    release_handle->d_bptr        = key->super.super.d_bptr;
+    release_handle->mapped_addr   = mapped_addr;
+    release_handle->cu_dev        = cuda_device;
+    release_handle->cache_enabled = cuda_ipc_rkey_handle->cache_enabled;
+    *release_handle_p             = release_handle;
 
     return UCS_OK;
 }
@@ -814,7 +840,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);
+                                 handle->cache_enabled);
     ucs_free(handle);
 }
 

src/uct/ib/mlx5/dv/ib_mlx5dv_md.c

@@ -3202,7 +3202,7 @@ uct_ib_mlx5_devx_md_open(struct ibv_device *ibv_device,
 
 static ucs_status_t
 uct_ib_md_mlx5_devx_md_mem_elem_pack(uct_md_h md, uct_mem_h memh,
-                                     uct_rkey_t rkey,
+                                     uct_rkey_t rkey, void *rkey_release_handle,
                                      uct_device_mem_elem_t *mem_elem_p,
                                      void **pack_resources_p)
 {

src/uct/rocm/ipc/rocm_ipc_md.c

@@ -126,6 +126,7 @@ uct_rocm_ipc_mem_dereg(uct_md_h md,
 
 static ucs_status_t
 uct_rocm_ipc_md_mem_elem_pack(uct_md_h md_h, uct_mem_h memh, uct_rkey_t rkey,
+                              void *rkey_handle,
                               uct_device_mem_elem_t *mem_elem_p,
                               void **release_handle_p)
 {

test/gtest/common/test_obj_size.cc

@@ -50,7 +50,7 @@ UCS_TEST_F(test_obj_size, size) {
     UCS_TEST_SKIP_R("Assert enabled");
 #else
     EXPECTED_SIZE(ucp_ep_t, 64);
-    EXPECTED_SIZE(ucp_ep_ext_t, 216);
+    EXPECTED_SIZE(ucp_ep_ext_t, 224);
 #if ENABLE_PARAMS_CHECK
     EXPECTED_SIZE(ucp_rkey_t, 24 + sizeof(ucp_ep_h));
 #else