wehub-resource-sync
299 lines
10 KiB
Python
299 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Unit tests for OffloadingSpecFactory.
|
|
|
|
These tests verify:
|
|
1. Pre-registration integrity — registered module paths can actually import
|
|
and yield correct OffloadingSpec subclasses (CI sentinel against file moves).
|
|
2. End-to-end factory → spec construction with real configs.
|
|
3. Downstream collaboration — build_metric_definitions delegation.
|
|
4. Error paths — unregistered specs, missing config, duplicate registration.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.config import KVTransferConfig
|
|
from vllm.v1.kv_cache_interface import (
|
|
FullAttentionSpec,
|
|
KVCacheConfig,
|
|
KVCacheGroupSpec,
|
|
KVCacheTensor,
|
|
)
|
|
from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec
|
|
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
|
|
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
|
|
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def restore_registry():
|
|
"""Save and restore OffloadingSpecFactory._registry between tests."""
|
|
original = dict(OffloadingSpecFactory._registry)
|
|
yield
|
|
OffloadingSpecFactory._registry = original
|
|
|
|
|
|
def _make_vllm_config(
|
|
spec_name: str | None = "CPUOffloadingSpec",
|
|
cpu_bytes_to_use: int | None = None,
|
|
store_threshold: int = 0,
|
|
extra_config: dict | None = None,
|
|
):
|
|
"""Build a real VllmConfig with kv_transfer_config set for offloading."""
|
|
from vllm.config import (
|
|
CacheConfig,
|
|
DeviceConfig,
|
|
ModelConfig,
|
|
SchedulerConfig,
|
|
VllmConfig,
|
|
)
|
|
|
|
model_config = ModelConfig(
|
|
model="facebook/opt-125m",
|
|
trust_remote_code=True,
|
|
dtype="float16",
|
|
seed=42,
|
|
)
|
|
scheduler_config = SchedulerConfig(
|
|
max_num_seqs=16,
|
|
max_num_batched_tokens=64,
|
|
max_model_len=10000,
|
|
enable_chunked_prefill=True,
|
|
is_encoder_decoder=model_config.is_encoder_decoder,
|
|
)
|
|
cache_config = CacheConfig(
|
|
block_size=16,
|
|
gpu_memory_utilization=0.9,
|
|
cache_dtype="auto",
|
|
enable_prefix_caching=True,
|
|
)
|
|
|
|
cfg = extra_config or {}
|
|
if cpu_bytes_to_use is not None:
|
|
cfg["cpu_bytes_to_use"] = cpu_bytes_to_use
|
|
cfg["spec_name"] = spec_name
|
|
if store_threshold > 0:
|
|
cfg["store_threshold"] = store_threshold
|
|
|
|
kv_transfer_config = KVTransferConfig(
|
|
kv_connector="OffloadingConnector",
|
|
kv_role="kv_both",
|
|
kv_connector_extra_config=cfg,
|
|
)
|
|
return VllmConfig(
|
|
scheduler_config=scheduler_config,
|
|
model_config=model_config,
|
|
cache_config=cache_config,
|
|
kv_transfer_config=kv_transfer_config,
|
|
device_config=DeviceConfig("cpu"),
|
|
)
|
|
|
|
|
|
def _make_kv_cache_config():
|
|
"""Build a minimal KVCacheConfig with one KV cache tensor."""
|
|
num_blocks = 16
|
|
num_kv_heads = 1
|
|
head_size = 1
|
|
dtype = torch.float32
|
|
page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
|
|
kv_tensor = KVCacheTensor(
|
|
size=num_blocks * page_size, shared_by=["layer"], block_stride=0
|
|
)
|
|
return KVCacheConfig(
|
|
num_blocks=num_blocks,
|
|
kv_cache_tensors=[kv_tensor],
|
|
kv_cache_groups=[
|
|
KVCacheGroupSpec(
|
|
["layer"],
|
|
FullAttentionSpec(
|
|
block_size=16,
|
|
num_kv_heads=num_kv_heads,
|
|
head_size=head_size,
|
|
dtype=dtype,
|
|
),
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pre-registration integrity (CI sentinel)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_pre_registered_specs_can_be_imported():
|
|
"""If someone moves cpu/spec.py but forgets to update factory.py, CI fails."""
|
|
for name in OffloadingSpecFactory._registry:
|
|
cls = OffloadingSpecFactory._registry[name]()
|
|
assert issubclass(cls, OffloadingSpec)
|
|
|
|
|
|
def test_cpu_spec_registered():
|
|
"""CPUOffloadingSpec is registered and importable."""
|
|
cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]()
|
|
assert cls is CPUOffloadingSpec
|
|
|
|
|
|
def test_tiering_spec_registered():
|
|
"""TieringOffloadingSpec is registered and importable."""
|
|
cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]()
|
|
assert cls is TieringOffloadingSpec
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Normal path — get_spec_cls
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_get_spec_cls_returns_registered_class():
|
|
"""Registered spec_name returns correct class."""
|
|
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
assert spec_cls is CPUOffloadingSpec
|
|
|
|
|
|
def test_get_spec_cls_default_to_cpu():
|
|
"""Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
|
|
config = _make_vllm_config(spec_name=None)
|
|
config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
assert spec_cls is CPUOffloadingSpec
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# End-to-end — create_spec
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_create_cpu_offloading_spec_end_to_end():
|
|
"""Full factory → spec construction with real VllmConfig/KVCacheConfig.
|
|
|
|
Verifies:
|
|
- cpu_bytes_to_use validation and num_blocks calculation
|
|
- block_size % hash_block_size assertion
|
|
- spec instance is CPUOffloadingSpec
|
|
"""
|
|
config = _make_vllm_config(cpu_bytes_to_use=65536)
|
|
kv_cache_config = _make_kv_cache_config()
|
|
spec = OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
|
assert isinstance(spec, CPUOffloadingSpec)
|
|
assert spec.num_blocks > 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dynamic import via spec_module_path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_dynamic_load_via_spec_module_path():
|
|
"""External spec loaded via spec_module_path.
|
|
|
|
This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec)
|
|
integrate with vLLM without being pre-registered in the factory.
|
|
The fallback path: registry miss → spec_module_path → importlib.import_module.
|
|
"""
|
|
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
|
|
# Delete from registry to force the dynamic import path
|
|
del OffloadingSpecFactory._registry["CPUOffloadingSpec"]
|
|
# spec_name not in registry → falls through to spec_module_path
|
|
config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
|
|
"vllm.v1.kv_offload.cpu.spec"
|
|
)
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
assert spec_cls is CPUOffloadingSpec
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Error paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_unregistered_spec_without_module_path_raises():
|
|
"""spec_name not in registry + no spec_module_path → ValueError."""
|
|
config = _make_vllm_config(spec_name="NonexistentSpec")
|
|
with pytest.raises(ValueError, match="Unsupported spec type"):
|
|
OffloadingSpecFactory.get_spec_cls(config)
|
|
|
|
# create_spec should also fail (calls get_spec_cls internally)
|
|
kv_cache_config = _make_kv_cache_config()
|
|
with pytest.raises(ValueError, match="Unsupported spec type"):
|
|
OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
|
|
|
|
|
def test_cpu_spec_missing_cpu_bytes_to_use_raises():
|
|
"""CPUOffloadingSpec requires cpu_bytes_to_use → Exception."""
|
|
config = _make_vllm_config(cpu_bytes_to_use=None)
|
|
config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
|
|
kv_cache_config = _make_kv_cache_config()
|
|
with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
|
|
OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
|
|
|
|
|
def test_duplicate_registration_raises():
|
|
"""register_spec with existing name → ValueError."""
|
|
with pytest.raises(ValueError, match="is already registered"):
|
|
OffloadingSpecFactory.register_spec(
|
|
"CPUOffloadingSpec", "some.module", "SomeClass"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Downstream collaboration — build_metric_definitions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_build_metric_definitions_below_threshold():
|
|
"""store_threshold < 2 keeps stores_skipped disabled."""
|
|
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
|
|
|
|
config = _make_vllm_config(store_threshold=1)
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
metrics = spec_cls.build_metric_definitions(
|
|
config.kv_transfer_config.kv_connector_extra_config
|
|
)
|
|
assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
|
|
assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics
|
|
|
|
|
|
def test_build_metric_definitions_allocation_size_histogram():
|
|
"""CPU allocation size is always reported as a histogram."""
|
|
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
|
|
|
|
config = _make_vllm_config(store_threshold=0)
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
metrics = spec_cls.build_metric_definitions(
|
|
config.kv_transfer_config.kv_connector_extra_config
|
|
)
|
|
metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE]
|
|
assert isinstance(metadata, OffloadingHistogramMetadata)
|
|
assert metadata.buckets == (
|
|
1,
|
|
4,
|
|
16,
|
|
64,
|
|
256,
|
|
1024,
|
|
4096,
|
|
16384,
|
|
65536,
|
|
262144,
|
|
)
|
|
|
|
|
|
def test_build_metric_definitions_returns_counter_at_threshold():
|
|
"""store_threshold >= 2 → returns stores_skipped counter definition."""
|
|
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
|
|
|
|
config = _make_vllm_config(store_threshold=2)
|
|
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
|
metrics = spec_cls.build_metric_definitions(
|
|
config.kv_transfer_config.kv_connector_extra_config
|
|
)
|
|
assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
|