A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Fix missing `clear_connector_metadata` (#25397)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Sep 23, 2025 at 02:10 UTC
090197034faf3b193c4467cedeb9281e3078892d
2 files changed
+60
tests/v1/kv_connector/unit/test_kv_connector_lifecyle.py
new
+59
@@ -0,0 +1,59 @@
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.shared_storage_connector import ( # noqa: E501
5
+ SharedStorageConnectorMetadata)
6
+from vllm.distributed.kv_transfer.kv_transfer_state import (
7
+ ensure_kv_transfer_initialized, get_kv_transfer_group)
8
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
9
+from vllm.v1.worker.kv_connector_model_runner_mixin import (
10
+ KVConnectorModelRunnerMixin)
11
+
12
+# Importing utils registers TestSharedStorageConnector with the factory
13
+from .utils import create_vllm_config
14
+
15
+
16
+def _make_empty_scheduler_output():
17
+ return SchedulerOutput(
18
+ scheduled_new_reqs=[],
19
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
20
+ num_scheduled_tokens={},
21
+ total_num_scheduled_tokens=0,
22
+ scheduled_spec_decode_tokens={},
23
+ scheduled_encoder_inputs={},
24
+ num_common_prefix_blocks=[],
25
+ finished_req_ids=set(),
26
+ free_encoder_mm_hashes=[],
27
+ structured_output_request_ids={},
28
+ grammar_bitmask=None,
29
+ kv_connector_metadata=SharedStorageConnectorMetadata(),
30
+ )
31
+
32
+
33
+def test_kv_connector_mixin_clears_metadata():
34
+ vllm_config = create_vllm_config()
35
+ vllm_config.kv_transfer_config.kv_connector = "TestSharedStorageConnector"
36
+ vllm_config.kv_transfer_config.kv_role = "kv_both"
37
+ vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = ("unit")
38
+
39
+ # Initialize the global connector instance
40
+ ensure_kv_transfer_initialized(vllm_config)
41
+
42
+ try:
43
+ # Minimal scheduler output with empty metadata; mixin should still
44
+ # bind/clear metadata even if no loads happen
45
+ scheduler_output = _make_empty_scheduler_output()
46
+
47
+ # Invoke the no-forward path which uses the mixin context manager
48
+ KVConnectorModelRunnerMixin.kv_connector_no_forward(
49
+ scheduler_output, vllm_config)
50
+
51
+ # Verify clear_connector_metadata was called on the connector
52
+ connector = get_kv_transfer_group()
53
+ assert connector._connector_metadata is None
54
+ # Test connector wrapper records method calls
55
+ assert connector.call_record.get("bind_connector_metadata", 0) == 1
56
+ assert connector.call_record.get("clear_connector_metadata", 0) == 1
57
+ finally:
58
+ # Ensure we clean up the global connector between tests
59
+ KVConnectorModelRunnerMixin.ensure_kv_transfer_shutdown()
vllm/v1/worker/kv_connector_model_runner_mixin.py
+1
@@ -123,6 +123,7 @@ class KVConnectorModelRunnerMixin:
123
124
output.kv_connector_stats = KVConnectorModelRunnerMixin.\
125
get_kv_connector_stats()
126
+ kv_connector.clear_connector_metadata()
127
128
@staticmethod
129
def get_kv_connector_stats() -> Optional[KVConnectorStats]: