diff --git a/lightllm/server/router/model_infer/mode_backend/pd/base_kv_move_manager.py b/lightllm/server/router/model_infer/mode_backend/pd/base_kv_move_manager.py index 1e170975a3..8326affeba 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/base_kv_move_manager.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/base_kv_move_manager.py @@ -1,12 +1,16 @@ import time import os -import signal import threading import queue import torch.multiprocessing as mp from typing import List, Dict, Union, Callable, Optional from lightllm.utils.log_utils import init_logger -from lightllm.server.pd_io_struct import PDChunckedTransTaskRet +from lightllm.server.pd_io_struct import ( + PDAbortReq, + PDChunckedTransTask, + PDChunckedTransTaskGroup, + PDChunckedTransTaskRet, +) from lightllm.server.core.objs import StartArgs from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer from .trans_process_obj import KVTransProcess @@ -14,6 +18,9 @@ logger = init_logger(__name__) +_MAX_TRANS_PROCESS_RESTARTS = 3 +_TRANS_PROCESS_RESTART_RESET_SECONDS = 60 + class BaseKVMoveManager: def __init__( @@ -33,8 +40,16 @@ def __init__( self.info_queue = info_queue self.ret_obj_queue = queue.Queue() + self.start_trans_process_func = start_trans_process_func + self.up_status_in_queue = up_status_in_queue self.kv_trans_processes: List[KVTransProcess] = [None] * self.node_world_size + self._trans_process_locks = [threading.Lock() for _ in range(self.node_world_size)] + self._pending_trans_tasks_by_device: List[Dict[str, List[PDChunckedTransTask]]] = [ + {} for _ in range(self.node_world_size) + ] + self._trans_process_restart_counts = [0] * self.node_world_size + self._trans_process_last_restart_at: List[Optional[float]] = [None] * self.node_world_size for device_id in range(self.node_world_size): trans_process = KVTransProcess() self.kv_trans_processes[device_id] = trans_process @@ -66,6 +81,37 @@ def __init__( def task_dispatcher_loop(self): raise NotImplementedError() + def dispatch_trans_task( + self, + device_id: int, + task: Union[PDChunckedTransTask, PDChunckedTransTaskGroup, PDAbortReq], + ) -> None: + while True: + with self._trans_process_locks[device_id]: + trans_process = self.kv_trans_processes[device_id] + if trans_process.is_trans_process_health(): + trans_process.task_in_queue.put(task) + self._register_pending_trans_tasks(device_id, task) + return + time.sleep(0.01) + + def _register_pending_trans_tasks( + self, + device_id: int, + task: Union[PDChunckedTransTask, PDChunckedTransTaskGroup, PDAbortReq], + ) -> None: + if isinstance(task, PDChunckedTransTaskGroup): + tasks = task.task_list + elif isinstance(task, PDChunckedTransTask): + tasks = [task] + else: + return + + pending_tasks = self._pending_trans_tasks_by_device[device_id] + for trans_task in tasks: + key = trans_task.createRetObj().get_key() + pending_tasks.setdefault(key, []).append(trans_task) + # ================================================================================== # 将收集到的传输返回信息，批量写回给推理进程，触发其进行相关管理信息的更新。 # ================================================================================== @@ -101,29 +147,100 @@ def _collect_return_objects(self): # 处理各个传输任务的返回的信息 # ================================================================================== @log_exception - def task_ret_handle_loop(self, trans_process: KVTransProcess): + def task_ret_handle_loop(self, trans_process: KVTransProcess) -> None: while True: - ret_obj: PDChunckedTransTaskRet = trans_process.task_out_queue.get() - self.ret_obj_queue.put(ret_obj) + ret_obj: Optional[PDChunckedTransTaskRet] = trans_process.task_out_queue.get() + if ret_obj is None: + return + device_id = trans_process.device_id + with self._trans_process_locks[device_id]: + if self.kv_trans_processes[device_id] is not trans_process: + continue + + pending_tasks = self._pending_trans_tasks_by_device[device_id] + ret_key = ret_obj.get_key() + tasks_for_key = pending_tasks.get(ret_key) + if not tasks_for_key: + continue + tasks_for_key.pop() + if not tasks_for_key: + del pending_tasks[ret_key] + self.ret_obj_queue.put(ret_obj) + + def _fail_pending_trans_tasks(self, device_id: int) -> None: + pending_tasks = self._pending_trans_tasks_by_device[device_id] + tasks_to_fail = [task for tasks in pending_tasks.values() for task in tasks] + pending_tasks.clear() + + for trans_task in tasks_to_fail: + trans_task.error_info = f"KV transfer worker for device {device_id} restarted before reporting task result" + self.ret_obj_queue.put(trans_task.createRetObj()) + + def _restart_trans_process(self, device_id: int) -> None: + with self._trans_process_locks[device_id]: + old_trans_process = self.kv_trans_processes[device_id] + if old_trans_process.process is not None: + if old_trans_process.process.is_alive(): + old_trans_process.killself() + old_trans_process.process.join(timeout=5) + if old_trans_process.process.is_alive(): + raise RuntimeError(f"KV transfer worker for device {device_id} did not stop before restart") + self._fail_pending_trans_tasks(device_id) + try: + old_trans_process.task_out_queue.put(None) + except (EOFError, OSError, ValueError): + logger.warning(f"Could not stop return handler for KV transfer worker on device {device_id}") + + trans_process = KVTransProcess() + if not trans_process.init_all( + device_id=device_id, + manager=self, + start_func=self.start_trans_process_func, + up_status_in_queue=self.up_status_in_queue, + ): + if trans_process.process is not None and trans_process.process.is_alive(): + trans_process.killself() + raise RuntimeError(f"KV transfer worker for device {device_id} failed to restart") + + if not trans_process.wait_until_ready(): + if trans_process.process is not None and trans_process.process.is_alive(): + trans_process.killself() + raise RuntimeError(f"KV transfer worker for device {device_id} did not become ready after restart") + + self.kv_trans_processes[device_id] = trans_process + threading.Thread(target=self.task_ret_handle_loop, args=(trans_process,), daemon=True).start() + self._trans_process_last_restart_at[device_id] = time.monotonic() + logger.warning(f"Restarted KV transfer worker for device {device_id}") # ================================================================================== # 定时检测传输进程的健康状态，出现问题拉崩整个系统触发重启 # ================================================================================== - def check_trans_process_loop(self): + def check_trans_process_loop(self) -> None: try: while True: for device_id in range(self.node_world_size): if not self.kv_trans_processes[device_id].is_trans_process_health(): - raise Exception(f"device_id {device_id} kv process is unhealth") - - time.sleep(10.0) + self._trans_process_restart_counts[device_id] += 1 + if self._trans_process_restart_counts[device_id] > _MAX_TRANS_PROCESS_RESTARTS: + raise RuntimeError(f"KV transfer worker for device {device_id} exceeded restart limit") + logger.warning(f"KV transfer worker for device {device_id} is unhealthy; restarting") + self._restart_trans_process(device_id) + elif ( + self._trans_process_last_restart_at[device_id] is not None + and time.monotonic() - self._trans_process_last_restart_at[device_id] + >= _TRANS_PROCESS_RESTART_RESET_SECONDS + ): + self._trans_process_restart_counts[device_id] = 0 + + time.sleep(1.0) except (BaseException, RuntimeError) as e: logger.exception(str(e)) for device_id in range(self.node_world_size): - self.kv_trans_processes[device_id].killself() + try: + self.kv_trans_processes[device_id].killself() + except Exception: + logger.exception(f"Failed to stop KV transfer worker for device {device_id}") - # 杀掉当前进程的父进程（router), 触发全局崩溃 - os.kill(os.getppid(), signal.SIGKILL) - os.kill(os.getpid(), signal.SIGKILL) - raise e + # The launcher supervises this manager and removes the node on exit. + os._exit(1) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_kv_move_manager.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_kv_move_manager.py index 6d567763fe..f3a20e4191 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_kv_move_manager.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_kv_move_manager.py @@ -10,7 +10,6 @@ from lightllm.server.pd_io_struct import PDChunckedTransTaskGroup, PDAbortReq from lightllm.server.core.objs import StartArgs from lightllm.utils.graceful_utils import graceful_registry -from ..trans_process_obj import KVTransProcess from ..base_kv_move_manager import BaseKVMoveManager from lightllm.utils.error_utils import log_exception from lightllm.utils.envs_utils import get_unique_server_name @@ -99,8 +98,7 @@ def task_dispatcher_loop(self): assert False, f"error obj {task_group}" try: - trans_process: KVTransProcess = self.kv_trans_processes[device_id] - trans_process.task_in_queue.put(task_group) + self.dispatch_trans_task(device_id, task_group) # A task group covers the whole request, so this log does not require pd_detail_log. if isinstance(task_group, PDChunckedTransTaskGroup) and logger.isEnabledFor(logging.DEBUG): logger.debug( diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index d3f2c0214d..d9f7f0053c 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -2,12 +2,17 @@ import errno import queue import pickle +import socket import threading +import time +import zlib from dataclasses import dataclass from typing import Dict, List, Optional import rpyc import torch +from rpyc.core.channel import Channel as RPyCChannel +from rpyc.core.stream import SocketStream from torch import Tensor from rpyc.utils.classic import obtain from rpyc.utils.server import ThreadedServer @@ -20,6 +25,25 @@ logger = init_logger(__name__) +_MAX_CONTROL_STORE_KEY_BYTES = 1024 +_MAX_CONTROL_STORE_VALUE_BYTES = 1 << 20 +_MAX_CONTROL_STORE_ENTRIES = 1024 +_MAX_CONTROL_STORE_BYTES = 16 << 20 +_MAX_CONTROL_RPC_FRAME_BYTES = 2 << 20 +_CONTROL_STORE_VALUE_TTL_SECONDS = 60 * 60 +_CONTROL_STORE_CLEANUP_INTERVAL_SECONDS = 60 + + +def _validate_control_store_key(key: str) -> None: + if not isinstance(key, str): + raise TypeError("NCCL control store keys must be strings") + if len(key.encode("utf-8")) > _MAX_CONTROL_STORE_KEY_BYTES: + raise ValueError(f"NCCL control store key exceeds {_MAX_CONTROL_STORE_KEY_BYTES} bytes") + + +def _control_store_now() -> float: + return time.monotonic() + @dataclass class NcclAgentMetadata: @@ -410,12 +434,59 @@ def exposed_push_notif(self, payload: bytes): self.channel.notif_queue.put(payload) return - def exposed_set_value(self, key: str, value: bytes): + def exposed_set_value(self, key: str, value: bytes) -> None: key = obtain(key) value = obtain(value) self.channel.add_store_value(key, value) return + def exposed_delete_value(self, key: str) -> None: + key = obtain(key) + self.channel.delete_store_value(key) + return + + +class _BoundedControlRPCChannel(RPyCChannel): + def recv(self) -> bytes: + header = self.stream.read(self.FRAME_HEADER.size) + length, compressed = self.FRAME_HEADER.unpack(header) + if length > _MAX_CONTROL_RPC_FRAME_BYTES: + raise _OversizedControlRPCFrame(f"frame exceeds {_MAX_CONTROL_RPC_FRAME_BYTES} bytes") + + data = self.stream.read(length + len(self.FLUSHER))[: -len(self.FLUSHER)] + if compressed: + decompressor = zlib.decompressobj() + data = decompressor.decompress(data, _MAX_CONTROL_RPC_FRAME_BYTES + 1) + if len(data) > _MAX_CONTROL_RPC_FRAME_BYTES or decompressor.unconsumed_tail or not decompressor.eof: + raise _OversizedControlRPCFrame(f"frame exceeds {_MAX_CONTROL_RPC_FRAME_BYTES} bytes") + return data + + +class _OversizedControlRPCFrame(ValueError): + pass + + +class _NcclControlServer(ThreadedServer): + def _serve_client(self, sock: socket.socket, credentials: Optional[object]) -> None: + addrinfo = sock.getpeername() + if credentials: + self.logger.info(f"welcome {addrinfo} ({credentials!r})") + else: + self.logger.info(f"welcome {addrinfo}") + try: + config = dict( + self.protocol_config, + credentials=credentials, + endpoints=(sock.getsockname(), addrinfo), + logger=self.logger, + ) + conn = self.service._connect(_BoundedControlRPCChannel(SocketStream(sock)), config) + self._handle_connection(conn) + except _OversizedControlRPCFrame as e: + self.logger.warning(f"Rejected NCCL control RPC: {e}") + finally: + self.logger.info(f"goodbye {addrinfo}") + class _NcclControlChannel: def __init__( @@ -426,16 +497,21 @@ def __init__( ): self.notif_queue: "queue.Queue[bytes]" = queue.Queue() self._store_values: Dict[str, bytes] = {} + self._store_value_expirations: Dict[str, float] = {} + self._store_value_bytes = 0 self._store_cond = threading.Condition() + self._store_cleanup_stop = threading.Event() self._conn_lock = threading.Lock() self._conns: Dict[tuple[str, str, int], rpyc.Connection] = {} self._server, self.port = self._start_server(host_ip, port_min, port_max) + self._store_cleanup_thread = threading.Thread(target=self._expire_store_values_loop, daemon=True) + self._store_cleanup_thread.start() - def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[ThreadedServer, int]: + def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[_NcclControlServer, int]: last_error = None for cur_port in range(port_min, port_max + 1): try: - server = ThreadedServer( + server = _NcclControlServer( _NcclControlService(self), hostname=host_ip, port=cur_port, @@ -457,7 +533,9 @@ def _start_server(self, host_ip: str, port_min: int, port_max: int) -> tuple[Thr logger.warning(f"Create NCCL RPyC control channel on {host_ip}:{cur_port} failed: {e}") raise RuntimeError(f"can not allocate NCCL control port in [{port_min}, {port_max}]") from last_error - def close(self): + def close(self) -> None: + self._store_cleanup_stop.set() + self._store_cleanup_thread.join(timeout=1) with self._conn_lock: for conn in self._conns.values(): try: @@ -465,21 +543,75 @@ def close(self): except Exception: pass self._conns.clear() + with self._store_cond: + self._store_values.clear() + self._store_value_expirations.clear() + self._store_value_bytes = 0 + self._store_cond.notify_all() self._server.close() return - def add_store_value(self, key: str, value: bytes): + def add_store_value(self, key: str, value: bytes) -> None: + _validate_control_store_key(key) + if not isinstance(value, bytes): + raise TypeError("NCCL control store values must be bytes") + if len(value) > _MAX_CONTROL_STORE_VALUE_BYTES: + raise ValueError(f"NCCL control store value exceeds {_MAX_CONTROL_STORE_VALUE_BYTES} bytes") + with self._store_cond: + now = _control_store_now() + self._expire_store_values(now) + previous_value = self._store_values.get(key) + previous_size = len(previous_value) if previous_value is not None else 0 + new_size = self._store_value_bytes - previous_size + len(value) + if previous_value is None and len(self._store_values) >= _MAX_CONTROL_STORE_ENTRIES: + raise ValueError(f"NCCL control store exceeds {_MAX_CONTROL_STORE_ENTRIES} entries") + if new_size > _MAX_CONTROL_STORE_BYTES: + raise ValueError(f"NCCL control store exceeds {_MAX_CONTROL_STORE_BYTES} bytes") + self._store_values[key] = value + self._store_value_expirations[key] = now + _CONTROL_STORE_VALUE_TTL_SECONDS + self._store_value_bytes = new_size + self._store_cond.notify_all() + return + + def delete_store_value(self, key: str) -> None: + _validate_control_store_key(key) + with self._store_cond: + self._expire_store_values(_control_store_now()) + self._remove_store_value(key) self._store_cond.notify_all() return + def _remove_store_value(self, key: str) -> Optional[bytes]: + value = self._store_values.pop(key, None) + self._store_value_expirations.pop(key, None) + if value is not None: + self._store_value_bytes -= len(value) + return value + + def _expire_store_values(self, now: float) -> bool: + expired_keys = [key for key, expires_at in self._store_value_expirations.items() if expires_at <= now] + for key in expired_keys: + self._remove_store_value(key) + return bool(expired_keys) + + def _expire_store_values_loop(self) -> None: + while not self._store_cleanup_stop.wait(_CONTROL_STORE_CLEANUP_INTERVAL_SECONDS): + with self._store_cond: + if self._expire_store_values(_control_store_now()): + self._store_cond.notify_all() + def wait_store_value(self, key: str, timeout: float = 30.0) -> bytes: + _validate_control_store_key(key) with self._store_cond: + self._expire_store_values(_control_store_now()) ok = self._store_cond.wait_for(lambda: key in self._store_values, timeout=timeout) if not ok: raise TimeoutError(f"wait timeout after {int(timeout * 1000)}ms, key: {key}") - return self._store_values.pop(key) + value = self._remove_store_value(key) + assert value is not None + return value def get_notifs(self) -> List[bytes]: notifs = [] @@ -498,6 +630,10 @@ def send_store_value(self, peer_name: str, host_ip: str, port: int, key: str, va self._call(peer_name, host_ip, port, "set_value", key, value) return + def send_delete_value(self, peer_name: str, host_ip: str, port: int, key: str) -> None: + self._call(peer_name, host_ip, port, "delete_value", key) + return + def _call(self, peer_name: str, host_ip: str, port: int, method: str, *args): conn_key = (peer_name, host_ip, port) with self._conn_lock: @@ -545,6 +681,16 @@ def set(self, key: str, value: bytes): def get(self, key: str) -> bytes: return self.transporter.control_channel.wait_store_value(self._recv_key(key)) + def delete_key(self, key: str) -> None: + remote_metadata = self.transporter._get_remote_metadata(self.remote_agent_name) + self.transporter.control_channel.send_delete_value( + self.remote_agent_name, + remote_metadata.host_ip, + remote_metadata.control_port, + self._send_key(key), + ) + return + def _send_key(self, key: str) -> str: return f"{self.transporter.agent_name}->{self.remote_agent_name}:{key}" diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_kv_move_manager.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_kv_move_manager.py index 00544c9525..e5b585d078 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_kv_move_manager.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_kv_move_manager.py @@ -9,7 +9,6 @@ from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask from lightllm.utils.graceful_utils import graceful_registry from lightllm.server.core.objs import StartArgs -from ..trans_process_obj import KVTransProcess from ..base_kv_move_manager import BaseKVMoveManager from lightllm.utils.error_utils import log_exception from lightllm.utils.envs_utils import get_unique_server_name @@ -81,8 +80,7 @@ def task_dispatcher_loop(self): raise TypeError(f"unsupported prefill kv move task type: {type(task)}") try: - trans_process: KVTransProcess = self.kv_trans_processes[device_id] - trans_process.task_in_queue.put(task) + self.dispatch_trans_task(device_id, task) if ( isinstance(task, PDChunckedTransTask) and self.args.pd_detail_log diff --git a/lightllm/utils/start_utils.py b/lightllm/utils/start_utils.py index 93bc95a436..d3f4046ece 100644 --- a/lightllm/utils/start_utils.py +++ b/lightllm/utils/start_utils.py @@ -66,10 +66,23 @@ def register_process_tree(self, root_process): except (psutil.NoSuchProcess, psutil.ZombieProcess): # A short-lived child may exit while the process tree is scanned. continue + except psutil.AccessDenied: + continue + + try: + process_cmdline = " ".join(process.cmdline()) + except (psutil.NoSuchProcess, psutil.ZombieProcess): + continue + except psutil.AccessDenied: + process_cmdline = process_name if not process_name.startswith("lightllm::"): continue + # KV move managers own recovery for their transfer workers. + if "::prefill_trans:" in process_cmdline or "::decode_trans:" in process_cmdline: + continue + if not process.is_running() or not is_process_active(process.pid): continue diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_control_store.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_control_store.py new file mode 100644 index 0000000000..859a4d65b5 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_control_store.py @@ -0,0 +1,221 @@ +import time +import threading +import socket +import zlib +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as nccl + + +class FakeFrameStream: + def __init__(self, payload, compressed=0): + self.parts = [ + nccl._BoundedControlRPCChannel.FRAME_HEADER.pack(len(payload), compressed), + payload + nccl._BoundedControlRPCChannel.FLUSHER, + ] + self.read_sizes = [] + + def read(self, size): + self.read_sizes.append(size) + return self.parts.pop(0) + + +def make_channel(): + channel = nccl._NcclControlChannel.__new__(nccl._NcclControlChannel) + channel._store_values = {} + channel._store_value_expirations = {} + channel._store_value_bytes = 0 + channel._store_cond = threading.Condition() + channel._store_cleanup_stop = threading.Event() + return channel + + +def test_control_rpc_frame_accepts_bounded_payload(): + stream = FakeFrameStream(zlib.compress(b"metadata"), compressed=1) + channel = nccl._BoundedControlRPCChannel(stream) + + assert channel.recv() == b"metadata" + + +def test_control_rpc_frame_rejects_oversized_wire_payload_before_read(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_RPC_FRAME_BYTES", 4) + stream = FakeFrameStream(b"large") + channel = nccl._BoundedControlRPCChannel(stream) + + with pytest.raises(ValueError, match="frame exceeds 4 bytes"): + channel.recv() + + assert stream.read_sizes == [nccl._BoundedControlRPCChannel.FRAME_HEADER.size] + + +def test_control_rpc_frame_rejects_oversized_decompressed_payload(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_RPC_FRAME_BYTES", 1024) + stream = FakeFrameStream(zlib.compress(b"x" * 10000), compressed=1) + channel = nccl._BoundedControlRPCChannel(stream) + + with pytest.raises(ValueError, match="frame exceeds 1024 bytes"): + channel.recv() + + +def test_control_server_handles_store_and_delete_rpcs(): + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + + control = nccl._NcclControlChannel("127.0.0.1", port, port) + try: + deadline = time.monotonic() + 2 + while not control._server.active and time.monotonic() < deadline: + time.sleep(0.01) + assert control._server.active + + control.send_store_value("peer", "127.0.0.1", port, "key", b"metadata") + assert control._store_values == {"key": b"metadata"} + + control.send_delete_value("peer", "127.0.0.1", port, "key") + assert control._store_values == {} + assert control._store_value_bytes == 0 + finally: + control.close() + + +def test_store_consumption_releases_accounting(): + channel = make_channel() + + channel.add_store_value("agent->peer:key", b"metadata") + + assert channel.wait_store_value("agent->peer:key") == b"metadata" + assert channel._store_values == {} + assert channel._store_value_bytes == 0 + + +def test_store_rejects_oversized_key_and_value_without_mutation(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_KEY_BYTES", 4) + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_VALUE_BYTES", 3) + channel = make_channel() + channel.add_store_value("safe", b"ok") + + with pytest.raises(ValueError, match="key exceeds"): + channel.add_store_value("unsafe", b"x") + with pytest.raises(ValueError, match="value exceeds"): + channel.add_store_value("next", b"long") + + assert channel._store_values == {"safe": b"ok"} + assert channel._store_value_bytes == 2 + + +def test_store_rejects_entry_limit_without_mutation(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_ENTRIES", 1) + channel = make_channel() + channel.add_store_value("first", b"") + + with pytest.raises(ValueError, match="exceeds 1 entries"): + channel.add_store_value("second", b"") + + assert channel._store_values == {"first": b""} + assert channel._store_value_bytes == 0 + + +def test_store_rejects_aggregate_overflow_and_preserves_existing_value(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_BYTES", 4) + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_VALUE_BYTES", 10) + channel = make_channel() + channel.add_store_value("first", b"abc") + channel.add_store_value("second", b"d") + + with pytest.raises(ValueError, match="exceeds 4 bytes"): + channel.add_store_value("first", b"abcd") + + assert channel._store_values == {"first": b"abc", "second": b"d"} + assert channel._store_value_bytes == 4 + + channel.add_store_value("first", b"a") + assert channel._store_value_bytes == 2 + + +def test_concurrent_store_writes_cannot_exceed_aggregate_limit(monkeypatch): + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_BYTES", 10) + channel = make_channel() + barrier = threading.Barrier(10) + + def add_value(index): + barrier.wait() + try: + channel.add_store_value(f"key-{index}", b"xx") + return True + except ValueError: + return False + + threads = [threading.Thread(target=add_value, args=(index,)) for index in range(10)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + + assert all(not thread.is_alive() for thread in threads) + assert len(channel._store_values) == 5 + assert channel._store_value_bytes == 10 + + +def test_expired_values_release_capacity(monkeypatch): + now = [100.0] + monkeypatch.setattr(nccl, "_control_store_now", lambda: now[0]) + monkeypatch.setattr(nccl, "_MAX_CONTROL_STORE_ENTRIES", 1) + monkeypatch.setattr(nccl, "_CONTROL_STORE_CLEANUP_INTERVAL_SECONDS", 0.01) + channel = make_channel() + channel.add_store_value("old", b"old") + + now[0] += nccl._CONTROL_STORE_VALUE_TTL_SECONDS + 1 + cleanup_thread = threading.Thread(target=channel._expire_store_values_loop) + cleanup_thread.start() + deadline = time.monotonic() + 1 + while time.monotonic() < deadline: + with channel._store_cond: + if not channel._store_values: + break + threading.Event().wait(timeout=0.01) + channel._store_cleanup_stop.set() + cleanup_thread.join(timeout=1) + + assert not cleanup_thread.is_alive() + assert channel._store_values == {} + assert channel._store_value_bytes == 0 + + +def test_explicit_delete_releases_capacity(): + channel = make_channel() + channel.add_store_value("key", b"value") + + channel.delete_store_value("key") + + assert channel._store_values == {} + assert channel._store_value_bytes == 0 + + +def test_exposed_control_service_validates_and_deletes_remote_values(): + channel = make_channel() + service = nccl._NcclControlService(channel) + + service.exposed_set_value("agent->peer:key", b"metadata") + assert channel._store_values == {"agent->peer:key": b"metadata"} + + service.exposed_delete_value("agent->peer:key") + assert channel._store_values == {} + assert channel._store_value_bytes == 0 + + +def test_store_delete_key_targets_remote_rendezvous_value(): + channel = Mock() + transporter = SimpleNamespace( + agent_name="local", + control_channel=channel, + _get_remote_metadata=Mock(return_value=SimpleNamespace(host_ip="127.0.0.2", control_port=4321)), + ) + store = nccl._NcclControlStore(transporter, "peer") + + store.delete_key("send_to/peer/0") + + channel.send_delete_value.assert_called_once_with("peer", "127.0.0.2", 4321, "local->peer:send_to/peer/0") diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py index bb732a0eba..d807929144 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py @@ -117,9 +117,10 @@ def test_prefill_kv_move_manager_routes_abort_to_target_device(): target_queue = queue.Queue() manager = PrefillKVMoveManager.__new__(PrefillKVMoveManager) manager.info_queue = _SequenceQueue(abort_req) + manager._trans_process_locks = [threading.Lock(), threading.Lock()] manager.kv_trans_processes = [ - SimpleNamespace(task_in_queue=queue.Queue()), - SimpleNamespace(task_in_queue=target_queue), + SimpleNamespace(task_in_queue=queue.Queue(), is_trans_process_health=lambda: True), + SimpleNamespace(task_in_queue=target_queue, is_trans_process_health=lambda: True), ] with pytest.raises(_StopLoop): diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_trans_process_failure.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_trans_process_failure.py new file mode 100644 index 0000000000..841e7fa061 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_trans_process_failure.py @@ -0,0 +1,228 @@ +import queue +import threading + +import pytest + +from lightllm.server.pd_io_struct import PDAbortReq, PDChunckedTransTask, PDChunckedTransTaskGroup +from lightllm.server.router.model_infer.mode_backend.pd import base_kv_move_manager as manager_module + + +class FakeTransProcess: + def __init__(self, healthy, init_result=True, ready=True): + self.healthy = healthy + self.init_result = init_result + self.ready = ready + self.killed = False + self.process = None + self.device_id = None + self.task_in_queue = queue.Queue() + self.task_out_queue = queue.Queue() + + def is_trans_process_health(self): + return self.healthy + + def killself(self): + self.killed = True + + def init_all(self, **kwargs): + self.init_args = kwargs + self.device_id = kwargs["device_id"] + return self.init_result + + def wait_until_ready(self): + return self.ready + + +def test_unhealthy_transfer_worker_exits_manager_without_signaling_router(monkeypatch): + manager = manager_module.BaseKVMoveManager.__new__(manager_module.BaseKVMoveManager) + manager.node_world_size = 2 + manager.kv_trans_processes = [FakeTransProcess(False), FakeTransProcess(True)] + manager._trans_process_locks = [threading.Lock(), threading.Lock()] + manager._pending_trans_tasks_by_device = [{}, {}] + manager._trans_process_restart_counts = [0, 0] + manager._trans_process_last_restart_at = [None, None] + manager.args = object() + manager.start_trans_process_func = object() + manager.up_status_in_queue = None + + def failed_worker(): + return FakeTransProcess(False, init_result=False) + + monkeypatch.setattr(manager_module, "KVTransProcess", failed_worker) + + def exit_manager(code): + raise SystemExit(code) + + monkeypatch.setattr(manager_module.os, "kill", lambda *args: pytest.fail("manager signaled a process")) + monkeypatch.setattr(manager_module.os, "_exit", exit_manager) + + with pytest.raises(SystemExit) as exc_info: + manager.check_trans_process_loop() + + assert exc_info.value.code == 1 + assert all(process.killed for process in manager.kv_trans_processes) + + +def test_unhealthy_transfer_worker_restarts_and_reopens_return_handler(monkeypatch): + manager = manager_module.BaseKVMoveManager.__new__(manager_module.BaseKVMoveManager) + old_process = FakeTransProcess(False) + new_process = FakeTransProcess(True) + new_process.task_out_queue.put(None) + manager.node_world_size = 1 + manager.kv_trans_processes = [old_process] + manager._trans_process_locks = [threading.Lock()] + manager._pending_trans_tasks_by_device = [{}] + manager._trans_process_restart_counts = [0] + manager._trans_process_last_restart_at = [None] + manager.ret_obj_queue = queue.Queue() + manager.args = object() + manager.start_trans_process_func = object() + manager.up_status_in_queue = None + monkeypatch.setattr(manager_module, "KVTransProcess", lambda: new_process) + + manager._restart_trans_process(0) + + assert manager.kv_trans_processes == [new_process] + assert new_process.init_args["device_id"] == 0 + assert new_process.init_args["manager"] is manager + assert old_process.task_out_queue.get_nowait() is None + assert manager._trans_process_last_restart_at[0] is not None + + +def _make_trans_task( + request_id: int, start_kv_index: int = 0, end_kv_index: int = 1, page_kind: str = "kv" +) -> PDChunckedTransTask: + return PDChunckedTransTask( + request_id=request_id, + start_kv_index=start_kv_index, + end_kv_index=end_kv_index, + time_out_secs=10, + pd_master_node_id=0, + prefill_dp_index=0, + decode_dp_index=0, + src_device_id=0, + dst_device_id=0, + mem_indexes=[0] * (end_kv_index - start_kv_index), + prefill_agent_name=None, + prefill_agent_metadata=None, + prefill_num_pages=None, + prefill_page_reg_desc=None, + decode_agent_name=None, + decode_agent_metadata=None, + decode_num_pages=None, + decode_page_reg_desc=None, + first_gen_token_id=None, + first_gen_token_logprob=None, + page_kind=page_kind, + ) + + +class NoReadQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + def get(self): + raise AssertionError("restart must not read the stopped worker's input queue") + + +def test_restarting_worker_fails_pending_tasks_without_reading_input_queue(monkeypatch): + manager = manager_module.BaseKVMoveManager.__new__(manager_module.BaseKVMoveManager) + old_process = FakeTransProcess(True) + old_process.device_id = 0 + old_process.task_in_queue = NoReadQueue() + new_process = FakeTransProcess(True) + tasks = [ + _make_trans_task(1), + _make_trans_task(2), + _make_trans_task(3), + _make_trans_task(4, start_kv_index=5, end_kv_index=5), + _make_trans_task(4, start_kv_index=5, end_kv_index=5, page_kind="att_state"), + ] + manager.node_world_size = 1 + manager.kv_trans_processes = [old_process] + manager._trans_process_locks = [threading.Lock()] + manager._pending_trans_tasks_by_device = [{}] + manager._trans_process_restart_counts = [0] + manager._trans_process_last_restart_at = [None] + manager.ret_obj_queue = queue.Queue() + manager.args = object() + manager.start_trans_process_func = object() + manager.up_status_in_queue = None + monkeypatch.setattr(manager_module, "KVTransProcess", lambda: new_process) + + manager.dispatch_trans_task(0, tasks[0]) + manager.dispatch_trans_task(0, PDChunckedTransTaskGroup(task_list=tasks[1:])) + manager.dispatch_trans_task(0, PDAbortReq(request_id=5, device_id=0)) + assert len(old_process.task_in_queue.items) == 3 + + old_process.healthy = False + old_process.task_out_queue.put(tasks[0].createRetObj()) + manager._restart_trans_process(0) + retry_task = _make_trans_task(1) + manager.dispatch_trans_task(0, retry_task) + manager.task_ret_handle_loop(old_process) + + returns = [manager.ret_obj_queue.get_nowait() for _ in tasks] + assert len(returns) == len(tasks) + assert sorted(ret.request_id for ret in returns) == [1, 2, 3, 4, 4] + assert all(ret.has_error for ret in returns) + assert all("restarted before reporting task result" in ret.error_info for ret in returns) + assert old_process.task_out_queue.empty() + assert len(manager._pending_trans_tasks_by_device[0][retry_task.createRetObj().get_key()]) == 1 + + new_process.task_out_queue.put(retry_task.createRetObj()) + new_process.task_out_queue.put(None) + retry_return = manager.ret_obj_queue.get(timeout=1) + assert retry_return.request_id == retry_task.request_id + assert not retry_return.has_error + assert not manager._pending_trans_tasks_by_device[0] + + +def test_worker_return_clears_pending_task_before_upload(): + manager = manager_module.BaseKVMoveManager.__new__(manager_module.BaseKVMoveManager) + trans_process = FakeTransProcess(True) + trans_process.device_id = 0 + manager.kv_trans_processes = [trans_process] + manager._trans_process_locks = [threading.Lock()] + manager._pending_trans_tasks_by_device = [{}] + manager.ret_obj_queue = queue.Queue() + task = _make_trans_task(1) + manager.dispatch_trans_task(0, task) + trans_process.task_out_queue.put(task.createRetObj()) + trans_process.task_out_queue.put(None) + manager.task_ret_handle_loop(trans_process) + + ret = manager.ret_obj_queue.get_nowait() + assert ret.request_id == 1 + assert not ret.has_error + assert not manager._pending_trans_tasks_by_device[0] + + +def test_dispatch_waits_for_a_healthy_replacement(): + checked_health = threading.Event() + + class UnhealthyTransProcess(FakeTransProcess): + def is_trans_process_health(self): + checked_health.set() + return False + + manager = manager_module.BaseKVMoveManager.__new__(manager_module.BaseKVMoveManager) + task = _make_trans_task(1) + old_process = UnhealthyTransProcess(False) + new_process = FakeTransProcess(True) + manager.kv_trans_processes = [old_process] + manager._trans_process_locks = [threading.Lock()] + manager._pending_trans_tasks_by_device = [{}] + dispatcher = threading.Thread(target=manager.dispatch_trans_task, args=(0, task)) + dispatcher.start() + + assert checked_health.wait(timeout=1) + with manager._trans_process_locks[0]: + manager.kv_trans_processes[0] = new_process + dispatcher.join(timeout=1) + + assert not dispatcher.is_alive() + assert new_process.task_in_queue.get_nowait() is task diff --git a/unit_tests/utils/test_start_utils.py b/unit_tests/utils/test_start_utils.py index 816e4000f7..c7ee59b87f 100644 --- a/unit_tests/utils/test_start_utils.py +++ b/unit_tests/utils/test_start_utils.py @@ -24,11 +24,14 @@ def wait(self, timeout=None): class FakeProcess: - def __init__(self, pid, running=True, children=None, name="process", exitcode=None, wait_timeout=False): + def __init__( + self, pid, running=True, children=None, name="process", cmdline=None, exitcode=None, wait_timeout=False + ): self.pid = pid self.running = running self._children = children or [] self._name = name + self._cmdline = cmdline or name self.exitcode = exitcode self.wait_timeout = wait_timeout @@ -41,6 +44,9 @@ def children(self, recursive=False): def name(self): return self._name + def cmdline(self): + return [self._cmdline] + def wait(self, timeout=None): if self.wait_timeout: raise start_utils.psutil.TimeoutExpired(timeout, pid=self.pid, name=self._name) @@ -109,20 +115,21 @@ def test_register_process_tree_adds_recursive_descendants(monkeypatch): } -def test_register_process_tree_filters_short_lived_helper_processes(monkeypatch): +def test_register_process_tree_leaves_transfer_workers_to_their_manager(monkeypatch): model_process = FakeProcess(pid=1001, name="lightllm::model_infer") compile_worker = FakeProcess(pid=1002, name="python") - pd_process = FakeProcess(pid=1003, name="lightllm::decode_trans") - router_process = FakeProcess(pid=1000, children=[model_process, compile_worker, pd_process]) + pd_process = FakeProcess(pid=1003, name="lightllm::node", cmdline="lightllm::node::decode_trans:Device0") + manager_process = FakeProcess(pid=1004, name="lightllm::pd_manager") + router_process = FakeProcess(pid=1000, children=[model_process, compile_worker, pd_process, manager_process]) process_manager = start_utils.SubmoduleManager() monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) process_manager.register_process_tree(router_process) - assert process_manager.processes == [model_process, pd_process] + assert process_manager.processes == [model_process, manager_process] assert process_manager.process_names == { model_process: "lightllm::model_infer", - pd_process: "lightllm::decode_trans", + manager_process: "lightllm::pd_manager", }