@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 149 lines 4.73 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import AsyncIterator, Iterator
5 from pathlib import Path
6 from unittest.mock import patch
7
8 import pytest
9
10 import vllm.envs as envs
11 from vllm.connections import HTTPConnection
12 from vllm.exceptions import VLLMValidationError
13 from vllm.multimodal.media import AudioMediaIO, MediaConnector
14 from vllm.utils.mem_constants import KiB_bytes, MiB_bytes
15
16 pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
17
18
19 class _SyncResponse:
20 def __init__(self, chunks: list[bytes], content_length: int | None = None):
21 self.headers = (
22 {} if content_length is None else {"Content-Length": str(content_length)}
23 )
24 self._chunks = chunks
25 self.iterated = 0
26
27 def __enter__(self):
28 return self
29
30 def __exit__(self, *_args):
31 return None
32
33 def raise_for_status(self) -> None:
34 return None
35
36 def iter_content(self, chunk_size: int) -> Iterator[bytes]:
37 assert chunk_size == 64 * KiB_bytes
38 for chunk in self._chunks:
39 self.iterated += 1
40 yield chunk
41
42
43 class _AsyncContent:
44 def __init__(self, chunks: list[bytes]):
45 self._chunks = chunks
46 self.iterated = 0
47
48 async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]:
49 assert chunk_size == 64 * KiB_bytes
50 for chunk in self._chunks:
51 self.iterated += 1
52 yield chunk
53
54
55 class _AsyncResponse:
56 def __init__(self, chunks: list[bytes], content_length: int | None = None):
57 self.content_length = content_length
58 self.content = _AsyncContent(chunks)
59
60 async def __aenter__(self):
61 return self
62
63 async def __aexit__(self, *_args):
64 return None
65
66 def raise_for_status(self) -> None:
67 return None
68
69
70 def test_audio_base64_rejects_before_decode(monkeypatch: pytest.MonkeyPatch):
71 monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1)
72 max_encoded_chars = 4 * ((MiB_bytes + 2) // 3)
73
74 with (
75 patch("vllm.multimodal.media.audio.pybase64.b64decode") as decode,
76 pytest.raises(VLLMValidationError, match="Maximum file size exceeded"),
77 ):
78 AudioMediaIO().load_base64("audio/wav", "A" * (max_encoded_chars + 1))
79
80 decode.assert_not_called()
81
82
83 def test_audio_load_bytes_rejects_oversized(monkeypatch: pytest.MonkeyPatch):
84 monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1)
85
86 with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"):
87 AudioMediaIO().load_bytes(b"\x00" * (MiB_bytes + 1))
88
89
90 def test_audio_load_file_rejects_oversized(
91 monkeypatch: pytest.MonkeyPatch, tmp_path: Path
92 ):
93 monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1)
94
95 oversized = tmp_path / "big.wav"
96 oversized.write_bytes(b"\x00" * (MiB_bytes + 1))
97
98 with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"):
99 AudioMediaIO().load_file(oversized)
100
101
102 def test_sync_http_reader_rejects_from_content_length_before_body_read():
103 response = _SyncResponse([b"A" * 10], content_length=MiB_bytes + 1)
104 connection = HTTPConnection()
105
106 with (
107 patch.object(connection, "get_response", return_value=response) as get_response,
108 pytest.raises(VLLMValidationError, match="Maximum file size exceeded"),
109 ):
110 connection.get_bytes("https://example.test/audio", max_bytes=MiB_bytes)
111
112 assert get_response.call_args.kwargs["stream"] is True
113 assert response.iterated == 0
114
115
116 @pytest.mark.asyncio
117 async def test_async_http_reader_stops_after_first_over_limit_chunk():
118 response = _AsyncResponse([b"A" * (64 * KiB_bytes), b"B", b"C" * 10])
119 connection = HTTPConnection()
120
121 async def get_async_response(*_args, **_kwargs):
122 return response
123
124 with (
125 patch.object(connection, "get_async_response", side_effect=get_async_response),
126 pytest.raises(VLLMValidationError, match="Maximum file size exceeded"),
127 ):
128 await connection.async_get_bytes(
129 "https://example.test/audio",
130 max_bytes=64 * KiB_bytes,
131 )
132
133 assert response.content.iterated == 2
134
135
136 @pytest.mark.asyncio
137 async def test_audio_connector_threads_byte_limit_to_http_reader(
138 monkeypatch: pytest.MonkeyPatch,
139 ):
140 monkeypatch.setattr(envs, "VLLM_MAX_AUDIO_CLIP_FILESIZE_MB", 1)
141
142 class _Connection:
143 async def async_get_bytes(self, _url: str, **kwargs):
144 assert kwargs["max_bytes"] == MiB_bytes
145 raise VLLMValidationError("Maximum file size exceeded")
146
147 connector = MediaConnector(connection=_Connection())
148 with pytest.raises(VLLMValidationError, match="Maximum file size exceeded"):
149 await connector.fetch_audio_async("https://example.test/audio")