Delay benchmark event loop startup

infinistore/benchmark.py

@@ -117,11 +117,6 @@ def start_loop(loop):
     loop.run_forever()
 
 
-loop = asyncio.new_event_loop()
-t = threading.Thread(target=start_loop, args=(loop,))
-t.start()
-
-
 def run(args):
     config = infinistore.ClientConfig(
         host_addr=args.server,
@@ -138,6 +133,12 @@ def run(args):
         config.connection_type = infinistore.TYPE_TCP
 
     conn = infinistore.InfinityConnection(config)
+    loop = None
+    t = None
+    if args.rdma:
+        loop = asyncio.new_event_loop()
+        t = threading.Thread(target=start_loop, args=(loop,))
+        t.start()
     try:
         conn.connect()
 
@@ -271,8 +272,9 @@ def run(args):
         assert torch.equal(src_tensor.cpu(), dst_tensor.cpu())
     finally:
         conn.close()
-        loop.call_soon_threadsafe(loop.stop)
-        t.join()
+        if loop is not None:
+            loop.call_soon_threadsafe(loop.stop)
+            t.join()
 
 
 if __name__ == "__main__":

tests/test_benchmark_cli.py

@@ -0,0 +1,28 @@
+import subprocess
+import sys
+from pathlib import Path
+
+
+def test_help_exits():
+    benchmark = Path(__file__).parents[1] / "infinistore" / "benchmark.py"
+    runner = """
+import runpy
+import sys
+import types
+
+sys.modules["infinistore"] = types.ModuleType("infinistore")
+sys.modules["torch"] = types.ModuleType("torch")
+benchmark = sys.argv[1]
+sys.argv = [benchmark, "--help"]
+runpy.run_path(benchmark, run_name="__main__")
+"""
+
+    result = subprocess.run(
+        [sys.executable, "-c", runner, str(benchmark)],
+        capture_output=True,
+        text=True,
+        timeout=5,
+    )
+
+    assert result.returncode == 0
+    assert "--rdma" in result.stdout