[Fix] Serialize FanOutCommunicator queueing calls with a FIFO-fair asyncio.Lock (#30606)
This commit is contained in:
@@ -2,8 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
from collections import deque
|
||||
from typing import Callable, Deque, Generic, List, Optional, TypeVar
|
||||
from typing import Callable, Generic, List, Optional, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
@@ -31,31 +30,26 @@ class FanOutCommunicator(Generic[T]):
|
||||
self._mode = mode
|
||||
self._result_event: Optional[asyncio.Event] = None
|
||||
self._result_values: Optional[List[T]] = None
|
||||
self._ready_queue: Deque[asyncio.Event] = deque()
|
||||
self._queueing_lock = asyncio.Lock()
|
||||
|
||||
assert mode in ["queueing", "watching"]
|
||||
|
||||
async def queueing_call(self, obj: T):
|
||||
ready_event = asyncio.Event()
|
||||
if self._result_event is not None or len(self._ready_queue) > 0:
|
||||
self._ready_queue.append(ready_event)
|
||||
await ready_event.wait()
|
||||
assert self._result_event is None
|
||||
assert self._result_values is None
|
||||
# asyncio.Lock is FIFO-fair: a new caller cannot acquire while earlier
|
||||
# callers are still waiting, so requests are strictly serialized in
|
||||
# arrival order. It also releases on exception/cancellation, so a
|
||||
# failed caller never blocks the callers queued behind it.
|
||||
async with self._queueing_lock:
|
||||
if obj is not None:
|
||||
self._send(obj)
|
||||
|
||||
if obj is not None:
|
||||
self._send(obj)
|
||||
self._result_event = asyncio.Event()
|
||||
self._result_values = []
|
||||
await self._result_event.wait()
|
||||
result_values = self._result_values
|
||||
self._result_event = self._result_values = None
|
||||
|
||||
self._result_event = asyncio.Event()
|
||||
self._result_values = []
|
||||
await self._result_event.wait()
|
||||
result_values = self._result_values
|
||||
self._result_event = self._result_values = None
|
||||
|
||||
if len(self._ready_queue) > 0:
|
||||
self._ready_queue.popleft().set()
|
||||
|
||||
return result_values
|
||||
return result_values
|
||||
|
||||
async def watching_call(self, obj):
|
||||
if self._result_event is None:
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Unit tests for FanOutCommunicator -- no server, no model loading."""
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
|
||||
from sglang.srt.managers.communicator import FanOutCommunicator
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestQueueingCall(CustomTestCase):
|
||||
def test_concurrent_caller_cannot_bypass_queue(self):
|
||||
"""A new caller arriving in the wakeup window must not overtake a
|
||||
queued caller (this interleaving used to raise AssertionError and
|
||||
return 500 on concurrent /server_info requests)."""
|
||||
|
||||
async def scenario():
|
||||
sent = []
|
||||
comm = FanOutCommunicator(send=sent.append, fan_out=1, mode="queueing")
|
||||
|
||||
# A in-flight, B queued behind it.
|
||||
task_a = asyncio.create_task(comm("A"))
|
||||
await asyncio.sleep(0)
|
||||
task_b = asyncio.create_task(comm("B"))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# Complete A, then create C before A's wakeup runs, so C's first
|
||||
# step lands between A's cleanup and B's wakeup.
|
||||
comm.handle_recv("resp-A")
|
||||
task_c = asyncio.create_task(comm("C"))
|
||||
|
||||
# Drive to completion: feed a response whenever one is in flight.
|
||||
tasks = [task_a, task_b, task_c]
|
||||
for _ in range(100):
|
||||
if all(t.done() for t in tasks):
|
||||
break
|
||||
if comm._result_event is not None and not comm._result_event.is_set():
|
||||
comm.handle_recv(f"resp-{len(sent)}")
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# All callers complete without error, in strict FIFO order.
|
||||
await asyncio.gather(*tasks)
|
||||
self.assertEqual(sent, ["A", "B", "C"])
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user