| 
 | 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()  | 
0 commit comments