Add tiering offloading metrics

tests/v1/kv_connector/unit/offloading_connector/test_metrics.py

@@ -25,6 +25,7 @@
     OffloadingHistogramMetadata,
 )
 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
+from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics
 
 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -101,10 +102,15 @@ def set(self, value):
 
 
 class _FakeVllmConfig:
-    def __init__(self, store_threshold: int = 2):
-        self.kv_transfer_config = SimpleNamespace(
-            kv_connector_extra_config={"store_threshold": store_threshold}
-        )
+    def __init__(
+        self,
+        store_threshold: int = 2,
+        extra_config: dict[str, Any] | None = None,
+    ):
+        config = {"store_threshold": store_threshold}
+        if extra_config is not None:
+            config.update(extra_config)
+        self.kv_transfer_config = SimpleNamespace(kv_connector_extra_config=config)
 
 
 def _spec_cls_with_metric_definitions(
@@ -679,6 +685,30 @@ def test_prom_metrics_uses_configured_manager_metrics():
     assert STORES_SKIPPED not in prom_metrics._offloading_metric_metadata
 
 
+def test_prom_metrics_registers_tiering_metrics_from_spec():
+    prom_metrics = OffloadPromMetrics(
+        vllm_config=_FakeVllmConfig(
+            store_threshold=0,
+            extra_config={
+                "spec_name": "TieringOffloadingSpec",
+                "secondary_tiers": [],
+            },
+        ),  # type: ignore[arg-type]
+        metric_types={
+            Gauge: _FakeMetric,
+            Counter: _FakeMetric,
+            Histogram: _FakeMetric,
+        },
+        labelnames=["model_name", "engine"],
+        per_engine_labelvalues={0: ["model", "0"]},
+    )
+
+    metric = prom_metrics._offloading_metric_defs[
+        TieringOffloadingMetrics.BLOCK_QUERIES
+    ]
+    assert metric.kwargs["labelnames"] == ["model_name", "engine", "tier"]
+
+
 def test_aggregate_into_empty_stats():
     """Aggregating non-empty stats into a fresh (empty) stats object works."""
     empty = OffloadingConnectorStats()

tests/v1/kv_offload/tiering/p2p/test_manager.py

@@ -16,7 +16,7 @@
 import pytest
 
 from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
-from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
+from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
 from vllm.v1.kv_offload.tiering.p2p.manager import (
     _UNBOUND_STORE_TIMEOUT_S,
@@ -86,12 +86,12 @@ def _job_metadata(
     keys: list[bytes] | None = None,
     block_ids: list[int] | None = None,
     kv_params: dict | None = None,
-) -> JobMetadata:
+) -> TransferJob:
     if keys is None:
         keys = [b"key1"]
     if block_ids is None:
         block_ids = list(range(len(keys)))
-    return JobMetadata(
+    return TransferJob(
         job_id=job_id,
         keys=keys,
         block_ids=np.array(block_ids),

tests/v1/kv_offload/tiering/p2p/test_sessions.py

@@ -24,7 +24,7 @@
     ReqContext,
     RequestOffloadingContext,
 )
-from vllm.v1.kv_offload.tiering.base import JobMetadata
+from vllm.v1.kv_offload.tiering.base import TransferJob
 from vllm.v1.kv_offload.tiering.p2p.session import (
     LoadResult,
     P2PSession,
@@ -256,13 +256,13 @@ def create_store_job(
         self,
         keys: Sequence[OffloadKey],
         ctx: ReqContext,
-    ) -> JobMetadata:
+    ) -> TransferJob:
         keys_list = list(keys)
         self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id))
         block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32)
         job_id = self._next_job_id
         self._next_job_id += 1
-        return JobMetadata(
+        return TransferJob(
             job_id=job_id,
             keys=keys_list,
             block_ids=block_ids,

tests/v1/kv_offload/tiering/test_fs_tier.py

@@ -35,7 +35,7 @@
     OffloadingModelConfig,
     OffloadingParallelConfig,
 )
-from vllm.v1.kv_offload.tiering.base import JobMetadata
+from vllm.v1.kv_offload.tiering.base import TransferJob
 from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
 from vllm.v1.kv_offload.tiering.fs.manager import (
     FileSystemTierManager,
@@ -105,10 +105,10 @@ def make_job(
     keys: list[OffloadKey],
     block_ids: list[int] | None = None,
     is_promotion: bool = False,
-) -> JobMetadata:
+) -> TransferJob:
     if block_ids is None:
         block_ids = list(range(len(keys)))
-    return JobMetadata(
+    return TransferJob(
         job_id=job_id,
         keys=keys,
         block_ids=np.array(block_ids, dtype=np.int64),

tests/v1/kv_offload/tiering/test_metrics.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+
+import numpy as np
+import pytest
+
+from vllm.v1.kv_offload.base import LookupResult, ReqContext, make_offload_key
+from vllm.v1.kv_offload.tiering.base import (
+    JobResult,
+    TieringOffloadingMetrics,
+    TransferJob,
+)
+from vllm.v1.kv_offload.tiering.manager import JobMetadata
+from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker
+
+_CTX = ReqContext(req_id="test")
+
+
+def to_keys(int_ids: Iterable[int]):
+    return [make_offload_key(str(i).encode(), 0) for i in int_ids]
+
+
+def test_tiering_metrics_tracker_records_lookup_metrics():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs", "p2p"],
+        num_primary_blocks=5,
+        primary_block_size=16,
+    )
+    tracker.on_new_request(_CTX)
+
+    key = to_keys([1])[0]
+    tracker.on_lookup(
+        _CTX,
+        key,
+        tracker.primary_tier_label,
+        LookupResult.MISS,
+        lookup_duration=0.0,
+    )
+    tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.MISS, 0.0)
+    tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0)
+    tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0)
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
+    assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
+    assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1
+    assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1
+    assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS]
+
+
+def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs"],
+        num_primary_blocks=5,
+        primary_block_size=16,
+    )
+    tracker.on_new_request(_CTX)
+    key = to_keys([1])[0]
+
+    tracker.on_lookup(
+        _CTX,
+        key,
+        tracker.primary_tier_label,
+        LookupResult.MISS,
+        lookup_duration=0.0,
+    )
+    tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0)
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
+    assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
+    assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1
+
+    tracker.on_request_allocated(_CTX)
+    tracker.on_lookup(
+        _CTX,
+        key,
+        tracker.primary_tier_label,
+        LookupResult.MISS,
+        lookup_duration=0.0,
+    )
+    tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0)
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    assert TieringOffloadingMetrics.BLOCK_QUERIES not in values
+    assert TieringOffloadingMetrics.BLOCK_HITS not in values
+
+
+def test_tiering_metrics_tracker_records_finished_job_metrics():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs"],
+        num_primary_blocks=5,
+        primary_block_size=16,
+    )
+    cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys(
+        range(4)
+    )
+    jobs = [
+        JobMetadata(
+            TransferJob(0, [cascade_key], np.array([0]), False, _CTX),
+            0,
+        ),
+        JobMetadata(
+            TransferJob(1, [promotion_key], np.array([1]), True, _CTX),
+            0,
+        ),
+        JobMetadata(
+            TransferJob(2, [failed_cascade_key], np.array([2]), False, _CTX),
+            0,
+        ),
+        JobMetadata(
+            TransferJob(3, [failed_promotion_key], np.array([3]), True, _CTX),
+            0,
+        ),
+    ]
+    for job in jobs:
+        tracker.on_job_registered(job)
+
+    tracker.on_job_finished(
+        jobs[0], JobResult(job_id=0, success=True, transfer_time=0.5)
+    )
+    tracker.on_job_finished(
+        jobs[1], JobResult(job_id=1, success=True, transfer_time=0.25)
+    )
+    tracker.on_job_finished(jobs[2], JobResult(job_id=2, success=False))
+    tracker.on_job_finished(jobs[3], JobResult(job_id=3, success=False))
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    label = ("1:fs",)
+    assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 16
+    assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5
+    assert values[TieringOffloadingMetrics.READ_BYTES][label] == 16
+    assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25
+    assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1
+    assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1
+    tracker.assert_idle()
+
+
+def test_tiering_metrics_tracker_records_partial_promotion_success_bytes():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs"],
+        num_primary_blocks=5,
+        primary_block_size=16,
+    )
+    keys = to_keys(range(3))
+    job = JobMetadata(
+        TransferJob(0, keys, np.array([0, 1, 2]), True, _CTX),
+        0,
+    )
+    tracker.on_job_registered(job)
+
+    tracker.on_job_finished(
+        job,
+        JobResult(
+            job_id=0,
+            success=False,
+            successful_keys=(keys[0], keys[2]),
+            transfer_time=0.5,
+        ),
+    )
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    label = ("1:fs",)
+    assert values[TieringOffloadingMetrics.READ_BYTES][label] == 32
+    assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.5
+    assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1
+    tracker.assert_idle()
+
+
+def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs", "p2p"],
+        num_primary_blocks=6,
+        primary_block_size=16,
+    )
+    fs_job = JobMetadata(
+        TransferJob(0, to_keys([0, 1]), np.array([0, 1]), False, _CTX),
+        0,
+    )
+    p2p_job = JobMetadata(
+        TransferJob(1, to_keys([2, 3, 4]), np.array([2, 3, 4]), True, _CTX),
+        1,
+    )
+    tracker.on_job_registered(fs_job)
+    tracker.on_job_registered(p2p_job)
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    fs_label = ("1:fs",)
+    p2p_label = ("2:p2p",)
+    assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][
+        fs_label
+    ] == pytest.approx(2 / 6)
+    assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 1
+    assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][
+        p2p_label
+    ] == pytest.approx(3 / 6)
+    assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 1
+
+
+def test_tiering_metrics_tracker_records_promotion_allocation_failures():
+    tracker = TieringMetricsTracker(
+        tier_types=["fs"],
+        num_primary_blocks=1,
+        primary_block_size=16,
+    )
+
+    tracker.on_promotion_allocation_failure()
+
+    stats = tracker.take_stats()
+    assert stats is not None
+    values = stats.data["data"]
+    assert values[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES][()] == 1

tests/v1/kv_offload/tiering/test_obj_tier.py

@@ -34,7 +34,7 @@
     OffloadingModelConfig,
     OffloadingParallelConfig,
 )
-from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
+from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
 from vllm.v1.kv_offload.tiering.manager import (
     CPUPrimaryTierOffloadingManager,
     TieringOffloadingManager,
@@ -105,10 +105,10 @@ def make_job(
     job_id: int,
     keys: list[OffloadKey],
     block_ids: list[int] | None = None,
-) -> JobMetadata:
+) -> TransferJob:
     if block_ids is None:
         block_ids = list(range(len(keys)))
-    return JobMetadata(
+    return TransferJob(
         job_id=job_id,
         keys=keys,
         block_ids=np.array(block_ids, dtype=np.int64),
@@ -201,6 +201,9 @@ def release_xfer_handle(self, handle):
     def release_dlist_handle(self, handle):
         pass
 
+    def get_xfer_telemetry(self, handle):
+        return SimpleNamespace(xferDuration=1000)
+
     def _query_memory(self, queries, mem_type, agent_name):
         return [object() if q[3] in self._stored_obj_keys else None for q in queries]