Fix CI

Author: Muqi1029

ai_infra_bench/sgl/cmp_bench.py

@@ -11,10 +11,10 @@
     check_str_list_str,
     check_values_in_features_metrics,
 )
+from ai_infra_bench.constants import FULL_DATA_JSON_PATH
 from ai_infra_bench.modes.cmp import cmp_export_table
 from ai_infra_bench.modes.gen import gen_export_csv, gen_run
 from ai_infra_bench.utils import (
-    FULL_DATA_JSON_PATH,
     kill_process_tree,
     maybe_create_labels,
     maybe_warmup,

ai_infra_bench/sgl/gen_bench.py

@@ -10,9 +10,9 @@
     check_server_labels,
     check_str_list_str,
 )
+from ai_infra_bench.constants import FULL_DATA_JSON_PATH
 from ai_infra_bench.modes.gen import gen_export_csv, gen_export_table, gen_plot, gen_run
 from ai_infra_bench.utils import (
-    FULL_DATA_JSON_PATH,
     kill_process_tree,
     maybe_create_labels,
     maybe_warmup,

ai_infra_bench/sgl/slo_bench.py

@@ -14,9 +14,9 @@
     check_server_labels,
     check_str_list_str,
 )
+from ai_infra_bench.constants import FULL_DATA_JSON_PATH
 from ai_infra_bench.modes.slo import slo_run
 from ai_infra_bench.utils import (
-    FULL_DATA_JSON_PATH,
     add_request_rate,
     kill_process_tree,
     maybe_create_labels,

examples/client_gen.py

@@ -1,4 +1,5 @@
 import os
+from typing import List
 
 from ai_infra_bench import client_gen
 
@@ -39,8 +40,8 @@
 		--max-concurrency {request_rate}
 		--num-prompt {num_prompt}
 """
-rate_lists = [1, 2, 4, 8]
-client_cmds = [
+rate_lists: List[int] = [1, 2, 4, 8]
+client_cmds: List[str] = [
     *[
         client_template.format(
             base_url=base_url,

test/modes/test_cmp.py

@@ -6,14 +6,17 @@
 from utils import check_output_content, input_features, output_metrics
 
 from ai_infra_bench.client import client_cmp
+from ai_infra_bench.constants import CSV_NAME, TABLE_NAME, WARMUP_FILE
 from ai_infra_bench.sgl import cmp_bench
-from ai_infra_bench.utils import CSV_NAME, TABLE_NAME, WARMUP_FILE, ServerAccessInfo
+from ai_infra_bench.utils import ServerAccessInfo
 
 
 class TestClientCmp(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.server_access_info = ServerAccessInfo(base_url="http://127.0.0.1:8000")
+        cls.server_access_infos = ServerAccessInfo(
+            base_url="http://127.0.0.1:8000", label="Test"
+        )
         cls.client_cmds = """
             python -m sglang.bench_serving \
                     --backend sglang-oai
@@ -28,10 +31,10 @@ def setUpClass(cls):
                     --num-prompt 40
         """
 
-    def run_single_cmd(self, server_access_info, client_cmds, **kwargs):
+    def run_single_cmd(self, server_access_infos, client_cmds, **kwargs):
         with tempfile.TemporaryDirectory() as output_dir:
             client_cmp(
-                server_access_info=server_access_info,
+                server_access_infos=server_access_infos,
                 client_cmds=client_cmds,
                 input_features=input_features,
                 output_metrics=output_metrics,
@@ -44,106 +47,85 @@ def run_single_cmd(self, server_access_info, client_cmds, **kwargs):
 
     ################## BASIC RUN ##########################
     def test_single_run(self):
-        self.run_single_cmd(self.server_access_info, self.client_cmds)
+        self.run_single_cmd(self.server_access_infos, self.client_cmds)
 
     def test_multiple_run(self):
-        self.run_single_cmd([self.server_access_info] * 2, [self.client_cmds] * 3)
+        self.run_single_cmd([self.server_access_infos] * 2, [self.client_cmds] * 3)
 
-        self.run_single_cmd([self.server_access_info] * 2, self.client_cmds)
+        self.run_single_cmd([self.server_access_infos] * 2, self.client_cmds)
 
-        self.run_single_cmd(self.server_access_info, [self.client_cmds] * 2)
+        self.run_single_cmd(self.server_access_infos, [self.client_cmds] * 2)
 
     def test_n_run(self):
-        self.run_single_cmd(self.server_access_info, self.client_cmds, n=3)
+        self.run_single_cmd(self.server_access_infos, self.client_cmds, n=3)
 
-        self.run_single_cmd([self.server_access_info] * 2, [self.client_cmds] * 3, n=3)
+        self.run_single_cmd([self.server_access_infos] * 2, [self.client_cmds] * 3, n=3)
 
-        self.run_single_cmd([self.server_access_info] * 2, self.client_cmds, n=3)
+        self.run_single_cmd([self.server_access_infos] * 2, self.client_cmds, n=3)
 
-        self.run_single_cmd(self.server_access_info, [self.client_cmds] * 2, n=3)
+        self.run_single_cmd(self.server_access_infos, [self.client_cmds] * 2, n=3)
 
         self.run_single_cmd(
-            self.server_access_info, [self.client_cmds] * 2, n=3, only_last=True
+            self.server_access_infos, [self.client_cmds] * 2, n=3, only_last=True
         )
 
     ################## LABEL SETTING ##########################
     def test_client_labels(self):
         self.run_single_cmd(
-            self.server_access_info, self.client_cmds, client_labels=["client1"]
+            self.server_access_infos, self.client_cmds, client_labels=["client1"]
         )
 
         self.run_single_cmd(
-            self.server_access_info, self.client_cmds, client_labels="client1"
+            self.server_access_infos, self.client_cmds, client_labels="client1"
         )
 
         self.run_single_cmd(
-            self.server_access_info, [self.client_cmds] * 2, client_labels="client1"
+            self.server_access_infos, [self.client_cmds] * 2, client_labels="client1"
         )
 
         self.run_single_cmd(
-            self.server_access_info,
+            self.server_access_infos,
             [self.client_cmds] * 2,
             client_labels=["client1"] * 2,
         )
 
     @unittest.expectedFailure
     def test_failed_client_labels(self):
         self.run_single_cmd(
-            self.server_access_info,
+            self.server_access_infos,
             [self.client_cmds] * 2,
             client_labels=["client1"] * 3,
         )
 
-    def test_server_labels(self):
-        self.run_single_cmd(
-            self.server_access_info, self.client_cmds, server_labels="server_label"
-        )
-        self.run_single_cmd(
-            self.server_access_info, self.client_cmds, server_labels=["server_label"]
-        )
-        self.run_single_cmd(
-            [self.server_access_info] * 2,
-            self.client_cmds,
-            server_labels=["server_label"] * 2,
-        )
-
-    @unittest.expectedFailure
-    def test_failed_server_labels(self):
-        self.run_single_cmd(
-            [self.server_access_info] * 2,
-            self.client_cmds,
-            server_labels=["server_label"] * 3,
-        )
-
     ################## EXPECTED FAIL ###########################
     @unittest.expectedFailure
     def test_failed_host(self):
         self.run_single_cmd(
-            self.server_access_info, self.client_cmds + " --host 127.0.0.1"
+            self.server_access_infos, self.client_cmds + " --host 127.0.0.1"
         )
 
     @unittest.expectedFailure
     def test_failed_port(self):
-        self.run_single_cmd(self.server_access_info, self.client_cmds + " --port 8888")
+        self.run_single_cmd(self.server_access_infos, self.client_cmds + " --port 8888")
 
     @unittest.expectedFailure
     def test_failed_base_url(self):
         self.run_single_cmd(
-            self.server_access_info,
+            self.server_access_infos,
             self.client_cmds + " --base-url http://127.0.0.1:8888",
         )
 
     @unittest.expectedFailure
     def test_failed_output_file(self):
         self.run_single_cmd(
-            self.server_access_info, self.client_cmds + " --output-file output.jsonl"
+            self.server_access_infos, self.client_cmds + " --output-file output.jsonl"
         )
 
     ################## DISABLE FEATURE ##########################
     def test_disable_warmup(self):
         with tempfile.TemporaryDirectory() as output_dir:
             client_cmp(
-                server_access_info=self.server_access_info,
+                server_access_infos=self.server_access_infos,
                 client_cmds=self.client_cmds,
                 input_features=input_features,
                 output_metrics=output_metrics,
@@ -155,7 +137,7 @@ def test_disable_warmup(self):
     def test_disable_csv(self):
         with tempfile.TemporaryDirectory() as output_dir:
             client_cmp(
-                server_access_info=self.server_access_info,
+                server_access_infos=self.server_access_infos,
                 client_cmds=self.client_cmds,
                 input_features=input_features,
                 output_metrics=output_metrics,
@@ -167,7 +149,7 @@ def test_disable_csv(self):
     def test_disable_md_table(self):
         with tempfile.TemporaryDirectory() as output_dir:
             client_cmp(
-                server_access_info=self.server_access_info,
+                server_access_infos=self.server_access_infos,
                 client_cmds=self.client_cmds,
                 input_features=input_features,
                 output_metrics=output_metrics,
@@ -179,7 +161,7 @@ def test_disable_md_table(self):
     def test_disable_plot(self):
         with tempfile.TemporaryDirectory() as output_dir:
             client_cmp(
-                server_access_info=self.server_access_info,
+                server_access_infos=self.server_access_infos,
                 client_cmds=self.client_cmds,
                 input_features=input_features,
                 output_metrics=output_metrics,

test/modes/test_gen.py

@@ -6,8 +6,8 @@
 from utils import check_output_content, input_features, output_metrics
 
 from ai_infra_bench.client import client_gen
+from ai_infra_bench.constants import CSV_NAME, TABLE_NAME, WARMUP_FILE
 from ai_infra_bench.sgl import gen_bench
-from ai_infra_bench.utils import CSV_NAME, TABLE_NAME, WARMUP_FILE
 
 
 class TestClientGen(unittest.TestCase):

test/modes/test_slo.py

@@ -30,8 +30,8 @@ def check_slo(item: Dict) -> bool:
 
 
 from ai_infra_bench.client import client_slo
+from ai_infra_bench.constants import CSV_NAME, FULL_DATA_JSON_PATH, TABLE_NAME
 from ai_infra_bench.sgl import slo_bench
-from ai_infra_bench.utils import CSV_NAME, FULL_DATA_JSON_PATH, TABLE_NAME
 
 
 class TestSGLSlo(unittest.TestCase):

test/modes/utils.py

@@ -1,6 +1,11 @@
 import os
 
-from ai_infra_bench.utils import CSV_NAME, FULL_DATA_JSON_PATH, TABLE_NAME, WARMUP_FILE
+from ai_infra_bench.constants import (
+    CSV_NAME,
+    FULL_DATA_JSON_PATH,
+    TABLE_NAME,
+    WARMUP_FILE,
+)
 
 input_features = [
     "random_input_len",