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 10, 06:30 UTC
[NIXL][Bugfix] Failure logging overhaul + early metadata free on failure (#32031)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jan 12, 2026 at 21:38 UTC
f8bd8394e3d06c18088d3f1a1944ca6001142400
2 files changed
+234
-28
tests/v1/kv_connector/unit/test_nixl_connector.py
+146
@@ -1705,6 +1705,8 @@ class FailingNixlWrapper(FakeNixlWrapper):
1705
self.fail_handshake = False
1706
self.fail_transfer_setup = False
1707
self.fail_send_notif = False
1708
+ self.fail_transfer_state = False # Returns "ERR" state
1709
+ self.fail_transfer_exception = False # Raises exception in check_xfer_state
1710
1711
def add_remote_agent(self, agent_metadata: bytes) -> str:
1712
if self.fail_handshake:
@@ -1739,6 +1741,150 @@ class FailingNixlWrapper(FakeNixlWrapper):
1741
raise RuntimeError("Simulated send_notif failure")
1742
return super().send_notif(agent_name, notif_msg)
1743
1744
+ def check_xfer_state(self, handle: int) -> str:
1745
+ if self.fail_transfer_exception:
1746
+ raise RuntimeError("Simulated check_xfer_state exception")
1747
+ if self.fail_transfer_state:
1748
+ return "ERR" # Bad transfer state
1749
+ return super().check_xfer_state(handle)
1750
+
1751
+
1752
+@patch(
1753
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1754
+ FailingNixlWrapper,
1755
+)
1756
+@pytest.mark.parametrize(
1757
+ "failure_type,wrapper_config,needs_get_finished",
1758
+ [
1759
+ ("transfer_setup_failed", {"fail_transfer_setup": True}, False),
1760
+ ("handshake_failed", {"fail_handshake": True}, False),
1761
+ ("notification_failed", {"fail_send_notif": True}, False),
1762
+ ("transfer_failed", {"fail_transfer_state": True}, True),
1763
+ ("transfer_exception", {"fail_transfer_exception": True}, True),
1764
+ ],
1765
+)
1766
+def test_transfer_failure_logging(
1767
+ default_vllm_config,
1768
+ dist_init,
1769
+ failure_type,
1770
+ wrapper_config,
1771
+ needs_get_finished,
1772
+):
1773
+ """Test that transfer failures are logged with structured context.
1774
+
1775
+ Run with `pytest -sv` to see the log output.
1776
+
1777
+ Covers failure types:
1778
+ - transfer_setup_failed: make_prepped_xfer fails
1779
+ - handshake_failed: add_remote_agent fails during request handshake
1780
+ - notification_failed: send_notif fails
1781
+ - transfer_failed: check_xfer_state returns bad state (e.g., "ERR")
1782
+ - transfer_exception: check_xfer_state raises exception
1783
+ """
1784
+ import logging
1785
+
1786
+ vllm_config = create_vllm_config()
1787
+
1788
+ connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1789
+ connector.connector_worker = FakeNixlConnectorWorker(
1790
+ vllm_config, connector.engine_id, hand_shake_latency=0.0
1791
+ )
1792
+
1793
+ # Configure FailingNixlWrapper to fail in the specified way
1794
+ for key, value in wrapper_config.items():
1795
+ setattr(connector.connector_worker.nixl_wrapper, key, value)
1796
+
1797
+ request_id = f"test_{failure_type}_req"
1798
+
1799
+ # For notification_failed, we need empty local blocks
1800
+ # (full cache hit path to trigger send_notif)
1801
+ local_blocks = [] if failure_type == "notification_failed" else [10, 11, 12]
1802
+ remote_blocks = [20, 21, 22]
1803
+
1804
+ metadata = NixlConnectorMetadata()
1805
+ metadata.add_new_req_to_recv(
1806
+ request_id=request_id,
1807
+ local_block_ids=local_blocks,
1808
+ kv_transfer_params={
1809
+ "remote_block_ids": remote_blocks,
1810
+ "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1811
+ "remote_request_id": f"prefill-{request_id}",
1812
+ "remote_host": "localhost",
1813
+ "remote_port": 1234,
1814
+ "remote_tp_size": 1,
1815
+ },
1816
+ )
1817
+ connector.bind_connector_metadata(metadata)
1818
+
1819
+ dummy_ctx = ForwardContext(
1820
+ no_compile_layers={},
1821
+ attn_metadata={},
1822
+ virtual_engine=0,
1823
+ )
1824
+
1825
+ # Capture logs from the nixl_connector logger specifically
1826
+ # vLLM loggers have propagate=False, so we need to capture directly
1827
+ nixl_logger = logging.getLogger(
1828
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1829
+ )
1830
+ captured_logs: list[logging.LogRecord] = []
1831
+
1832
+ class LogCapture(logging.Handler):
1833
+ def emit(self, record):
1834
+ captured_logs.append(record)
1835
+
1836
+ handler = LogCapture()
1837
+ handler.setLevel(logging.ERROR)
1838
+ nixl_logger.addHandler(handler)
1839
+
1840
+ try:
1841
+ connector.start_load_kv(dummy_ctx)
1842
+ # Process the ready_requests queue (for async handshake)
1843
+ connector.bind_connector_metadata(NixlConnectorMetadata())
1844
+ # Wait for async handshake to complete
1845
+ time.sleep(0.2)
1846
+ connector.start_load_kv(dummy_ctx)
1847
+
1848
+ # For transfer_failed/transfer_exception, the error happens in
1849
+ # get_finished() when checking transfer state
1850
+ if needs_get_finished:
1851
+ connector.get_finished(finished_req_ids=set())
1852
+ finally:
1853
+ nixl_logger.removeHandler(handler)
1854
+
1855
+ # Print logs for manual comparison between commits
1856
+ error_logs = [r for r in captured_logs if r.levelno >= logging.ERROR]
1857
+ print("\n" + "=" * 60)
1858
+ print(f"CAPTURED ERROR LOGS for {failure_type}:")
1859
+ print("=" * 60)
1860
+ for i, record in enumerate(error_logs):
1861
+ print(f"\n--- Log {i + 1} ---")
1862
+ print(f"Message: {record.message}")
1863
+ print("=" * 60 + "\n")
1864
+
1865
+ assert len(error_logs) >= 1, f"Expected at least one error log for {failure_type}"
1866
+
1867
+ # Verify structured logging output (new format)
1868
+ # Check that at least one log matches the expected format
1869
+ all_messages = [r.message for r in error_logs]
1870
+ combined_logs = "\n".join(all_messages)
1871
+
1872
+ assert any("NIXL transfer failure" in msg for msg in all_messages), (
1873
+ f"Expected structured log format with 'NIXL transfer failure' prefix "
1874
+ f"for {failure_type}. Got: {all_messages}"
1875
+ )
1876
+ assert any("failure_type" in msg for msg in all_messages), (
1877
+ f"Expected 'failure_type' in logs. Got: {all_messages}"
1878
+ )
1879
+ assert any("Context:" in msg for msg in all_messages), (
1880
+ f"Expected 'Context:' in logs. Got: {all_messages}"
1881
+ )
1882
+ # Check that the expected failure_type appears in at least one log
1883
+ # Note: handshake_failed also triggers handshake_setup_failed
1884
+ assert failure_type in combined_logs or (
1885
+ failure_type == "handshake_failed" and "handshake_setup_failed" in combined_logs
1886
+ ), f"Expected '{failure_type}' in logs. Got: {all_messages}"
1887
+
1888
1889
@patch(
1890
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+88
-28
@@ -1152,6 +1152,50 @@ class NixlConnectorWorker:
1152
assert self.use_host_buffer
1153
self.copy_blocks = copy_operation
1154
1155
+ def _log_failure(
1156
+ self,
1157
+ failure_type: str,
1158
+ req_id: str | None,
1159
+ msg: str = "",
1160
+ error: Exception | None = None,
1161
+ meta: ReqMeta | None = None,
1162
+ **extra_context,
1163
+ ):
1164
+ """Log transfer failure with structured context for easier debugging."""
1165
+ context: dict[str, Any] = {
1166
+ "failure_type": failure_type,
1167
+ "request_id": req_id,
1168
+ "engine_id": self.engine_id,
1169
+ }
1170
+ if meta is None and req_id is not None:
1171
+ # Try to get metadata from in progress transfers when not provided
1172
+ meta = self._recving_metadata.get(req_id)
1173
+
1174
+ if meta and meta.remote:
1175
+ context.update(
1176
+ {
1177
+ "remote_engine_id": meta.remote.engine_id,
1178
+ "remote_request_id": meta.remote.request_id,
1179
+ "remote_host": meta.remote.host,
1180
+ "remote_port": meta.remote.port,
1181
+ "num_local_blocks": len(meta.local_block_ids),
1182
+ "num_remote_blocks": len(meta.remote.block_ids),
1183
+ "local_block_ids_sample": meta.local_block_ids[:10],
1184
+ }
1185
+ )
1186
+
1187
+ context.update(extra_context)
1188
+ if msg:
1189
+ failure_type = f"{failure_type}. {msg}"
1190
+
1191
+ logger.error(
1192
+ "NIXL transfer failure: %s | Context: %s",
1193
+ failure_type,
1194
+ context,
1195
+ exc_info=error is not None,
1196
+ stacklevel=2,
1197
+ )
1198
+
1199
def _background_nixl_handshake(
1200
self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
1201
):
@@ -1173,8 +1217,13 @@ class NixlConnectorWorker:
1217
del self._handshake_futures[eid]
1218
try:
1219
self._remote_agents[eid] = f.result()
1176
- except Exception:
1177
- logger.exception("Handshake with %s failed", eid)
1220
+ except Exception as e:
1221
+ self._log_failure(
1222
+ failure_type="handshake_setup_failed",
1223
+ req_id=None,
1224
+ error=e,
1225
+ remote_engine_id=eid,
1226
+ )
1227
1228
fut.add_done_callback(done_callback)
1229
@@ -1184,10 +1233,13 @@ class NixlConnectorWorker:
1233
# check if handshake succeeded
1234
f.result()
1235
self._ready_requests.put(entry)
1187
- except Exception:
1236
+ except Exception as e:
1237
# handshake failed - mark blocks as invalid
1189
- logger.exception(
1190
- "Handshake failed for request %s, marking blocks as invalid", req_id
1238
+ self._log_failure(
1239
+ failure_type="handshake_failed",
1240
+ req_id=req_id,
1241
+ error=e,
1242
+ meta=meta,
1243
)
1244
if req_meta := self._recving_metadata.get(req_id):
1245
self._invalid_block_ids.update(req_meta.local_block_ids)
@@ -1941,18 +1993,19 @@ class NixlConnectorWorker:
1993
in_progress.append(handle)
1994
continue
1995
else:
1944
- logger.error(
1945
- "NIXL transfer failed for request %s with state "
1946
- "%s. Marking blocks as invalid.",
1947
- req_id,
1948
- xfer_state,
1996
+ self._log_failure(
1997
+ failure_type="transfer_failed",
1998
+ msg="Marking blocks as invalid",
1999
+ req_id=req_id,
2000
+ xfer_state=xfer_state,
2001
)
2002
self._handle_failed_transfer(req_id, handle)
1951
- except Exception:
1952
- logger.exception(
1953
- "NIXL transfer exception for request %s. "
1954
- "Marking blocks as invalid.",
1955
- req_id,
2003
+ except Exception as e:
2004
+ self._log_failure(
2005
+ failure_type="transfer_exception",
2006
+ msg="Marking blocks as invalid",
2007
+ req_id=req_id,
2008
+ error=e,
2009
)
2010
self._handle_failed_transfer(req_id, handle)
2011
@@ -1973,9 +2026,9 @@ class NixlConnectorWorker:
2026
req_id: The request ID.
2027
handle: The transfer handle.
2028
"""
1976
- if meta := self._recving_metadata.pop(req_id, None):
2029
+ # Use .get() here as the metadata cleanup is handled by get_finished()
2030
+ if meta := self._recving_metadata.get(req_id):
2031
self._invalid_block_ids.update(meta.local_block_ids)
1978
- self._recving_metadata.pop(req_id, None)
2032
self.nixl_wrapper.release_xfer_handle(handle)
2033
self.xfer_stats.record_failed_transfer()
2034
@@ -2150,12 +2203,16 @@ class NixlConnectorWorker:
2203
agent_name = self._remote_agents[dst_engine_id][remote_rank]
2204
try:
2205
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
2153
- except Exception:
2154
- logger.exception(
2155
- "NIXL send_notif failed for request %s: "
2156
- "P worker blocks will be freed after timeout. "
2206
+ except Exception as e:
2207
+ self._log_failure(
2208
+ failure_type="notification_failed",
2209
+ msg="P worker blocks will be freed after timeout. "
2210
"This may indicate network issues.",
2158
- request_id,
2211
+ req_id=request_id,
2212
+ error=e,
2213
+ dst_engine_id=dst_engine_id,
2214
+ remote_rank=remote_rank,
2215
+ remote_agent_name=agent_name,
2216
)
2217
self.xfer_stats.record_failed_notification()
2218
return
@@ -2240,13 +2297,16 @@ class NixlConnectorWorker:
2297
2298
# Use handle to check completion in future step().
2299
self._recving_transfers[request_id].append(handle)
2243
- except Exception:
2244
- logger.exception(
2245
- "NIXL transfer setup/initiation failed for request %s. "
2246
- "Marking blocks as invalid.",
2247
- request_id,
2248
- )
2300
+ except Exception as e:
2301
# mark all (logical) blocks for this request as invalid
2302
+ self._log_failure(
2303
+ failure_type="transfer_setup_failed",
2304
+ req_id=request_id,
2305
+ msg="Marking blocks as invalid",
2306
+ error=e,
2307
+ dst_engine_id=dst_engine_id,
2308
+ remote_rank=remote_rank,
2309
+ )
2310
if meta := self._recving_metadata.get(request_id):
2311
self._invalid_block_ids.update(meta.local_block_ids)
2312
self.xfer_stats.record_failed_transfer()