[Bugfix][KVConnector][Mooncake] Close MooncakeDistributedStore on connector teardown (#45206)

Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

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()

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

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
     # ============================================================

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