fix bug for pd reformater.

Author: hiworldwzj

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_impl.py

@@ -35,7 +35,7 @@ def init_custom(self):
             ranks.append(i + self.global_dp_rank * self.dp_world_size)
 
         self.lock_nccl_group = dist.new_group(ranks=ranks, backend="gloo")
-        logger.info(f"lock_nccl_group ranks {self.lock_nccl_group.get_rank()}")
+        logger.info(f"lock_nccl_group ranks {dist.get_rank(self.lock_nccl_group)}")
 
         from .decode_infer_rpyc import PDDecodeInferRpcServer
 

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_kv_move_manager.py

@@ -416,8 +416,13 @@ def exposed_request_data_transfer(self, tasks: List[KVMoveTask]) -> List[Optiona
 
             trans_obj = self.get_trans_obj(tasks[0])
             assert trans_obj is not None
-
-            id_to_test_range = {task.group_request_id: random.shuffle(list(range(self.dp_size_in_node))) for task in tasks}
+            
+            id_to_test_range = {}
+            for task in tasks:
+                test_dp_indexes = list(range(self.dp_size_in_node))
+                random.shuffle(test_dp_indexes)
+                id_to_test_range[task.group_request_id] = test_dp_indexes
+                                
             id_has_result = {}
             for test_index in range(self.dp_size_in_node):
                 dp_tasks = [[] for _ in range(self.dp_size_in_node)]

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_impl.py

@@ -35,7 +35,7 @@ def init_custom(self):
             ranks.append(i + self.global_dp_rank * self.dp_world_size)
 
         self.lock_nccl_group = dist.new_group(ranks=ranks, backend="gloo")
-        logger.info(f"lock_nccl_group ranks {self.lock_nccl_group.get_rank()}")
+        logger.info(f"lock_nccl_group ranks {dist.get_rank(self.lock_nccl_group)}")
 
         from .prefill_infer_rpyc import PDPrefillInferRpcServer