From 1f76fc6e3f6f95e823e350330e575e573f4bb3ee Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Mon, 25 Nov 2024 16:02:03 -0800 Subject: [PATCH] [router] Rust e2e test (#2184) --- .github/workflows/pr-test-rust.yml | 25 +++++++- rust/py_test/run_suite.py | 19 ++++++ rust/py_test/test_launch_router.py | 66 ++++++++++++++++++++ rust/py_test/test_launch_server.py | 99 ++++++++++++++++++++++++++++++ rust/src/tree.rs | 2 +- 5 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 rust/py_test/run_suite.py create mode 100644 rust/py_test/test_launch_router.py create mode 100644 rust/py_test/test_launch_server.py diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 5dcf706a5..92bd986a0 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -40,8 +40,31 @@ jobs: cd rust/ cargo test + e2e-rust: + if: github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request' + runs-on: 1-gpu-runner + steps: + - name: Checkout code + uses: actions/checkout@v3 + + - name: Install rust dependencies + run: | + bash scripts/ci_install_rust.sh + + - name: Build python binding + run: | + source "$HOME/.cargo/env" + cd rust + pip install setuptools-rust wheel build + python3 -m build + pip install dist/*.whl + - name: Run e2e test + run: | + cd rust/py_test + python3 run_suite.py + finish: - needs: [unit-test-rust] + needs: [unit-test-rust, e2e-rust] runs-on: ubuntu-latest steps: - name: Finish diff --git a/rust/py_test/run_suite.py b/rust/py_test/run_suite.py new file mode 100644 index 000000000..0f012b751 --- /dev/null +++ b/rust/py_test/run_suite.py @@ -0,0 +1,19 @@ +import argparse +import glob + +from sglang.test.test_utils import run_unittest_files + +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.", + ) + args = arg_parser.parse_args() + + files = glob.glob("**/test_*.py", recursive=True) + + exit_code = run_unittest_files(files, args.timeout_per_file) + exit(exit_code) diff --git a/rust/py_test/test_launch_router.py b/rust/py_test/test_launch_router.py new file mode 100644 index 000000000..787c091bf --- /dev/null +++ b/rust/py_test/test_launch_router.py @@ -0,0 +1,66 @@ +import multiprocessing +import time +import unittest +from types import SimpleNamespace + + +def terminate_process(process: multiprocessing.Process, timeout: float = 1.0) -> None: + """Terminate a process gracefully, with forced kill as fallback. + + Args: + process: The process to terminate + timeout: Seconds to wait for graceful termination before forcing kill + """ + if not process.is_alive(): + return + + process.terminate() + process.join(timeout=timeout) + if process.is_alive(): + process.kill() # Force kill if terminate didn't work + process.join() + + +class TestLaunchRouter(unittest.TestCase): + def test_launch_router_no_exception(self): + + # Create SimpleNamespace with default arguments + args = SimpleNamespace( + worker_urls=["http://localhost:8000"], + host="127.0.0.1", + port=30000, + policy="cache_aware", + cache_threshold=0.5, + balance_abs_threshold=32, + balance_rel_threshold=1.0001, + eviction_interval=60, + max_tree_size=2**24, + verbose=False, + ) + + def run_router(): + try: + from sglang_router.launch_router import launch_router + + router = launch_router(args) + if router is None: + return 1 + return 0 + except Exception as e: + print(e) + return 1 + + # Start router in separate process + process = multiprocessing.Process(target=run_router) + try: + process.start() + # Wait 3 seconds + time.sleep(3) + # Process is still running means router started successfully + self.assertTrue(process.is_alive()) + finally: + terminate_process(process) + + +if __name__ == "__main__": + unittest.main() diff --git a/rust/py_test/test_launch_server.py b/rust/py_test/test_launch_server.py new file mode 100644 index 000000000..7fdaea6b1 --- /dev/null +++ b/rust/py_test/test_launch_server.py @@ -0,0 +1,99 @@ +import subprocess +import time +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import kill_child_process +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, +) + + +def popen_launch_router( + model: str, + base_url: str, + dp_size: int, + timeout: float, +): + """ + Launch the router server process. + + Args: + model: Model path/name + base_url: Server base URL + dp_size: Data parallel size + timeout: Server launch timeout + """ + _, host, port = base_url.split(":") + host = host[2:] + + command = [ + "python3", + "-m", + "sglang_router.launch_server", + "--model-path", + model, + "--host", + host, + "--port", + port, + "--dp", + str(dp_size), # Convert dp_size to string + ] + + # Use current environment + env = None + + process = subprocess.Popen(command, stdout=None, stderr=None, env=env) + + start_time = time.time() + with requests.Session() as session: + while time.time() - start_time < timeout: + try: + response = session.get(f"{base_url}/health") + if response.status_code == 200: + return process + except requests.RequestException: + pass + time.sleep(10) + + raise TimeoutError("Server failed to start within the timeout period.") + + +class TestEvalAccuracyMini(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_router( + cls.model, + cls.base_url, + dp_size=1, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + + @classmethod + def tearDownClass(cls): + kill_child_process(cls.process.pid, include_self=True) + + def test_mmlu(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmlu", + num_examples=64, + num_threads=32, + temperature=0.1, + ) + + metrics = run_eval(args) + self.assertGreaterEqual(metrics["score"], 0.65) + + +if __name__ == "__main__": + unittest.main() diff --git a/rust/src/tree.rs b/rust/src/tree.rs index 022c897c9..fcec5f578 100644 --- a/rust/src/tree.rs +++ b/rust/src/tree.rs @@ -656,7 +656,7 @@ mod tests { assert_eq!( tree.get_smallest_tenant(), "tenant2", - "Expected tenant2 to be smallest with 3 characters" + "Expected tenant2 to be smallest with 3 characters." ); // Insert overlapping data for tenant3 and tenant4 to test equal counts