fix(config): support TCPStore startup in daemon workers

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)
 

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

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)

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()

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())

unit_tests/server/test_pd_start_args.py

@@ -1,11 +1,59 @@
 import json
+from pathlib import Path
 
 import pytest
 
 from lightllm.server.api_start import _launch_subprocesses
 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 (