@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 258 lines 7.27 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import asyncio
5
6 import pytest
7 from starlette.responses import StreamingResponse
8
9 from vllm.entrypoints.serve.utils.sse_keep_alive import (
10 SSE_KEEP_ALIVE_COMMENT,
11 with_sse_keep_alive,
12 )
13
14 # --- Disabled path ----------------------------------------------------------
15
16
17 def test_disabled_returns_same_generator():
18 """For 0, negative, and NaN intervals the exact same generator object must
19 be returned (identity), proving no wrapper or per-chunk overhead when
20 disabled."""
21
22 async def gen():
23 yield "x"
24
25 for disabled in (0.0, -5.0, float("nan"), float("inf")):
26 g = gen()
27 assert with_sse_keep_alive(g, disabled) is g
28
29
30 # --- Keep-alive emission ----------------------------------------------------
31
32
33 @pytest.mark.asyncio
34 async def test_keep_alive_during_silent_start():
35 """A silent start (queue wait / prefill) emits keep-alive comments before
36 the first real chunk."""
37 release = asyncio.Event()
38
39 async def gen():
40 await release.wait() # simulates queue wait / prefill
41 yield "data: first\n\n"
42
43 out: list[str] = []
44 async for item in with_sse_keep_alive(gen(), 0.01):
45 out.append(item)
46 if item == SSE_KEEP_ALIVE_COMMENT:
47 release.set()
48 if item == "data: first\n\n":
49 break
50
51 assert SSE_KEEP_ALIVE_COMMENT in out
52 assert out.index(SSE_KEEP_ALIVE_COMMENT) < out.index("data: first\n\n")
53
54
55 @pytest.mark.asyncio
56 async def test_no_keep_alive_when_streaming_is_fast():
57 """If chunks arrive faster than the interval, no comment is emitted and all
58 data is forwarded."""
59
60 async def gen():
61 for chunk in ["a", "b", "c"]:
62 yield chunk
63
64 out = [c async for c in with_sse_keep_alive(gen(), 0.5)]
65 assert out == ["a", "b", "c"]
66 assert SSE_KEEP_ALIVE_COMMENT not in out
67
68
69 @pytest.mark.asyncio
70 async def test_keep_alive_between_chunks_does_not_drop_data():
71 """Keep-alives interleaved with real chunks must never drop or reorder
72 data."""
73 release_b = asyncio.Event()
74 release_c = asyncio.Event()
75
76 async def gen():
77 yield "data: a\n\n"
78 await release_b.wait()
79 yield "data: b\n\n"
80 await release_c.wait()
81 yield "data: c\n\n"
82
83 out: list[str] = []
84 async for item in with_sse_keep_alive(gen(), 0.01):
85 out.append(item)
86 if item == SSE_KEEP_ALIVE_COMMENT and not release_b.is_set():
87 release_b.set()
88 elif item == "data: b\n\n" and not release_c.is_set():
89 release_c.set()
90 elif item == "data: c\n\n":
91 break
92
93 assert SSE_KEEP_ALIVE_COMMENT in out
94 data_chunks = [c for c in out if c != SSE_KEEP_ALIVE_COMMENT]
95 assert data_chunks == ["data: a\n\n", "data: b\n\n", "data: c\n\n"]
96
97
98 @pytest.mark.asyncio
99 async def test_empty_generator_returns_immediately():
100 async def gen():
101 if False:
102 yield "never"
103
104 out = [c async for c in with_sse_keep_alive(gen(), 0.01)]
105 assert out == []
106
107
108 # --- Errors and cancellation ------------------------------------------------
109
110
111 @pytest.mark.asyncio
112 async def test_exception_propagates():
113 async def gen():
114 yield "data: before\n\n"
115 raise RuntimeError("boom")
116
117 wrapped = with_sse_keep_alive(gen(), 0.01)
118 assert await anext(wrapped) == "data: before\n\n"
119 with pytest.raises(RuntimeError, match="boom"):
120 await anext(wrapped)
121
122
123 @pytest.mark.asyncio
124 async def test_upstream_cancelled_error_propagates():
125 async def gen():
126 yield "data: before\n\n"
127 raise asyncio.CancelledError("upstream cancelled")
128
129 wrapped = with_sse_keep_alive(gen(), 0.01)
130 assert await anext(wrapped) == "data: before\n\n"
131 with pytest.raises(asyncio.CancelledError):
132 await anext(wrapped)
133
134
135 @pytest.mark.asyncio
136 async def test_aclose_closes_upstream():
137 """Closing a wrapper that is mid-stream must complete promptly and close
138 the upstream generator (regression for the earlier leaked-task hang)."""
139 finalized = asyncio.Event()
140
141 async def gen():
142 try:
143 yield "chunk-0"
144 await asyncio.Event().wait() # block forever
145 finally:
146 finalized.set()
147
148 wrapped = with_sse_keep_alive(gen(), 10.0)
149 assert await anext(wrapped) == "chunk-0"
150 await asyncio.wait_for(wrapped.aclose(), timeout=2.0)
151 assert finalized.is_set()
152
153
154 @pytest.mark.asyncio
155 async def test_client_disconnect_cancels_upstream():
156 """Cancelling the consumer mid-stream must cancel the pending anext and
157 close the upstream generator so its cleanup runs."""
158 finalized = asyncio.Event()
159
160 async def gen():
161 try:
162 for i in range(1000):
163 yield f"chunk-{i}"
164 await asyncio.Event().wait()
165 finally:
166 finalized.set()
167
168 wrapped = with_sse_keep_alive(gen(), 10.0)
169 assert await anext(wrapped) == "chunk-0"
170 task = asyncio.current_task()
171 assert task is not None
172 task.cancel()
173 with pytest.raises(asyncio.CancelledError):
174 await anext(wrapped)
175 assert finalized.is_set()
176
177
178 # --- Response-level integration (the shipped StreamingResponse path) --------
179
180
181 @pytest.mark.asyncio
182 async def test_streaming_response_emits_comment():
183 """A keep-alive comment must reach the client through the real
184 StreamingResponse used by the routers."""
185 release = asyncio.Event()
186
187 async def gen():
188 await release.wait()
189 yield "data: hi\n\n"
190
191 response = StreamingResponse(
192 content=with_sse_keep_alive(gen(), 0.01),
193 media_type="text/event-stream",
194 )
195 body: list[bytes] = []
196 started = asyncio.Event()
197
198 async def send(message):
199 if message["type"] == "http.response.start":
200 started.set()
201 elif message["type"] == "http.response.body":
202 body.append(message["body"])
203 if SSE_KEEP_ALIVE_COMMENT.encode() in message["body"]:
204 release.set()
205
206 async def receive():
207 await asyncio.Event().wait()
208
209 await response(
210 {"type": "http", "asgi": {"spec_version": "2.3"}},
211 receive,
212 send,
213 )
214 data = b"".join(body).decode()
215 assert SSE_KEEP_ALIVE_COMMENT in data
216 assert "data: hi" in data
217
218
219 @pytest.mark.asyncio
220 async def test_streaming_response_disconnect_closes_upstream():
221 """On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette
222 cancels streaming and the wrapper's finally must close the upstream
223 generator."""
224 finalized = asyncio.Event()
225
226 async def gen():
227 try:
228 await asyncio.Event().wait()
229 yield "never"
230 finally:
231 finalized.set()
232
233 response = StreamingResponse(
234 content=with_sse_keep_alive(gen(), 10.0),
235 media_type="text/event-stream",
236 )
237 started = asyncio.Event()
238
239 async def receive():
240 await started.wait()
241 return {"type": "http.disconnect"}
242
243 async def send(message):
244 if message["type"] == "http.response.start":
245 started.set()
246
247 await response(
248 {"type": "http", "asgi": {"spec_version": "2.3"}},
249 receive,
250 send,
251 )
252 assert finalized.is_set()
253 pending = [
254 t
255 for t in asyncio.all_tasks()
256 if t is not asyncio.current_task() and not t.done()
257 ]
258 assert pending == []