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
[Bugfix] Fix `_reqs_to_process` leak on abort (#26012)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Oct 4, 2025 at 13:39 UTC
2a6dc67eb520ddb9c4138d8b35ed6fe6226997fb
2 files changed
+82
-2
tests/v1/kv_connector/unit/test_nixl_connector.py
+66
@@ -33,6 +33,7 @@ from vllm.platforms.interface import Platform
33
from vllm.sampling_params import SamplingParams
34
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
35
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
36
+from vllm.v1.request import RequestStatus
37
38
from .utils import create_request, create_scheduler, create_vllm_config
39
@@ -1023,3 +1024,68 @@ def test_shutdown_cleans_up_resources(dist_init):
1024
assert mock_dereg.call_count == 2
1025
mock_dereg.assert_any_call("desc1")
1026
mock_dereg.assert_any_call("desc2")
1027
+
1028
+
1029
+@patch(
1030
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
1031
+ FakeNixlWrapper)
1032
+def test_aborted_request_removed_from_worker_in_batch(dist_init):
1033
+ """
1034
+ Create and schedule a request so that P adds it to in-batch tracking via
1035
+ the real scheduler, then simulate an abort (request not in next scheduler
1036
+ iteration) and verify the worker no longer tracks it as in-batch.
1037
+ """
1038
+ vllm_config = create_vllm_config()
1039
+
1040
+ scheduler = create_scheduler(vllm_config)
1041
+ # KVConnector Worker in P
1042
+ connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1043
+ connector.connector_worker = FakeNixlConnectorWorker(vllm_config,
1044
+ connector.engine_id,
1045
+ hand_shake_latency=0)
1046
+
1047
+ # Create a request that triggers do_remote_decode so that
1048
+ # the scheduler adds it to reqs_in_batch
1049
+ req = create_request(request_id=1, do_remote_decode=True, max_tokens=1)
1050
+ scheduler.add_request(req)
1051
+
1052
+ # First scheduling pass - examinate build_connector_meta output
1053
+ sched_out = scheduler.schedule()
1054
+ kv_meta = sched_out.kv_connector_metadata
1055
+ assert kv_meta is not None
1056
+ assert isinstance(kv_meta, NixlConnectorMetadata)
1057
+ assert req.request_id in kv_meta.reqs_in_batch
1058
+
1059
+ #### Model Runner start ####
1060
+ # Bind scheduler-produced metadata and start worker processing.
1061
+ connector.bind_connector_metadata(kv_meta)
1062
+
1063
+ dummy_ctx = ForwardContext(
1064
+ no_compile_layers={},
1065
+ attn_metadata={},
1066
+ virtual_engine=0,
1067
+ )
1068
+ connector.start_load_kv(dummy_ctx)
1069
+
1070
+ # Ensure it was tracked by the worker
1071
+ assert req.request_id in connector.connector_worker._reqs_to_process
1072
+
1073
+ #### Model Runner end ####
1074
+
1075
+ # Abort request - request_finished call in connector scheduler
1076
+ scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
1077
+ # Second scheduling pass - build metadata with aborted request
1078
+ sched_out2 = scheduler.schedule()
1079
+ kv_meta2 = sched_out2.kv_connector_metadata
1080
+ assert kv_meta2 is not None
1081
+ assert isinstance(kv_meta2, NixlConnectorMetadata)
1082
+ assert req.request_id not in kv_meta2.reqs_in_batch
1083
+
1084
+ # Bind empty/abort metadata and run worker step
1085
+ #### Model Runner start ####
1086
+ connector.bind_connector_metadata(kv_meta2)
1087
+ connector.start_load_kv(dummy_ctx)
1088
+
1089
+ # After abort, the worker should not keep tracking it as "in-batch"
1090
+ assert req.request_id not in connector.connector_worker._reqs_to_process
1091
+ #### Model Runner end ####
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+16
-2
@@ -113,6 +113,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
113
self.reqs_to_save: dict[ReqId, ReqMeta] = {}
114
self.reqs_to_send: dict[ReqId, float] = {}
115
self.reqs_in_batch: set[ReqId] = set()
116
+ self.reqs_not_processed: set[ReqId] = set()
117
118
def add_new_req(
119
self,
@@ -287,6 +288,9 @@ class NixlConnectorScheduler:
288
# Reqs to send and their expiration time
289
self._reqs_need_send: dict[ReqId, float] = {}
290
self._reqs_in_batch: set[ReqId] = set()
291
+ # Reqs to remove from processed set because they're not to send after
292
+ # remote prefill or aborted.
293
+ self._reqs_not_processed: set[ReqId] = set()
294
295
def get_num_new_matched_tokens(
296
self, request: "Request",
@@ -401,11 +405,13 @@ class NixlConnectorScheduler:
405
406
meta.reqs_to_send = self._reqs_need_send
407
meta.reqs_in_batch = self._reqs_in_batch
408
+ meta.reqs_not_processed = self._reqs_not_processed
409
410
# Clear the list once workers start the transfers
411
self._reqs_need_recv.clear()
412
self._reqs_need_save.clear()
413
self._reqs_in_batch = set()
414
+ self._reqs_not_processed = set()
415
self._reqs_need_send = {}
416
417
return meta
@@ -439,8 +445,12 @@ class NixlConnectorScheduler:
445
params["do_remote_prefill"] = False
446
return False, None
447
442
- if (not params.get("do_remote_decode")
443
- or request.status != RequestStatus.FINISHED_LENGTH_CAPPED):
448
+ if not params.get("do_remote_decode"):
449
+ return False, None
450
+ if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
451
+ # Also include the case of a P/D Prefill request with immediate
452
+ # block free (eg abort). Stop tracking this request.
453
+ self._reqs_not_processed.add(request.request_id)
454
return False, None
455
456
# TODO: check whether block_ids actually ever be 0. If not we could
@@ -1234,6 +1244,10 @@ class NixlConnectorWorker:
1244
for req_id in metadata.reqs_in_batch:
1245
self._reqs_to_process.add(req_id)
1246
1247
+ # Remove all requests that are not to be processed (eg aborted).
1248
+ for req_id in metadata.reqs_not_processed:
1249
+ self._reqs_to_process.discard(req_id)
1250
+
1251
# Add to requests that are waiting to be read and track expiration.
1252
for req_id, expiration_time in metadata.reqs_to_send.items():
1253
if req_id in self._reqs_to_process: