Enhance retract test (page cases, long output cases) (#12781)

python/sglang/test/test_utils.py

@@ -1685,7 +1685,7 @@ class CustomTestCase(unittest.TestCase):
         )
 
     def setUp(self):
-        print(f"[Test Method] {self._testMethodName}", flush=True)
+        print(f"[CI Test Method] {self.__class__.__name__}.{self._testMethodName}")
 
 
 def dump_bench_raw_result(

python/sglang/utils.py

@@ -26,6 +26,8 @@ from IPython.display import HTML, display
 from pydantic import BaseModel
 from tqdm import tqdm
 
+from sglang.srt.environ import envs
+
 logger = logging.getLogger(__name__)
 
 
@@ -357,9 +359,7 @@ def download_and_cache_file(url: str, filename: Optional[str] = None):
 
 
 def is_in_ci() -> bool:
-    import os
+    return envs.SGLANG_IS_IN_CI.get()
-
-    return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("true", "1")
 
 
 def print_highlight(html_content: str):

test/srt/test_retract_decode.py

@@ -2,6 +2,8 @@ import time
 import unittest
 from types import SimpleNamespace
 
+import requests
+
 from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
 from sglang.test.run_eval import run_eval
@@ -12,51 +14,31 @@ from sglang.test.test_utils import (
     CustomTestCase,
     popen_launch_server,
 )
+from sglang.utils import is_in_ci
 
 
 class TestRetractDecode(CustomTestCase):
+    """python -m unittest test_retract_decode.TestRetractDecode"""
+
+    other_args = []
+
     @classmethod
     def setUpClass(cls):
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.base_url = DEFAULT_URL_FOR_TEST
-        with envs.SGLANG_TEST_RETRACT.override(True):
+        launch_args = ["--chunked-prefill-size", "128"] + cls.other_args
-            cls.process = popen_launch_server(
-                cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
-            )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_mmlu(self):
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="mmlu",
-            num_examples=64,
-            num_threads=32,
-        )
-
-        metrics = run_eval(args)
-        self.assertGreaterEqual(metrics["score"], 0.65)
-        time.sleep(1)  # wait for mem check
-
-        assert self.process.poll() is None, "Server crashed during test"
-
-
-class TestRetractDecodeChunkCache(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
-        cls.base_url = DEFAULT_URL_FOR_TEST
         with envs.SGLANG_TEST_RETRACT.override(True):
             cls.process = popen_launch_server(
                 cls.model,
                 cls.base_url,
                 timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-                other_args=["--disable-radix-cache", "--chunked-prefill-size", 128],
+                other_args=launch_args,
             )
 
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
     def test_mmlu(self):
         args = SimpleNamespace(
             base_url=self.base_url,
@@ -72,9 +54,63 @@ class TestRetractDecodeChunkCache(CustomTestCase):
 
         assert self.process.poll() is None, "Server crashed during test"
 
+
+class TestRetractDecodePaged(TestRetractDecode):
+    """python -m unittest test_retract_decode.TestRetractDecodePaged"""
+
+    other_args = ["--page-size", "16"]
+
+
+class TestRetractDecodeChunkCache(TestRetractDecode):
+    """python -m unittest test_retract_decode.TestRetractDecodeChunkCache"""
+
+    other_args = ["--disable-radix-cache"]
+
+
+class TestRetractDecodeChunkCachePaged(TestRetractDecode):
+    """python -m unittest test_retract_decode.TestRetractDecodeChunkCachePaged"""
+
+    other_args = ["--disable-radix-cache", "--page-size", "16"]
+
+
+@unittest.skipIf(is_in_ci(), "Skipped in CI due to long runtime")
+class TestRetractDecodeLongOutput(CustomTestCase):
+    """python -m unittest test_retract_decode.TestRetractDecodeLongOutput"""
+
+    other_args = []
+
     @classmethod
-    def tearDownClass(cls):
+    def setUpClass(cls):
-        kill_process_tree(cls.process.pid)
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        launch_args = [
+            "--chunked-prefill-size",
+            "128",
+            "--page-size",
+            "16",
+        ] + cls.other_args
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=launch_args,
+        )
+
+    def test_long_output_retract(self):
+        data = {
+            "input_ids": [[233 + i] * 1234 for i in range(256)],
+            "sampling_params": {"max_new_tokens": 90000, "ignore_eos": True},
+        }
+        res = requests.post(f"{self.base_url}/generate", json=data)
+        assert res.status_code == 200, f"Request failed: {res.status_code}"
+        assert self.process.poll() is None, "Server crashed during test"
+
+
+@unittest.skipIf(is_in_ci(), "Skipped in CI due to long runtime")
+class TestRetractDecodeLongOutputChunkCache(TestRetractDecodeLongOutput):
+    """python -m unittest test_retract_decode.TestRetractDecodeLongOutputChunkCache"""
+
+    other_args = ["--disable-radix-cache"]
 
 
 if __name__ == "__main__":