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, 05:30 UTC
[Bugfix] Fix MultiConnector stats reconstruction across process boundaries (#27366)
Signed-off-by: Kourosh Hakhamaneshi <Kourosh@anyscale.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
kourosh hakhamaneshi committed
Oct 24, 2025 at 10:08 UTC
7e1d697b568d4da27a013e2f3e4016c2e23d3fc1
3 files changed
+421
-5
tests/v1/kv_connector/unit/test_multi_connector.py
+368
@@ -4,9 +4,22 @@ import filecmp
4
import shutil
5
import tempfile
6
from pathlib import Path
7
+from typing import Any
8
+
9
+import pytest
10
11
from vllm import LLM, SamplingParams
12
from vllm.config import KVTransferConfig
13
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
14
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
15
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
16
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
17
+ MultiConnector,
18
+ MultiKVConnectorStats,
19
+)
20
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
21
+ NixlKVConnectorStats,
22
+)
23
24
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
25
@@ -19,6 +32,27 @@ PROMPTS = [
32
SAMPLING_PARAMS = SamplingParams(temperature=0, max_tokens=20)
33
34
35
+# Test connector with custom stats for testing MultiConnector
36
+class MockConnectorStats(KVConnectorStats):
37
+ """Mock stats class for testing."""
38
+
39
+ pass
40
+
41
+
42
+class MockConnector(KVConnectorBase_V1):
43
+ """Mock connector that implements build_kv_connector_stats for testing."""
44
+
45
+ @classmethod
46
+ def build_kv_connector_stats(
47
+ cls, data: dict[str, Any] | None = None
48
+ ) -> KVConnectorStats | None:
49
+ return MockConnectorStats(data=data) if data is not None else None
50
+
51
+
52
+# Register the mock connector
53
+KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
54
+
55
+
56
# Helper function to compare directories recursively
57
def _compare_directories(dir1: Path, dir2: Path) -> bool:
58
"""Compares two directories recursively for identical content."""
@@ -225,3 +259,337 @@ def test_engine_id_conflict():
259
assert ids[0] != ids[1], (
260
f"Engine IDs should be different for different configs. Got {ids}"
261
)
262
+
263
+
264
+class TestMultiConnectorStats:
265
+ """Tests for MultiConnector stats reconstruction and operations."""
266
+
267
+ def test_build_kv_connector_stats_with_none(self):
268
+ """Test that build_kv_connector_stats returns empty stats when given None."""
269
+ stats = MultiConnector.build_kv_connector_stats(data=None)
270
+
271
+ assert stats is not None
272
+ assert isinstance(stats, MultiKVConnectorStats)
273
+ assert len(stats.data) == 0
274
+ assert stats.is_empty()
275
+
276
+ def test_build_kv_connector_stats_with_empty_dict(self):
277
+ """Test that build_kv_connector_stats returns empty stats with empty dict."""
278
+ stats = MultiConnector.build_kv_connector_stats(data={})
279
+
280
+ assert stats is not None
281
+ assert isinstance(stats, MultiKVConnectorStats)
282
+ assert len(stats.data) == 0
283
+ assert stats.is_empty()
284
+
285
+ def test_build_kv_connector_stats_reconstructs_nixl_stats(self):
286
+ """Test that NixlConnector stats are properly reconstructed with
287
+ correct data."""
288
+ serialized_data = {
289
+ "NixlConnector": {
290
+ "data": {
291
+ "transfer_duration": [1.5, 2.3],
292
+ "post_duration": [0.1, 0.2],
293
+ "bytes_transferred": [1024, 2048],
294
+ "num_descriptors": [10, 20],
295
+ "num_failed_transfers": [],
296
+ "num_failed_notifications": [],
297
+ }
298
+ }
299
+ }
300
+
301
+ stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
302
+
303
+ assert "NixlConnector" in stats.data
304
+ nixl_stats = stats.data["NixlConnector"]
305
+ assert isinstance(nixl_stats, NixlKVConnectorStats)
306
+ assert nixl_stats.data["transfer_duration"] == [1.5, 2.3]
307
+ assert nixl_stats.data["post_duration"] == [0.1, 0.2]
308
+ assert nixl_stats.data["bytes_transferred"] == [1024, 2048]
309
+ assert nixl_stats.data["num_descriptors"] == [10, 20]
310
+
311
+ def test_build_kv_connector_stats_with_multiple_connectors(self):
312
+ """Test reconstruction with multiple connector types that have custom stats."""
313
+ serialized_data = {
314
+ "NixlConnector": {
315
+ "data": {
316
+ "transfer_duration": [1.5],
317
+ "post_duration": [0.1],
318
+ "bytes_transferred": [1024],
319
+ "num_descriptors": [10],
320
+ "num_failed_transfers": [],
321
+ "num_failed_notifications": [],
322
+ }
323
+ },
324
+ "MockConnector": {"data": {"mock_field": [1, 2, 3]}},
325
+ }
326
+
327
+ stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
328
+
329
+ assert stats is not None
330
+ assert isinstance(stats, MultiKVConnectorStats)
331
+ # Both connectors should be reconstructed
332
+ assert len(stats.data) == 2
333
+ assert "NixlConnector" in stats.data
334
+ assert "MockConnector" in stats.data
335
+ assert isinstance(stats.data["NixlConnector"], NixlKVConnectorStats)
336
+ assert isinstance(stats.data["MockConnector"], MockConnectorStats)
337
+ # Verify data is preserved
338
+ assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]}
339
+
340
+ def test_build_kv_connector_stats_raises_error_for_unknown_connector(self):
341
+ """Test that unknown connectors raise an error."""
342
+ serialized_data = {
343
+ "UnknownConnector": {"data": {"some_field": [1, 2, 3]}},
344
+ "NixlConnector": {
345
+ "data": {
346
+ "transfer_duration": [1.5],
347
+ "post_duration": [0.1],
348
+ "bytes_transferred": [1024],
349
+ "num_descriptors": [10],
350
+ "num_failed_transfers": [],
351
+ "num_failed_notifications": [],
352
+ }
353
+ },
354
+ }
355
+
356
+ with pytest.raises(
357
+ ValueError, match="Connector 'UnknownConnector' is not registered."
358
+ ):
359
+ MultiConnector.build_kv_connector_stats(data=serialized_data)
360
+
361
+ def test_build_kv_connector_stats_with_already_instantiated_objects(self):
362
+ """Test that already-instantiated stats objects are preserved (same process)."""
363
+ # This simulates the in-process case where stats are not serialized
364
+ nixl_stats = NixlKVConnectorStats(
365
+ data={
366
+ "transfer_duration": [1.5],
367
+ "post_duration": [0.1],
368
+ "bytes_transferred": [1024],
369
+ "num_descriptors": [10],
370
+ "num_failed_transfers": [],
371
+ "num_failed_notifications": [],
372
+ }
373
+ )
374
+ mock_stats = MockConnectorStats(data={"mock_field": [1, 2, 3]})
375
+
376
+ data_with_objects = {
377
+ "NixlConnector": nixl_stats,
378
+ "MockConnector": mock_stats,
379
+ }
380
+
381
+ stats = MultiConnector.build_kv_connector_stats(data=data_with_objects)
382
+
383
+ assert stats is not None
384
+ assert isinstance(stats, MultiKVConnectorStats)
385
+ assert len(stats.data) == 2
386
+ # Verify objects are preserved as-is
387
+ assert stats.data["NixlConnector"] is nixl_stats
388
+ assert stats.data["MockConnector"] is mock_stats
389
+
390
+ def test_build_kv_connector_stats_with_mixed_objects_and_dicts(self):
391
+ """Test handling mixed already-instantiated and serialized stats."""
392
+ # This can happen during transition or partial serialization
393
+ nixl_stats = NixlKVConnectorStats(
394
+ data={
395
+ "transfer_duration": [1.5],
396
+ "post_duration": [0.1],
397
+ "bytes_transferred": [1024],
398
+ "num_descriptors": [10],
399
+ "num_failed_transfers": [],
400
+ "num_failed_notifications": [],
401
+ }
402
+ )
403
+
404
+ mixed_data = {
405
+ "NixlConnector": nixl_stats, # Already instantiated
406
+ "MockConnector": {"data": {"mock_field": [1, 2, 3]}}, # Serialized
407
+ }
408
+
409
+ stats = MultiConnector.build_kv_connector_stats(data=mixed_data)
410
+
411
+ assert stats is not None
412
+ assert isinstance(stats, MultiKVConnectorStats)
413
+ assert len(stats.data) == 2
414
+ # Instantiated object preserved
415
+ assert stats.data["NixlConnector"] is nixl_stats
416
+ # Serialized object reconstructed
417
+ assert isinstance(stats.data["MockConnector"], MockConnectorStats)
418
+ assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]}
419
+
420
+ def test_build_kv_connector_stats_skips_connectors_without_custom_stats(self):
421
+ """Test that connectors without custom stats (return None) are skipped."""
422
+ # SharedStorageConnector doesn't override build_kv_connector_stats,
423
+ # so it returns None and should be skipped
424
+ serialized_data = {
425
+ "NixlConnector": {
426
+ "data": {
427
+ "transfer_duration": [1.5],
428
+ "post_duration": [0.1],
429
+ "bytes_transferred": [1024],
430
+ "num_descriptors": [10],
431
+ "num_failed_transfers": [],
432
+ "num_failed_notifications": [],
433
+ }
434
+ },
435
+ "SharedStorageConnector": {"data": {"some_field": [1, 2, 3]}},
436
+ }
437
+
438
+ stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
439
+
440
+ assert stats is not None
441
+ assert isinstance(stats, MultiKVConnectorStats)
442
+ # Only NixlConnector should be reconstructed
443
+ assert len(stats.data) == 1
444
+ assert "NixlConnector" in stats.data
445
+ assert isinstance(stats.data["NixlConnector"], NixlKVConnectorStats)
446
+ # SharedStorageConnector should be skipped (returns None)
447
+ assert "SharedStorageConnector" not in stats.data
448
+
449
+ def test_build_kv_connector_stats_handles_malformed_data(self):
450
+ """Test that malformed data raises appropriate errors."""
451
+ serialized_data = {
452
+ "NixlConnector": {"wrong_field": {"transfer_duration": [1.5]}}
453
+ }
454
+
455
+ with pytest.raises(AssertionError, match="Expected a dict with a 'data' field"):
456
+ MultiConnector.build_kv_connector_stats(data=serialized_data)
457
+
458
+ def test_aggregate_same_connector(self):
459
+ """Test aggregating stats from the same connector type."""
460
+ stats1 = MultiKVConnectorStats(
461
+ data={
462
+ "NixlConnector": NixlKVConnectorStats(
463
+ data={
464
+ "transfer_duration": [1.0],
465
+ "post_duration": [0.1],
466
+ "bytes_transferred": [1024],
467
+ "num_descriptors": [10],
468
+ "num_failed_transfers": [],
469
+ "num_failed_notifications": [],
470
+ }
471
+ )
472
+ }
473
+ )
474
+
475
+ stats2 = MultiKVConnectorStats(
476
+ data={
477
+ "NixlConnector": NixlKVConnectorStats(
478
+ data={
479
+ "transfer_duration": [2.0],
480
+ "post_duration": [0.2],
481
+ "bytes_transferred": [2048],
482
+ "num_descriptors": [20],
483
+ "num_failed_transfers": [],
484
+ "num_failed_notifications": [],
485
+ }
486
+ )
487
+ }
488
+ )
489
+
490
+ result = stats1.aggregate(stats2)
491
+
492
+ assert result is stats1 # Should return self
493
+ assert "NixlConnector" in result.data
494
+ nixl_stats = result.data["NixlConnector"]
495
+ assert nixl_stats.data["transfer_duration"] == [1.0, 2.0]
496
+ assert nixl_stats.data["post_duration"] == [0.1, 0.2]
497
+ assert nixl_stats.data["bytes_transferred"] == [1024, 2048]
498
+ assert nixl_stats.data["num_descriptors"] == [10, 20]
499
+
500
+ def test_aggregate_new_connector(self):
501
+ """Test aggregating stats when a new connector type appears."""
502
+ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
503
+ KVConnectorStats,
504
+ )
505
+
506
+ stats1 = MultiKVConnectorStats(
507
+ data={
508
+ "NixlConnector": NixlKVConnectorStats(
509
+ data={
510
+ "transfer_duration": [1.0],
511
+ "post_duration": [0.1],
512
+ "bytes_transferred": [1024],
513
+ "num_descriptors": [10],
514
+ "num_failed_transfers": [],
515
+ "num_failed_notifications": [],
516
+ }
517
+ )
518
+ }
519
+ )
520
+
521
+ stats2 = MultiKVConnectorStats(
522
+ data={"SharedStorageConnector": KVConnectorStats(data={"field": [1, 2]})}
523
+ )
524
+
525
+ result = stats1.aggregate(stats2)
526
+
527
+ assert "NixlConnector" in result.data
528
+ assert "SharedStorageConnector" in result.data
529
+
530
+ def test_reduce(self):
531
+ """Test that reduce() correctly reduces all nested connector stats."""
532
+ stats = MultiKVConnectorStats(
533
+ data={
534
+ "NixlConnector": NixlKVConnectorStats(
535
+ data={
536
+ "transfer_duration": [1.0, 2.0],
537
+ "post_duration": [0.1, 0.2],
538
+ "bytes_transferred": [1024, 2048],
539
+ "num_descriptors": [10, 20],
540
+ "num_failed_transfers": [],
541
+ "num_failed_notifications": [],
542
+ }
543
+ )
544
+ }
545
+ )
546
+
547
+ reduced = stats.reduce()
548
+
549
+ assert "NixlConnector" in reduced
550
+ assert isinstance(reduced["NixlConnector"], dict)
551
+ # Check that the stats were reduced (should have aggregated values)
552
+ assert "Num successful transfers" in reduced["NixlConnector"]
553
+ assert reduced["NixlConnector"]["Num successful transfers"] == 2
554
+
555
+ def test_reset(self):
556
+ """Test that reset() resets all nested connector stats."""
557
+ stats = MultiKVConnectorStats(
558
+ data={
559
+ "NixlConnector": NixlKVConnectorStats(
560
+ data={
561
+ "transfer_duration": [1.0, 2.0],
562
+ "post_duration": [0.1, 0.2],
563
+ "bytes_transferred": [1024, 2048],
564
+ "num_descriptors": [10, 20],
565
+ "num_failed_transfers": [],
566
+ "num_failed_notifications": [],
567
+ }
568
+ )
569
+ }
570
+ )
571
+
572
+ assert not stats.is_empty()
573
+
574
+ stats.reset()
575
+
576
+ # After reset, stats should be empty
577
+ assert stats.is_empty()
578
+ nixl_stats = stats.data["NixlConnector"]
579
+ assert len(nixl_stats.data["transfer_duration"]) == 0
580
+
581
+ def test_is_empty_with_multiple_connectors(self):
582
+ """Test is_empty() returns correct value with multiple connectors."""
583
+ # All empty
584
+ stats = MultiKVConnectorStats(
585
+ data={
586
+ "NixlConnector": NixlKVConnectorStats(data={}),
587
+ }
588
+ )
589
+ # Initialize empty stats
590
+ stats.data["NixlConnector"].reset()
591
+ assert stats.is_empty()
592
+
593
+ # One non-empty
594
+ stats.data["NixlConnector"].data["transfer_duration"].append(1.0)
595
+ assert not stats.is_empty()
vllm/distributed/kv_transfer/kv_connector/factory.py
+18
@@ -66,6 +66,24 @@ class KVConnectorFactory:
66
# We build separately to enforce strict separation
67
return connector_cls(config, role)
68
69
+ @classmethod
70
+ def get_connector_class_by_name(
71
+ cls, connector_name: str
72
+ ) -> type[KVConnectorBaseType]:
73
+ """Get a registered connector class by name.
74
+
75
+ Raises ValueError if the connector is not registered.
76
+
77
+ Args:
78
+ connector_name: Name of the registered connector.
79
+
80
+ Returns:
81
+ The connector class.
82
+ """
83
+ if connector_name not in cls._registry:
84
+ raise ValueError(f"Connector '{connector_name}' is not registered.")
85
+ return cls._registry[connector_name]()
86
+
87
@classmethod
88
def get_connector_class(
89
cls, kv_transfer_config: "KVTransferConfig"
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+35
-5
@@ -324,11 +324,41 @@ class MultiConnector(KVConnectorBase_V1):
324
def build_kv_connector_stats(
325
cls, data: dict[str, Any] | None = None
326
) -> KVConnectorStats | None:
327
- return (
328
- MultiKVConnectorStats(data=data)
329
- if data is not None
330
- else MultiKVConnectorStats()
331
- )
327
+ if data is None:
328
+ return MultiKVConnectorStats()
329
+
330
+ # data is a dict mapping connector name to their stats data.
331
+ # The stats data can be either:
332
+ # 1. Already-instantiated KVConnectorStats objects (same process)
333
+ # 2. Serialized dicts (cross-process after serialization)
334
+ # We need to reconstruct proper KVConnectorStats objects from dicts
335
+ reconstructed_data = {}
336
+ for connector_name, stats_value in data.items():
337
+ # If already a KVConnectorStats object, use it directly
338
+ if isinstance(stats_value, KVConnectorStats):
339
+ reconstructed_data[connector_name] = stats_value
340
+ continue
341
+
342
+ # Otherwise, reconstruct from serialized dict
343
+ # Get the connector class to reconstruct its stats
344
+ connector_cls = KVConnectorFactory.get_connector_class_by_name(
345
+ connector_name
346
+ )
347
+
348
+ # stats_value is the serialized dataclass which contains {'data': {...}}
349
+ # We need to extract the inner 'data' field to avoid double-nesting
350
+ assert isinstance(stats_value, dict) and "data" in stats_value, (
351
+ f"Expected a dict with a 'data' field, got {stats_value}"
352
+ )
353
+ inner_data = stats_value["data"]
354
+
355
+ # Use the connector's build_kv_connector_stats to reconstruct
356
+ if reconstructed_stats := connector_cls.build_kv_connector_stats(
357
+ data=inner_data
358
+ ):
359
+ reconstructed_data[connector_name] = reconstructed_stats
360
+
361
+ return MultiKVConnectorStats(data=reconstructed_data)
362
363
def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:
364
# Group connector stats by connector type.