[Bugfix] Catch asyncio.TimeoutError so duplex serving starts on Python 3.10 (#7412)

Signed-off-by: Tianyao Wu <rayroy31@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>

tests/engine/duplex/test_session_manager.py

@@ -51,6 +51,7 @@ from vllm_omni.engine.duplex.plugin import (
     PcmAppendReservation,
 )
 from vllm_omni.engine.duplex.session import manager as session_manager_module
+from vllm_omni.engine.duplex.session.context import DuplexSessionTasks
 from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession
 from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity
 from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
@@ -1612,3 +1613,34 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo
     finally:
         shutdown.set()
         await asyncio.wait_for(task, timeout=5.0)
+
+
+# --------------------------------------------------------------------------- #
+# Cancellation waits                                                          #
+# --------------------------------------------------------------------------- #
+
+
+async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> None:
+    """The wait after cancelling is allowed to time out; that is not an error.
+
+    On Python 3.10 the timeout arrives as ``asyncio.TimeoutError``, which is a
+    different class from the builtin ``TimeoutError`` until 3.11.
+    """
+
+    async def outlives_the_first_cancel() -> bool:
+        try:
+            await asyncio.sleep(5)
+        except asyncio.CancelledError:
+            await asyncio.sleep(1)
+            raise
+        return True
+
+    tasks = DuplexSessionTasks()
+    task = asyncio.create_task(outlives_the_first_cancel())
+    await asyncio.sleep(0)
+    tasks.track_append_task(task, epoch=0, final=False, response_bound=False)
+    tasks.append_tail = task
+
+    assert await tasks.cancel_append_tasks(timeout_s=0.05) is True
+    assert tasks.append_tail is None
+    await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0)

tests/engine/duplex/test_session_runner.py

@@ -771,6 +771,41 @@ async def test_cancel_response_aborts_the_stage_request_and_reports_playback() -
         await close_harness(h)
 
 
+@pytest.mark.asyncio
+async def test_cancel_response_absorbs_a_response_task_that_outlives_the_wait() -> None:
+    """The 0.25 s wait after cancelling the response task may time out.
+
+    On Python 3.10 that timeout is ``asyncio.TimeoutError``, not the builtin
+    ``TimeoutError``, and it must not stop the cancel from completing.
+    """
+
+    async def outlives_the_first_cancel() -> None:
+        try:
+            await asyncio.sleep(5)
+        except asyncio.CancelledError:
+            await asyncio.sleep(1)
+            raise
+
+    h = await open_harness()
+    task = asyncio.create_task(outlives_the_first_cancel())
+    try:
+        await h.run(append_audio())
+        request_id = h.stage0_request_id()
+        await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello"))
+        h.runner.tasks.active_response_task = task
+
+        events = await h.run(commands.CancelResponse())
+        # The runner is still inside that wait when the mailbox goes quiet.
+        events += await h.settle(idle_s=0.6, timeout_s=5.0)
+        assert find(events, "response.done").status == "cancelled"
+        assert "error" not in types(events)
+        assert h.session.epoch == 1
+    finally:
+        task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+        await close_harness(h)
+
+
 @pytest.mark.asyncio
 async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None:
     h = await open_harness()

vllm_omni/engine/duplex/session/context.py

@@ -141,7 +141,8 @@ class DuplexSessionTasks:
             task.cancel()
         try:
             await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s)
-        except TimeoutError:
+        # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
+        except (TimeoutError, asyncio.TimeoutError):
             pass
         if cancelled_tail is not None and self.append_tail is cancelled_tail:
             self.append_tail = None

vllm_omni/engine/duplex/session/manager.py

@@ -699,7 +699,9 @@ class DuplexSessionManager:
         while not shutdown_event.is_set():
             try:
                 await asyncio.wait_for(shutdown_event.wait(), timeout=interval)
-            except TimeoutError:
+            # The timeout is this loop's tick. asyncio.TimeoutError is not the builtin
+            # TimeoutError before Python 3.11, so catch both or the tick escapes the loop.
+            except (TimeoutError, asyncio.TimeoutError):
                 try:
                     await self.reap_expired()
                 except Exception:

vllm_omni/engine/duplex/session/runner.py

@@ -1486,7 +1486,8 @@ class DuplexSessionRunner:
             active_task.cancel()
             try:
                 await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25)
-            except TimeoutError:
+            # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
+            except (TimeoutError, asyncio.TimeoutError):
                 pass
         if notify:
             self.emit(

vllm_omni/entrypoints/duplex/websocket.py

@@ -28,7 +28,8 @@ async def receive_text_with_timeout(websocket: WebSocket, timeout_s: float | Non
         return await websocket.receive_text()
     try:
         return await asyncio.wait_for(websocket.receive_text(), timeout=timeout_s)
-    except TimeoutError:
+    # asyncio.TimeoutError is not the builtin TimeoutError before Python 3.11.
+    except (TimeoutError, asyncio.TimeoutError):
         return None