diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index de37afc9a93..b1b4df53635 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -17,18 +17,29 @@ from unittest.mock import MagicMock import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) from vllm.v1.kv_offload.base import ( + OffloadingCounterMetadata, OffloadKey, OffloadPolicy, ReqContext, RequestOffloadingContext, make_offload_key, ) +from vllm.v1.kv_offload.tiering.base import ( + JobMetadata, + JobResult, + SecondaryTierManager, +) from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, TieringOffloadingManager, ) +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec _CTX = ReqContext(req_id="test") _MOCK_OFFLOADING_SPEC = MagicMock() @@ -63,6 +74,102 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None: return count +class MetricsSecondaryTierManager(SecondaryTierManager): + """Test-only secondary tier that declares and emits one labeled metric.""" + + MY_TIER_METRIC = "my_tier_metric" + + @classmethod + def build_metric_definitions(cls, extra_config): + return { + cls.MY_TIER_METRIC: OffloadingCounterMetadata( + documentation="Number of bytes served by the test tier.", + labelnames=("tier",), + ) + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.stats: OffloadingConnectorStats | None = None + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + return False + + def submit_store(self, job_metadata: JobMetadata) -> None: + return + + def submit_load(self, job_metadata: JobMetadata) -> None: + return + + def get_finished_jobs(self) -> Iterable[JobResult]: + return () + + def drain_jobs(self) -> None: + return + + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + return RequestOffloadingContext() + + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.stats + self.stats = None + return stats + + +def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): + monkeypatch.setitem( + SecondaryTierFactory._registry, + "test_metrics", + lambda: MetricsSecondaryTierManager, + ) + + metrics = TieringOffloadingSpec.build_metric_definitions( + {"secondary_tiers": [{"type": "test_metrics"}]} + ) + + metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC] + assert metadata.documentation == "Number of bytes served by the test tier." + assert metadata.labelnames == ("tier",) + + +def test_tiering_manager_aggregates_secondary_stats(): + mock_region = _mock_mmap_region(5) + primary_tier = CPUPrimaryTierOffloadingManager( + num_blocks=5, mmap_region=mock_region + ) + secondary_tier = MetricsSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="test_metrics", + ) + secondary_stats = OffloadingConnectorStats() + secondary_stats.increase_counter( + MetricsSecondaryTierManager.MY_TIER_METRIC, 7, ("test_metrics",) + ) + secondary_tier.stats = secondary_stats + manager = TieringOffloadingManager( + primary_tier=primary_tier, + secondary_tiers=[secondary_tier], + ) + + stats = manager.get_stats() + + assert stats is not None + assert ( + stats.data["data"][MetricsSecondaryTierManager.MY_TIER_METRIC][ + ("test_metrics",) + ] + == 7 + ) + + # The primary tier's cache-usage gauge is always reported, so get_stats() + # never returns None, but the secondary tier has nothing new to report + # once its stats have been consumed. + second_stats = manager.get_stats() + assert second_stats is not None + assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"] + + class TestExampleSecondaryTierManager: """Tests for ExampleSecondaryTierManager implementation.""" diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 87481603f53..c7927572491 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -7,13 +7,21 @@ Abstract interfaces and data types for the secondary tiering layer. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext +from vllm.v1.kv_offload.base import ( + OffloadingMetricMetadata, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) if TYPE_CHECKING: + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, + ) from vllm.v1.kv_offload.base import OffloadingSpec # Type alias for job IDs used in async transfer tracking @@ -228,3 +236,14 @@ class SecondaryTierManager(ABC): def shutdown(self) -> None: """Release resources held by this tier (threads, connections, etc.).""" return + + @classmethod + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + """Return Prometheus metric definitions emitted by this tier.""" + return {} + + def get_stats(self) -> "OffloadingConnectorStats | None": + """Return and reset metric observations collected by this tier.""" + return None diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index be703a03b3d..ed69de9b27e 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -31,19 +31,9 @@ class SecondaryTierFactory: primary_kv_view: memoryview, offloading_spec: "OffloadingSpec", ) -> SecondaryTierManager: + tier_cls = cls.get_tier_class(tier_config) config = tier_config.copy() - - tier_type = config.pop("type", None) - if not tier_type: - raise ValueError("Secondary tier configuration must include 'type'") - - if tier_type not in cls._registry: - raise ValueError( - f"Unknown secondary tier type: {tier_type!r}. " - f"Supported types: {list(cls._registry)}" - ) - - tier_cls = cls._registry[tier_type]() + tier_type = config.pop("type") return tier_cls( offloading_spec=offloading_spec, primary_kv_view=primary_kv_view, @@ -51,6 +41,18 @@ class SecondaryTierFactory: **config, ) + @classmethod + def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]: + tier_type = tier_config.get("type") + if not tier_type: + raise ValueError("Secondary tier configuration must include 'type'") + if tier_type not in cls._registry: + raise ValueError( + f"Unknown secondary tier type: {tier_type!r}. " + f"Supported types: {list(cls._registry)}" + ) + return cls._registry[tier_type]() + SecondaryTierFactory.register_tier( "example", diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index d13e1f1eea5..abed61a3e48 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -27,6 +27,9 @@ from dataclasses import dataclass, field import numpy as np from typing_extensions import override +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( LoadStoreSpec, @@ -628,6 +631,24 @@ class TieringOffloadingManager(OffloadingManager): self._request_level_tiers.clear() self._processed_jobs_this_step = False + @override + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.primary_tier.get_stats() + + if stats is not None and stats.is_empty(): + stats = None + + for tier in self.secondary_tiers: + tier_stats = tier.get_stats() + if tier_stats is None or tier_stats.is_empty(): + continue + if stats is None: + stats = tier_stats + else: + stats.aggregate(tier_stats) + + return stats + @override def shutdown(self) -> None: """Shutdown all tiers and release resources.""" diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index e9dd68c44f6..f4a44a4a8a9 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -31,13 +31,19 @@ Example configuration: } """ +from typing import Any + import torch from typing_extensions import override from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.kv_offload.base import CanonicalKVCaches, OffloadingManager +from vllm.v1.kv_offload.base import ( + CanonicalKVCaches, + OffloadingManager, + OffloadingMetricMetadata, +) from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec @@ -65,6 +71,22 @@ class TieringOffloadingSpec(CPUOffloadingSpec): BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + @classmethod + @override + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + metrics = super().build_metric_definitions(extra_config) + secondary_tier_configs = extra_config.get("secondary_tiers", []) + if not isinstance(secondary_tier_configs, list): + raise ValueError("secondary_tiers must be a list of tier configurations") + + for tier_config in secondary_tier_configs: + assert isinstance(tier_config, dict) + tier_cls = SecondaryTierFactory.get_tier_class(tier_config) + metrics.update(tier_cls.build_metric_definitions(tier_config)) + return metrics + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): super().__init__(vllm_config, kv_cache_config) # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it