101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
|
||
|
|
from vllm_omni.diffusion.cache.cachedit.runtime import (
|
||
|
|
CacheDiTRequestSpec,
|
||
|
|
RequestScopedCacheDiTRuntime,
|
||
|
|
)
|
||
|
|
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||
|
|
|
||
|
|
|
||
|
|
def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
|
||
|
|
return CacheDiTRequestSpec(
|
||
|
|
installation_key=key,
|
||
|
|
cache_config=DiffusionCacheConfig(),
|
||
|
|
num_inference_steps=steps,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _recording_backend(events):
|
||
|
|
class FakeBackend:
|
||
|
|
def __init__(self, config):
|
||
|
|
events.append(("create", config))
|
||
|
|
self.enabled = False
|
||
|
|
|
||
|
|
def enable(self, pipeline):
|
||
|
|
events.append(("enable", pipeline))
|
||
|
|
self.enabled = True
|
||
|
|
|
||
|
|
def refresh(self, pipeline, num_inference_steps):
|
||
|
|
events.append(("refresh", pipeline, num_inference_steps))
|
||
|
|
|
||
|
|
def disable(self, pipeline):
|
||
|
|
events.append(("disable", pipeline))
|
||
|
|
self.enabled = False
|
||
|
|
|
||
|
|
def is_enabled(self):
|
||
|
|
return self.enabled
|
||
|
|
|
||
|
|
return FakeBackend
|
||
|
|
|
||
|
|
|
||
|
|
def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
|
||
|
|
events = []
|
||
|
|
FakeBackend = _recording_backend(events)
|
||
|
|
monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
|
||
|
|
pipeline = object()
|
||
|
|
config = DiffusionCacheConfig()
|
||
|
|
runtime = RequestScopedCacheDiTRuntime(pipeline)
|
||
|
|
startup_backend = FakeBackend(config)
|
||
|
|
startup_backend.enabled = True
|
||
|
|
|
||
|
|
runtime.adopt(startup_backend, installation_key="generic")
|
||
|
|
runtime.prepare(_spec("generic", 50)) # omitted: keep startup cache
|
||
|
|
runtime.prepare(None) # lossless: uninstall
|
||
|
|
runtime.prepare(None) # repeated lossless: no-op
|
||
|
|
runtime.prepare(_spec("high", 40)) # high: install H3 profile
|
||
|
|
runtime.prepare(_spec("high", 30)) # repeated high: refresh only
|
||
|
|
runtime.prepare(_spec("generic", 20)) # omitted: restore startup profile
|
||
|
|
|
||
|
|
assert [event[0] for event in events] == [
|
||
|
|
"create",
|
||
|
|
"refresh",
|
||
|
|
"disable",
|
||
|
|
"create",
|
||
|
|
"enable",
|
||
|
|
"refresh",
|
||
|
|
"refresh",
|
||
|
|
"disable",
|
||
|
|
"create",
|
||
|
|
"enable",
|
||
|
|
"refresh",
|
||
|
|
]
|
||
|
|
assert events[1][-1] == 50
|
||
|
|
assert events[5][-1] == 40
|
||
|
|
assert events[6][-1] == 30
|
||
|
|
assert events[10][-1] == 20
|
||
|
|
assert runtime.is_enabled
|
||
|
|
|
||
|
|
|
||
|
|
def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
|
||
|
|
events = []
|
||
|
|
monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
|
||
|
|
pipeline = object()
|
||
|
|
runtime = RequestScopedCacheDiTRuntime(pipeline)
|
||
|
|
|
||
|
|
runtime.prepare(_spec("high", 40))
|
||
|
|
runtime.prepare(None)
|
||
|
|
|
||
|
|
assert [event[0] for event in events] == [
|
||
|
|
"create",
|
||
|
|
"enable",
|
||
|
|
"refresh",
|
||
|
|
"disable",
|
||
|
|
]
|
||
|
|
assert events[2][-1] == 40
|
||
|
|
assert not runtime.is_enabled
|