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.engine_score_mixin import EngineScoreMixin
|
||||||
from sglang.srt.entrypoints.EngineBase import EngineBase
|
from sglang.srt.entrypoints.EngineBase import EngineBase
|
||||||
from sglang.srt.managers.data_parallel_controller import (
|
from sglang.srt.managers.data_parallel_controller import (
|
||||||
|
SCHEDULER_PIDS_ARG,
|
||||||
run_data_parallel_controller_process,
|
run_data_parallel_controller_process,
|
||||||
)
|
)
|
||||||
from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
|
from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
|
||||||
@@ -115,6 +116,7 @@ class SchedulerInitResult:
|
|||||||
"""Result from launching schedulers."""
|
"""Result from launching schedulers."""
|
||||||
|
|
||||||
scheduler_infos: List[Dict[str, Any]]
|
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_ready: Callable[[], None] = lambda: None
|
||||||
wait_for_completion: Callable[[], None] = lambda: None
|
wait_for_completion: Callable[[], None] = lambda: None
|
||||||
engine_info_bootstrap_server: Optional[Any] = None
|
engine_info_bootstrap_server: Optional[Any] = None
|
||||||
@@ -242,6 +244,10 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
self.loop = asyncio.new_event_loop()
|
self.loop = asyncio.new_event_loop()
|
||||||
asyncio.set_event_loop(self.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(
|
def _resolve_routed_dp_rank(
|
||||||
self,
|
self,
|
||||||
routed_dp_rank: Optional[int],
|
routed_dp_rank: Optional[int],
|
||||||
@@ -606,11 +612,17 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
proc.start()
|
proc.start()
|
||||||
scheduler_procs.append(proc)
|
scheduler_procs.append(proc)
|
||||||
|
|
||||||
|
all_child_pids = [proc.pid for proc in scheduler_procs]
|
||||||
scheduler_infos = []
|
scheduler_infos = []
|
||||||
|
|
||||||
def wait_for_ready():
|
def wait_for_ready():
|
||||||
infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs)
|
infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs)
|
||||||
scheduler_infos.extend(infos)
|
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():
|
def wait_for_completion():
|
||||||
for proc in scheduler_procs:
|
for proc in scheduler_procs:
|
||||||
@@ -623,6 +635,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
return (
|
return (
|
||||||
SchedulerInitResult(
|
SchedulerInitResult(
|
||||||
scheduler_infos=scheduler_infos,
|
scheduler_infos=scheduler_infos,
|
||||||
|
all_child_pids=all_child_pids,
|
||||||
wait_for_ready=wait_for_ready,
|
wait_for_ready=wait_for_ready,
|
||||||
wait_for_completion=wait_for_completion,
|
wait_for_completion=wait_for_completion,
|
||||||
),
|
),
|
||||||
@@ -733,6 +746,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
detoken_proc.start()
|
detoken_proc.start()
|
||||||
|
scheduler_init_result.all_child_pids.append(detoken_proc.pid)
|
||||||
|
|
||||||
# Init tokenizer manager first, as the bootstrap server is initialized here
|
# Init tokenizer manager first, as the bootstrap server is initialized here
|
||||||
if server_args.tokenizer_worker_num == 1:
|
if server_args.tokenizer_worker_num == 1:
|
||||||
|
|||||||
@@ -66,6 +66,8 @@ from sglang.utils import TypeBasedDispatcher, get_exception_traceback
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
SCHEDULER_PIDS_ARG = "scheduler_pids"
|
||||||
|
|
||||||
|
|
||||||
class LoadBalanceMethod(Enum):
|
class LoadBalanceMethod(Enum):
|
||||||
"""Load balance method."""
|
"""Load balance method."""
|
||||||
@@ -617,11 +619,15 @@ def run_data_parallel_controller_process(
|
|||||||
controller = DataParallelController(
|
controller = DataParallelController(
|
||||||
server_args, port_args, run_scheduler_process_func
|
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(
|
pipe_writer.send(
|
||||||
{
|
{
|
||||||
"status": "ready",
|
"status": "ready",
|
||||||
"max_total_num_tokens": controller.max_total_num_tokens,
|
"max_total_num_tokens": controller.max_total_num_tokens,
|
||||||
"max_req_input_len": controller.max_req_input_len,
|
"max_req_input_len": controller.max_req_input_len,
|
||||||
|
SCHEDULER_PIDS_ARG: scheduler_pids,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
if server_args.node_rank == 0:
|
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