From 953ce2d614a5d715593632785d0a4018f86c19bf Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 30 Jul 2026 14:50:16 +0800 Subject: [PATCH] fix: clean up shared memory on process exit --- lightllm/server/api_start.py | 11 ++- lightllm/utils/auto_shm_cleanup.py | 49 +++++------ lightllm/utils/kv_cache_utils.py | 7 +- lightllm/utils/start_utils.py | 102 ++++++++++++++-------- unit_tests/utils/test_auto_shm_cleanup.py | 71 +++++++++++++++ 5 files changed, 172 insertions(+), 68 deletions(-) create mode 100644 unit_tests/utils/test_auto_shm_cleanup.py diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index f74a7210ad..873ecb6e4e 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -29,6 +29,10 @@ auto_set_fused_shared_experts, ) from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args +from lightllm.utils.auto_shm_cleanup import ( + mark_registered_sysv_shm_for_deletion, + register_sysv_shm_for_cleanup, +) logger = init_logger(__name__) @@ -40,7 +44,7 @@ def signal_handler(sig, frame): if http_server_process: kill_recursive(http_server_process) - process_manager.terminate_all_processes() + process_manager.terminate_all_processes(graceful=False) logger.info("All processes have been forcefully terminated.") sys.exit(0) elif sig == signal.SIGTERM: @@ -140,10 +144,13 @@ def _launch_subprocesses(args: StartArgs): if args.enable_cpu_cache: # 生成一个用于创建cpu kv cache的共享内存id。 args.cpu_kv_cache_shm_id = uuid.uuid1().int % 123456789 + register_sysv_shm_for_cleanup(args.cpu_kv_cache_shm_id) if args.enable_multimodal: args.multi_modal_cache_shm_id = uuid.uuid1().int % 123456789 + register_sysv_shm_for_cleanup(args.multi_modal_cache_shm_id) + setup_signal_handlers(None, process_manager) # 调度参数的自动设置, 人工设置则听人工的 if args.router_token_ratio is None: if args.run_mode in ["normal"]: @@ -470,6 +477,8 @@ def _launch_subprocesses(args: StartArgs): ], ) + mark_registered_sysv_shm_for_deletion() + return process_manager diff --git a/lightllm/utils/auto_shm_cleanup.py b/lightllm/utils/auto_shm_cleanup.py index 2417fef085..0c8b2e9736 100644 --- a/lightllm/utils/auto_shm_cleanup.py +++ b/lightllm/utils/auto_shm_cleanup.py @@ -1,11 +1,9 @@ -import os import ctypes import atexit import signal import threading -import psutil from multiprocessing import shared_memory -from typing import Set, Optional +from typing import Optional from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -21,8 +19,7 @@ def __init__(self): self.libc = None self._init_libc() # System V - self.registered_shm_keys = [] - self.registered_shm_ids = [] + self.registered_sysv_shms = {} # POSIX self.registered_posix_shm_names = [] self.signal_handlers_registered = False @@ -52,33 +49,26 @@ def register_signal_handlers(self): def _signal_cleanup_handler(self, signum, frame): self._cleanup() - parent = psutil.Process(os.getpid()) - # 递归拿到所有子进程并终止 - for ch in parent.children(recursive=True): - ch.kill() + raise SystemExit(128 + signum) - def _cleanup(self): - """清理：System V 执行 IPC_RMID，POSIX 执行 unlink。""" - removed_sysv = 0 + def mark_registered_sysv_shm_for_deletion(self): IPC_RMID = 0 - for shmid in self.registered_shm_ids: + marked_sysv = 0 + for key, shmid in self.registered_sysv_shms.items(): try: + if shmid is None: + shmid = self.libc.shmget(key, 0, 0) if self.libc.shmctl(shmid, IPC_RMID, None) == 0: - removed_sysv += 1 - except Exception as e: - logger.warning(f"cleanup: shmid {shmid} clean failed, reason: {e}") - pass - for key in self.registered_shm_keys: - shmid = self.libc.shmget(key, 0, 0) - try: - if shmid >= 0 and self.libc.shmctl(shmid, IPC_RMID, None) == 0: - removed_sysv += 1 + marked_sysv += 1 except Exception as e: logger.warning(f"cleanup: shmid {shmid} clean failed, reason: {e}") - pass - if removed_sysv: - logger.info(f"cleanup: removed {removed_sysv} System V shm segments") + if marked_sysv: + logger.info(f"cleanup: marked {marked_sysv} System V shm segments for deletion") + return marked_sysv + def _cleanup(self): + """清理：System V 执行 IPC_RMID，POSIX 执行 unlink。""" + self.mark_registered_sysv_shm_for_deletion() removed_posix = 0 for name in self.registered_posix_shm_names: try: @@ -103,9 +93,8 @@ def _cleanup(self): def register_sysv_shm(self, key: int, shmid: Optional[int] = None): """注册 System V 共享内存。""" - self.registered_shm_keys.append(key) - if shmid is not None: - self.registered_shm_ids.append(shmid) + if key not in self.registered_sysv_shms or shmid is not None: + self.registered_sysv_shms[key] = shmid return def register_posix_shm(self, name: str): @@ -131,5 +120,9 @@ def register_sysv_shm_for_cleanup(key: int, shmid: Optional[int] = None): get_auto_cleanup().register_sysv_shm(key, shmid) +def mark_registered_sysv_shm_for_deletion(): + return get_auto_cleanup().mark_registered_sysv_shm_for_deletion() + + def register_posix_shm_for_cleanup(name: str): get_auto_cleanup().register_posix_shm(name) diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 06c4211ebf..7763b5ebb0 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -1,6 +1,7 @@ import torch import ctypes import dataclasses +import errno import os import xxhash import threading @@ -195,7 +196,9 @@ def _get_default_hugepage_size() -> int: pass return 2 * 1024 * 1024 # fallback 2MB - shmflg = 0o666 | 0o1000 # 权限和 IPC_CREAT 标志 + IPC_CREAT = 0o1000 + IPC_EXCL = 0o2000 + shmflg = 0o666 | IPC_CREAT | IPC_EXCL if use_hugetlb: # 向上对齐到大页大小 huge_sz = _get_default_hugepage_size() @@ -213,6 +216,8 @@ def _get_default_hugepage_size() -> int: hugepages_num = (size_to_alloc + 1024 * 1024 * 1024 - 1) // (1024 * 1024 * 1024) if shmid < 0: err = ctypes.get_errno() + if err == errno.EEXIST: + raise RuntimeError(f"System V shared memory key {key} already exists") if use_hugetlb: raise Exception( f"shmget with SHM_HUGETLB failed (errno={err}). Falling back to regular pages." diff --git a/lightllm/utils/start_utils.py b/lightllm/utils/start_utils.py index 372b7e1cfa..07caf359d8 100644 --- a/lightllm/utils/start_utils.py +++ b/lightllm/utils/start_utils.py @@ -7,6 +7,8 @@ class SubmoduleManager: + _TERMINATE_TIMEOUT = 5 + def __init__(self): self.processes = [] @@ -15,58 +17,82 @@ def start_submodule_processes(self, start_funcs=[], start_args=[]): pipe_readers = [] processes = [] - for start_func, start_arg in zip(start_funcs, start_args): - pipe_reader, pipe_writer = mp.Pipe(duplex=False) - process = mp.Process( - target=start_func, - args=start_arg + (pipe_writer,), - ) - process.start() - pipe_readers.append(pipe_reader) - processes.append(process) - - # Wait for all processes to initialize - for index, pipe_reader in enumerate(pipe_readers): - init_state = pipe_reader.recv() - if init_state != "init ok": - logger.error(f"init func {start_funcs[index].__name__} : {str(init_state)}") - for proc in processes: - proc.kill() - sys.exit(1) - else: + try: + for start_func, start_arg in zip(start_funcs, start_args): + pipe_reader, pipe_writer = mp.Pipe(duplex=False) + process = mp.Process( + target=start_func, + args=start_arg + (pipe_writer,), + ) + process.start() + pipe_readers.append(pipe_reader) + processes.append(process) + self.processes.append(process) + + for index, pipe_reader in enumerate(pipe_readers): + init_state = pipe_reader.recv() + if init_state != "init ok": + raise RuntimeError(f"init func {start_funcs[index].__name__} : {str(init_state)}") logger.info(f"init func {start_funcs[index].__name__} : {str(init_state)}") - assert all([proc.is_alive() for proc in processes]) - self.processes.extend(processes) + assert all([proc.is_alive() for proc in processes]) + except BaseException: + self.terminate_all_processes(graceful=False) + raise return - def terminate_all_processes(self): - from lightllm.utils.envs_utils import get_env_start_args - - def kill_recursive(proc): + def _terminate_processes(self, processes, graceful): + process_by_pid = {} + for proc in processes: + if proc.pid is None: + continue try: parent = psutil.Process(proc.pid) - children = parent.children(recursive=True) - for child in children: - logger.info(f"Killing child process {child.pid}") - child.kill() - logger.info(f"Killing parent process {proc.pid}") - parent.kill() + process_by_pid[parent.pid] = parent + for child in parent.children(recursive=True): + process_by_pid[child.pid] = child except psutil.NoSuchProcess: - logger.warning(f"Process {proc.pid} does not exist.") + continue + + process_tree = list(process_by_pid.values()) + for process in reversed(process_tree): + try: + if graceful: + process.terminate() + else: + process.kill() + except psutil.NoSuchProcess: + pass + + if graceful: + _, alive = psutil.wait_procs(process_tree, timeout=self._TERMINATE_TIMEOUT) + for process in alive: + try: + process.kill() + except psutil.NoSuchProcess: + pass + psutil.wait_procs(alive, timeout=self._TERMINATE_TIMEOUT) + + for proc in processes: + if proc.pid is not None: + proc.join(timeout=1) + + def terminate_all_processes(self, graceful=True): + from lightllm.utils.envs_utils import get_env_start_args - for proc in self.processes: - if proc.is_alive(): - kill_recursive(proc) - proc.join() + self._terminate_processes(self.processes, graceful) + self.processes.clear() # recover the gpu compute mode - is_enable_mps = get_env_start_args().enable_mps + try: + is_enable_mps = get_env_start_args().enable_mps + except (AttributeError, KeyError): + is_enable_mps = False if is_enable_mps: from lightllm.utils.device_utils import stop_mps stop_mps() - logger.info("All processes terminated gracefully.") + logger.info("All processes terminated.") def start_submodule_processes(start_funcs=[], start_args=[]): diff --git a/unit_tests/utils/test_auto_shm_cleanup.py b/unit_tests/utils/test_auto_shm_cleanup.py new file mode 100644 index 0000000000..c178e4f1da --- /dev/null +++ b/unit_tests/utils/test_auto_shm_cleanup.py @@ -0,0 +1,71 @@ +import pytest + +from lightllm.utils.auto_shm_cleanup import AutoShmCleanup +from lightllm.utils.start_utils import SubmoduleManager + + +class FakeLibc: + def __init__(self): + self.shmget_calls = [] + self.shmctl_calls = [] + + def shmget(self, key, size, flags): + self.shmget_calls.append((key, size, flags)) + return 5678 + + def shmctl(self, shmid, command, buffer): + self.shmctl_calls.append((shmid, command, buffer)) + return 0 + + +@pytest.mark.parametrize("shmid", [None, 5678]) +def test_mark_sysv_shm(shmid): + cleanup = AutoShmCleanup.__new__(AutoShmCleanup) + cleanup.libc = FakeLibc() + cleanup.registered_sysv_shms = {} + cleanup.register_sysv_shm(1234, shmid) + + assert cleanup.mark_registered_sysv_shm_for_deletion() == 1 + assert cleanup.libc.shmget_calls == ([(1234, 0, 0)] if shmid is None else []) + assert cleanup.libc.shmctl_calls == [(5678, 0, None)] + + +def test_processes_get_term_before_kill(monkeypatch): + events = [] + + class PsutilProcess: + def __init__(self, pid, children=None): + self.pid = pid + self._children = children or [] + + def children(self, recursive): + return self._children + + def terminate(self): + events.append(("term", self.pid)) + + def kill(self): + events.append(("kill", self.pid)) + + child = PsutilProcess(11) + parent = PsutilProcess(10, [child]) + processes = {10: parent, 11: child} + wait_count = 0 + + def wait_procs(procs, timeout): + nonlocal wait_count + wait_count += 1 + return ([], [parent]) if wait_count == 1 else (procs, []) + + monkeypatch.setattr("lightllm.utils.start_utils.psutil.Process", lambda pid: processes[pid]) + monkeypatch.setattr("lightllm.utils.start_utils.psutil.wait_procs", wait_procs) + + class Process: + pid = 10 + + def join(self, timeout): + events.append(("join", timeout)) + + SubmoduleManager()._terminate_processes([Process()], graceful=True) + + assert events == [("term", 11), ("term", 10), ("kill", 10), ("join", 1)]