diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 0c4fc805a..75d33977b 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -369,6 +369,11 @@ class Engine(EngineScoreMixin, EngineBase): video_data: Optional[MultimodalDataInputFormat] = None, # See GenerateReqInput.mm_hashes / async_generate for the contract. mm_hashes: Optional[Union[List[str], List[List[str]]]] = None, + # SHA-256 identities for the original media contents. See + # GenerateReqInput.mm_content_hashes. + mm_content_hashes: Optional[ + Union[List[Optional[str]], List[List[Optional[str]]]] + ] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, @@ -413,6 +418,7 @@ class Engine(EngineScoreMixin, EngineBase): audio_data=audio_data, video_data=video_data, mm_hashes=mm_hashes, + mm_content_hashes=mm_content_hashes, cache_salt=cache_salt, return_logprob=return_logprob, logprob_start_len=logprob_start_len, @@ -478,6 +484,9 @@ class Engine(EngineScoreMixin, EngineBase): # that compute their own per-image hash for routing decisions and need # sglang's prefix-cache key to align. See GenerateReqInput.mm_hashes. mm_hashes: Optional[Union[List[str], List[List[str]]]] = None, + mm_content_hashes: Optional[ + Union[List[Optional[str]], List[List[Optional[str]]]] + ] = None, return_logprob: Optional[Union[List[bool], bool]] = False, logprob_start_len: Optional[Union[List[int], int]] = None, top_logprobs_num: Optional[Union[List[int], int]] = None, @@ -522,6 +531,7 @@ class Engine(EngineScoreMixin, EngineBase): audio_data=audio_data, video_data=video_data, mm_hashes=mm_hashes, + mm_content_hashes=mm_content_hashes, cache_salt=cache_salt, return_logprob=return_logprob, logprob_start_len=logprob_start_len, @@ -1260,6 +1270,9 @@ class Engine(EngineScoreMixin, EngineBase): kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60) finally: if isinstance(self.tokenizer_manager, TokenizerManager): + mm_processor = getattr(self.tokenizer_manager, "mm_processor", None) + if mm_processor is not None: + mm_processor.shutdown() self.tokenizer_manager.cuda_vmm_feature_transport.shutdown() def __enter__(self): diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 8c8fd7888..5652d63fd 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -529,6 +529,14 @@ class ChatCompletionMessageContentImageURL(BaseModel): detail: Optional[Literal["auto", "low", "high"]] = "auto" max_dynamic_patch: Optional[int] = None min_dynamic_patch: Optional[int] = None + content_hash: Optional[str] = None + + @field_validator("content_hash") + @classmethod + def validate_content_hash(cls, value: Optional[str]) -> Optional[str]: + from sglang.srt.multimodal.cache import parse_content_hash + + return parse_content_hash(value) class ChatCompletionMessageContentVideoURL(BaseModel): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 7d9c94a5d..736642c7f 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -196,6 +196,12 @@ class GenerateReqInput: # sglang's prefix-cache key to align. When unset, behavior is unchanged # (sglang hashes the processor feature tensor). mm_hashes: Optional[Union[List[str], List[List[str]]]] = None + # Optional `sha256:<64-hex>` identities for the original media contents. Unlike + # mm_hashes, these identify processor inputs and never replace the + # processor-output feature hash used by the embedding/prefix cache. + mm_content_hashes: Optional[ + Union[List[Optional[str]], List[List[Optional[str]]]] + ] = None # Whether to extract and process audio from video inputs. use_audio_in_video: bool = False # The sampling_params. See descriptions below. @@ -514,6 +520,7 @@ class GenerateReqInput: self._normalize_rid(num) self._normalize_lora_paths(num) self._normalize_image_data(num) + self._normalize_mm_hashes(num) self._normalize_video_data(num) self._normalize_audio_data(num) self._normalize_sampling_params(num) @@ -595,6 +602,41 @@ class GenerateReqInput: self.image_data = wrapped_images * self.parallel_sample_num self.modalities = ["image"] * num + def _normalize_mm_hashes(self, num): + """Align per-media hashes with normalized batched image inputs.""" + for field_name in ("mm_hashes", "mm_content_hashes"): + hashes = getattr(self, field_name) + if hashes is None: + setattr(self, field_name, [None] * num) + continue + if not isinstance(hashes, list): + raise ValueError(f"{field_name} must be a list") + if len(hashes) != self.batch_size: + raise ValueError( + f"The length of {field_name} should equal the batch size" + ) + + normalized = [] + for request_index, request_hashes in enumerate(hashes): + images = self.image_data[request_index] + image_count = len(images or []) + if isinstance(request_hashes, list): + per_request = request_hashes + elif image_count == 1: + per_request = [request_hashes] + else: + raise ValueError( + f"{field_name}[{request_index}] must be a list with one " + "entry per image" + ) + if len(per_request) != image_count: + raise ValueError( + f"{field_name}[{request_index}] has {len(per_request)} " + f"entries for {image_count} images" + ) + normalized.append(per_request) + setattr(self, field_name, normalized * self.parallel_sample_num) + def _normalize_video_data(self, num): """Normalize video data for batch processing.""" if self.video_data is None: @@ -817,6 +859,12 @@ class GenerateReqInput: image_data=self.image_data[i], video_data=self.video_data[i], audio_data=self.audio_data[i], + mm_hashes=self.mm_hashes[i] if self.mm_hashes is not None else None, + mm_content_hashes=( + self.mm_content_hashes[i] + if self.mm_content_hashes is not None + else None + ), sampling_params=self.sampling_params[i], return_logprob=self.return_logprob[i], logprob_start_len=self.logprob_start_len[i], diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 7cfd98b8f..c724a53d3 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -297,9 +297,12 @@ class TokenizerControlMixin: self: TokenizerManager, timeout_s: Optional[float] = None ) -> FlushCacheReqOutput: self.auto_create_handle_loop() - return ( + result = ( await self.flush_cache_communicator(FlushCacheReqInput(timeout_s=timeout_s)) )[0] + if result.success and self.mm_processor is not None: + self.mm_processor.clear_preprocess_cache() + return result async def clear_hicache_storage(self: TokenizerManager) -> ClearHiCacheReqOutput: """Clear the hierarchical cache storage.""" @@ -460,6 +463,8 @@ class TokenizerControlMixin: results = await self.update_weights_from_distributed_communicator(obj) success, message = FanOutCommunicator.merge_results(results) + if success and obj.flush_cache and self.mm_processor is not None: + self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." @@ -521,6 +526,8 @@ class TokenizerControlMixin: results = await self.update_weights_from_tensor_communicator(obj) success, message = FanOutCommunicator.merge_results(results) + if success and obj.flush_cache and self.mm_processor is not None: + self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." @@ -556,6 +563,8 @@ class TokenizerControlMixin: logger.error(error_msg) success, message = False, error_msg + if success and obj.flush_cache and self.mm_processor is not None: + self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index cd0c33fab..e60fc053e 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1046,6 +1046,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): obj.audio_data = [obj.audio_data] if contains_mm_input: self._validate_mm_limits(obj) + self._normalize_mm_content_hashes(obj) mm_inputs = None mm_processor_input = ( @@ -1150,6 +1151,37 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): obj, input_text, input_ids, input_embeds, mm_inputs, token_type_ids ) + @staticmethod + def _normalize_mm_content_hashes(obj: GenerateReqInput) -> None: + """Merge Native/OpenAI content identities and validate their alignment.""" + from sglang.srt.multimodal.cache import parse_content_hash + from sglang.srt.utils import ImageData + + images = obj.image_data or [] + explicit = obj.mm_content_hashes + inline = [ + image.content_hash if isinstance(image, ImageData) else None + for image in images + ] + if explicit is None and not any(inline): + return + if explicit is None: + explicit = inline + if len(explicit) != len(images): + raise ValueError( + f"mm_content_hashes has {len(explicit)} entries for " + f"{len(images)} images" + ) + + normalized = [] + for index, (provided, embedded) in enumerate(zip(explicit, inline)): + provided = parse_content_hash(provided) + embedded = parse_content_hash(embedded) + if provided is not None and embedded is not None and provided != embedded: + raise ValueError(f"Conflicting content hashes for image_data[{index}]") + normalized.append(provided or embedded) + obj.mm_content_hashes = normalized + def _validate_one_request( self, obj: Union[GenerateReqInput, EmbeddingReqInput], input_ids: List[int] ) -> None: @@ -2011,6 +2043,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): await self._wait_for_model_update_from_disk(obj) ) + if success and obj.flush_cache and self.mm_processor is not None: + self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: self._update_weight_version_if_provided(obj.weight_version) message += f" Weight version updated to {obj.weight_version}." diff --git a/python/sglang/srt/multimodal/cache/__init__.py b/python/sglang/srt/multimodal/cache/__init__.py new file mode 100644 index 000000000..f5a0b9ce5 --- /dev/null +++ b/python/sglang/srt/multimodal/cache/__init__.py @@ -0,0 +1,29 @@ +"""Content-addressed caches used by multimodal preprocessing.""" + +from sglang.srt.multimodal.cache.identity import ( + CONTENT_HASH_PREFIX, + MediaSnapshot, + build_artifact_key, + build_feature_hash, + build_processor_fingerprint, + parse_content_hash, + snapshot_media, +) +from sglang.srt.multimodal.cache.preprocess_cache import ( + CacheLookup, + MultimodalPreprocessCache, + estimate_cache_size_bytes, +) + +__all__ = [ + "CONTENT_HASH_PREFIX", + "CacheLookup", + "MediaSnapshot", + "MultimodalPreprocessCache", + "build_artifact_key", + "build_feature_hash", + "build_processor_fingerprint", + "estimate_cache_size_bytes", + "parse_content_hash", + "snapshot_media", +] diff --git a/python/sglang/srt/multimodal/cache/identity.py b/python/sglang/srt/multimodal/cache/identity.py new file mode 100644 index 000000000..a53c64bd9 --- /dev/null +++ b/python/sglang/srt/multimodal/cache/identity.py @@ -0,0 +1,329 @@ +"""Stable identities for multimodal inputs and processor artifacts.""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import struct +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Mapping, Optional +from urllib.parse import unquote, urlparse + +import numpy as np +import torch +import transformers +from PIL import Image + +CONTENT_HASH_PREFIX = "sha256:" +_SHA256_HEX_LENGTH = 64 + + +def parse_content_hash(value: Optional[str]) -> Optional[str]: + """Validate and normalize a public content digest.""" + if value is None: + return None + if not isinstance(value, str) or not value.startswith(CONTENT_HASH_PREFIX): + raise ValueError("content_hash must use the form 'sha256:<64 hex digits>'") + digest = value[len(CONTENT_HASH_PREFIX) :] + if len(digest) != _SHA256_HEX_LENGTH: + raise ValueError("content_hash must contain exactly 64 SHA-256 hex digits") + try: + bytes.fromhex(digest) + except ValueError as exc: + raise ValueError("content_hash contains non-hexadecimal characters") from exc + return CONTENT_HASH_PREFIX + digest.lower() + + +def _digest_bytes(payload: bytes) -> str: + return CONTENT_HASH_PREFIX + hashlib.sha256(payload).hexdigest() + + +def _hash_parts(*parts: bytes) -> str: + hasher = hashlib.sha256() + for part in parts: + hasher.update(len(part).to_bytes(8, "big")) + hasher.update(part) + return CONTENT_HASH_PREFIX + hasher.hexdigest() + + +@dataclass(frozen=True) +class MediaSnapshot: + """An immutable-enough media snapshot paired with its strict identity.""" + + data: Any + content_digest: str + size_bytes: int + source: str + + +def _snapshot_pil(image: Image.Image) -> MediaSnapshot: + snapshot = image.copy() + snapshot.load() + payload = snapshot.tobytes() + palette = snapshot.palette.tobytes() if snapshot.palette is not None else b"" + palette_mode = ( + snapshot.palette.mode.encode() if snapshot.palette is not None else b"" + ) + transparency = snapshot.info.get("transparency") + if transparency is None: + transparency_payload = b"none" + elif isinstance(transparency, bytes): + transparency_payload = b"bytes:" + transparency + else: + transparency_payload = ( + f"{type(transparency).__name__}:{transparency!r}".encode() + ) + digest = _hash_parts( + b"pil", + snapshot.mode.encode(), + json.dumps(snapshot.size).encode(), + palette_mode, + palette, + transparency_payload, + payload, + ) + return MediaSnapshot(snapshot, digest, len(payload), "pil") + + +def _snapshot_tensor(tensor: torch.Tensor) -> MediaSnapshot: + snapshot = tensor.detach().to("cpu").contiguous().clone() + payload = snapshot.view(torch.uint8).numpy().tobytes() + digest = _hash_parts( + b"torch", + str(snapshot.dtype).encode(), + json.dumps(list(snapshot.shape)).encode(), + payload, + ) + return MediaSnapshot(snapshot, digest, len(payload), "tensor") + + +def _snapshot_ndarray(array: np.ndarray) -> MediaSnapshot: + snapshot = np.ascontiguousarray(array).copy() + payload = snapshot.view(np.uint8).tobytes() + digest = _hash_parts( + b"numpy", + snapshot.dtype.str.encode(), + json.dumps(list(snapshot.shape)).encode(), + payload, + ) + return MediaSnapshot(snapshot, digest, len(payload), "ndarray") + + +def _read_media_bytes(media: str | bytes) -> bytes: + if isinstance(media, bytes): + return bytes(media) + + from sglang.srt.utils import get_image_bytes, image_extension_names + + if media.startswith("file://"): + media = unquote(urlparse(media).path) + elif media.startswith(("http://", "https://", "data:")): + return get_image_bytes(media) + # ``load_image`` accepts relative local paths only by image extension. + # Match that contract instead of probing arbitrary base64 as a filename. + if media.lower().endswith(image_extension_names) and Path(media).is_file(): + return Path(media).read_bytes() + return get_image_bytes(media) + + +def snapshot_media(media: Any) -> MediaSnapshot: + """Snapshot media and hash exactly what will be handed to the decoder. + + Paths and URLs are deliberately not identities. They are resolved to bytes + on every untrusted lookup, so changing their contents produces a cache miss. + """ + from sglang.srt.utils import ImageData + + if isinstance(media, ImageData): + media = media.url + elif isinstance(media, Mapping) and "format" not in media: + if "url" in media: + media = media["url"] + elif "image" in media: + media = media["image"] + + if isinstance(media, (str, bytes)): + payload = _read_media_bytes(media) + return MediaSnapshot(payload, _digest_bytes(payload), len(payload), "bytes") + if isinstance(media, Image.Image): + return _snapshot_pil(media) + if isinstance(media, torch.Tensor): + return _snapshot_tensor(media) + if isinstance(media, np.ndarray): + return _snapshot_ndarray(media) + raise TypeError(f"Unsupported media identity input: {type(media).__name__}") + + +def _qualified_type_name(value: Any) -> str: + value_type = type(value) + return f"{value_type.__module__}.{value_type.__qualname__}" + + +def _canonical_sort_key(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _canonicalize(value: Any) -> Any: + """Encode cache-key inputs without collapsing distinct Python values.""" + if dataclasses.is_dataclass(value): + return { + "type": "dataclass", + "class": _qualified_type_name(value), + "fields": [ + [field.name, _canonicalize(getattr(value, field.name))] + for field in dataclasses.fields(value) + ], + } + if isinstance(value, Enum): + return { + "type": "enum", + "class": _qualified_type_name(value), + "value": _canonicalize(value.value), + } + if isinstance(value, Path): + return {"type": "path", "value": str(value)} + if value is None: + return {"type": "none"} + if isinstance(value, bool): + return {"type": "bool", "value": value} + if isinstance(value, int): + return {"type": "int", "value": str(value)} + if isinstance(value, float): + return {"type": "float64", "bits": struct.pack("!d", value).hex()} + if isinstance(value, str): + return {"type": "str", "value": value} + if isinstance(value, (bytes, bytearray, memoryview)): + return {"type": "bytes", "value": bytes(value).hex()} + if isinstance(value, torch.dtype): + return {"type": "torch_dtype", "value": str(value)} + if isinstance(value, torch.Tensor): + snapshot = value.detach().to("cpu").contiguous() + payload = snapshot.view(torch.uint8).numpy().tobytes() + return { + "type": "torch_tensor", + "dtype": str(snapshot.dtype), + "shape": list(snapshot.shape), + "digest": _digest_bytes(payload), + } + if isinstance(value, np.generic): + scalar = np.asarray(value) + return { + "type": "numpy_scalar", + "dtype": scalar.dtype.str, + "value": scalar.tobytes().hex(), + } + if isinstance(value, np.ndarray): + snapshot = np.ascontiguousarray(value) + return { + "type": "numpy_array", + "dtype": snapshot.dtype.str, + "shape": list(snapshot.shape), + "digest": _digest_bytes(snapshot.view(np.uint8).tobytes()), + } + if isinstance(value, Mapping): + items = [ + [_canonicalize(key), _canonicalize(item)] for key, item in value.items() + ] + items.sort(key=lambda pair: _canonical_sort_key(pair[0])) + return { + "type": "mapping", + "items": items, + } + if isinstance(value, list): + return {"type": "list", "items": [_canonicalize(item) for item in value]} + if isinstance(value, tuple): + return {"type": "tuple", "items": [_canonicalize(item) for item in value]} + if isinstance(value, (set, frozenset)): + items = [_canonicalize(item) for item in value] + items.sort(key=_canonical_sort_key) + return { + "type": "frozenset" if isinstance(value, frozenset) else "set", + "items": items, + } + raise ValueError( + "Unsupported value in multimodal cache identity: " + f"{_qualified_type_name(value)}" + ) + + +def _canonical_json(value: Any) -> bytes: + return json.dumps( + _canonicalize(value), sort_keys=True, separators=(",", ":") + ).encode() + + +def build_artifact_key( + content_digest: str, + *, + modality: str, + processor_fingerprint: str, + preprocess_kwargs: Optional[Mapping[str, Any]] = None, +) -> str: + """Build the cache key for a processor artifact.""" + content_digest = parse_content_hash(content_digest) + payload = { + "content_digest": content_digest, + "modality": modality, + "processor_fingerprint": processor_fingerprint, + "preprocess_kwargs": preprocess_kwargs or {}, + } + return _digest_bytes(_canonical_json(payload)) + + +def build_feature_hash(artifact_key: str, processor_output_hash: int) -> int: + """Namespace a processor-output hash by its complete artifact identity.""" + artifact_key = parse_content_hash(artifact_key) + if ( + isinstance(processor_output_hash, bool) + or not isinstance(processor_output_hash, int) + or processor_output_hash < 0 + ): + raise ValueError("processor_output_hash must be a non-negative integer") + output_hash_bytes = processor_output_hash.to_bytes( + max(1, (processor_output_hash.bit_length() + 7) // 8), + byteorder="big", + signed=False, + ) + digest = _hash_parts( + b"multimodal-feature-v1", + bytes.fromhex(artifact_key[len(CONTENT_HASH_PREFIX) :]), + output_hash_bytes, + ) + return int.from_bytes( + bytes.fromhex(digest[len(CONTENT_HASH_PREFIX) :])[:8], + byteorder="big", + signed=False, + ) + + +def build_processor_fingerprint( + processor: Any, + hf_config: Any, + server_args: Any, + *, + extra: Optional[Mapping[str, Any]] = None, +) -> str: + """Fingerprint preprocessing choices that can change processor output.""" + processor_payload = ( + processor.preprocess_fingerprint_payload() + if hasattr(processor, "preprocess_fingerprint_payload") + else {} + ) + payload = { + "transformers": transformers.__version__, + "processor_class": f"{type(processor).__module__}.{type(processor).__qualname__}", + "model_type": getattr(hf_config, "model_type", None), + "architectures": getattr(hf_config, "architectures", None), + "model_revision": getattr(server_args, "revision", None), + "tokenizer_revision": getattr(server_args, "tokenizer_revision", None), + "disable_fast_image_processor": getattr( + server_args, "disable_fast_image_processor", False + ), + "mm_process_config": getattr(server_args, "mm_process_config", None) or {}, + "processor": processor_payload, + "extra": extra or {}, + } + return _digest_bytes(_canonical_json(payload)) diff --git a/python/sglang/srt/multimodal/cache/preprocess_cache.py b/python/sglang/srt/multimodal/cache/preprocess_cache.py new file mode 100644 index 000000000..c8e48cfdf --- /dev/null +++ b/python/sglang/srt/multimodal/cache/preprocess_cache.py @@ -0,0 +1,259 @@ +"""Bounded CPU cache and single-flight coordination for MM preprocessing.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import dataclasses +import sys +import threading +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Generic, Optional, TypeVar + +import numpy as np +import torch +from PIL import Image + +K = TypeVar("K") +V = TypeVar("V") + + +@dataclass(frozen=True) +class CacheLookup(Generic[V]): + value: V + hit: bool + joined: bool = False + + +@dataclass +class _Entry(Generic[V]): + value: V + size_bytes: int + + +def estimate_cache_size_bytes(value: Any) -> Optional[int]: + """Estimate owned CPU bytes, returning None for GPU-backed artifacts.""" + seen: set[int] = set() + + def visit(item: Any) -> Optional[int]: + if item is None or isinstance(item, (bool, int, float)): + return sys.getsizeof(item) + item_id = id(item) + if item_id in seen: + return 0 + seen.add(item_id) + + if isinstance(item, torch.Tensor): + if item.device.type != "cpu": + return None + return item.untyped_storage().nbytes() + if isinstance(item, np.ndarray): + return int(item.nbytes) + if isinstance(item, Image.Image): + return len(item.tobytes()) + if isinstance(item, (bytes, bytearray, memoryview)): + return len(item) + if isinstance(item, str): + return len(item.encode()) + if dataclasses.is_dataclass(item): + return visit(dataclasses.asdict(item)) + if isinstance(item, dict): + total = 0 + for key, child in item.items(): + key_size = visit(key) + child_size = visit(child) + if key_size is None or child_size is None: + return None + total += key_size + child_size + return total + if isinstance(item, (list, tuple, set)): + total = 0 + for child in item: + child_size = visit(child) + if child_size is None: + return None + total += child_size + return total + return sys.getsizeof(item) + + return visit(value) + + +class MultimodalPreprocessCache(Generic[K, V]): + """Thread-safe byte-accounted LRU with per-key async single-flight.""" + + def __init__(self, max_size_bytes: int, max_entries: int = 8192): + if max_size_bytes < 0: + raise ValueError("max_size_bytes must be non-negative") + if max_entries <= 0: + raise ValueError("max_entries must be positive") + self.max_size_bytes = max_size_bytes + self.max_entries = max_entries + self._entries: OrderedDict[K, _Entry[V]] = OrderedDict() + self._inflight: dict[K, tuple[concurrent.futures.Future[V], int]] = {} + self._background_tasks: set[asyncio.Task] = set() + self._lock = threading.Lock() + self._generation = 0 + self.current_size_bytes = 0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + self.singleflight_joins = 0 + + @property + def enabled(self) -> bool: + return self.max_size_bytes > 0 + + def __len__(self) -> int: + with self._lock: + return len(self._entries) + + def __contains__(self, key: K) -> bool: + with self._lock: + return key in self._entries + + def get(self, key: K) -> Optional[V]: + with self._lock: + entry = self._entries.get(key) + if entry is None: + self.misses += 1 + return None + self._entries.move_to_end(key) + self.hits += 1 + return entry.value + + def put( + self, + key: K, + value: V, + size_bytes: Optional[int] = None, + *, + _generation: Optional[int] = None, + ) -> bool: + if not self.enabled: + return False + if size_bytes is None: + size_bytes = estimate_cache_size_bytes(value) + if size_bytes is None or size_bytes < 0 or size_bytes > self.max_size_bytes: + return False + + with self._lock: + if _generation is not None and _generation != self._generation: + return False + old = self._entries.pop(key, None) + if old is not None: + self.current_size_bytes -= old.size_bytes + while self._entries and ( + self.current_size_bytes + size_bytes > self.max_size_bytes + or len(self._entries) >= self.max_entries + ): + _, evicted = self._entries.popitem(last=False) + self.current_size_bytes -= evicted.size_bytes + self.evictions += 1 + self._entries[key] = _Entry(value=value, size_bytes=size_bytes) + self.current_size_bytes += size_bytes + return True + + def pop(self, key: K) -> Optional[V]: + with self._lock: + entry = self._entries.pop(key, None) + if entry is None: + return None + self.current_size_bytes -= entry.size_bytes + return entry.value + + def clear(self) -> None: + with self._lock: + self._entries.clear() + self.current_size_bytes = 0 + # Let active requests finish, but prevent work started before this + # flush from repopulating the cache afterwards. + self._generation += 1 + + async def get_or_compute( + self, + key: K, + compute: Callable[[], Awaitable[V]], + *, + size_bytes: Optional[Callable[[V], Optional[int]]] = None, + ) -> CacheLookup[V]: + cached = self.get(key) + if cached is not None: + return CacheLookup(cached, hit=True) + + with self._lock: + inflight = self._inflight.get(key) + if inflight is None or inflight[1] != self._generation: + future = concurrent.futures.Future() + generation = self._generation + self._inflight[key] = (future, generation) + owner = True + else: + future, generation = inflight + self.singleflight_joins += 1 + owner = False + + if owner: + self.create_background_task( + self._compute_owned_value( + key, future, generation, compute, size_bytes=size_bytes + ) + ) + + # The cache owns the shared computation. Cancelling either its first + # caller or a later joiner ends only that caller's local await. + value = await asyncio.shield(asyncio.wrap_future(future)) + return CacheLookup(value, hit=False, joined=not owner) + + async def _compute_owned_value( + self, + key: K, + future: concurrent.futures.Future[V], + generation: int, + compute: Callable[[], Awaitable[V]], + *, + size_bytes: Optional[Callable[[V], Optional[int]]], + ) -> None: + try: + value = await compute() + measured = size_bytes(value) if size_bytes is not None else None + self.put(key, value, measured, _generation=generation) + future.set_result(value) + except BaseException as exc: + future.set_exception(exc) + # Retrieve the exception locally when no waiter joined, avoiding a + # noisy "Future exception was never retrieved" warning. + future.exception() + finally: + with self._lock: + if self._inflight.get(key) == (future, generation): + self._inflight.pop(key, None) + + def _background_task_done(self, task: asyncio.Task) -> None: + with self._lock: + self._background_tasks.discard(task) + try: + task.exception() + except asyncio.CancelledError: + pass + + def create_background_task(self, awaitable: Awaitable[Any]) -> asyncio.Task: + """Keep shared cache work alive independently of one caller task.""" + task = asyncio.create_task(awaitable) + with self._lock: + self._background_tasks.add(task) + task.add_done_callback(self._background_task_done) + return task + + def stats(self) -> dict[str, int]: + with self._lock: + return { + "entries": len(self._entries), + "size_bytes": self.current_size_bytes, + "hits": self.hits, + "misses": self.misses, + "evictions": self.evictions, + "singleflight_joins": self.singleflight_joins, + "inflight": len(self._inflight), + } diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index f7ffde730..61339c39b 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -19,6 +19,10 @@ from sglang.srt.managers.schedule_batch import ( MultimodalInputFormat, MultimodalProcessorOutput, ) +from sglang.srt.multimodal.cache import ( + MultimodalPreprocessCache, + build_processor_fingerprint, +) from sglang.srt.multimodal.processors.executor import MultimodalProcessorExecutor from sglang.srt.multimodal.transport.cuda_ipc import ( MM_FEATURE_CACHE_SIZE, @@ -185,6 +189,7 @@ class BaseMultimodalProcessor(ABC): preserve_processor_input_ids = False auto_mm_processor_worker_num = 1 auto_mm_io_worker_num = 4 + auto_mm_preprocess_cache_size_mb = 0 supports_mm_processor_concurrency = False def __init__( @@ -219,6 +224,41 @@ class BaseMultimodalProcessor(ABC): self.video_config = mm_process_config.get("video", {}) self.audio_config = mm_process_config.get("audio", {}) + requested_cache_mb = getattr( + self.server_args, "mm_preprocess_cache_size_mb", None + ) + total_cache_mb = ( + self.auto_mm_preprocess_cache_size_mb + if requested_cache_mb is None + else requested_cache_mb + ) + tokenizer_worker_num = max( + int(getattr(self.server_args, "tokenizer_worker_num", 1)), 1 + ) + worker_cache_bytes = total_cache_mb * 1024 * 1024 // tokenizer_worker_num + self.mm_preprocess_cache = MultimodalPreprocessCache( + max_size_bytes=worker_cache_bytes, + max_entries=8192, + ) + self.trust_mm_content_hashes = bool( + getattr(self.server_args, "trust_mm_content_hashes", False) + ) + self.processor_fingerprint = ( + build_processor_fingerprint(self, hf_config, server_args) + if self.mm_preprocess_cache.enabled + else None + ) + if self.mm_preprocess_cache.enabled: + logger.info( + "Multimodal preprocess cache enabled for %s: %d MiB total " + "(%d MiB per tokenizer worker), at most 8192 entries; " + "caller content hashes are %s.", + type(self).__name__, + total_cache_mb, + worker_cache_bytes // (1024 * 1024), + "trusted" if self.trust_mm_content_hashes else "verified", + ) + # Resolve tokenizer: some processors (e.g. InternVL) pass a tokenizer # directly as _processor rather than a processor that wraps a tokenizer. if hasattr(self._processor, "tokenizer"): @@ -375,6 +415,30 @@ class BaseMultimodalProcessor(ABC): def keep_mm_features_on_device(self) -> bool: return self.mm_feature_transport in ("cuda_ipc", "cuda_vmm") + def preprocess_fingerprint_payload(self) -> dict[str, Any]: + """Stable processor choices that may change per-media artifacts.""" + return { + "wrapper_class": ( + f"{type(self._processor).__module__}." + f"{type(self._processor).__qualname__}" + ), + "gpu_image_decode": self.gpu_image_decode, + "image_config": self.image_config, + "video_config": self.video_config, + "audio_config": self.audio_config, + } + + def clear_preprocess_cache(self) -> None: + self.mm_preprocess_cache.clear() + + def shutdown(self) -> None: + """Release executor resources and cached CPU artifacts.""" + self.clear_preprocess_cache() + self.io_executor.shutdown(wait=False, cancel_futures=True) + self.cpu_executor.shutdown(wait=False, cancel_futures=True) + if self.mm_processor_executor is not None: + self.mm_processor_executor.shutdown() + def compute_mrope_positions(self, input_ids, mm_items): """Compute M-RoPE positions from expanded input_ids and multimodal items. @@ -938,7 +1002,6 @@ class BaseMultimodalProcessor(ABC): discard_alpha_channel: bool = True, audio_sample_rate: Optional[int] = None, ) -> BaseMultiModalProcessorOutput: - BaseMultimodalProcessor.validate_mm_data(image_data, video_data, audio_data) input_ids = prompt if isinstance(prompt, list) else None diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py index 6b1170757..a5686f906 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k3.py +++ b/python/sglang/srt/multimodal/processors/kimi_k3.py @@ -290,6 +290,7 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): precompute_hash_before_cpu_transfer = True auto_mm_processor_worker_num = 2 auto_mm_io_worker_num = 16 + auto_mm_preprocess_cache_size_mb = 256 supports_mm_processor_concurrency = True preserve_processor_input_ids = True diff --git a/python/sglang/srt/parser/jinja_template_utils.py b/python/sglang/srt/parser/jinja_template_utils.py index c187219dc..f22222fd9 100644 --- a/python/sglang/srt/parser/jinja_template_utils.py +++ b/python/sglang/srt/parser/jinja_template_utils.py @@ -168,6 +168,7 @@ def process_content_for_template_format( url=image_obj["url"], detail=image_obj.get("detail") or "auto", max_dynamic_patch=mdp, + content_hash=image_obj.get("content_hash"), ) ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 932011cee..ba4c66432 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2769,6 +2769,21 @@ class ServerArgs: "environment override when this argument is 0.", NS("mm"), ] = 0 + mm_preprocess_cache_size_mb: A[ + Optional[int], + "CPU memory budget for content-addressed multimodal preprocessing " + "artifacts. Unset selects a model-specific default (256 MiB for " + "Kimi-K3); 0 disables the cache. The budget is divided across " + "tokenizer workers and does not reserve GPU memory.", + NS("mm"), + ] = None + trust_mm_content_hashes: A[ + bool, + "Trust caller-provided multimodal SHA-256 content hashes. This can " + "skip reading media on a hot metadata-cache hit; only enable it when " + "the caller guarantees that hashes identify immutable media bytes.", + NS("mm"), + ] = False limit_mm_data_per_request: A[ Optional[Union[str, Dict[str, int]]], Arg( @@ -4008,6 +4023,11 @@ class ServerArgs: def _handle_multimodal(self): """Validate mm_process_config structure before model loading.""" + if ( + self.mm_preprocess_cache_size_mb is not None + and self.mm_preprocess_cache_size_mb < 0 + ): + raise ValueError("mm_preprocess_cache_size_mb must be non-negative") if self.mm_process_config is not None: if not isinstance(self.mm_process_config, dict): raise TypeError( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 466f72ea5..690d4f629 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1634,6 +1634,7 @@ class ImageData: detail: Optional[Literal["auto", "low", "high"]] = "auto" max_dynamic_patch: Optional[int] = None preprocess_kwargs: Optional[Dict] = None + content_hash: Optional[str] = None @dataclass diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index 33cb1c379..9b0cd6868 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -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", diff --git a/test/registered/unit/managers/test_mm_hashes.py b/test/registered/unit/managers/test_mm_hashes.py index 85d888542..943586e34 100644 --- a/test/registered/unit/managers/test_mm_hashes.py +++ b/test/registered/unit/managers/test_mm_hashes.py @@ -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) diff --git a/test/registered/unit/managers/test_mm_process_config.py b/test/registered/unit/managers/test_mm_process_config.py index 7f9c8b18c..3e665f9e0 100644 --- a/test/registered/unit/managers/test_mm_process_config.py +++ b/test/registered/unit/managers/test_mm_process_config.py @@ -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" diff --git a/test/registered/unit/multimodal/test_preprocess_cache.py b/test/registered/unit/multimodal/test_preprocess_cache.py new file mode 100644 index 000000000..8b9fcfe59 --- /dev/null +++ b/test/registered/unit/multimodal/test_preprocess_cache.py @@ -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() diff --git a/test/registered/unit/parser/test_jinja_template_utils.py b/test/registered/unit/parser/test_jinja_template_utils.py index 79b43e977..85fcd474a 100644 --- a/test/registered/unit/parser/test_jinja_template_utils.py +++ b/test/registered/unit/parser/test_jinja_template_utils.py @@ -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 = {