Skip to content

Commit ec009ab

Browse files
committed
[UT] add ut for NONE_HASH
Signed-off-by: Andy Xie <andy.xning@gmail.com>
1 parent 5b2dcbf commit ec009ab

File tree

1 file changed

+31
-10
lines changed

1 file changed

+31
-10
lines changed

tests/v1/core/test_kv_cache_utils.py

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# SPDX-License-Identifier: Apache-2.0
2+
import importlib
23

34
import pytest
45
import torch
@@ -10,8 +11,7 @@
1011
from vllm.v1.core.kv_cache_manager import KVCacheManager
1112
# disable yapf here as it formats differently than isort such that both fail
1213
# yapf: disable
13-
from vllm.v1.core.kv_cache_utils import (NONE_HASH, BlockHashType,
14-
FreeKVCacheBlockQueue, KVCacheBlock,
14+
from vllm.v1.core.kv_cache_utils import (FreeKVCacheBlockQueue, KVCacheBlock,
1515
PrefixCachingMetrics,
1616
estimate_max_model_len,
1717
generate_block_hash_extra_keys,
@@ -62,13 +62,29 @@ def new_kv_cache_spec(block_size=16,
6262
use_mla=use_mla)
6363

6464

65-
def test_none_hash():
66-
assert NONE_HASH is not None
67-
assert isinstance(NONE_HASH, int)
68-
assert NONE_HASH != 0
65+
def test_none_hash(monkeypatch):
66+
import vllm.v1.core.kv_cache_utils
67+
68+
# case 1: PYTHONHASHSEED is not set, use random
69+
with monkeypatch.context() as m:
70+
m.delenv('PYTHONHASHSEED', raising=False)
71+
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
72+
assert reloaded_kv_cache_utils.NONE_HASH is not None
73+
assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
74+
assert reloaded_kv_cache_utils.NONE_HASH != 0
75+
76+
# case 2: PYTHONHASHSEED is set, use the seed
77+
with monkeypatch.context() as m:
78+
m.setenv('PYTHONHASHSEED', 'python hash seed')
79+
reloaded_kv_cache_utils = importlib.reload(vllm.v1.core.kv_cache_utils)
80+
assert reloaded_kv_cache_utils.NONE_HASH is not None
81+
assert isinstance(reloaded_kv_cache_utils.NONE_HASH, int)
82+
assert sha256('python hash seed') == reloaded_kv_cache_utils.NONE_HASH
6983

7084

7185
def test_kv_cache_block():
86+
import vllm.v1.core.kv_cache_utils
87+
7288
# Test KVCacheBlock initialization
7389
block = KVCacheBlock(block_id=0)
7490
assert block.block_id == 0
@@ -82,7 +98,8 @@ def test_kv_cache_block():
8298
assert block.ref_cnt == 0
8399

84100
# Test block hash setting and resetting
85-
block_hash = BlockHashType(hash_value=123, token_ids=(1, 2, 3))
101+
block_hash = vllm.v1.core.kv_cache_utils.BlockHashType(hash_value=123,
102+
token_ids=(1, 2, 3))
86103
block.block_hash = block_hash
87104
assert block.block_hash == block_hash
88105

@@ -256,13 +273,14 @@ def test_generate_block_hash_extra_keys_cache_salt():
256273

257274
@pytest.mark.parametrize("hash_fn", [sha256, hash])
258275
def test_hash_block_tokens(hash_fn):
276+
import vllm.v1.core.kv_cache_utils
259277
parent_block_hash = 123
260278
curr_block_token_ids = (1, 2, 3)
261279
extra_keys = ("key1", "key2")
262280

263281
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
264282
curr_block_token_ids, extra_keys)
265-
assert isinstance(block_hash, BlockHashType)
283+
assert isinstance(block_hash, vllm.v1.core.kv_cache_utils.BlockHashType)
266284
assert block_hash.hash_value == hash_fn(
267285
(parent_block_hash, curr_block_token_ids, extra_keys))
268286
assert block_hash.token_ids == curr_block_token_ids
@@ -271,6 +289,7 @@ def test_hash_block_tokens(hash_fn):
271289

272290
@pytest.mark.parametrize("hash_fn", [sha256, hash])
273291
def test_hash_request_tokens(hash_fn):
292+
import vllm.v1.core.kv_cache_utils
274293
request = make_request(
275294
request_id=0,
276295
prompt_token_ids=[_ for _ in range(6)],
@@ -285,8 +304,10 @@ def test_hash_request_tokens(hash_fn):
285304
block_hashes = hash_request_tokens(hash_fn, block_size, request)
286305

287306
assert len(block_hashes) == 2
288-
assert isinstance(block_hashes[0], BlockHashType)
289-
assert isinstance(block_hashes[1], BlockHashType)
307+
assert isinstance(block_hashes[0],
308+
vllm.v1.core.kv_cache_utils.BlockHashType)
309+
assert isinstance(block_hashes[1],
310+
vllm.v1.core.kv_cache_utils.BlockHashType)
290311

291312
# Check the first block
292313
assert block_hashes[0].token_ids == (0, 1, 2)

0 commit comments

Comments
 (0)