Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
0
third_party/vllm/tests/multimodal/media/__init__.py
vendored
Normal file
0
third_party/vllm/tests/multimodal/media/__init__.py
vendored
Normal file
84
third_party/vllm/tests/multimodal/media/test_audio.py
vendored
Normal file
84
third_party/vllm/tests/multimodal/media/test_audio.py
vendored
Normal file
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm.multimodal.media import AudioMediaIO
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
ASSETS_DIR = Path(__file__).parent.parent / "assets"
|
||||
assert ASSETS_DIR.exists()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_audio():
|
||||
return np.array([0.0, 0.1, 0.2, 0.3, 0.4], dtype=float)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_audio_bytes():
|
||||
return b"FAKEAUDIOBYTES"
|
||||
|
||||
|
||||
def test_audio_media_io_load_bytes(dummy_audio_bytes):
|
||||
audio_io = AudioMediaIO()
|
||||
with patch("librosa.load") as mock_load:
|
||||
mock_load.return_value = (np.array([0.1, 0.2]), 16000)
|
||||
out = audio_io.load_bytes(dummy_audio_bytes)
|
||||
mock_load.assert_called_once()
|
||||
assert isinstance(out[0], np.ndarray)
|
||||
assert out[1] == 16000
|
||||
|
||||
|
||||
def test_audio_media_io_load_base64(dummy_audio_bytes):
|
||||
audio_io = AudioMediaIO()
|
||||
encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8")
|
||||
with patch.object(AudioMediaIO, "load_bytes") as mock_load_bytes:
|
||||
mock_load_bytes.return_value = (np.array([0.1, 0.2]), 16000)
|
||||
out = audio_io.load_base64("audio/wav", encoded)
|
||||
mock_load_bytes.assert_called_once()
|
||||
assert isinstance(out[0], np.ndarray)
|
||||
assert out[1] == 16000
|
||||
|
||||
|
||||
def test_audio_media_io_load_file():
|
||||
audio_io = AudioMediaIO()
|
||||
path = Path("/fake/path.wav")
|
||||
with patch("librosa.load") as mock_load:
|
||||
mock_load.return_value = (np.array([0.1, 0.2]), 16000)
|
||||
out = audio_io.load_file(path)
|
||||
mock_load.assert_called_once_with(path, sr=None)
|
||||
assert isinstance(out[0], np.ndarray)
|
||||
assert out[1] == 16000
|
||||
|
||||
|
||||
def test_audio_media_io_encode_base64(dummy_audio):
|
||||
audio_io = AudioMediaIO()
|
||||
media = (dummy_audio, 16000)
|
||||
with patch("soundfile.write") as mock_write:
|
||||
|
||||
def write_to_buffer(buffer, *_args, **_kwargs):
|
||||
buffer.write(b"dummy_wav_data")
|
||||
|
||||
mock_write.side_effect = write_to_buffer
|
||||
|
||||
out = audio_io.encode_base64(media)
|
||||
decoded = base64.b64decode(out)
|
||||
assert decoded == b"dummy_wav_data"
|
||||
mock_write.assert_called_once()
|
||||
|
||||
|
||||
def test_audio_media_io_from_video(video_assets):
|
||||
audio_io = AudioMediaIO()
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
audio, sr = audio_io.load_bytes(f.read())
|
||||
audio_ref, sr_ref = librosa.load(video_path, sr=None)
|
||||
assert sr == sr_ref
|
||||
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
|
||||
45
third_party/vllm/tests/multimodal/media/test_base.py
vendored
Normal file
45
third_party/vllm/tests/multimodal/media/test_base.py
vendored
Normal file
@@ -0,0 +1,45 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm.multimodal.media import MediaWithBytes
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
ASSETS_DIR = Path(__file__).parent.parent / "assets"
|
||||
assert ASSETS_DIR.exists()
|
||||
|
||||
|
||||
def test_media_with_bytes_pickle_roundtrip():
|
||||
"""Regression test for pickle/unpickle of MediaWithBytes.
|
||||
|
||||
Verifies that MediaWithBytes can be pickled and unpickled without
|
||||
RecursionError. See: https://github.com/vllm-project/vllm/issues/30818
|
||||
"""
|
||||
original_image = Image.open(ASSETS_DIR / "image1.png").convert("RGB")
|
||||
original_bytes = b"test_bytes_data"
|
||||
|
||||
wrapper = MediaWithBytes(media=original_image, original_bytes=original_bytes)
|
||||
|
||||
# Verify attribute delegation works before pickling
|
||||
assert wrapper.width == original_image.width
|
||||
assert wrapper.height == original_image.height
|
||||
assert wrapper.mode == original_image.mode
|
||||
|
||||
# Pickle and unpickle (this would cause RecursionError before the fix)
|
||||
pickled = pickle.dumps(wrapper)
|
||||
unpickled = pickle.loads(pickled)
|
||||
|
||||
# Verify the unpickled object works correctly
|
||||
assert unpickled.original_bytes == original_bytes
|
||||
assert unpickled.media.width == original_image.width
|
||||
assert unpickled.media.height == original_image.height
|
||||
|
||||
# Verify attribute delegation works after unpickling
|
||||
assert unpickled.width == original_image.width
|
||||
assert unpickled.height == original_image.height
|
||||
assert unpickled.mode == original_image.mode
|
||||
377
third_party/vllm/tests/multimodal/media/test_connector.py
vendored
Normal file
377
third_party/vllm/tests/multimodal/media/test_connector.py
vendored
Normal file
@@ -0,0 +1,377 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image, ImageChops
|
||||
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.inputs import PlaceholderRange
|
||||
from vllm.multimodal.media import MediaConnector
|
||||
|
||||
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
|
||||
TEST_IMAGE_ASSETS = [
|
||||
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
|
||||
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
|
||||
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
|
||||
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
|
||||
]
|
||||
|
||||
TEST_VIDEO_URLS = [
|
||||
"https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4",
|
||||
"https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def url_images(local_asset_server) -> dict[str, Image.Image]:
|
||||
return {
|
||||
image_url: local_asset_server.get_image_asset(image_url)
|
||||
for image_url in TEST_IMAGE_ASSETS
|
||||
}
|
||||
|
||||
|
||||
def get_supported_suffixes() -> tuple[str, ...]:
|
||||
# We should at least test the file types mentioned in GPT-4 with Vision
|
||||
OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif")
|
||||
|
||||
# Additional file types that are supported by us
|
||||
EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff")
|
||||
|
||||
return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
|
||||
|
||||
|
||||
def _image_equals(a: Image.Image, b: Image.Image) -> bool:
|
||||
return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_fetch_image_http(image_url: str):
|
||||
connector = MediaConnector()
|
||||
|
||||
image_sync = connector.fetch_image(image_url)
|
||||
image_async = await connector.fetch_image_async(image_url)
|
||||
assert _image_equals(image_sync, image_async)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
|
||||
@pytest.mark.parametrize("suffix", get_supported_suffixes())
|
||||
async def test_fetch_image_base64(
|
||||
url_images: dict[str, Image.Image], raw_image_url: str, suffix: str
|
||||
):
|
||||
connector = MediaConnector(
|
||||
# Domain restriction should not apply to data URLs.
|
||||
allowed_media_domains=[
|
||||
"www.bogotobogo.com",
|
||||
"github.com",
|
||||
]
|
||||
)
|
||||
url_image = url_images[raw_image_url]
|
||||
|
||||
try:
|
||||
mime_type = Image.MIME[Image.registered_extensions()[suffix]]
|
||||
except KeyError:
|
||||
try:
|
||||
mime_type = mimetypes.types_map[suffix]
|
||||
except KeyError:
|
||||
pytest.skip("No MIME type")
|
||||
|
||||
with NamedTemporaryFile(suffix=suffix) as f:
|
||||
try:
|
||||
url_image.save(f.name)
|
||||
except Exception as e:
|
||||
if e.args[0] == "cannot write mode RGBA as JPEG":
|
||||
pytest.skip("Conversion not supported")
|
||||
|
||||
raise
|
||||
|
||||
base64_image = base64.b64encode(f.read()).decode("utf-8")
|
||||
data_url = f"data:{mime_type};base64,{base64_image}"
|
||||
|
||||
data_image_sync = connector.fetch_image(data_url)
|
||||
if _image_equals(url_image, Image.open(f)):
|
||||
assert _image_equals(url_image, data_image_sync)
|
||||
else:
|
||||
pass # Lossy format; only check that image can be opened
|
||||
|
||||
data_image_async = await connector.fetch_image_async(data_url)
|
||||
assert _image_equals(data_image_sync, data_image_async)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_fetch_image_local_files(image_url: str):
|
||||
connector = MediaConnector()
|
||||
|
||||
with TemporaryDirectory() as temp_dir:
|
||||
local_connector = MediaConnector(allowed_local_media_path=temp_dir)
|
||||
|
||||
origin_image = connector.fetch_image(image_url)
|
||||
origin_image.save(
|
||||
os.path.join(temp_dir, os.path.basename(image_url)),
|
||||
quality=100,
|
||||
icc_profile=origin_image.info.get("icc_profile"),
|
||||
)
|
||||
|
||||
image_async = await local_connector.fetch_image_async(
|
||||
f"file://{temp_dir}/{os.path.basename(image_url)}"
|
||||
)
|
||||
image_sync = local_connector.fetch_image(
|
||||
f"file://{temp_dir}/{os.path.basename(image_url)}"
|
||||
)
|
||||
# Check that the images are equal
|
||||
assert not ImageChops.difference(image_sync, image_async).getbbox()
|
||||
|
||||
with pytest.raises(ValueError, match="must be a subpath"):
|
||||
await local_connector.fetch_image_async(
|
||||
f"file://{temp_dir}/../{os.path.basename(image_url)}"
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="Cannot load local files"):
|
||||
await connector.fetch_image_async(
|
||||
f"file://{temp_dir}/../{os.path.basename(image_url)}"
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="must be a subpath"):
|
||||
local_connector.fetch_image(
|
||||
f"file://{temp_dir}/../{os.path.basename(image_url)}"
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="Cannot load local files"):
|
||||
connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True)
|
||||
async def test_fetch_image_local_files_with_space_in_name(image_url: str):
|
||||
connector = MediaConnector()
|
||||
|
||||
with TemporaryDirectory() as temp_dir:
|
||||
local_connector = MediaConnector(allowed_local_media_path=temp_dir)
|
||||
|
||||
origin_image = connector.fetch_image(image_url)
|
||||
filename = "file name with space.jpg"
|
||||
origin_image.save(
|
||||
os.path.join(temp_dir, filename),
|
||||
quality=100,
|
||||
icc_profile=origin_image.info.get("icc_profile"),
|
||||
)
|
||||
|
||||
try:
|
||||
image_async = await local_connector.fetch_image_async(
|
||||
f"file://{temp_dir}/{filename}"
|
||||
)
|
||||
image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}")
|
||||
except FileNotFoundError as e:
|
||||
pytest.fail("Failed to fetch image with space in name: {}".format(e))
|
||||
# Check that the images are equal
|
||||
assert not ImageChops.difference(image_sync, image_async).getbbox()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_image_error_conversion():
|
||||
connector = MediaConnector()
|
||||
broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK"
|
||||
|
||||
# PIL.UnidentifiedImageError should be converted to ValueError
|
||||
with pytest.raises(ValueError):
|
||||
await connector.fetch_image_async(broken_img)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
connector.fetch_image(broken_img)
|
||||
|
||||
|
||||
@pytest.mark.flaky(reruns=3, reruns_delay=5)
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
|
||||
@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
|
||||
async def test_fetch_video_http(video_url: str, num_frames: int):
|
||||
connector = MediaConnector(
|
||||
media_io_kwargs={
|
||||
"video": {
|
||||
"num_frames": num_frames,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
video_sync, metadata_sync = connector.fetch_video(video_url)
|
||||
video_async, metadata_async = await connector.fetch_video_async(video_url)
|
||||
except (TimeoutError, asyncio.TimeoutError) as e:
|
||||
pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
|
||||
|
||||
assert np.array_equal(video_sync, video_async)
|
||||
assert metadata_sync == metadata_async
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
|
||||
@pytest.mark.parametrize("max_duration", [1, 60, 1800])
|
||||
@pytest.mark.parametrize("requested_fps", [2, 24])
|
||||
async def test_fetch_video_http_with_dynamic_loader(
|
||||
video_url: str,
|
||||
max_duration: int,
|
||||
requested_fps: int,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic")
|
||||
connector = MediaConnector(
|
||||
media_io_kwargs={
|
||||
"video": {
|
||||
"max_duration": max_duration,
|
||||
"requested_fps": requested_fps,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
video_sync, metadata_sync = connector.fetch_video(video_url)
|
||||
video_async, metadata_async = await connector.fetch_video_async(video_url)
|
||||
|
||||
assert np.array_equal(video_sync, video_async)
|
||||
assert metadata_sync == metadata_async
|
||||
assert metadata_sync["video_backend"] == "opencv_dynamic"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"is_embed,start_idx,end_idx,expected",
|
||||
[
|
||||
(None, 2, 4, (2, 4)),
|
||||
(
|
||||
torch.tensor([False, True, False, True, True]),
|
||||
3,
|
||||
5,
|
||||
(1, 3),
|
||||
),
|
||||
(
|
||||
torch.tensor([False, True, False, True, True]),
|
||||
0,
|
||||
2,
|
||||
(0, 1),
|
||||
),
|
||||
(
|
||||
torch.tensor([True, False, True, False]),
|
||||
2,
|
||||
2,
|
||||
(1, 1),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_placeholder_range_get_embeds_indices_in_range(
|
||||
is_embed, start_idx, end_idx, expected
|
||||
):
|
||||
length = len(is_embed) if is_embed is not None else 5
|
||||
pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
|
||||
assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"offset,is_embed,expected",
|
||||
[
|
||||
(0, None, [(0, 4)]),
|
||||
(
|
||||
2,
|
||||
torch.tensor([False, True, False, True, True]),
|
||||
[(3, 3), (5, 6)],
|
||||
),
|
||||
(0, torch.tensor([True, True, True, True]), [(0, 3)]),
|
||||
(0, torch.tensor([False, False, False, False]), []),
|
||||
],
|
||||
)
|
||||
def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
|
||||
length = len(is_embed) if is_embed is not None else 5
|
||||
pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
|
||||
assert pr.extract_embeds_range() == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
|
||||
@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
|
||||
async def test_allowed_media_domains(video_url: str, num_frames: int):
|
||||
connector = MediaConnector(
|
||||
media_io_kwargs={
|
||||
"video": {
|
||||
"num_frames": num_frames,
|
||||
}
|
||||
},
|
||||
allowed_media_domains=[
|
||||
"www.bogotobogo.com",
|
||||
"github.com",
|
||||
],
|
||||
)
|
||||
|
||||
video_sync, metadata_sync = connector.fetch_video(video_url)
|
||||
video_async, metadata_async = await connector.fetch_video_async(video_url)
|
||||
assert np.array_equal(video_sync, video_async)
|
||||
assert metadata_sync == metadata_async
|
||||
|
||||
disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png"
|
||||
with pytest.raises(ValueError):
|
||||
_, _ = connector.fetch_video(disallowed_url)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_, _ = await connector.fetch_video_async(disallowed_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_bypass_backslash_in_url(local_asset_server):
|
||||
"""Verify that backslash-@ URL parsing confusion cannot bypass the
|
||||
allowed_media_domains check (GHSA-v359-jj2v-j536).
|
||||
|
||||
urllib3.parse_url() and aiohttp/yarl disagree on how to parse a
|
||||
backslash before ``@``. urllib3 treats ``\\`` as part of the path
|
||||
(encoding it as ``%5C``), while yarl treats it as a userinfo
|
||||
separator, changing the effective host. The fix normalises the URL
|
||||
through urllib3 *before* handing it to aiohttp so both layers agree.
|
||||
"""
|
||||
port = local_asset_server.port
|
||||
asset = TEST_IMAGE_ASSETS[0]
|
||||
|
||||
# Craft the bypass payload: urllib3 sees host=127.0.0.1, but an
|
||||
# un-patched aiohttp would see host=example.com.
|
||||
bypass_url = f"http://127.0.0.1:{port}\\@example.com/{asset}"
|
||||
|
||||
connector = MediaConnector(
|
||||
allowed_media_domains=["127.0.0.1"],
|
||||
)
|
||||
|
||||
# After the fix the request is made to 127.0.0.1 (the local asset
|
||||
# server) using the normalised URL. The normalised path will be
|
||||
# /%5C@example.com/<asset> which won't match any file the server
|
||||
# knows about, so we expect an HTTP error — but crucially NOT a
|
||||
# successful fetch from example.com.
|
||||
with pytest.raises(requests.exceptions.HTTPError):
|
||||
connector.fetch_image(bypass_url)
|
||||
|
||||
with pytest.raises(aiohttp.ClientResponseError):
|
||||
await connector.fetch_image_async(bypass_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ssrf_bypass_backslash_disallowed_domain():
|
||||
"""The reverse direction: even when the *attacker-controlled* host
|
||||
appears in the urllib3-parsed hostname position the allowlist must
|
||||
still block it.
|
||||
"""
|
||||
# urllib3.parse_url sees host=example.com which is NOT in the
|
||||
# allowlist, so this must be rejected before any request is made.
|
||||
bypass_url = "https://example.com\\@safe.example.org/image.png"
|
||||
|
||||
connector = MediaConnector(
|
||||
allowed_media_domains=["safe.example.org"],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="allowed domains"):
|
||||
connector.fetch_image(bypass_url)
|
||||
|
||||
with pytest.raises(ValueError, match="allowed domains"):
|
||||
await connector.fetch_image_async(bypass_url)
|
||||
133
third_party/vllm/tests/multimodal/media/test_image.py
vendored
Normal file
133
third_party/vllm/tests/multimodal/media/test_image.py
vendored
Normal file
@@ -0,0 +1,133 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm.multimodal.media import ImageMediaIO
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
ASSETS_DIR = Path(__file__).parent.parent / "assets"
|
||||
assert ASSETS_DIR.exists()
|
||||
|
||||
|
||||
def test_image_media_io_rgba_custom_background(tmp_path):
|
||||
"""Test RGBA to RGB conversion with custom background colors."""
|
||||
# Create a simple RGBA image with transparent and opaque pixels
|
||||
rgba_image = Image.new("RGBA", (10, 10), (255, 0, 0, 255)) # Red with full opacity
|
||||
|
||||
# Make top-left quadrant transparent
|
||||
for i in range(5):
|
||||
for j in range(5):
|
||||
rgba_image.putpixel((i, j), (0, 0, 0, 0)) # Fully transparent
|
||||
|
||||
# Save the test image to tmp_path
|
||||
test_image_path = tmp_path / "test_rgba.png"
|
||||
rgba_image.save(test_image_path)
|
||||
|
||||
# Test 1: Default white background (backward compatibility)
|
||||
image_io_default = ImageMediaIO()
|
||||
converted_default = image_io_default.load_file(test_image_path)
|
||||
default_numpy = np.array(converted_default)
|
||||
|
||||
# Check transparent pixels are white
|
||||
assert default_numpy[0][0][0] == 255 # R
|
||||
assert default_numpy[0][0][1] == 255 # G
|
||||
assert default_numpy[0][0][2] == 255 # B
|
||||
# Check opaque pixels remain red
|
||||
assert default_numpy[5][5][0] == 255 # R
|
||||
assert default_numpy[5][5][1] == 0 # G
|
||||
assert default_numpy[5][5][2] == 0 # B
|
||||
|
||||
# Test 2: Custom black background via kwargs
|
||||
image_io_black = ImageMediaIO(rgba_background_color=(0, 0, 0))
|
||||
converted_black = image_io_black.load_file(test_image_path)
|
||||
black_numpy = np.array(converted_black)
|
||||
|
||||
# Check transparent pixels are black
|
||||
assert black_numpy[0][0][0] == 0 # R
|
||||
assert black_numpy[0][0][1] == 0 # G
|
||||
assert black_numpy[0][0][2] == 0 # B
|
||||
# Check opaque pixels remain red
|
||||
assert black_numpy[5][5][0] == 255 # R
|
||||
assert black_numpy[5][5][1] == 0 # G
|
||||
assert black_numpy[5][5][2] == 0 # B
|
||||
|
||||
# Test 3: Custom blue background via kwargs (as list)
|
||||
image_io_blue = ImageMediaIO(rgba_background_color=[0, 0, 255])
|
||||
converted_blue = image_io_blue.load_file(test_image_path)
|
||||
blue_numpy = np.array(converted_blue)
|
||||
|
||||
# Check transparent pixels are blue
|
||||
assert blue_numpy[0][0][0] == 0 # R
|
||||
assert blue_numpy[0][0][1] == 0 # G
|
||||
assert blue_numpy[0][0][2] == 255 # B
|
||||
|
||||
# Test 4: Test with load_bytes method
|
||||
with open(test_image_path, "rb") as f:
|
||||
image_data = f.read()
|
||||
|
||||
image_io_green = ImageMediaIO(rgba_background_color=(0, 255, 0))
|
||||
converted_green = image_io_green.load_bytes(image_data)
|
||||
green_numpy = np.array(converted_green)
|
||||
|
||||
# Check transparent pixels are green
|
||||
assert green_numpy[0][0][0] == 0 # R
|
||||
assert green_numpy[0][0][1] == 255 # G
|
||||
assert green_numpy[0][0][2] == 0 # B
|
||||
|
||||
|
||||
def test_image_media_io_rgba_background_color_validation():
|
||||
"""Test that invalid rgba_background_color values are properly rejected."""
|
||||
|
||||
# Test invalid types
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color="255,255,255")
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=255)
|
||||
|
||||
# Test wrong number of elements
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(255, 255))
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(255, 255, 255, 255))
|
||||
|
||||
# Test non-integer values
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(255.0, 255.0, 255.0))
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(255, "255", 255))
|
||||
|
||||
# Test out of range values
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(256, 255, 255))
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="rgba_background_color must be a list or tuple"
|
||||
):
|
||||
ImageMediaIO(rgba_background_color=(255, -1, 255))
|
||||
|
||||
# Test that valid values work
|
||||
ImageMediaIO(rgba_background_color=(0, 0, 0)) # Should not raise
|
||||
ImageMediaIO(rgba_background_color=[255, 255, 255]) # Should not raise
|
||||
ImageMediaIO(rgba_background_color=(128, 128, 128)) # Should not raise
|
||||
237
third_party/vllm/tests/multimodal/media/test_video.py
vendored
Normal file
237
third_party/vllm/tests/multimodal/media/test_video.py
vendored
Normal file
@@ -0,0 +1,237 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm.assets.base import get_vllm_public_assets
|
||||
from vllm.assets.video import video_to_ndarrays, video_to_pil_images_list
|
||||
from vllm.multimodal.media import ImageMediaIO, VideoMediaIO
|
||||
from vllm.multimodal.video import VIDEO_LOADER_REGISTRY, VideoLoader
|
||||
|
||||
from ..utils import cosine_similarity, create_video_from_image, normalize_image
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
ASSETS_DIR = Path(__file__).parent.parent / "assets"
|
||||
assert ASSETS_DIR.exists()
|
||||
|
||||
|
||||
@VIDEO_LOADER_REGISTRY.register("assert_10_frames_1_fps")
|
||||
class Assert10Frames1FPSVideoLoader(VideoLoader):
|
||||
@classmethod
|
||||
def load_bytes(
|
||||
cls, data: bytes, num_frames: int = -1, fps: float = -1.0, **kwargs
|
||||
) -> npt.NDArray:
|
||||
assert num_frames == 10, "bad num_frames"
|
||||
assert fps == 1.0, "bad fps"
|
||||
return FAKE_OUTPUT_2
|
||||
|
||||
|
||||
def test_video_media_io_kwargs(monkeypatch: pytest.MonkeyPatch):
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_VIDEO_LOADER_BACKEND", "assert_10_frames_1_fps")
|
||||
imageio = ImageMediaIO()
|
||||
|
||||
# Verify that different args pass/fail assertions as expected.
|
||||
videoio = VideoMediaIO(imageio, **{"num_frames": 10, "fps": 1.0})
|
||||
_ = videoio.load_bytes(b"test")
|
||||
|
||||
videoio = VideoMediaIO(
|
||||
imageio, **{"num_frames": 10, "fps": 1.0, "not_used": "not_used"}
|
||||
)
|
||||
_ = videoio.load_bytes(b"test")
|
||||
|
||||
with pytest.raises(AssertionError, match="bad num_frames"):
|
||||
videoio = VideoMediaIO(imageio, **{})
|
||||
_ = videoio.load_bytes(b"test")
|
||||
|
||||
with pytest.raises(AssertionError, match="bad num_frames"):
|
||||
videoio = VideoMediaIO(imageio, **{"num_frames": 9, "fps": 1.0})
|
||||
_ = videoio.load_bytes(b"test")
|
||||
|
||||
with pytest.raises(AssertionError, match="bad fps"):
|
||||
videoio = VideoMediaIO(imageio, **{"num_frames": 10, "fps": 2.0})
|
||||
_ = videoio.load_bytes(b"test")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_color", [True, False])
|
||||
@pytest.mark.parametrize("fourcc, ext", [("mp4v", "mp4"), ("XVID", "avi")])
|
||||
def test_opencv_video_io_colorspace(tmp_path, is_color: bool, fourcc: str, ext: str):
|
||||
"""
|
||||
Test all functions that use OpenCV for video I/O return RGB format.
|
||||
Both RGB and grayscale videos are tested.
|
||||
"""
|
||||
image_path = get_vllm_public_assets(
|
||||
filename="stop_sign.jpg", s3_prefix="vision_model_images"
|
||||
)
|
||||
image = Image.open(image_path)
|
||||
|
||||
if not is_color:
|
||||
image_path = f"{tmp_path}/test_grayscale_image.png"
|
||||
image = image.convert("L")
|
||||
image.save(image_path)
|
||||
# Convert to gray RGB for comparison
|
||||
image = image.convert("RGB")
|
||||
video_path = f"{tmp_path}/test_RGB_video.{ext}"
|
||||
create_video_from_image(
|
||||
image_path,
|
||||
video_path,
|
||||
num_frames=2,
|
||||
is_color=is_color,
|
||||
fourcc=fourcc,
|
||||
)
|
||||
|
||||
frames = video_to_ndarrays(video_path)
|
||||
for frame in frames:
|
||||
sim = cosine_similarity(
|
||||
normalize_image(np.array(frame)), normalize_image(np.array(image))
|
||||
)
|
||||
assert np.sum(np.isnan(sim)) / sim.size < 0.001
|
||||
assert np.nanmean(sim) > 0.99
|
||||
|
||||
pil_frames = video_to_pil_images_list(video_path)
|
||||
for frame in pil_frames:
|
||||
sim = cosine_similarity(
|
||||
normalize_image(np.array(frame)), normalize_image(np.array(image))
|
||||
)
|
||||
assert np.sum(np.isnan(sim)) / sim.size < 0.001
|
||||
assert np.nanmean(sim) > 0.99
|
||||
|
||||
io_frames, _ = VideoMediaIO(ImageMediaIO()).load_file(Path(video_path))
|
||||
for frame in io_frames:
|
||||
sim = cosine_similarity(
|
||||
normalize_image(np.array(frame)), normalize_image(np.array(image))
|
||||
)
|
||||
assert np.sum(np.isnan(sim)) / sim.size < 0.001
|
||||
assert np.nanmean(sim) > 0.99
|
||||
|
||||
|
||||
NUM_FRAMES = 10
|
||||
FAKE_OUTPUT_1 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
|
||||
FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3)
|
||||
|
||||
|
||||
@VIDEO_LOADER_REGISTRY.register("test_video_backend_override_1")
|
||||
class TestVideoBackendOverride1(VideoLoader):
|
||||
"""Test loader that returns FAKE_OUTPUT_1 to verify backend selection."""
|
||||
|
||||
@classmethod
|
||||
def load_bytes(
|
||||
cls, data: bytes, num_frames: int = -1, **kwargs
|
||||
) -> tuple[npt.NDArray, dict]:
|
||||
return FAKE_OUTPUT_1, {"video_backend": "test_video_backend_override_1"}
|
||||
|
||||
|
||||
@VIDEO_LOADER_REGISTRY.register("test_video_backend_override_2")
|
||||
class TestVideoBackendOverride2(VideoLoader):
|
||||
"""Test loader that returns FAKE_OUTPUT_2 to verify backend selection."""
|
||||
|
||||
@classmethod
|
||||
def load_bytes(
|
||||
cls, data: bytes, num_frames: int = -1, **kwargs
|
||||
) -> tuple[npt.NDArray, dict]:
|
||||
return FAKE_OUTPUT_2, {"video_backend": "test_video_backend_override_2"}
|
||||
|
||||
|
||||
def test_video_media_io_backend_kwarg_override(monkeypatch: pytest.MonkeyPatch):
|
||||
"""
|
||||
Test that video_backend kwarg can override the VLLM_VIDEO_LOADER_BACKEND
|
||||
environment variable.
|
||||
|
||||
This allows users to dynamically select a different video backend
|
||||
via --media-io-kwargs without changing the global env var, which is
|
||||
useful when plugins set a default backend but a specific request
|
||||
needs a different one.
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
# Set the env var to one backend
|
||||
m.setenv("VLLM_VIDEO_LOADER_BACKEND", "test_video_backend_override_1")
|
||||
|
||||
imageio = ImageMediaIO()
|
||||
|
||||
# Without video_backend kwarg, should use env var backend
|
||||
videoio_default = VideoMediaIO(imageio, num_frames=10)
|
||||
frames_default, metadata_default = videoio_default.load_bytes(b"test")
|
||||
np.testing.assert_array_equal(frames_default, FAKE_OUTPUT_1)
|
||||
assert metadata_default["video_backend"] == "test_video_backend_override_1"
|
||||
|
||||
# With video_backend kwarg, should override env var
|
||||
videoio_override = VideoMediaIO(
|
||||
imageio, num_frames=10, video_backend="test_video_backend_override_2"
|
||||
)
|
||||
frames_override, metadata_override = videoio_override.load_bytes(b"test")
|
||||
np.testing.assert_array_equal(frames_override, FAKE_OUTPUT_2)
|
||||
assert metadata_override["video_backend"] == "test_video_backend_override_2"
|
||||
|
||||
|
||||
def test_video_media_io_backend_kwarg_not_passed_to_loader(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""
|
||||
Test that video_backend kwarg is consumed by VideoMediaIO and NOT passed
|
||||
through to the underlying video loader's load_bytes method.
|
||||
|
||||
This ensures the kwarg is properly popped from kwargs before forwarding.
|
||||
"""
|
||||
|
||||
@VIDEO_LOADER_REGISTRY.register("test_reject_video_backend_kwarg")
|
||||
class RejectVideoBackendKwargLoader(VideoLoader):
|
||||
"""Test loader that fails if video_backend is passed through."""
|
||||
|
||||
@classmethod
|
||||
def load_bytes(
|
||||
cls, data: bytes, num_frames: int = -1, **kwargs
|
||||
) -> tuple[npt.NDArray, dict]:
|
||||
# This should never receive video_backend in kwargs
|
||||
if "video_backend" in kwargs:
|
||||
raise AssertionError(
|
||||
"video_backend should be consumed by VideoMediaIO, "
|
||||
"not passed to loader"
|
||||
)
|
||||
return FAKE_OUTPUT_1, {"received_kwargs": list(kwargs.keys())}
|
||||
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_VIDEO_LOADER_BACKEND", "test_reject_video_backend_kwarg")
|
||||
|
||||
imageio = ImageMediaIO()
|
||||
|
||||
# Even when video_backend is provided, it should NOT be passed to loader
|
||||
videoio = VideoMediaIO(
|
||||
imageio,
|
||||
num_frames=10,
|
||||
video_backend="test_reject_video_backend_kwarg",
|
||||
other_kwarg="should_pass_through",
|
||||
)
|
||||
|
||||
# This should NOT raise AssertionError
|
||||
frames, metadata = videoio.load_bytes(b"test")
|
||||
np.testing.assert_array_equal(frames, FAKE_OUTPUT_1)
|
||||
# Verify other kwargs are still passed through
|
||||
assert "other_kwarg" in metadata["received_kwargs"]
|
||||
|
||||
|
||||
def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch):
|
||||
"""
|
||||
Test that when video_backend kwarg is None or not provided,
|
||||
VideoMediaIO falls back to VLLM_VIDEO_LOADER_BACKEND env var.
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_VIDEO_LOADER_BACKEND", "test_video_backend_override_2")
|
||||
|
||||
imageio = ImageMediaIO()
|
||||
|
||||
# Explicit None should fall back to env var
|
||||
videoio_none = VideoMediaIO(imageio, num_frames=10, video_backend=None)
|
||||
frames_none, metadata_none = videoio_none.load_bytes(b"test")
|
||||
np.testing.assert_array_equal(frames_none, FAKE_OUTPUT_2)
|
||||
assert metadata_none["video_backend"] == "test_video_backend_override_2"
|
||||
|
||||
# Not providing video_backend should also fall back to env var
|
||||
videoio_missing = VideoMediaIO(imageio, num_frames=10)
|
||||
frames_missing, metadata_missing = videoio_missing.load_bytes(b"test")
|
||||
np.testing.assert_array_equal(frames_missing, FAKE_OUTPUT_2)
|
||||
assert metadata_missing["video_backend"] == "test_video_backend_override_2"
|
||||
Reference in New Issue
Block a user