UCP/DEVICE: Add retry loops in perftest and gtest

Author: michal-shalev

src/tools/perf/cuda/ucp_cuda_kernel.cu

@@ -152,8 +152,15 @@ private:
         params.num_elements = count;
         params.elements     = elems;
 
-        ucs_status_t status = ucp_device_mem_list_create(perf.ucp.ep, &params,
-                                                         &m_params.mem_list);
+        ucs_status_t status;
+        do {
+            status = ucp_device_mem_list_create(perf.ucp.ep, &params,
+                                                &m_params.mem_list);
+            if (status == UCS_ERR_NOT_CONNECTED) {
+                ucp_worker_progress(perf.ucp.worker);
+            }
+        } while (status == UCS_ERR_NOT_CONNECTED);
+
         if (status != UCS_OK) {
             throw std::runtime_error("Failed to create memory list");
         }

test/gtest/ucp/test_ucp_device.cc

@@ -144,9 +144,17 @@ test_ucp_device::mem_list::mem_list(entity &sender, entity &receiver,
     params.num_elements = count;
     params.elements     = elems.data();
 
-    // Create memory list
-    ASSERT_UCS_OK(
-            ucp_device_mem_list_create(sender.ep(), &params, &m_mem_list_h));
+    // Create memory list (with retry on connection)
+    ucs_status_t status;
+    do {
+        status = ucp_device_mem_list_create(sender.ep(), &params, &m_mem_list_h);
+        if (status != UCS_ERR_NOT_CONNECTED) {
+            break;
+        }
+        sender.progress();
+        receiver.progress();
+    } while (status == UCS_ERR_NOT_CONNECTED);
+    ASSERT_UCS_OK(status);
 }
 
 test_ucp_device::mem_list::~mem_list()