Support SHA256 as hash function in prefix caching (#15297)

Signed-off-by: Marko Rosenmueller <5467316+dr75@users.noreply.github.com>
This commit is contained in:
marko
2025-03-26 19:11:28 +01:00
committed by GitHub
parent 35fad35a48
commit 27df5199d9
12 changed files with 214 additions and 71 deletions

View File

@@ -2,6 +2,8 @@
# ruff: noqa
import asyncio
import hashlib
import pickle
import socket
from collections.abc import AsyncIterator
from unittest.mock import patch
@@ -14,7 +16,8 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.utils import (FlexibleArgumentParser, MemorySnapshot,
PlaceholderModule, StoreBoolean, bind_kv_cache,
deprecate_kwargs, get_open_port, memory_profiling,
merge_async_iterators, supports_kw, swap_dict_values)
merge_async_iterators, sha256, supports_kw,
swap_dict_values)
from .utils import create_new_process_for_each_test, error_on_warning
@@ -476,3 +479,21 @@ def test_swap_dict_values(obj, key1, key2):
assert obj[key1] == original_obj[key2]
else:
assert key1 not in obj
@pytest.mark.parametrize("input", [(), ("abc", ), (None, ),
(None, bool, [1, 2, 3])])
@pytest.mark.parametrize("output", [0, 1, 2])
def test_sha256(input: tuple, output: int):
hash = sha256(input)
assert hash is not None
assert isinstance(hash, int)
assert hash != 0
bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
assert hash == int.from_bytes(hashlib.sha256(bytes).digest(), byteorder="big")
# hashing again, returns the same value
assert hash == sha256(input)
# hashing different input, returns different value
assert hash != sha256(input + (1, ))