diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index df487ec2ccaa..cc73b1444955 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +import importlib import pytest import torch @@ -10,8 +11,7 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager # disable yapf here as it formats differently than isort such that both fail # yapf: disable -from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType, - FreeKVCacheBlockQueue, KVCacheBlock, +from vllm.v1.core.kv_cache_utils import (FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics, estimate_max_model_len, generate_block_hash_extra_keys, @@ -62,13 +62,29 @@ def new_kv_cache_spec(block_size=16, use_mla=use_mla) -def test_none_hash(): - assert NONE_HASH is not None - assert isinstance(NONE_HASH, int) - assert NONE_HASH != 0 +def test_none_hash(monkeypatch): + import vllm.v1.core.kv_cache_utils + + # case 1: PYTHONHASHSEED is not set, use random + with monkeypatch.context() as m: + m.delenv('PYTHONHASHSEED', raising=False) + reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils) + assert reloaded_kv_cache_utils.NONE_HASH is not None + assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int) + assert reloaded_kv_cache_utils.NONE_HASH != 0 + + # case 2: PYTHONHASHSEED is set, use the seed + with monkeypatch.context() as m: + m.setenv('PYTHONHASHSEED', 'python hash seed') + reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils) + assert reloaded_kv_cache_utils.NONE_HASH is not None + assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int) + assert sha256('python hash seed') == reloaded_kv_cache_utils.NONE_HASH def test_kv_cache_block(): + import vllm.v1.core.kv_cache_utils + # Test KVCacheBlock initialization block = KVCacheBlock(block_id=0) assert block.block_id == 0 @@ -82,7 +98,8 @@ def test_kv_cache_block(): assert block.ref_cnt == 0 # Test block hash setting and resetting - block_hash = BlockHashType(hash_value=123, token_ids=(1, 2, 3)) + block_hash = vllm.v1.core.kv_cache_utils.BlockHashType(hash_value=123, + token_ids=(1, 2, 3)) block.block_hash = block_hash assert block.block_hash == block_hash @@ -256,13 +273,14 @@ def test_generate_block_hash_extra_keys_cache_salt(): @pytest.mark.parametrize("hash_fn", [sha256, hash]) def test_hash_block_tokens(hash_fn): + import vllm.v1.core.kv_cache_utils parent_block_hash = 123 curr_block_token_ids = (1, 2, 3) extra_keys = ("key1", "key2") block_hash = hash_block_tokens(hash_fn, parent_block_hash, curr_block_token_ids, extra_keys) - assert isinstance(block_hash, BlockHashType) + assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHashType) assert block_hash.hash_value == hash_fn( (parent_block_hash, curr_block_token_ids, extra_keys)) assert block_hash.token_ids == curr_block_token_ids @@ -271,6 +289,7 @@ def test_hash_block_tokens(hash_fn): @pytest.mark.parametrize("hash_fn", [sha256, hash]) def test_hash_request_tokens(hash_fn): + import vllm.v1.core.kv_cache_utils request = make_request( request_id=0, prompt_token_ids=[_ for _ in range(6)], @@ -285,8 +304,10 @@ def test_hash_request_tokens(hash_fn): block_hashes = hash_request_tokens(hash_fn, block_size, request) assert len(block_hashes) == 2 - assert isinstance(block_hashes[0], BlockHashType) - assert isinstance(block_hashes[1], BlockHashType) + assert isinstance(block_hashes[0], + vllm.v1.core.kv_cache_utils.BlockHashType) + assert isinstance(block_hashes[1], + vllm.v1.core.kv_cache_utils.BlockHashType) # Check the first block assert block_hashes[0].token_ids == (0, 1, 2)