From 9a938df64e5e296e42524fc0a67956923806197d Mon Sep 17 00:00:00 2001 From: AlexHuang Date: Mon, 22 Jun 2026 22:45:04 +0800 Subject: [PATCH] [Test][KV Offloading] Add unit tests for OffloadingSpecFactory and SecondaryTierFactory (#46355) Signed-off-by: Alex --- tests/v1/kv_offload/__init__.py | 0 tests/v1/kv_offload/test_factory.py | 272 ++++++++++++++++++++ tests/v1/kv_offload/tiering/__init__.py | 0 tests/v1/kv_offload/tiering/test_factory.py | 152 +++++++++++ 4 files changed, 424 insertions(+) create mode 100644 tests/v1/kv_offload/__init__.py create mode 100644 tests/v1/kv_offload/test_factory.py create mode 100644 tests/v1/kv_offload/tiering/__init__.py create mode 100644 tests/v1/kv_offload/tiering/test_factory.py diff --git a/tests/v1/kv_offload/__init__.py b/tests/v1/kv_offload/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py new file mode 100644 index 0000000000..543ee44330 --- /dev/null +++ b/tests/v1/kv_offload/test_factory.py @@ -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 diff --git a/tests/v1/kv_offload/tiering/__init__.py b/tests/v1/kv_offload/tiering/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py new file mode 100644 index 0000000000..2b65baa020 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -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")