diff --git a/components/src/dynamo/sglang/request_handlers/cancellation.py b/components/src/dynamo/sglang/request_handlers/cancellation.py index 063e26d3fe1e..e709884bd687 100644 --- a/components/src/dynamo/sglang/request_handlers/cancellation.py +++ b/components/src/dynamo/sglang/request_handlers/cancellation.py @@ -20,17 +20,15 @@ _CANCELLATION_REGISTRATION_WAIT_TIMEOUT_S = 30.0 # Leave most of the drain window available after an ordered abort is submitted. _CANCELLATION_DISPATCH_WAIT_TIMEOUT_S = 0.25 - - -async def _next_stream_item(iterator: AsyncIterator[Any]) -> Any: - return await anext(iterator) +# Ends a stream without storing an exception (and its traceback cycle) in the reply. +_STREAM_END = object() def _consume_detached_task(task: asyncio.Task[Any]) -> None: """Consume a detached cancellation task result without surfacing closure.""" try: task.result() - except (asyncio.CancelledError, StopAsyncIteration): + except asyncio.CancelledError: pass except Exception: logging.exception("Detached SGLang task failed during cancellation") @@ -110,48 +108,63 @@ async def _stream_until_cancelled( ) -> AsyncGenerator[Any, None]: """Drain briefly after abort, then stop a stream that did not terminate.""" iterator = aiter(stream_source) - drain_deadline: float | None = None - while True: - next_item = asyncio.create_task(_next_stream_item(iterator)) - try: - if drain_deadline is None: - done, _ = await asyncio.wait( - (next_item, cancellation_task), - return_when=asyncio.FIRST_COMPLETED, - ) - if cancellation_task in done: - cancellation_task.result() - drain_deadline = ( - asyncio.get_running_loop().time() - + _CANCELLATION_DRAIN_TIMEOUT_S - ) - if next_item in done: - try: - yield next_item.result() - except StopAsyncIteration: - return - continue - - assert drain_deadline is not None - remaining = drain_deadline - asyncio.get_running_loop().time() - if remaining <= 0: - logging.warning( - "Timed out draining SGLang stream after cancellation" - ) + loop = asyncio.get_running_loop() + # Carries the consumer's reply future to the reader. + request: asyncio.Future[asyncio.Future[Any]] = loop.create_future() + reply: asyncio.Future[Any] | None = None + interrupted = False + timer: asyncio.TimerHandle | None = None + + async def read_on_request() -> None: + nonlocal request + while True: + pending = await request + request = loop.create_future() + try: + item = await anext(iterator, _STREAM_END) + except (Exception, asyncio.CancelledError) as exc: + if pending.done(): + raise # surfaced by _consume_detached_task + pending.set_exception(exc) return - done, _ = await asyncio.wait((next_item,), timeout=remaining) - if next_item not in done: - logging.warning( - "Timed out draining SGLang stream after cancellation" - ) + if pending.done(): # the consumer stopped waiting return - try: - yield next_item.result() - except StopAsyncIteration: + pending.set_result(item) + if item is _STREAM_END: return - finally: - if not next_item.done(): - _cancel_and_detach(next_item) + + def interrupt() -> None: + nonlocal interrupted + interrupted = True + if reply is not None and not reply.done(): + reply.set_result(None) + + def on_cancellation(task: asyncio.Task[Any]) -> None: + nonlocal timer + if task.cancelled() or task.exception() is not None: + interrupt() + else: + timer = loop.call_later(_CANCELLATION_DRAIN_TIMEOUT_S, interrupt) + + reader = asyncio.create_task(read_on_request()) + cancellation_task.add_done_callback(on_cancellation) + try: + while not interrupted: + reply = loop.create_future() + request.set_result(reply) + item = await reply + if item is _STREAM_END: + return + if interrupted: + break + yield item + cancellation_task.result() # re-raise a monitor failure + logging.warning("Timed out draining SGLang stream after cancellation") + finally: + cancellation_task.remove_done_callback(on_cancellation) + if timer is not None: + timer.cancel() + _cancel_and_detach(reader) @asynccontextmanager async def _wait_for_signal(self, context: Context) -> AsyncGenerator[bool, None]: @@ -433,7 +446,23 @@ async def _cancellation_monitor( finally: request_id = self._resolved_request_id(request_id_future) ordered_abort_task = None + owner_cancelled: asyncio.CancelledError | None = None try: + if not cancellation_task.done() and ( + context.is_stopped() + or ( + self.shutdown_event is not None and self.shutdown_event.is_set() + ) + ): + # Let a signalled monitor send its abort before any fallback abort. + try: + await asyncio.wait( + (cancellation_task,), + timeout=_CANCELLATION_DISPATCH_WAIT_TIMEOUT_S, + ) + except asyncio.CancelledError as exc: + # Finish the cleanup below, then propagate. + owner_cancelled = exc if not cancellation_task.done(): logging.debug( "Cancelling cancellation monitor task for SGLang Request ID %s, Context: %s", @@ -457,3 +486,5 @@ async def _cancellation_monitor( if self.shutdown_event and self.shutdown_event.is_set(): raise EngineShutdown("Engine was shut down during token generation") + if owner_cancelled is not None: + raise owner_cancelled diff --git a/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py b/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py index 5298d6c29b6c..536c218b3d83 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py +++ b/components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py @@ -43,13 +43,63 @@ async def set_shutdown_when_cancelled(*_args): handler._handle_cancellation = set_shutdown_when_cancelled request_id_future = asyncio.get_running_loop().create_future() request_id_future.set_result("sglang-request-id") - context = SimpleNamespace(id=lambda: "request-id") + context = SimpleNamespace(id=lambda: "request-id", is_stopped=lambda: False) with pytest.raises(EngineShutdown, match="shut down during token generation"): async with handler._cancellation_monitor(request_id_future, context): await asyncio.sleep(0) +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_cancellation_monitor_cleanup_finishes_when_owner_cancelled( + monkeypatch, +): + monkeypatch.setattr( + "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DISPATCH_WAIT_TIMEOUT_S", + 5, + ) + handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler) + handler.shutdown_event = None + handler._abort_tasks = set() + abort_calls = [] + handler.engine = SimpleNamespace( + tokenizer_manager=SimpleNamespace( + abort_request=lambda rid, abort_all: abort_calls.append(rid) + ) + ) + monitor_cancelled = asyncio.Event() + + async def monitor_never_finishes(*_args): + try: + await asyncio.Future() + except asyncio.CancelledError: + monitor_cancelled.set() + raise + + handler._handle_cancellation = monitor_never_finishes + request_id_future = asyncio.get_running_loop().create_future() + request_id_future.set_result("sglang-request-id") + context = SimpleNamespace(id=lambda: "request-id", is_stopped=lambda: True) + in_cleanup = asyncio.Event() + + async def owner(): + async with handler._cancellation_monitor( + request_id_future, context, request_ids={"sglang-request-id"} + ): + in_cleanup.set() + + owner_task = asyncio.create_task(owner()) + await in_cleanup.wait() + await asyncio.sleep(0) # the owner is now waiting on the monitor + owner_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await owner_task + assert monitor_cancelled.is_set() + assert abort_calls == ["sglang-request-id"] + + @pytest.mark.asyncio @pytest.mark.timeout(5) async def test_cancellation_monitor_keeps_pending_ordered_abort_alive( @@ -607,8 +657,9 @@ async def buffered_stream(): @pytest.mark.asyncio @pytest.mark.timeout(5) +@pytest.mark.parametrize("cancel_during_read", [False, True]) async def test_cancelled_stream_drain_does_not_wait_for_resistant_iterator( - decode_cancellation_case, monkeypatch, caplog + decode_cancellation_case, monkeypatch, caplog, cancel_during_read ): caplog.set_level(logging.WARNING) monkeypatch.setattr( @@ -627,8 +678,11 @@ async def resistant_stream(): iterator_finished.set() yield {} - cancellation_task = asyncio.create_task(asyncio.sleep(0)) - await cancellation_task + if cancel_during_read: + cancellation_task = asyncio.create_task(asyncio.sleep(0.005)) + else: + cancellation_task = asyncio.create_task(asyncio.sleep(0)) + await cancellation_task started = asyncio.get_running_loop().time() outputs = await _collect( decode_cancellation_case.handler._stream_until_cancelled(