mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-01-23 18:54:33 +08:00
[BugFix] Fix MultiConnector test after HMA changes (#19291)
Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
parent
94ecee6282
commit
aad30bd306
@ -60,7 +60,8 @@ class TestSharedStorageConnector(SharedStorageConnector):
|
||||
if isinstance(arg, int):
|
||||
to_log.append(str(arg))
|
||||
elif isinstance(arg, KVCacheBlocks):
|
||||
to_log.append(f"num_blocks={len(arg.blocks)}")
|
||||
to_log.append(
|
||||
f"num_blocks={[len(b) for b in arg.blocks]}")
|
||||
|
||||
# Log the event as a line to the file
|
||||
try:
|
||||
@ -176,7 +177,7 @@ def test_multi_shared_storage_connector_consistency():
|
||||
# on each connector in turn.
|
||||
assert events["storage1-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
|
||||
]
|
||||
assert events["storage1-WORKER"][:5] == [
|
||||
'register_kv_caches', 'bind_connector_metadata', 'start_load_kv',
|
||||
@ -184,7 +185,7 @@ def test_multi_shared_storage_connector_consistency():
|
||||
]
|
||||
assert events["storage2-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
|
||||
]
|
||||
assert events["storage2-WORKER"][:5] == [
|
||||
'register_kv_caches', 'bind_connector_metadata', 'start_load_kv',
|
||||
@ -205,11 +206,11 @@ def test_multi_shared_storage_connector_consistency():
|
||||
# chosen).
|
||||
assert events["storage1-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=7 96', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta'
|
||||
]
|
||||
assert events["storage2-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
|
||||
]
|
||||
|
||||
# Delete storage1 connector state
|
||||
@ -229,11 +230,11 @@ def test_multi_shared_storage_connector_consistency():
|
||||
# blocks for the second connector.
|
||||
assert events["storage1-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=0 0', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta'
|
||||
]
|
||||
assert events["storage2-SCHEDULER"][:3] == [
|
||||
'get_num_new_matched_tokens 0',
|
||||
'update_state_after_alloc num_blocks=7 96', 'build_connector_meta'
|
||||
'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta'
|
||||
]
|
||||
|
||||
# Clean up
|
||||
|
||||
@ -153,6 +153,7 @@ class MultiConnector(KVConnectorBase_V1):
|
||||
num_external_tokens: int):
|
||||
chosen_connector = self._requests_to_connector.get(
|
||||
request.request_id, -1)
|
||||
empty_blocks = blocks.new_empty()
|
||||
for i, c in enumerate(self._connectors):
|
||||
if i == chosen_connector:
|
||||
# Forward call to the chosen connector (if any).
|
||||
@ -160,8 +161,7 @@ class MultiConnector(KVConnectorBase_V1):
|
||||
num_external_tokens)
|
||||
else:
|
||||
# Call with empty blocks for other connectors.
|
||||
c.update_state_after_alloc(request,
|
||||
KVCacheBlocks.create_empty(), 0)
|
||||
c.update_state_after_alloc(request, empty_blocks, 0)
|
||||
|
||||
def build_connector_meta(
|
||||
self,
|
||||
|
||||
@ -61,6 +61,10 @@ class KVCacheBlocks:
|
||||
if block.block_hash is None
|
||||
]
|
||||
|
||||
def new_empty(self) -> "KVCacheBlocks":
|
||||
"""Creates a new KVCacheBlocks instance with no blocks."""
|
||||
return KVCacheBlocks([[] for _ in range(len(self.blocks))])
|
||||
|
||||
|
||||
class KVCacheManager:
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user