Skip to content

Commit 864bbe3

Browse files
NickLuccheyewentao256
authored andcommitted
[Bugfix] Fix missing clear_connector_metadata (#25397)
Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com>
1 parent e97cf2e commit 864bbe3

File tree

2 files changed

+60
-0
lines changed

2 files changed

+60
-0
lines changed
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ def _get_kv_connector_output(
123123

124124
output.kv_connector_stats = KVConnectorModelRunnerMixin.\
125125
get_kv_connector_stats()
126+
kv_connector.clear_connector_metadata()
126127

127128
@staticmethod
128129
def get_kv_connector_stats() -> Optional[KVConnectorStats]:

0 commit comments

Comments
 (0)