| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for FlexKVConnectorV1. |
| 4 | |
| 5 | These tests mock the ``flexkv`` package so they can run without a real FlexKV |
| 6 | installation. They verify: |
| 7 | |
| 8 | 1. That ``FlexKVConnectorV1`` raises a helpful ``ImportError`` when FlexKV is |
| 9 | not installed. |
| 10 | 2. That all public methods are correctly delegated to the underlying |
| 11 | ``FlexKVConnectorV1Impl``. |
| 12 | """ |
| 13 | |
| 14 | import sys |
| 15 | import types |
| 16 | from unittest.mock import MagicMock, patch |
| 17 | |
| 18 | import pytest |
| 19 | import torch |
| 20 | |
| 21 | from vllm.config import KVTransferConfig, VllmConfig |
| 22 | from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole |
| 23 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 24 | |
| 25 | from .utils import create_vllm_config |
| 26 | |
| 27 | # --------------------------------------------------------------------------- |
| 28 | # Helpers |
| 29 | # --------------------------------------------------------------------------- |
| 30 | |
| 31 | |
| 32 | def _make_vllm_config( |
| 33 | kv_connector: str = "FlexKVConnectorV1", |
| 34 | kv_role: str = "kv_both", |
| 35 | ) -> VllmConfig: |
| 36 | """Return a minimal VllmConfig with a KVTransferConfig attached.""" |
| 37 | vllm_config = create_vllm_config(block_size=16, max_num_batched_tokens=512) |
| 38 | vllm_config.kv_transfer_config = KVTransferConfig( |
| 39 | kv_connector=kv_connector, |
| 40 | kv_role=kv_role, |
| 41 | ) |
| 42 | return vllm_config |
| 43 | |
| 44 | |
| 45 | def _make_kv_cache_config() -> KVCacheConfig: |
| 46 | return MagicMock(spec=KVCacheConfig) |
| 47 | |
| 48 | |
| 49 | def _make_flexkv_module( |
| 50 | impl_mock: MagicMock, |
| 51 | ) -> tuple[types.ModuleType, types.ModuleType]: |
| 52 | """Build a fake ``flexkv`` package hierarchy that returns *impl_mock* |
| 53 | when ``FlexKVConnectorV1Impl`` is instantiated.""" |
| 54 | flexkv_mod = types.ModuleType("flexkv") |
| 55 | integration_mod = types.ModuleType("flexkv.integration") |
| 56 | vllm_mod = types.ModuleType("flexkv.integration.vllm") |
| 57 | adapter_mod = types.ModuleType("flexkv.integration.vllm.vllm_v1_adapter") |
| 58 | |
| 59 | # Make FlexKVConnectorV1Impl() return our mock instance. |
| 60 | # The "# type: ignore" markers below are needed because ModuleType does |
| 61 | # not declare these attributes statically; they are set dynamically. |
| 62 | FlexKVConnectorV1ImplCls = MagicMock(return_value=impl_mock) |
| 63 | adapter_mod.FlexKVConnectorV1Impl = FlexKVConnectorV1ImplCls # type: ignore |
| 64 | |
| 65 | flexkv_mod.integration = integration_mod # type: ignore |
| 66 | integration_mod.vllm = vllm_mod # type: ignore |
| 67 | vllm_mod.vllm_v1_adapter = adapter_mod # type: ignore |
| 68 | |
| 69 | return flexkv_mod, adapter_mod |
| 70 | |
| 71 | |
| 72 | def _install_flexkv_mock(impl_mock: MagicMock): |
| 73 | """Insert fake flexkv modules into sys.modules and return a context that |
| 74 | cleans them up afterwards.""" |
| 75 | flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock) |
| 76 | mods = { |
| 77 | "flexkv": flexkv_mod, |
| 78 | "flexkv.integration": flexkv_mod.integration, |
| 79 | "flexkv.integration.vllm": flexkv_mod.integration.vllm, |
| 80 | "flexkv.integration.vllm.vllm_v1_adapter": adapter_mod, |
| 81 | } |
| 82 | return patch.dict(sys.modules, mods) |
| 83 | |
| 84 | |
| 85 | def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock): |
| 86 | """Instantiate FlexKVConnectorV1 with faked flexkv modules.""" |
| 87 | from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import ( |
| 88 | FlexKVConnectorV1, |
| 89 | ) |
| 90 | |
| 91 | with _install_flexkv_mock(impl_mock): |
| 92 | connector = FlexKVConnectorV1( |
| 93 | vllm_config=vllm_config, |
| 94 | role=KVConnectorRole.WORKER, |
| 95 | kv_cache_config=_make_kv_cache_config(), |
| 96 | ) |
| 97 | return connector |
| 98 | |
| 99 | |
| 100 | # --------------------------------------------------------------------------- |
| 101 | # Tests |
| 102 | # --------------------------------------------------------------------------- |
| 103 | |
| 104 | |
| 105 | class TestFlexKVConnectorImportError: |
| 106 | """FlexKVConnectorV1 should fail with a helpful message when flexkv is |
| 107 | absent.""" |
| 108 | |
| 109 | def test_import_error_message(self): |
| 110 | from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import ( |
| 111 | FlexKVConnectorV1, |
| 112 | ) |
| 113 | |
| 114 | # Ensure flexkv is NOT in sys.modules |
| 115 | for key in list(sys.modules): |
| 116 | if key.startswith("flexkv"): |
| 117 | del sys.modules[key] |
| 118 | |
| 119 | with pytest.raises(ImportError, match="(?i)flexkv") as exc_info: |
| 120 | FlexKVConnectorV1( |
| 121 | vllm_config=_make_vllm_config(), |
| 122 | role=KVConnectorRole.WORKER, |
| 123 | kv_cache_config=_make_kv_cache_config(), |
| 124 | ) |
| 125 | |
| 126 | assert "https://github.com/taco-project/FlexKV" in str(exc_info.value) |
| 127 | |
| 128 | |
| 129 | class TestFlexKVConnectorDelegation: |
| 130 | """All public API methods should be forwarded to the impl.""" |
| 131 | |
| 132 | @pytest.fixture() |
| 133 | def connector_and_impl(self): |
| 134 | impl = MagicMock() |
| 135 | cfg = _make_vllm_config() |
| 136 | connector = _build_connector(cfg, impl) |
| 137 | return connector, impl |
| 138 | |
| 139 | def test_shutdown(self, connector_and_impl): |
| 140 | connector, impl = connector_and_impl |
| 141 | connector.shutdown() |
| 142 | impl.shutdown.assert_called_once() |
| 143 | |
| 144 | def test_start_load_kv(self, connector_and_impl): |
| 145 | connector, impl = connector_and_impl |
| 146 | ctx = MagicMock() |
| 147 | connector.start_load_kv(ctx, extra_arg="x") |
| 148 | impl.start_load_kv.assert_called_once_with(ctx, extra_arg="x") |
| 149 | |
| 150 | def test_save_kv_layer(self, connector_and_impl): |
| 151 | connector, impl = connector_and_impl |
| 152 | kv_layer = torch.zeros(4, 4) |
| 153 | attn_meta = MagicMock() |
| 154 | connector.save_kv_layer("layer_0", kv_layer, attn_meta) |
| 155 | impl.save_kv_layer.assert_called_once_with("layer_0", kv_layer, attn_meta) |
| 156 | |
| 157 | def test_wait_for_save(self, connector_and_impl): |
| 158 | connector, impl = connector_and_impl |
| 159 | connector.wait_for_save() |
| 160 | impl.wait_for_save.assert_called_once() |
| 161 | |
| 162 | def test_get_finished(self, connector_and_impl): |
| 163 | connector, impl = connector_and_impl |
| 164 | impl.get_finished.return_value = ({"req1"}, None) |
| 165 | result = connector.get_finished({"req1"}) |
| 166 | impl.get_finished.assert_called_once_with({"req1"}) |
| 167 | assert result == ({"req1"}, None) |
| 168 | |
| 169 | def test_register_kv_caches(self, connector_and_impl): |
| 170 | connector, impl = connector_and_impl |
| 171 | kv_caches = {"layer_0": torch.zeros(1)} |
| 172 | connector.register_kv_caches(kv_caches) |
| 173 | impl.register_kv_caches.assert_called_once_with(kv_caches) |
| 174 | |
| 175 | def test_get_num_new_matched_tokens(self, connector_and_impl): |
| 176 | connector, impl = connector_and_impl |
| 177 | req = MagicMock() |
| 178 | impl.get_num_new_matched_tokens.return_value = (10, False) |
| 179 | result = connector.get_num_new_matched_tokens(req, 5) |
| 180 | impl.get_num_new_matched_tokens.assert_called_once_with(req, 5) |
| 181 | assert result == (10, False) |
| 182 | |
| 183 | def test_update_state_after_alloc(self, connector_and_impl): |
| 184 | connector, impl = connector_and_impl |
| 185 | req = MagicMock() |
| 186 | blocks = MagicMock() |
| 187 | connector.update_state_after_alloc(req, blocks, 4) |
| 188 | impl.update_state_after_alloc.assert_called_once_with(req, blocks, 4) |
| 189 | |
| 190 | def test_build_connector_meta(self, connector_and_impl): |
| 191 | connector, impl = connector_and_impl |
| 192 | sched_out = MagicMock() |
| 193 | connector.build_connector_meta(sched_out) |
| 194 | impl.build_connector_meta.assert_called_once_with(sched_out) |
| 195 | |
| 196 | def test_update_connector_output(self, connector_and_impl): |
| 197 | connector, impl = connector_and_impl |
| 198 | out = MagicMock() |
| 199 | connector.update_connector_output(out) |
| 200 | impl.update_connector_output.assert_called_once_with(out) |
| 201 | |
| 202 | def test_request_finished(self, connector_and_impl): |
| 203 | connector, impl = connector_and_impl |
| 204 | req = MagicMock() |
| 205 | impl.request_finished.return_value = (True, {"key": "val"}) |
| 206 | result = connector.request_finished(req, [1, 2, 3]) |
| 207 | impl.request_finished.assert_called_once_with(req, [1, 2, 3]) |
| 208 | assert result == (True, {"key": "val"}) |
| 209 | |
| 210 | def test_take_events(self, connector_and_impl): |
| 211 | connector, impl = connector_and_impl |
| 212 | impl.take_events.return_value = iter([]) |
| 213 | list(connector.take_events()) |
| 214 | impl.take_events.assert_called_once() |
| 215 | |
| 216 | def test_get_kv_connector_stats(self, connector_and_impl): |
| 217 | connector, impl = connector_and_impl |
| 218 | impl.get_kv_connector_stats.return_value = None |
| 219 | result = connector.get_kv_connector_stats() |
| 220 | impl.get_kv_connector_stats.assert_called_once() |
| 221 | assert result is None |
| 222 | |
| 223 | def test_get_block_ids_with_load_errors(self, connector_and_impl): |
| 224 | connector, impl = connector_and_impl |
| 225 | impl.get_block_ids_with_load_errors.return_value = {7, 8} |
| 226 | result = connector.get_block_ids_with_load_errors() |
| 227 | assert result == {7, 8} |
| 228 | |
| 229 | def test_wait_for_layer_load(self, connector_and_impl): |
| 230 | connector, impl = connector_and_impl |
| 231 | connector.wait_for_layer_load("layer_0") |
| 232 | impl.wait_for_layer_load.assert_called_once_with("layer_0") |