@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c01f6e525f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 00:30 UTC

[CI] Fix mypy for `vllm/v1/core` and `vllm/v1/engine` (#27108)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Oct 30, 2025 at 07:32 UTC c01f6e525f457133cfb00127a89c09e5247e563c

12 files changed +91 -61

tools/pre_commit/mypy.py

+13 -1

@@ -36,12 +36,15 @@ FILES = [

36 "vllm/transformers_utils",

37 "vllm/triton_utils",

38 "vllm/usage",

39 + "vllm/v1/core",

40 + "vllm/v1/engine",

41 ]

43 # After fixing errors resulting from changing follow_imports

44 # from "skip" to "silent", move the following directories to FILES

45 SEPARATE_GROUPS = [

46 "tests",

47 + # v0 related

48 "vllm/attention",

49 "vllm/compilation",

50 "vllm/engine",

@@ -50,7 +53,16 @@ SEPARATE_GROUPS = [

53 "vllm/model_executor",

54 "vllm/plugins",

55 "vllm/worker",

53 - "vllm/v1",

56 + # v1 related

57 + "vllm/v1/attention",

58 + "vllm/v1/executor",

59 + "vllm/v1/kv_offload",

60 + "vllm/v1/metrics",

61 + "vllm/v1/pool",

62 + "vllm/v1/sample",

63 + "vllm/v1/spec_decode",

64 + "vllm/v1/structured_output",

65 + "vllm/v1/worker",

66 ]

68 # TODO(woosuk): Include the code from Megatron and HuggingFace.

vllm/config/vllm.py

+4 -5

@@ -84,7 +84,9 @@ class VllmConfig:

84 default_factory=StructuredOutputsConfig

85 )

86 """Structured outputs configuration."""

87 - observability_config: ObservabilityConfig | None = None

87 + observability_config: ObservabilityConfig = Field(

88 + default_factory=ObservabilityConfig

89 + )

90 """Observability configuration."""

91 quant_config: QuantizationConfig | None = None

92 """Quantization configuration."""

@@ -170,10 +172,7 @@ class VllmConfig:

172 vllm_factors.append(self.structured_outputs_config.compute_hash())

173 else:

174 vllm_factors.append("None")

173 - if self.observability_config:

174 - vllm_factors.append(self.observability_config.compute_hash())

175 - else:

176 - vllm_factors.append("None")

175 + vllm_factors.append(self.observability_config.compute_hash())

176 if self.quant_config:

177 pass # should be captured by model_config.quantization

178 if self.compilation_config:

vllm/engine/protocol.py

@@ -77,6 +77,7 @@ class EngineClient(ABC):

77 lora_request: LoRARequest | None = None,

78 trace_headers: Mapping[str, str] | None = None,

79 priority: int = 0,

80 + truncate_prompt_tokens: int | None = None,

81 tokenization_kwargs: dict[str, Any] | None = None,

82 ) -> AsyncGenerator[PoolingRequestOutput, None]:

83 """Generate outputs for a request from a pooling model."""

vllm/v1/core/sched/scheduler.py

+9 -7

@@ -167,7 +167,7 @@ class Scheduler(SchedulerInterface):

167 self.kv_cache_manager = KVCacheManager(

168 kv_cache_config=kv_cache_config,

169 max_model_len=self.max_model_len,

170 - enable_caching=self.cache_config.enable_prefix_caching,

170 + enable_caching=bool(self.cache_config.enable_prefix_caching),

171 use_eagle=self.use_eagle,

172 log_stats=self.log_stats,

173 enable_kv_cache_events=self.enable_kv_cache_events,

@@ -407,13 +407,13 @@ class Scheduler(SchedulerInterface):

407

408 # Get externally-cached tokens if using a KVConnector.

409 if self.connector is not None:

410 - num_external_computed_tokens, load_kv_async = (

410 + ext_tokens, load_kv_async = (

411 self.connector.get_num_new_matched_tokens(

412 request, num_new_local_computed_tokens

413 )

414 )

415

416 - if num_external_computed_tokens is None:

416 + if ext_tokens is None:

417 # The request cannot be scheduled because

418 # the KVConnector couldn't determine

419 # the number of matched tokens.

@@ -421,6 +421,8 @@ class Scheduler(SchedulerInterface):

421 skipped_waiting_requests.prepend_request(request)

422 continue

423

424 + num_external_computed_tokens = ext_tokens

425 +

426 # Total computed tokens (local + external).

427 num_computed_tokens = (

428 num_new_local_computed_tokens + num_external_computed_tokens

@@ -905,13 +907,13 @@ class Scheduler(SchedulerInterface):

907

908 outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)

909 spec_decoding_stats: SpecDecodingStats | None = None

908 - kv_connector_stats = (

910 + kv_connector_stats: KVConnectorStats | None = (

911 kv_connector_output.kv_connector_stats if kv_connector_output else None

912 )

913 if kv_connector_stats and self.connector:

912 - stats = self.connector.get_kv_connector_stats()

913 - if stats:

914 - kv_connector_stats = kv_connector_stats.aggregate(stats)

914 + kv_stats = self.connector.get_kv_connector_stats()

915 + if kv_stats:

916 + kv_connector_stats = kv_connector_stats.aggregate(kv_stats)

917

918 failed_kv_load_req_ids = None

919 if kv_connector_output and kv_connector_output.invalid_block_ids:

vllm/v1/engine/async_llm.py

+13 -8

@@ -6,7 +6,7 @@ import socket

6 import time

7 from collections.abc import AsyncGenerator, Iterable, Mapping

8 from copy import copy

9 -from typing import Any

9 +from typing import Any, cast

11 import numpy as np

12 import torch

@@ -131,10 +131,9 @@ class AsyncLLM(EngineClient):

131 self.output_processor = OutputProcessor(

132 self.tokenizer, log_stats=self.log_stats

133 )

134 - if self.observability_config.otlp_traces_endpoint is not None:

135 - tracer = init_tracer(

136 - "vllm.llm_engine", self.observability_config.otlp_traces_endpoint

137 - )

134 + endpoint = self.observability_config.otlp_traces_endpoint

135 + if endpoint is not None:

136 + tracer = init_tracer("vllm.llm_engine", endpoint)

137 self.output_processor.tracer = tracer

138

139 # EngineCore (starts the engine in background process).

@@ -266,7 +265,9 @@ class AsyncLLM(EngineClient):

265 if engine_core := getattr(self, "engine_core", None):

266 engine_core.shutdown()

267

269 - cancel_task_threadsafe(getattr(self, "output_handler", None))

268 + handler = getattr(self, "output_handler", None)

269 + if handler is not None:

270 + cancel_task_threadsafe(handler)

271

272 async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

273 return await self.engine_core.get_supported_tasks_async()

@@ -314,7 +315,10 @@ class AsyncLLM(EngineClient):

315 priority,

316 data_parallel_rank,

317 )

317 - prompt_text = prompt if isinstance(prompt, str) else prompt.get("prompt")

318 + if isinstance(prompt, str):

319 + prompt_text = prompt

320 + elif isinstance(prompt, Mapping):

321 + prompt_text = cast(str | None, prompt.get("prompt"))

322

323 if is_pooling or params.n == 1:

324 await self._add_request(request, prompt_text, None, 0, queue)

@@ -436,6 +440,7 @@ class AsyncLLM(EngineClient):

440 # Note: both OutputProcessor and EngineCore handle their

441 # own request cleanup based on finished.

442 finished = out.finished

443 + assert isinstance(out, RequestOutput)

444 yield out

445

446 # If the request is disconnected by the client, generate()

@@ -653,7 +658,7 @@ class AsyncLLM(EngineClient):

658 return self.tokenizer

659

660 async def is_tracing_enabled(self) -> bool:

656 - return self.observability_config.otlp_traces_endpoint is not None

661 + return self.observability_config.otlp_traces_endpoint is not None # type: ignore

662

663 async def do_log_stats(self) -> None:

664 if self.logger_manager:

vllm/v1/engine/core.py

@@ -1075,6 +1075,7 @@ class DPEngineCoreProc(EngineCoreProc):

1075 local_dp_rank = vllm_config.parallel_config.data_parallel_rank_local

1076

1077 assert dp_size > 1

1078 + assert local_dp_rank is not None

1079 assert 0 <= local_dp_rank <= dp_rank < dp_size

1080

1081 if vllm_config.kv_transfer_config is not None:

vllm/v1/engine/core_client.py

+7 -7

@@ -385,10 +385,11 @@ class BackgroundResources:

385 with contextlib.suppress(Exception):

386 task.cancel()

387

388 - if in_loop(loop):

389 - close_sockets_and_tasks()

390 - elif loop and not loop.is_closed():

391 - loop.call_soon_threadsafe(close_sockets_and_tasks)

388 + if loop is not None:

389 + if in_loop(loop):

390 + close_sockets_and_tasks()

391 + elif not loop.is_closed():

392 + loop.call_soon_threadsafe(close_sockets_and_tasks)

393 else:

394 # Loop has been closed, try to clean up directly.

395 del tasks

@@ -1044,6 +1045,7 @@ class DPAsyncMPClient(AsyncMPClient):

1045 return

1046

1047 assert self.stats_update_address is not None

1048 + stats_addr: str = self.stats_update_address

1049 assert len(self.engine_ranks_managed) > 0

1050 # NOTE: running and waiting counts are all global from

1051 # the Coordinator include all global EngineCores. This

@@ -1054,9 +1056,7 @@ class DPAsyncMPClient(AsyncMPClient):

1056

1057 async def run_engine_stats_update_task():

1058 with (

1057 - make_zmq_socket(

1058 - self.ctx, self.stats_update_address, zmq.XSUB, linger=0

1059 - ) as socket,

1059 + make_zmq_socket(self.ctx, stats_addr, zmq.XSUB, linger=0) as socket,

1060 make_zmq_socket(

1061 self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=False, linger=0

1062 ) as first_req_rcv_socket,

vllm/v1/engine/detokenizer.py

+10 -3

@@ -69,14 +69,21 @@ class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC):

69 # Stop strings

70 params = request.sampling_params

71 assert params is not None

72 - self.stop = stop = params.stop

72 + stop_list: list[str]

73 + if params.stop is None:

74 + stop_list = []

75 + elif isinstance(params.stop, str):

76 + stop_list = [params.stop]

77 + else:

78 + stop_list = params.stop

79 + self.stop = stop_list

80 self.min_tokens = params.min_tokens

81 self.include_stop_str_in_output = params.include_stop_str_in_output

83 # Number of chars to hold back when stop strings are to be excluded

84 # from streamed output.

78 - if stop and not self.include_stop_str_in_output:

79 - self.stop_buffer_length = max(len(s) for s in stop) - 1

85 + if self.stop and not self.include_stop_str_in_output:

86 + self.stop_buffer_length = max(len(s) for s in self.stop) - 1

87 else:

88 self.stop_buffer_length = 0

89 self._last_output_text_offset: int = 0

vllm/v1/engine/llm_engine.py

+9 -7

@@ -4,7 +4,7 @@

4 import time

5 from collections.abc import Callable, Mapping

6 from copy import copy

7 -from typing import Any

7 +from typing import Any, cast

9 import torch.nn as nn

10 from typing_extensions import TypeVar

@@ -112,10 +112,9 @@ class LLMEngine:

112 self.output_processor = OutputProcessor(

113 self.tokenizer, log_stats=self.log_stats

114 )

115 - if self.observability_config.otlp_traces_endpoint is not None:

116 - tracer = init_tracer(

117 - "vllm.llm_engine", self.observability_config.otlp_traces_endpoint

118 - )

115 + endpoint = self.observability_config.otlp_traces_endpoint

116 + if endpoint is not None:

117 + tracer = init_tracer("vllm.llm_engine", endpoint)

118 self.output_processor.tracer = tracer

119

120 # EngineCore (gets EngineCoreRequests and gives EngineCoreOutputs)

@@ -259,7 +258,10 @@ class LLMEngine:

258 trace_headers,

259 priority,

260 )

262 - prompt_text = prompt if isinstance(prompt, str) else prompt.get("prompt")

261 + if isinstance(prompt, str):

262 + prompt_text = prompt

263 + elif isinstance(prompt, Mapping):

264 + prompt_text = cast(str | None, prompt.get("prompt"))

265

266 n = params.n if isinstance(params, SamplingParams) else 1

267

@@ -285,7 +287,7 @@ class LLMEngine:

287 # Add the request to EngineCore.

288 self.engine_core.add_request(child_request)

289

288 - def step(self) -> list[RequestOutput] | list[PoolingRequestOutput]:

290 + def step(self) -> list[RequestOutput | PoolingRequestOutput]:

291 if self.should_execute_dummy_batch:

292 self.should_execute_dummy_batch = False

293 self.engine_core.execute_dummy_batch()

vllm/v1/engine/output_processor.py

+8 -2

@@ -44,10 +44,16 @@ class RequestOutputCollector:

44 if self.output is None or isinstance(output, Exception):

45 self.output = output

46 self.ready.set()

47 - elif isinstance(self.output, (RequestOutput, PoolingRequestOutput)):

47 + elif isinstance(self.output, RequestOutput) and isinstance(

48 + output, RequestOutput

49 + ):

50 # This ensures that request outputs with different request indexes

51 # (if n > 1) do not override each other.

52 self.output.add(output, aggregate=self.aggregate)

53 + elif isinstance(self.output, PoolingRequestOutput) and isinstance(

54 + output, PoolingRequestOutput

55 + ):

56 + self.output = output

58 async def get(self) -> RequestOutput | PoolingRequestOutput:

59 """Get operation blocks on put event."""

@@ -408,7 +414,7 @@ class OutputProcessor:

414 within the loop below.

415 """

416

411 - request_outputs: list[RequestOutput] | list[PoolingRequestOutput] = []

417 + request_outputs: list[RequestOutput | PoolingRequestOutput] = []

418 reqs_to_abort: list[str] = []

419 for engine_core_output in engine_core_outputs:

420 req_id = engine_core_output.request_id

vllm/v1/engine/parallel_sampling.py

+2 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from copy import copy

5 -from typing import Optional

5 +from typing import Optional, cast

7 from vllm.outputs import CompletionOutput

8 from vllm.sampling_params import RequestOutputKind, SamplingParams

@@ -37,7 +37,7 @@ class ParentRequest:

38 self.child_requests = set()

39 self.output_aggregator = (

40 - [None] * sampling_params.n

40 + [cast(CompletionOutput, None)] * sampling_params.n

41 if (sampling_params.output_kind == RequestOutputKind.FINAL_ONLY)

42 else []

43 )

vllm/v1/engine/processor.py

+14 -19

@@ -3,7 +3,7 @@

4 import time

5 from collections.abc import Mapping

6 -from typing import Any, Literal

6 +from typing import Any, Literal, cast

8 from vllm.config import VllmConfig

9 from vllm.inputs import ProcessorInputs, PromptType, SingletonInputs

@@ -208,9 +208,9 @@ class Processor:

208 enc = prompt.get("encoder_prompt")

209 dec = prompt.get("decoder_prompt")

210 if enc is not None:

211 - _validate_single_prompt(enc)

211 + _validate_single_prompt(cast(dict | str, enc))

212 if dec is not None:

213 - _validate_single_prompt(dec)

213 + _validate_single_prompt(cast(dict | str, dec))

214 else:

215 _validate_single_prompt(prompt) # type: ignore[arg-type]

216

@@ -332,7 +332,7 @@ class Processor:

332 if not mm_data:

333 return None

334

335 - mm_uuids: MultiModalUUIDDict = {}

335 + mm_uuids: dict[str, list[str | None] | str] = {}

336 for modality, data in mm_data.items():

337 n = len(data) if isinstance(data, list) else 1

338 mm_uuids[modality] = [f"{request_id}-{modality}-{i}" for i in range(n)]

@@ -384,7 +384,9 @@ class Processor:

384 # if provided.

385 self._validate_multi_modal_uuids(prompt)

386 if isinstance(prompt, dict):

387 - mm_uuids = prompt.get("multi_modal_uuids")

387 + mm_uuids = cast(

388 + MultiModalUUIDDict | None, prompt.get("multi_modal_uuids")

389 + )

390 else:

391 mm_uuids = None

392

@@ -410,20 +412,13 @@ class Processor:

412 encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)

413 self._validate_model_inputs(encoder_inputs, decoder_inputs)

414

413 - # Mypy does not always properly infer the types of some elements of

414 - # discriminated unions of TypedDicts, because of how it handles

415 - # inheritance of TypedDict. If we explicitly extract the items we want

416 - # we can avoid type errors from using `dict.get` later in the method.

417 - prompt_token_ids = (

418 - decoder_inputs["prompt_token_ids"]

419 - if decoder_inputs["type"] != "embeds"

420 - else None

421 - )

422 - prompt_embeds = (

423 - decoder_inputs["prompt_embeds"]

424 - if decoder_inputs["type"] == "embeds"

425 - else None

426 - )

415 + # Mypy can be conservative for TypedDict unions; normalize access.

416 + if decoder_inputs["type"] == "embeds":

417 + prompt_token_ids = None

418 + prompt_embeds = decoder_inputs["prompt_embeds"]

419 + else:

420 + prompt_token_ids = decoder_inputs["prompt_token_ids"]

421 + prompt_embeds = None

422

423 sampling_params = None

424 pooling_params = None