main
py 660 lines 22.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import asyncio
5 import atexit
6 import contextlib
7 import hashlib
8 import os
9 import tempfile
10 import time
11 from concurrent.futures import ThreadPoolExecutor
12 from pathlib import Path
13 from typing import Any, TypeVar
14 from urllib.request import url2pathname
15
16 import aiohttp
17 import numpy as np
18 import numpy.typing as npt
19 import requests
20 import torch
21 from PIL import Image, UnidentifiedImageError
22 from urllib3.util import Url, parse_url
23
24 import vllm.envs as envs
25 from vllm.connections import (
26 HTTPConnection,
27 MediaDownloadSizeExceededError,
28 global_http_connection,
29 )
30 from vllm.exceptions import VLLMUnprocessableEntityError, VLLMValidationError
31 from vllm.logger import init_logger
32 from vllm.multimodal.video import get_video_loader_backend_for_processor
33 from vllm.utils.registry import ExtensionManager
34
35 from .audio import AudioEmbeddingMediaIO, AudioMediaIO
36 from .base import MediaIO, MediaWithBytes
37 from .image import ImageEmbeddingMediaIO, ImageMediaIO
38 from .video import VideoEmbeddingMediaIO, VideoMediaIO
39
40 logger = init_logger(__name__)
41
42 _M = TypeVar("_M")
43
44 global_thread_pool = ThreadPoolExecutor(
45 max_workers=envs.VLLM_MEDIA_LOADING_THREAD_COUNT
46 )
47 atexit.register(global_thread_pool.shutdown)
48
49 MEDIA_CONNECTOR_REGISTRY = ExtensionManager()
50
51 MODALITY_IO_MAP: dict[str, type[MediaIO]] = {
52 "audio": AudioMediaIO,
53 "image": ImageMediaIO,
54 "video": VideoMediaIO,
55 }
56
57
58 def _wrap_media_fetch_error(
59 url: str, exc: Exception
60 ) -> VLLMUnprocessableEntityError | Exception:
61 """Convert media fetch exceptions to VLLMUnprocessableEntityError.
62
63 This handles HTTP errors that indicate the media resource is invalid
64 (4xx responses except 408/429, malformed URLs) and converts them to a
65 422 Unprocessable Entity error instead of 500.
66
67 Transient errors (5xx, 408, 429, DNS failures, connection errors,
68 timeouts) are returned as-is to allow retry logic to handle them
69 appropriately.
70
71 Returns:
72 VLLMUnprocessableEntityError for permanent client errors (4xx except
73 408/429, invalid URL)
74 Original exception for transient errors (5xx, 408, 429, network blips)
75 or other exceptions
76
77 """
78 if isinstance(exc, VLLMValidationError):
79 return exc
80
81 if isinstance(exc, aiohttp.ClientResponseError):
82 if exc.status in (408, 429):
83 return exc
84 if exc.status < 500:
85 return VLLMUnprocessableEntityError(
86 f"Failed to fetch media from URL: HTTP {exc.status} error",
87 parameter="image_url",
88 value=url,
89 )
90 return exc
91
92 if isinstance(exc, requests.exceptions.HTTPError):
93 if exc.response is not None:
94 status_code = exc.response.status_code
95 if status_code in (408, 429):
96 return exc
97 if status_code < 500:
98 return VLLMUnprocessableEntityError(
99 f"Failed to fetch media from URL: HTTP {status_code} error",
100 parameter="image_url",
101 value=url,
102 )
103 return exc
104
105 if isinstance(exc, requests.exceptions.InvalidURL):
106 return VLLMUnprocessableEntityError(
107 "Failed to fetch media from URL: Invalid URL format",
108 parameter="image_url",
109 value=url,
110 )
111
112 if isinstance(exc, MediaDownloadSizeExceededError):
113 return VLLMUnprocessableEntityError(
114 f"Failed to fetch media from URL: {exc}",
115 parameter="image_url",
116 value=url,
117 )
118
119 if isinstance(exc, ValueError):
120 return VLLMUnprocessableEntityError(
121 "Failed to fetch media from URL: Invalid URL",
122 parameter="image_url",
123 value=url,
124 )
125 return exc
126
127
128 def merge_media_io_kwargs(
129 defaults: dict[str, dict[str, Any]] | None,
130 overrides: dict[str, dict[str, Any]] | None,
131 ) -> dict[str, dict[str, Any]] | None:
132 """Merge config-level and per-request media_io_kwargs per modality.
133
134 Each modality key is merged using the corresponding MediaIO subclass's
135 ``merge_kwargs``, which may apply modality-specific logic (e.g.
136 VideoMediaIO clears cross-dependent fps/num_frames fields).
137 """
138 if not defaults and not overrides:
139 return None
140 all_keys = set(defaults or {}) | set(overrides or {})
141 merged = {}
142 for key in all_keys:
143 io_cls = MODALITY_IO_MAP.get(key, MediaIO)
144 merged[key] = io_cls.merge_kwargs(
145 (defaults or {}).get(key),
146 (overrides or {}).get(key),
147 )
148 return merged or None
149
150
151 @MEDIA_CONNECTOR_REGISTRY.register("http")
152 class MediaConnector:
153 """Configuration values can be user-provided either by --media-io-kwargs or
154 by the runtime API field "media_io_kwargs". Ensure proper validation and
155 error handling.
156 """
157
158 def __init__(
159 self,
160 media_io_kwargs: dict[str, dict[str, Any]] | None = None,
161 connection: HTTPConnection = global_http_connection,
162 *,
163 allowed_local_media_path: str = "",
164 allowed_media_domains: list[str] | None = None,
165 ) -> None:
166 """Args:
167 media_io_kwargs: Additional args passed to process media
168 inputs, keyed by modalities. For example,
169 to set num_frames for video, set
170 `--media-io-kwargs '{"video":{"num_frames":40}}'`
171 connection: HTTP connection client to download media contents.
172 allowed_local_media_path: A local directory to load media files from.
173 allowed_media_domains: If set, only media URLs that belong to this
174 domain can be used for multi-modal inputs.
175
176 """
177 super().__init__()
178
179 self.media_io_kwargs: dict[str, dict[str, Any]] = (
180 media_io_kwargs if media_io_kwargs else {}
181 )
182 self.connection = connection
183
184 if allowed_local_media_path:
185 allowed_local_media_path_ = Path(allowed_local_media_path).resolve()
186
187 if not allowed_local_media_path_.exists():
188 raise ValueError(
189 "Invalid `--allowed-local-media-path`: The path "
190 f"{allowed_local_media_path_} does not exist."
191 )
192 if not allowed_local_media_path_.is_dir():
193 raise ValueError(
194 "Invalid `--allowed-local-media-path`: The path "
195 f"{allowed_local_media_path_} must be a directory."
196 )
197 else:
198 allowed_local_media_path_ = None
199
200 self.allowed_local_media_path = allowed_local_media_path_
201 if allowed_media_domains is None:
202 allowed_media_domains = []
203 self.allowed_media_domains = allowed_media_domains
204
205 # Media download cache (opt-in via VLLM_MEDIA_CACHE)
206 self._media_cache_dir: str | None = None
207 self._media_cache_max_bytes: int = 0
208 self._media_cache_ttl_secs: float = 0
209 media_cache = envs.VLLM_MEDIA_CACHE
210 if media_cache:
211 try:
212 os.makedirs(media_cache, exist_ok=True)
213 # Verify the directory is writable before enabling caching
214 with tempfile.NamedTemporaryFile(dir=media_cache, delete=True):
215 pass
216 self._media_cache_dir = media_cache
217 self._media_cache_max_bytes = (
218 envs.VLLM_MEDIA_CACHE_MAX_SIZE_MB * 1024 * 1024
219 )
220 self._media_cache_ttl_secs = envs.VLLM_MEDIA_CACHE_TTL_HOURS * 3600
221 logger.info(
222 "Media cache enabled at %s (max %d MB, TTL %s hours)",
223 media_cache,
224 envs.VLLM_MEDIA_CACHE_MAX_SIZE_MB,
225 envs.VLLM_MEDIA_CACHE_TTL_HOURS,
226 )
227 except OSError:
228 logger.warning(
229 "VLLM_MEDIA_CACHE path %s is not writable, media caching disabled",
230 media_cache,
231 )
232
233 def _get_cached_bytes(self, url: str) -> bytes | None:
234 """Return cached bytes for a URL, or None if not cached/expired."""
235 if not self._media_cache_dir:
236 return None
237 cache_path = self._media_cache_path(url)
238 # Check TTL
239 try:
240 age = time.time() - cache_path.stat().st_mtime
241 except OSError:
242 return None
243 if age > self._media_cache_ttl_secs:
244 cache_path.unlink(missing_ok=True)
245 return None
246 # Touch mtime for LRU ordering
247 try:
248 cache_path.touch()
249 return cache_path.read_bytes()
250 except OSError:
251 return None
252
253 def _put_cached_bytes(self, url: str, data: bytes) -> None:
254 """Store downloaded bytes and evict if over budget."""
255 if not self._media_cache_dir:
256 return
257 cache_path = self._media_cache_path(url)
258 # Atomic write via temp file + rename
259 tmp_path = None
260 try:
261 with tempfile.NamedTemporaryFile(
262 mode="wb", dir=self._media_cache_dir, delete=False
263 ) as tmp_file:
264 tmp_file.write(data)
265 tmp_path = tmp_file.name
266 os.rename(tmp_path, str(cache_path))
267 except OSError:
268 # Another process beat us or disk issue
269 if tmp_path is not None:
270 with contextlib.suppress(OSError):
271 os.remove(tmp_path)
272 return
273 self._maybe_evict(exclude=cache_path)
274
275 def _maybe_evict(self, exclude: Path | None = None) -> None:
276 """Evict expired entries first, then LRU until under size limit."""
277 cache_dir = Path(self._media_cache_dir) # type: ignore[arg-type]
278 entries = []
279 expired = []
280 total_size = 0
281 now = time.time()
282 for f in cache_dir.iterdir():
283 if f.name.startswith("."):
284 continue
285 try:
286 stat = f.stat()
287 except OSError:
288 continue
289 age = now - stat.st_mtime
290 if age > self._media_cache_ttl_secs:
291 expired.append(f)
292 continue
293 total_size += stat.st_size
294 # Never evict the file we just wrote
295 if exclude is not None and f.name == exclude.name:
296 continue
297 entries.append((stat.st_mtime, stat.st_size, f))
298
299 # Evict items according to LRU policy
300 entries.sort(key=lambda e: e[0], reverse=True)
301 while total_size > self._media_cache_max_bytes and entries:
302 mtime, size, f = entries.pop()
303 expired.append(f)
304 total_size -= size
305
306 for f in expired:
307 f.unlink(missing_ok=True)
308
309 def _media_cache_path(self, url: str) -> Path:
310 url_hash = hashlib.sha256(url.encode()).hexdigest()[:20]
311 ext = Path(url.split("?", 1)[0]).suffix or ""
312 return Path(self._media_cache_dir) / f"{url_hash}{ext}" # type: ignore[arg-type]
313
314 def _load_data_url(
315 self,
316 url: str,
317 media_io: MediaIO[_M],
318 ) -> _M: # type: ignore[type-var]
319 # Format per RFC 2397:
320 # data:[<mediatype>][;<param>=<value>]*[;base64],<data>
321 data_spec, sep, data = url[5:].partition(",")
322 if not sep:
323 msg = f"Invalid data URL {url[:32]!r}: missing ',' separator."
324 raise ValueError(msg)
325
326 media_type, sep, encoding = data_spec.rpartition(";")
327 if not sep or encoding != "base64":
328 msg = "Only base64 data URLs are supported for now."
329 raise NotImplementedError(msg)
330
331 media_type = media_type.partition(";")[0]
332 return media_io.load_base64(media_type, data)
333
334 def _load_file_url(
335 self,
336 url_spec: Url,
337 media_io: MediaIO[_M],
338 ) -> _M: # type: ignore[type-var]
339 allowed_local_media_path = self.allowed_local_media_path
340 if allowed_local_media_path is None:
341 raise RuntimeError(
342 "Cannot load local files without `--allowed-local-media-path`."
343 )
344
345 url_spec_path = url_spec.path or ""
346 url_spec_netloc = url_spec.netloc or ""
347 filepath = Path(url2pathname(url_spec_netloc + url_spec_path))
348 if allowed_local_media_path not in filepath.resolve().parents:
349 raise ValueError(
350 f"The file path {filepath} must be a subpath "
351 f"of `--allowed-local-media-path {allowed_local_media_path}`."
352 )
353
354 return media_io.load_file(filepath)
355
356 def _assert_url_in_allowed_media_domains(self, url_spec: Url) -> None:
357 if (
358 self.allowed_media_domains
359 and url_spec.hostname not in self.allowed_media_domains
360 ):
361 raise ValueError(
362 f"The URL must be from one of the allowed domains: "
363 f"{self.allowed_media_domains}. Input URL domain: "
364 f"{url_spec.hostname}"
365 )
366
367 def load_from_url(
368 self,
369 url: str,
370 media_io: MediaIO[_M],
371 *,
372 fetch_timeout: int | None = None,
373 ) -> _M: # type: ignore[type-var]
374 if url[:5].lower() == "data:":
375 return self._load_data_url(url, media_io)
376
377 url_spec = parse_url(url)
378
379 if url_spec.scheme and url_spec.scheme.startswith("http"):
380 self._assert_url_in_allowed_media_domains(url_spec)
381 max_bytes = media_io.get_max_bytes()
382
383 cached = self._get_cached_bytes(url)
384 if cached is not None:
385 return media_io.load_bytes(cached)
386
387 connection = self.connection
388 try:
389 data = connection.get_bytes(
390 url_spec.url,
391 timeout=fetch_timeout,
392 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
393 max_bytes=max_bytes,
394 )
395 except Exception as e:
396 wrapped = _wrap_media_fetch_error(url, e)
397 if isinstance(wrapped, VLLMUnprocessableEntityError):
398 raise wrapped from e
399 raise
400
401 self._put_cached_bytes(url, data)
402 return media_io.load_bytes(data)
403
404 if url_spec.scheme == "file":
405 return self._load_file_url(url_spec, media_io)
406
407 msg = "The URL must be either a HTTP, data or file URL."
408 raise ValueError(msg)
409
410 async def load_from_url_async(
411 self,
412 url: str,
413 media_io: MediaIO[_M],
414 *,
415 fetch_timeout: int | None = None,
416 ) -> _M:
417 loop = asyncio.get_running_loop()
418
419 if url[:5].lower() == "data:":
420 future = loop.run_in_executor(
421 global_thread_pool, self._load_data_url, url, media_io
422 )
423 return await future
424
425 url_spec = parse_url(url)
426
427 if url_spec.scheme and url_spec.scheme.startswith("http"):
428 self._assert_url_in_allowed_media_domains(url_spec)
429 max_bytes = media_io.get_max_bytes()
430
431 cached = await loop.run_in_executor(
432 global_thread_pool, self._get_cached_bytes, url
433 )
434 if cached is not None:
435 future = loop.run_in_executor(
436 global_thread_pool, media_io.load_bytes, cached
437 )
438 return await future
439
440 connection = self.connection
441 try:
442 data = await connection.async_get_bytes(
443 url_spec.url,
444 timeout=fetch_timeout,
445 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
446 max_bytes=max_bytes,
447 )
448 except Exception as e:
449 wrapped = _wrap_media_fetch_error(url, e)
450 if isinstance(wrapped, VLLMUnprocessableEntityError):
451 raise wrapped from e
452 raise
453
454 await loop.run_in_executor(
455 global_thread_pool, self._put_cached_bytes, url, data
456 )
457 future = loop.run_in_executor(global_thread_pool, media_io.load_bytes, data)
458 return await future
459
460 if url_spec.scheme == "file":
461 future = loop.run_in_executor(
462 global_thread_pool, self._load_file_url, url_spec, media_io
463 )
464 return await future
465 msg = "The URL must be either a HTTP, data or file URL."
466 raise ValueError(msg)
467
468 def fetch_audio(
469 self,
470 audio_url: str,
471 ) -> tuple[np.ndarray, int | float]:
472 """Load audio from a URL."""
473 audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
474
475 return self.load_from_url(
476 audio_url,
477 audio_io,
478 fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
479 )
480
481 async def fetch_audio_async(
482 self,
483 audio_url: str,
484 ) -> tuple[np.ndarray, int | float]:
485 """Asynchronously fetch audio from a URL."""
486 audio_io = AudioMediaIO(**self.media_io_kwargs.get("audio", {}))
487
488 return await self.load_from_url_async(
489 audio_url,
490 audio_io,
491 fetch_timeout=envs.VLLM_AUDIO_FETCH_TIMEOUT,
492 )
493
494 def fetch_image(
495 self,
496 image_url: str,
497 *,
498 image_mode: str | None = "RGB",
499 ) -> Image.Image:
500 """Load a PIL image from an HTTP or base64 data URL.
501
502 By default, the image is converted into RGB format. Set
503 `media_io_kwargs={"image": {"image_mode": None}}` to keep the
504 original image mode (e.g. preserving the alpha channel).
505 """
506 image_io = ImageMediaIO(
507 **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
508 )
509
510 try:
511 return self.load_from_url(
512 image_url,
513 image_io,
514 fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
515 )
516 except UnidentifiedImageError as e:
517 # convert to ValueError to be properly caught upstream
518 raise ValueError(str(e)) from e
519
520 async def fetch_image_async(
521 self,
522 image_url: str,
523 *,
524 image_mode: str | None = "RGB",
525 ) -> Image.Image:
526 """Asynchronously load a PIL image from an HTTP or base64 data URL.
527
528 By default, the image is converted into RGB format. Set
529 `media_io_kwargs={"image": {"image_mode": None}}` to keep the
530 original image mode (e.g. preserving the alpha channel).
531 """
532 image_io = ImageMediaIO(
533 **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
534 )
535
536 try:
537 return await self.load_from_url_async(
538 image_url,
539 image_io,
540 fetch_timeout=envs.VLLM_IMAGE_FETCH_TIMEOUT,
541 )
542 except UnidentifiedImageError as e:
543 # convert to ValueError to be properly caught upstream
544 raise ValueError(str(e)) from e
545
546 def fetch_video(
547 self,
548 video_url: str,
549 *,
550 image_mode: str | None = "RGB",
551 video_processor: str | None = None,
552 ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
553 """Load video from an HTTP or base64 data URL."""
554 image_io = ImageMediaIO(
555 **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
556 )
557 video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
558 if "video_backend" not in video_io_kwargs and (
559 video_backend := get_video_loader_backend_for_processor(video_processor)
560 ):
561 video_io_kwargs["video_backend"] = video_backend
562 video_io = VideoMediaIO(image_io, **video_io_kwargs)
563
564 return self.load_from_url(
565 video_url,
566 video_io,
567 fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
568 )
569
570 async def fetch_video_async(
571 self,
572 video_url: str,
573 *,
574 image_mode: str | None = "RGB",
575 video_processor: str | None = None,
576 ) -> MediaWithBytes[tuple[npt.NDArray, dict[str, Any]]]:
577 """Asynchronously load video from an HTTP or base64 data URL.
578
579 By default, the image is converted into RGB format. Set
580 `media_io_kwargs={"image": {"image_mode": None}}` to keep the
581 original image mode (e.g. preserving the alpha channel).
582 """
583 image_io = ImageMediaIO(
584 **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
585 )
586 video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
587 if "video_backend" not in video_io_kwargs and (
588 video_backend := get_video_loader_backend_for_processor(video_processor)
589 ):
590 video_io_kwargs["video_backend"] = video_backend
591 video_io = VideoMediaIO(image_io, **video_io_kwargs)
592
593 return await self.load_from_url_async(
594 video_url,
595 video_io,
596 fetch_timeout=envs.VLLM_VIDEO_FETCH_TIMEOUT,
597 )
598
599 def fetch_image_embedding(
600 self,
601 data: str,
602 ) -> torch.Tensor:
603 """Load image embedding from a URL."""
604 image_embedding_io = ImageEmbeddingMediaIO()
605
606 return image_embedding_io.load_base64("", data)
607
608 async def fetch_image_embedding_async(
609 self,
610 data: str,
611 ) -> torch.Tensor:
612 """Asynchronously load image embedding from a URL."""
613 image_embedding_io = ImageEmbeddingMediaIO()
614 loop = asyncio.get_running_loop()
615
616 return await loop.run_in_executor(
617 global_thread_pool, image_embedding_io.load_base64, "", data
618 )
619
620 def fetch_audio_embedding(
621 self,
622 data: str,
623 ) -> torch.Tensor:
624 """Load audio embedding from a URL."""
625 audio_embedding_io = AudioEmbeddingMediaIO()
626
627 return audio_embedding_io.load_base64("", data)
628
629 async def fetch_audio_embedding_async(
630 self,
631 data: str,
632 ) -> torch.Tensor:
633 """Asynchronously load audio embedding from a URL."""
634 audio_embedding_io = AudioEmbeddingMediaIO()
635 loop = asyncio.get_running_loop()
636
637 return await loop.run_in_executor(
638 global_thread_pool, audio_embedding_io.load_base64, "", data
639 )
640
641 def fetch_video_embedding(
642 self,
643 data: str,
644 ) -> torch.Tensor:
645 """Load video embedding from a URL."""
646 video_embedding_io = VideoEmbeddingMediaIO()
647
648 return video_embedding_io.load_base64("", data)
649
650 async def fetch_video_embedding_async(
651 self,
652 data: str,
653 ) -> torch.Tensor:
654 """Asynchronously load video embedding from a URL."""
655 video_embedding_io = VideoEmbeddingMediaIO()
656 loop = asyncio.get_running_loop()
657
658 return await loop.run_in_executor(
659 global_thread_pool, video_embedding_io.load_base64, "", data
660 )