Extract dumper and prefill delayer tests common utils (#18857)

test/registered/scheduler/test_prefill_delayer.py

@@ -3,7 +3,6 @@ import os
 import re
 import time
 import unittest
-from collections import defaultdict
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import List, Optional
@@ -11,7 +10,6 @@ from typing import List, Optional
 import openai
 import requests
 import torch
-import torch.multiprocessing as mp
 
 from sglang.bench_serving import run_benchmark
 from sglang.srt.managers.prefill_delayer import PrefillDelayer
@@ -25,6 +23,7 @@ from sglang.test.test_utils import (
     CustomTestCase,
     get_benchmark_args,
     popen_launch_server,
+    run_distributed_test,
 )
 
 register_cuda_ci(
@@ -54,13 +53,8 @@ class NegotiateTestCase:
     expected_reason: str
 
 
-def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
-    torch.distributed.init_process_group(
-        backend="gloo",
-        init_method=f"tcp://127.0.0.1:{port}",
-        world_size=world_size,
-        rank=rank,
-    )
+def _run_negotiate_test(rank, test_cases):
+    world_size = torch.distributed.get_world_size()
     cpu_group = torch.distributed.new_group(backend="gloo")
 
     for case in test_cases:
@@ -83,9 +77,10 @@ def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
                 token_usage=call.token_usage[rank],
             )
 
-        results_queue.put((rank, case.name, result.output_allow, result.output_reason))
-
-    torch.distributed.destroy_process_group()
+        assert (result.output_allow, result.output_reason) == (
+            case.expected_allow,
+            case.expected_reason,
+        ), f"Case {case.name} rank {rank}"
 
 
 _NEGOTIATE_TEST_CASES = [
@@ -210,38 +205,12 @@ _NEGOTIATE_TEST_CASES = [
 
 class TestPrefillDelayerNegotiate(unittest.TestCase):
     def test_negotiate(self):
-        world_size = 4
-        test_cases = _NEGOTIATE_TEST_CASES
-
-        ctx = mp.get_context("spawn")
-        results_queue = ctx.Queue()
-        port = 29500 + os.getpid() % 1000
-
-        processes = []
-        for rank in range(world_size):
-            p = ctx.Process(
-                target=_run_negotiate_test,
-                args=(rank, world_size, test_cases, results_queue, port),
-            )
-            p.start()
-            processes.append(p)
-
-        for p in processes:
-            p.join()
-
-        results = defaultdict(dict)
-        for _ in range(world_size * len(test_cases)):
-            rank, case_name, output_allow, output_reason = results_queue.get()
-            results[case_name][rank] = (output_allow, output_reason)
-
-        for case in test_cases:
-            for rank in range(world_size):
-                output_allow, output_reason = results[case.name][rank]
-                self.assertEqual(
-                    (output_allow, output_reason),
-                    (case.expected_allow, case.expected_reason),
-                    f"Case {case.name} rank {rank}",
-                )
+        run_distributed_test(
+            _run_negotiate_test,
+            world_size=4,
+            backend="gloo",
+            test_cases=_NEGOTIATE_TEST_CASES,
+        )
 
 
 # ============================ E2E Tests ============================