Add missing endpoint error handler to listener's endpoints

cpp/src/endpoint.cpp

@@ -83,14 +83,12 @@ std::shared_ptr<Endpoint> createEndpointFromHostname(std::shared_ptr<Worker> wor
   if (worker == nullptr || worker->getHandle() == nullptr)
     throw ucxx::Error("Worker not initialized");
 
-  ucp_ep_params_t params = {.field_mask = UCP_EP_PARAM_FIELD_FLAGS | UCP_EP_PARAM_FIELD_SOCK_ADDR |
-                                          UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE |
-                                          UCP_EP_PARAM_FIELD_ERR_HANDLER,
-                            .flags = UCP_EP_PARAMS_FLAGS_CLIENT_SERVER};
   auto info              = ucxx::utils::get_addrinfo(ipAddress.c_str(), port);
-
-  params.sockaddr.addrlen = info->ai_addrlen;
-  params.sockaddr.addr    = info->ai_addr;
+  ucp_ep_params_t params = {.field_mask = UCP_EP_PARAM_FIELD_FLAGS | UCP_EP_PARAM_FIELD_SOCK_ADDR,
+                            .flags      = UCP_EP_PARAMS_FLAGS_CLIENT_SERVER,
+                            .sockaddr   = {.addr = info->ai_addr, .addrlen = info->ai_addrlen}};
+  if (endpointErrorHandling)
+    params.field_mask |= UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE | UCP_EP_PARAM_FIELD_ERR_HANDLER;
 
   return std::shared_ptr<Endpoint>(new Endpoint(worker, &params, endpointErrorHandling));
 }
@@ -103,10 +101,11 @@ std::shared_ptr<Endpoint> createEndpointFromConnRequest(std::shared_ptr<Listener
     throw ucxx::Error("Worker not initialized");
 
   ucp_ep_params_t params = {
-    .field_mask = UCP_EP_PARAM_FIELD_FLAGS | UCP_EP_PARAM_FIELD_CONN_REQUEST |
-                  UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE | UCP_EP_PARAM_FIELD_ERR_HANDLER,
+    .field_mask   = UCP_EP_PARAM_FIELD_FLAGS | UCP_EP_PARAM_FIELD_CONN_REQUEST,
     .flags        = UCP_EP_PARAMS_FLAGS_NO_LOOPBACK,
     .conn_request = connRequest};
+  if (endpointErrorHandling)
+    params.field_mask |= UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE | UCP_EP_PARAM_FIELD_ERR_HANDLER;
 
   return std::shared_ptr<Endpoint>(new Endpoint(listener, &params, endpointErrorHandling));
 }
@@ -120,10 +119,10 @@ std::shared_ptr<Endpoint> createEndpointFromWorkerAddress(std::shared_ptr<Worker
   if (address == nullptr || address->getHandle() == nullptr || address->getLength() == 0)
     throw ucxx::Error("Address not initialized");
 
-  ucp_ep_params_t params = {.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS |
-                                          UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE |
-                                          UCP_EP_PARAM_FIELD_ERR_HANDLER,
-                            .address = address->getHandle()};
+  ucp_ep_params_t params = {.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS,
+                            .address    = address->getHandle()};
+  if (endpointErrorHandling)
+    params.field_mask |= UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE | UCP_EP_PARAM_FIELD_ERR_HANDLER;
 
   return std::shared_ptr<Endpoint>(new Endpoint(worker, &params, endpointErrorHandling));
 }
@@ -142,7 +141,7 @@ void Endpoint::close()
   ucxx_debug("Endpoint %p canceled %lu requests", _handle, canceled);
 
   // Close the endpoint
-  unsigned closeMode = UCP_EP_CLOSE_MODE_FORCE;
+  unsigned closeMode = UCP_EP_CLOSE_MODE_FLUSH;
   if (_endpointErrorHandling && _callbackData->status != UCS_OK) {
     // We force close endpoint if endpoint error handling is enabled and
     // the endpoint status is not UCS_OK

python/ucxx/_lib/libucxx.pyx

@@ -1210,7 +1210,7 @@ cdef void _listener_callback(ucp_conn_request_h conn_request, void *args) with g
         cb_data['cb_func'](
             (
                 cb_data['listener'].create_endpoint_from_conn_request(
-                    int(<uintptr_t>conn_request), True
+                    int(<uintptr_t>conn_request), cb_data['endpoint_error_handling']
                 ) if 'listener' in cb_data else
                 int(<uintptr_t>conn_request)
             ),
@@ -1242,6 +1242,7 @@ cdef class UCXListener():
             cls,
             UCXWorker worker,
             uint16_t port,
+            bint endpoint_error_handling,
             cb_func,
             tuple cb_args=None,
             dict cb_kwargs=None,
@@ -1260,6 +1261,7 @@ cdef class UCXListener():
             "cb_func": cb_func,
             "cb_args": cb_args,
             "cb_kwargs": cb_kwargs,
+            "endpoint_error_handling": endpoint_error_handling,
         }
 
         with nogil:

python/ucxx/_lib/tests/test_cancel.py

@@ -27,7 +27,7 @@ def _listener_handler(conn_request):
         ep[0] = listener.create_endpoint_from_conn_request(conn_request, True)
 
     listener = ucx_api.UCXListener.create(
-        worker=worker, port=0, cb_func=_listener_handler
+        worker=worker, port=0, endpoint_error_handling=True, cb_func=_listener_handler
     )
     queue.put(listener.port)
 

python/ucxx/_lib/tests/test_endpoint.py

@@ -42,7 +42,7 @@ def _listener_handler(conn_request):
         listener_finished[0] = True
 
     listener = ucx_api.UCXListener.create(
-        worker=worker, port=0, cb_func=_listener_handler
+        worker=worker, port=0, endpoint_error_handling=True, cb_func=_listener_handler
     )
     queue.put(listener.port)
 

python/ucxx/_lib/tests/test_listener.py

@@ -12,7 +12,7 @@ def _listener_handler(conn_request):
         pass
 
     listener = ucx_api.UCXListener.create(
-        worker=worker, port=0, cb_func=_listener_handler
+        worker=worker, port=0, endpoint_error_handling=True, cb_func=_listener_handler
     )
 
     assert isinstance(listener.ip, str) and listener.ip

python/ucxx/_lib/tests/test_probe.py

@@ -36,7 +36,7 @@ def _listener_handler(conn_request):
         )
 
     listener = ucx_api.UCXListener.create(
-        worker=worker, port=0, cb_func=_listener_handler
+        worker=worker, port=0, endpoint_error_handling=True, cb_func=_listener_handler
     )
     queue.put(listener.port)
 

python/ucxx/_lib/tests/test_server_client.py

@@ -66,7 +66,7 @@ def _listener_handler(conn_request):
         ep[0] = listener.create_endpoint_from_conn_request(conn_request, True)
 
     listener = ucx_api.UCXListener.create(
-        worker=worker, port=0, cb_func=_listener_handler
+        worker=worker, port=0, endpoint_error_handling=True, cb_func=_listener_handler
     )
     put_queue.put(listener.port)
 

python/ucxx/_lib_async/application_context.py

@@ -219,6 +219,7 @@ def create_listener(
             ucx_api.UCXListener.create(
                 worker=self.worker,
                 port=port,
+                endpoint_error_handling=endpoint_error_handling,
                 cb_func=_listener_handler,
                 cb_args=(
                     loop,

python/ucxx/benchmarks/backends/ucxx_core.py

@@ -145,7 +145,10 @@ def _listener_handler(conn_request):
             ep = listener.create_endpoint_from_conn_request(conn_request, True)
 
         listener = ucx_api.UCXListener.create(
-            worker=worker, port=self.args.port or 0, cb_func=_listener_handler
+            worker=worker,
+            port=self.args.port or 0,
+            endpoint_error_handling=True,
+            cb_func=_listener_handler,
         )
         self.queue.put(listener.port)