mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-10 12:05:48 +08:00
[Bugfix] Fix missing clear_connector_metadata (#25397)
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
parent
f31ff87460
commit
090197034f
59
tests/v1/kv_connector/unit/test_kv_connector_lifecyle.py
Normal file
59
tests/v1/kv_connector/unit/test_kv_connector_lifecyle.py
Normal file
@ -0,0 +1,59 @@
|
|||||||
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||||
|
|
||||||
|
from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa: E501
|
||||||
|
SharedStorageConnectorMetadata)
|
||||||
|
from vllm.distributed.kv_transfer.kv_transfer_state import (
|
||||||
|
ensure_kv_transfer_initialized, get_kv_transfer_group)
|
||||||
|
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
|
||||||
|
from vllm.v1.worker.kv_connector_model_runner_mixin import (
|
||||||
|
KVConnectorModelRunnerMixin)
|
||||||
|
|
||||||
|
# Importing utils registers TestSharedStorageConnector with the factory
|
||||||
|
from .utils import create_vllm_config
|
||||||
|
|
||||||
|
|
||||||
|
def _make_empty_scheduler_output():
|
||||||
|
return SchedulerOutput(
|
||||||
|
scheduled_new_reqs=[],
|
||||||
|
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||||
|
num_scheduled_tokens={},
|
||||||
|
total_num_scheduled_tokens=0,
|
||||||
|
scheduled_spec_decode_tokens={},
|
||||||
|
scheduled_encoder_inputs={},
|
||||||
|
num_common_prefix_blocks=[],
|
||||||
|
finished_req_ids=set(),
|
||||||
|
free_encoder_mm_hashes=[],
|
||||||
|
structured_output_request_ids={},
|
||||||
|
grammar_bitmask=None,
|
||||||
|
kv_connector_metadata=SharedStorageConnectorMetadata(),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_kv_connector_mixin_clears_metadata():
|
||||||
|
vllm_config = create_vllm_config()
|
||||||
|
vllm_config.kv_transfer_config.kv_connector = "TestSharedStorageConnector"
|
||||||
|
vllm_config.kv_transfer_config.kv_role = "kv_both"
|
||||||
|
vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = ("unit")
|
||||||
|
|
||||||
|
# Initialize the global connector instance
|
||||||
|
ensure_kv_transfer_initialized(vllm_config)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Minimal scheduler output with empty metadata; mixin should still
|
||||||
|
# bind/clear metadata even if no loads happen
|
||||||
|
scheduler_output = _make_empty_scheduler_output()
|
||||||
|
|
||||||
|
# Invoke the no-forward path which uses the mixin context manager
|
||||||
|
KVConnectorModelRunnerMixin.kv_connector_no_forward(
|
||||||
|
scheduler_output, vllm_config)
|
||||||
|
|
||||||
|
# Verify clear_connector_metadata was called on the connector
|
||||||
|
connector = get_kv_transfer_group()
|
||||||
|
assert connector._connector_metadata is None
|
||||||
|
# Test connector wrapper records method calls
|
||||||
|
assert connector.call_record.get("bind_connector_metadata", 0) == 1
|
||||||
|
assert connector.call_record.get("clear_connector_metadata", 0) == 1
|
||||||
|
finally:
|
||||||
|
# Ensure we clean up the global connector between tests
|
||||||
|
KVConnectorModelRunnerMixin.ensure_kv_transfer_shutdown()
|
||||||
@ -123,6 +123,7 @@ class KVConnectorModelRunnerMixin:
|
|||||||
|
|
||||||
output.kv_connector_stats = KVConnectorModelRunnerMixin.\
|
output.kv_connector_stats = KVConnectorModelRunnerMixin.\
|
||||||
get_kv_connector_stats()
|
get_kv_connector_stats()
|
||||||
|
kv_connector.clear_connector_metadata()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def get_kv_connector_stats() -> Optional[KVConnectorStats]:
|
def get_kv_connector_stats() -> Optional[KVConnectorStats]:
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user