[BugFix] Make PD work with Ray

vllm/executor/ray_utils.py

@@ -17,6 +17,7 @@
 from vllm.utils import get_ip
 from vllm.worker.worker_base import WorkerWrapperBase
 
+
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
     from vllm.v1.outputs import ModelRunnerOutput

vllm/v1/executor/ray_distributed_executor.py

@@ -1,14 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from concurrent.futures import Future
-from typing import Union
+from collections import defaultdict
+from concurrent.futures import CancelledError, Future
+from typing import Optional, Sequence, Union, cast
 
 from vllm.executor.ray_distributed_executor import (  # noqa
     RayDistributedExecutor as RayDistributedExecutorV0)
 from vllm.v1.executor.abstract import Executor
 from vllm.v1.outputs import ModelRunnerOutput
+from vllm.logger import init_logger, logger
 
+init_logger("vllm")
 
 class FutureWrapper(Future):
     """A wrapper around a Ray output reference to meet the interface
@@ -28,6 +31,19 @@ def result(self, timeout=None):
 class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
     """Ray distributed executor using Ray Compiled Graphs."""
 
+    def _init_executor(self) -> None:
+        super()._init_executor()
+
+        # KV connector setup
+        self.has_connector = self.vllm_config.kv_transfer_config is not None
+
+        # Complete transfer tracker. Used by to track finished requests
+        # [req_id -> n_finished_workers]
+        self._recv_remaining_count = defaultdict[str,
+                                                 int](lambda: self.parallel_config.world_size)
+        self._send_remaining_count = defaultdict[str,
+                                                 int](lambda: self.parallel_config.world_size)
+
     @property
     def max_concurrent_batches(self) -> int:
         """Ray distributed executor supports pipeline parallelism,
@@ -37,6 +53,80 @@ def max_concurrent_batches(self) -> int:
             return 2
         return self.parallel_config.pipeline_parallel_size
 
+    def _aggregate_workers_output(
+            self, outputs: list[ModelRunnerOutput]) -> ModelRunnerOutput:
+        # aggregate finished_sending, finished_recving from all workers
+
+        finished_sending = set[str]()
+        finished_recving = set[str]()
+        for output in outputs:
+            # update finished_sending
+            for req_id in output.finished_sending or []:
+                new_count = self._send_remaining_count[req_id] - 1
+                if new_count == 0:
+                    # got response from all workers, report back to scheduler
+                    finished_sending.add(req_id)
+                    del self._send_remaining_count[req_id]
+                else:
+                    self._send_remaining_count[req_id] = new_count
+
+            # update finished_recving
+            for req_id in output.finished_recving or []:
+                new_count = self._recv_remaining_count[req_id] - 1
+                if new_count == 0:
+                    # got response from all workers, report back to scheduler
+                    finished_recving.add(req_id)
+                    del self._recv_remaining_count[req_id]
+                else:
+                    self._recv_remaining_count[req_id] = new_count
+
+        # select output of the worker specified by output_rank
+        output = outputs[0]
+
+        # set the aggregated finished_sending / finished_recving
+        if finished_sending:
+            output.finished_sending = finished_sending
+        if finished_recving:
+            output.finished_recving = finished_recving
+
+        return output
+
+    def _async_aggregate_workers_output(
+        self, output_futures: Sequence[Union[Future[ModelRunnerOutput], FutureWrapper]]
+    ) -> Future[ModelRunnerOutput]:
+        """Takes a list of futures and returns a single future which resolves
+        to the respective list of outputs."""
+        result_future: Future[ModelRunnerOutput] = Future()
+
+        outputs: list[Optional[ModelRunnerOutput]] = [None
+                                                      ] * len(output_futures)
+
+        def make_callback(idx):
+
+            def callback(fut):
+                if result_future.done():
+                    return
+
+                try:
+                    outputs[idx] = fut.result()
+                except CancelledError:
+                    result_future.cancel()
+                except Exception as e:
+                    result_future.set_exception(e)
+
+                # Check if all outputs are ready
+                if all(outputs):
+                    result_future.set_result(
+                        self._aggregate_workers_output(
+                            cast(list[ModelRunnerOutput], outputs)))
+
+            return callback
+
+        for i, output_future in enumerate(output_futures):
+            output_future.add_done_callback(make_callback(i))
+
+        return result_future
+
     def execute_model(
         self,
         scheduler_output,
@@ -55,10 +145,22 @@ def execute_model(
 
         refs = self.forward_dag.execute(scheduler_output)  # type: ignore
 
-        # When PP is not used, we block here until the result is available.
-        if self.max_concurrent_batches == 1:
-            return refs[0].get()
+        if not self.has_connector:
+            # get output only from a single worker (output_rank)
+            # When PP is not used, we block here until the result is available.
+            if self.max_concurrent_batches == 1:
+                return refs[0].get()
 
-        # When PP is used, we return a FutureWrapper immediately so that
-        # the scheduler can yield to the next batch.
-        return FutureWrapper(refs[0])
+            # When PP is used, we return a FutureWrapper immediately so that
+            # the scheduler can yield to the next batch.
+            return FutureWrapper(refs[0])
+
+        # get output from all workers when connector is present
+        if self.max_concurrent_batches == 1:
+            # Block and get results from all workers
+            outputs = [ref.get() for ref in refs]
+            return self._aggregate_workers_output(outputs)
+        else:
+            # Return a future that will aggregate outputs from all workers
+            output_futures = [FutureWrapper(ref) for ref in refs]
+            return self._async_aggregate_workers_output(output_futures)

vllm/v1/worker/gpu_model_runner.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import copy
 import gc
 import time
 import weakref
@@ -1233,6 +1234,8 @@ def _pool(
         hidden_states: torch.Tensor,
         num_scheduled_tokens: int,
         num_scheduled_tokens_np: np.ndarray,
+        finished_sending: Optional[set[str]],
+        finished_recving: Optional[set[str]],
     ) -> ModelRunnerOutput:
         assert self.input_batch.num_reqs ==\
             len(self.input_batch.pooling_params), \
@@ -1267,6 +1270,8 @@ def _pool(
             logprobs=None,
             prompt_logprobs_dict={},
             pooler_output=pooler_output,
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
         )
 
     @torch.inference_mode()
@@ -1277,12 +1282,12 @@ def execute_model(
     ) -> Union[ModelRunnerOutput, IntermediateTensors]:
         self._update_states(scheduler_output)
         if not scheduler_output.total_num_scheduled_tokens:
-            if has_kv_transfer_group():
-                with set_forward_context(None, self.vllm_config):
-                    self.maybe_setup_kv_connector(scheduler_output)
+            if not has_kv_transfer_group():
+                # Return empty ModelRunnerOutput if there's no work to do.
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            return self.kv_connector_no_forward(scheduler_output)
 
-            # Return empty ModelRunnerOutput if there's no work to do.
-            return EMPTY_MODEL_RUNNER_OUTPUT
 
         # Prepare the decoder inputs.
         (attn_metadata, attention_cuda_graphs, logits_indices,
@@ -1375,6 +1380,8 @@ def execute_model(
             )
 
             self.maybe_wait_for_kv_save()
+            finished_sending, finished_recving = (
+                self.get_finished_kv_transfers(scheduler_output))
 
         if self.use_aux_hidden_state_outputs:
             hidden_states, aux_hidden_states = model_output
@@ -1400,7 +1407,7 @@ def execute_model(
         else:
             if self.input_batch.pooling_params:
                 return self._pool(hidden_states, num_scheduled_tokens,
-                                  num_scheduled_tokens_np)
+                                  num_scheduled_tokens_np, finished_sending, finished_recving)
 
             sample_hidden_states = hidden_states[logits_indices]
             logits = self.model.compute_logits(sample_hidden_states, None)
@@ -1694,6 +1701,31 @@ def maybe_wait_for_kv_save() -> None:
         if has_kv_transfer_group():
             get_kv_transfer_group().wait_for_save()
 
+    @staticmethod
+    def get_finished_kv_transfers(
+        scheduler_output: "SchedulerOutput",
+    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        if has_kv_transfer_group():
+            return get_kv_transfer_group().get_finished(
+                scheduler_output.finished_req_ids)
+        return None, None
+
+    def kv_connector_no_forward(
+            self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
+        # KV send/recv even if no work to do.
+        with set_forward_context(None, self.vllm_config):
+            self.maybe_setup_kv_connector(scheduler_output)
+            finished_sending, finished_recving = (
+                self.get_finished_kv_transfers(scheduler_output))
+
+        if not finished_sending and not finished_recving:
+            return EMPTY_MODEL_RUNNER_OUTPUT
+
+        output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
+        output.finished_sending = finished_sending
+        output.finished_recving = finished_recving
+        return output
+
     def propose_ngram_draft_token_ids(
         self,
         sampled_token_ids: list[list[int]],

vllm/v1/worker/gpu_worker.py

@@ -332,22 +332,6 @@ def execute_model(
             output = EMPTY_MODEL_RUNNER_OUTPUT
 
         assert isinstance(output, ModelRunnerOutput)
-        if has_kv_transfer_group():
-            finished_sending, finished_recving = (
-                get_kv_transfer_group().get_finished(
-                    scheduler_output.finished_req_ids))
-            if finished_sending or finished_recving:
-                if output is EMPTY_MODEL_RUNNER_OUTPUT:
-                    output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
-                output.finished_sending = finished_sending
-                output.finished_recving = finished_recving
-
-            # Clear KVConnector state for this step.
-            get_kv_transfer_group().clear_connector_metadata()
-
-            # with a connector, the scheduler expects output from all workers
-            return output
-
         # return output only from the driver worker
         return output if self.is_driver_worker else None