[mypy] Enable type checking for more directories (#29674)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
@@ -565,7 +565,7 @@ class KVConnectorBase_V1(ABC):
|
||||
vllm_config: "VllmConfig",
|
||||
metric_types: dict[type["PromMetric"], type["PromMetricT"]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
) -> Optional["KVConnectorPromMetrics"]:
|
||||
"""
|
||||
Create a KVConnectorPromMetrics subclass which should register
|
||||
|
||||
@@ -806,7 +806,7 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
vllm_config: "VllmConfig",
|
||||
metric_types: dict[type["PromMetric"], type["PromMetricT"]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
) -> Optional["KVConnectorPromMetrics"]:
|
||||
"""
|
||||
Create a KVConnectorPromMetrics subclass which should register
|
||||
|
||||
@@ -52,13 +52,13 @@ class KVConnectorStats:
|
||||
|
||||
|
||||
class KVConnectorLogging:
|
||||
def __init__(self, kv_tranfer_config: KVTransferConfig):
|
||||
def __init__(self, kv_transfer_config: KVTransferConfig | None):
|
||||
# This should be called on frontend process.
|
||||
assert not has_kv_transfer_group()
|
||||
# Instantiate the connector's stats class.
|
||||
if kv_tranfer_config and kv_tranfer_config.kv_connector:
|
||||
if kv_transfer_config and kv_transfer_config.kv_connector:
|
||||
self.connector_cls = KVConnectorFactory.get_connector_class(
|
||||
kv_tranfer_config
|
||||
kv_transfer_config
|
||||
)
|
||||
self.reset()
|
||||
|
||||
@@ -120,7 +120,7 @@ class KVConnectorPromMetrics:
|
||||
vllm_config: VllmConfig,
|
||||
metric_types: dict[type[PromMetric], type[PromMetricT]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
):
|
||||
self._kv_transfer_config = vllm_config.kv_transfer_config
|
||||
self._gauge_cls = metric_types[Gauge]
|
||||
@@ -129,7 +129,7 @@ class KVConnectorPromMetrics:
|
||||
self._labelnames = labelnames
|
||||
self._per_engine_labelvalues = per_engine_labelvalues
|
||||
|
||||
def make_per_engine(self, metric: PromMetric) -> PromMetric:
|
||||
def make_per_engine(self, metric: PromMetric) -> dict[int, PromMetric]:
|
||||
"""
|
||||
Create a per-engine child of a prometheus_client.Metric with
|
||||
the appropriate labels set. The parent metric must be created
|
||||
@@ -165,7 +165,7 @@ class KVConnectorPrometheus:
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
):
|
||||
self.prom_metrics: KVConnectorPromMetrics | None = None
|
||||
kv_transfer_config = vllm_config.kv_transfer_config
|
||||
|
||||
@@ -85,7 +85,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
|
||||
vllm_config: "VllmConfig",
|
||||
metric_types: dict[type[PromMetric], type[PromMetricT]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
prom_metrics: dict[str, KVConnectorPromMetrics],
|
||||
):
|
||||
super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
|
||||
@@ -434,7 +434,7 @@ class MultiConnector(KVConnectorBase_V1):
|
||||
vllm_config: "VllmConfig",
|
||||
metric_types: dict[type["PromMetric"], type["PromMetricT"]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
) -> KVConnectorPromMetrics:
|
||||
prom_metrics: dict[str, KVConnectorPromMetrics] = {}
|
||||
for connector_cls, temp_config in cls._get_connector_classes_and_configs(
|
||||
|
||||
@@ -288,7 +288,7 @@ class NixlConnector(KVConnectorBase_V1):
|
||||
vllm_config: VllmConfig,
|
||||
metric_types: dict[type[PromMetric], type[PromMetricT]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
) -> KVConnectorPromMetrics:
|
||||
return NixlPromMetrics(
|
||||
vllm_config, metric_types, labelnames, per_engine_labelvalues
|
||||
@@ -2345,9 +2345,9 @@ class NixlKVConnectorStats(KVConnectorStats):
|
||||
return {
|
||||
"Num successful transfers": n,
|
||||
"Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
|
||||
"P90 xfer time (ms)": round(np.percentile(xfer_time, 90) * 1e3, 3),
|
||||
"P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
|
||||
"Avg post time (ms)": round(post_time.mean() * 1e3, 3),
|
||||
"P90 post time (ms)": round(np.percentile(post_time, 90) * 1e3, 3),
|
||||
"P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),
|
||||
"Avg MB per transfer": round(avg_mb, 3),
|
||||
"Throughput (MB/s)": round(throughput_mb_s, 3),
|
||||
"Avg number of descriptors": round(descs.mean(), 1),
|
||||
@@ -2364,7 +2364,7 @@ class NixlPromMetrics(KVConnectorPromMetrics):
|
||||
vllm_config: VllmConfig,
|
||||
metric_types: dict[type[PromMetric], type[PromMetricT]],
|
||||
labelnames: list[str],
|
||||
per_engine_labelvalues: dict[int, list[str]],
|
||||
per_engine_labelvalues: dict[int, list[object]],
|
||||
):
|
||||
super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user