From 77a6b2a6b06f7a897ccb57c40caf153e3224ea0a Mon Sep 17 00:00:00 2001
From: ToLiveAndLove <59090766+ToLiveAndLove@users.noreply.github.com>
Date: Fri, 18 Sep 2026 10:51:18 +0800
Subject: [PATCH] Delay benchmark event loop startup

---
 infinistore/benchmark.py    | 16 +++++++++-------
 tests/test_benchmark_cli.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 7 deletions(-)
 create mode 100644 tests/test_benchmark_cli.py

diff --git a/infinistore/benchmark.py b/infinistore/benchmark.py
index 780375a..17e27f5 100644
--- a/infinistore/benchmark.py
+++ b/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__":
diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py
new file mode 100644
index 0000000..353d946
--- /dev/null
+++ b/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