wehub-resource-sync
139 lines
3.7 KiB
Python
139 lines
3.7 KiB
Python
import pytest
|
|
|
|
from mlc_llm.serve import PagedRadixTree
|
|
|
|
# category "runtime_module"
|
|
pytestmark = [pytest.mark.unittest]
|
|
|
|
|
|
def test_add():
|
|
prt = PagedRadixTree()
|
|
prt.add(0)
|
|
assert list(prt.get(0)) == []
|
|
prt.add(1)
|
|
assert list(prt.get(1)) == []
|
|
|
|
|
|
def test_remove():
|
|
prt = PagedRadixTree()
|
|
capacity = prt.free_capacity()
|
|
prt.add(0)
|
|
prt.remove(0)
|
|
prt.add(0)
|
|
prt.extend(0, [1 for _ in range(200)])
|
|
prt.remove(0)
|
|
assert prt.free_capacity() == capacity
|
|
|
|
prt.add(1)
|
|
prt.extend(1, [1 for _ in range(200)])
|
|
capacity = prt.free_capacity()
|
|
prt.add(2)
|
|
prt.extend(2, [1 for _ in range(100)] + [2 for _ in range(100)])
|
|
prt.remove(2)
|
|
assert prt.free_capacity() == capacity
|
|
|
|
prt.add(3)
|
|
prt.extend(3, [1 for _ in range(200)])
|
|
prt.remove(3)
|
|
assert prt.free_capacity() == capacity
|
|
|
|
prt.add(4)
|
|
prt.add(5)
|
|
prt.add(6)
|
|
assert prt.free_capacity() == capacity
|
|
prt.remove(4)
|
|
assert prt.free_capacity() == capacity
|
|
prt.remove(5)
|
|
assert prt.free_capacity() == capacity
|
|
prt.remove(6)
|
|
assert prt.free_capacity() == capacity
|
|
|
|
|
|
def test_extend():
|
|
prt = PagedRadixTree()
|
|
L = prt.free_capacity() // 64
|
|
H = L // 2
|
|
Q = L // 4
|
|
seq_id = 0
|
|
for start_pos in [0, H, L, L + H]:
|
|
for length in [Q, L - H, L, 2 * L - H, 2 * L]:
|
|
prt.add(seq_id)
|
|
if start_pos:
|
|
tokens_1 = [seq_id for _ in range(start_pos)]
|
|
prt.extend(seq_id, tokens_1)
|
|
assert list(prt.get(seq_id)) == tokens_1
|
|
else:
|
|
tokens_1 = []
|
|
tokens_2 = [seq_id for _ in range(length)]
|
|
prt.extend(seq_id, tokens_2)
|
|
assert list(prt.get(seq_id)) == tokens_1 + tokens_2
|
|
seq_id += 1
|
|
|
|
|
|
def test_fork():
|
|
prt = PagedRadixTree()
|
|
L = prt.free_capacity() // 64
|
|
H = L // 2
|
|
Q = L // 4
|
|
seq_id = 0
|
|
length_list = [Q, H, L, L + Q, L + H, L * 2]
|
|
for p_idx in range(1, len(length_list)):
|
|
for c_idx in range(0, p_idx + 1):
|
|
prt.add(seq_id)
|
|
tokens = [seq_id for _ in range(length_list[p_idx])]
|
|
prt.extend(seq_id, tokens)
|
|
prt.fork(seq_id + 1, seq_id, length_list[c_idx])
|
|
assert list(prt.get(seq_id + 1)) == tokens[: length_list[c_idx]]
|
|
seq_id += 2
|
|
|
|
|
|
def test_fork_2():
|
|
prt = PagedRadixTree()
|
|
prt.add(0)
|
|
prt.extend(0, [0, 1, 2, 3])
|
|
prt.fork(1, 0, 3)
|
|
prt.extend(1, [4])
|
|
prt.fork(2, 0, 3)
|
|
prt.extend(2, [5])
|
|
assert prt.match([0, 1, 2, 4]) == (4, (1,))
|
|
assert prt.match([0, 1, 2, 5]) == (4, (2,))
|
|
|
|
|
|
def test_rollback():
|
|
prt = PagedRadixTree()
|
|
L = prt.free_capacity() // 64
|
|
H = L // 2
|
|
Q = L // 4
|
|
seq_id = 0
|
|
for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
|
|
for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
|
|
if length > start_pos:
|
|
continue
|
|
prt.add(seq_id)
|
|
tokens = [seq_id for _ in range(start_pos)]
|
|
prt.extend(seq_id, tokens)
|
|
prt.rollback(seq_id, length)
|
|
assert list(prt.get(seq_id)) == tokens[:-length]
|
|
seq_id += 1
|
|
|
|
for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
|
|
for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
|
|
if length > start_pos:
|
|
continue
|
|
prt.add(seq_id)
|
|
tokens = [seq_id for _ in range(start_pos)]
|
|
prt.extend(seq_id, tokens)
|
|
prt.fork(seq_id + 1, seq_id, start_pos)
|
|
prt.rollback(seq_id + 1, length)
|
|
assert list(prt.get(seq_id + 1)) == tokens[:-length]
|
|
seq_id += 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_add()
|
|
test_remove()
|
|
test_extend()
|
|
test_fork()
|
|
test_fork_2()
|
|
test_rollback()
|