squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Pipeline Parallel: Guard for KeyErrors at request abort (#6587)

Browse Source 
Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>
This commit is contained in:
Travis Johnson 2024-07-19 20:18:19 -06:00 committed by GitHub
parent 7bd82002ae
commit 3f8d42c81f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/engine/async_llm_engine.py
View File

@ -131,7 +131,10 @@ class RequestTracker:
"""Process a request output from the engine.""" """Process a request output from the engine."""
request_id = request_output.request_id request_id = request_output.request_id
self._request_streams[request_id].put(request_output) # Guard against a KeyError which can occur if the request was aborted
# while the output was generated
if (stream := self._request_streams.get(request_id)) is not None:
stream.put(request_output)
if request_output.finished: if request_output.finished:
if verbose: if verbose:
logger.info("Finished request %s.", request_id) logger.info("Finished request %s.", request_id)

6
vllm/engine/output_processor/single_step.py
View File

@ -90,7 +90,11 @@ class SingleStepOutputProcessor(SequenceGroupOutputProcessor):
for parent_seq in parent_seqs for parent_seq in parent_seqs
} }
for sample in samples: for sample in samples:
parent_child_dict[sample.parent_seq_id].append(sample) # Guard against a KeyError which can occur if the request was
# aborted while the output was generated
if (child_list :=
parent_child_dict.get(sample.parent_seq_id)) is not None:
child_list.append(sample)
# List of (child, parent) # List of (child, parent)
child_seqs: List[Tuple[Sequence, Sequence]] = [] child_seqs: List[Tuple[Sequence, Sequence]] = []