A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][KVConnector][Mooncake] Close MooncakeDistributedStore on connector teardown (#45206)
Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Dao007forever committed
Jun 11, 2026 at 14:33 UTC
8a91228dbe363d1d113deb2a82e289429130dd01
4 files changed
+125
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+63
@@ -614,3 +614,66 @@ def test_lookup_key_server_reset_skips_drain_when_no_send_thread():
614
615
assert call_order == ["remove_all"]
616
assert sent == [protocol.RESP_OK]
617
+
618
+
619
+def test_shutdown_closes_worker_store():
620
+ vllm_config = _make_vllm_config()
621
+ kv_cache_config = _make_kv_cache_config()
622
+
623
+ with (
624
+ set_current_vllm_config(vllm_config),
625
+ patch(
626
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
627
+ "connector.MooncakeStoreWorker"
628
+ ) as mock_worker_cls,
629
+ ):
630
+ connector = mooncake_store_connector.MooncakeStoreConnector(
631
+ vllm_config, KVConnectorRole.WORKER, kv_cache_config
632
+ )
633
+
634
+ worker = mock_worker_cls.return_value
635
+ connector.shutdown()
636
+
637
+ worker.close.assert_called_once_with()
638
+
639
+
640
+def test_del_invokes_shutdown_and_closes_store():
641
+ vllm_config = _make_vllm_config()
642
+ kv_cache_config = _make_kv_cache_config()
643
+
644
+ with (
645
+ set_current_vllm_config(vllm_config),
646
+ patch(
647
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
648
+ "connector.MooncakeStoreWorker"
649
+ ) as mock_worker_cls,
650
+ ):
651
+ connector = mooncake_store_connector.MooncakeStoreConnector(
652
+ vllm_config, KVConnectorRole.WORKER, kv_cache_config
653
+ )
654
+
655
+ worker = mock_worker_cls.return_value
656
+ # __del__ is the GC backstop; it must route through shutdown() -> close().
657
+ connector.__del__()
658
+
659
+ worker.close.assert_called_once_with()
660
+
661
+
662
+def test_shutdown_scheduler_role_is_noop():
663
+ vllm_config = _make_vllm_config()
664
+ kv_cache_config = _make_kv_cache_config()
665
+
666
+ with (
667
+ set_current_vllm_config(vllm_config),
668
+ patch(
669
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
670
+ "connector.MooncakeStoreScheduler"
671
+ ),
672
+ ):
673
+ connector = mooncake_store_connector.MooncakeStoreConnector(
674
+ vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
675
+ )
676
+
677
+ # Scheduler role holds no store handle, so shutdown must be a safe no-op.
678
+ assert connector.connector_worker is None
679
+ connector.shutdown()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+31
@@ -1558,3 +1558,34 @@ def test_lookup_records_mooncake_metrics():
1558
assert isinstance(stats, MooncakeStoreConnectorStats)
1559
assert len(stats.data["lookup_exists"]) == 1
1560
assert stats.data["lookup_exists"][0]["num_keys"] == 2
1561
+
1562
+
1563
+def test_store_worker_close_releases_store():
1564
+ worker = _make_bare_worker()
1565
+ store = worker.store
1566
+
1567
+ worker.close()
1568
+
1569
+ store.close.assert_called_once_with()
1570
+ assert worker.store is None
1571
+
1572
+
1573
+def test_store_worker_close_is_idempotent():
1574
+ worker = _make_bare_worker()
1575
+ store = worker.store
1576
+
1577
+ worker.close()
1578
+ worker.close()
1579
+
1580
+ # Second call short-circuits because store was already released.
1581
+ store.close.assert_called_once_with()
1582
+
1583
+
1584
+def test_store_worker_close_swallows_store_errors():
1585
+ worker = _make_bare_worker()
1586
+ worker.store.close.side_effect = RuntimeError("boom")
1587
+
1588
+ # A failure tearing down the store must not propagate out of close().
1589
+ worker.close()
1590
+
1591
+ assert worker.store is None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+15
@@ -153,6 +153,21 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
153
else:
154
self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config)
155
156
+ def shutdown(self):
157
+ """Release connector resources on teardown.
158
+
159
+ Closes the worker's MooncakeDistributedStore handle so its
160
+ TransferEngine and RDMA registrations are released. Invoked from the
161
+ engine's explicit shutdown path and as a backstop from ``__del__``;
162
+ a no-op on the scheduler role, which holds no store handle.
163
+ """
164
+ worker = getattr(self, "connector_worker", None)
165
+ if worker is not None:
166
+ worker.close()
167
+
168
+ def __del__(self):
169
+ self.shutdown()
170
+
171
# ============================================================
172
# Scheduler-side methods
173
# ============================================================
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+16
@@ -1426,6 +1426,22 @@ class MooncakeStoreWorker:
1426
return self.kv_send_thread.get_kv_events()
1427
return []
1428
1429
+ def close(self) -> None:
1430
+ """Release the MooncakeDistributedStore handle on teardown.
1431
+
1432
+ Closing the store frees its TransferEngine, the registered RDMA
1433
+ buffers, and the connection to the master server. Idempotent so it is
1434
+ safe to call from both the explicit shutdown path and ``__del__``.
1435
+ """
1436
+ store = getattr(self, "store", None)
1437
+ if store is None:
1438
+ return
1439
+ self.store = None
1440
+ try:
1441
+ store.close()
1442
+ except Exception as e:
1443
+ logger.warning("Error closing MooncakeDistributedStore: %s", e)
1444
+
1445
1446
# ============================================================
1447
# Lookup Key Server