wehub-resource-sync
530 lines
18 KiB
Python
530 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""
|
|
Unit tests for prompt-embedding cache for diffusion pipelines.
|
|
|
|
Covers the module :mod:`vllm_omni.diffusion.cache.prompt_embed_cache`:
|
|
|
|
* :class:`PromptEmbedCache` LRU semantics, stats and thread-safety.
|
|
* :func:`_hashable` safe-input classification (scalars, torch dtype/device,
|
|
numpy scalars, nested containers, tensors/PIL bypass).
|
|
* :func:`install_prompt_embed_cache` / :func:`uninstall_prompt_embed_cache`
|
|
behaviour: caching, bypass on tensors, bypass on precomputed embeds,
|
|
idempotent install, restore on uninstall, detachment of cached tensors.
|
|
* :func:`resolve_prompt_embed_cache_config` env-var overrides.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.cache.prompt_embed_cache import (
|
|
_CACHE_MISS,
|
|
_NOT_HASHABLE,
|
|
PromptEmbedCache,
|
|
_build_key,
|
|
_detach_output,
|
|
_hashable,
|
|
install_prompt_embed_cache,
|
|
resolve_prompt_embed_cache_config,
|
|
uninstall_prompt_embed_cache,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _hashable
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestHashable:
|
|
"""Tests for the ``_hashable`` helper."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[None, "prompt", "", 0, 1, -5, 3.14, True, False, b"bytes"],
|
|
)
|
|
def test_scalars_are_returned_unchanged(self, value):
|
|
assert _hashable(value) == value
|
|
|
|
def test_torch_device_is_normalized(self):
|
|
out = _hashable(torch.device("cpu"))
|
|
assert out == ("__torch_device__", "cpu")
|
|
assert hash(out) is not None # must actually be hashable
|
|
|
|
def test_torch_dtype_is_normalized(self):
|
|
out = _hashable(torch.float32)
|
|
assert out[0] == "__torch_dtype__"
|
|
assert "float32" in out[1]
|
|
assert hash(out) is not None
|
|
|
|
def test_numpy_scalar_is_unwrapped(self):
|
|
np = pytest.importorskip("numpy")
|
|
assert _hashable(np.int64(7)) == 7
|
|
assert _hashable(np.float32(1.5)) == pytest.approx(1.5)
|
|
|
|
def test_numpy_ndarray_is_not_hashable(self):
|
|
np = pytest.importorskip("numpy")
|
|
assert _hashable(np.array([1, 2, 3])) is _NOT_HASHABLE
|
|
|
|
def test_nested_list_and_tuple(self):
|
|
out = _hashable(["a", ("b", 1), [2, 3]])
|
|
assert out[0] == "__seq__"
|
|
# fully hashable
|
|
assert hash(out) is not None
|
|
|
|
def test_dict_is_normalized_by_sorted_keys(self):
|
|
a = _hashable({"b": 1, "a": 2})
|
|
b = _hashable({"a": 2, "b": 1})
|
|
assert a == b
|
|
assert hash(a) is not None
|
|
|
|
def test_tensor_is_not_hashable(self):
|
|
assert _hashable(torch.zeros(2)) is _NOT_HASHABLE
|
|
|
|
def test_list_containing_tensor_is_not_hashable(self):
|
|
assert _hashable(["a", torch.zeros(1)]) is _NOT_HASHABLE
|
|
|
|
def test_dict_containing_tensor_is_not_hashable(self):
|
|
assert _hashable({"x": torch.zeros(1)}) is _NOT_HASHABLE
|
|
|
|
def test_custom_object_is_not_hashable(self):
|
|
class Obj:
|
|
pass
|
|
|
|
assert _hashable(Obj()) is _NOT_HASHABLE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _detach_output
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDetachOutput:
|
|
def test_detaches_tensor(self):
|
|
t = torch.zeros(2, requires_grad=True)
|
|
out = _detach_output(t)
|
|
assert isinstance(out, torch.Tensor)
|
|
assert not out.requires_grad
|
|
|
|
def test_detaches_inside_tuple(self):
|
|
t = torch.ones(1, requires_grad=True)
|
|
out = _detach_output((t, None))
|
|
assert isinstance(out, tuple)
|
|
assert not out[0].requires_grad
|
|
assert out[1] is None
|
|
|
|
def test_detaches_inside_list_and_dict(self):
|
|
t = torch.ones(1, requires_grad=True)
|
|
out = _detach_output({"a": [t]})
|
|
assert not out["a"][0].requires_grad
|
|
|
|
def test_preserves_namedtuple_type(self):
|
|
from collections import namedtuple
|
|
|
|
NT = namedtuple("NT", ["x", "y"])
|
|
t = torch.ones(1, requires_grad=True)
|
|
out = _detach_output(NT(x=t, y=1))
|
|
assert isinstance(out, NT)
|
|
assert not out.x.requires_grad
|
|
assert out.y == 1
|
|
|
|
def test_passthrough_for_non_tensor(self):
|
|
assert _detach_output("hi") == "hi"
|
|
assert _detach_output(None) is None
|
|
assert _detach_output(5) == 5
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PromptEmbedCache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPromptEmbedCache:
|
|
def test_rejects_non_positive_max_size(self):
|
|
with pytest.raises(ValueError):
|
|
PromptEmbedCache(max_size=0)
|
|
with pytest.raises(ValueError):
|
|
PromptEmbedCache(max_size=-1)
|
|
|
|
def test_miss_returns_sentinel_and_bumps_miss_counter(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
assert cache.get("k") is _CACHE_MISS
|
|
assert cache.stats()["misses"] == 1
|
|
assert cache.stats()["hits"] == 0
|
|
|
|
def test_put_then_get_returns_value_and_bumps_hit_counter(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("k", 42)
|
|
assert cache.get("k") == 42
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
def test_none_is_a_cacheable_value(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("k", None)
|
|
# Distinguishable from a miss.
|
|
assert cache.get("k") is None
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
def test_lru_eviction_keeps_recently_used(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("a", 1)
|
|
cache.put("b", 2)
|
|
# Access "a" to mark it MRU.
|
|
assert cache.get("a") == 1
|
|
cache.put("c", 3) # should evict "b"
|
|
assert cache.get("b") is _CACHE_MISS
|
|
assert cache.get("a") == 1
|
|
assert cache.get("c") == 3
|
|
assert cache.stats()["size"] == 2
|
|
|
|
def test_get_refreshes_recency(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("a", 1)
|
|
cache.put("b", 2)
|
|
# Hit on "a" → "b" becomes LRU.
|
|
cache.get("a")
|
|
cache.put("c", 3)
|
|
assert cache.get("b") is _CACHE_MISS
|
|
assert cache.get("a") == 1
|
|
|
|
def test_put_overwrites_existing_key(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("k", 1)
|
|
cache.put("k", 2)
|
|
assert cache.get("k") == 2
|
|
assert cache.stats()["size"] == 1
|
|
|
|
def test_clear_empties_store_but_keeps_counters(self):
|
|
cache = PromptEmbedCache(max_size=2)
|
|
cache.put("k", 1)
|
|
cache.get("k")
|
|
cache.clear()
|
|
assert cache.stats()["size"] == 0
|
|
# clear only drops entries, counters reflect prior activity.
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
def test_stats_reports_expected_fields(self):
|
|
cache = PromptEmbedCache(max_size=4)
|
|
stats = cache.stats()
|
|
assert set(stats.keys()) == {"size", "max_size", "hits", "misses", "bypassed"}
|
|
assert stats["max_size"] == 4
|
|
assert stats["size"] == 0
|
|
|
|
def test_thread_safety_under_concurrent_writes(self):
|
|
cache = PromptEmbedCache(max_size=1024)
|
|
|
|
def worker(start: int) -> None:
|
|
for i in range(start, start + 200):
|
|
cache.put(f"k{i}", i)
|
|
cache.get(f"k{i}")
|
|
|
|
threads = [threading.Thread(target=worker, args=(i * 200,)) for i in range(4)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
stats = cache.stats()
|
|
# 4 * 200 = 800 writes, each followed by a successful read.
|
|
assert stats["hits"] == 800
|
|
assert stats["size"] == 800
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _build_key
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_sig(fn):
|
|
import inspect
|
|
|
|
return inspect.signature(fn)
|
|
|
|
|
|
class TestBuildKey:
|
|
def test_positional_and_keyword_collide(self):
|
|
def encode(prompt, device=None):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
k1 = _build_key(sig, "m", ("hello",), {"device": "cpu"})
|
|
k2 = _build_key(sig, "m", (), {"prompt": "hello", "device": "cpu"})
|
|
assert k1 == k2
|
|
assert k1 is not None
|
|
|
|
def test_different_prompts_produce_different_keys(self):
|
|
def encode(prompt):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
assert _build_key(sig, "m", ("a",), {}) != _build_key(sig, "m", ("b",), {})
|
|
|
|
def test_different_model_tags_produce_different_keys(self):
|
|
def encode(prompt):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
assert _build_key(sig, "m1", ("a",), {}) != _build_key(sig, "m2", ("a",), {})
|
|
|
|
def test_tensor_argument_bypasses(self):
|
|
def encode(prompt, extra=None):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
assert _build_key(sig, "m", ("hi",), {"extra": torch.zeros(1)}) is None
|
|
|
|
def test_precomputed_prompt_embeds_bypasses(self):
|
|
def encode(prompt=None, prompt_embeds=None):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
# None → not a bypass.
|
|
assert _build_key(sig, "m", (), {"prompt": "hi"}) is not None
|
|
# A non-None precomputed embed triggers bypass even if it would
|
|
# otherwise be "hashable".
|
|
assert _build_key(sig, "m", (), {"prompt": "hi", "prompt_embeds": "x"}) is None
|
|
|
|
def test_invalid_binding_returns_none(self):
|
|
def encode(prompt):
|
|
return None
|
|
|
|
sig = _make_sig(encode)
|
|
# Unknown kwarg → TypeError on bind → bypass.
|
|
assert _build_key(sig, "m", (), {"nope": 1}) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# install_prompt_embed_cache / uninstall_prompt_embed_cache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakePipeline:
|
|
"""Minimal pipeline with a signature-bearing ``encode_prompt``."""
|
|
|
|
def __init__(self):
|
|
self.call_count = 0
|
|
|
|
def encode_prompt(
|
|
self,
|
|
prompt: str,
|
|
device=None,
|
|
num_images_per_prompt: int = 1,
|
|
negative_prompt: str | None = None,
|
|
prompt_embeds=None,
|
|
negative_prompt_embeds=None,
|
|
):
|
|
self.call_count += 1
|
|
# Return a fresh tensor so identity changes between real calls.
|
|
return torch.tensor([float(self.call_count)])
|
|
|
|
|
|
class TestInstallAndUninstall:
|
|
def test_returns_none_when_pipeline_has_no_encode_prompt(self):
|
|
class Empty:
|
|
pass
|
|
|
|
assert install_prompt_embed_cache(Empty()) is None
|
|
|
|
def test_cache_hits_on_identical_args(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe, max_size=4)
|
|
assert cache is not None
|
|
|
|
out1 = pipe.encode_prompt("cat", device="cpu")
|
|
out2 = pipe.encode_prompt("cat", device="cpu")
|
|
# Second call must be served from cache (wrapped fn not invoked).
|
|
assert pipe.call_count == 1
|
|
# Cache returns the same tensor object.
|
|
assert torch.equal(out1, out2)
|
|
stats = cache.stats()
|
|
assert stats["hits"] == 1
|
|
assert stats["misses"] == 1
|
|
assert stats["bypassed"] == 0
|
|
|
|
def test_positional_and_keyword_share_cache_slot(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe)
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt(prompt="cat")
|
|
assert pipe.call_count == 1
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
def test_different_prompts_miss(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe)
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt("dog")
|
|
assert pipe.call_count == 2
|
|
assert cache.stats()["misses"] == 2
|
|
|
|
def test_tensor_argument_bypasses_and_increments_counter(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe)
|
|
# Put a tensor into a non-precomputed-embed slot to trigger bypass.
|
|
pipe.encode_prompt("cat", device=torch.device("cpu"), num_images_per_prompt=1)
|
|
|
|
# Now pass a fake tensor via a harmless positional-like call: use
|
|
# the negative_prompt slot with an unhashable object.
|
|
class Unhashable:
|
|
pass
|
|
|
|
pipe.encode_prompt("cat", negative_prompt=Unhashable()) # type: ignore[arg-type]
|
|
pipe.encode_prompt("cat", negative_prompt=Unhashable()) # type: ignore[arg-type]
|
|
assert cache.stats()["bypassed"] == 2
|
|
# All three calls actually executed the underlying function.
|
|
assert pipe.call_count == 3
|
|
|
|
def test_precomputed_embeds_bypass_cache(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe)
|
|
pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
|
|
pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
|
|
assert cache.stats()["bypassed"] == 2
|
|
assert cache.stats()["hits"] == 0
|
|
assert pipe.call_count == 2
|
|
|
|
def test_install_is_idempotent(self):
|
|
pipe = _FakePipeline()
|
|
c1 = install_prompt_embed_cache(pipe, max_size=4)
|
|
c2 = install_prompt_embed_cache(pipe, max_size=16)
|
|
assert c1 is c2
|
|
# The second call must not re-wrap (max_size stays from first install).
|
|
assert c1.max_size == 4
|
|
|
|
def test_disabled_cache_passes_through(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe, enabled=False)
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt("cat")
|
|
assert pipe.call_count == 2
|
|
assert cache.stats()["hits"] == 0
|
|
assert cache.stats()["misses"] == 0
|
|
|
|
def test_toggle_enabled_flag_takes_effect(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe, enabled=True)
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt("cat")
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
cache.enabled = False
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt("cat")
|
|
# Cache-bypassed calls go through; hit count unchanged.
|
|
assert cache.stats()["hits"] == 1
|
|
|
|
def test_cached_tensor_is_detached(self):
|
|
class GradPipeline:
|
|
def encode_prompt(self, prompt):
|
|
return torch.ones(1, requires_grad=True)
|
|
|
|
pipe = GradPipeline()
|
|
install_prompt_embed_cache(pipe)
|
|
# First call populates the cache.
|
|
pipe.encode_prompt("cat")
|
|
# Second call returns the cached tensor which must be detached.
|
|
cached = pipe.encode_prompt("cat")
|
|
assert isinstance(cached, torch.Tensor)
|
|
assert not cached.requires_grad
|
|
|
|
def test_lru_eviction_through_wrapper(self):
|
|
pipe = _FakePipeline()
|
|
cache = install_prompt_embed_cache(pipe, max_size=2)
|
|
pipe.encode_prompt("a")
|
|
pipe.encode_prompt("b")
|
|
pipe.encode_prompt("c") # evicts "a"
|
|
pipe.encode_prompt("a") # miss again → re-run
|
|
assert pipe.call_count == 4
|
|
assert cache.stats()["size"] == 2
|
|
|
|
def test_uninstall_restores_original(self):
|
|
pipe = _FakePipeline()
|
|
original = pipe.encode_prompt
|
|
install_prompt_embed_cache(pipe)
|
|
assert pipe.encode_prompt is not original
|
|
uninstall_prompt_embed_cache(pipe)
|
|
# The bound method reference will differ, but identity of underlying
|
|
# function matches and wrapper attributes are gone.
|
|
assert not hasattr(pipe, "_prompt_embed_cache") or pipe._prompt_embed_cache is None
|
|
# After uninstall calls go straight to the real function.
|
|
pipe.encode_prompt("cat")
|
|
pipe.encode_prompt("cat")
|
|
assert pipe.call_count == 2
|
|
|
|
def test_uninstall_on_unwrapped_pipeline_is_noop(self):
|
|
class Empty:
|
|
def encode_prompt(self, prompt):
|
|
return None
|
|
|
|
pipe = Empty()
|
|
uninstall_prompt_embed_cache(pipe) # must not raise
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# resolve_prompt_embed_cache_config
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_env(monkeypatch):
|
|
monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raising=False)
|
|
monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", raising=False)
|
|
|
|
|
|
class TestResolveConfig:
|
|
def test_defaults_when_nothing_set(self, clean_env):
|
|
enable, size = resolve_prompt_embed_cache_config()
|
|
assert enable is False
|
|
assert size == 32
|
|
|
|
def test_explicit_args_used(self, clean_env):
|
|
enable, size = resolve_prompt_embed_cache_config(enable=True, max_size=8)
|
|
assert enable is True
|
|
assert size == 8
|
|
|
|
@pytest.mark.parametrize("raw", ["1", "true", "TRUE", "yes", "on"])
|
|
def test_env_enable_truthy(self, clean_env, monkeypatch, raw):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
|
|
enable, _ = resolve_prompt_embed_cache_config(enable=False)
|
|
assert enable is True
|
|
|
|
@pytest.mark.parametrize("raw", ["0", "false", "FALSE", "no", "off"])
|
|
def test_env_enable_falsy_overrides_true(self, clean_env, monkeypatch, raw):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
|
|
enable, _ = resolve_prompt_embed_cache_config(enable=True)
|
|
assert enable is False
|
|
|
|
def test_env_enable_invalid_leaves_arg_untouched(self, clean_env, monkeypatch):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "garbage")
|
|
enable, _ = resolve_prompt_embed_cache_config(enable=True)
|
|
assert enable is True
|
|
|
|
def test_env_size_overrides_arg(self, clean_env, monkeypatch):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "7")
|
|
_, size = resolve_prompt_embed_cache_config(max_size=99)
|
|
assert size == 7
|
|
|
|
def test_env_size_zero_or_negative_ignored(self, clean_env, monkeypatch):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "0")
|
|
_, size = resolve_prompt_embed_cache_config(max_size=12)
|
|
assert size == 12
|
|
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "-5")
|
|
_, size = resolve_prompt_embed_cache_config(max_size=12)
|
|
assert size == 12
|
|
|
|
def test_env_size_non_integer_ignored(self, clean_env, monkeypatch):
|
|
monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "notanint")
|
|
_, size = resolve_prompt_embed_cache_config(max_size=12)
|
|
assert size == 12
|
|
|
|
def test_max_size_none_falls_back_to_default(self, clean_env):
|
|
_, size = resolve_prompt_embed_cache_config(enable=True, max_size=None)
|
|
assert size == 32
|