[Metrics] [KVConnector] Add Offloading Connector metrics

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -16,6 +16,7 @@
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
     OffloadingConnector,
     OffloadingConnectorMetadata,
+    OffloadingConnectorStats,
 )
 from vllm.forward_context import ForwardContext
 from vllm.utils.hashing import sha256
@@ -86,7 +87,14 @@ def complete_jobs(self, job_ids: set[int]) -> None:
             if job_id in self.waiting_jobs:
                 self.waiting_jobs.remove(job_id)
                 self.completed_jobs.append(job_id)
-                self.completed_transfers.append((job_id, True))
+                result = TransferResult(
+                    job_id=job_id,
+                    success=True,
+                    transfer_size=None,
+                    transfer_time=None,
+                    transfer_type=None,
+                )
+                self.completed_transfers.append(result)
 
     def wait(self, job_ids: set[int]) -> None:
         self.flushed_jobs |= job_ids
@@ -720,3 +728,144 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
 
     # second request will use the GPU prefix cache
     assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
+
+
+class TestOffloadingConnectorStats:
+    """Tests for OffloadingConnector stats reconstruction and operations."""
+
+    def test_build_kv_connector_stats_with_none(self):
+        """Test that build_kv_connector_stats returns empty stats when given None."""
+        stats = OffloadingConnector.build_kv_connector_stats(data=None)
+
+        assert stats is not None
+        assert isinstance(stats, OffloadingConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_with_empty_dict(self):
+        """Test that build_kv_connector_stats returns empty stats with empty dict."""
+        stats = OffloadingConnector.build_kv_connector_stats(data={})
+
+        assert stats is not None
+        assert isinstance(stats, OffloadingConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_reconstructs_offload_stats(self):
+        """Test that OffloadingConnector stats are properly reconstructed with
+        correct data."""
+        serialized_data = {
+            "CPU_to_GPU": [
+                {"op_size": 16, "op_time": 1.0},
+                {"op_size": 8, "op_time": 0.5},
+            ],
+            "GPU_to_CPU": [
+                {"op_size": 1, "op_time": 0.1},
+                {"op_size": 2, "op_time": 0.2},
+            ],
+        }
+
+        stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
+
+        offload_connector_stats = stats
+        assert isinstance(offload_connector_stats, OffloadingConnectorStats)
+        assert offload_connector_stats.data["CPU_to_GPU"] == [
+            {"op_size": 16, "op_time": 1.0},
+            {"op_size": 8, "op_time": 0.5},
+        ]
+        assert offload_connector_stats.data["GPU_to_CPU"] == [
+            {"op_size": 1, "op_time": 0.1},
+            {"op_size": 2, "op_time": 0.2},
+        ]
+
+    def test_aggregate_same_connector(self):
+        """Test aggregating stats from the same connector type."""
+        stats1 = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 16, "op_time": 1.0},
+                    {"op_size": 8, "op_time": 0.5},
+                ],
+                "GPU_to_CPU": [
+                    {"op_size": 1, "op_time": 0.1},
+                    {"op_size": 2, "op_time": 0.2},
+                ],
+            }
+        )
+
+        stats2 = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+            }
+        )
+
+        result = stats1.aggregate(stats2)
+
+        assert result is stats1  # Should return self
+        offload_connector_stats = result
+        assert offload_connector_stats.data["CPU_to_GPU"] == [
+            {"op_size": 16, "op_time": 1.0},
+            {"op_size": 8, "op_time": 0.5},
+            {"op_size": 3, "op_time": 0.2},
+            {"op_size": 7, "op_time": 0.9},
+        ]
+        assert offload_connector_stats.data["GPU_to_CPU"] == [
+            {"op_size": 1, "op_time": 0.1},
+            {"op_size": 2, "op_time": 0.2},
+            {"op_size": 16, "op_time": 2},
+        ]
+
+    def test_reduce(self):
+        """Test that reduce() correctly reduces all nested connector stats."""
+        stats = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 16, "op_time": 1.0},
+                    {"op_size": 8, "op_time": 0.5},
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [
+                    {"op_size": 1, "op_time": 0.1},
+                    {"op_size": 2, "op_time": 0.2},
+                    {"op_size": 16, "op_time": 2},
+                ],
+            }
+        )
+
+        reduced = stats.reduce()
+
+        assert isinstance(reduced, dict)
+        # Check that the stats were reduced (should have aggregated values)
+        assert "CPU_to_GPU_total_bytes" in reduced
+        assert "CPU_to_GPU_total_time" in reduced
+        assert "GPU_to_CPU_total_bytes" in reduced
+        assert "GPU_to_CPU_total_time" in reduced
+        assert reduced["CPU_to_GPU_total_bytes"] == 34
+        assert reduced["CPU_to_GPU_total_time"] == 2.6
+        assert reduced["GPU_to_CPU_total_time"] == 2.3
+        assert reduced["GPU_to_CPU_total_bytes"] == 19
+
+    def test_reset(self):
+        """Test that reset() resets all nested connector stats."""
+        offload_connector_stats = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+            }
+        )
+
+        assert not offload_connector_stats.is_empty()
+
+        offload_connector_stats.reset()
+
+        # After reset, stats should be empty
+        assert offload_connector_stats.is_empty()
+        assert len(offload_connector_stats.data) == 0

tests/v1/kv_offload/test_cpu_gpu.py

@@ -168,15 +168,30 @@ def test_transfer(
     orig_dst_caches = [x.clone() for x in handler.dst_tensors]
 
     # call transfer function
+    start_time = time.time()
     assert handler.transfer_async(1, (src_spec, dst_spec))
-    assert set({x[0] for x in handler._transfers}) == {1}
+    assert set({x.job_id for x in handler._transfers}) == {1}
 
     # wait for transfer to complete
     end_time = time.time() + 10
     while time.time() < end_time:
         finished = handler.get_finished()
         if finished:
-            assert finished == [(1, True)]
+            assert finished[0].job_id == 1
+            assert finished[0].success
+            assert (
+                finished[0].transfer_type == ("GPU", "CPU")
+                if gpu_to_cpu
+                else ("CPU", "GPU")
+            )
+            assert (
+                finished[0].transfer_size
+                == handler.total_block_size_in_bytes
+                * handler.dst_block_size_factor
+                * len(dst_blocks)
+            )
+            assert finished[0].transfer_time > 0
+            assert finished[0].transfer_time < (time.time() - start_time)
             break
         time.sleep(0.1)
 

vllm/distributed/kv_transfer/kv_connector/v1/metrics.py

@@ -124,7 +124,7 @@ def __init__(
         self._counter_cls = metric_types[Counter]
         self._histogram_cls = metric_types[Histogram]
         self._labelnames = labelnames
-        self._per_engine_labelvalues = per_engine_labelvalues
+        self.per_engine_labelvalues = per_engine_labelvalues
 
     def make_per_engine(self, metric: PromMetric) -> dict[int, PromMetric]:
         """
@@ -134,7 +134,7 @@ def make_per_engine(self, metric: PromMetric) -> dict[int, PromMetric]:
         """
         return {
             idx: metric.labels(*labelvalues)
-            for idx, labelvalues in self._per_engine_labelvalues.items()
+            for idx, labelvalues in self.per_engine_labelvalues.items()
         }
 
     def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):