Add argument to enable/disable endpoint error handling in benchmarks

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.
rapidsai · Nov 13, 2023 · 006d014 · 006d014
1 parent 4464ca9
commit 006d014
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/python/ucxx/benchmarks/backends/ucxx_async.py b/python/ucxx/benchmarks/backends/ucxx_async.py
@@ -91,7 +91,11 @@ async def server_handler(ep):
             await ep.close()
             lf.close()
 
-        lf = ucxx.create_listener(server_handler, port=self.args.port)
+        lf = ucxx.create_listener(
+            server_handler,
+            port=self.args.port,
+            endpoint_error_handling=self.args.error_handling,
+        )
         self.queue.put(lf.port)
 
         while not lf.closed():
@@ -126,7 +130,11 @@ async def run(self):
 
         register_am_allocators(self.args)
 
-        ep = await ucxx.create_endpoint(self.server_address, self.port)
+        ep = await ucxx.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")

diff --git a/python/ucxx/benchmarks/backends/ucxx_core.py b/python/ucxx/benchmarks/backends/ucxx_core.py
@@ -142,7 +142,9 @@ def run(self):
 
         def _listener_handler(conn_request):
             global ep
-            ep = listener.create_endpoint_from_conn_request(conn_request, True)
+            ep = listener.create_endpoint_from_conn_request(
+                conn_request, endpoint_error_handling=self.args.error_handling
+            )
 
         listener = ucx_api.UCXListener.create(
             worker=worker, port=self.args.port or 0, cb_func=_listener_handler
@@ -236,7 +238,7 @@ def run(self):
             worker,
             self.server_address,
             self.port,
-            endpoint_error_handling=True,
+            endpoint_error_handling=self.args.error_handling,
         )
 
         # Wireup before starting to transfer data

diff --git a/python/ucxx/benchmarks/send_recv.py b/python/ucxx/benchmarks/send_recv.py
@@ -322,6 +322,12 @@ def parse_args():
         help="Only applies to 'ucxx-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()