[Bugfix] Fix np.float64 leaking into the KV transfer metrics log (#57068)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Sep 16, 2026 at 20:09 UTC
fc8132a5e523294ea69decb5e7cbf50e9d3e135e
7 files changed
+63
-14
tests/v1/kv_connector/unit/test_hf3fs_connector.py
+4
@@ -8,6 +8,7 @@
8
import os
9
from unittest.mock import MagicMock
10
11
+import numpy as np
12
import pytest
13
import torch
14
@@ -218,6 +219,9 @@ class TestHF3FSKVConnectorStats:
219
assert result["Num save task success"] == pytest.approx(2.0, rel=0.01)
220
assert result["Num save task failed"] == pytest.approx(0.0, rel=0.01)
221
assert result["Avg save duration (ms)"] == pytest.approx(2000.0, rel=0.01)
222
+ # Reduced values must be plain Python scalars so CLI logging renders
223
+ # them without numpy reprs (eg np.float64(...)).
224
+ assert all(not isinstance(v, np.generic) for v in result.values())
225
226
def test_clone_and_reset(self, hf3fs_stats):
227
"""clone_and_reset() returns a copy with data and resets the original."""
tests/v1/kv_connector/unit/test_mooncake_stats.py
+5
@@ -4,6 +4,8 @@
4
import threading
5
from unittest.mock import MagicMock
6
7
+import numpy as np
8
+
9
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
10
MooncakeConnector,
11
MooncakeConnectorWorker,
@@ -40,6 +42,9 @@ def test_record_transfer_and_reduce():
42
assert reduced["Num failed transfers"] == 0
43
assert reduced["Num failed recvs"] == 0
44
assert reduced["Num KV expired reqs"] == 0
45
+ # Reduced values must be plain Python scalars so CLI logging renders
46
+ # them without numpy reprs (eg np.float64(...)).
47
+ assert all(not isinstance(v, np.generic) for v in reduced.values())
48
49
50
def test_record_failures_keeps_stats_non_empty():
tests/v1/kv_connector/unit/test_multi_connector.py
+34
-1
@@ -20,7 +20,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
20
SupportsHMA,
21
supports_hma,
22
)
23
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
23
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
24
+ KVConnectorLogging,
25
+ KVConnectorStats,
26
+)
27
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
28
MultiConnector,
29
MultiKVConnectorPromMetrics,
@@ -817,6 +820,36 @@ class TestMultiConnectorStats:
820
assert "Num successful transfers" in reduced["NixlConnector"]
821
assert reduced["NixlConnector"]["Num successful transfers"] == 2
822
823
+ def test_log_renders_plain_scalars(self):
824
+ """Reduced stats must render as plain numbers in the CLI log, not
825
+ numpy reprs (eg np.float64(2.338))."""
826
+ stats = MultiKVConnectorStats(
827
+ data={
828
+ "NixlConnector": NixlKVConnectorStats(
829
+ data={
830
+ "transfer_duration": [1.0, 2.0],
831
+ "post_duration": [0.1, 0.2],
832
+ "bytes_transferred": [1024, 2048],
833
+ "num_descriptors": [10, 20],
834
+ "num_failed_transfers": [],
835
+ "num_failed_notifications": [],
836
+ }
837
+ )
838
+ }
839
+ )
840
+
841
+ kv_logging = KVConnectorLogging(kv_transfer_config=None)
842
+ kv_logging.transfer_stats_accumulator = stats
843
+ log_records: list[tuple] = []
844
+ kv_logging.log(log_fn=lambda *args: log_records.append(args))
845
+
846
+ assert len(log_records) == 1
847
+ msg = log_records[0][1]
848
+ assert "np.float" not in msg
849
+ assert "'Avg xfer time (ms)': 1500.0" in msg
850
+ # The accumulator is reset after logging.
851
+ assert kv_logging.transfer_stats_accumulator is None
852
+
853
def test_reset(self):
854
"""Test that reset() resets all nested connector stats."""
855
stats = MultiKVConnectorStats(
tests/v1/kv_connector/unit/test_nixl_connector.py
+3
@@ -1676,6 +1676,9 @@ def test_kv_connector_stats_aggregation():
1676
assert cli_stats["Avg xfer time (ms)"] == 1500.0
1677
assert cli_stats["Avg post time (ms)"] == 1500.0
1678
assert cli_stats["Avg number of descriptors"] == 1.5
1679
+ # Reduced values must be plain Python scalars so CLI logging renders
1680
+ # them without numpy reprs (eg np.float64(...)).
1681
+ assert all(not isinstance(v, np.generic) for v in cli_stats.values())
1682
1683
1684
def test_multi_kv_connector_stats_aggregation():
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
+8
-4
@@ -1075,10 +1075,14 @@ class HF3FSKVConnectorStats(KVConnectorStats):
1075
"Num save task failed": num_failed_save,
1076
"Num load task success": num_success_load,
1077
"Num load task failed": num_failed_load,
1078
- "Avg save duration (ms)": round(save_duration.mean() * 1e3, 3),
1079
- "P90 save duration (ms)": round(np.percentile(save_duration, 90) * 1e3, 3),
1080
- "Avg load duration (ms)": round(load_duration.mean() * 1e3, 3),
1081
- "P90 load duration (ms)": round(np.percentile(load_duration, 90) * 1e3, 3),
1078
+ "Avg save duration (ms)": round(save_duration.mean().item() * 1e3, 3),
1079
+ "P90 save duration (ms)": round(
1080
+ np.percentile(save_duration, 90).item() * 1e3, 3
1081
+ ),
1082
+ "Avg load duration (ms)": round(load_duration.mean().item() * 1e3, 3),
1083
+ "P90 load duration (ms)": round(
1084
+ np.percentile(load_duration, 90).item() * 1e3, 3
1085
+ ),
1086
}
1087
1088
def is_empty(self) -> bool:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py
+4
-4
@@ -126,16 +126,16 @@ class MooncakeKVConnectorStats(KVConnectorStats):
126
avg_mb = total_mb / n
127
total_time_seconds = xfer_time.sum()
128
throughput_mb_s = (
129
- total_mb / total_time_seconds if total_time_seconds > 0 else 0.0
129
+ (total_mb / total_time_seconds).item() if total_time_seconds > 0 else 0.0
130
)
131
132
return {
133
"Num successful transfers": n,
134
- "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
134
+ "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3),
135
"P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
136
- "Avg MB per transfer": round(avg_mb, 3),
136
+ "Avg MB per transfer": round(avg_mb.item(), 3),
137
"Throughput (MB/s)": round(throughput_mb_s, 3),
138
- "Avg number of descriptors": round(descs.mean(), 1),
138
+ "Avg number of descriptors": round(descs.mean().item(), 1),
139
"Num failed transfers": num_failed_transfers,
140
"Num failed recvs": num_failed_recvs,
141
"Num KV expired reqs": num_kv_expired_reqs,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
+5
-5
@@ -115,13 +115,13 @@ class NixlKVConnectorStats(KVConnectorStats):
115
116
return {
117
"Num successful transfers": n,
118
- "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
118
+ "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3),
119
"P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
120
- "Avg post time (ms)": round(post_time.mean() * 1e3, 3),
120
+ "Avg post time (ms)": round(post_time.mean().item() * 1e3, 3),
121
"P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3),
122
- "Avg MB per transfer": round(avg_mb, 3),
123
- "Throughput (MB/s)": round(throughput_mb_s, 3),
124
- "Avg number of descriptors": round(descs.mean(), 1),
122
+ "Avg MB per transfer": round(avg_mb.item(), 3),
123
+ "Throughput (MB/s)": round(throughput_mb_s.item(), 3),
124
+ "Avg number of descriptors": round(descs.mean().item(), 1),
125
}
126
127
@property