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

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
2 changed files with 9 additions and 2 deletions

View File

@@ -131,7 +131,10 @@ class RequestTracker:
"""Process a request output from the engine."""
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 verbose:
logger.info("Finished request %s.", request_id)