[Bugfix][KVConnector] Make ExampleHiddenStatesConnector abort-safe (#56841)
Signed-off-by: ys2025-AI <yuanshuaien@163.com>
ys2025-AI committed
Sep 19, 2026 at 05:53 UTC
6aab78f4777c945b02123af330ff351f72785c0f
2 files changed
+93
-9
tests/v1/kv_connector/unit/test_hidden_states_connector.py
+59
@@ -169,3 +169,62 @@ def test_hidden_state_group_isolated_from_packed_mixed_page_groups():
169
assert (
170
ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == hidden_group_ids[0]
171
)
172
+
173
+
174
+# ---- abort-path robustness --------------------------------------------------
175
+
176
+
177
+def _bare_connector() -> ExampleHiddenStatesConnector:
178
+ """Instance with scheduler-side state but bypassing ``__init__`` (no engine)."""
179
+ conn = ExampleHiddenStatesConnector.__new__(ExampleHiddenStatesConnector)
180
+ conn._request_filenames = {}
181
+ conn._pending_saves = {}
182
+ conn._lock_fds = {}
183
+ conn._cache_kv_group_id = 1
184
+ conn._connector_metadata = None
185
+ conn._req_copy_events = {}
186
+ conn._accumulated_finished_req_ids = set()
187
+ return conn
188
+
189
+
190
+def test_get_finished_count_is_one():
191
+ # Only TP rank 0 writes, so KVOutputAggregator must expect a single
192
+ # finished_sending notification per request (not the TP world size).
193
+ assert _bare_connector().get_finished_count() == 1
194
+
195
+
196
+def test_request_finished_is_noop_for_never_scheduled_request():
197
+ # A request aborted while still queued never reaches build_connector_meta,
198
+ # so no filename was recorded. request_finished must not raise KeyError.
199
+ conn = _bare_connector()
200
+ request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None)
201
+ assert conn.request_finished(request, []) == (False, None)
202
+
203
+
204
+def test_request_finished_all_groups_handles_missing_group():
205
+ # Guard against indexing a nonexistent per-group block table.
206
+ conn = _bare_connector()
207
+ conn._cache_kv_group_id = 2
208
+ request = SimpleNamespace(request_id="cmpl-aborted", kv_transfer_params=None)
209
+ assert conn.request_finished_all_groups(request, ([], [])) == (False, None)
210
+
211
+
212
+def test_get_finished_does_not_report_untracked_request():
213
+ # A never-scheduled aborted request has no copy event. get_finished must not
214
+ # report it as done_sending, or the scheduler asserts it is still tracked.
215
+ conn = _bare_connector()
216
+ assert conn.get_finished({"cmpl-aborted"}) == (None, None)
217
+ assert conn._accumulated_finished_req_ids == set()
218
+
219
+
220
+def test_get_finished_reports_tracked_completed_request():
221
+ conn = _bare_connector()
222
+
223
+ class _DoneEvent:
224
+ def query(self) -> bool:
225
+ return True
226
+
227
+ conn._req_copy_events["cmpl-done"] = _DoneEvent()
228
+ done_sending, done_recving = conn.get_finished({"cmpl-done"})
229
+ assert done_sending == {"cmpl-done"}
230
+ assert done_recving is None
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+34
-9
@@ -103,6 +103,18 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
103
Must be used in conjunction with the `extract_hidden_states` spec decoding method.
104
"""
105
106
+ def get_finished_count(self) -> int:
107
+ """Number of ``finished_sending`` notifications expected per request.
108
+
109
+ Only TP rank 0 writes hidden states to disk, so each request yields a
110
+ single notification. The default (``None``) makes ``KVOutputAggregator``
111
+ wait for every worker in the world size; since ranks 1..N-1 never submit
112
+ a copy, the request would never be reported as finished and its blocks
113
+ would be leaked forever (the scheduler then spins on
114
+ ``has_finished_requests()``).
115
+ """
116
+ return 1
117
+
118
@classmethod
119
def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
120
"""Index of the KV cache group holding the extracted hidden states.
@@ -517,7 +529,13 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
529
the hidden states from the KV cache.
530
"""
531
req_id = request.request_id
520
- filename = self._request_filenames.pop(req_id)
532
+ filename = self._request_filenames.pop(req_id, None)
533
+ if filename is None:
534
+ # Request was aborted before it was ever scheduled, so
535
+ # build_connector_meta never recorded a filename for it. There are
536
+ # no hidden states to save and no blocks worth delaying.
537
+ self._pending_saves.pop(req_id, None)
538
+ return False, None
539
kv_params = request.kv_transfer_params or {}
540
if kv_params.get("include_output_tokens", False):
541
# Exclude the final token — it was the model's output, never an
@@ -566,15 +584,18 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
584
done_sending: set[str] = set()
585
for req_id in list(self._accumulated_finished_req_ids):
586
event = self._req_copy_events.get(req_id)
569
- if event is None or event.query():
570
- self._req_copy_events.pop(req_id, None)
587
+ if event is not None and not event.query():
588
+ continue # DtoH copy still in flight
589
+ self._req_copy_events.pop(req_id, None)
590
+ self._accumulated_finished_req_ids.discard(req_id)
591
+ # Only report requests we actually tracked for sending.
592
+ if event is not None:
593
done_sending.add(req_id)
572
- self._accumulated_finished_req_ids.discard(req_id)
573
- # Clean up any leftover lock fds (e.g. aborted requests
574
- # that never went through _submit_async_write).
575
- lock_fd = self._lock_fds.pop(req_id, None)
576
- if lock_fd is not None:
577
- os.close(lock_fd)
594
+ # Clean up any leftover lock fds (e.g. aborted requests that never
595
+ # went through _submit_async_write).
596
+ lock_fd = self._lock_fds.pop(req_id, None)
597
+ if lock_fd is not None:
598
+ os.close(lock_fd)
599
600
return done_sending or None, None
601
@@ -583,6 +604,10 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
604
request: "Request",
605
block_ids: tuple[list[int], ...],
606
) -> tuple[bool, dict[str, Any] | None]:
607
+ if not 0 <= self._cache_kv_group_id < len(block_ids):
608
+ # No per-group block table (e.g. a request that was never
609
+ # scheduled); treat as a no-op instead of indexing out of bounds.
610
+ return False, None
611
return self.request_finished(request, block_ids[self._cache_kv_group_id])
612
613
@classmethod