main
py 529 lines 15.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import asyncio
5 import functools
6 import time
7 from collections.abc import Callable, Coroutine, Mapping, MutableMapping
8 from pathlib import Path
9 from typing import Any, ParamSpec, TypeVar
10
11 import aiohttp
12 import requests
13 from urllib3.util import parse_url
14
15 import vllm.envs as envs
16 from vllm.exceptions import VLLMValidationError
17 from vllm.logger import init_logger
18 from vllm.utils.mem_constants import KiB_bytes, MiB_bytes
19 from vllm.version import __version__ as VLLM_VERSION
20
21 logger = init_logger(__name__)
22
23 _P = ParamSpec("_P")
24 _T = TypeVar("_T")
25
26 # Multiplier applied to timeout and sleep on each retry attempt.
27 # Attempt N uses: base_timeout * (_RETRY_BACKOFF_FACTOR ** N) for the
28 # per-attempt timeout and sleeps _RETRY_BACKOFF_FACTOR ** N seconds.
29 _RETRY_BACKOFF_FACTOR = 4
30 _RESPONSE_READ_CHUNK_SIZE = 64 * KiB_bytes
31
32
33 class HTTPResponseSizeExceededError(VLLMValidationError):
34 """Raised when an HTTP response exceeds a caller-supplied byte ceiling."""
35
36 def __init__(self, max_bytes: int, received_bytes: int) -> None:
37 self.max_bytes = max_bytes
38 self.received_bytes = received_bytes
39 super().__init__(
40 "Maximum file size exceeded: HTTP response exceeded maximum size "
41 f"of {max_bytes} bytes (received at least {received_bytes} bytes)",
42 parameter="audio_filesize_mb",
43 value=received_bytes / MiB_bytes,
44 )
45
46
47 class MediaDownloadSizeExceededError(ValueError):
48 """Raised when a remote media download exceeds the configured byte limit."""
49
50 def __init__(self, max_bytes: int, received_bytes: int | None = None) -> None:
51 self.max_bytes = max_bytes
52 self.received_bytes = received_bytes
53 max_size_mb = max_bytes / MiB_bytes
54 message = (
55 "Remote media download exceeds "
56 f"VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB={max_size_mb:g}"
57 )
58 if received_bytes is not None:
59 message += f" (received at least {received_bytes} bytes)"
60 super().__init__(message)
61
62
63 def _get_media_download_limit_bytes() -> int | None:
64 max_size_mb = envs.VLLM_MAX_MEDIA_DOWNLOAD_SIZE_MB
65 if max_size_mb <= 0:
66 return None
67 return max_size_mb * MiB_bytes
68
69
70 def _get_response_size_limit(max_bytes: int | None) -> tuple[int | None, bool]:
71 media_max_bytes = _get_media_download_limit_bytes()
72 if media_max_bytes is not None and (
73 max_bytes is None or media_max_bytes <= max_bytes
74 ):
75 return media_max_bytes, True
76 return max_bytes, False
77
78
79 def _raise_size_exceeded(
80 max_bytes: int, received_bytes: int, media_download_limit: bool
81 ) -> None:
82 if media_download_limit:
83 raise MediaDownloadSizeExceededError(max_bytes, received_bytes)
84 raise HTTPResponseSizeExceededError(max_bytes, received_bytes)
85
86
87 def _check_content_length(
88 content_length: str | int | None,
89 max_bytes: int | None,
90 media_download_limit: bool,
91 ) -> None:
92 if content_length is None or max_bytes is None:
93 return
94 try:
95 parsed_length = int(content_length)
96 except (TypeError, ValueError):
97 return
98 if parsed_length > max_bytes:
99 _raise_size_exceeded(max_bytes, parsed_length, media_download_limit)
100
101
102 def _append_chunk(
103 chunks: list[bytes],
104 chunk: bytes,
105 received_bytes: int,
106 max_bytes: int | None,
107 media_download_limit: bool,
108 ) -> int:
109 if not chunk:
110 return received_bytes
111 received_bytes += len(chunk)
112 if max_bytes is not None and received_bytes > max_bytes:
113 _raise_size_exceeded(max_bytes, received_bytes, media_download_limit)
114 chunks.append(chunk)
115 return received_bytes
116
117
118 def _read_response_bytes(response: requests.Response, max_bytes: int | None) -> bytes:
119 max_bytes, media_download_limit = _get_response_size_limit(max_bytes)
120 if max_bytes is None:
121 return response.content
122
123 _check_content_length(
124 response.headers.get("Content-Length"),
125 max_bytes,
126 media_download_limit,
127 )
128 chunks: list[bytes] = []
129 received_bytes = 0
130 for chunk in response.iter_content(_RESPONSE_READ_CHUNK_SIZE):
131 received_bytes = _append_chunk(
132 chunks, chunk, received_bytes, max_bytes, media_download_limit
133 )
134 return b"".join(chunks)
135
136
137 async def _async_read_response_bytes(
138 response: aiohttp.ClientResponse, max_bytes: int | None
139 ) -> bytes:
140 max_bytes, media_download_limit = _get_response_size_limit(max_bytes)
141 if max_bytes is None:
142 return await response.read()
143
144 _check_content_length(response.content_length, max_bytes, media_download_limit)
145 chunks: list[bytes] = []
146 received_bytes = 0
147 async for chunk in response.content.iter_chunked(_RESPONSE_READ_CHUNK_SIZE):
148 received_bytes = _append_chunk(
149 chunks, chunk, received_bytes, max_bytes, media_download_limit
150 )
151 return b"".join(chunks)
152
153
154 def _is_retryable(exc: Exception) -> bool:
155 """Return True for transient errors that are worth retrying.
156
157 Retryable:
158 - Timeouts (aiohttp, requests, stdlib)
159 - Connection-level failures (refused, reset, DNS)
160 - Server errors (5xx) -- includes S3 503 SlowDown
161 Not retryable:
162 - Client errors (4xx) -- bad URL, auth, not-found
163 - Programming errors (ValueError, TypeError, ...)
164 """
165 # Timeouts
166 if isinstance(
167 exc,
168 (
169 TimeoutError,
170 asyncio.TimeoutError,
171 requests.exceptions.Timeout,
172 aiohttp.ServerTimeoutError,
173 ),
174 ):
175 return True
176 # Connection-level failures
177 if isinstance(
178 exc,
179 (
180 ConnectionError,
181 aiohttp.ClientConnectionError,
182 requests.exceptions.ConnectionError,
183 ),
184 ):
185 return True
186 # aiohttp server-side disconnects
187 if isinstance(exc, aiohttp.ServerDisconnectedError):
188 return True
189 # requests 5xx -- raise_for_status() throws HTTPError
190 if (
191 isinstance(exc, requests.exceptions.HTTPError)
192 and exc.response is not None
193 and exc.response.status_code >= 500
194 ):
195 return True
196 # aiohttp 5xx -- raise_for_status() throws ClientResponseError
197 return isinstance(exc, aiohttp.ClientResponseError) and exc.status >= 500
198
199
200 def _log_retry(
201 args: tuple,
202 kwargs: dict,
203 attempt: int,
204 max_retries: int,
205 attempt_timeout: float | None,
206 exc: Exception,
207 backoff: float,
208 base_timeout: float | None,
209 ) -> None:
210 # args[0] is `self` (bound method), args[1] is the URL
211 url = args[1] if len(args) > 1 else kwargs.get("url")
212 timeout_info = (
213 f"timeout={attempt_timeout:.3f}s" if base_timeout is not None else "no timeout"
214 )
215 next_timeout = (
216 f" with timeout={base_timeout * (_RETRY_BACKOFF_FACTOR ** (attempt + 1)):.3f}s"
217 if base_timeout is not None
218 else ""
219 )
220 logger.warning(
221 "HTTP fetch failed for %s (attempt %d/%d, %s): %s -- retrying in %.3fs%s",
222 url,
223 attempt + 1,
224 max_retries,
225 timeout_info,
226 exc,
227 backoff,
228 next_timeout,
229 )
230
231
232 def _sync_retry(
233 fn: Callable[_P, _T],
234 ) -> Callable[_P, _T]:
235 """Add retry logic with exponential backoff to a sync method.
236
237 The decorated method must accept ``timeout`` as a keyword argument.
238 The decorator replaces it with a per-attempt timeout that grows by
239 ``_RETRY_BACKOFF_FACTOR`` on each retry so transient slowness on busy
240 hosts is absorbed.
241 """
242
243 @functools.wraps(fn)
244 def wrapper(*args: Any, **kwargs: Any) -> _T:
245 base_timeout: float | None = kwargs.get("timeout")
246 max_retries = max(envs.VLLM_MEDIA_FETCH_MAX_RETRIES, 1)
247
248 for attempt in range(max_retries):
249 attempt_timeout = (
250 base_timeout * (_RETRY_BACKOFF_FACTOR**attempt)
251 if base_timeout is not None
252 else None
253 )
254 kwargs["timeout"] = attempt_timeout
255 try:
256 return fn(*args, **kwargs)
257 except Exception as e:
258 if not _is_retryable(e) or attempt + 1 >= max_retries:
259 raise
260 backoff = _RETRY_BACKOFF_FACTOR**attempt
261 _log_retry(
262 args,
263 kwargs,
264 attempt,
265 max_retries,
266 attempt_timeout,
267 e,
268 backoff,
269 base_timeout,
270 )
271 time.sleep(backoff)
272
273 raise AssertionError("unreachable")
274
275 return wrapper # type: ignore[return-value]
276
277
278 def _async_retry(
279 fn: Callable[_P, Coroutine[Any, Any, _T]],
280 ) -> Callable[_P, Coroutine[Any, Any, _T]]:
281 """Add retry logic with exponential backoff to an async method.
282
283 The decorated method must accept ``timeout`` as a keyword argument.
284 The decorator replaces it with a per-attempt timeout that grows by
285 ``_RETRY_BACKOFF_FACTOR`` on each retry so transient slowness on busy
286 hosts is absorbed.
287 """
288
289 @functools.wraps(fn)
290 async def wrapper(*args: Any, **kwargs: Any) -> _T:
291 base_timeout: float | None = kwargs.get("timeout")
292 max_retries = max(envs.VLLM_MEDIA_FETCH_MAX_RETRIES, 1)
293
294 for attempt in range(max_retries):
295 attempt_timeout = (
296 base_timeout * (_RETRY_BACKOFF_FACTOR**attempt)
297 if base_timeout is not None
298 else None
299 )
300 kwargs["timeout"] = attempt_timeout
301 try:
302 return await fn(*args, **kwargs)
303 except Exception as e:
304 if not _is_retryable(e) or attempt + 1 >= max_retries:
305 raise
306 backoff = _RETRY_BACKOFF_FACTOR**attempt
307 _log_retry(
308 args,
309 kwargs,
310 attempt,
311 max_retries,
312 attempt_timeout,
313 e,
314 backoff,
315 base_timeout,
316 )
317 await asyncio.sleep(backoff)
318
319 raise AssertionError("unreachable")
320
321 return wrapper # type: ignore[return-value]
322
323
324 class HTTPConnection:
325 """Helper class to send HTTP requests."""
326
327 def __init__(self, *, reuse_client: bool = True) -> None:
328 super().__init__()
329
330 self.reuse_client = reuse_client
331
332 self._sync_client: requests.Session | None = None
333 self._async_client: aiohttp.ClientSession | None = None
334
335 def get_sync_client(self) -> requests.Session:
336 if self._sync_client is None or not self.reuse_client:
337 self._sync_client = requests.Session()
338
339 return self._sync_client
340
341 # NOTE: We intentionally use an async function even though it is not
342 # required, so that the client is only accessible inside async event loop
343 async def get_async_client(self) -> aiohttp.ClientSession:
344 if self._async_client is None or not self.reuse_client:
345 self._async_client = aiohttp.ClientSession(trust_env=True)
346
347 return self._async_client
348
349 def _validate_http_url(self, url: str):
350 parsed_url = parse_url(url)
351
352 if parsed_url.scheme not in ("http", "https"):
353 raise ValueError(
354 "Invalid HTTP URL: A valid HTTP URL must have scheme 'http' or 'https'."
355 )
356
357 def _headers(self, **extras: str) -> MutableMapping[str, str]:
358 return {"User-Agent": f"vLLM/{VLLM_VERSION}", **extras}
359
360 def get_response(
361 self,
362 url: str,
363 *,
364 stream: bool = False,
365 timeout: float | None = None,
366 extra_headers: Mapping[str, str] | None = None,
367 allow_redirects: bool = True,
368 ):
369 self._validate_http_url(url)
370
371 client = self.get_sync_client()
372 extra_headers = extra_headers or {}
373
374 return client.get(
375 url,
376 headers=self._headers(**extra_headers),
377 stream=stream,
378 timeout=timeout,
379 allow_redirects=allow_redirects,
380 )
381
382 async def get_async_response(
383 self,
384 url: str,
385 *,
386 timeout: float | None = None,
387 extra_headers: Mapping[str, str] | None = None,
388 allow_redirects: bool = True,
389 ):
390 self._validate_http_url(url)
391
392 client = await self.get_async_client()
393 extra_headers = extra_headers or {}
394
395 return client.get(
396 url,
397 headers=self._headers(**extra_headers),
398 timeout=timeout,
399 allow_redirects=allow_redirects,
400 )
401
402 @_sync_retry
403 def get_bytes(
404 self,
405 url: str,
406 *,
407 timeout: float | None = None,
408 allow_redirects: bool = True,
409 max_bytes: int | None = None,
410 ) -> bytes:
411 with self.get_response(
412 url,
413 stream=True,
414 timeout=timeout,
415 allow_redirects=allow_redirects,
416 ) as r:
417 r.raise_for_status()
418
419 return _read_response_bytes(r, max_bytes)
420
421 @_async_retry
422 async def async_get_bytes(
423 self,
424 url: str,
425 *,
426 timeout: float | None = None,
427 allow_redirects: bool = True,
428 max_bytes: int | None = None,
429 ) -> bytes:
430 async with await self.get_async_response(
431 url,
432 timeout=timeout,
433 allow_redirects=allow_redirects,
434 ) as r:
435 r.raise_for_status()
436
437 return await _async_read_response_bytes(r, max_bytes)
438
439 def get_text(self, url: str, *, timeout: float | None = None) -> str:
440 with self.get_response(url, timeout=timeout) as r:
441 r.raise_for_status()
442
443 return r.text
444
445 async def async_get_text(
446 self,
447 url: str,
448 *,
449 timeout: float | None = None,
450 ) -> str:
451 async with await self.get_async_response(url, timeout=timeout) as r:
452 r.raise_for_status()
453
454 return await r.text()
455
456 def get_json(self, url: str, *, timeout: float | None = None) -> str:
457 with self.get_response(url, timeout=timeout) as r:
458 r.raise_for_status()
459
460 return r.json()
461
462 async def async_get_json(
463 self,
464 url: str,
465 *,
466 timeout: float | None = None,
467 ) -> str:
468 async with await self.get_async_response(url, timeout=timeout) as r:
469 r.raise_for_status()
470
471 return await r.json()
472
473 @_sync_retry
474 def download_file(
475 self,
476 url: str,
477 save_path: Path,
478 *,
479 timeout: float | None = None,
480 chunk_size: int = 128,
481 ) -> Path:
482 try:
483 with self.get_response(url, timeout=timeout) as r:
484 r.raise_for_status()
485
486 with save_path.open("wb") as f:
487 for chunk in r.iter_content(chunk_size):
488 f.write(chunk)
489
490 return save_path
491 except Exception:
492 # Clean up partial downloads before retrying or propagating
493 if save_path.exists():
494 save_path.unlink()
495 raise
496
497 @_async_retry
498 async def async_download_file(
499 self,
500 url: str,
501 save_path: Path,
502 *,
503 timeout: float | None = None,
504 chunk_size: int = 128,
505 ) -> Path:
506 try:
507 async with await self.get_async_response(
508 url,
509 timeout=timeout,
510 ) as r:
511 r.raise_for_status()
512
513 with save_path.open("wb") as f:
514 async for chunk in r.content.iter_chunked(chunk_size):
515 f.write(chunk)
516
517 return save_path
518 except Exception:
519 # Clean up partial downloads before retrying or propagating
520 if save_path.exists():
521 save_path.unlink()
522 raise
523
524
525 global_http_connection = HTTPConnection()
526 """
527 The global [`HTTPConnection`][vllm.connections.HTTPConnection] instance used
528 by vLLM.
529 """