[1/N] CI refactor: introduce CI register. (#13345)

test/lang/run_suite.py

@@ -1,36 +0,0 @@
-import argparse
-import glob
-
-from sglang.test.test_utils import TestFile, run_unittest_files
-
-suites = {
-    "per-commit": [
-        TestFile("test_srt_backend.py"),
-    ],
-}
-
-
-if __name__ == "__main__":
-    arg_parser = argparse.ArgumentParser()
-    arg_parser.add_argument(
-        "--timeout-per-file",
-        type=int,
-        default=1000,
-        help="The time limit for running one file in seconds.",
-    )
-    arg_parser.add_argument(
-        "--suite",
-        type=str,
-        default=list(suites.keys())[0],
-        choices=list(suites.keys()) + ["all"],
-        help="The suite to run",
-    )
-    args = arg_parser.parse_args()
-
-    if args.suite == "all":
-        files = glob.glob("**/test_*.py", recursive=True)
-    else:
-        files = suites[args.suite]
-
-    exit_code = run_unittest_files(files, args.timeout_per_file)
-    exit(exit_code)

test/per_commit/test_srt_backend.py

@@ -1,12 +1,7 @@
-"""
-Usage:
-python3 -m unittest test_srt_backend.TestSRTBackend.test_gen_min_new_tokens
-python3 -m unittest test_srt_backend.TestSRTBackend.test_hellaswag_select
-"""
-
 import unittest
 
 import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_programs import (
     test_decode_int,
     test_decode_json_regex,
@@ -24,6 +19,8 @@ from sglang.test.test_programs import (
 )
 from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
 
+register_cuda_ci(estimation_time=80, ci_stage="stage-a-test-1")
+
 
 class TestSRTBackend(CustomTestCase):
     backend = None

test/run_suite.py

@@ -0,0 +1,39 @@
+import glob
+from typing import List
+
+from sglang.test.ci.ci_register import CIRegistry, HWBackend, collect_tests
+from sglang.test.ci.ci_utils import TestFile, run_unittest_files
+
+LABEL_MAPPING = {HWBackend.CUDA: ["stage-a-test-1"]}
+
+
+def _filter_tests(
+    ci_tests: List[CIRegistry], hw: HWBackend, suite: str
+) -> List[CIRegistry]:
+    ci_tests = [t for t in ci_tests if t.backend == hw]
+    ret = []
+    for t in ci_tests:
+        assert t.stage in LABEL_MAPPING[hw], f"Unknown stage {t.stage} for backend {hw}"
+        if t.stage == suite:
+            ret.append(t)
+    return ret
+
+
+def run_per_commit(hw: HWBackend, suite: str):
+    files = glob.glob("per_commit/**/*.py", recursive=True)
+    ci_tests = _filter_tests(collect_tests(files), hw, suite)
+    test_files = [TestFile(t.filename, t.estimation_time) for t in ci_tests]
+
+    run_unittest_files(
+        test_files,
+        timeout_per_file=1200,
+        continue_on_error=False,
+    )
+
+
+def main():
+    run_per_commit(HWBackend.CUDA, "stage-a-test-1")
+
+
+if __name__ == "__main__":
+    main()

test/srt/run_suite.py

@@ -1,16 +1,8 @@
 import argparse
 import glob
-from dataclasses import dataclass
 from pathlib import Path
 
-from sglang.test.test_utils import run_unittest_files
-
-
-@dataclass
-class TestFile:
-    name: str
-    estimated_time: float = 60
-
+from sglang.test.ci.ci_utils import TestFile, run_unittest_files
 
 # NOTE: please sort the test cases alphabetically by the test file name
 suites = {
@@ -619,7 +611,7 @@ def _sanity_check_suites(suites):
     )
 
 
-if __name__ == "__main__":
+def main():
     arg_parser = argparse.ArgumentParser()
     arg_parser.add_argument(
         "--timeout-per-file",
@@ -667,3 +659,7 @@ if __name__ == "__main__":
 
     exit_code = run_unittest_files(files, args.timeout_per_file, args.continue_on_error)
     exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()