update

Author: shihaobai

lightllm/server/api_start.py

@@ -66,9 +66,9 @@ def normal_or_p_d_start(args):
     if args.enable_mps:
         from lightllm.utils.device_utils import enable_mps, set_gpu_exclusive_mode
 
-        enable_mps()
         for i in range(args.tp):
             set_gpu_exclusive_mode(gpu_index=i)
+        enable_mps()
 
     if args.run_mode not in ["normal", "prefill", "decode"]:
         return

lightllm/server/visualserver/manager.py

@@ -55,7 +55,7 @@ async def wait_to_model_ready(self):
         for dp_rank_id in range(self.vit_dp):
             tp_ports_each_dp = self.visual_model_rpc_ports[dp_rank_id]
             for tp_rank_id in range(self.vit_tp):
-                device_id = dp_rank_id * self.vit_tp + tp_rank_id
+                device_id = self.args.visual_gpu_ids[dp_rank_id * self.vit_tp + tp_rank_id]
                 rpc_model = await start_model_process(
                     port=tp_ports_each_dp[tp_rank_id], vit_tp=self.vit_tp, device_id=device_id
                 )

lightllm/utils/start_utils.py

@@ -43,9 +43,6 @@ def start_submodule_processes(self, start_funcs=[], start_args=[]):
     def terminate_all_processes(self):
         from lightllm.utils.envs_utils import get_env_start_args
 
-        is_enable_mps = get_env_start_args().enable_mps
-        world_size = get_env_start_args().tp
-
         def kill_recursive(proc):
             try:
                 parent = psutil.Process(proc.pid)
@@ -62,6 +59,10 @@ def kill_recursive(proc):
             if proc.is_alive():
                 kill_recursive(proc)
                 proc.join()
+
+        # recover the gpu compute mode
+        is_enable_mps = get_env_start_args().enable_mps
+        world_size = get_env_start_args().tp
         if is_enable_mps:
             from lightllm.utils.device_utils import stop_mps, set_gpu_default_mode