fix: malformed KV events for NVIDIA Dynamo (#13488)
Signed-off-by: PeaBrane <yanrpei@gmail.com> Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
Yan Ru Pei
co-authored by
ishandhanani
Zhiqiang Xie
ishandhanani
Zhiqiang Xie
parent
a102a0507a
commit
f446b51c41
+70
-155
@@ -77,166 +77,81 @@ class TestKvEvents(CustomTestCase):
|
||||
},
|
||||
)
|
||||
|
||||
# Expected events. These may be dependent on model used (meta-llama/Llama-3.2-1B-Instruct)
|
||||
expected_events = [
|
||||
# <begin> The capital city of France is
|
||||
BlockStored(
|
||||
block_hashes=[-6650323075460941099],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[128000, 791, 6864, 3363, 315, 9822, 374],
|
||||
block_size=7,
|
||||
lora_id=None,
|
||||
),
|
||||
# Paris. The Eiffel Tower
|
||||
BlockStored(
|
||||
block_hashes=[-7584018293207282755],
|
||||
parent_block_hash=-6650323075460941099,
|
||||
token_ids=[12366, 13, 578, 469, 3168, 301, 22703],
|
||||
block_size=7,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockStored(
|
||||
block_hashes=[-8753497827991233192],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[0],
|
||||
block_size=1,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-6650323075460941099]),
|
||||
# <begin> The capital
|
||||
BlockStored(
|
||||
block_hashes=[-2697055055087824455],
|
||||
parent_block_hash=5740354900026072187,
|
||||
token_ids=[128000, 791, 6864],
|
||||
block_size=3,
|
||||
lora_id=None,
|
||||
),
|
||||
# city of France is
|
||||
BlockStored(
|
||||
block_hashes=[-7505627135785778022],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[3363, 315, 9822, 374],
|
||||
block_size=4,
|
||||
lora_id=None,
|
||||
),
|
||||
# of France is
|
||||
BlockStored(
|
||||
block_hashes=[-3861108700662737012],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[315, 9822, 374],
|
||||
block_size=3,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-7584018293207282755]),
|
||||
BlockRemoved(block_hashes=[-8753497827991233192]),
|
||||
BlockRemoved(block_hashes=[-7505627135785778022]),
|
||||
# Paris. The Eiffel Tower is located in Paris. The Eiffel Tower is a famous landmark in Paris
|
||||
BlockStored(
|
||||
block_hashes=[-3064341286825792715],
|
||||
parent_block_hash=-3861108700662737012,
|
||||
token_ids=[
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
469,
|
||||
3168,
|
||||
301,
|
||||
22703,
|
||||
374,
|
||||
7559,
|
||||
304,
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
469,
|
||||
3168,
|
||||
301,
|
||||
22703,
|
||||
374,
|
||||
264,
|
||||
11495,
|
||||
38350,
|
||||
304,
|
||||
12366,
|
||||
],
|
||||
block_size=23,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-3861108700662737012]),
|
||||
# of
|
||||
BlockStored(
|
||||
block_hashes=[6115672085296369592],
|
||||
parent_block_hash=-2697055055087824455,
|
||||
token_ids=[315],
|
||||
block_size=1,
|
||||
lora_id=None,
|
||||
),
|
||||
# France is
|
||||
BlockStored(
|
||||
block_hashes=[4208810872343132234],
|
||||
parent_block_hash=6115672085296369592,
|
||||
token_ids=[9822, 374],
|
||||
block_size=2,
|
||||
lora_id=None,
|
||||
),
|
||||
# Spain is
|
||||
BlockStored(
|
||||
block_hashes=[1675819893649989955],
|
||||
parent_block_hash=6115672085296369592,
|
||||
token_ids=[18157, 374],
|
||||
block_size=2,
|
||||
lora_id=None,
|
||||
),
|
||||
BlockRemoved(block_hashes=[-3064341286825792715]),
|
||||
# Madrid. The capital of France is Paris. The capital of Italy is Rome. The capital of Spain is Madrid.
|
||||
BlockStored(
|
||||
block_hashes=[-8505834929190027295],
|
||||
parent_block_hash=1675819893649989955,
|
||||
token_ids=[
|
||||
25048,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
9822,
|
||||
374,
|
||||
12366,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
15704,
|
||||
374,
|
||||
22463,
|
||||
13,
|
||||
578,
|
||||
6864,
|
||||
315,
|
||||
18157,
|
||||
374,
|
||||
25048,
|
||||
13,
|
||||
],
|
||||
block_size=23,
|
||||
lora_id=None,
|
||||
),
|
||||
]
|
||||
|
||||
# Get events
|
||||
events = []
|
||||
start = time.time()
|
||||
max_wait_s = 5
|
||||
while (
|
||||
len(events) < len(expected_events)
|
||||
and (time.time() - start) < max_wait_s
|
||||
):
|
||||
_, seq_bytes, payload = sub.recv_multipart()
|
||||
event_batch = decoder.decode(payload)
|
||||
for event in event_batch.events:
|
||||
events.append(event)
|
||||
min_events_expected = 5 # Expect at least some events
|
||||
|
||||
for expected in expected_events:
|
||||
self.assertIn(expected, events)
|
||||
while (
|
||||
len(events) < min_events_expected and (time.time() - start) < max_wait_s
|
||||
):
|
||||
if sub.poll(timeout=100): # 100ms timeout
|
||||
_, seq_bytes, payload = sub.recv_multipart()
|
||||
event_batch = decoder.decode(payload)
|
||||
for event in event_batch.events:
|
||||
events.append(event)
|
||||
|
||||
# Verify we received events
|
||||
self.assertGreater(
|
||||
len(events), 0, "Should have received at least one KV cache event"
|
||||
)
|
||||
|
||||
# Track which blocks were stored and removed
|
||||
stored_blocks = {} # hash -> BlockStored event
|
||||
removed_hashes = set()
|
||||
|
||||
# Validate event structure and relationships
|
||||
for event in events:
|
||||
self.assertIsInstance(
|
||||
event,
|
||||
(BlockStored, BlockRemoved, AllBlocksCleared),
|
||||
f"Event should be a KV cache event, got {type(event)}",
|
||||
)
|
||||
|
||||
if isinstance(event, BlockStored):
|
||||
# Validate BlockStored structure
|
||||
self.assertIsInstance(event.block_hashes, list)
|
||||
self.assertEqual(
|
||||
len(event.block_hashes), 1, "Should have one hash per block"
|
||||
)
|
||||
self.assertIsInstance(event.token_ids, list)
|
||||
self.assertEqual(
|
||||
event.block_size,
|
||||
len(event.token_ids),
|
||||
"block_size should match token_ids length",
|
||||
)
|
||||
self.assertIsNone(
|
||||
event.lora_id, "lora_id should be None for basic test"
|
||||
)
|
||||
|
||||
# Store this block for later validation
|
||||
block_hash = event.block_hashes[0]
|
||||
stored_blocks[block_hash] = event
|
||||
|
||||
# If parent_block_hash is set, verify it was stored earlier
|
||||
if event.parent_block_hash is not None:
|
||||
# Parent should either be in stored_blocks or could be from a previous request
|
||||
pass # Don't strictly enforce this as root blocks may have synthetic parents
|
||||
|
||||
elif isinstance(event, BlockRemoved):
|
||||
# Validate BlockRemoved structure
|
||||
self.assertIsInstance(event.block_hashes, list)
|
||||
self.assertEqual(
|
||||
len(event.block_hashes), 1, "Should have one hash per block"
|
||||
)
|
||||
removed_hashes.add(event.block_hashes[0])
|
||||
|
||||
# Verify we got both BlockStored and BlockRemoved events
|
||||
self.assertGreater(
|
||||
len(stored_blocks), 0, "Should have at least one BlockStored event"
|
||||
)
|
||||
# BlockRemoved events may not always occur in this short test, so just check if they do occur
|
||||
# that they reference previously stored blocks
|
||||
for removed_hash in removed_hashes:
|
||||
# It's OK if the removed block wasn't in our stored_blocks
|
||||
# (it could have been stored before we started listening)
|
||||
pass
|
||||
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
Reference in New Issue
Block a user