perf(sglang): stop racing every stream chunk against the cancellation task

components/src/dynamo/sglang/request_handlers/cancellation.py

@@ -108,9 +108,70 @@ async def _stream_until_cancelled(
         stream_source: AsyncIterator[Any],
         cancellation_task: asyncio.Task[Any],
     ) -> AsyncGenerator[Any, None]:
-        """Drain briefly after abort, then stop a stream that did not terminate."""
+        """Drain briefly after abort, then stop a stream that did not terminate.
+
+        Until cancellation fires, chunks are awaited directly: racing every chunk
+        against the cancellation task costs a Task plus asyncio.wait() per token.
+        Cancellation arms a timer that interrupts the in-flight read at the drain
+        deadline (at once if the monitor failed); after that, reads race the
+        deadline as before.
+        """
         iterator = aiter(stream_source)
+        loop = asyncio.get_running_loop()
         drain_deadline: float | None = None
+        if not cancellation_task.done():
+            reader: asyncio.Task[Any] | None = None
+            interrupted = False
+            timer: asyncio.TimerHandle | None = None
+
+            def interrupt() -> None:
+                nonlocal interrupted
+                interrupted = True
+                if reader is not None:
+                    reader.cancel()
+
+            def on_cancellation(task: asyncio.Task[Any]) -> None:
+                nonlocal drain_deadline, timer
+                if task.cancelled():
+                    return
+                if task.exception() is not None:
+                    interrupt()
+                    return
+                drain_deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_S
+                timer = loop.call_later(_CANCELLATION_DRAIN_TIMEOUT_S, interrupt)
+
+            cancellation_task.add_done_callback(on_cancellation)
+            try:
+                while not cancellation_task.done():
+                    reader = asyncio.current_task()
+                    try:
+                        item = await anext(iterator)
+                    except StopAsyncIteration:
+                        return
+                    except asyncio.CancelledError:
+                        if not interrupted:
+                            raise
+                        # Withdraw our own cancel request; another one still
+                        # pending (the consumer went away too) wins.
+                        current = asyncio.current_task()
+                        uncancel = getattr(current, "uncancel", None)
+                        if uncancel is not None and uncancel() > 0:
+                            raise
+                        cancellation_task.result()  # re-raise a monitor failure
+                        logging.warning(
+                            "Timed out draining SGLang stream after cancellation"
+                        )
+                        return
+                    finally:
+                        reader = None
+                    yield item
+            finally:
+                cancellation_task.remove_done_callback(on_cancellation)
+                if timer is not None:
+                    timer.cancel()
+            cancellation_task.result()  # re-raise a monitor failure
+            if drain_deadline is None:
+                drain_deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_S
         while True:
             next_item = asyncio.create_task(_next_stream_item(iterator))
             try:
@@ -121,10 +182,7 @@ async def _stream_until_cancelled(
                     )
                     if cancellation_task in done:
                         cancellation_task.result()
-                        drain_deadline = (
-                            asyncio.get_running_loop().time()
-                            + _CANCELLATION_DRAIN_TIMEOUT_S
-                        )
+                        drain_deadline = loop.time() + _CANCELLATION_DRAIN_TIMEOUT_S
                     if next_item in done:
                         try:
                             yield next_item.result()
@@ -133,7 +191,7 @@ async def _stream_until_cancelled(
                         continue
 
                 assert drain_deadline is not None
-                remaining = drain_deadline - asyncio.get_running_loop().time()
+                remaining = drain_deadline - loop.time()
                 if remaining <= 0:
                     logging.warning(
                         "Timed out draining SGLang stream after cancellation"
@@ -434,6 +492,18 @@ async def _cancellation_monitor(
             request_id = self._resolved_request_id(request_id_future)
             ordered_abort_task = None
             try:
+                if not cancellation_task.done() and (
+                    context.is_stopped()
+                    or (
+                        self.shutdown_event is not None and self.shutdown_event.is_set()
+                    )
+                ):
+                    # The signal fired and the stream ended first: let the monitor
+                    # finish ordering its abort rather than racing a fallback abort.
+                    await asyncio.wait(
+                        (cancellation_task,),
+                        timeout=_CANCELLATION_DISPATCH_WAIT_TIMEOUT_S,
+                    )
                 if not cancellation_task.done():
                     logging.debug(
                         "Cancelling cancellation monitor task for SGLang Request ID %s, Context: %s",

components/src/dynamo/sglang/tests/test_sglang_decode_cancellation.py

@@ -43,7 +43,7 @@ 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):
@@ -644,6 +644,98 @@ async def resistant_stream():
     assert cancellation_seen.is_set()
 
 
+@pytest.mark.asyncio
+@pytest.mark.timeout(5)
+async def test_stream_reads_chunks_directly_until_cancellation(
+    decode_cancellation_case, monkeypatch
+):
+    async def no_race(*_args, **_kwargs):
+        raise AssertionError("chunks must not race cancellation before it fires")
+
+    monkeypatch.setattr(
+        "dynamo.sglang.request_handlers.cancellation.asyncio.wait", no_race
+    )
+
+    async def chunks():
+        for i in range(3):
+            await asyncio.sleep(0)
+            yield {"chunk": i}
+
+    cancellation_task = asyncio.create_task(asyncio.Event().wait())
+    try:
+        outputs = await _collect(
+            decode_cancellation_case.handler._stream_until_cancelled(
+                chunks(), cancellation_task
+            )
+        )
+    finally:
+        cancellation_task.cancel()
+        await asyncio.gather(cancellation_task, return_exceptions=True)
+
+    assert outputs == [{"chunk": 0}, {"chunk": 1}, {"chunk": 2}]
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(5)
+async def test_cancellation_during_read_is_bounded_by_drain_deadline(
+    decode_cancellation_case, monkeypatch, caplog
+):
+    caplog.set_level(logging.WARNING)
+    monkeypatch.setattr(
+        "dynamo.sglang.request_handlers.cancellation._CANCELLATION_DRAIN_TIMEOUT_S",
+        0.01,
+    )
+    read_cancelled = asyncio.Event()
+
+    async def stalls_after_first_chunk():
+        yield {"chunk": 0}
+        try:
+            await asyncio.Future()
+        except asyncio.CancelledError:
+            read_cancelled.set()
+            raise
+        yield {"chunk": 1}
+
+    cancellation_task = asyncio.create_task(asyncio.sleep(0.01))
+    consumer = asyncio.create_task(
+        _collect(
+            decode_cancellation_case.handler._stream_until_cancelled(
+                stalls_after_first_chunk(), cancellation_task
+            )
+        )
+    )
+    outputs = await asyncio.wait_for(consumer, timeout=1)
+
+    assert outputs == [{"chunk": 0}]
+    assert read_cancelled.is_set()
+    assert "Timed out draining SGLang stream after cancellation" in caplog.messages
+    # The interrupt withdraws its own cancel request.
+    assert consumer.cancelling() == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(5)
+async def test_monitor_failure_during_read_is_raised(decode_cancellation_case):
+    async def stalled_stream():
+        await asyncio.Future()
+        yield {}
+
+    async def failing_monitor():
+        await asyncio.sleep(0.01)
+        raise EngineShutdown("Engine was shut down during token generation")
+
+    cancellation_task = asyncio.create_task(failing_monitor())
+    with pytest.raises(EngineShutdown, match="shut down during token generation"):
+        await asyncio.wait_for(
+            _collect(
+                decode_cancellation_case.handler._stream_until_cancelled(
+                    stalled_stream(), cancellation_task
+                )
+            ),
+            timeout=0.5,
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.timeout(5)
 async def test_decode_cancellation_drains_buffered_empty_chunk(