From 41bd744671d1c2c0c44505b9c3a26e131daf52bb Mon Sep 17 00:00:00 2001
From: Nils Mechtel <nils.mech@gmail.com>
Date: Sat, 19 Oct 2024 00:04:15 +0200
Subject: [PATCH] stress test for sam service

---
 test/stress_test.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 test/stress_test.py

diff --git a/test/stress_test.py b/test/stress_test.py
new file mode 100644
index 0000000..df92826
--- /dev/null
+++ b/test/stress_test.py
@@ -0,0 +1,36 @@
+from hypha_rpc import connect_to_server
+import numpy as np
+import asyncio
+import argparse
+
+
+SERVER_URL = "https://hypha.aicell.io"
+WORKSPACE_NAME = "bioimageio-colab"
+CLIENT_ID = "kubernetes"
+SERVICE_ID = "sam"
+SID = f"{WORKSPACE_NAME}/{CLIENT_ID}:{SERVICE_ID}"
+
+
+async def run_client(client_id: int, image: np.ndarray, point_coordinates: list, point_labels: list):
+    print(f"Client {client_id} started")
+    client = await connect_to_server({"server_url": SERVER_URL, "method_timeout": 10})
+    segment_svc = await client.get_service(SID)
+    await segment_svc.segment(model_name="vit_b", image=image, point_coordinates=point_coordinates, point_labels=point_labels)
+
+
+async def stress_test(num_clients: int):
+    image=np.random.rand(256, 256)
+    point_coordinates=[[128, 128]]
+    point_labels=[1]
+    tasks = [
+        run_client(client_id, image, point_coordinates, point_labels)
+        for client_id in range(num_clients)
+    ]
+    await asyncio.gather(*tasks)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clients", type=int, default=50)
+    args = parser.parse_args()
+
+    asyncio.run(stress_test(args.num_clients))
\ No newline at end of file