Use staged communicator for connection establishment

Mellanox · Oct 26, 2023 · 4ccb98a · 4ccb98a
1 parent 7845726
commit 4ccb98a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/ucx_plugin.c b/src/ucx_plugin.c
@@ -403,7 +403,7 @@ static void ucx_request_init(ucx_comm_t *comm) {
 ncclResult_t nccl_ucx_connect(int dev, void *handle, void **send_comm, ncclNetDeviceHandle_t** sendDevComm) {
   ucx_listen_handle_t  *recv_handle = (ucx_listen_handle_t*)handle;
   struct ncclUCXCommStage* stage = &recv_handle->stage;
-  ucx_comm_t           *comm;
+  ucx_comm_t           *comm = stage->comm;
   ucp_address_t        *my_addr;
   size_t               local_addr_len;
   enum ncclSocketState conState;

diff --git a/src/ucx_rma_plugin.c b/src/ucx_rma_plugin.c
@@ -440,7 +440,7 @@ ncclResult_t nccl_ucx_rma_connect(int dev, void *handle, void **send_comm, ncclN
 {
   ucx_rma_listen_handle_t  *recv_handle = (ucx_rma_listen_handle_t*)handle;
   struct ncclUCXCommStage* stage = &recv_handle->stage;
-  nccl_ucx_rma_send_comm_t *comm;
+  nccl_ucx_rma_send_comm_t *comm = stage->comm;
   ucp_mem_map_params_t     mmap_params;
   size_t                   rkey_buf_size;
   void                     *rkey_buf;