[P/D] [NixlConnector] kv load recovery integration

[wseaton]

Purpose

This integrates nixl_connector with additional scheduler features exposed in #19330 for retrying requests that have failed blocks.

This PR also includes a small bugfix where if P crashes during zmq handshake, the D node's request status would get stuck in WAITING_FOR_REMOTE_KV forever.

Test Plan

For integration testing, tested injecting faults using a vllm process instrumented with https://github.com/wseaton/ucx-fault-injector/, which forces nixl exceptions to be thrown during transfer.

Logs

1. KV Load Failure Recovery after failed metadata handshake (decode does local prefill):

(EngineCore_DP0 pid=3317003) ERROR 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:820]   File "zmq/backend/cython/_zmq.py", line 186, in zmq.backend.cython._zmq._check_rc
(EngineCore_DP0 pid=3317003) ERROR 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:820]     raise Again(errno)
(EngineCore_DP0 pid=3317003) ERROR 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:820]     ^^^^^^^^^^^
(EngineCore_DP0 pid=3317003) ERROR 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:820] zmq.error.Again: Resource temporarily unavailable
(EngineCore_DP0 pid=3317003) ERROR 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:832] Handshake failed for request cmpl-6e8a849f-7ec0-4abe-b4c7-38bce04d53d1-0, marking blocks as invalid
(EngineCore_DP0 pid=3317003) WARNING 10-07 17:11:06 [v1/core/sched/scheduler.py:1511] Recovered from KV load failure: 1 request(s) rescheduled (64 tokens affected).
(EngineCore_DP0 pid=3317003) DEBUG 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:349] NIXLConnector get_num_new_matched_tokens: num_computed_tokens=0, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': 'c1ce8a2a-f52d-42fe-8b54-4faa2d6ade51', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3317003) DEBUG 10-07 17:11:06 [distributed/.../v1/nixl_connector.py:369] NIXLConnector update_state_after_alloc: num_external_tokens=0, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': 'c1ce8a2a-f52d-42fe-8b54-4faa2d6ade51', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3317003) DEBUG 10-07 17:11:09 [distributed/.../v1/nixl_connector.py:477] NIXLConnector request_finished, request_status=FINISHED_LENGTH_CAPPED, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': 'c1ce8a2a-f52d-42fe-8b54-4faa2d6ade51', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3317003) DEBUG 10-07 17:11:09 [v1/engine/core.py:810] EngineCore waiting for work.
(APIServer pid=3316606) INFO 10-07 17:11:09 [v1/metrics/loggers.py:128] Engine 000: Avg prompt throughput: 3.6 tokens/s, Avg generation throughput: 7.2 tokens/s, Running: 1 reqs, Waiting: 0 reqs, GPU KV cache usage: 0.0%, Prefix cache hit rate: 0.0%
(APIServer pid=3316606) INFO:     127.0.0.1:40116 - "POST /v1/completions HTTP/1.1" 200 OK

2. Load failure recovery during block pull (decode does local prefill)

[1759857217.590230](EngineCore_DP0 pid=3320874) ERROR 10-07 17:13:37 [distributed/.../v1/nixl_connector.py:1536] NIXL transfer setup/initiation failed for request cmpl-8f11141d-6a13-42f4-aaf7-548dd83a8952-0: NIXL_ERR_REMOTE_DISCONNECT. Marking blocks as invalid.
(EngineCore_DP0 pid=3320874) WARNING 10-07 17:13:37 [v1/core/sched/scheduler.py:1511] Recovered from KV load failure: 1 request(s) rescheduled (64 tokens affected).
(EngineCore_DP0 pid=3320874) DEBUG 10-07 17:13:37 [distributed/.../v1/nixl_connector.py:349] NIXLConnector get_num_new_matched_tokens: num_computed_tokens=0, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': '37ae0670-4572-46d6-8f06-4de596379492', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3320874) DEBUG 10-07 17:13:37 [distributed/.../v1/nixl_connector.py:369] NIXLConnector update_state_after_alloc: num_external_tokens=0, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': '37ae0670-4572-46d6-8f06-4de596379492', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3320874) DEBUG 10-07 17:13:38 [distributed/.../v1/nixl_connector.py:477] NIXLConnector request_finished, request_status=FINISHED_LENGTH_CAPPED, kv_transfer_params={'do_remote_prefill': False, 'do_remote_decode': False, 'remote_block_ids': [1, 2, 3, 4], 'remote_engine_id': '37ae0670-4572-46d6-8f06-4de596379492', 'remote_host': 'localhost', 'remote_port': 45307, 'tp_size': 1}
(EngineCore_DP0 pid=3320874) DEBUG 10-07 17:13:38 [v1/engine/core.py:810] EngineCore waiting for work.
(APIServer pid=3320562) INFO:     127.0.0.1:44080 - "POST /v1/completions HTTP/1.1" 200 OK

Future Work

Make this behavior opt-out via a global configuration option, and then enable aborting in the API server for the fail path, since this results in locall prefills on the decode node as the failure recovery mechanism.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @wseaton.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces failure recovery mechanisms for the NIXL KV connector. It adds error handling for transfer initiation failures and failures during block reads. When a failure occurs, the affected KV cache blocks are marked as invalid, and this information is propagated to the scheduler for retrying the request. The changes also include adding statistics for failed transfers and notifications, and rate-limiting for some log messages to prevent spam.

The overall approach is sound and significantly improves the robustness of the NIXL connector. However, I've found a critical issue where failed blocks are not reported correctly when use_host_buffer is disabled, which would prevent failure recovery in that configuration. I've left a comment with details on the issue and a suggested fix.

[wseaton]

@njhill this needs a manual merge, had to rebase because of formatting changes 😬