9e6bcda3ac · tduncan/Nvidia-Optimization-Stack-VLLM

tools/pre_commit/mypy.py

+6 -6

@@ -27,19 +27,24 @@ FILES = [

27 "vllm/*.py",

28 "vllm/assets",

29 "vllm/distributed",

30 + "vllm/engine",

31 "vllm/entrypoints",

32 "vllm/executor",

33 "vllm/inputs",

34 "vllm/logging_utils",

35 "vllm/multimodal",

36 "vllm/platforms",

37 + "vllm/plugins",

38 "vllm/transformers_utils",

39 "vllm/triton_utils",

40 "vllm/usage",

41 "vllm/utils",

42 + "vllm/worker",

43 "vllm/v1/core",

44 "vllm/v1/engine",

45 + "vllm/v1/metrics",

46 "vllm/v1/pool",

47 + "vllm/v1/sample",

48 "vllm/v1/worker",

49 ]

50

@@ -50,24 +55,19 @@ SEPARATE_GROUPS = [

55 # v0 related

56 "vllm/attention",

57 "vllm/compilation",

53 - "vllm/engine",

54 - "vllm/inputs",

58 "vllm/lora",

59 "vllm/model_executor",

57 - "vllm/plugins",

58 - "vllm/worker",

60 # v1 related

61 "vllm/v1/attention",

62 "vllm/v1/executor",

63 "vllm/v1/kv_offload",

63 - "vllm/v1/metrics",

64 - "vllm/v1/sample",

64 "vllm/v1/spec_decode",

65 "vllm/v1/structured_output",

66 ]

67

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

69 EXCLUDE = [

70 + "vllm/engine/arg_utils.py",

71 "vllm/model_executor/parallel_utils",

72 "vllm/model_executor/models",

73 "vllm/model_executor/layers/fla/ops",

vllm/distributed/kv_transfer/kv_connector/v1/base.py

+1 -1

@@ -565,7 +565,7 @@ class KVConnectorBase_V1(ABC):

565 vllm_config: "VllmConfig",

566 metric_types: dict[type["PromMetric"], type["PromMetricT"]],

567 labelnames: list[str],

568 - per_engine_labelvalues: dict[int, list[str]],

568 + per_engine_labelvalues: dict[int, list[object]],

569 ) -> Optional["KVConnectorPromMetrics"]:

570 """

571 Create a KVConnectorPromMetrics subclass which should register

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py

+1 -1

@@ -806,7 +806,7 @@ class LMCacheMPConnector(KVConnectorBase_V1):

806 vllm_config: "VllmConfig",

807 metric_types: dict[type["PromMetric"], type["PromMetricT"]],

808 labelnames: list[str],

809 - per_engine_labelvalues: dict[int, list[str]],

809 + per_engine_labelvalues: dict[int, list[object]],

810 ) -> Optional["KVConnectorPromMetrics"]:

811 """

812 Create a KVConnectorPromMetrics subclass which should register

vllm/distributed/kv_transfer/kv_connector/v1/metrics.py

+6 -6

@@ -52,13 +52,13 @@ class KVConnectorStats:

52

53

54 class KVConnectorLogging:

55 - def __init__(self, kv_tranfer_config: KVTransferConfig):

55 + def __init__(self, kv_transfer_config: KVTransferConfig | None):

56 # This should be called on frontend process.

57 assert not has_kv_transfer_group()

58 # Instantiate the connector's stats class.

59 - if kv_tranfer_config and kv_tranfer_config.kv_connector:

59 + if kv_transfer_config and kv_transfer_config.kv_connector:

60 self.connector_cls = KVConnectorFactory.get_connector_class(

61 - kv_tranfer_config

61 + kv_transfer_config

62 )

63 self.reset()

64

@@ -120,7 +120,7 @@ class KVConnectorPromMetrics:

120 vllm_config: VllmConfig,

121 metric_types: dict[type[PromMetric], type[PromMetricT]],

122 labelnames: list[str],

123 - per_engine_labelvalues: dict[int, list[str]],

123 + per_engine_labelvalues: dict[int, list[object]],

124 ):

125 self._kv_transfer_config = vllm_config.kv_transfer_config

126 self._gauge_cls = metric_types[Gauge]

@@ -129,7 +129,7 @@ class KVConnectorPromMetrics:

129 self._labelnames = labelnames

130 self._per_engine_labelvalues = per_engine_labelvalues

131

132 - def make_per_engine(self, metric: PromMetric) -> PromMetric:

132 + def make_per_engine(self, metric: PromMetric) -> dict[int, PromMetric]:

133 """

134 Create a per-engine child of a prometheus_client.Metric with

135 the appropriate labels set. The parent metric must be created

@@ -165,7 +165,7 @@ class KVConnectorPrometheus:

165 self,

166 vllm_config: VllmConfig,

167 labelnames: list[str],

168 - per_engine_labelvalues: dict[int, list[str]],

168 + per_engine_labelvalues: dict[int, list[object]],

169 ):

170 self.prom_metrics: KVConnectorPromMetrics | None = None

171 kv_transfer_config = vllm_config.kv_transfer_config

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+2 -2

@@ -85,7 +85,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):

85 vllm_config: "VllmConfig",

86 metric_types: dict[type[PromMetric], type[PromMetricT]],

87 labelnames: list[str],

88 - per_engine_labelvalues: dict[int, list[str]],

88 + per_engine_labelvalues: dict[int, list[object]],

89 prom_metrics: dict[str, KVConnectorPromMetrics],

90 ):

91 super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)

@@ -434,7 +434,7 @@ class MultiConnector(KVConnectorBase_V1):

434 vllm_config: "VllmConfig",

435 metric_types: dict[type["PromMetric"], type["PromMetricT"]],

436 labelnames: list[str],

437 - per_engine_labelvalues: dict[int, list[str]],

437 + per_engine_labelvalues: dict[int, list[object]],

438 ) -> KVConnectorPromMetrics:

439 prom_metrics: dict[str, KVConnectorPromMetrics] = {}

440 for connector_cls, temp_config in cls._get_connector_classes_and_configs(

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+4 -4

@@ -288,7 +288,7 @@ class NixlConnector(KVConnectorBase_V1):

288 vllm_config: VllmConfig,

289 metric_types: dict[type[PromMetric], type[PromMetricT]],

290 labelnames: list[str],

291 - per_engine_labelvalues: dict[int, list[str]],

291 + per_engine_labelvalues: dict[int, list[object]],

292 ) -> KVConnectorPromMetrics:

293 return NixlPromMetrics(

294 vllm_config, metric_types, labelnames, per_engine_labelvalues

@@ -2345,9 +2345,9 @@ class NixlKVConnectorStats(KVConnectorStats):

2345 return {

2346 "Num successful transfers": n,

2347 "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),

2348 - "P90 xfer time (ms)": round(np.percentile(xfer_time, 90) * 1e3, 3),

2348 + "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),

2349 "Avg post time (ms)": round(post_time.mean() * 1e3, 3),

2350 - "P90 post time (ms)": round(np.percentile(post_time, 90) * 1e3, 3),

2350 + "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),

2351 "Avg MB per transfer": round(avg_mb, 3),

2352 "Throughput (MB/s)": round(throughput_mb_s, 3),

2353 "Avg number of descriptors": round(descs.mean(), 1),

@@ -2364,7 +2364,7 @@ class NixlPromMetrics(KVConnectorPromMetrics):

2364 vllm_config: VllmConfig,

2365 metric_types: dict[type[PromMetric], type[PromMetricT]],

2366 labelnames: list[str],

2367 - per_engine_labelvalues: dict[int, list[str]],

2367 + per_engine_labelvalues: dict[int, list[object]],

2368 ):

2369 super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)

2370

vllm/engine/arg_utils.py

+3 -1

@@ -1954,7 +1954,9 @@ class EngineArgs:

1954 self.enable_prefix_caching = False

1955

1956 def _set_default_max_num_seqs_and_batched_tokens_args(

1957 - self, usage_context: UsageContext, model_config: ModelConfig

1957 + self,

1958 + usage_context: UsageContext | None,

1959 + model_config: ModelConfig,

1960 ):

1961 world_size = self.pipeline_parallel_size * self.tensor_parallel_size

1962 (

vllm/transformers_utils/config.py

+2 -2

@@ -614,12 +614,12 @@ def _maybe_remap_hf_config_attrs(config: PretrainedConfig) -> PretrainedConfig:

614

615 def maybe_override_with_speculators(

616 model: str,

617 - tokenizer: str,

617 + tokenizer: str | None,

618 trust_remote_code: bool,

619 revision: str | None = None,

620 vllm_speculative_config: dict[str, Any] | None = None,

621 **kwargs,

622 -) -> tuple[str, str, dict[str, Any] | None]:

622 +) -> tuple[str, str | None, dict[str, Any] | None]:

623 """

624 Resolve model configuration when speculators are detected.

625

vllm/triton_utils/__init__.py

+2 -1

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from typing import TYPE_CHECKING

4

5 from vllm.triton_utils.importing import (

6 HAS_TRITON,

@@ -7,7 +8,7 @@ from vllm.triton_utils.importing import (

8 TritonPlaceholder,

9 )

10

10 -if HAS_TRITON:

11 +if TYPE_CHECKING or HAS_TRITON:

12 import triton

13 import triton.language as tl

14 import triton.language.extra.libdevice as tldevice

vllm/v1/metrics/loggers.py

+4 -4

@@ -104,8 +104,8 @@ class LoggingStatLogger(StatLoggerBase):

104 self.mm_caching_metrics = CachingMetrics()

105

106 self.spec_decoding_logging = SpecDecodingLogging()

107 - kv_tranfer_config = self.vllm_config.kv_transfer_config

108 - self.kv_connector_logging = KVConnectorLogging(kv_tranfer_config)

107 + kv_transfer_config = self.vllm_config.kv_transfer_config

108 + self.kv_connector_logging = KVConnectorLogging(kv_transfer_config)

109 self.last_prompt_throughput: float = 0.0

110 self.last_generation_throughput: float = 0.0

111 self.engine_is_idle = False

@@ -380,7 +380,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase):

380 model_name = vllm_config.model_config.served_model_name

381 max_model_len = vllm_config.model_config.max_model_len

382

383 - per_engine_labelvalues: dict[int, list[str]] = {

383 + per_engine_labelvalues: dict[int, list[object]] = {

384 idx: [model_name, str(idx)] for idx in engine_indexes

385 }

386

@@ -1052,7 +1052,7 @@ PromMetric: TypeAlias = Gauge | Counter | Histogram

1052

1053

1054 def make_per_engine(

1055 - metric: PromMetric, engine_idxs: list[int], model_name: str

1055 + metric: PromMetric, engine_idxs: list[int], model_name: object

1056 ) -> dict[int, PromMetric]:

1057 return {idx: metric.labels(model_name, str(idx)) for idx in engine_idxs}

1058

vllm/v1/sample/logits_processor/__init__.py

+1 -1

@@ -313,7 +313,7 @@ class AdapterLogitsProcessor(LogitsProcessor):

313 if (len(inspect.signature(req_lp).parameters) == 3)

314 else [output_ids]

315 )

316 - return partial(req_lp, *args)

316 + return partial(req_lp, *args) # type: ignore[misc]

317 return None

318

319 def update_state(self, batch_update: BatchUpdate | None):

vllm/v1/spec_decode/metrics.py

+3 -2

@@ -144,7 +144,7 @@ class SpecDecodingProm:

144 self,

145 speculative_config: SpeculativeConfig | None,

146 labelnames: list[str],

147 - per_engine_labelvalues: dict[int, list[str]],

147 + per_engine_labelvalues: dict[int, list[object]],

148 ):

149 self.spec_decoding_enabled = speculative_config is not None

150 if not self.spec_decoding_enabled:

@@ -215,7 +215,8 @@ class SpecDecodingProm:

215

216

217 def make_per_engine(

218 - counter: prometheus_client.Counter, per_engine_labelvalues: dict[int, list[str]]

218 + counter: prometheus_client.Counter,

219 + per_engine_labelvalues: dict[int, list[object]],

220 ):

221 """Create a counter for each label value."""

222 return {

[mypy] Enable type checking for more directories (#29674)