@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 677424c7ac

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 19:45 UTC

[Core][CI] Add opt-in media URL caching via VLLM_MEDIA_CACHE (#37123)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed Mar 30, 2026 at 06:58 UTC 677424c7acd9fb7477294017c99f798588002d4f

3 files changed +266

tests/multimodal/media/test_connector.py

+112

@@ -4,6 +4,8 @@

4 import asyncio

5 import mimetypes

6 import os

7 +import shutil

8 +import time

9 from tempfile import NamedTemporaryFile, TemporaryDirectory

11 import aiohttp

@@ -375,3 +377,113 @@ async def test_ssrf_bypass_backslash_disallowed_domain():

377

378 with pytest.raises(ValueError, match="allowed domains"):

379 await connector.fetch_image_async(bypass_url)

380 +

381 +

382 +def _make_cached_connector(cache_dir, *, max_mb=10, ttl_hours=24):

383 + """Create a MediaConnector with caching enabled via monkeypatched internals.

384 +

385 + We bypass __init__'s env-var path and wire up the cache fields directly

386 + so tests don't depend on environment variables. URLs in these tests are

387 + only used as cache keys (hashed to derive filenames); no HTTP requests

388 + are made.

389 + """

390 + connector = MediaConnector()

391 + connector._media_cache_dir = cache_dir

392 + connector._media_cache_max_bytes = max_mb * 1024 * 1024

393 + connector._media_cache_ttl_secs = ttl_hours * 3600

394 + return connector

395 +

396 +

397 +def test_cache_put_and_get():

398 + """Basic round-trip: put bytes, get them back."""

399 + with TemporaryDirectory() as cache_dir:

400 + connector = _make_cached_connector(cache_dir)

401 + url = "https://example.com/image.png"

402 + data = b"fake-image-bytes"

403 +

404 + connector._put_cached_bytes(url, data)

405 + cached = connector._get_cached_bytes(url)

406 + assert cached == data

407 +

408 +

409 +def test_cache_ttl_expiry():

410 + """Entries older than TTL are evicted on read."""

411 + with TemporaryDirectory() as cache_dir:

412 + connector = _make_cached_connector(cache_dir, ttl_hours=24)

413 + url = "https://example.com/old.png"

414 + data = b"old-data"

415 +

416 + connector._put_cached_bytes(url, data)

417 +

418 + # Backdate the file's mtime so it appears expired

419 + cache_path = connector._media_cache_path(url)

420 + expired_time = time.time() - (25 * 3600) # 25 hours ago

421 + os.utime(cache_path, (expired_time, expired_time))

422 +

423 + assert connector._get_cached_bytes(url) is None

424 + assert not cache_path.exists()

425 +

426 +

427 +def test_cache_lru_eviction():

428 + """Oldest entries are evicted when cache exceeds size budget."""

429 + with TemporaryDirectory() as cache_dir:

430 + # Set a very small max size: 100 bytes

431 + connector = _make_cached_connector(cache_dir, max_mb=0)

432 + connector._media_cache_max_bytes = 100

433 +

434 + # Write three 50-byte entries (total 150 > 100 budget)

435 + urls = [f"https://example.com/{i}.png" for i in range(3)]

436 + for i, url in enumerate(urls):

437 + connector._put_cached_bytes(url, b"x" * 50)

438 + # Stagger mtime so eviction order is deterministic

439 + path = connector._media_cache_path(url)

440 + os.utime(path, (time.time() + i, time.time() + i))

441 +

442 + # The oldest entry (urls[0]) should have been evicted

443 + assert connector._get_cached_bytes(urls[0]) is None

444 + # The newest entries should still be present

445 + assert connector._get_cached_bytes(urls[2]) == b"x" * 50

446 +

447 +

448 +def test_cache_ttl_eviction_during_write():

449 + """_maybe_evict removes expired files even if under size budget."""

450 + with TemporaryDirectory() as cache_dir:

451 + connector = _make_cached_connector(cache_dir, ttl_hours=1)

452 + url_old = "https://example.com/stale.png"

453 + url_new = "https://example.com/fresh.png"

454 +

455 + connector._put_cached_bytes(url_old, b"stale")

456 + # Backdate old entry past TTL

457 + old_path = connector._media_cache_path(url_old)

458 + expired_time = time.time() - (2 * 3600)

459 + os.utime(old_path, (expired_time, expired_time))

460 +

461 + # Writing a new entry triggers _maybe_evict

462 + connector._put_cached_bytes(url_new, b"fresh")

463 +

464 + assert not old_path.exists()

465 + assert connector._get_cached_bytes(url_new) == b"fresh"

466 +

467 +

468 +def test_put_cached_bytes_missing_dir():

469 + """_put_cached_bytes does not crash when the cache dir disappears."""

470 + with TemporaryDirectory() as cache_dir:

471 + connector = _make_cached_connector(cache_dir)

472 + # Remove the directory to simulate it disappearing at runtime

473 + shutil.rmtree(cache_dir)

474 +

475 + # Should not raise (graceful degradation)

476 + connector._put_cached_bytes("https://example.com/x.png", b"data")

477 +

478 +

479 +def test_get_cached_bytes_file_deleted_before_read():

480 + """_get_cached_bytes returns None if the file vanishes mid-read."""

481 + with TemporaryDirectory() as cache_dir:

482 + connector = _make_cached_connector(cache_dir)

483 + url = "https://example.com/vanish.png"

484 +

485 + connector._put_cached_bytes(url, b"data")

486 + # Delete the file to simulate concurrent eviction

487 + connector._media_cache_path(url).unlink()

488 +

489 + assert connector._get_cached_bytes(url) is None

vllm/envs.py

+19

@@ -64,6 +64,9 @@ if TYPE_CHECKING:

64 VLLM_IMAGE_FETCH_TIMEOUT: int = 5

65 VLLM_VIDEO_FETCH_TIMEOUT: int = 30

66 VLLM_AUDIO_FETCH_TIMEOUT: int = 10

67 + VLLM_MEDIA_CACHE: str = ""

68 + VLLM_MEDIA_CACHE_MAX_SIZE_MB: int = 5120

69 + VLLM_MEDIA_CACHE_TTL_HOURS: float = 24

70 VLLM_MEDIA_FETCH_MAX_RETRIES: int = 3

71 VLLM_MEDIA_URL_ALLOW_REDIRECTS: bool = True

72 VLLM_MEDIA_LOADING_THREAD_COUNT: int = 8

@@ -776,6 +779,19 @@ environment_variables: dict[str, Callable[[], Any]] = {

779 "VLLM_AUDIO_FETCH_TIMEOUT": lambda: int(

780 os.getenv("VLLM_AUDIO_FETCH_TIMEOUT", "10")

781 ),

782 + # Directory for caching media downloads (images, video, audio fetched

783 + # from URLs during inference). Empty string disables caching.

784 + "VLLM_MEDIA_CACHE": lambda: os.getenv("VLLM_MEDIA_CACHE", ""),

785 + # Maximum cache size in MB. When exceeded, least-recently-used entries

786 + # are evicted. Default is 5120 (5 GB).

787 + "VLLM_MEDIA_CACHE_MAX_SIZE_MB": lambda: int(

788 + os.getenv("VLLM_MEDIA_CACHE_MAX_SIZE_MB", "5120")

789 + ),

790 + # Time-to-live in hours for cached media files. Entries older than this

791 + # are evicted regardless of cache size. Default is 24 hours.

792 + "VLLM_MEDIA_CACHE_TTL_HOURS": lambda: float(

793 + os.getenv("VLLM_MEDIA_CACHE_TTL_HOURS", "24")

794 + ),

795 # Maximum number of retries for fetching media (images, audio, video)

796 # from URLs. Each retry quadruples the timeout. Default is 3.

797 "VLLM_MEDIA_FETCH_MAX_RETRIES": lambda: int(

@@ -1777,6 +1793,9 @@ def compile_factors() -> dict[str, object]:

1793 "VLLM_IMAGE_FETCH_TIMEOUT",

1794 "VLLM_VIDEO_FETCH_TIMEOUT",

1795 "VLLM_AUDIO_FETCH_TIMEOUT",

1796 + "VLLM_MEDIA_CACHE",

1797 + "VLLM_MEDIA_CACHE_MAX_SIZE_MB",

1798 + "VLLM_MEDIA_CACHE_TTL_HOURS",

1799 "VLLM_MEDIA_FETCH_MAX_RETRIES",

1800 "VLLM_MEDIA_URL_ALLOW_REDIRECTS",

1801 "VLLM_MEDIA_LOADING_THREAD_COUNT",

vllm/multimodal/media/connector.py

+135

@@ -3,6 +3,11 @@

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

@@ -16,6 +21,7 @@ from urllib3.util import Url, parse_url

22 import vllm.envs as envs

23 from vllm.connections import HTTPConnection, global_http_connection

24 +from vllm.logger import init_logger

25 from vllm.utils.registry import ExtensionManager

27 from .audio import AudioEmbeddingMediaIO, AudioMediaIO

@@ -23,6 +29,8 @@ from .base import MediaIO

29 from .image import ImageEmbeddingMediaIO, ImageMediaIO

30 from .video import VideoMediaIO

32 +logger = init_logger(__name__)

33 +

34 _M = TypeVar("_M")

36 global_thread_pool = ThreadPoolExecutor(

@@ -116,6 +124,115 @@ class MediaConnector:

124 allowed_media_domains = []

125 self.allowed_media_domains = allowed_media_domains

126

127 + # Media download cache (opt-in via VLLM_MEDIA_CACHE)

128 + self._media_cache_dir: str | None = None

129 + self._media_cache_max_bytes: int = 0

130 + self._media_cache_ttl_secs: float = 0

131 + media_cache = envs.VLLM_MEDIA_CACHE

132 + if media_cache:

133 + try:

134 + os.makedirs(media_cache, exist_ok=True)

135 + # Verify the directory is writable before enabling caching

136 + with tempfile.NamedTemporaryFile(dir=media_cache, delete=True):

137 + pass

138 + self._media_cache_dir = media_cache

139 + self._media_cache_max_bytes = (

140 + envs.VLLM_MEDIA_CACHE_MAX_SIZE_MB * 1024 * 1024

141 + )

142 + self._media_cache_ttl_secs = envs.VLLM_MEDIA_CACHE_TTL_HOURS * 3600

143 + logger.info(

144 + "Media cache enabled at %s (max %d MB, TTL %s hours)",

145 + media_cache,

146 + envs.VLLM_MEDIA_CACHE_MAX_SIZE_MB,

147 + envs.VLLM_MEDIA_CACHE_TTL_HOURS,

148 + )

149 + except OSError:

150 + logger.warning(

151 + "VLLM_MEDIA_CACHE path %s is not writable, media caching disabled",

152 + media_cache,

153 + )

154 +

155 + def _get_cached_bytes(self, url: str) -> bytes | None:

156 + """Return cached bytes for a URL, or None if not cached/expired."""

157 + if not self._media_cache_dir:

158 + return None

159 + cache_path = self._media_cache_path(url)

160 + # Check TTL

161 + try:

162 + age = time.time() - cache_path.stat().st_mtime

163 + except OSError:

164 + return None

165 + if age > self._media_cache_ttl_secs:

166 + cache_path.unlink(missing_ok=True)

167 + return None

168 + # Touch mtime for LRU ordering

169 + try:

170 + cache_path.touch()

171 + return cache_path.read_bytes()

172 + except OSError:

173 + return None

174 +

175 + def _put_cached_bytes(self, url: str, data: bytes) -> None:

176 + """Store downloaded bytes and evict if over budget."""

177 + if not self._media_cache_dir:

178 + return

179 + cache_path = self._media_cache_path(url)

180 + # Atomic write via temp file + rename

181 + tmp_path = None

182 + try:

183 + with tempfile.NamedTemporaryFile(

184 + mode="wb", dir=self._media_cache_dir, delete=False

185 + ) as tmp_file:

186 + tmp_file.write(data)

187 + tmp_path = tmp_file.name

188 + os.rename(tmp_path, str(cache_path))

189 + except OSError:

190 + # Another process beat us or disk issue

191 + if tmp_path is not None:

192 + with contextlib.suppress(OSError):

193 + os.remove(tmp_path)

194 + return

195 + self._maybe_evict(exclude=cache_path)

196 +

197 + def _maybe_evict(self, exclude: Path | None = None) -> None:

198 + """Evict expired entries first, then LRU until under size limit."""

199 + cache_dir = Path(self._media_cache_dir) # type: ignore[arg-type]

200 + entries = []

201 + expired = []

202 + total_size = 0

203 + now = time.time()

204 + for f in cache_dir.iterdir():

205 + if f.name.startswith("."):

206 + continue

207 + try:

208 + stat = f.stat()

209 + except OSError:

210 + continue

211 + age = now - stat.st_mtime

212 + if age > self._media_cache_ttl_secs:

213 + expired.append(f)

214 + continue

215 + total_size += stat.st_size

216 + # Never evict the file we just wrote

217 + if exclude is not None and f.name == exclude.name:

218 + continue

219 + entries.append((stat.st_mtime, stat.st_size, f))

220 +

221 + # Evict items according to LRU policy

222 + entries.sort(key=lambda e: e[0], reverse=True)

223 + while total_size > self._media_cache_max_bytes and entries:

224 + mtime, size, f = entries.pop()

225 + expired.append(f)

226 + total_size -= size

227 +

228 + for f in expired:

229 + f.unlink(missing_ok=True)

230 +

231 + def _media_cache_path(self, url: str) -> Path:

232 + url_hash = hashlib.sha256(url.encode()).hexdigest()[:20]

233 + ext = Path(url.split("?", 1)[0]).suffix or ""

234 + return Path(self._media_cache_dir) / f"{url_hash}{ext}" # type: ignore[arg-type]

235 +

236 def _load_data_url(

237 self,

238 url_spec: Url,

@@ -178,6 +295,10 @@ class MediaConnector:

295 if url_spec.scheme and url_spec.scheme.startswith("http"):

296 self._assert_url_in_allowed_media_domains(url_spec)

297

298 + cached = self._get_cached_bytes(url)

299 + if cached is not None:

300 + return media_io.load_bytes(cached)

301 +

302 connection = self.connection

303 data = connection.get_bytes(

304 url_spec.url,

@@ -185,6 +306,7 @@ class MediaConnector:

306 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,

307 )

308

309 + self._put_cached_bytes(url, data)

310 return media_io.load_bytes(data)

311

312 if url_spec.scheme == "data":

@@ -209,12 +331,25 @@ class MediaConnector:

331 if url_spec.scheme and url_spec.scheme.startswith("http"):

332 self._assert_url_in_allowed_media_domains(url_spec)

333

334 + cached = await loop.run_in_executor(

335 + global_thread_pool, self._get_cached_bytes, url

336 + )

337 + if cached is not None:

338 + future = loop.run_in_executor(

339 + global_thread_pool, media_io.load_bytes, cached

340 + )

341 + return await future

342 +

343 connection = self.connection

344 data = await connection.async_get_bytes(

345 url_spec.url,

346 timeout=fetch_timeout,

347 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,

348 )

349 +

350 + await loop.run_in_executor(

351 + global_thread_pool, self._put_cached_bytes, url, data

352 + )

353 future = loop.run_in_executor(global_thread_pool, media_io.load_bytes, data)

354 return await future

355