diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e0e6038d60..46b11f5034 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -410,7 +410,8 @@ def _launch_subprocesses(args: StartArgs): ports_to_check = [args.port] if args.dp == 1 and args.nnodes > 1: ports_to_check.extend([args.multinode_httpmanager_port, args.multinode_router_gloo_port]) - if args.node_rank == 0 and args.nccl_port is not None: + # ConfigServer owns and replaces its TCPStore; a surviving store must not block rank-0 restart. + if args.node_rank == 0 and args.nccl_port is not None and not args.use_config_server_to_init_nccl: ports_to_check.append(args.nccl_port) validate_ports(ports_to_check) diff --git a/lightllm/server/config_server/api_http.py b/lightllm/server/config_server/api_http.py index 285e82e6bf..fc20115578 100644 --- a/lightllm/server/config_server/api_http.py +++ b/lightllm/server/config_server/api_http.py @@ -3,7 +3,7 @@ import base64 import pickle import setproctitle -import multiprocessing as mp +import subprocess from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Query from threading import Lock from typing import Dict, List @@ -156,7 +156,7 @@ async def allocate_global_unique_multimodal_id_range(): return {"start_id": start_id, "end_id": end_id} -global_store_port_to_process: Dict[int, mp.Process] = {} +global_store_port_to_process: Dict[int, subprocess.Popen] = {} global_store_port_to_client_states: Dict[int, List[bool]] = {} global_store_port_lock = asyncio.Lock() @@ -192,7 +192,7 @@ async def http_start_tcp_store_server( logger.warning(f"tcp store server {tcp_store_port} already started, kill and restart it") process = global_store_port_to_process[tcp_store_port] process.kill() - process.join() + process.wait() global_store_port_to_process[tcp_store_port] = start_tcp_store_server( args.config_server_host, tcp_store_port diff --git a/lightllm/server/config_server/nccl_tcp_store.py b/lightllm/server/config_server/nccl_tcp_store.py index 91f779e612..300f0a7c68 100644 --- a/lightllm/server/config_server/nccl_tcp_store.py +++ b/lightllm/server/config_server/nccl_tcp_store.py @@ -1,8 +1,10 @@ -import psutil +import argparse +import os +import subprocess +import sys import time import setproctitle import torch.distributed as dist -import torch.multiprocessing as mp from lightllm.utils.log_utils import init_logger from lightllm.utils.process_check import start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name @@ -14,20 +16,27 @@ def start_tcp_store_server(nccl_store_host, nccl_store_port): """ start a process to run a TCPStore server. """ - process = mp.Process( - target=_start_tcp_store_server, - args=(nccl_store_host, nccl_store_port), - daemon=True, + # Hypercorn workers are multiprocessing daemons and cannot spawn mp.Process children. + return subprocess.Popen( + args=[ + sys.executable, + "-m", + "lightllm.server.config_server.nccl_tcp_store", + "--host", + nccl_store_host, + "--port", + str(nccl_store_port), + "--parent-pid", + str(os.getpid()), + ] ) - process.start() - return process -def _start_tcp_store_server(nccl_store_host, nccl_store_port): +def _start_tcp_store_server(nccl_store_host, nccl_store_port, parent_pid: int | None = None): """ start a TCPStore server. """ - start_parent_check_thread() + start_parent_check_thread(parent_pid=parent_pid) setproctitle.setproctitle(f"lightllm::{get_unique_server_name()}::nccl_tcp_store") try: @@ -58,6 +67,16 @@ def _start_tcp_store_server(nccl_store_host, nccl_store_port): logger.info(f"[Server] TCPStore start: {nccl_store_host}:{nccl_store_port} keys num: {keys_num}") time.sleep(20) - except Exception as e: - logger.warning(str(e)) - logger.info(f"TCPStore server {nccl_store_host}:{nccl_store_port} start failed, retrying ...") + except Exception: + logger.exception(f"TCPStore server {nccl_store_host}:{nccl_store_port} start failed") + raise + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", required=True) + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--parent-pid", type=int, required=True) + args = parser.parse_args() + if os.getppid() == args.parent_pid: + _start_tcp_store_server(nccl_store_host=args.host, nccl_store_port=args.port, parent_pid=args.parent_pid) diff --git a/lightllm/utils/process_check.py b/lightllm/utils/process_check.py index d92eae439d..f8bd4c50b9 100644 --- a/lightllm/utils/process_check.py +++ b/lightllm/utils/process_check.py @@ -4,6 +4,7 @@ import threading import psutil import signal +from typing import Optional from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -42,10 +43,11 @@ def kill_child_processes(parent_pid): logger.warning(f"kill pid {child.pid} failed {str(e)}") -def check_parent_alive(): +def check_parent_alive(parent_pid: Optional[int] = None): # A late-starting worker may already have been reparented after its parent exited. - parent = mp.parent_process() - parent_pid = parent.pid if parent is not None else os.getppid() + if parent_pid is None: + parent = mp.parent_process() + parent_pid = parent.pid if parent is not None else os.getppid() while True: if not is_process_active(parent_pid): logger.warning(f"parent is dead, kill self {os.getpid()}") @@ -56,9 +58,9 @@ def check_parent_alive(): return -def start_parent_check_thread(): +def start_parent_check_thread(parent_pid: Optional[int] = None): """ 检测父进程是否健康，如果出现问题，清理退出所有进程 """ - thread = threading.Thread(target=check_parent_alive, daemon=True) + thread = threading.Thread(target=check_parent_alive, kwargs={"parent_pid": parent_pid}, daemon=True) thread.start() diff --git a/unit_tests/server/test_config_server_tcp_store.py b/unit_tests/server/test_config_server_tcp_store.py new file mode 100644 index 0000000000..15248c0e7b --- /dev/null +++ b/unit_tests/server/test_config_server_tcp_store.py @@ -0,0 +1,111 @@ +import asyncio +from datetime import timedelta +import multiprocessing as mp +from multiprocessing.connection import Connection +import os +import socket +import subprocess +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import psutil +import pytest +import torch.distributed as dist + +from lightllm.server.config_server import nccl_tcp_store +from lightllm.utils import process_check + + +def _run_store(connection: Connection, port: int) -> None: + process = nccl_tcp_store.start_tcp_store_server(nccl_store_host="127.0.0.1", nccl_store_port=port) + try: + connection.send(process.pid) + connection.recv() + finally: + process.terminate() + process.wait(timeout=10) + connection.close() + + +@pytest.mark.skipif(sys.platform != "linux", reason="requires fork") +@pytest.mark.parametrize("daemon", [False, True]) +def test_tcp_store_serves_clients_from_a_daemon_worker(daemon: bool, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(name="LIGHTLLM_UNIQUE_SERVICE_NAME_ID", value="tcp_store_regression") + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + context = mp.get_context(method="fork") + reader, writer = context.Pipe() + worker = context.Process(target=_run_store, args=(writer, port), daemon=daemon) + worker.start() + server_pid = None + try: + assert reader.poll(timeout=10), "worker did not start TCPStore" + server_pid = reader.recv() + client = dist.TCPStore( + host_name="127.0.0.1", port=port, is_master=False, wait_for_workers=False, timeout=timedelta(seconds=20) + ) + client.set("roundtrip", "value") + assert client.get("roundtrip") == b"value" + del client + reader.send("stop") + worker.join(timeout=10) + assert worker.exitcode == 0 + assert not process_check.is_process_active(pid=server_pid) + finally: + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + if server_pid is not None and process_check.is_process_active(pid=server_pid): + psutil.Process(pid=server_pid).kill() + worker.close() + reader.close() + writer.close() + + +def test_restart_reaps_previous_subprocess(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(target=process_check, name="start_parent_check_thread", value=lambda: None) + from lightllm.server.config_server import api_http + + old = Mock(spec=subprocess.Popen) + new = Mock(spec=subprocess.Popen) + start = Mock(return_value=new) + monkeypatch.setattr( + target=api_http, name="get_env_start_args", value=lambda: SimpleNamespace(config_server_host="host") + ) + monkeypatch.setattr(target=api_http, name="start_tcp_store_server", value=start) + monkeypatch.setattr(target=api_http, name="global_store_port_to_process", value={29501: old}) + monkeypatch.setattr(target=api_http, name="global_store_port_to_client_states", value={}) + monkeypatch.setattr(target=api_http, name="global_store_port_lock", value=asyncio.Lock()) + + result = asyncio.run(main=api_http.http_start_tcp_store_server(tcp_store_port=29501, rank_id=0, world_size=1)) + + assert result == {"status": "ok"} + old.kill.assert_called_once_with() + old.wait.assert_called_once_with() + start.assert_called_once_with("host", 29501) + assert api_http.global_store_port_to_process == {29501: new} + assert not api_http.global_store_port_to_client_states + + +def test_tcp_store_propagates_start_failure(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(target=nccl_tcp_store, name="start_parent_check_thread", value=lambda parent_pid: None) + monkeypatch.setattr(target=nccl_tcp_store, name="get_unique_server_name", value=lambda: "test") + monkeypatch.setattr(target=dist, name="TCPStore", value=Mock(side_effect=RuntimeError("port already in use"))) + with pytest.raises(expected_exception=RuntimeError, match="port already in use"): + nccl_tcp_store._start_tcp_store_server(nccl_store_host="127.0.0.1", nccl_store_port=29501, parent_pid=123) + + +def test_parent_monitor_uses_launch_parent_after_reparenting(monkeypatch: pytest.MonkeyPatch) -> None: + active = Mock(return_value=False) + kill_children = Mock() + kill = Mock(side_effect=SystemExit) + monkeypatch.setattr(target=process_check, name="is_process_active", value=active) + monkeypatch.setattr(target=process_check, name="kill_child_processes", value=kill_children) + monkeypatch.setattr(target=process_check.os, name="kill", value=kill) + monkeypatch.setattr(target=process_check.os, name="getppid", value=Mock(side_effect=AssertionError("reparented"))) + with pytest.raises(expected_exception=SystemExit): + process_check.check_parent_alive(parent_pid=123) + active.assert_called_once_with(123) + kill_children.assert_called_once_with(os.getpid()) diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py index 3109fe11b9..ac11d8b1cf 100644 --- a/unit_tests/server/test_pd_start_args.py +++ b/unit_tests/server/test_pd_start_args.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import pytest @@ -6,6 +7,53 @@ from lightllm.server.core.objs.start_args_type import StartArgs +@pytest.mark.parametrize("config_owned", [False, True]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_startup_checks_only_locally_owned_rendezvous_ports( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, config_owned: bool, rank: int +) -> None: + from lightllm.server import api_start + + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + "auto_set_page_size", + "auto_set_response_parsers", + "auto_configure_allreduce_flags_from_args", + ): + monkeypatch.setattr(target=api_start, name=name, value=lambda args: None) + (tmp_path / "config.json").write_text(data=json.dumps(obj={"model_type": "llama"})) + args = StartArgs( + run_mode="prefill", + model_dir=str(tmp_path), + tp=2, + dp=2, + nnodes=2, + node_rank=rank, + use_config_server_to_init_nccl=config_owned, + config_server_host="127.0.0.1", + nccl_host="127.0.0.1", + port=8019, + nccl_port=29501, + max_req_total_len=8192, + eos_id=[2], + data_type="bfloat16", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + def checked(ports: list[int]) -> None: + assert ports == ([8019, 29501] if rank == 0 and not config_owned else [8019]) + raise RuntimeError("port validation reached") + + monkeypatch.setattr(target=api_start, name="validate_ports", value=checked) + with pytest.raises(expected_exception=RuntimeError, match="port validation reached"): + api_start._launch_subprocesses(args=args) + + @pytest.fixture def startup_args(monkeypatch): for name in (