[VLM] add content-addressed preprocessing cache infrastructure (#34398)
This commit is contained in:
@@ -369,6 +369,11 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
video_data: Optional[MultimodalDataInputFormat] = None,
|
video_data: Optional[MultimodalDataInputFormat] = None,
|
||||||
# See GenerateReqInput.mm_hashes / async_generate for the contract.
|
# See GenerateReqInput.mm_hashes / async_generate for the contract.
|
||||||
mm_hashes: Optional[Union[List[str], List[List[str]]]] = None,
|
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,
|
return_logprob: Optional[Union[List[bool], bool]] = False,
|
||||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||||
top_logprobs_num: 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,
|
audio_data=audio_data,
|
||||||
video_data=video_data,
|
video_data=video_data,
|
||||||
mm_hashes=mm_hashes,
|
mm_hashes=mm_hashes,
|
||||||
|
mm_content_hashes=mm_content_hashes,
|
||||||
cache_salt=cache_salt,
|
cache_salt=cache_salt,
|
||||||
return_logprob=return_logprob,
|
return_logprob=return_logprob,
|
||||||
logprob_start_len=logprob_start_len,
|
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
|
# that compute their own per-image hash for routing decisions and need
|
||||||
# sglang's prefix-cache key to align. See GenerateReqInput.mm_hashes.
|
# sglang's prefix-cache key to align. See GenerateReqInput.mm_hashes.
|
||||||
mm_hashes: Optional[Union[List[str], List[List[str]]]] = None,
|
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,
|
return_logprob: Optional[Union[List[bool], bool]] = False,
|
||||||
logprob_start_len: Optional[Union[List[int], int]] = None,
|
logprob_start_len: Optional[Union[List[int], int]] = None,
|
||||||
top_logprobs_num: 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,
|
audio_data=audio_data,
|
||||||
video_data=video_data,
|
video_data=video_data,
|
||||||
mm_hashes=mm_hashes,
|
mm_hashes=mm_hashes,
|
||||||
|
mm_content_hashes=mm_content_hashes,
|
||||||
cache_salt=cache_salt,
|
cache_salt=cache_salt,
|
||||||
return_logprob=return_logprob,
|
return_logprob=return_logprob,
|
||||||
logprob_start_len=logprob_start_len,
|
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)
|
kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60)
|
||||||
finally:
|
finally:
|
||||||
if isinstance(self.tokenizer_manager, TokenizerManager):
|
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()
|
self.tokenizer_manager.cuda_vmm_feature_transport.shutdown()
|
||||||
|
|
||||||
def __enter__(self):
|
def __enter__(self):
|
||||||
|
|||||||
@@ -529,6 +529,14 @@ class ChatCompletionMessageContentImageURL(BaseModel):
|
|||||||
detail: Optional[Literal["auto", "low", "high"]] = "auto"
|
detail: Optional[Literal["auto", "low", "high"]] = "auto"
|
||||||
max_dynamic_patch: Optional[int] = None
|
max_dynamic_patch: Optional[int] = None
|
||||||
min_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):
|
class ChatCompletionMessageContentVideoURL(BaseModel):
|
||||||
|
|||||||
@@ -196,6 +196,12 @@ class GenerateReqInput:
|
|||||||
# sglang's prefix-cache key to align. When unset, behavior is unchanged
|
# sglang's prefix-cache key to align. When unset, behavior is unchanged
|
||||||
# (sglang hashes the processor feature tensor).
|
# (sglang hashes the processor feature tensor).
|
||||||
mm_hashes: Optional[Union[List[str], List[List[str]]]] = None
|
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.
|
# Whether to extract and process audio from video inputs.
|
||||||
use_audio_in_video: bool = False
|
use_audio_in_video: bool = False
|
||||||
# The sampling_params. See descriptions below.
|
# The sampling_params. See descriptions below.
|
||||||
@@ -514,6 +520,7 @@ class GenerateReqInput:
|
|||||||
self._normalize_rid(num)
|
self._normalize_rid(num)
|
||||||
self._normalize_lora_paths(num)
|
self._normalize_lora_paths(num)
|
||||||
self._normalize_image_data(num)
|
self._normalize_image_data(num)
|
||||||
|
self._normalize_mm_hashes(num)
|
||||||
self._normalize_video_data(num)
|
self._normalize_video_data(num)
|
||||||
self._normalize_audio_data(num)
|
self._normalize_audio_data(num)
|
||||||
self._normalize_sampling_params(num)
|
self._normalize_sampling_params(num)
|
||||||
@@ -595,6 +602,41 @@ class GenerateReqInput:
|
|||||||
self.image_data = wrapped_images * self.parallel_sample_num
|
self.image_data = wrapped_images * self.parallel_sample_num
|
||||||
self.modalities = ["image"] * 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):
|
def _normalize_video_data(self, num):
|
||||||
"""Normalize video data for batch processing."""
|
"""Normalize video data for batch processing."""
|
||||||
if self.video_data is None:
|
if self.video_data is None:
|
||||||
@@ -817,6 +859,12 @@ class GenerateReqInput:
|
|||||||
image_data=self.image_data[i],
|
image_data=self.image_data[i],
|
||||||
video_data=self.video_data[i],
|
video_data=self.video_data[i],
|
||||||
audio_data=self.audio_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],
|
sampling_params=self.sampling_params[i],
|
||||||
return_logprob=self.return_logprob[i],
|
return_logprob=self.return_logprob[i],
|
||||||
logprob_start_len=self.logprob_start_len[i],
|
logprob_start_len=self.logprob_start_len[i],
|
||||||
|
|||||||
@@ -297,9 +297,12 @@ class TokenizerControlMixin:
|
|||||||
self: TokenizerManager, timeout_s: Optional[float] = None
|
self: TokenizerManager, timeout_s: Optional[float] = None
|
||||||
) -> FlushCacheReqOutput:
|
) -> FlushCacheReqOutput:
|
||||||
self.auto_create_handle_loop()
|
self.auto_create_handle_loop()
|
||||||
return (
|
result = (
|
||||||
await self.flush_cache_communicator(FlushCacheReqInput(timeout_s=timeout_s))
|
await self.flush_cache_communicator(FlushCacheReqInput(timeout_s=timeout_s))
|
||||||
)[0]
|
)[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:
|
async def clear_hicache_storage(self: TokenizerManager) -> ClearHiCacheReqOutput:
|
||||||
"""Clear the hierarchical cache storage."""
|
"""Clear the hierarchical cache storage."""
|
||||||
@@ -460,6 +463,8 @@ class TokenizerControlMixin:
|
|||||||
results = await self.update_weights_from_distributed_communicator(obj)
|
results = await self.update_weights_from_distributed_communicator(obj)
|
||||||
|
|
||||||
success, message = FanOutCommunicator.merge_results(results)
|
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:
|
if success and obj.weight_version is not None:
|
||||||
self._update_weight_version_if_provided(obj.weight_version)
|
self._update_weight_version_if_provided(obj.weight_version)
|
||||||
message += f" Weight version updated to {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)
|
results = await self.update_weights_from_tensor_communicator(obj)
|
||||||
|
|
||||||
success, message = FanOutCommunicator.merge_results(results)
|
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:
|
if success and obj.weight_version is not None:
|
||||||
self._update_weight_version_if_provided(obj.weight_version)
|
self._update_weight_version_if_provided(obj.weight_version)
|
||||||
message += f" Weight version updated to {obj.weight_version}."
|
message += f" Weight version updated to {obj.weight_version}."
|
||||||
@@ -556,6 +563,8 @@ class TokenizerControlMixin:
|
|||||||
logger.error(error_msg)
|
logger.error(error_msg)
|
||||||
success, message = False, 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:
|
if success and obj.weight_version is not None:
|
||||||
self._update_weight_version_if_provided(obj.weight_version)
|
self._update_weight_version_if_provided(obj.weight_version)
|
||||||
message += f" Weight version updated to {obj.weight_version}."
|
message += f" Weight version updated to {obj.weight_version}."
|
||||||
|
|||||||
@@ -1046,6 +1046,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
obj.audio_data = [obj.audio_data]
|
obj.audio_data = [obj.audio_data]
|
||||||
if contains_mm_input:
|
if contains_mm_input:
|
||||||
self._validate_mm_limits(obj)
|
self._validate_mm_limits(obj)
|
||||||
|
self._normalize_mm_content_hashes(obj)
|
||||||
|
|
||||||
mm_inputs = None
|
mm_inputs = None
|
||||||
mm_processor_input = (
|
mm_processor_input = (
|
||||||
@@ -1150,6 +1151,37 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
obj, input_text, input_ids, input_embeds, mm_inputs, token_type_ids
|
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(
|
def _validate_one_request(
|
||||||
self, obj: Union[GenerateReqInput, EmbeddingReqInput], input_ids: List[int]
|
self, obj: Union[GenerateReqInput, EmbeddingReqInput], input_ids: List[int]
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -2011,6 +2043,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
await self._wait_for_model_update_from_disk(obj)
|
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:
|
if success and obj.weight_version is not None:
|
||||||
self._update_weight_version_if_provided(obj.weight_version)
|
self._update_weight_version_if_provided(obj.weight_version)
|
||||||
message += f" Weight version updated to {obj.weight_version}."
|
message += f" Weight version updated to {obj.weight_version}."
|
||||||
|
|||||||
+29
@@ -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",
|
||||||
|
]
|
||||||
+329
@@ -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))
|
||||||
@@ -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),
|
||||||
|
}
|
||||||
@@ -19,6 +19,10 @@ from sglang.srt.managers.schedule_batch import (
|
|||||||
MultimodalInputFormat,
|
MultimodalInputFormat,
|
||||||
MultimodalProcessorOutput,
|
MultimodalProcessorOutput,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.multimodal.cache import (
|
||||||
|
MultimodalPreprocessCache,
|
||||||
|
build_processor_fingerprint,
|
||||||
|
)
|
||||||
from sglang.srt.multimodal.processors.executor import MultimodalProcessorExecutor
|
from sglang.srt.multimodal.processors.executor import MultimodalProcessorExecutor
|
||||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||||
MM_FEATURE_CACHE_SIZE,
|
MM_FEATURE_CACHE_SIZE,
|
||||||
@@ -185,6 +189,7 @@ class BaseMultimodalProcessor(ABC):
|
|||||||
preserve_processor_input_ids = False
|
preserve_processor_input_ids = False
|
||||||
auto_mm_processor_worker_num = 1
|
auto_mm_processor_worker_num = 1
|
||||||
auto_mm_io_worker_num = 4
|
auto_mm_io_worker_num = 4
|
||||||
|
auto_mm_preprocess_cache_size_mb = 0
|
||||||
supports_mm_processor_concurrency = False
|
supports_mm_processor_concurrency = False
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -219,6 +224,41 @@ class BaseMultimodalProcessor(ABC):
|
|||||||
self.video_config = mm_process_config.get("video", {})
|
self.video_config = mm_process_config.get("video", {})
|
||||||
self.audio_config = mm_process_config.get("audio", {})
|
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
|
# Resolve tokenizer: some processors (e.g. InternVL) pass a tokenizer
|
||||||
# directly as _processor rather than a processor that wraps a tokenizer.
|
# directly as _processor rather than a processor that wraps a tokenizer.
|
||||||
if hasattr(self._processor, "tokenizer"):
|
if hasattr(self._processor, "tokenizer"):
|
||||||
@@ -375,6 +415,30 @@ class BaseMultimodalProcessor(ABC):
|
|||||||
def keep_mm_features_on_device(self) -> bool:
|
def keep_mm_features_on_device(self) -> bool:
|
||||||
return self.mm_feature_transport in ("cuda_ipc", "cuda_vmm")
|
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):
|
def compute_mrope_positions(self, input_ids, mm_items):
|
||||||
"""Compute M-RoPE positions from expanded input_ids and multimodal items.
|
"""Compute M-RoPE positions from expanded input_ids and multimodal items.
|
||||||
|
|
||||||
@@ -938,7 +1002,6 @@ class BaseMultimodalProcessor(ABC):
|
|||||||
discard_alpha_channel: bool = True,
|
discard_alpha_channel: bool = True,
|
||||||
audio_sample_rate: Optional[int] = None,
|
audio_sample_rate: Optional[int] = None,
|
||||||
) -> BaseMultiModalProcessorOutput:
|
) -> BaseMultiModalProcessorOutput:
|
||||||
|
|
||||||
BaseMultimodalProcessor.validate_mm_data(image_data, video_data, audio_data)
|
BaseMultimodalProcessor.validate_mm_data(image_data, video_data, audio_data)
|
||||||
|
|
||||||
input_ids = prompt if isinstance(prompt, list) else None
|
input_ids = prompt if isinstance(prompt, list) else None
|
||||||
|
|||||||
@@ -290,6 +290,7 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
|
|||||||
precompute_hash_before_cpu_transfer = True
|
precompute_hash_before_cpu_transfer = True
|
||||||
auto_mm_processor_worker_num = 2
|
auto_mm_processor_worker_num = 2
|
||||||
auto_mm_io_worker_num = 16
|
auto_mm_io_worker_num = 16
|
||||||
|
auto_mm_preprocess_cache_size_mb = 256
|
||||||
supports_mm_processor_concurrency = True
|
supports_mm_processor_concurrency = True
|
||||||
preserve_processor_input_ids = True
|
preserve_processor_input_ids = True
|
||||||
|
|
||||||
|
|||||||
@@ -168,6 +168,7 @@ def process_content_for_template_format(
|
|||||||
url=image_obj["url"],
|
url=image_obj["url"],
|
||||||
detail=image_obj.get("detail") or "auto",
|
detail=image_obj.get("detail") or "auto",
|
||||||
max_dynamic_patch=mdp,
|
max_dynamic_patch=mdp,
|
||||||
|
content_hash=image_obj.get("content_hash"),
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -2769,6 +2769,21 @@ class ServerArgs:
|
|||||||
"environment override when this argument is 0.",
|
"environment override when this argument is 0.",
|
||||||
NS("mm"),
|
NS("mm"),
|
||||||
] = 0
|
] = 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[
|
limit_mm_data_per_request: A[
|
||||||
Optional[Union[str, Dict[str, int]]],
|
Optional[Union[str, Dict[str, int]]],
|
||||||
Arg(
|
Arg(
|
||||||
@@ -4008,6 +4023,11 @@ class ServerArgs:
|
|||||||
|
|
||||||
def _handle_multimodal(self):
|
def _handle_multimodal(self):
|
||||||
"""Validate mm_process_config structure before model loading."""
|
"""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 self.mm_process_config is not None:
|
||||||
if not isinstance(self.mm_process_config, dict):
|
if not isinstance(self.mm_process_config, dict):
|
||||||
raise TypeError(
|
raise TypeError(
|
||||||
|
|||||||
@@ -1634,6 +1634,7 @@ class ImageData:
|
|||||||
detail: Optional[Literal["auto", "low", "high"]] = "auto"
|
detail: Optional[Literal["auto", "low", "high"]] = "auto"
|
||||||
max_dynamic_patch: Optional[int] = None
|
max_dynamic_patch: Optional[int] = None
|
||||||
preprocess_kwargs: Optional[Dict] = None
|
preprocess_kwargs: Optional[Dict] = None
|
||||||
|
content_hash: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ from typing import List, Optional
|
|||||||
from pydantic import BaseModel, Field, ValidationError
|
from pydantic import BaseModel, Field, ValidationError
|
||||||
|
|
||||||
from sglang.srt.entrypoints.openai.protocol import (
|
from sglang.srt.entrypoints.openai.protocol import (
|
||||||
|
ChatCompletionMessageContentImageURL,
|
||||||
ChatCompletionRequest,
|
ChatCompletionRequest,
|
||||||
ChatCompletionResponse,
|
ChatCompletionResponse,
|
||||||
ChatCompletionResponseChoice,
|
ChatCompletionResponseChoice,
|
||||||
@@ -125,6 +126,17 @@ class TestChatCompletionRequest(unittest.TestCase):
|
|||||||
self.assertFalse(request.stream) # default
|
self.assertFalse(request.stream) # default
|
||||||
self.assertEqual(request.tool_choice, "none") # default when no tools
|
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):
|
def test_sampling_param_build(self):
|
||||||
req = ChatCompletionRequest(
|
req = ChatCompletionRequest(
|
||||||
model="x",
|
model="x",
|
||||||
|
|||||||
@@ -43,6 +43,32 @@ class TestMmHashesContract(CustomTestCase):
|
|||||||
req = GenerateReqInput(text="hi")
|
req = GenerateReqInput(text="hi")
|
||||||
self.assertIsNone(req.mm_hashes)
|
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):
|
def test_set_pad_value_honors_preset_hash(self):
|
||||||
"""set_pad_value() must use a pre-set hash without recomputing."""
|
"""set_pad_value() must use a pre-set hash without recomputing."""
|
||||||
item = MultimodalDataItem(modality=Modality.IMAGE, hash=0xDEADBEEF)
|
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_process_config = mm_process_config
|
||||||
server_args.mm_processor_worker_num = mm_processor_worker_num
|
server_args.mm_processor_worker_num = mm_processor_worker_num
|
||||||
server_args.mm_io_worker_num = mm_io_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()
|
hf_config = MagicMock()
|
||||||
mock_hf_processor = MagicMock()
|
mock_hf_processor = MagicMock()
|
||||||
@@ -767,6 +770,9 @@ class TestDoubleBosGuard(CustomTestCase):
|
|||||||
server_args.mm_io_worker_num = 0
|
server_args.mm_io_worker_num = 0
|
||||||
server_args.mm_feature_transport = "cpu"
|
server_args.mm_feature_transport = "cpu"
|
||||||
server_args.disable_fast_image_processor = True
|
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 = MagicMock()
|
||||||
mock_hf_processor.__class__.__name__ = "TestProcessor"
|
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["content"], expected_content)
|
||||||
self.assertEqual(result["role"], "user")
|
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):
|
def test_process_content_string_format(self):
|
||||||
"""Test content processing for string format."""
|
"""Test content processing for string format."""
|
||||||
msg_dict = {
|
msg_dict = {
|
||||||
|
|||||||
Reference in New Issue
Block a user