diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index e082d8643..06b57b63d 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -55,6 +55,7 @@ from sglang.srt.entrypoints.engine_info_bootstrap_server import ( from sglang.srt.entrypoints.engine_score_mixin import EngineScoreMixin from sglang.srt.entrypoints.EngineBase import EngineBase from sglang.srt.managers.data_parallel_controller import ( + SCHEDULER_PIDS_ARG, run_data_parallel_controller_process, ) from sglang.srt.managers.detokenizer_manager import run_detokenizer_process @@ -115,6 +116,7 @@ class SchedulerInitResult: """Result from launching schedulers.""" scheduler_infos: List[Dict[str, Any]] + all_child_pids: List[int] = dataclasses.field(default_factory=list) wait_for_ready: Callable[[], None] = lambda: None wait_for_completion: Callable[[], None] = lambda: None engine_info_bootstrap_server: Optional[Any] = None @@ -242,6 +244,10 @@ class Engine(EngineScoreMixin, EngineBase): self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) + def get_all_child_pids(self) -> List[int]: + """Returns a list of all child process PIDs.""" + return self._scheduler_init_result.all_child_pids + def _resolve_routed_dp_rank( self, routed_dp_rank: Optional[int], @@ -606,11 +612,17 @@ class Engine(EngineScoreMixin, EngineBase): proc.start() scheduler_procs.append(proc) + all_child_pids = [proc.pid for proc in scheduler_procs] scheduler_infos = [] def wait_for_ready(): infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs) scheduler_infos.extend(infos) + # For dp_size > 1, collect child scheduler PIDs from the DP controller + if server_args.dp_size > 1: + for info in infos: + if SCHEDULER_PIDS_ARG in info: + all_child_pids.extend(info[SCHEDULER_PIDS_ARG]) def wait_for_completion(): for proc in scheduler_procs: @@ -623,6 +635,7 @@ class Engine(EngineScoreMixin, EngineBase): return ( SchedulerInitResult( scheduler_infos=scheduler_infos, + all_child_pids=all_child_pids, wait_for_ready=wait_for_ready, wait_for_completion=wait_for_completion, ), @@ -733,6 +746,7 @@ class Engine(EngineScoreMixin, EngineBase): ), ) detoken_proc.start() + scheduler_init_result.all_child_pids.append(detoken_proc.pid) # Init tokenizer manager first, as the bootstrap server is initialized here if server_args.tokenizer_worker_num == 1: diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 1b71b1853..29b454a45 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -66,6 +66,8 @@ from sglang.utils import TypeBasedDispatcher, get_exception_traceback logger = logging.getLogger(__name__) +SCHEDULER_PIDS_ARG = "scheduler_pids" + class LoadBalanceMethod(Enum): """Load balance method.""" @@ -617,11 +619,15 @@ def run_data_parallel_controller_process( controller = DataParallelController( server_args, port_args, run_scheduler_process_func ) + scheduler_pids = [ + proc.pid for proc in controller.scheduler_procs if proc is not None + ] pipe_writer.send( { "status": "ready", "max_total_num_tokens": controller.max_total_num_tokens, "max_req_input_len": controller.max_req_input_len, + SCHEDULER_PIDS_ARG: scheduler_pids, } ) if server_args.node_rank == 0: diff --git a/test/registered/core/test_engine_child_pids.py b/test/registered/core/test_engine_child_pids.py new file mode 100644 index 000000000..7762ef79e --- /dev/null +++ b/test/registered/core/test_engine_child_pids.py @@ -0,0 +1,90 @@ +""" +Unit tests for Engine.get_all_child_pids(). + +Verifies that launching an Engine exposes the PIDs of all child processes +(schedulers, detokenizer) and that those PIDs correspond to live processes. + +Usage: + python -m unittest test_engine_child_pids -v +""" + +import os +import unittest + +import psutil + +import sglang as sgl +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + CustomTestCase, +) + +register_cuda_ci(est_time=60, suite="stage-b-test-1-gpu-small") + + +class TestEngineChildPids(CustomTestCase): + + def test_get_all_child_pids_returns_live_pids(self): + engine = sgl.Engine( + model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + random_seed=42, + ) + try: + pids = engine.get_all_child_pids() + + self.assertIsInstance(pids, list) + self.assertGreater(len(pids), 0, "Expected at least one child PID") + + for pid in pids: + self.assertIsInstance(pid, int) + self.assertTrue( + psutil.pid_exists(pid), + f"PID {pid} does not correspond to a running process", + ) + + current_proc = psutil.Process(os.getpid()) + child_pids = {c.pid for c in current_proc.children(recursive=True)} + for pid in pids: + self.assertIn( + pid, + child_pids, + f"PID {pid} is not a child of the current process", + ) + finally: + engine.shutdown() + + def test_child_pids_include_scheduler_and_detokenizer(self): + engine = sgl.Engine( + model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + random_seed=42, + ) + try: + pids = engine.get_all_child_pids() + # dp_size=1 gives one scheduler + one detokenizer = at least 2 PIDs + self.assertGreaterEqual( + len(pids), + 2, + "Expected at least 2 child PIDs (scheduler + detokenizer)", + ) + finally: + engine.shutdown() + + def test_child_pids_no_duplicates(self): + engine = sgl.Engine( + model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + random_seed=42, + ) + try: + pids = engine.get_all_child_pids() + self.assertEqual( + len(pids), + len(set(pids)), + f"Duplicate PIDs found: {pids}", + ) + finally: + engine.shutdown() + + +if __name__ == "__main__": + unittest.main()