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] Handle case where async utility call is cancelled (#22996)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Yinghai Lu <yinghai@thinkingmachines.ai>
Nick Hill committed
Aug 15, 2025 at 16:38 UTC
b9dc9d260762cdd98946ac6bae88e298ed28b055
2 files changed
+38
-7
tests/v1/engine/test_engine_core_client.py
+23
-1
@@ -121,8 +121,13 @@ async def loop_until_fully_done_async(client: EngineCoreClient, outputs: dict):
121
122
123
# Dummy utility function to monkey-patch into engine core.
124
-def echo(self, msg: str, err_msg: Optional[str] = None) -> str:
124
+def echo(self,
125
+ msg: str,
126
+ err_msg: Optional[str] = None,
127
+ sleep: Optional[float] = None) -> str:
128
print(f"echo util function called: {msg}, {err_msg}")
129
+ if sleep is not None:
130
+ time.sleep(sleep)
131
if err_msg is not None:
132
raise ValueError(err_msg)
133
return msg
@@ -289,6 +294,23 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
294
await core_client.call_utility_async("echo", None, "help!")
295
296
assert str(e_info.value) == "Call to echo method failed: help!"
297
+
298
+ # Test that cancelling the utility call doesn't destabilize the
299
+ # engine.
300
+ util_task = asyncio.create_task(
301
+ core_client.call_utility_async("echo", "testarg2", None,
302
+ 0.5)) # sleep for 0.5 sec
303
+ await asyncio.sleep(0.05)
304
+ cancelled = util_task.cancel()
305
+ assert cancelled
306
+
307
+ # Ensure client is still functional. The engine runs utility
308
+ # methods in a single thread so this request won't be processed
309
+ # until the cancelled sleeping one is complete.
310
+ result = await asyncio.wait_for(core_client.call_utility_async(
311
+ "echo", "testarg3"),
312
+ timeout=1.0)
313
+ assert result == "testarg3"
314
finally:
315
client.shutdown()
316
vllm/v1/engine/core_client.py
+15
-6
@@ -574,13 +574,22 @@ class MPClient(EngineCoreClient):
574
575
def _process_utility_output(output: UtilityOutput,
576
utility_results: dict[int, AnyFuture]):
577
- """Set the result from a utility method in the waiting future"""
577
+ """Set the result from a utility method in the waiting future."""
578
future = utility_results.pop(output.call_id)
579
- if output.failure_message is not None:
580
- future.set_exception(Exception(output.failure_message))
581
- else:
582
- assert output.result is not None
583
- future.set_result(output.result.result)
579
+ failure_message = output.failure_message
580
+ try:
581
+ if failure_message is not None:
582
+ future.set_exception(Exception(failure_message))
583
+ else:
584
+ assert output.result is not None
585
+ future.set_result(output.result.result)
586
+ except asyncio.InvalidStateError:
587
+ # This can happen if the future is cancelled due to the
588
+ # original calling task being cancelled.
589
+ if failure_message is not None:
590
+ logger.error(
591
+ "Cancelled call to utility method failed "
592
+ "with error: %s", failure_message)
593
594
595
class SyncMPClient(MPClient):