[fix]0915-fix-rpyc-cost

Author: SangChengC

lightllm/models/vit/model.py

@@ -1,4 +1,5 @@
 import os
+import time
 import json
 import torch
 from lightllm.models.vit.layer_infer.pre_layer_infer import ViTPreLayerInfer
@@ -179,10 +180,7 @@ def encode(self, images: List[ImageItem]):
         for i, img in enumerate(images):
             if isinstance(img, ImageItem):
                 uuids.append(img.uuid)
-                if self.remote_vit:
-                    image_data = img._preload_data
-                else:
-                    image_data = read_shm(get_shm_name_data(img.uuid))
+                image_data = read_shm(get_shm_name_data(img.uuid))
                 image_data = Image.open(BytesIO(image_data))
                 t = self.load_image_func(image_data, max_num=img.extra_params["image_patch_max_num"])
                 img_tensors.append(t)

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -81,7 +81,7 @@ def _clear(self, free_max_count: int):
                     # if self.args.run_mode == "visual":
                     #     free_afs(get_shm_name_embed(id), self.args.image_embed_dir)
                     # elif not self.args.enable_remote_vit:
-                    if not self.args.run_mode == "visual":
+                    if not self.args.enable_remote_vit and self.args.run_mode != "visual":
                         free_shm(get_shm_name_embed(id))
                 del self._md5_to_record[record.md5sum]
                 del self._records[id]

lightllm/server/httpserver/manager.py

@@ -4,6 +4,7 @@
 import asyncio
 import uvloop
 import rpyc
+import socket
 import time
 import copy
 import hashlib
@@ -84,6 +85,7 @@ def __init__(
         self.enable_multimodal = enable_multimodal
         if self.enable_multimodal:
             self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
+            self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
             # 初始化VIT连接管理器
             from lightllm.server.visualserver.vit_connect import VITConnectionManager
 

lightllm/server/visualserver/manager.py

@@ -4,13 +4,15 @@
 import asyncio
 import uvloop
 import rpyc
+import socket
 import pickle
 import hashlib
 import datetime
 import inspect
 from fastapi import Request
 from ..tokenizer import get_tokenizer
 from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes
+from lightllm.server.embed_cache.utils import get_shm_name_data, create_shm
 from lightllm.server.core.objs import ShmReqManager
 from lightllm.server.core.objs import SamplingParams
 from lightllm.server.core.objs import Req, FinishStatus
@@ -63,6 +65,7 @@ def _setup_connections(self):
             self.send_to_next_module = context.socket(zmq.PUSH)  # router or audio server (if --enable_multimodal_audio)
             self.send_to_next_module.connect(f"{self.args.zmq_mode}127.0.0.1:{self.next_module_port}")
         self.cache_client = rpyc.connect("localhost", self.cache_port, config={"allow_pickle": True})
+        self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
     async def wait_to_model_ready(self):
         visual_dp = self.args.visual_dp
@@ -100,7 +103,6 @@ async def infer_imgs(self, images: List[ImageItem]):
                 for vit_tp_rank in range(self.args.visual_tp):
                     task = asyncio.create_task(self.model_rpcs[vit_dp_rank][vit_tp_rank].encode(assigned_images))
                     tasks.append(task)
-
         await asyncio.gather(*tasks)
         return
 
@@ -162,19 +164,34 @@ async def _recv_reqs(self):
             # ]
             logger.info(f"Receive req {recv_req.group_req_id}, image_count:{len(recv_req.multimodal_params.images)}")
             uuids = [img.uuid for img in recv_req.multimodal_params.images]
-            already_embed = self.cache_client.root.get_items_embed(uuids)
+            already_embed = await asyncio.to_thread(self.cache_client.root.get_items_embed, uuids)
             if all(already_embed):
                 return None
+
+            uuids = []
             token_nums = []
+            datas = []
             for img, embed in zip(recv_req.multimodal_params.images, already_embed):
                 if not embed:
                     uuids.append(img.uuid)
                     token_nums.append(img.token_num)
+                    datas.append(img._preload_data)
+                    img.free()
             while True:
-                records = self.cache_client.root.alloc(uuids, token_nums)
+                records = await asyncio.to_thread(self.cache_client.root.alloc, uuids, token_nums)
                 if records is not None:
                     break
-                await asyncio.sleep(0.1)
+                await asyncio.sleep(0.01)
+            ready_flags = obtain(self.cache_client.root.get_items_data(uuids))
+            update_data_ids = []
+
+            for uid, ready, data in zip(uuids, ready_flags, datas):
+                if not ready:
+                    create_shm(get_shm_name_data(uid), data)
+                    update_data_ids.append(uid)
+
+            if update_data_ids:
+                await asyncio.to_thread(self.cache_client.root.set_items_data, update_data_ids)
             return recv_req
         else:
             return self.vit_receiver.recv_pyobj(zmq.NOBLOCK)
@@ -193,7 +210,8 @@ async def loop_for_netio_req(self):
                         self.waiting_reqs.append(recv_req)
                     else:
                         assert False, f"Error Req Inf {recv_req}"
-                self.visual_recv_max_count = min(self.visual_recv_max_count * 1.3, 256)
+                    await asyncio.sleep(0)
+                self.visual_recv_max_count = min(int(self.visual_recv_max_count * 1.3), 256)
             except zmq.ZMQError:
                 # 当队列已经开始清空的时候，将一次接受数量下调
                 self.visual_recv_max_count = 64
@@ -217,21 +235,11 @@ async def loop_for_fwd_visual_only(self):
                         images_need_infer.append(img)
 
                         if len(images_need_infer) == self.infer_batch_size:
-                            _t0 = time.perf_counter()
                             await self.infer_imgs(images_need_infer)
-                            logger.info(
-                                f"[visual] batch infer complete, image_count: {len(images_need_infer)}, "
-                                f"elapsed_time {(time.perf_counter()-_t0) * 1000}ms"
-                            )
                             images_need_infer = []
 
                     if len(images_need_infer) > 0:
-                        _t1 = time.perf_counter()
                         await self.infer_imgs(images_need_infer)
-                        logger.info(
-                            f"[visual] batch infer complete, image_count:{len(images_need_infer)}, "
-                            f"elapsed_time {(time.perf_counter()-_t1) * 1000}ms"
-                        )
                         images_need_infer = []
                     # 在这里release这个image，ref-1
                     logger.info(f"req-id {visual_req.group_req_id} has been release ok")

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -2,6 +2,7 @@
 import numpy as np
 import rpyc
 import torch
+import time
 import inspect
 from datetime import timedelta
 from typing import Dict, List, Tuple
@@ -30,6 +31,11 @@
 from lightllm.utils.dist_utils import init_vision_distributed_env
 from lightllm.utils.graceful_utils import graceful_registry
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.log_utils import init_logger
+import pickle
+import socket
+
+logger = init_logger(__name__)
 
 
 class VisualModelRpcServer(rpyc.Service):
@@ -48,10 +54,12 @@ def exposed_init_model(self, kvargs):
         max_batch_size = min(self.args.visual_infer_batch_size // self.args.visual_dp, 1)
         remote_vit = True if self.args.run_mode == "visual" else False
 
+        self.image_embed_dir = self.args.image_embed_dir
         self.dp_rank_id = kvargs["dp_rank_id"]
         self.tp_rank_id = kvargs["tp_rank_id"]
         kvargs["vit_rank_id"] = self.dp_rank_id * self.args.visual_tp + self.tp_rank_id
         self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
+        self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
         init_vision_distributed_env(kvargs)
         model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir)
@@ -109,19 +117,18 @@ def forward(self, images: List[ImageItem]):
     def exposed_encode(self, images: List[ImageItem]):
         images = obtain(images)
         all_img_embeds, uuids, valid_ids = self.forward(images)
-        all_img_embeds = all_img_embeds.to(torch.device("cpu"))
-
+        all_img_embeds = all_img_embeds.to(torch.device("cpu"), non_blocking=True)
         if self.tp_rank_id == 0:
             # ready_flags = obtain(self.cache_client.root.get_items_embed(uuids))
             ids_to_set = []
-            for i, img in enumerate(images):
+            for i in range(len(images)):
                 # if ready:
                 #     continue
                 uid = uuids[i]
                 start, end = valid_ids[i]
                 cur_embed_bytes = tensor2bytes(all_img_embeds[start:end])
                 if self.args.run_mode == "visual":
-                    create_afs(get_shm_name_embed(uid), cur_embed_bytes, self.args.image_embed_dir)
+                    create_afs(get_shm_name_embed(uid), cur_embed_bytes, self.image_embed_dir)
                 else:
                     create_shm(get_shm_name_embed(uid), cur_embed_bytes)
                 ids_to_set.append(uid)
@@ -131,11 +138,13 @@ def exposed_encode(self, images: List[ImageItem]):
 
 
 class VisualModelRpcClient:
-    def __init__(self, model_rpc, vit_tp, rpc_server_process=None):
-        self.model: VisualModelRpcServer = model_rpc
+    def __init__(self, conn, vit_tp, rpc_server_process=None):
+        self.conn = conn
+        self.model: VisualModelRpcServer = conn.root
         self.vit_tp = vit_tp
         self.rpc_server_process = rpc_server_process
         self.use_rpc = True
+        self._bg = rpyc.BgServingThread(self.conn)
         if self.use_rpc:
 
             def async_wrap(f):
@@ -176,7 +185,13 @@ def _init_env(port, device_id):
     # 注册graceful 退出的处理
     graceful_registry(inspect.currentframe().f_code.co_name)
 
-    t = ThreadedServer(VisualModelRpcServer(), port=port, protocol_config={"allow_pickle": True})
+    auth = lambda sock: (sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) or (sock, None))
+    t = ThreadedServer(
+        VisualModelRpcServer(),
+        port=port,
+        protocol_config={"allow_pickle": True},
+        authenticator=auth,
+    )
     t.start()
     return
 
@@ -197,6 +212,7 @@ async def start_model_process(port, vit_tp, device_id):
     while repeat_count < 20:
         try:
             con = rpyc.connect("localhost", port, config={"allow_pickle": True})
+            con._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
             break
         except BaseException:
             await asyncio.sleep(1)
@@ -205,4 +221,4 @@ async def start_model_process(port, vit_tp, device_id):
         raise Exception("init rpc env error!")
 
     assert proc.is_alive()
-    return VisualModelRpcClient(con.root, vit_tp, rpc_server_process=proc)
+    return VisualModelRpcClient(con, vit_tp, rpc_server_process=proc)