add visual start

Author: shihaobai

lightllm/server/api_cli.py

@@ -7,7 +7,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--run_mode",
         type=str,
-        choices=["normal", "prefill", "decode", "pd_master", "config_server", "visual_only"],
+        choices=["normal", "prefill", "decode", "pd_master", "config_server", "visual"],
         default="normal",
         help="""set run mode, normal is started for a single server, prefill decode pd_master is for pd split run mode,
                 config_server is for pd split mode used to register pd_master node, and get pd_master node list,
@@ -529,6 +529,12 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=6379,
         help="The port number for the redis service in config_server mode.",
     )
+    parser.add_argument(
+        "--redis_evict_fraction",
+        type=float,
+        default=0.3,
+        help="The evict fraction for the redis service in config_server mode.",
+    )
     parser.add_argument(
         "--start_redis",
         action="store_true",

lightllm/server/api_http.py

@@ -93,7 +93,7 @@ def set_args(self, args):
                 args,
                 metric_port=args.metric_port,
             )
-        elif args.run_mode == "visual_only":
+        elif args.run_mode == "visual":
             self.metric_client = MetricClient(args.metric_port)
         elif args.run_mode == "llm_only":
             init_tokenizer(args)  # for openai api
@@ -160,7 +160,7 @@ def get_model_name():
 @app.get("/health", summary="Check server health")
 @app.head("/health", summary="Check server health")
 async def healthcheck(request: Request):
-    if g_objs.args.run_mode in ["pd_master", "visual_only"]:
+    if g_objs.args.run_mode in ["pd_master", "visual"]:
         return JSONResponse({"message": "Ok"}, status_code=200)
 
     if os.environ.get("DEBUG_HEALTHCHECK_RETURN_FAIL") == "true":
@@ -367,7 +367,7 @@ async def startup_event():
     logger.info("server start up")
     loop = asyncio.get_event_loop()
     g_objs.set_args(get_env_start_args())
-    if g_objs.args.run_mode != "visual_only":
+    if g_objs.args.run_mode != "visual":
         loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
     return

lightllm/server/api_server.py

@@ -5,13 +5,13 @@
     torch.multiprocessing.set_start_method("spawn")  # this code will not be ok for settings to fork to subprocess
     parser = make_argument_parser()
     args = parser.parse_args()
-    from .api_start import pd_master_start, normal_or_p_d_start, visual_only_start, config_server_start, llm_only_start
+    from .api_start import pd_master_start, normal_or_p_d_start, visual_start, config_server_start
 
     if args.run_mode == "pd_master":
         pd_master_start(args)
     elif args.run_mode == "config_server":
         config_server_start(args)
-    elif args.run_mode == "visual_only":
-        visual_only_start(args)
+    elif args.run_mode == "visual":
+        visual_start(args)
     else:
         normal_or_p_d_start(args)

lightllm/server/api_start.py

@@ -57,7 +57,8 @@ def signal_handler(sig, frame):
     signal.signal(signal.SIGINT, signal_handler)
 
     logger.info(f"start process pid {os.getpid()}")
-    logger.info(f"http server pid {http_server_process.pid}")
+    if http_server_process:
+        logger.info(f"http server pid {http_server_process.pid}")
     return
 
 
@@ -72,7 +73,7 @@ def check_and_set_args(args):
 
         enable_mps()
 
-    if args.run_mode not in ["normal", "prefill", "decode", "llm_only", "visual_only"]:
+    if args.run_mode not in ["normal", "prefill", "decode", "llm_only", "visual"]:
         return
 
     assert args.zmq_mode in ["tcp://", "ipc:///tmp/"]
@@ -420,11 +421,9 @@ def pd_master_start(args):
     http_server_process.wait()
 
 
-def visual_only_start(args):
+def visual_start(args):
     check_and_set_args(args)
-    if args.run_mode != "visual_only":
-        return
-    already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.port]
+    already_uesd_ports = args.visual_nccl_ports + [args.nccl_port, args.remote_vit_port]
     can_use_ports = alloc_can_use_network_port(
         num=5 + args.visual_dp * args.visual_tp, used_nccl_ports=already_uesd_ports
     )
@@ -437,6 +436,7 @@ def visual_only_start(args):
         metric_port,
     ) = can_use_ports[0:5]
     can_use_ports = can_use_ports[5:]
+    print(cache_port)
 
     visual_model_tp_ports = []
     for _ in range(args.visual_dp):
@@ -456,13 +456,6 @@ def visual_only_start(args):
 
     set_env_start_args(args)
 
-    process_manager.start_submodule_processes(
-        start_funcs=[
-            start_metric_manager,
-        ],
-        start_args=[(metric_port, args)],
-    )
-
     from .visualserver.manager import start_visual_process
 
     process_manager.start_submodule_processes(
@@ -476,58 +469,18 @@ def visual_only_start(args):
             start_visual_process,
         ],
         start_args=[
-            (args, audio_port, visual_port, cache_port, visual_model_tp_ports),
+            (args, router_port, visual_port, cache_port, visual_model_tp_ports),
         ],
     )
-    if args.enable_multimodal_audio:
-        from .audioserver.manager import start_audio_process
-
-        process_manager.start_submodule_processes(
-            start_funcs=[
-                start_audio_process,
-            ],
-            start_args=[
-                (args, router_port, audio_port, cache_port),
-            ],
-        )
-
-    # 启动 gunicorn
-    command = [
-        "gunicorn",
-        "--workers",
-        f"{args.httpserver_workers}",
-        "--worker-class",
-        "uvicorn.workers.UvicornWorker",
-        "--bind",
-        f"{args.host}:{args.port}",
-        "--log-level",
-        "info",
-        "--access-logfile",
-        "-",
-        "--error-logfile",
-        "-",
-        "lightllm.server.api_http:app",
-        "--timeout",
-        f"{get_lightllm_gunicorn_time_out_seconds()}",
-        "--keep-alive",
-        f"{get_lightllm_gunicorn_keep_alive()}",
-    ]
-
-    # 启动子进程
-    http_server_process = subprocess.Popen(command)
-
-    if "s3://" in args.model_dir:
-        from lightllm.utils.petrel_helper import s3_model_clear
-
-        s3_model_clear(args.model_dir)
-
-    if args.health_monitor:
-        from lightllm.server.health_monitor.manager import start_health_check_process
-
-        process_manager.start_submodule_processes(start_funcs=[start_health_check_process], start_args=[(args,)])
-    setup_signal_handlers(http_server_process, process_manager)
-    http_server_process.wait()
-    return
+    setup_signal_handlers(None, process_manager)
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        logger.info("Received keyboard interrupt, shutting down...")
+        process_manager.terminate_all_processes()
+        logger.info("All processes have been terminated gracefully.")
+        sys.exit(0)
 
 
 def config_server_start(args):

lightllm/server/core/objs/req.py

@@ -161,30 +161,6 @@ def init(
 
         self.post_init()
 
-    def init_visual_only(
-        self,
-        request_id: int,
-    ):
-        # 只是为了有更好的编码辅助类型提示
-        self.index_in_shm_mem: int = self.index_in_shm_mem
-        self.ref_count: int = self.ref_count
-
-        self.request_id = request_id
-        self.group_req_id = convert_sub_id_to_group_id(request_id)
-        self.is_paused = False
-        self.finish_status = FinishStatus()
-        self.is_aborted = False
-        self.router_aborted = False
-        self.shm_infer_released = False
-        self.shm_cur_kv_len = 0
-        self.shm_cur_output_len = 0
-        self.candetoken_out_len = 0
-        self.prompt_cache_len = 0
-        self.finish_token_index = -1
-        self.can_released_mark = False
-
-        self.post_init()
-
     def post_init(self):
         # 子类继承进行一些额外的初始化操作
         pass

lightllm/server/embed_cache/impl/memory_cache_with_redis.py

@@ -21,7 +21,7 @@ def __init__(self, args) -> None:
         self.redis_cache = EmbedRefCountRedis(
             redis_url=redis_url,
             capacity=args.cache_capacity,
-            evict_fraction=args.evict_fraction,
+            evict_fraction=args.redis_evict_fraction,
             image_embed_dir=args.image_embed_dir,
         )
         # 这里之所以把cache * 2是因为，在分离模式下，cache 服务只是为了更新redis状态，以及维护图片cache的 token_id

lightllm/server/embed_cache/manager.py

@@ -55,7 +55,7 @@ def exposed_get_items_embed(self, ids: list[int]) -> list[bool]:
 
 
 def get_cache_manager(args):
-    if args.enable_remote_vit:
+    if args.enable_remote_vit or args.run_mode == "visual":
         return MemoryCacheWithRedis(args)
     else:
         return InMemoryCache(args)

lightllm/server/httpserver/manager.py

@@ -13,7 +13,7 @@
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 from typing import Union, List, Tuple, Dict, Optional
-from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes, VisualOnlyReqIndexes
+from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes
 from fastapi import Request
 from ..tokenizer import get_tokenizer
 from ..pd_io_struct import NodeRole
@@ -143,7 +143,7 @@ async def _alloc_resource(self, items, md5sums, token_nums, datas):
                 uid_list.append(rec["id"])
 
             # If enable the vit/audio-llm disaggregation, no need to cache the data in the memory of the server
-            if self.args.run_mode == "llm_only":
+            if self.enable_remote_vit:
                 return
 
             ready_flags = obtain(self.cache_client.root.get_items_data(uid_list))
@@ -304,41 +304,6 @@ async def _initialize_multimodal_metadata(
             img.uuid = int(md5sum, 16)
             img.token_num = token_num
 
-    # async def get_image_embeding(
-    #     self,
-    #     sampling_params: SamplingParams,
-    #     multimodal_params: MultimodalParams,
-    #     request: Request,
-    #     is_health_req: bool = False,
-    # ) -> Tuple[int, str, dict, FinishStatus]:
-
-    #     request_headers = request.headers if request is not None else {}
-    #     group_request_id = self.alloc_req_id(sampling_params, is_health_req)
-
-    #     try:
-    #         await multimodal_params.verify_and_preload(request)
-    #         image_count = len(multimodal_params.images)
-    #         # 记录请求到达的相关信息
-    #         await self._log_req_header_for_visual_only(request_headers, group_request_id, image_count)
-    #         assert (
-    #             len(multimodal_params.images + multimodal_params.audios) <= self.args.cache_capacity
-    #         ), "too many multimodal items!"
-
-    #         await self._initialize_multimodal_metadata(multimodal_params, sampling_params)
-
-    #         visual_req_status = VisualOnlyReqIndexes(group_req_id=group_request_id, multimodal_params=multimodal_params)
-
-    #         self.send_to_visual.send_pyobj(
-    #             visual_req_status,
-    #             protocol=pickle.HIGHEST_PROTOCOL,
-    #         )
-
-    #     except Exception as e:
-    #         logger.error(f"group_request_id: {group_request_id} has exception {str(e)}")
-    #         await self.abort(group_request_id, multimodal_params)
-    #         raise e
-    #     return
-
     async def generate(
         self,
         prompt: Union[str, List[int]],

lightllm/server/visualserver/manager.py

@@ -10,7 +10,7 @@
 import inspect
 from fastapi import Request
 from ..tokenizer import get_tokenizer
-from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes, VisualOnlyReqIndexes
+from lightllm.server.core.objs.io_objs.group_req import GroupReqIndexes
 from lightllm.server.core.objs import ShmReqManager
 from lightllm.server.core.objs import SamplingParams
 from lightllm.server.core.objs import Req, FinishStatus
@@ -41,9 +41,8 @@ def __init__(
         visual_model_rpc_ports,
     ):
         self.args = args
-        self.remote_vit = args.enable_remote_vit
+        self.remote_vit = args.enable_remote_vit or args.run_mode == "visual"
         self.cache_port = cache_port
-        self.memory_cache = MemoryCacheWithRedis(args)
         self.waiting_reqs: List[GroupReqIndexes] = []
         self.infer_batch_size = args.visual_infer_batch_size
         self.trust_remote_code = args.trust_remote_code
@@ -53,29 +52,33 @@ def __init__(
     def _setup_connections(self):
         context = zmq.Context(2)
         if self.remote_vit:
-            self.recv_from_httpserver.bind(f"tcp://*:{self.args.remote_vit_port}")
+            self.recv_from_remote_llm = context.socket(zmq.PULL)
+            self.recv_from_remote_llm.bind(f"tcp://*:{self.args.remote_vit_port}")
         else:
+            self.recv_from_httpserver = context.socket(zmq.PULL)
             self.recv_from_httpserver.bind(f"{self.args.zmq_mode}127.0.0.1:{self.visual_port}")
             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})
 
     async def wait_to_model_ready(self):
         # 待完成，需要读取config_server来起多个vit
-        self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(self.vit_dp)]
+        visual_dp = self.args.visual_dp
+        visual_tp = self.args.visual_tp
+        self.model_rpcs: List[List[VisualModelRpcClient]] = [[] for _ in range(visual_dp)]
 
-        for dp_rank_id in range(self.args.visual_dp):
+        for dp_rank_id in range(visual_dp):
             tp_ports_each_dp = self.visual_model_rpc_ports[dp_rank_id]
-            for tp_rank_id in range(self.args.visual_tp):
-                device_id = self.args.visual_gpu_ids[dp_rank_id * self.args.visual_tp + tp_rank_id]
+            for tp_rank_id in range(visual_tp):
+                device_id = self.args.visual_gpu_ids[dp_rank_id * visual_tp + tp_rank_id]
                 rpc_model = await start_model_process(
-                    port=tp_ports_each_dp[tp_rank_id], vit_tp=self.args.visual_tp, device_id=device_id
+                    port=tp_ports_each_dp[tp_rank_id], vit_tp=visual_tp, device_id=device_id
                 )
                 self.model_rpcs[dp_rank_id].append(rpc_model)
 
         init_model_ret = []
-        for dp_rank_id in range(self.args.visual_dp):  # async init model process
-            for tp_rank_id in range(self.args.visual_tp):
+        for dp_rank_id in range(visual_dp):  # async init model process
+            for tp_rank_id in range(visual_tp):
                 kvargs = {
                     "tp_rank_id": tp_rank_id,
                     "dp_rank_id": dp_rank_id,

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -40,20 +40,18 @@ def exposed_init_model(self, kvargs):
 
         self.args = get_env_start_args()
 
-        weight_dir = (self.args.model_dir,)
-        cache_port = (self.args.cache_port,)
-        data_type = (self.args.data_type,)
-        quant_type = (self.args.vit_quant_type,)
-        quant_cfg = (self.args.vit_quant_cfg,)
-        max_batch_size = (min(self.args.visual_infer_batch_size // self.args.visual_dp, 1),)
+        weight_dir = self.args.model_dir
+        cache_port = self.args.cache_port
+        data_type = self.args.data_type
+        quant_type = self.args.vit_quant_type
+        quant_cfg = self.args.vit_quant_cfg
+        max_batch_size = min(self.args.visual_infer_batch_size // self.args.visual_dp, 1)
 
         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
-
-        if self.args.run_mode != "visual_only":
-            self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
-        self.visual_only = True if self.args.run_mode == "visual_only" else False
+        print(cache_port)
+        self.cache_client = rpyc.connect("localhost", cache_port, config={"allow_pickle": True})
 
         init_vision_distributed_env(kvargs)
         model_cfg, _ = PretrainedConfig.get_config_dict(weight_dir)