rdma: move idpool init to endpoint refcnt+ stage

The idpool is freed when refcnt goes to zero, so it should be initialized again when refcnt increases. Signed-off-by: Eric Raut <[email protected]>
aws · Feb 21, 2024 · 213172a · 213172a
1 parent f931085
commit 213172a
Showing 1 changed file with 16 additions and 16 deletions.
diff --git a/src/nccl_ofi_rdma.c b/src/nccl_ofi_rdma.c
@@ -5570,22 +5570,6 @@ static int get_ep(nccl_net_ofi_device_t *base_dev,
 		/* Initialize number of rail */
 		ep->num_rails = num_rails;
 
-		/* Initialize endpoint ID pool */
-		ep->comm_idpool = malloc(sizeof(nccl_ofi_idpool_t));
-		if (OFI_UNLIKELY(ep->comm_idpool == NULL)) {
-			ret = ncclSystemError;
-			NCCL_OFI_WARN("Unable to allocate rdma endpoint ID pool");
-			goto unlock;
-		}
-
-		ret = nccl_ofi_idpool_init(ep->comm_idpool, device->num_comm_ids);
-		if (OFI_UNLIKELY(ret != 0)) {
-			ret = ncclSystemError;
-			free(ep->comm_idpool);
-			ep->comm_idpool = NULL;
-			goto unlock;
-		}
-
 		/* Initialize reference count */
 		ep->ref_cnt = 0;
 
@@ -5625,6 +5609,22 @@ static int get_ep(nccl_net_ofi_device_t *base_dev,
 			goto unlock;
 		}
 
+		/* Initialize endpoint ID pool */
+		ep->comm_idpool = malloc(sizeof(nccl_ofi_idpool_t));
+		if (OFI_UNLIKELY(ep->comm_idpool == NULL)) {
+			ret = ncclSystemError;
+			NCCL_OFI_WARN("Unable to allocate rdma endpoint ID pool");
+			goto unlock;
+		}
+
+		ret = nccl_ofi_idpool_init(ep->comm_idpool, device->num_comm_ids);
+		if (OFI_UNLIKELY(ret != 0)) {
+			ret = ncclSystemError;
+			free(ep->comm_idpool);
+			ep->comm_idpool = NULL;
+			goto unlock;
+		}
+
 		ret = init_rail_ofi_resources(device, ep);
 		if (ret != 0) {
 			goto unlock;