[VLM] add content-addressed preprocessing cache infrastructure (#34398)
This commit is contained in:
@@ -19,6 +19,7 @@ from typing import List, Optional
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import (
|
||||
ChatCompletionMessageContentImageURL,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionResponse,
|
||||
ChatCompletionResponseChoice,
|
||||
@@ -125,6 +126,17 @@ class TestChatCompletionRequest(unittest.TestCase):
|
||||
self.assertFalse(request.stream) # default
|
||||
self.assertEqual(request.tool_choice, "none") # default when no tools
|
||||
|
||||
def test_image_content_hash_validation(self):
|
||||
digest = "sha256:" + "AB" * 32
|
||||
image = ChatCompletionMessageContentImageURL(
|
||||
url="https://example.com/image.jpg", content_hash=digest
|
||||
)
|
||||
self.assertEqual(image.content_hash, digest.lower())
|
||||
with self.assertRaises(ValidationError):
|
||||
ChatCompletionMessageContentImageURL(
|
||||
url="https://example.com/image.jpg", content_hash="not-a-hash"
|
||||
)
|
||||
|
||||
def test_sampling_param_build(self):
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
|
||||
@@ -43,6 +43,32 @@ class TestMmHashesContract(CustomTestCase):
|
||||
req = GenerateReqInput(text="hi")
|
||||
self.assertIsNone(req.mm_hashes)
|
||||
|
||||
def test_content_hashes_are_distinct_from_feature_hashes(self):
|
||||
content_hash = "sha256:" + "ab" * 32
|
||||
req = GenerateReqInput(
|
||||
text="hi",
|
||||
image_data=["http://example.com/img.png"],
|
||||
mm_hashes=["deadbeef"],
|
||||
mm_content_hashes=[content_hash],
|
||||
)
|
||||
self.assertEqual(req.mm_hashes, ["deadbeef"])
|
||||
self.assertEqual(req.mm_content_hashes, [content_hash])
|
||||
|
||||
def test_batched_hashes_follow_each_request(self):
|
||||
req = GenerateReqInput(
|
||||
text=["one", "two"],
|
||||
image_data=[["a"], ["b", "c"]],
|
||||
mm_hashes=["01", ["02", "03"]],
|
||||
mm_content_hashes=[
|
||||
["sha256:" + "11" * 32],
|
||||
["sha256:" + "22" * 32, "sha256:" + "33" * 32],
|
||||
],
|
||||
)
|
||||
req.normalize_batch_and_arguments()
|
||||
self.assertEqual(req[0].mm_hashes, ["01"])
|
||||
self.assertEqual(req[1].mm_hashes, ["02", "03"])
|
||||
self.assertEqual(len(req[1].mm_content_hashes), 2)
|
||||
|
||||
def test_set_pad_value_honors_preset_hash(self):
|
||||
"""set_pad_value() must use a pre-set hash without recomputing."""
|
||||
item = MultimodalDataItem(modality=Modality.IMAGE, hash=0xDEADBEEF)
|
||||
|
||||
@@ -84,6 +84,9 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
|
||||
server_args.mm_process_config = mm_process_config
|
||||
server_args.mm_processor_worker_num = mm_processor_worker_num
|
||||
server_args.mm_io_worker_num = mm_io_worker_num
|
||||
server_args.mm_preprocess_cache_size_mb = None
|
||||
server_args.tokenizer_worker_num = 1
|
||||
server_args.trust_mm_content_hashes = False
|
||||
|
||||
hf_config = MagicMock()
|
||||
mock_hf_processor = MagicMock()
|
||||
@@ -767,6 +770,9 @@ class TestDoubleBosGuard(CustomTestCase):
|
||||
server_args.mm_io_worker_num = 0
|
||||
server_args.mm_feature_transport = "cpu"
|
||||
server_args.disable_fast_image_processor = True
|
||||
server_args.mm_preprocess_cache_size_mb = None
|
||||
server_args.tokenizer_worker_num = 1
|
||||
server_args.trust_mm_content_hashes = False
|
||||
|
||||
mock_hf_processor = MagicMock()
|
||||
mock_hf_processor.__class__.__name__ = "TestProcessor"
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from sglang.srt.multimodal.cache import (
|
||||
MultimodalPreprocessCache,
|
||||
build_artifact_key,
|
||||
build_feature_hash,
|
||||
build_processor_fingerprint,
|
||||
estimate_cache_size_bytes,
|
||||
parse_content_hash,
|
||||
snapshot_media,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestMediaIdentity(unittest.TestCase):
|
||||
def test_hash_format_is_strict_and_normalized(self):
|
||||
digest = "AB" * 32
|
||||
self.assertEqual(
|
||||
parse_content_hash(f"sha256:{digest}"), f"sha256:{digest.lower()}"
|
||||
)
|
||||
for invalid in (
|
||||
"",
|
||||
digest,
|
||||
"md5:" + digest,
|
||||
"sha256:1234",
|
||||
"sha256:" + "z" * 64,
|
||||
):
|
||||
with self.subTest(invalid=invalid), self.assertRaises(ValueError):
|
||||
parse_content_hash(invalid)
|
||||
|
||||
def test_same_bytes_have_same_identity_across_input_forms(self):
|
||||
# Keep the encoded data URL above common filesystem filename limits;
|
||||
# probing it as a local path must not raise ENAMETOOLONG.
|
||||
payload = b"strict-media-identity" * 32
|
||||
data_url = (
|
||||
"data:application/octet-stream;base64," + base64.b64encode(payload).decode()
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "image.png"
|
||||
path.write_bytes(payload)
|
||||
snapshots = [
|
||||
snapshot_media(payload),
|
||||
snapshot_media(data_url),
|
||||
snapshot_media(str(path)),
|
||||
]
|
||||
self.assertEqual(len({item.content_digest for item in snapshots}), 1)
|
||||
self.assertTrue(all(item.data == payload for item in snapshots))
|
||||
|
||||
def test_wrapped_image_input_snapshots_the_image_not_the_wrapper(self):
|
||||
image = Image.new("RGB", (2, 2), (1, 2, 3))
|
||||
direct = snapshot_media(image)
|
||||
wrapped = snapshot_media({"type": "image", "image": image})
|
||||
self.assertEqual(direct.content_digest, wrapped.content_digest)
|
||||
|
||||
def test_same_path_with_new_contents_misses(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "image.png"
|
||||
path.write_bytes(b"first")
|
||||
first = snapshot_media(str(path))
|
||||
path.write_bytes(b"second")
|
||||
second = snapshot_media(str(path))
|
||||
self.assertNotEqual(first.content_digest, second.content_digest)
|
||||
|
||||
def test_relative_local_path_uses_file_bytes(self):
|
||||
payload = b"relative-image-bytes"
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "image.png"
|
||||
path.write_bytes(payload)
|
||||
previous = Path.cwd()
|
||||
try:
|
||||
os.chdir(directory)
|
||||
snapshot = snapshot_media("image.png")
|
||||
finally:
|
||||
os.chdir(previous)
|
||||
|
||||
self.assertEqual(snapshot.data, payload)
|
||||
self.assertEqual(
|
||||
snapshot.content_digest, snapshot_media(payload).content_digest
|
||||
)
|
||||
|
||||
def test_same_url_with_new_contents_misses(self):
|
||||
with patch(
|
||||
"sglang.srt.utils.get_image_bytes", side_effect=[b"first", b"second"]
|
||||
):
|
||||
first = snapshot_media("https://example.com/image.png")
|
||||
second = snapshot_media("https://example.com/image.png")
|
||||
self.assertNotEqual(first.content_digest, second.content_digest)
|
||||
|
||||
def test_pil_and_noncontiguous_tensor_are_snapshotted(self):
|
||||
image = Image.new("RGBA", (3, 2), (1, 2, 3, 4))
|
||||
first = snapshot_media(image)
|
||||
image.putpixel((0, 0), (9, 9, 9, 9))
|
||||
self.assertNotEqual(first.content_digest, snapshot_media(image).content_digest)
|
||||
|
||||
tensor = torch.arange(24, dtype=torch.uint8).reshape(2, 3, 4).transpose(1, 2)
|
||||
tensor_snapshot = snapshot_media(tensor)
|
||||
self.assertTrue(tensor_snapshot.data.is_contiguous())
|
||||
self.assertTrue(torch.equal(tensor_snapshot.data, tensor))
|
||||
|
||||
same_bytes_new_shape = tensor.contiguous().reshape(2, 2, 6)
|
||||
self.assertNotEqual(
|
||||
tensor_snapshot.content_digest,
|
||||
snapshot_media(same_bytes_new_shape).content_digest,
|
||||
)
|
||||
self.assertNotEqual(
|
||||
snapshot_media(torch.tensor([1], dtype=torch.int32)).content_digest,
|
||||
snapshot_media(torch.tensor([1], dtype=torch.int64)).content_digest,
|
||||
)
|
||||
|
||||
def test_pil_palette_and_transparency_are_part_of_identity(self):
|
||||
first = Image.new("P", (2, 2), color=0)
|
||||
second = first.copy()
|
||||
first.putpalette([255, 0, 0] + [0, 0, 0] * 255)
|
||||
second.putpalette([0, 255, 0] + [0, 0, 0] * 255)
|
||||
self.assertNotEqual(
|
||||
snapshot_media(first).content_digest,
|
||||
snapshot_media(second).content_digest,
|
||||
)
|
||||
|
||||
second.putpalette(first.getpalette())
|
||||
first.info["transparency"] = 0
|
||||
second.info["transparency"] = 1
|
||||
self.assertNotEqual(
|
||||
snapshot_media(first).content_digest,
|
||||
snapshot_media(second).content_digest,
|
||||
)
|
||||
|
||||
def test_artifact_key_includes_processor_and_kwargs(self):
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
base = build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor-a",
|
||||
preprocess_kwargs={"antialias": True},
|
||||
)
|
||||
self.assertNotEqual(
|
||||
base,
|
||||
build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor-b",
|
||||
preprocess_kwargs={"antialias": True},
|
||||
),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
base,
|
||||
build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor-a",
|
||||
preprocess_kwargs={"antialias": False},
|
||||
),
|
||||
)
|
||||
|
||||
def test_artifact_key_canonicalization_is_type_preserving(self):
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
|
||||
def key(kwargs):
|
||||
return build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor",
|
||||
preprocess_kwargs=kwargs,
|
||||
)
|
||||
|
||||
# These pairs used to collapse to the same JSON representation. A
|
||||
# processor is allowed to distinguish them, so sharing an artifact
|
||||
# would be a correctness bug rather than a harmless cache collision.
|
||||
self.assertNotEqual(key({1: "value"}), key({"1": "value"}))
|
||||
self.assertNotEqual(key({"value": [1, 2]}), key({"value": (1, 2)}))
|
||||
self.assertNotEqual(key({"value": 1}), key({"value": True}))
|
||||
self.assertNotEqual(
|
||||
key({"value": np.array([1, 2], dtype=np.int32)}),
|
||||
key({"value": np.array([1, 3], dtype=np.int32)}),
|
||||
)
|
||||
self.assertEqual(
|
||||
key({"first": 1, "second": 2}),
|
||||
key({"second": 2, "first": 1}),
|
||||
)
|
||||
|
||||
def test_artifact_key_rejects_lossy_unknown_values(self):
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
with self.assertRaisesRegex(ValueError, "Unsupported value"):
|
||||
build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor",
|
||||
preprocess_kwargs={"value": object()},
|
||||
)
|
||||
|
||||
def test_processor_fingerprint_changes_with_output_affecting_config(self):
|
||||
class Processor:
|
||||
def __init__(self, backend):
|
||||
self.backend = backend
|
||||
|
||||
def preprocess_fingerprint_payload(self):
|
||||
return {"backend": self.backend, "antialias": True}
|
||||
|
||||
config = SimpleNamespace(model_type="vlm", architectures=["VLM"])
|
||||
args = SimpleNamespace(
|
||||
revision="model-revision",
|
||||
tokenizer_revision="tokenizer-revision",
|
||||
disable_fast_image_processor=False,
|
||||
mm_process_config={"image": {"max_pixels": 1024}},
|
||||
)
|
||||
base = build_processor_fingerprint(Processor("gpu"), config, args)
|
||||
|
||||
changed_backend = build_processor_fingerprint(Processor("cpu"), config, args)
|
||||
changed_args = SimpleNamespace(
|
||||
**{
|
||||
**vars(args),
|
||||
"mm_process_config": {"image": {"max_pixels": 2048}},
|
||||
}
|
||||
)
|
||||
changed_config = build_processor_fingerprint(
|
||||
Processor("gpu"), config, changed_args
|
||||
)
|
||||
self.assertNotEqual(base, changed_backend)
|
||||
self.assertNotEqual(base, changed_config)
|
||||
|
||||
def test_feature_hash_includes_artifact_and_processor_output(self):
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
first = build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor-a",
|
||||
)
|
||||
second = build_artifact_key(
|
||||
digest,
|
||||
modality="image",
|
||||
processor_fingerprint="processor-b",
|
||||
)
|
||||
self.assertNotEqual(build_feature_hash(first, 1), build_feature_hash(second, 1))
|
||||
self.assertNotEqual(build_feature_hash(first, 1), build_feature_hash(first, 2))
|
||||
self.assertIsInstance(build_feature_hash(first, 1 << 128), int)
|
||||
with self.assertRaises(ValueError):
|
||||
build_feature_hash(first, -1)
|
||||
|
||||
|
||||
class TestMultimodalPreprocessCache(unittest.TestCase):
|
||||
def test_byte_and_entry_bounded_lru(self):
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=6, max_entries=2)
|
||||
self.assertTrue(cache.put("a", b"aaa"))
|
||||
self.assertTrue(cache.put("b", b"bbb"))
|
||||
self.assertEqual(cache.get("a"), b"aaa")
|
||||
self.assertTrue(cache.put("c", b"ccc"))
|
||||
self.assertNotIn("b", cache)
|
||||
self.assertIn("a", cache)
|
||||
self.assertIn("c", cache)
|
||||
self.assertEqual(cache.current_size_bytes, 6)
|
||||
|
||||
def test_gpu_backed_values_are_not_implicitly_copied(self):
|
||||
if not torch.cuda.is_available():
|
||||
self.skipTest("CUDA is not available")
|
||||
value = torch.zeros(1, device="cuda")
|
||||
cache = MultimodalPreprocessCache[str, torch.Tensor](max_size_bytes=1024)
|
||||
self.assertIsNone(estimate_cache_size_bytes(value))
|
||||
self.assertFalse(cache.put("gpu", value))
|
||||
|
||||
def test_async_singleflight(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
calls = 0
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def compute():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
started.set()
|
||||
await release.wait()
|
||||
return b"artifact"
|
||||
|
||||
first = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await started.wait()
|
||||
second = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
owner, joiner = await asyncio.gather(first, second)
|
||||
self.assertEqual(calls, 1)
|
||||
self.assertFalse(owner.hit)
|
||||
self.assertTrue(joiner.joined)
|
||||
self.assertEqual(cache.get("key"), b"artifact")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_cancelled_singleflight_joiner_does_not_cancel_owner(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def compute():
|
||||
started.set()
|
||||
await release.wait()
|
||||
return b"artifact"
|
||||
|
||||
owner = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await started.wait()
|
||||
joiner = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await asyncio.sleep(0)
|
||||
joiner.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await joiner
|
||||
|
||||
release.set()
|
||||
result = await owner
|
||||
self.assertEqual(result.value, b"artifact")
|
||||
self.assertEqual(cache.get("key"), b"artifact")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_cancelled_singleflight_owner_does_not_cancel_joiner(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def compute():
|
||||
started.set()
|
||||
await release.wait()
|
||||
return b"artifact"
|
||||
|
||||
owner = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await started.wait()
|
||||
joiner = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await asyncio.sleep(0)
|
||||
owner.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await owner
|
||||
|
||||
release.set()
|
||||
result = await joiner
|
||||
self.assertEqual(result.value, b"artifact")
|
||||
self.assertTrue(result.joined)
|
||||
self.assertEqual(cache.get("key"), b"artifact")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_clear_does_not_repopulate_from_inflight_work(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def compute():
|
||||
started.set()
|
||||
await release.wait()
|
||||
return b"old-generation"
|
||||
|
||||
task = asyncio.create_task(cache.get_or_compute("key", compute))
|
||||
await started.wait()
|
||||
cache.clear()
|
||||
release.set()
|
||||
self.assertEqual((await task).value, b"old-generation")
|
||||
self.assertNotIn("key", cache)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_clear_starts_a_new_singleflight_generation(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def compute_old():
|
||||
started.set()
|
||||
await release.wait()
|
||||
return b"old"
|
||||
|
||||
async def compute_new():
|
||||
return b"new"
|
||||
|
||||
old_task = asyncio.create_task(cache.get_or_compute("key", compute_old))
|
||||
await started.wait()
|
||||
cache.clear()
|
||||
new_result = await cache.get_or_compute("key", compute_new)
|
||||
release.set()
|
||||
old_result = await old_task
|
||||
|
||||
self.assertEqual(old_result.value, b"old")
|
||||
self.assertEqual(new_result.value, b"new")
|
||||
self.assertEqual(cache.get("key"), b"new")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -138,6 +138,31 @@ class TestTemplateContentFormatDetection(CustomTestCase):
|
||||
self.assertEqual(result["content"], expected_content)
|
||||
self.assertEqual(result["role"], "user")
|
||||
|
||||
def test_process_content_preserves_image_content_hash(self):
|
||||
content_hash = "sha256:" + "ab" * 32
|
||||
image_data = []
|
||||
result = process_content_for_template_format(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "http://example.com/image.jpg",
|
||||
"content_hash": content_hash,
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
"openai",
|
||||
image_data,
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
)
|
||||
self.assertEqual(result["content"], [{"type": "image"}])
|
||||
self.assertEqual(image_data[0].content_hash, content_hash)
|
||||
|
||||
def test_process_content_string_format(self):
|
||||
"""Test content processing for string format."""
|
||||
msg_dict = {
|
||||
|
||||
Reference in New Issue
Block a user