| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for HF3FS KV Connector high-level components: |
| 4 | - TestHf3fsMockClient : file-backed mock client I/O correctness |
| 5 | - TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation |
| 6 | """ |
| 7 | |
| 8 | import os |
| 9 | from unittest.mock import MagicMock |
| 10 | |
| 11 | import numpy as np |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_connector import ( |
| 16 | HF3FSKVConnectorStats, |
| 17 | ) |
| 18 | from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils.hf3fs_mock_client import ( |
| 19 | Hf3fsClient as MockHf3fsClient, |
| 20 | ) |
| 21 | |
| 22 | # --------------------------------------------------------------------------- |
| 23 | # Helpers |
| 24 | # --------------------------------------------------------------------------- |
| 25 | |
| 26 | |
| 27 | @pytest.fixture |
| 28 | def hf3fs_stats(): |
| 29 | """Fresh HF3FSKVConnectorStats instance.""" |
| 30 | return HF3FSKVConnectorStats() |
| 31 | |
| 32 | |
| 33 | def _make_cuda_event(): |
| 34 | """Return a real CUDA event when available, otherwise a MagicMock.""" |
| 35 | if torch.cuda.is_available(): |
| 36 | return torch.cuda.Event() |
| 37 | return MagicMock() |
| 38 | |
| 39 | |
| 40 | # =========================================================================== |
| 41 | # TestHf3fsMockClient |
| 42 | # =========================================================================== |
| 43 | |
| 44 | |
| 45 | class TestHf3fsMockClient: |
| 46 | """Tests for hf3fs_mock_client.Hf3fsClient (file-backend mock).""" |
| 47 | |
| 48 | def test_init_creates_file(self, tmp_path): |
| 49 | """Initializing the client should create the backing file.""" |
| 50 | path = str(tmp_path / "test_file") |
| 51 | client = MockHf3fsClient(path=path, size=4096, bytes_per_page=512, entries=4) |
| 52 | assert os.path.exists(path), "Backing file should be created on init" |
| 53 | assert os.path.getsize(path) == 4096 |
| 54 | client.close() |
| 55 | |
| 56 | @pytest.mark.parametrize( |
| 57 | "dtype, bytes_per_page", |
| 58 | [ |
| 59 | (torch.float32, 512), |
| 60 | (torch.float16, 256), |
| 61 | (torch.bfloat16, 256), |
| 62 | ], |
| 63 | ids=["float32", "float16", "bfloat16"], |
| 64 | ) |
| 65 | def test_batch_write_and_read_dtype(self, tmp_path, dtype, bytes_per_page): |
| 66 | """Write a tensor of the given dtype and verify round-trip correctness.""" |
| 67 | path = str(tmp_path / f"rw_{dtype}") |
| 68 | client = MockHf3fsClient( |
| 69 | path=path, size=bytes_per_page * 8, bytes_per_page=bytes_per_page, entries=4 |
| 70 | ) |
| 71 | elem_size = torch.tensor([], dtype=dtype).element_size() |
| 72 | numel = bytes_per_page // elem_size |
| 73 | tensor_write = torch.arange(numel, dtype=dtype) |
| 74 | event = _make_cuda_event() |
| 75 | |
| 76 | results = client.batch_write([0], [tensor_write], event) |
| 77 | assert results == [bytes_per_page], f"Write should succeed, got {results}" |
| 78 | |
| 79 | tensor_read = torch.zeros(numel, dtype=dtype) |
| 80 | results = client.batch_read([0], [tensor_read]) |
| 81 | assert results == [bytes_per_page], f"Read should succeed, got {results}" |
| 82 | assert torch.equal(tensor_write, tensor_read), ( |
| 83 | "Read tensor should match written tensor" |
| 84 | ) |
| 85 | client.close() |
| 86 | |
| 87 | def test_batch_read_empty_file_returns_error(self, tmp_path): |
| 88 | """Reading out-of-bounds offset should return -1.""" |
| 89 | bytes_per_page = 128 |
| 90 | size = bytes_per_page * 4 |
| 91 | path = str(tmp_path / "empty_read") |
| 92 | client = MockHf3fsClient( |
| 93 | path=path, size=size, bytes_per_page=bytes_per_page, entries=4 |
| 94 | ) |
| 95 | numel = bytes_per_page // 4 |
| 96 | tensor_read = torch.zeros(numel, dtype=torch.float32) |
| 97 | results = client.batch_read([size], [tensor_read]) # offset == size => OOB |
| 98 | assert results[0] == -1, "Out-of-bounds read should return -1" |
| 99 | client.close() |
| 100 | |
| 101 | def test_batch_write_out_of_bounds_returns_error(self, tmp_path): |
| 102 | """Writing at an offset beyond file size should return -1.""" |
| 103 | bytes_per_page = 128 |
| 104 | size = bytes_per_page * 4 |
| 105 | path = str(tmp_path / "oob_write") |
| 106 | client = MockHf3fsClient( |
| 107 | path=path, size=size, bytes_per_page=bytes_per_page, entries=4 |
| 108 | ) |
| 109 | numel = bytes_per_page // 4 |
| 110 | tensor = torch.ones(numel, dtype=torch.float32) |
| 111 | event = _make_cuda_event() |
| 112 | results = client.batch_write([size], [tensor], event) # OOB offset |
| 113 | assert results[0] == -1, "Out-of-bounds write should return -1" |
| 114 | client.close() |
| 115 | |
| 116 | def test_multiple_tensors_rw(self, tmp_path): |
| 117 | """Write multiple tensors at different offsets, then read all back.""" |
| 118 | bytes_per_page = 128 |
| 119 | n = 4 |
| 120 | path = str(tmp_path / "multi_rw") |
| 121 | client = MockHf3fsClient( |
| 122 | path=path, |
| 123 | size=bytes_per_page * n * 2, |
| 124 | bytes_per_page=bytes_per_page, |
| 125 | entries=8, |
| 126 | ) |
| 127 | tensors_write = [ |
| 128 | torch.full((bytes_per_page // 4,), float(i), dtype=torch.float32) |
| 129 | for i in range(n) |
| 130 | ] |
| 131 | offsets = [i * bytes_per_page for i in range(n)] |
| 132 | event = _make_cuda_event() |
| 133 | |
| 134 | results = client.batch_write(offsets, tensors_write, event) |
| 135 | assert all(r == bytes_per_page for r in results) |
| 136 | |
| 137 | tensors_read = [ |
| 138 | torch.zeros(bytes_per_page // 4, dtype=torch.float32) for _ in range(n) |
| 139 | ] |
| 140 | results = client.batch_read(offsets, tensors_read) |
| 141 | assert all(r == bytes_per_page for r in results) |
| 142 | |
| 143 | for i, (tw, tr) in enumerate(zip(tensors_write, tensors_read)): |
| 144 | assert torch.allclose(tw, tr), f"Tensor {i} mismatch after round-trip" |
| 145 | client.close() |
| 146 | |
| 147 | def test_flush_and_close_no_error(self, tmp_path): |
| 148 | """flush() and close() should not raise exceptions.""" |
| 149 | path = str(tmp_path / "flush_close") |
| 150 | client = MockHf3fsClient(path=path, size=1024, bytes_per_page=128, entries=4) |
| 151 | client.flush() |
| 152 | client.close() |
| 153 | |
| 154 | |
| 155 | # =========================================================================== |
| 156 | # TestHF3FSKVConnectorStats |
| 157 | # =========================================================================== |
| 158 | |
| 159 | |
| 160 | class TestHF3FSKVConnectorStats: |
| 161 | """Tests for HF3FSKVConnectorStats metric collection and aggregation.""" |
| 162 | |
| 163 | def test_initial_is_empty(self, hf3fs_stats): |
| 164 | """Fresh stats object should report is_empty() == True.""" |
| 165 | assert hf3fs_stats.is_empty() is True |
| 166 | |
| 167 | @pytest.mark.parametrize( |
| 168 | "task_type, duration_key", |
| 169 | [ |
| 170 | ("Saved", "save_duration"), |
| 171 | ("Loaded", "load_duration"), |
| 172 | ], |
| 173 | ids=["save", "load"], |
| 174 | ) |
| 175 | def test_record_success_duration(self, hf3fs_stats, task_type, duration_key): |
| 176 | """Recording a successful task should update duration list and total count.""" |
| 177 | hf3fs_stats.record_success_task_duration(task_type, 0.5) |
| 178 | assert not hf3fs_stats.is_empty() |
| 179 | assert len(hf3fs_stats.data[duration_key]) == 1 |
| 180 | assert hf3fs_stats.data[duration_key][0] == pytest.approx(0.5) |
| 181 | assert hf3fs_stats.data["num_transfer_task"] == 1 |
| 182 | |
| 183 | @pytest.mark.parametrize( |
| 184 | "task_type, failed_key", |
| 185 | [ |
| 186 | ("Saved", "num_failed_save"), |
| 187 | ("Loaded", "num_failed_load"), |
| 188 | ], |
| 189 | ids=["save", "load"], |
| 190 | ) |
| 191 | def test_record_failed_task(self, hf3fs_stats, task_type, failed_key): |
| 192 | """Recording a failed task should increment the corresponding counter.""" |
| 193 | hf3fs_stats.record_failed_task_count(task_type) |
| 194 | assert hf3fs_stats.data[failed_key] == 1 |
| 195 | assert hf3fs_stats.data["num_transfer_task"] == 1 |
| 196 | |
| 197 | def test_aggregate_two_stats(self): |
| 198 | """aggregate() should merge save/load duration lists and sum counters.""" |
| 199 | stats1 = HF3FSKVConnectorStats() |
| 200 | stats1.record_success_task_duration("Saved", 0.1) |
| 201 | stats1.record_success_task_duration("Loaded", 0.2) |
| 202 | |
| 203 | stats2 = HF3FSKVConnectorStats() |
| 204 | stats2.record_success_task_duration("Saved", 0.3) |
| 205 | stats2.record_failed_task_count("Loaded") |
| 206 | |
| 207 | stats1.aggregate(stats2) |
| 208 | assert stats1.data["save_duration"] == pytest.approx([0.1, 0.3]) |
| 209 | assert stats1.data["load_duration"] == pytest.approx([0.2]) |
| 210 | assert stats1.data["num_failed_load"] == 1 |
| 211 | assert stats1.data["num_transfer_task"] == 4 |
| 212 | |
| 213 | def test_reduce_with_data(self): |
| 214 | """reduce() computes correct averages when data is present.""" |
| 215 | stats = HF3FSKVConnectorStats() |
| 216 | stats.record_success_task_duration("Saved", 1.0) |
| 217 | stats.record_success_task_duration("Saved", 3.0) |
| 218 | result = stats.reduce() |
| 219 | assert result["Num save task success"] == pytest.approx(2.0, rel=0.01) |
| 220 | assert result["Num save task failed"] == pytest.approx(0.0, rel=0.01) |
| 221 | assert result["Avg save duration (ms)"] == pytest.approx(2000.0, rel=0.01) |
| 222 | # Reduced values must be plain Python scalars so CLI logging renders |
| 223 | # them without numpy reprs (eg np.float64(...)). |
| 224 | assert all(not isinstance(v, np.generic) for v in result.values()) |
| 225 | |
| 226 | def test_clone_and_reset(self, hf3fs_stats): |
| 227 | """clone_and_reset() returns a copy with data and resets the original.""" |
| 228 | hf3fs_stats.record_success_task_duration("Saved", 0.7) |
| 229 | hf3fs_stats.record_success_task_duration("Loaded", 0.4) |
| 230 | |
| 231 | clone = hf3fs_stats.clone_and_reset() |
| 232 | assert clone.data["num_transfer_task"] == 2 |
| 233 | assert hf3fs_stats.is_empty() |