[KV Offloading] Add tiering metric plumbing (#45959)

Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com>
Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>

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."""
 

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

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",

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."""

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