mirror of
https://github.com/vllm-project/vllm.git
synced 2025-10-20 14:53:52 +08:00
[Bugfix] Fix missing clear_connector_metadata
(#25397)
Signed-off-by: NickLucche <nlucches@redhat.com>
This commit is contained in:
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.\
|
||||
get_kv_connector_stats()
|
||||
kv_connector.clear_connector_metadata()
|
||||
|
||||
@staticmethod
|
||||
def get_kv_connector_stats() -> Optional[KVConnectorStats]:
|
||||
|
Reference in New Issue
Block a user