5940590855 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/entrypoints/openai/completion/test_shutdown.py

+12 -2

@@ -26,6 +26,9 @@ _SERVER_STARTUP_TIMEOUT = 120

26 _PROCESS_EXIT_TIMEOUT = 15

27 _SHUTDOWN_DETECTION_TIMEOUT = 10

28 _CHILD_CLEANUP_TIMEOUT = 10

29 +_INFLIGHT_REQUEST_START_TIMEOUT = 5

30 +_INFLIGHT_REQUEST_POLL_INTERVAL = 0.1

31 +_ABORT_CLIENT_TIMEOUT = 3

32

33

34 def _get_child_pids(parent_pid: int) -> list[int]:

@@ -71,6 +74,7 @@ class ShutdownState:

74 requests_after_sigterm: int = 0

75 aborted_requests: int = 0

76 connection_errors: int = 0

77 + inflight_requests: int = 0

78 stop_requesting: bool = False

79 errors: list[str] = field(default_factory=list)

80

@@ -86,6 +90,7 @@ async def _concurrent_request_loop(

90 async def single_request():

91 while not state.stop_requesting:

92 try:

93 + state.inflight_requests += 1

94 response = await client.completions.create(

95 model=MODEL_NAME,

96 prompt="Write a story: ",

@@ -110,6 +115,8 @@ async def _concurrent_request_loop(

115 except Exception as e:

116 state.errors.append(f"Unexpected error: {e}")

117 break

118 + finally:

119 + state.inflight_requests -= 1

120 await asyncio.sleep(0.01)

121

122 tasks = [asyncio.create_task(single_request()) for _ in range(concurrency)]

@@ -392,7 +399,7 @@ async def test_abort_timeout_fails_inflight_requests():

399 ]

400

401 with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:

395 - client = remote_server.get_async_client()

402 + client = remote_server.get_async_client(timeout=_ABORT_CLIENT_TIMEOUT)

403 proc = remote_server.proc

404 child_pids = _get_child_pids(proc.pid)

405

@@ -403,7 +410,10 @@ async def test_abort_timeout_fails_inflight_requests():

410 _concurrent_request_loop(client, state, sigterm_sent, concurrency=10)

411 )

412

406 - await asyncio.sleep(0.5)

413 + deadline = time.time() + _INFLIGHT_REQUEST_START_TIMEOUT

414 + while state.inflight_requests == 0 and time.time() < deadline:

415 + await asyncio.sleep(_INFLIGHT_REQUEST_POLL_INTERVAL)

416 + assert state.inflight_requests > 0

417

418 proc.send_signal(signal.SIGTERM)

419 sigterm_sent.set()

tests/entrypoints/openai/test_openai_schema.py

+3 -1

@@ -118,7 +118,9 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):

118 # the default filtered-vs-good ratio. The filter is intentional, so

119 # suppress the health check rather than drop the filter — dropping it

120 # exposes pre-existing server bugs out of scope here.

121 - suppress_health_check=[HealthCheck.filter_too_much],

121 + # The same nested schema can also trip Hypothesis' entropy budget while

122 + # generating large-but-valid request bodies before vLLM is called.

123 + suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],

124 )

125 def test_openapi_stateless(case: Case):

126 key = (

tests/entrypoints/serve/disagg/test_generate_stream.py

+9

@@ -69,10 +69,16 @@ class MockParallelConfig:

69 _api_process_rank: int = 0

70

71

72 +@dataclass

73 +class MockSchedulerConfig:

74 + max_num_seqs: int = 128

75 +

76 +

77 @dataclass

78 class MockVllmConfig:

79 model_config: MockModelConfig

80 parallel_config: MockParallelConfig

81 + scheduler_config: MockSchedulerConfig = field(default_factory=MockSchedulerConfig)

82

83

84 def _build_renderer(model_config: MockModelConfig):

@@ -149,6 +155,9 @@ def _mock_engine() -> MagicMock:

155 engine = MagicMock(spec=AsyncLLM)

156 engine.errored = False

157 engine.model_config = MockModelConfig()

158 + engine.vllm_config = MockVllmConfig(

159 + engine.model_config, parallel_config=MockParallelConfig()

160 + )

161 engine.input_processor = MagicMock()

162 engine.renderer = _build_renderer(engine.model_config)

163 return engine

tests/entrypoints/weight_transfer/test_weight_transfer_llm.py

+4 -4

@@ -63,8 +63,8 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo

63 last_init_info: MockInitInfo | None = None

64 last_update_info: MockUpdateInfo | None = None

65

66 - def __init__(self, config, parallel_config):

67 - super().__init__(config, parallel_config)

66 + def __init__(self, config, parallel_config, model):

67 + super().__init__(config, parallel_config, model)

68 # Reset tracking on init

69 MockWeightTransferEngine.init_transfer_engine_called = False

70 MockWeightTransferEngine.receive_weights_called = False

@@ -95,9 +95,9 @@ class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo

95 pass

96

97

98 -def mock_create_engine(config, parallel_config):

98 +def mock_create_engine(config, parallel_config, model):

99 """Mock factory function that returns our mock engine."""

100 - return MockWeightTransferEngine(config, parallel_config)

100 + return MockWeightTransferEngine(config, parallel_config, model)

101

102

103 # --- Tests ---

vllm/entrypoints/chat_utils.py

+19 -3

@@ -1818,12 +1818,28 @@ def _postprocess_messages(messages: list[ConversationMessage]) -> None:

1818 continue

1819

1820 for item in tool_calls:

1821 + if not isinstance(item, dict):

1822 + raise VLLMValidationError(

1823 + "assistant tool_calls entries must be objects.",

1824 + parameter="tool_calls",

1825 + )

1826 +

1827 + function = item.get("function")

1828 + if item.get("type", "function") != "function" or not isinstance(

1829 + function, dict

1830 + ):

1831 + raise VLLMValidationError(

1832 + "chat completions only support assistant tool_calls "

1833 + "of type 'function'.",

1834 + parameter="tool_calls",

1835 + )

1836 +

1837 # if arguments is None or empty string, set to {}

1822 - if content := item["function"].get("arguments"):

1838 + if content := function.get("arguments"):

1839 if not isinstance(content, (dict, list)):

1824 - item["function"]["arguments"] = json.loads(content)

1840 + function["arguments"] = json.loads(content)

1841 else:

1826 - item["function"]["arguments"] = {}

1842 + function["arguments"] = {}

1843

1844

1845 def parse_chat_messages(

vllm/entrypoints/openai/api_server.py

+1 -1

@@ -151,7 +151,7 @@ async def build_async_engine_client_from_engine_args(

151 yield async_llm

152 finally:

153 if async_llm:

154 - async_llm.shutdown()

154 + async_llm.shutdown(timeout=vllm_config.shutdown_timeout)

155

156

157 def build_app(

vllm/entrypoints/openai/chat_completion/batch_serving.py

+6 -2

@@ -218,8 +218,6 @@ class OpenAIServingChatBatch(OpenAIServingChat):

218 ``check_batch_mode`` validator, so neither needs to be handled here.

219 """

220 created_time = int(time.time())

221 - role = self.get_chat_request_role(request) # type: ignore[arg-type]

222 -

221 final_results: dict[int, RequestOutput] = {}

222 try:

223 async for prompt_idx, res in merge_async_iterators(*generators):

@@ -275,6 +273,12 @@ class OpenAIServingChatBatch(OpenAIServingChat):

273 reasoning = None

274 content = output.text

275

276 + role = (

277 + self.response_role

278 + if request.add_generation_prompt

279 + else request.messages[prompt_idx][-1]["role"]

280 + )

281 +

282 message = ChatMessage(role=role, reasoning=reasoning, content=content)

283

284 if request.echo:

vllm/entrypoints/openai/chat_completion/protocol.py

+3 -1

@@ -909,7 +909,9 @@ class BatchChatCompletionRequest(OpenAIBaseModel):

909 - The ``n`` parameter must be 1 (or omitted).

910 """

911

912 - messages: list[list[ChatCompletionMessageParam]] = Field(..., min_length=1)

912 + messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (

913 + Field(..., min_length=1)

914 + )

915 model: str | None = None

916

917 # Shared sampling / generation fields — mirror ChatCompletionRequest.

vllm/entrypoints/serve/disagg/protocol.py

+1 -1

@@ -68,7 +68,7 @@ class GenerateRequest(BaseModel):

68 "through out the inference process and return in response."

69 ),

70 )

71 - token_ids: list[int]

71 + token_ids: list[int] = Field(min_length=1)

72 """The token ids to generate text from."""

73

74 @field_validator("token_ids")

vllm/entrypoints/serve/disagg/serving.py

+13 -1

@@ -8,6 +8,7 @@ import time

8 from collections.abc import AsyncGenerator

9 from collections.abc import Sequence as GenericSequence

10

11 +import msgspec

12 import numpy as np

13 import pybase64 as base64

14 from fastapi import Request

@@ -125,6 +126,18 @@ class ServingTokens(OpenAIServing):

126 if raw_request:

127 raw_request.state.request_metadata = request_metadata

128

129 + sampling_params = request.sampling_params

130 + max_num_seqs = self.engine_client.vllm_config.scheduler_config.max_num_seqs

131 + if sampling_params.n > max_num_seqs:

132 + return self.create_error_response(

133 + f"sampling_params.n must be at most the server's max_num_seqs "

134 + f"({max_num_seqs}), got {sampling_params.n}."

135 + )

136 + try:

137 + msgspec.msgpack.encode(sampling_params)

138 + except (OverflowError, TypeError, ValueError) as e:

139 + return self.create_error_response(e)

140 +

141 engine_input: EngineInput

142 if features := request.features:

143 # Convert PlaceholderRangeInfo → PlaceholderRange per modality.

@@ -164,7 +177,6 @@ class ServingTokens(OpenAIServing):

177

178 # Schedule the request and get the result generator.

179 result_generator: AsyncGenerator[RequestOutput, None] | None = None

167 - sampling_params = request.sampling_params

180

181 # Apply server-side ``max_tokens`` defaulting when the client did

182 # not set it, matching the OpenAI-compat endpoints. ``SamplingParams``

vllm/v1/utils.py

+3 -4

@@ -474,10 +474,9 @@ def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None:

474 timeout: Maximum time in seconds to wait for graceful shutdown

475 """

476 if timeout is None:

477 - timeout = 0.0

478 -

479 - # Allow at least 5 seconds for remaining procs to terminate.

480 - timeout = max(timeout, 5.0)

477 + # Keep a small grace period for best-effort cleanup paths that do not

478 + # have a user-configured shutdown timeout.

479 + timeout = 5.0

480

481 # Shutdown the process.

482 for proc in procs:

[ROCm][CI] Stabilize 400 error return code for invalid schema inputs (#43016)