iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Test][KV Offloading] Add unit tests for OffloadingSpecFactory and SecondaryTierFactory (#46355)

Browse Source 
Signed-off-by: Alex <alex.tech.lab@outlook.com>
This commit is contained in:
AlexHuang
2026-06-22 17:45:04 +03:00
committed by GitHub
parent 3da4a1b124
commit 9a938df64e
4 changed files with 424 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_offload/__init__.py
View File
tests/v1/kv_offload/test_factory.py
+272
View File
@@ -0,0 +1,272 @@
# 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 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_empty_below_threshold():
"""store_threshold < 2 → only base metric (no stores_skipped)."""
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
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
tests/v1/kv_offload/tiering/__init__.py
View File
tests/v1/kv_offload/tiering/test_factory.py
+152
View File
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for SecondaryTierFactory.
These tests verify:
1. Pre-registration integrity — registered tier module paths can import
and yield correct SecondaryTierManager subclasses (CI sentinel).
2. Multi-tier creation via factory with correct tier_type propagation.
3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
"""
from unittest.mock import MagicMock
import pytest
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def restore_registry():
"""Save and restore SecondaryTierFactory._registry between tests."""
original = dict(SecondaryTierFactory._registry)
yield
SecondaryTierFactory._registry = original
def _make_mock_args():
"""Build common mock args for create_secondary_tier."""
return MagicMock(), MagicMock() # primary_kv_view, offloading_spec
# ---------------------------------------------------------------------------
# Pre-registration integrity (CI sentinel)
# ---------------------------------------------------------------------------
def test_pre_registered_tiers_can_be_imported():
"""CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager."""
for tier_type in SecondaryTierFactory._registry:
cls = SecondaryTierFactory._registry[tier_type]()
assert issubclass(cls, SecondaryTierManager)
def test_example_tier_registered():
"""Example tier is registered."""
cls = SecondaryTierFactory._registry["example"]()
assert cls is ExampleSecondaryTierManager
# ---------------------------------------------------------------------------
# Normal path — create_secondary_tier
# ---------------------------------------------------------------------------
def test_create_tier_from_registry():
"""Registered tier_type creates instance with correct tier_type."""
primary_kv_view, offloading_spec = _make_mock_args()
tier_config = {"type": "example"}
tier = SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
)
assert isinstance(tier, SecondaryTierManager)
assert tier.tier_type == "example"
def test_create_multiple_tiers():
"""Multiple tier configs can be created with correct tier_types."""
primary_kv_view, offloading_spec = _make_mock_args()
configs = [
{"type": "example", "custom_param": 1},
{"type": "example", "custom_param": 2},
]
tiers = [
SecondaryTierFactory.create_secondary_tier(
cfg.copy(), primary_kv_view, offloading_spec
)
for cfg in configs
]
assert len(tiers) == 2
assert all(tier.tier_type == "example" for tier in tiers)
assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers)
def test_register_new_tier_type():
"""Verify that new tier types can be registered and created.
This is how external projects add custom secondary tiers
(e.g., llm-d FS backend was upstreamed as "fs" tier via this mechanism).
"""
# Register a new tier type (reuse example manager for simplicity)
SecondaryTierFactory.register_tier(
"custom_tier",
"vllm.v1.kv_offload.tiering.example.manager",
"ExampleSecondaryTierManager",
)
primary_kv_view, offloading_spec = _make_mock_args()
tier = SecondaryTierFactory.create_secondary_tier(
{"type": "custom_tier", "custom_param": 99},
primary_kv_view,
offloading_spec,
)
assert tier.tier_type == "custom_tier"
assert isinstance(tier, ExampleSecondaryTierManager)
# ---------------------------------------------------------------------------
# Error paths
# ---------------------------------------------------------------------------
def test_missing_tier_type_raises():
"""tier_config without 'type' → ValueError."""
primary_kv_view, offloading_spec = _make_mock_args()
tier_config: dict[str, str] = {}
with pytest.raises(ValueError, match="must include 'type'"):
SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
)
def test_unknown_tier_type_raises():
"""Unrecognized tier_type → ValueError with supported types list."""
primary_kv_view, offloading_spec = _make_mock_args()
tier_config = {"type": "nonexistent_tier"}
with pytest.raises(
ValueError,
match=r"Unknown secondary tier type.*Supported types:",
):
SecondaryTierFactory.create_secondary_tier(
tier_config, primary_kv_view, offloading_spec
)
def test_duplicate_registration_raises():
"""register_tier with existing type → ValueError."""
with pytest.raises(ValueError, match="is already registered"):
SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")