| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( |
| 5 | HiSparseConnector, |
| 6 | ) |
| 7 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( |
| 8 | HiSparseKVConnectorStats, |
| 9 | ) |
| 10 | |
| 11 | |
| 12 | def test_is_empty_on_fresh_stats(): |
| 13 | stats = HiSparseKVConnectorStats() |
| 14 | assert stats.is_empty() |
| 15 | |
| 16 | |
| 17 | def test_record_snapshot_and_reduce(): |
| 18 | stats = HiSparseKVConnectorStats() |
| 19 | stats.record_snapshot(hits=7, misses=3, host_to_device_bytes=48) |
| 20 | stats.record_snapshot(hits=5, misses=1, host_to_device_bytes=16) |
| 21 | assert not stats.is_empty() |
| 22 | |
| 23 | reduced = stats.reduce() |
| 24 | assert reduced["HiSparse hot-buffer hits"] == 12 |
| 25 | assert reduced["HiSparse hot-buffer misses"] == 4 |
| 26 | assert reduced["HiSparse host-to-device bytes"] == 64 |
| 27 | |
| 28 | |
| 29 | def test_aggregate_extends_snapshot_deltas(): |
| 30 | first = HiSparseKVConnectorStats() |
| 31 | first.record_snapshot(hits=7, misses=3, host_to_device_bytes=48) |
| 32 | second = HiSparseKVConnectorStats() |
| 33 | second.record_snapshot(hits=5, misses=1, host_to_device_bytes=16) |
| 34 | |
| 35 | first.aggregate(second) |
| 36 | |
| 37 | assert first.data == { |
| 38 | "cache_hits": [7, 5], |
| 39 | "cache_misses": [3, 1], |
| 40 | "host_to_device_bytes": [48, 16], |
| 41 | } |
| 42 | |
| 43 | |
| 44 | def test_aggregate_skips_empty_stats(): |
| 45 | stats = HiSparseKVConnectorStats() |
| 46 | stats.record_snapshot(hits=2, misses=1, host_to_device_bytes=16) |
| 47 | |
| 48 | stats.aggregate(HiSparseKVConnectorStats()) |
| 49 | |
| 50 | assert stats.data["cache_hits"] == [2] |
| 51 | |
| 52 | |
| 53 | def test_build_kv_connector_stats_round_trip(): |
| 54 | stats = HiSparseKVConnectorStats() |
| 55 | stats.record_snapshot(hits=12, misses=4, host_to_device_bytes=64) |
| 56 | payload = stats.to_dict() |
| 57 | |
| 58 | rebuilt = HiSparseConnector.build_kv_connector_stats(data=payload) |
| 59 | |
| 60 | assert rebuilt is not None |
| 61 | assert isinstance(rebuilt, HiSparseKVConnectorStats) |
| 62 | assert rebuilt.reduce() == stats.reduce() |