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
[P/D] Support NIXL connector to disconnect during a clean shutdown (#24423)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>
Chauncey committed
Sep 23, 2025 at 22:08 UTC
f05a4f0e345bbfd4a7cb3f421bd9412e1cc53e74
2 files changed
+80
-7
tests/v1/kv_connector/unit/test_nixl_connector.py
+52
@@ -60,6 +60,9 @@ class FakeNixlWrapper:
60
def register_memory(self, descs, backends) -> None:
61
pass
62
63
+ def deregister_memory(self, descs) -> None:
64
+ pass
65
+
66
def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
67
return [str(uuid.uuid4()) for _ in blocks_data]
68
@@ -86,6 +89,12 @@ class FakeNixlWrapper:
89
def release_xfer_handle(self, handle: int) -> None:
90
pass
91
92
+ def release_dlist_handle(self, handle: int) -> None:
93
+ pass
94
+
95
+ def remove_remote_agent(self, agent: str) -> None:
96
+ pass
97
+
98
def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
99
pass
100
@@ -905,3 +914,46 @@ def test_kv_buffer_to_nixl_memory_types(dist_init, kv_buffer_device,
914
# Verify get_reg_descs was called with the correct memory_type
915
assert connector.connector_worker.kv_buffer_device == kv_buffer_device
916
assert connector.connector_worker.nixl_memory_type == nixl_memory_type
917
+
918
+
919
+@patch(
920
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
921
+ FakeNixlWrapper)
922
+def test_shutdown_cleans_up_resources(dist_init):
923
+ """Test that shutdown() properly cleans up all resources."""
924
+ vllm_config = create_vllm_config()
925
+
926
+ worker = NixlConnectorWorker(vllm_config,
927
+ vllm_config.kv_transfer_config.engine_id)
928
+ nixl_wrapper = worker.nixl_wrapper
929
+
930
+ with patch.object(worker, '_handshake_initiation_executor') as mock_exec, \
931
+ patch.object(worker, '_nixl_handshake_listener_t') as mock_listener, \
932
+ patch.object(nixl_wrapper, 'release_xfer_handle') as mock_rel_xfer, \
933
+ patch.object(nixl_wrapper, 'release_dlist_handle') as mock_rel_dlist, \
934
+ patch.object(nixl_wrapper, 'remove_remote_agent') as mock_rem_agent, \
935
+ patch.object(nixl_wrapper, 'deregister_memory') as mock_dereg:
936
+
937
+ worker._recving_transfers = {"req1": [(123, time.perf_counter())]}
938
+ worker.src_xfer_side_handle = 456
939
+ worker.dst_xfer_side_handles = {"engine1": 789}
940
+ worker._remote_agents = {"engine1": {0: "agent1"}}
941
+ worker._registered_descs = ["desc1", "desc2"]
942
+
943
+ worker.shutdown()
944
+
945
+ # Test idempotency
946
+ worker.shutdown()
947
+ worker.shutdown()
948
+
949
+ mock_exec.shutdown.assert_called_with(wait=False)
950
+ mock_listener.join.assert_called_once_with(timeout=0)
951
+
952
+ mock_rel_xfer.assert_called_once_with(123)
953
+ assert mock_rel_dlist.call_count == 2
954
+ mock_rel_dlist.assert_any_call(456) # src handle
955
+ mock_rel_dlist.assert_any_call(789) # dst handle
956
+ mock_rem_agent.assert_called_once_with("agent1")
957
+ assert mock_dereg.call_count == 2
958
+ mock_dereg.assert_any_call("desc1")
959
+ mock_dereg.assert_any_call("desc2")
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+28
-7
@@ -250,6 +250,10 @@ class NixlConnector(KVConnectorBase_V1):
250
self.connector_worker.copy_blocks:
251
self.connector_worker.save_kv_to_host(self._connector_metadata)
252
253
+ def shutdown(self):
254
+ if self.connector_worker is not None:
255
+ self.connector_worker.shutdown()
256
+
257
258
class NixlConnectorScheduler:
259
"""Implementation of Scheduler side methods"""
@@ -586,13 +590,6 @@ class NixlConnectorWorker:
590
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
591
self.xfer_stats = NixlKVConnectorStats()
592
589
- def __del__(self):
590
- """Cleanup background threads on destruction."""
591
- if executor := getattr(self, "_handshake_initiation_executor", None):
592
- executor.shutdown(wait=False)
593
- if listener_t := getattr(self, "_nixl_handshake_listener_t", None):
594
- listener_t.join(timeout=0)
595
-
593
@staticmethod
594
def _nixl_handshake_listener(metadata: NixlAgentMetadata,
595
ready_event: threading.Event, base_port: int,
@@ -1346,6 +1343,30 @@ class NixlConnectorWorker:
1343
return self.xfer_stats.clone_and_reset()
1344
return None
1345
1346
+ def shutdown(self):
1347
+ """Shutdown the connector worker."""
1348
+ self._handshake_initiation_executor.shutdown(wait=False)
1349
+ if self._nixl_handshake_listener_t is not None:
1350
+ self._nixl_handshake_listener_t.join(timeout=0)
1351
+ self._nixl_handshake_listener_t = None
1352
+ for handles in self._recving_transfers.values():
1353
+ for handle, _ in handles:
1354
+ self.nixl_wrapper.release_xfer_handle(handle)
1355
+ self._recving_transfers.clear()
1356
+ if self.src_xfer_side_handle:
1357
+ self.nixl_wrapper.release_dlist_handle(self.src_xfer_side_handle)
1358
+ self.src_xfer_side_handle = 0
1359
+ for dst_xfer_side_handle in self.dst_xfer_side_handles.values():
1360
+ self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle)
1361
+ self.dst_xfer_side_handles.clear()
1362
+ for remote_agents in self._remote_agents.values():
1363
+ for agent_name in remote_agents.values():
1364
+ self.nixl_wrapper.remove_remote_agent(agent_name)
1365
+ self._remote_agents.clear()
1366
+ for desc in self._registered_descs:
1367
+ self.nixl_wrapper.deregister_memory(desc)
1368
+ self._registered_descs.clear()
1369
+
1370
1371
@contextlib.contextmanager
1372
def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: