fix(cache): send router notifications on a dedicated thread

lightllm/server/multi_level_kv_cache/manager.py

@@ -16,7 +16,7 @@
 from lightllm.utils.graceful_utils import graceful_registry
 from .cpu_cache_client import CpuKvCacheClient
 from lightllm.utils.log_utils import init_logger
-from lightllm.utils.process_check import start_parent_check_thread
+from lightllm.utils.process_check import install_fatal_thread_excepthook, start_parent_check_thread
 from lightllm.utils.envs_utils import get_unique_server_name
 from lightllm.utils.shm_port_args import get_shm_port_args
 
@@ -30,13 +30,11 @@ def __init__(
     ):
         self.args: StartArgs = args
         ports = get_shm_port_args()
-        context = zmq.Context(2)
-        self.zmq_recv_socket = context.socket(zmq.PULL)
+        self.context = zmq.Context(2)
+        self.zmq_recv_socket = self.context.socket(zmq.PULL)
         self.zmq_recv_socket.bind(f"{args.zmq_mode}127.0.0.1:{ports.multi_level_kv_cache_port}")
 
-        self.send_to_router = context.socket(zmq.PUSH)
-        self.send_to_router.connect(f"{args.zmq_mode}127.0.0.1:{ports.router_port}")
-        logger.info(f"send_to_router sendhwm {self.send_to_router.getsockopt(zmq.SNDHWM)}")
+        self.router_addr = f"{args.zmq_mode}127.0.0.1:{ports.router_port}"
         self.cpu_cache_client = CpuKvCacheClient(only_create_meta_data=False, init_shm_data=True)
         self.shm_req_manager = ShmReqManager()
         self.only_cpu_cache_enable = args.enable_cpu_cache and not args.enable_disk_cache
@@ -45,6 +43,11 @@ def __init__(
         # 控制进行 cpu cache 页面匹配的时间，超过时间则不再匹配，直接转发。
         self.cpu_cache_time_out = 0.5
         self.recv_queue = Queue(maxsize=1024)
+        self.send_to_router_queue = Queue(maxsize=args.running_max_req_size)
+        self.send_to_router_init = concurrent.futures.Future()
+        self.send_to_router_thread = threading.Thread(target=self.send_loop, daemon=True)
+        self.send_to_router_thread.start()
+        self.send_to_router_init.result()
         self.cpu_cache_thread = threading.Thread(target=self.cpu_cache_hanle_loop, daemon=True)
         self.cpu_cache_thread.start()
 
@@ -144,7 +147,7 @@ def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes
         # 超时时，放弃进行 cache page 的匹配。
         current_time = time.time()
         if current_time - start_time >= self.cpu_cache_time_out:
-            self.send_to_router.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+            self.send_to_router_queue.put(group_req_indexes)
             logger.warning(
                 f"cache matching time out {current_time - start_time}s, "
                 f"group_req_id: {group_req_indexes.group_req_id}"
@@ -211,9 +214,23 @@ def _handle_group_req_multi_cache_match(self, group_req_indexes: GroupReqIndexes
         for req in reqs:
             self.shm_req_manager.put_back_req_obj(req)
 
-        self.send_to_router.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+        self.send_to_router_queue.put(group_req_indexes)
         return
 
+    def send_loop(self):
+        try:
+            with self.context.socket(zmq.PUSH) as send_to_router:
+                send_to_router.connect(self.router_addr)
+                logger.debug(f"send_to_router sendhwm {send_to_router.getsockopt(zmq.SNDHWM)}")
+                self.send_to_router_init.set_result(None)
+                while True:
+                    group_req_indexes = self.send_to_router_queue.get()
+                    send_to_router.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+        except Exception as e:
+            if self.send_to_router_init.done():
+                raise
+            self.send_to_router_init.set_exception(e)
+
     def recv_loop(self):
         try:
             recv_max_count = 128
@@ -228,7 +245,7 @@ def recv_loop(self):
                         recv_objs.append(recv_obj)
 
                         start_time = recv_obj.time_mark
-                        logger.info(
+                        logger.debug(
                             f"multi_level_kv_cache recive group req id {recv_obj.group_req_id} "
                             f"cost time {time.time() - start_time} s"
                         )
@@ -254,6 +271,7 @@ def start_multi_level_kv_cache_manager(args, pipe_writer):
     # 注册graceful 退出的处理
     graceful_registry(inspect.currentframe().f_code.co_name)
     setproctitle.setproctitle(f"lightllm::{get_unique_server_name()}::multi_level_kv_cache")
+    install_fatal_thread_excepthook()
     start_parent_check_thread()
 
     try:

unit_tests/server/multi_level_kv_cache/test_manager_socket_ownership.py

@@ -0,0 +1,166 @@
+import pickle
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock
+
+import pytest
+import zmq
+
+from lightllm.server.core.objs.io_objs import GroupReqIndexes
+from lightllm.server.multi_level_kv_cache import manager as manager_impl
+from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager
+
+
+class StopLoop(BaseException):
+    pass
+
+
+@pytest.fixture
+def manager(monkeypatch):
+    stop = object()
+    socket_threads = []
+    sent = Queue()
+    errors = Queue()
+    send_started = threading.Event()
+    send_allowed = threading.Event()
+    send_allowed.set()
+    send_socket = MagicMock()
+    send_socket.__enter__.return_value = send_socket
+    send_socket.connect.side_effect = lambda address: socket_threads.append(threading.get_ident())
+    send_socket.__exit__.side_effect = lambda *args: socket_threads.append(threading.get_ident())
+    recv_socket = Mock()
+    context = Mock()
+
+    def socket(kind):
+        if kind == zmq.PULL:
+            return recv_socket
+        assert kind == zmq.PUSH
+        socket_threads.append(threading.get_ident())
+        return send_socket
+
+    def send(group, protocol):
+        socket_threads.append(threading.get_ident())
+        assert protocol == pickle.HIGHEST_PROTOCOL
+        send_started.set()
+        assert send_allowed.wait(5)
+        sent.put(group)
+
+    class StoppableQueue(Queue):
+        def get(self, *args, **kwargs):
+            item = super().get(*args, **kwargs)
+            if item is stop:
+                raise StopLoop()
+            return item
+
+    send_loop = MultiLevelKVCacheManager.send_loop
+
+    def run_sender(self):
+        try:
+            send_loop(self)
+        except StopLoop:
+            pass
+        except BaseException as exc:
+            errors.put(exc)
+
+    context.socket.side_effect = socket
+    send_socket.send_pyobj.side_effect = send
+    monkeypatch.setattr(manager_impl.zmq, "Context", lambda _: context)
+    monkeypatch.setattr(manager_impl, "Queue", StoppableQueue)
+    monkeypatch.setattr(
+        manager_impl, "get_shm_port_args", lambda: SimpleNamespace(multi_level_kv_cache_port=1, router_port=2)
+    )
+    monkeypatch.setattr(manager_impl, "CpuKvCacheClient", Mock())
+    monkeypatch.setattr(manager_impl, "ShmReqManager", Mock())
+    monkeypatch.setattr(MultiLevelKVCacheManager, "cpu_cache_hanle_loop", lambda self: None)
+    monkeypatch.setattr(MultiLevelKVCacheManager, "send_loop", run_sender)
+    manager = MultiLevelKVCacheManager(
+        SimpleNamespace(zmq_mode="tcp://", enable_cpu_cache=True, enable_disk_cache=False, running_max_req_size=2)
+    )
+    try:
+        yield manager, sent, send_started, send_allowed
+    finally:
+        send_allowed.set()
+        manager.send_to_router_queue.put(stop, timeout=5)
+        manager.send_to_router_thread.join(timeout=5)
+        manager.cpu_cache_thread.join(timeout=5)
+        manager.executor.shutdown()
+        assert not manager.send_to_router_thread.is_alive()
+        assert errors.empty(), list(errors.queue)
+        assert set(socket_threads) == {manager.send_to_router_thread.ident}
+        send_socket.connect.assert_called_once_with("tcp://127.0.0.1:2")
+        send_socket.__exit__.assert_called_once()
+
+
+@pytest.mark.parametrize("expired", [False, True])
+def test_concurrent_workers_leave_socket_sends_to_owner_thread(manager, expired):
+    manager, sent, _, _ = manager
+    groups = [GroupReqIndexes(i, None, [], time.time()) for i in range(100)]
+    start_time = time.time() - 10 if expired else time.time()
+
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        list(executor.map(lambda group: manager._handle_group_req_multi_cache_match(group, start_time), groups))
+
+    assert sorted(sent.get(timeout=5).group_req_id for _ in groups) == list(range(100))
+    assert sent.empty()
+    assert manager.send_to_router_thread.ident != threading.get_ident()
+
+
+def test_send_while_receiver_sleeps(manager, monkeypatch):
+    manager, sent, _, _ = manager
+    group = GroupReqIndexes(0, None, [], time.time())
+    receiver_sleeping = threading.Event()
+    manager.zmq_recv_socket.recv_pyobj.side_effect = zmq.Again()
+
+    def worker():
+        assert receiver_sleeping.wait(5)
+        manager._handle_group_req_multi_cache_match(group, time.time())
+
+    def sleep(interval):
+        assert interval == 0.01
+        receiver_sleeping.set()
+        assert sent.get(timeout=5) is group
+        raise StopLoop()
+
+    monkeypatch.setattr(manager_impl.time, "sleep", sleep)
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(worker)
+        with pytest.raises(StopLoop):
+            manager.recv_loop()
+        future.result(timeout=5)
+
+
+def test_receive_under_send_backpressure(manager, monkeypatch):
+    manager, sent, send_started, send_allowed = manager
+    groups = [GroupReqIndexes(i, None, [], time.time()) for i in range(5)]
+    send_allowed.clear()
+    manager.send_to_router_queue.put(groups[0])
+    assert send_started.wait(5)
+    assert manager.send_to_router_queue.maxsize == manager.args.running_max_req_size
+    manager.send_to_router_queue.put(groups[1])
+    manager.send_to_router_queue.put(groups[2])
+    worker_publishing = threading.Event()
+    put = manager.send_to_router_queue.put
+
+    def publish(group, *args, **kwargs):
+        worker_publishing.set()
+        put(group, *args, **kwargs)
+
+    monkeypatch.setattr(manager.send_to_router_queue, "put", publish)
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(manager._handle_group_req_multi_cache_match, groups[3], time.time() - 10)
+        try:
+            assert worker_publishing.wait(5)
+            assert manager.send_to_router_queue.full()
+            assert not future.done()
+            manager.zmq_recv_socket.recv_pyobj.side_effect = [groups[4], zmq.Again(), StopLoop()]
+            with pytest.raises(StopLoop):
+                manager.recv_loop()
+            assert manager.recv_queue.get_nowait() is groups[4]
+            assert sent.empty()
+        finally:
+            send_allowed.set()
+        future.result(timeout=5)
+    assert [sent.get(timeout=5).group_req_id for _ in range(4)] == list(range(4))

unit_tests/server/multi_level_kv_cache/test_manager_startup.py

@@ -0,0 +1,115 @@
+import errno
+import multiprocessing as mp
+import subprocess
+import sys
+import threading
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+import zmq
+
+from lightllm.server.multi_level_kv_cache import manager as manager_impl
+from lightllm.utils.start_utils import SubmoduleManager
+
+
+def _start_cache(stage, commands, pipe_writer):
+    context = Mock()
+    send_socket = MagicMock()
+    send_socket.__enter__.return_value = send_socket
+    if stage == "connect":
+        send_socket.connect.side_effect = zmq.ZMQError(errno.EMFILE)
+    send_socket.send_pyobj.side_effect = zmq.ZMQError(errno.EIO)
+
+    def socket(kind):
+        if kind == zmq.PUSH:
+            if stage == "socket":
+                raise zmq.ZMQError(errno.EMFILE)
+            return send_socket
+        return Mock()
+
+    def recv_loop(self):
+        command = commands.recv()
+        if command == "send":
+            self.send_to_router_queue.put(None)
+            commands.recv()
+
+    context.socket.side_effect = socket
+    context_factory = manager_impl.zmq.Context if stage == "ready" else lambda _: context
+    thread_start = threading.Thread.start
+
+    def start_thread(thread):
+        wait_for_fault = stage in ("socket", "connect") and thread._target.__name__ == "send_loop"
+        thread_start(thread)
+        if wait_for_fault:
+            thread.join(timeout=5)
+
+    with patch.multiple(
+        manager_impl,
+        CpuKvCacheClient=Mock(),
+        ShmReqManager=Mock(),
+        get_shm_port_args=lambda: SimpleNamespace(multi_level_kv_cache_port=1, router_port=2),
+        get_unique_server_name=lambda: "cache-startup-test",
+        start_parent_check_thread=lambda: None,
+    ), patch.multiple(
+        manager_impl.MultiLevelKVCacheManager,
+        cpu_cache_hanle_loop=lambda self: None,
+        recv_loop=recv_loop,
+    ), patch.object(
+        manager_impl.zmq, "Context", context_factory
+    ), patch.object(
+        threading.Thread, "start", start_thread
+    ):
+        manager_impl.start_multi_level_kv_cache_manager(
+            SimpleNamespace(
+                zmq_mode="inproc://",
+                enable_cpu_cache=True,
+                enable_disk_cache=False,
+                running_max_req_size=2,
+            ),
+            pipe_writer,
+        )
+
+
+def _run_launcher(stage):
+    mp.set_start_method("spawn")
+    commands, controller = mp.Pipe(duplex=False)
+    try:
+        SubmoduleManager().start_submodule_processes(start_funcs=[_start_cache], start_args=[(stage, commands)])
+        assert stage in ("ready", "send")
+        child = mp.active_children()[0]
+        controller.send("send" if stage == "send" else "stop")
+        child.join(timeout=5)
+        assert child.exitcode == (1 if stage == "send" else 0)
+    finally:
+        for child in mp.active_children():
+            child.kill()
+            child.join(timeout=5)
+        commands.close()
+        controller.close()
+
+
+@pytest.mark.parametrize("stage", ["socket", "connect", "ready", "send"])
+def test_sender_startup(stage):
+    result = subprocess.run(
+        [sys.executable, "-m", "unit_tests.server.multi_level_kv_cache.test_manager_startup", stage],
+        cwd=Path(__file__).resolve().parents[3],
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+    output = result.stdout + result.stderr
+    if stage in ("socket", "connect"):
+        assert result.returncode == 1, output
+        assert "init func _start_cache : Too many open files" in output, output
+        assert "init func _start_cache : init ok" not in output, output
+    else:
+        assert result.returncode == 0, output
+        assert "init func _start_cache : init ok" in output, output
+        if stage == "send":
+            assert "fatal background thread failure" in output, output
+
+
+if __name__ == "__main__":
+    _run_launcher(sys.argv[1])