Fix the wrong KVAggregator finished count cause dead loop, adopt vllm changes (#1009)

Author: mrjunwan-lang

tpu_inference/distributed/tpu_connector.py

@@ -190,6 +190,10 @@ def request_finished(
         assert self.connector_scheduler is not None
         return self.connector_scheduler.request_finished(request, block_ids)
 
+    def get_finished_count(self) -> int:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_finished_count()
+
     ############################################################
     # Worker Side Methods
     ############################################################
@@ -280,7 +284,7 @@ def get_num_new_matched_tokens(
                   because TPU pulls KV cache in a blocking way.
 
         """
-        if self.is_producer:
+        if self.is_producer or not request.kv_transfer_params:
             return 0, False
 
         assert num_computed_tokens % self.block_size == 0
@@ -345,7 +349,9 @@ def update_state_after_alloc(self, request: "Request",
                 remote_host=params["remote_host"],
                 remote_port=params["remote_port"],
             )
-        logger.info(f"Scheduler -->  reqs_to_load={self.reqs_to_load}")
+        logger.info(
+            f"TPUConnector Scheduler update_state_after_alloc -->  reqs_to_load={self.reqs_to_load}"
+        )
 
     def build_connector_meta(self) -> TPUConnectorMetadata:
         """
@@ -365,6 +371,12 @@ def build_connector_meta(self) -> TPUConnectorMetadata:
 
         return meta
 
+    def get_finished_count(self) -> int:
+        """
+        Return how many workers need pull the kv cache and report back.
+        """
+        return len(self.kv_ip) if isinstance(self.kv_ip, list) else 1
+
     def request_finished(
         self,
         request: "Request",

tpu_inference/executors/ray_distributed_executor.py

@@ -6,13 +6,13 @@
 import vllm.envs as envs
 from ray.util.placement_group import PlacementGroup
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
 from vllm.multimodal.inputs import MultiModalKwargs
 from vllm.platforms import current_platform
 from vllm.ray.ray_env import get_env_vars_to_copy
 from vllm.sequence import VLLM_TOKEN_ID_ARRAY_TYPE
 from vllm.utils.network_utils import (get_distributed_init_method, get_ip,
                                       get_open_port)
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.executor.ray_distributed_executor import \
     RayDistributedExecutor as RayDistributedExecutorV1
 from vllm.v1.executor.ray_executor import RayWorkerMetaData
@@ -101,10 +101,10 @@ def _init_executor(self) -> None:
 
         self.pp_locks: Optional[List[asyncio.Lock]] = None
 
+        self.scheduler_output: SchedulerOutput | None = None
+
         # KV connector setup
         self.has_connector = self.vllm_config.kv_transfer_config is not None
-        self.kv_output_aggregator = KVOutputAggregator(
-            self.parallel_config.world_size)
         if self.has_connector:
             ip_port = self.collective_rpc("get_node_kv_ip_port")
             for item in ip_port:
@@ -229,7 +229,7 @@ def _init_workers_ray(self, placement_group: "PlacementGroup",
         for each, ip in zip(worker_metadata, worker_ips):
             each.ip = ip
 
-        logger.debug("workers: %s", worker_metadata)
+        logger.debug(f"Initialized worker_metadata: {worker_metadata}")
 
         ip_counts: Dict[str, int] = {}
         for ip in worker_ips:
@@ -256,6 +256,9 @@ def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
         start_rank = 0
         for i, item in enumerate(sorted_worker_metadata):
             item.adjusted_rank = i + start_rank
+        logger.info(
+            f"Initialized sorted worker_metadata: {sorted_worker_metadata}")
+
         self.workers = [item.worker for item in sorted_worker_metadata]
         rerank_mapping = {
             item.created_rank: item.adjusted_rank
@@ -353,3 +356,8 @@ def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
                     assert len(self.pp_tp_workers[pp_rank]) == tp_rank
                     assert pp_rank < len(self.pp_tp_workers)
                     self.pp_tp_workers[pp_rank].append(self.workers[rank])
+
+    # Ray executor do not need handshake metadata
+    # as we pass the kv_parameters through proxy server
+    def get_kv_connector_handshake_metadata(self) -> None:
+        pass