[CI] Fix the issue of unit test hanging (#1211)

test/srt/models/test_embedding_models.py

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import multiprocessing as mp
 import unittest
 
 import torch
@@ -71,4 +72,9 @@ class TestEmbeddingModels(unittest.TestCase):
 
 
 if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
     unittest.main(warnings="ignore")

test/srt/models/test_generation_models.py

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import multiprocessing as mp
 import unittest
 
 import torch
@@ -108,13 +109,6 @@ class TestGenerationModels(unittest.TestCase):
         ), f"Not all ROUGE-L scores are greater than {rouge_threshold}"
 
     def test_prefill_logits_and_output_strs(self):
-        import multiprocessing as mp
-
-        try:
-            mp.set_start_method("spawn")
-        except RuntimeError:
-            pass
-
         for (
             model,
             tp_size,
@@ -137,4 +131,9 @@ class TestGenerationModels(unittest.TestCase):
 
 
 if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
     unittest.main(warnings="ignore")

test/srt/run_suite.py

@@ -1,5 +1,6 @@
 import argparse
 import glob
+import multiprocessing as mp
 
 from sglang.test.test_utils import run_unittest_files
 
@@ -54,5 +55,10 @@ if __name__ == "__main__":
     else:
         files = suites[args.suite]
 
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
     exit_code = run_unittest_files(files, args.timeout_per_file)
     exit(exit_code)