Expose child process PIDs from Engine for health check support (#23320)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user