[Priority merge] NewRequestData parameter introduced in vllm upstream (#245)

Temporary hack until the parameter makes it to a new release version.
Needs to be merged first for the tests on the other PRs to pass.

(PS: this was actually the error after fixing the merge conflict in PR
#240, which had nothing to do with the conflict)

---------

Signed-off-by: Sophie du Couédic <[email protected]>
Co-authored-by: Yannick Schnider <[email protected]>

Author: sducouedic

tests/spyre_util.py

@@ -1,4 +1,3 @@
-import inspect
 import math
 import os
 import subprocess
@@ -540,35 +539,24 @@ def create_random_request(
         request_id: int, num_tokens: int,
         sampling_params: SamplingParams) -> EngineCoreRequest:
 
-    # Temporary until 'data_parallel_rank' parameter makes it to
-    # a release version in vllm
-    if "data_parallel_rank" in [
-            x[0] for x in inspect.getmembers(EngineCoreRequest)
-    ]:
-        return EngineCoreRequest(
-            request_id=str(request_id),
-            prompt_token_ids=[request_id] * num_tokens,
-            mm_inputs=None,
-            mm_hashes=None,
-            mm_placeholders=None,
-            sampling_params=sampling_params,
-            eos_token_id=None,
-            arrival_time=0,
-            lora_request=None,
-            cache_salt=None,
-            data_parallel_rank=None,
-        )
-    else:
-        return EngineCoreRequest(request_id=str(request_id),
-                                 prompt_token_ids=[request_id] * num_tokens,
-                                 mm_inputs=None,
-                                 mm_hashes=None,
-                                 mm_placeholders=None,
-                                 sampling_params=sampling_params,
-                                 eos_token_id=None,
-                                 arrival_time=0,
-                                 lora_request=None,
-                                 cache_salt=None)
+    # Temporary until these parameters make it to a release version in vllm
+    extra_kwargs: dict[str, Any] = {}
+    if "data_parallel_rank" in EngineCoreRequest.__annotations__:
+        extra_kwargs["data_parallel_rank"] = None
+    if "pooling_params" in EngineCoreRequest.__annotations__:
+        extra_kwargs["pooling_params"] = None
+
+    return EngineCoreRequest(request_id=str(request_id),
+                             prompt_token_ids=[request_id] * num_tokens,
+                             mm_inputs=None,
+                             mm_hashes=None,
+                             mm_placeholders=None,
+                             sampling_params=sampling_params,
+                             eos_token_id=None,
+                             arrival_time=0,
+                             lora_request=None,
+                             cache_salt=None,
+                             **extra_kwargs)
 
 
 def skip_unsupported_tp_size(size: int):

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -2,7 +2,7 @@
 from collections import deque
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 from torch import nn
@@ -430,6 +430,12 @@ def execute_model(
 
         t0 = time.time()
 
+        # TODO temporary until 'pooler_output' makes it to a release version
+        # in vllm
+        extra_kwargs: dict[str, Any] = {}
+        if "pooler_output" in ModelRunnerOutput.__dataclass_fields__:
+            extra_kwargs["pooler_output"] = None
+
         # TODO: change to EMPTY_MODEL_RUNNER_OUTPUT, right now this
         # will be a breaking change, or clumsy to make retrocompatible
         # with conditional import
@@ -442,6 +448,7 @@ def execute_model(
                 spec_token_ids=None,
                 logprobs=None,
                 prompt_logprobs_dict={},
+                **extra_kwargs,
             )
 
         self._update_states(scheduler_output)
@@ -490,6 +497,7 @@ def execute_model(
                 req_id: None
                 for req_id in self.input_batch.req_id_to_index
             },  # TODO(wallas?): prompt logprobs too
+            **extra_kwargs,
         )
         return model_output
 
@@ -937,21 +945,27 @@ def execute_model(
 
         t0 = time.time()
 
+        # TODO temporary until 'pooler_output' makes it to a release version
+        # in vllm
+        extra_kwargs: dict[str, Any] = {}
+        if "pooler_output" in CBSpyreModelRunnerOutput.__dataclass_fields__:
+            extra_kwargs["pooler_output"] = None
+
         self._update_states(scheduler_output)
         # TODO: change to EMPTY_MODEL_RUNNER_OUTPUT, right now this
         # will be a breaking change, or clumsy to make retrocompatible
         # with conditional import
         if not scheduler_output.total_num_scheduled_tokens:
+
             # Return empty ModelRunnerOuptut if there's no work to do.
-            return CBSpyreModelRunnerOutput(
-                req_ids=[],
-                req_id_to_index={},
-                sampled_token_ids=[],
-                spec_token_ids=None,
-                logprobs=None,
-                prompt_logprobs_dict={},
-                tkv=0,
-            )
+            return CBSpyreModelRunnerOutput(req_ids=[],
+                                            req_id_to_index={},
+                                            sampled_token_ids=[],
+                                            spec_token_ids=None,
+                                            logprobs=None,
+                                            prompt_logprobs_dict={},
+                                            tkv=0,
+                                            **extra_kwargs)
 
         model_input = self.prepare_model_input(scheduler_output)
 
@@ -1037,5 +1051,6 @@ def execute_model(
                                   for req_id in req_ids
                                   },  # TODO(wallas?): prompt logprobs too
             tkv=self.tkv,
+            **extra_kwargs,
         )
         return model_output

vllm_spyre/v1/worker/spyre_worker.py

@@ -5,7 +5,7 @@
 import platform
 import signal
 import time
-from typing import Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
 import torch
 import torch.distributed as dist
@@ -322,6 +322,11 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
+        # TODO temporary until 'pooling_params' makes it to a release version
+        # in vllm
+        extra_kwargs: dict[str, Any] = {}
+        if "pooling_params" in NewRequestData.__dataclass_fields__:
+            extra_kwargs["pooling_params"] = None
         dummy_requests = [
             NewRequestData(
                 req_id="warmup-%d" % (i),
@@ -333,7 +338,7 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
                 block_ids=[0],  # not actually used
                 num_computed_tokens=0,
                 lora_request=None,
-            ) for i in range(batch_size)
+                **extra_kwargs) for i in range(batch_size)
         ]
 
         for i, req in enumerate(dummy_requests):
@@ -487,6 +492,12 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size, prompt_len))]
 
+        # TODO temporary until 'pooling_params' makes it to a release version
+        # in vllm
+        extra_kwargs: dict[str, Any] = {}
+        if "pooling_params" in NewRequestData.__dataclass_fields__:
+            extra_kwargs["pooling_params"] = None
+
         # Set up dummy requests for prefill steps
         dummy_requests = [
             NewRequestData(
@@ -499,7 +510,7 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
                 block_ids=[0],
                 num_computed_tokens=0,
                 lora_request=None,
-            ) for i in range(batch_size)
+                **extra_kwargs) for i in range(batch_size)
         ]
 
         # Set up dummy cached_requests for decode steps