[BugFix] Fix MultiConnector test after HMA changes (#19291)

Signed-off-by: Nick Hill <nhill@redhat.com>
This commit is contained in:
Nick Hill 2025-06-06 13:16:24 -07:00 committed by GitHub
parent 94ecee6282
commit aad30bd306
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 14 additions and 9 deletions

View File

@ -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

View File

@ -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,

View File

@ -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: