@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 192 lines 5.85 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import threading
5 from collections.abc import Generator
6 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
7 from unittest.mock import AsyncMock, patch
8
9 import pytest
10
11 import vllm.envs as envs
12 from vllm.connections import HTTPConnection, HTTPResponseSizeExceededError
13
14 _ONE_MIB = 1024 * 1024
15 _SMALL_BODY = b"a" * (_ONE_MIB // 2)
16 _LARGE_BODY = b"b" * (2 * _ONE_MIB)
17 _CHUNK_SIZE = 64 * 1024
18
19
20 class _MediaHandler(BaseHTTPRequestHandler):
21 def do_GET(self) -> None:
22 body = _LARGE_BODY if "large" in self.path else _SMALL_BODY
23 self.send_response(200)
24 self.send_header("Content-Type", "application/octet-stream")
25 if self.path != "/chunked-large":
26 self.send_header("Content-Length", str(len(body)))
27 self.end_headers()
28
29 try:
30 for offset in range(0, len(body), _CHUNK_SIZE):
31 self.wfile.write(body[offset : offset + _CHUNK_SIZE])
32 except (BrokenPipeError, ConnectionResetError):
33 self.close_connection = True
34
35 def log_message(self, fmt: str, *args: object) -> None:
36 return
37
38
39 @pytest.fixture
40 def local_media_server() -> Generator[str, None, None]:
41 server = ThreadingHTTPServer(("127.0.0.1", 0), _MediaHandler)
42 thread = threading.Thread(target=server.serve_forever, daemon=True)
43 thread.start()
44 try:
45 yield f"http://127.0.0.1:{server.server_port}"
46 finally:
47 server.shutdown()
48 thread.join()
49 server.server_close()
50
51
52 class _SyncResponse:
53 def __init__(self, chunks: list[bytes]) -> None:
54 self.headers: dict[str, str] = {}
55 self._chunks = chunks
56 self.content_accessed = False
57 self.iterated = 0
58
59 @property
60 def content(self) -> bytes:
61 self.content_accessed = True
62 raise AssertionError("bounded reads must not access response.content")
63
64 def iter_content(self, chunk_size: int):
65 del chunk_size
66 for chunk in self._chunks:
67 self.iterated += 1
68 yield chunk
69
70 def raise_for_status(self) -> None:
71 return
72
73 def __enter__(self):
74 return self
75
76 def __exit__(self, exc_type, exc_value, traceback):
77 return False
78
79
80 class _AsyncContent:
81 def __init__(self, chunks: list[bytes]) -> None:
82 self._chunks = chunks
83 self.iterated = 0
84
85 async def iter_chunked(self, chunk_size: int):
86 del chunk_size
87 for chunk in self._chunks:
88 self.iterated += 1
89 yield chunk
90
91
92 class _AsyncResponse:
93 def __init__(self, chunks: list[bytes]) -> None:
94 self.content = _AsyncContent(chunks)
95 self.content_length = None
96 self.read = AsyncMock(
97 side_effect=AssertionError("bounded reads must not call response.read()")
98 )
99
100 def raise_for_status(self) -> None:
101 return
102
103 async def __aenter__(self):
104 return self
105
106 async def __aexit__(self, exc_type, exc_value, traceback):
107 return False
108
109
110 def test_get_bytes_with_limit_streams_instead_of_materializing() -> None:
111 response = _SyncResponse([b"abcd", b"efgh", b"ijkl"])
112 connection = HTTPConnection()
113
114 with (
115 patch.object(
116 connection,
117 "get_response",
118 return_value=response,
119 ) as get_response,
120 pytest.raises(HTTPResponseSizeExceededError, match="maximum size"),
121 ):
122 connection.get_bytes("http://example.com/media", max_bytes=4)
123
124 get_response.assert_called_once_with(
125 "http://example.com/media",
126 stream=True,
127 timeout=None,
128 allow_redirects=True,
129 )
130 assert response.content_accessed is False
131 assert response.iterated == 2
132
133
134 @pytest.mark.asyncio
135 async def test_async_get_bytes_with_limit_streams_instead_of_materializing() -> None:
136 response = _AsyncResponse([b"abcd", b"efgh", b"ijkl"])
137 connection = HTTPConnection()
138
139 with (
140 patch.object(
141 connection,
142 "get_async_response",
143 new=AsyncMock(return_value=response),
144 ),
145 pytest.raises(HTTPResponseSizeExceededError, match="maximum size"),
146 ):
147 await connection.async_get_bytes("http://example.com/media", max_bytes=4)
148
149 response.read.assert_not_awaited()
150 assert response.content.iterated == 2
151
152
153 def test_get_bytes_rejects_content_length_over_limit(
154 monkeypatch: pytest.MonkeyPatch, local_media_server: str
155 ) -> None:
156 monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB", 1, raising=False)
157 connection = HTTPConnection(reuse_client=False)
158
159 with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB"):
160 connection.get_bytes(f"{local_media_server}/large")
161
162
163 @pytest.mark.asyncio
164 async def test_async_get_bytes_rejects_stream_over_limit_without_content_length(
165 monkeypatch: pytest.MonkeyPatch, local_media_server: str
166 ) -> None:
167 monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB", 1, raising=False)
168 connection = HTTPConnection(reuse_client=False)
169
170 try:
171 with pytest.raises(ValueError, match="VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB"):
172 await connection.async_get_bytes(f"{local_media_server}/chunked-large")
173 finally:
174 if connection._async_client is not None:
175 await connection._async_client.close()
176
177
178 @pytest.mark.asyncio
179 async def test_get_bytes_allows_body_within_limit(
180 monkeypatch: pytest.MonkeyPatch, local_media_server: str
181 ) -> None:
182 monkeypatch.setattr(envs, "VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB", 1, raising=False)
183 connection = HTTPConnection(reuse_client=False)
184
185 assert connection.get_bytes(f"{local_media_server}/small") == _SMALL_BODY
186
187 try:
188 actual = await connection.async_get_bytes(f"{local_media_server}/small")
189 assert actual == _SMALL_BODY
190 finally:
191 if connection._async_client is not None:
192 await connection._async_client.close()