From 8a91228dbe363d1d113deb2a82e289429130dd01 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Thu, 11 Jun 2026 14:33:48 -0700 Subject: [PATCH] [Bugfix][KVConnector][Mooncake] Close MooncakeDistributedStore on connector teardown (#45206) Signed-off-by: Dao Le Co-authored-by: Claude --- .../unit/test_mooncake_store_connector.py | 63 +++++++++++++++++++ .../unit/test_mooncake_store_worker.py | 31 +++++++++ .../v1/mooncake/store/connector.py | 15 +++++ .../kv_connector/v1/mooncake/store/worker.py | 16 +++++ 4 files changed, 125 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 69593011db..d3992b02b6 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -614,3 +614,66 @@ def test_lookup_key_server_reset_skips_drain_when_no_send_thread(): assert call_order == ["remove_all"] assert sent == [protocol.RESP_OK] + + +def test_shutdown_closes_worker_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + connector.shutdown() + + worker.close.assert_called_once_with() + + +def test_del_invokes_shutdown_and_closes_store(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + # __del__ is the GC backstop; it must route through shutdown() -> close(). + connector.__del__() + + worker.close.assert_called_once_with() + + +def test_shutdown_scheduler_role_is_noop(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreScheduler" + ), + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + + # Scheduler role holds no store handle, so shutdown must be a safe no-op. + assert connector.connector_worker is None + connector.shutdown() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 8cd5e6e535..1130a7d6a7 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1558,3 +1558,34 @@ def test_lookup_records_mooncake_metrics(): assert isinstance(stats, MooncakeStoreConnectorStats) assert len(stats.data["lookup_exists"]) == 1 assert stats.data["lookup_exists"][0]["num_keys"] == 2 + + +def test_store_worker_close_releases_store(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + + store.close.assert_called_once_with() + assert worker.store is None + + +def test_store_worker_close_is_idempotent(): + worker = _make_bare_worker() + store = worker.store + + worker.close() + worker.close() + + # Second call short-circuits because store was already released. + store.close.assert_called_once_with() + + +def test_store_worker_close_swallows_store_errors(): + worker = _make_bare_worker() + worker.store.close.side_effect = RuntimeError("boom") + + # A failure tearing down the store must not propagate out of close(). + worker.close() + + assert worker.store is None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 14d4b381a3..d53cd13c2e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -153,6 +153,21 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): else: self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config) + def shutdown(self): + """Release connector resources on teardown. + + Closes the worker's MooncakeDistributedStore handle so its + TransferEngine and RDMA registrations are released. Invoked from the + engine's explicit shutdown path and as a backstop from ``__del__``; + a no-op on the scheduler role, which holds no store handle. + """ + worker = getattr(self, "connector_worker", None) + if worker is not None: + worker.close() + + def __del__(self): + self.shutdown() + # ============================================================ # Scheduler-side methods # ============================================================ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 9c3ac83e06..105762ccfc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1426,6 +1426,22 @@ class MooncakeStoreWorker: return self.kv_send_thread.get_kv_events() return [] + def close(self) -> None: + """Release the MooncakeDistributedStore handle on teardown. + + Closing the store frees its TransferEngine, the registered RDMA + buffers, and the connection to the master server. Idempotent so it is + safe to call from both the explicit shutdown path and ``__del__``. + """ + store = getattr(self, "store", None) + if store is None: + return + self.store = None + try: + store.close() + except Exception as e: + logger.warning("Error closing MooncakeDistributedStore: %s", e) + # ============================================================ # Lookup Key Server