Enhance the test case for chunked prefill (#1785)

test/srt/test_large_max_new_tokens.py

@@ -1,3 +1,7 @@
+"""
+python3 -m unittest test_large_max_new_tokens.TestLargeMaxNewTokens.test_chat_completion
+"""
+
 import os
 import unittest
 from concurrent.futures import ThreadPoolExecutor
@@ -20,6 +24,10 @@ class TestLargeMaxNewTokens(unittest.TestCase):
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.base_url = DEFAULT_URL_FOR_TEST
         cls.api_key = "sk-123456"
+
+        cls.stdout = open("stdout.txt", "w")
+        cls.stderr = open("stderr.txt", "w")
+
         cls.process = popen_launch_server(
             cls.model,
             cls.base_url,
@@ -27,7 +35,7 @@ class TestLargeMaxNewTokens(unittest.TestCase):
             api_key=cls.api_key,
             other_args=("--max-total-token", "1024", "--context-len", "8192"),
             env={"SGLANG_CLIP_MAX_NEW_TOKENS": "256", **os.environ},
-            return_stdout_stderr=True,
+            return_stdout_stderr=(cls.stdout, cls.stderr),
         )
         cls.base_url += "/v1"
         cls.tokenizer = get_tokenizer(DEFAULT_MODEL_NAME_FOR_TEST)
@@ -35,6 +43,10 @@ class TestLargeMaxNewTokens(unittest.TestCase):
     @classmethod
     def tearDownClass(cls):
         kill_child_process(cls.process.pid)
+        cls.stdout.close()
+        cls.stderr.close()
+        os.remove("stdout.txt")
+        os.remove("stderr.txt")
 
     def run_chat_completion(self):
         client = openai.Client(api_key=self.api_key, base_url=self.base_url)
@@ -56,16 +68,21 @@ class TestLargeMaxNewTokens(unittest.TestCase):
 
         futures = []
         with ThreadPoolExecutor(num_requests) as executor:
+            # Send multiple requests
             for i in range(num_requests):
                 futures.append(executor.submit(self.run_chat_completion))
 
-            all_requests_running = False
-            for line in iter(self.process.stderr.readline, ""):
-                line = str(line)
-                print(line, end="")
-                if f"#running-req: {num_requests}" in line:
-                    all_requests_running = True
-                    break
+            # Ensure that they are running concurrently
+            pt = 0
+            while pt >= 0:
+                lines = open("stderr.txt").readlines()
+                for line in lines[pt:]:
+                    print(line, end="", flush=True)
+                    if f"#running-req: {num_requests}" in line:
+                        all_requests_running = True
+                        pt = -1
+                        break
+                    pt += 1
 
         assert all_requests_running