[VLM] add content-addressed preprocessing cache infrastructure (#34398)

This commit is contained in:
Mick
2026-08-13 14:09:07 +08:00
committed by GitHub
parent 035c622a14
commit e8c7dddfa0
18 changed files with 1288 additions and 2 deletions
@@ -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 = {