Add argument to enable/disable endpoint error handling in benchmarks (#…

…1007)

Some transports (e.g., `sm`) do not support endpoint error handling. To be able to benchmark them, add a new argument to benchmarks to enable or disable endpoint error handling.

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)

URL: #1007

ucp/benchmarks/backends/ucp_async.py

@@ -81,7 +81,11 @@ async def server_handler(ep):
             await ep.close()
             lf.close()
 
-        lf = ucp.create_listener(server_handler, port=self.args.port)
+        lf = ucp.create_listener(
+            server_handler,
+            port=self.args.port,
+            endpoint_error_handling=self.args.error_handling,
+        )
         self.queue.put(lf.port)
 
         while not lf.closed():
@@ -114,7 +118,11 @@ async def run(self):
 
         register_am_allocators(self.args)
 
-        ep = await ucp.create_endpoint(self.server_address, self.port)
+        ep = await ucp.create_endpoint(
+            self.server_address,
+            self.port,
+            endpoint_error_handling=self.args.error_handling,
+        )
 
         if self.args.enable_am:
             msg = xp.arange(self.args.n_bytes, dtype="u1")

ucp/benchmarks/backends/ucp_core.py

@@ -125,7 +125,7 @@ def _listener_handler(conn_request, msg):
             self.ep = ucx_api.UCXEndpoint.create_from_conn_request(
                 worker,
                 conn_request,
-                endpoint_error_handling=True,
+                endpoint_error_handling=self.args.error_handling,
             )
 
             # Wireup before starting to transfer data
@@ -229,7 +229,7 @@ def run(self):
             worker,
             self.server_address,
             self.port,
-            endpoint_error_handling=True,
+            endpoint_error_handling=self.args.error_handling,
         )
 
         send_msg = xp.arange(self.args.n_bytes, dtype="u1")

ucp/benchmarks/send_recv.py

@@ -324,6 +324,12 @@ def parse_args():
         help="Only applies to 'ucp-core' backend: number of maximum outstanding "
         "operations, see --delay-progress. (Default: 32)",
     )
+    parser.add_argument(
+        "--error-handling",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Enable endpoint error handling.",
+    )
 
     args = parser.parse_args()