vlm: cache kimi-k3 per-image processor artifacts (#34404)

This commit is contained in:
Mick
2026-08-16 19:51:13 +08:00
committed by GitHub
parent 968b355f12
commit 3d3194f6c3
20 changed files with 2617 additions and 276 deletions
@@ -19,6 +19,7 @@ from sglang.srt.disaggregation.encode_receiver import (
EmbeddingData,
MMReceiverHTTP,
MultiModalEmbeddingData,
_encoder_media_item,
_select_mm_processor_prompt,
)
from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
@@ -27,8 +28,10 @@ from sglang.srt.managers.tokenizer_manager import (
_reject_missing_dispatched_encoder_embedding,
)
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
from sglang.srt.multimodal.cache import snapshot_media
from sglang.srt.multimodal.encoder_preprocessing import (
LOCAL_PREPROCESSED_KEY,
EncoderMediaProcessorConfig,
EncoderPreprocessOutput,
get_encoder_preprocessed_items,
hash_raw_encoder_item,
@@ -41,6 +44,7 @@ from sglang.srt.multimodal.kimi_k3_image_processing import (
)
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import resolve_encoder_transfer_backend
from sglang.srt.utils import ImageData
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -208,6 +212,11 @@ def _encoder(model_type="kimi_k3"):
vision_config=SimpleNamespace(merge_kernel_size=(2, 2))
)
)
encoder.encoder_media_processor_config = (
KimiK3ForConditionalGeneration.encoder_media_processor_config
if model_type == "kimi_k3"
else EncoderMediaProcessorConfig()
)
return encoder
@@ -296,6 +305,19 @@ def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items():
assert deferred.image_std == [0.5, 0.5, 0.5]
def test_kimi_k3_epd_preserves_verified_content_identity():
image = Image.new("RGB", (8, 6), color=(1, 2, 3))
digest = "sha256:" + "ab" * 32
output = prepare_kimi_k3_encoder_inputs(
[{"type": "image", "image": image, "content_hash": digest}],
_kimi_k3_image_processor(),
)
item = get_encoder_preprocessed_items(output)[0]
assert item.model_specific_data["content_digest"] == digest
def test_kimi_k3_epd_model_preprocessor_receives_image_processor():
image = Image.new("RGB", (8, 6), color=(1, 2, 3))
image_processor = _kimi_k3_image_processor()
@@ -458,6 +480,67 @@ def test_kimi_k3_epd_selects_matching_jpeg_decode_mode(
load.assert_called_once_with(b"jpeg", expected_decode_mode)
def test_kimi_k3_epd_verifies_content_hash_before_decode():
payload = b"jpeg"
digest = snapshot_media(payload).content_digest
expected = torch.zeros((3, 2, 3), dtype=torch.uint8)
encoder = _encoder()
encoder.use_image_processor_gpu = False
with patch(
"sglang.srt.disaggregation.encode_server.load_image",
return_value=(expected, None),
) as load:
output = encoder._load_single_item(
{"url": payload, "content_hash": digest}, Modality.IMAGE
)
assert output == {
"type": "image",
"image": expected,
"content_hash": digest,
}
load.assert_called_once_with(payload, False)
def test_epd_receiver_keeps_content_hash_aligned_with_image():
digest = "sha256:" + "cd" * 32
receiver = MMReceiverHTTP.__new__(MMReceiverHTTP)
request = SimpleNamespace(
image_data=[
ImageData(
url="image",
detail="high",
max_dynamic_patch=12,
preprocess_kwargs={"crop": False},
content_hash=digest,
)
],
video_data=None,
audio_data=None,
mm_content_hashes=[digest],
)
assert receiver._extract_url_data(request) == [
{
"url": "image",
"modality": Modality.IMAGE,
"detail": "high",
"max_dynamic_patch": 12,
"preprocess_kwargs": {"crop": False},
"content_hash": digest,
}
]
assert _encoder_media_item(receiver._extract_url_data(request)[0]) == {
"url": "image",
"detail": "high",
"max_dynamic_patch": 12,
"preprocess_kwargs": {"crop": False},
"content_hash": digest,
}
def test_kimi_k3_epd_aggregates_original_image_sizes_in_part_order():
first = EmbeddingData(
req_id="request",
@@ -127,12 +127,14 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("SGLANG_IO_WORKERS", None)
with patch.object(
BaseMultimodalProcessor, "auto_mm_processor_worker_num", 4
), patch.object(
BaseMultimodalProcessor, "auto_mm_io_worker_num", 16
), patch.object(
BaseMultimodalProcessor, "supports_mm_processor_concurrency", True
with (
patch.object(
BaseMultimodalProcessor, "auto_mm_processor_worker_num", 4
),
patch.object(BaseMultimodalProcessor, "auto_mm_io_worker_num", 16),
patch.object(
BaseMultimodalProcessor, "supports_mm_processor_concurrency", True
),
):
proc = self._make_processor({})
try:
@@ -172,14 +174,18 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
def _server_args(mm_feature_transport):
return SimpleNamespace(
mm_feature_transport=mm_feature_transport,
image_processor_backend="auto",
disable_fast_image_processor=False,
skip_tokenizer_init=False,
mm_process_config={},
mm_preprocess_cache_size_mb=0,
trust_mm_content_hashes=False,
mm_processor_worker_num=0,
mm_io_worker_num=0,
tokenizer_worker_num=1,
base_gpu_id=2,
tp_size=8,
rl_on_policy_target=None,
allowed_media_domains=[],
media_url_max_file_size_mb=64,
)
@@ -195,9 +201,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
# transport policy must still resolve from the instance's ServerArgs.
from sglang.srt.multimodal.processors import base_processor
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True), patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
with (
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True),
patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
),
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
):
processor = base_processor.BaseMultimodalProcessor(
hf_config=MagicMock(),
server_args=self._server_args("cuda_ipc"),
@@ -213,9 +223,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
def test_cuda_ipc_pool_handle_cache_can_be_disabled(self):
from sglang.srt.multimodal.processors import base_processor
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(False), patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
with (
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(False),
patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
),
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
):
processor = base_processor.BaseMultimodalProcessor(
hf_config=MagicMock(),
server_args=self._server_args("cuda_ipc"),
@@ -230,9 +244,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
def test_cpu_transport_does_not_allocate_ipc_pool(self):
from sglang.srt.multimodal.processors import base_processor
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True), patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
with (
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True),
patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
),
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
):
processor = base_processor.BaseMultimodalProcessor(
hf_config=MagicMock(),
server_args=self._server_args("cpu"),
@@ -251,9 +269,12 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
hf_processor = self._processor()
feature = torch.empty(1, device="meta")
hf_processor.return_value = {"pixel_values": feature}
with patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
with (
patch.object(
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
),
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
):
processor = base_processor.BaseMultimodalProcessor(
hf_config=MagicMock(),
server_args=self._server_args("cuda_vmm"),
@@ -365,9 +386,10 @@ class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase):
BaseMultimodalProcessor,
)
with patch.object(
BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
with (
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
):
processor = BaseMultimodalProcessor()
processor.precompute_hash_before_cpu_transfer = enabled
processor.use_cuda_ipc = False
@@ -409,9 +431,10 @@ class TestMultimodalProcessorConcurrency(unittest.IsolatedAsyncioTestCase):
MultimodalProcessorExecutor,
)
with patch.object(
BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
with (
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
):
processor = BaseMultimodalProcessor()
processor.mm_processor_executor = MultimodalProcessorExecutor(
@@ -438,9 +461,10 @@ class TestMultimodalProcessorConcurrency(unittest.IsolatedAsyncioTestCase):
BaseMultimodalProcessor,
)
with patch.object(
BaseMultimodalProcessor, "__abstractmethods__", set()
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
with (
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
):
processor = BaseMultimodalProcessor()
processor.mm_processor_executor = None
+558 -19
View File
@@ -1,7 +1,14 @@
"""CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring."""
import asyncio
import base64
import functools
import io
import pickle
import tempfile
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, patch
@@ -18,11 +25,26 @@ from sglang.srt.managers.schedule_batch import (
MultimodalInputs,
MultimodalProcessorOutput,
)
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
from sglang.srt.models.kimi_k25 import (
KimiK25ForConditionalGeneration,
mm_projection_auto,
)
from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger
from sglang.srt.multimodal.cache import (
MultimodalPreprocessCache,
resolve_multimodal_item_hash,
snapshot_media,
)
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
KimiK3DeferredPreprocessing,
)
from sglang.srt.multimodal.media_artifacts.kimi_k3 import (
KimiK3ImagePreprocessArtifact,
KimiK3PreprocessConfig,
KimiK3ResizeConfig,
)
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
@@ -45,6 +67,8 @@ from sglang.srt.multimodal.transport.cuda_ipc import (
CudaIpcTensorTransportProxy,
)
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import ImageData
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
@@ -320,8 +344,9 @@ def test_dp_helper_supports_moonvit3d_packed_embeddings_on_tp1():
# The IPC consumer count asks for the *configured* TP size (matching
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
with get_context().override_server_args(tp_size=1), get_parallel().override(
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
with (
get_context().override_server_args(tp_size=1),
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
):
output = run_dp_sharded_mrope_vision_model(
tower, pixel_values, [[1, 2, 2]], rope_type="rope_2d_packed"
@@ -338,8 +363,9 @@ def test_dp_helper_can_lazily_load_kimi_features_on_tp1():
# The IPC consumer count asks for the *configured* TP size (matching
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
with get_context().override_server_args(tp_size=1), get_parallel().override(
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
with (
get_context().override_server_args(tp_size=1),
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
):
output = run_dp_sharded_mrope_vision_model(
tower,
@@ -490,8 +516,9 @@ def test_kimi_non_dp_keeps_grid_thws_on_the_host():
# The IPC consumer count asks for the *configured* TP size (matching
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
with get_context().override_server_args(tp_size=1), get_parallel().override(
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
with (
get_context().override_server_args(tp_size=1),
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
):
model.get_image_feature(items)
@@ -582,6 +609,30 @@ class _HFProcessor:
)
class _AnySizeTokenizer:
def encode(self, text, allowed_special=None):
if text.startswith("<|media_begin|>image "):
return [10, 11]
if text == "<|media_end|>":
return [14]
return []
def _k3_preprocess_config(
*, patch_size=14, in_patch_limit=16384
) -> KimiK3PreprocessConfig:
return KimiK3PreprocessConfig(
patch_size=patch_size,
merge_kernel_size=2,
in_patch_limit=in_patch_limit,
patch_limit_on_one_side=512,
fixed_output_tokens=None,
image_mean=(0.5, 0.5, 0.5),
image_std=(0.5, 0.5, 0.5),
transparent_bg_config=None,
)
@pytest.mark.parametrize(
("processor_cls", "wrapper_cls"),
[
@@ -592,13 +643,17 @@ class _HFProcessor:
def test_kimi_processor_workers_clone_the_gpu_wrapper(processor_cls, wrapper_cls):
server_args = SimpleNamespace(
mm_feature_transport="cpu",
image_processor_backend="auto",
disable_fast_image_processor=False,
skip_tokenizer_init=False,
mm_process_config={},
mm_io_worker_num=0,
mm_processor_worker_num=0,
tokenizer_worker_num=1,
mm_preprocess_cache_size_mb=0,
trust_mm_content_hashes=False,
base_gpu_id=0,
rl_on_policy_target=None,
allowed_media_domains=[],
media_url_max_file_size_mb=64,
)
@@ -615,6 +670,12 @@ def test_kimi_processor_workers_clone_the_gpu_wrapper(processor_cls, wrapper_cls
assert isinstance(processor._processor, wrapper_cls)
assert isinstance(worker_processor, wrapper_cls)
assert worker_processor is not processor._processor
if processor_cls is KimiK3ImageProcessor:
fingerprint_config = processor.preprocess_fingerprint_payload()[
"wrapped_processor"
]
assert isinstance(fingerprint_config, KimiK3PreprocessConfig)
assert fingerprint_config.patch_size == 14
finally:
processor.mm_processor_executor.shutdown()
processor.io_executor.shutdown()
@@ -690,18 +751,498 @@ def test_kimi_k3_epd_rebuild_uses_the_same_media_contract():
)
def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
KimiK3DeferredPreprocessing,
def _cached_k3_artifact(content_digest, artifact_key, value=1):
return KimiK3ImagePreprocessArtifact(
content_digest=content_digest,
artifact_key=artifact_key,
feature_hash=123,
original_size=(1536, 1024),
resize_config=KimiK3ResizeConfig(
num_tokens=3,
new_width=6,
new_height=2,
pad_width=0,
pad_height=0,
),
grid_thw=(1, 2, 6),
feature=torch.full((12, 2), value, dtype=torch.float32),
)
def test_kimi_k3_cached_artifact_is_composed_per_prompt():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_tokens = SimpleNamespace(image_token_id=99)
processor._tokenizer = _Tokenizer()
processor.mm_feature_transport = "cpu"
processor.use_cuda_ipc = False
artifact = _cached_k3_artifact("sha256:" + "ab" * 32, "artifact")
first = processor.compose_request([1, 99, 2], [artifact])
second = processor.compose_request([3, 4, 99, 5], [artifact])
assert first.input_ids != second.input_ids
assert first.mm_items[0].offsets == [(3, 5)]
assert second.mm_items[0].offsets == [(4, 6)]
assert first.mm_items[0].hash == second.mm_items[0].hash == 123
torch.testing.assert_close(first.mm_items[0].feature, second.mm_items[0].feature)
def test_kimi_k3_cached_deferred_artifact_has_model_contract():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_feature_transport = "cpu"
feature = torch.zeros((3, 2, 2), dtype=torch.uint8)
artifact = processor._make_artifact(
content_digest="sha256:" + "ab" * 32,
artifact_key="sha256:" + "cd" * 32,
original_size=(2, 2),
resize_config={
"num_tokens": 1,
"new_width": 2,
"new_height": 2,
"pad_width": 0,
"pad_height": 0,
},
grid_thw=(1, 1, 1),
feature=feature,
deferred=KimiK3DeferredPreprocessing(
backend="gpu",
image_mean=[0.5, 0.5, 0.5],
image_std=[0.5, 0.5, 0.5],
transparent_bg_config=None,
resize_config={
"num_tokens": 1,
"new_width": 2,
"new_height": 2,
"pad_width": 0,
"pad_height": 0,
},
),
)
config = artifact.deferred
assert config.backend == "gpu"
assert config.resize_config["new_width"] == 2
def test_kimi_k3_normal_cache_path_connects_real_producer_to_model_consumer():
hf_processor = _HFProcessor()
hf_processor.tokenizer = _AnySizeTokenizer()
hf_config = SimpleNamespace(
media_placeholder_token_id=42,
to_dict=lambda: {
"model_type": "kimi_k3",
"architectures": ["KimiK3ForConditionalGeneration"],
},
)
server_args = ServerArgs(
model_path="dummy",
mm_feature_transport="cpu",
mm_process_config={},
mm_io_worker_num=1,
mm_processor_worker_num=0,
tokenizer_worker_num=1,
mm_preprocess_cache_size_mb=1,
)
processor = KimiK3ImageProcessor(
hf_config=hf_config,
server_args=server_args,
_processor=hf_processor,
transport_mode=None,
)
image = Image.new("RGB", (28, 28), color=(1, 2, 3))
encoded_image = io.BytesIO()
image.save(encoded_image, format="PNG")
image_data = ImageData(
url="data:image/png;base64,"
+ base64.b64encode(encoded_image.getvalue()).decode()
)
request = SimpleNamespace(video_data=None, mm_content_hashes=None)
class _Tower(nn.Module):
device = torch.device("cpu")
patch_size = 14
def __init__(self):
super().__init__()
self.patch_embed = SimpleNamespace(
proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32))
)
def forward(self, pixel_values, _grid_thws):
return pixel_values
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
nn.Module.__init__(model)
model.use_data_parallel = False
model.vision_tower = _Tower()
model.mm_projector = _Projector()
try:
with (
patch(
"sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True
),
patch.object(
processor,
"prepare_artifact_batch",
wraps=processor.prepare_artifact_batch,
) as prepare_artifacts,
):
cold = asyncio.run(
processor.process_mm_data_async([image_data], [1, 42, 2], request)
)
hot = asyncio.run(
processor.process_mm_data_async([image_data], [3, 42, 4], request)
)
cold_items = pickle.loads(pickle.dumps(cold.mm_items))
hot_items = pickle.loads(pickle.dumps(hot.mm_items))
with (
patch("sglang.srt.models.kimi_k3.configured_tp_size", return_value=1),
patch(
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
return_value=(
torch.ones((4, 3), dtype=torch.float32),
torch.tensor([[1, 2, 2]], dtype=torch.int64),
),
),
):
cold_features = model.get_image_feature(cold_items)
hot_features = model.get_image_feature(hot_items)
finally:
processor.shutdown()
assert prepare_artifacts.call_count == 1
assert cold.mm_items[0].hash == hot.mm_items[0].hash
assert cold.mm_items[0].offsets == hot.mm_items[0].offsets == [(3, 3)]
assert (
cold_items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY].backend == "gpu"
)
torch.testing.assert_close(cold_features, hot_features)
def test_kimi_k3_model_accepts_mixed_cached_eager_and_deferred_artifacts():
class _Tower(nn.Module):
device = torch.device("cpu")
patch_size = 2
def __init__(self):
super().__init__()
self.patch_embed = SimpleNamespace(
proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32))
)
def forward(self, pixel_values, _grid_thws):
return pixel_values
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
nn.Module.__init__(model)
model.use_data_parallel = False
model.vision_tower = _Tower()
model.mm_projector = _Projector()
eager = _image_item(torch.ones((1, 3)), [[1, 1, 1]])
deferred = _image_item(torch.zeros((3, 2, 2), dtype=torch.uint8), [[1, 1, 1]])
deferred.model_specific_data[DEFERRED_PREPROCESSING_KEY] = (
KimiK3DeferredPreprocessing(
backend="gpu",
image_mean=[0.5, 0.5, 0.5],
image_std=[0.5, 0.5, 0.5],
transparent_bg_config=None,
resize_config={
"num_tokens": 1,
"new_width": 2,
"new_height": 2,
"pad_width": 0,
"pad_height": 0,
},
)
)
with (
patch("sglang.srt.models.kimi_k3.configured_tp_size", return_value=1),
patch(
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
return_value=(torch.full((1, 3), 2.0), torch.tensor([[1, 1, 1]])),
),
):
output = model.get_image_feature([eager, deferred])
torch.testing.assert_close(output, torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]))
def test_kimi_k3_trusted_hot_hit_skips_media_read():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
processor.trust_mm_content_hashes = True
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
processor.io_executor = ThreadPoolExecutor(max_workers=1)
digest = "sha256:" + "ab" * 32
key = processor._artifact_key(digest, "unread-source")
artifact = _cached_k3_artifact(digest, key)
processor.mm_preprocess_cache.put(key, artifact)
try:
with patch(
"sglang.srt.multimodal.media_artifacts.base.snapshot_media",
side_effect=AssertionError("trusted cache hit must not read media"),
):
result = asyncio.run(
processor.prepare_media_artifacts(
["unread-source"],
content_hashes=[digest],
)
)
finally:
processor.io_executor.shutdown()
assert result == [artifact]
def test_kimi_k3_default_media_options_share_one_artifact_key():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
digest = "sha256:" + "ab" * 32
keys = {
processor._artifact_key(digest, "image.png"),
processor._artifact_key(digest, ImageData(url="image.png")),
processor._artifact_key(digest, {"url": "image.png", "detail": "auto"}),
}
assert len(keys) == 1
def test_kimi_k3_output_affecting_media_options_do_not_share_artifacts():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
digest = "sha256:" + "ab" * 32
base = processor._artifact_key(digest, ImageData(url="image.png"))
assert base != processor._artifact_key(
digest, ImageData(url="image.png", detail="low")
)
assert base != processor._artifact_key(
digest, ImageData(url="image.png", max_dynamic_patch=4)
)
assert base != processor._artifact_key(
digest,
ImageData(url="image.png", preprocess_kwargs={"max_pixels": 1024}),
)
assert base != processor._artifact_key(
digest,
{"url": "image.png", "future_model_option": "new-behavior"},
)
def test_kimi_k3_rejects_changed_feature_hash_for_same_artifact():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
processor.trust_mm_content_hashes = False
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
processor.mm_processor_executor = None
processor.io_executor = ThreadPoolExecutor(max_workers=2)
image = Image.new("RGB", (2, 2), color=(1, 2, 3))
digest = snapshot_media(image).content_digest
key = processor._artifact_key(digest, image)
old = replace(_cached_k3_artifact(digest, key), feature=None)
new = replace(_cached_k3_artifact(digest, key), feature_hash=old.feature_hash + 1)
processor.mm_preprocess_cache.put(key, old)
async def prepare(_entries):
return [new]
processor._run_preprocess_and_build_artifact_batch = prepare
try:
with pytest.raises(ValueError, match="feature hash changed"):
asyncio.run(processor.prepare_media_artifacts([image]))
finally:
processor.io_executor.shutdown()
assert key not in processor.mm_preprocess_cache
def test_kimi_k3_artifact_and_data_item_share_hash_resolution():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_feature_transport = "cpu"
processor.mm_tokens = SimpleNamespace(image_token_id=99)
processor._tokenizer = _Tokenizer()
processor.use_cuda_ipc = False
feature = torch.zeros((4, 3), dtype=torch.float32)
digest = "sha256:" + "ab" * 32
artifact = processor._make_artifact(
content_digest=digest,
artifact_key="sha256:" + "01" * 32,
original_size=(2, 2),
resize_config={
"num_tokens": 1,
"new_width": 2,
"new_height": 2,
"pad_width": 0,
"pad_height": 0,
},
grid_thw=(1, 1, 1),
feature=feature,
)
direct_item = MultimodalDataItem(modality=Modality.IMAGE, feature=feature)
direct_item.set_pad_value()
composed_item = processor.compose_request([1, 99, 2], [artifact]).mm_items[0]
expected_hash = resolve_multimodal_item_hash(
existing_hash=direct_item.hash,
namespace=artifact.artifact_key,
)
expected_item = MultimodalDataItem(modality=Modality.IMAGE, hash=expected_hash)
expected_item.set_pad_value()
assert artifact.feature_hash == composed_item.hash == expected_hash
assert composed_item.pad_value == expected_item.pad_value
def test_kimi_k3_untrusted_path_change_is_a_cache_miss():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
processor.trust_mm_content_hashes = False
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
processor.mm_processor_executor = None
processor.io_executor = ThreadPoolExecutor(max_workers=2)
async def prepare(entries):
return [
_cached_k3_artifact(
entry.content_digest,
entry.artifact_key,
entry.media.getpixel((0, 0))[0],
)
for entry in entries
]
processor._run_preprocess_and_build_artifact_batch = prepare
try:
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "mutable.png"
Image.new("RGB", (2, 2), color=(1, 0, 0)).save(path)
first = asyncio.run(processor.prepare_media_artifacts([str(path)]))[0]
Image.new("RGB", (2, 2), color=(2, 0, 0)).save(path)
second = asyncio.run(processor.prepare_media_artifacts([str(path)]))[0]
finally:
processor.io_executor.shutdown()
assert first.content_digest != second.content_digest
assert first.artifact_key != second.artifact_key
assert first.feature[0, 0].item() == 1
assert second.feature[0, 0].item() == 2
def test_kimi_k3_partial_hits_deduplicate_misses_and_preserve_order():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
processor.trust_mm_content_hashes = False
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
processor.mm_processor_executor = None
processor.io_executor = ThreadPoolExecutor(max_workers=4)
cached_image = Image.new("RGB", (2, 2), color=(1, 0, 0))
missed_image = Image.new("RGB", (2, 2), color=(2, 0, 0))
cached_digest = snapshot_media(cached_image).content_digest
missed_digest = snapshot_media(missed_image).content_digest
cached_key = processor._artifact_key(cached_digest, cached_image)
cached = _cached_k3_artifact(cached_digest, cached_key, value=1)
processor.mm_preprocess_cache.put(cached_key, cached)
batches = []
async def prepare(entries):
batches.append(entries)
return [
replace(
_cached_k3_artifact(
entry.content_digest,
entry.artifact_key,
entry.media.getpixel((0, 0))[0],
),
feature_hash=456,
)
for entry in entries
]
processor._run_preprocess_and_build_artifact_batch = prepare
try:
artifacts = asyncio.run(
processor.prepare_media_artifacts(
[cached_image, missed_image, missed_image, cached_image],
)
)
finally:
processor.io_executor.shutdown()
assert len(batches) == 1
assert len(batches[0]) == 1
assert batches[0][0].content_digest == missed_digest
assert batches[0][0].artifact_key == processor._artifact_key(
missed_digest, missed_image
)
assert [artifact.content_digest for artifact in artifacts] == [
cached_digest,
missed_digest,
missed_digest,
cached_digest,
]
assert artifacts[0] is artifacts[3]
assert artifacts[1] is artifacts[2]
def test_kimi_k3_cancelled_artifact_owner_does_not_fail_joiner():
processor = object.__new__(KimiK3ImageProcessor)
processor.processor_fingerprint = "processor"
processor.trust_mm_content_hashes = False
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
processor.mm_processor_executor = None
processor.io_executor = ThreadPoolExecutor(max_workers=2)
processor._preprocess_metrics_callback = None
image = Image.new("RGB", (2, 2), color=(1, 2, 3))
started = asyncio.Event()
release = asyncio.Event()
async def prepare(entries):
started.set()
await release.wait()
return [
_cached_k3_artifact(
entry.content_digest,
entry.artifact_key,
entry.media.getpixel((0, 0))[0],
)
for entry in entries
]
processor._run_preprocess_and_build_artifact_batch = prepare
async def run():
owner = asyncio.create_task(processor.prepare_media_artifacts([image]))
await started.wait()
joiner = asyncio.create_task(processor.prepare_media_artifacts([image]))
await asyncio.sleep(0)
owner.cancel()
with pytest.raises(asyncio.CancelledError):
await owner
release.set()
artifacts = await joiner
assert len(artifacts) == 1
assert artifacts[0].feature[0, 0].item() == 1
try:
asyncio.run(run())
finally:
processor.io_executor.shutdown()
def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_tokens = SimpleNamespace(image_token_id=99)
processor.mm_feature_transport = "cpu"
processor.use_cuda_ipc = False
processor._processor = SimpleNamespace(
_patch_size=2,
preprocess_config=_k3_preprocess_config(patch_size=2),
prepare_deferred=Mock(
return_value=(
torch.tensor([[1, 99, 99, 2, 99, 3]]),
@@ -776,11 +1317,7 @@ def test_kimi_k3_defers_only_when_raw_transport_is_smaller(
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_feature_transport = "cpu"
processor._processor = SimpleNamespace(
_patch_size=14,
_merge_kernel_size=2,
_in_patch_limit=in_patch_limit,
_patch_limit_on_one_side=512,
_fixed_output_tokens=None,
preprocess_config=_k3_preprocess_config(in_patch_limit=in_patch_limit),
)
image = torch.zeros(image_shape, dtype=torch.uint8)
@@ -817,7 +1354,7 @@ def test_kimi_k3_eager_preprocessing_preserves_float_tensor_support():
assert output.shape == (3, 4, 4)
@pytest.mark.parametrize("transport", ["cuda_ipc", "fabric"])
@pytest.mark.parametrize("transport", ["cuda_ipc", "cuda_vmm"])
def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_feature_transport = transport
@@ -831,6 +1368,7 @@ def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
def test_kimi_k3_rejects_silently_dropped_images():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_tokens = Mock()
processor.mm_preprocess_cache = MultimodalPreprocessCache(0)
processor.load_mm_data = AsyncMock(return_value=SimpleNamespace(images=[object()]))
with pytest.raises(ValueError, match="expected 2, loaded 1"):
@@ -838,13 +1376,14 @@ def test_kimi_k3_rejects_silently_dropped_images():
processor.process_mm_data_async(
image_data=["image-1", "image-2"],
input_text="<|media_pad|><|media_pad|>",
request_obj=SimpleNamespace(video_data=None),
request_obj=SimpleNamespace(video_data=None, mm_content_hashes=None),
)
)
def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
processor = object.__new__(KimiK3ImageProcessor)
processor.mm_preprocess_cache = MultimodalPreprocessCache(0)
processor.mm_feature_transport = "cpu"
processor.mm_tokens = SimpleNamespace(image_token_id=99)
processor.mm_feature_transport = "cuda_ipc"
@@ -863,7 +1402,7 @@ def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
processor.process_mm_data_async(
image_data=["image-1", "image-2"],
input_text=[1, 99, 2, 99, 3],
request_obj=SimpleNamespace(video_data=None),
request_obj=SimpleNamespace(video_data=None, mm_content_hashes=None),
)
)
@@ -69,8 +69,11 @@ def make_processor(config, image_processor_cls=None):
model_impl="sglang",
keep_mm_feature_on_device=False,
mm_feature_transport="cpu",
image_processor_backend="auto",
disable_fast_image_processor=True,
skip_tokenizer_init=False,
mm_preprocess_cache_size_mb=0,
trust_mm_content_hashes=False,
# Read by NativeMmHost._use_feature_shm (single-rank fixture → the
# inline zero-copy transport, like the 1-GPU e2e).
tp_size=1,
@@ -80,6 +83,7 @@ def make_processor(config, image_processor_cls=None):
mm_processor_worker_num=1,
tokenizer_worker_num=1,
base_gpu_id=0,
rl_on_policy_target=None,
allowed_media_domains=[],
media_url_max_file_size_mb=64,
)
@@ -0,0 +1,238 @@
import asyncio
import unittest
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from typing import Optional
from sglang.srt.managers.schedule_batch import Modality
from sglang.srt.multimodal.cache import MultimodalPreprocessCache, snapshot_media
from sglang.srt.multimodal.media_artifacts import (
MediaArtifactCacheMixin,
MediaArtifactInput,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@dataclass(frozen=True)
class _Artifact:
content_digest: str
artifact_key: str
feature_hash: int
feature: Optional[bytes]
@property
def has_feature(self) -> bool:
return self.feature is not None
def cache_value(self):
return self
def cache_size_items(self):
return (
self.content_digest,
self.artifact_key,
self.feature_hash,
self.feature,
)
@dataclass(frozen=True)
class _FutureMediaInput:
url: str
content_hash: Optional[str] = None
frame_sampling: int = 2
class _Processor(MediaArtifactCacheMixin):
artifact_modality = Modality.IMAGE
artifact_option_defaults = {"detail": "auto", "frame_sampling": 2}
def __init__(self):
self.processor_fingerprint = "processor"
self.trust_mm_content_hashes = False
self.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
self.mm_processor_executor = None
self.io_executor = ThreadPoolExecutor(max_workers=4)
self.batches = []
def decode_media_snapshot(self, snapshot, modality):
self.assert_artifact_modality(modality)
return snapshot.data
@staticmethod
def assert_artifact_modality(modality):
if modality != Modality.IMAGE:
raise AssertionError(f"unexpected modality: {modality}")
def prepare_artifact_batch(
self, entries: list[MediaArtifactInput]
) -> list[_Artifact]:
self.batches.append(entries)
return [
_Artifact(
content_digest=entry.content_digest,
artifact_key=entry.artifact_key,
feature_hash=int(entry.content_digest[-16:], 16),
feature=entry.media,
)
for entry in entries
]
def close(self):
self.io_executor.shutdown()
class TestMediaArtifactProcessor(unittest.TestCase):
def test_unknown_model_option_is_part_of_artifact_identity(self):
processor = _Processor()
digest = snapshot_media(b"image").content_digest
try:
base = processor._artifact_key(digest, _FutureMediaInput(url="image.png"))
self.assertEqual(
base,
processor._artifact_key(
digest,
{
"url": "image.png",
"content_hash": digest,
"frame_sampling": 2,
},
),
)
self.assertNotEqual(
base,
processor._artifact_key(
digest,
{"url": "image.png", "future_model_knob": "different"},
),
)
self.assertNotEqual(
base,
processor._artifact_key(
digest,
_FutureMediaInput(url="image.png"),
modality=Modality.VIDEO,
),
)
finally:
processor.close()
def test_non_image_models_can_override_identity_and_decode_hooks(self):
class _VideoProcessor(_Processor):
artifact_modality = Modality.VIDEO
def snapshot_media_source(self, source, modality):
self.assert_video_modality(modality)
return snapshot_media(source.encode())
def decode_media_snapshot(self, snapshot, modality):
self.assert_video_modality(modality)
return snapshot.data
@staticmethod
def assert_video_modality(modality):
if modality != Modality.VIDEO:
raise AssertionError(f"unexpected modality: {modality}")
processor = _VideoProcessor()
try:
artifacts = asyncio.run(processor.prepare_media_artifacts(["clip.mp4"]))
finally:
processor.close()
self.assertEqual(len(artifacts), 1)
self.assertEqual(
artifacts[0].content_digest, snapshot_media(b"clip.mp4").content_digest
)
def test_partial_hits_and_duplicate_misses_are_shared_by_contract(self):
processor = _Processor()
first_digest = snapshot_media(b"first").content_digest
first_key = processor._artifact_key(first_digest, b"first")
first = _Artifact(first_digest, first_key, 1, b"first")
processor.mm_preprocess_cache.put(first_key, first)
try:
artifacts = asyncio.run(
processor.prepare_media_artifacts(
[b"first", b"second", b"second", b"first"]
)
)
finally:
processor.close()
self.assertEqual(len(processor.batches), 1)
self.assertEqual(len(processor.batches[0]), 1)
self.assertEqual(
[artifact.content_digest for artifact in artifacts],
[
first_digest,
snapshot_media(b"second").content_digest,
snapshot_media(b"second").content_digest,
first_digest,
],
)
self.assertIs(artifacts[0], artifacts[3])
self.assertIs(artifacts[1], artifacts[2])
def test_adapter_cannot_change_validated_artifact_identity(self):
processor = _Processor()
async def wrong_identity(entries):
artifact = processor.prepare_artifact_batch(entries)[0]
return [replace(artifact, artifact_key="sha256:" + "0" * 64)]
processor._run_preprocess_and_build_artifact_batch = wrong_identity
try:
with self.assertRaisesRegex(ValueError, "changed the media artifact key"):
asyncio.run(processor.prepare_media_artifacts([b"image"]))
finally:
processor.close()
def test_trusted_hit_uses_identity_without_reading_source(self):
processor = _Processor()
processor.trust_mm_content_hashes = True
digest = snapshot_media(b"cached").content_digest
key = processor._artifact_key(digest, "unread-source")
artifact = _Artifact(digest, key, 1, b"cached")
processor.mm_preprocess_cache.put(key, artifact)
try:
artifacts = asyncio.run(
processor.prepare_media_artifacts(
["unread-source"], content_hashes=[digest]
)
)
finally:
processor.close()
self.assertEqual(artifacts, [artifact])
self.assertEqual(processor.batches, [])
def test_cached_artifact_must_match_content_identity(self):
processor = _Processor()
digest = snapshot_media(b"fresh").content_digest
key = processor._artifact_key(digest, b"fresh")
processor.mm_preprocess_cache.put(
key,
_Artifact(
snapshot_media(b"stale").content_digest,
key,
1,
b"stale",
),
)
try:
artifacts = asyncio.run(processor.prepare_media_artifacts([b"fresh"]))
finally:
processor.close()
self.assertEqual(artifacts[0].content_digest, digest)
self.assertEqual(artifacts[0].feature, b"fresh")
self.assertEqual(len(processor.batches), 1)
if __name__ == "__main__":
unittest.main()
@@ -4,22 +4,24 @@ 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.managers.schedule_batch import Modality, MultimodalDataItem
from sglang.srt.multimodal.cache import (
CacheMiss,
MultimodalPreprocessCache,
build_artifact_key,
build_feature_hash,
build_processor_fingerprint,
estimate_cache_size_bytes,
parse_content_hash,
resolve_multimodal_item_hash,
snapshot_media,
)
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
@@ -209,21 +211,25 @@ class TestMediaIdentity(unittest.TestCase):
def preprocess_fingerprint_payload(self):
return {"backend": self.backend, "antialias": True}
config = SimpleNamespace(model_type="vlm", architectures=["VLM"])
args = SimpleNamespace(
class Config:
def to_dict(self):
return {"model_type": "vlm", "architectures": ["VLM"]}
config = Config()
args = ServerArgs(
model_path="dummy",
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_args = ServerArgs(
model_path="dummy",
revision="model-revision",
disable_fast_image_processor=False,
mm_process_config={"image": {"max_pixels": 2048}},
)
changed_config = build_processor_fingerprint(
Processor("gpu"), config, changed_args
@@ -231,7 +237,7 @@ class TestMediaIdentity(unittest.TestCase):
self.assertNotEqual(base, changed_backend)
self.assertNotEqual(base, changed_config)
def test_feature_hash_includes_artifact_and_processor_output(self):
def test_item_hash_namespace_covers_identity_and_processor_output(self):
digest = snapshot_media(b"image").content_digest
first = build_artifact_key(
digest,
@@ -243,11 +249,25 @@ class TestMediaIdentity(unittest.TestCase):
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)
self.assertNotEqual(
resolve_multimodal_item_hash(existing_hash=1, namespace=first),
resolve_multimodal_item_hash(existing_hash=1, namespace=second),
)
self.assertNotEqual(
resolve_multimodal_item_hash(existing_hash=1, namespace=first),
resolve_multimodal_item_hash(existing_hash=2, namespace=first),
)
with self.assertRaises(ValueError):
build_feature_hash(first, -1)
resolve_multimodal_item_hash(existing_hash=-1, namespace=first)
def test_multimodal_data_item_uses_shared_feature_hash(self):
feature = torch.arange(12, dtype=torch.float32).reshape(4, 3)
expected = resolve_multimodal_item_hash(feature=feature)
item = MultimodalDataItem(modality=Modality.IMAGE, feature=feature)
item.set_pad_value()
self.assertEqual(item.hash, expected)
class TestMultimodalPreprocessCache(unittest.TestCase):
@@ -262,6 +282,31 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
self.assertIn("c", cache)
self.assertEqual(cache.current_size_bytes, 6)
def test_compatible_lookup_is_atomic_and_does_not_count_bypass_as_miss(self):
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
cache.put("key", b"metadata-only")
self.assertIsNone(cache.get_if_present("key", lambda value: False))
self.assertEqual((cache.hits, cache.misses), (0, 0))
self.assertEqual(
cache.get_if_present("key", lambda value: value.startswith(b"metadata")),
b"metadata-only",
)
self.assertEqual((cache.hits, cache.misses), (1, 0))
def test_claimed_miss_rejects_an_incompatible_racing_entry(self):
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
cache.put("key", b"metadata-only")
miss = cache.lookup_or_claim_many(
["key"], predicate=lambda key, value: value == b"full-feature"
)[0]
self.assertIsInstance(miss, CacheMiss)
self.assertTrue(miss.should_compute)
self.assertNotIn("key", cache)
self.assertEqual(cache.current_size_bytes, 0)
def test_gpu_backed_values_are_not_implicitly_copied(self):
if not torch.cuda.is_available():
self.skipTest("CUDA is not available")
@@ -370,6 +415,61 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
asyncio.run(run())
def test_lookup_or_claim_many_batches_owned_and_joined_misses(self):
async def run():
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
results = cache.lookup_or_claim_many(["a", "b", "a"])
misses_to_compute = [
item
for item in results
if isinstance(item, CacheMiss) and item.should_compute
]
self.assertEqual([item.key for item in misses_to_compute], ["a", "b"])
cache.complete_miss(misses_to_compute[0], b"value-a")
cache.complete_miss(misses_to_compute[1], b"value-b")
self.assertEqual(await cache.wait_for_miss(results[2]), b"value-a")
self.assertEqual(cache.get("b"), b"value-b")
asyncio.run(run())
def test_cancelled_miss_waiter_does_not_cancel_computing_caller(self):
async def run():
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
computing_miss = cache.lookup_or_claim_many(["key"])[0]
waiting_miss = cache.lookup_or_claim_many(["key"])[0]
self.assertTrue(computing_miss.should_compute)
self.assertFalse(waiting_miss.should_compute)
waiter = asyncio.create_task(cache.wait_for_miss(waiting_miss))
await asyncio.sleep(0)
waiter.cancel()
with self.assertRaises(asyncio.CancelledError):
await waiter
cache.complete_miss(computing_miss, b"artifact")
self.assertEqual(computing_miss.future.result(), b"artifact")
self.assertEqual(cache.get("key"), b"artifact")
asyncio.run(run())
def test_disabled_cache_does_not_join_or_retain(self):
async def run():
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=0)
misses = cache.lookup_or_claim_many(["a", "a"])
self.assertTrue(
all(
isinstance(item, CacheMiss) and item.should_compute
for item in misses
)
)
for item in misses:
cache.complete_miss(item, b"value")
self.assertEqual(len(cache), 0)
self.assertEqual(cache.stats()["singleflight_joins"], 0)
asyncio.run(run())
def test_clear_starts_a_new_singleflight_generation(self):
async def run():
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
@@ -397,6 +497,20 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
asyncio.run(run())
def test_clear_starts_a_new_cache_miss_generation(self):
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
old = cache.lookup_or_claim_many(["key"])[0]
cache.clear()
new = cache.lookup_or_claim_many(["key"])[0]
self.assertTrue(old.should_compute)
self.assertTrue(new.should_compute)
self.assertIsNot(old.future, new.future)
cache.complete_miss(old, b"old")
self.assertNotIn("key", cache)
cache.complete_miss(new, b"new")
self.assertEqual(cache.get("key"), b"new")
if __name__ == "__main__":
unittest.main()