File size: 3,521 Bytes
8efb28e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | import os
import tempfile
import pytest
from utils import *
server = ServerPreset.tinyllama2()
class LogReader:
def __init__(self, path):
self.path = path
self.pos = 0
def drain(self):
with open(self.path) as f:
f.seek(self.pos)
content = f.read()
self.pos = f.tell()
return content
@pytest.fixture(autouse=True)
def create_server():
global server
server = ServerPreset.tinyllama2()
server.n_slots = 2
server.n_predict = 4
server.temperature = 0.0
server.server_slots = True
server.cache_ram = 100
server.kv_unified = True
server.debug = True
fd, server.log_path = tempfile.mkstemp(suffix='.log')
os.close(fd)
yield
LONG_PROMPT = (
"Once upon a time in a land far away, there lived a brave knight "
"who traveled across mountains and rivers to find the legendary "
"golden sword hidden deep within the enchanted forest of whispers. "
"He met many creatures along the way including dragons and fairies "
"and wizards who helped him on his noble quest to save the kingdom."
)
# idle slot cleared on launch should restore from cache-ram
def test_clear_and_restore():
global server
server.start()
log = LogReader(server.log_path)
# verify feature is enabled
assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log.drain()
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
original_prompt_n = res.body["timings"]["prompt_n"]
# Slot 0 is the only slot with KV — should NOT be cleared
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
# Launching slot 1 clears idle slot 0
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" in log.drain()
# Re-send same prompt — should restore from cache-ram
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"cache_prompt": True,
})
assert res.status_code == 200
assert "updating prompt cache" in log.drain()
assert res.body["timings"]["cache_n"] > 0
assert res.body["timings"]["prompt_n"] < original_prompt_n
# Follow-up — slot 0 kept its KV, no clearing needed
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT + " The knight finally reached the castle gates.",
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
def test_disabled_with_flag():
global server
server.no_cache_idle_slots = True
server.start()
log = LogReader(server.log_path)
# Feature should not be enabled
assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" not in log.drain()
res = server.make_request("POST", "/completion", data={
"prompt": LONG_PROMPT,
"id_slot": 0,
"cache_prompt": True,
})
assert res.status_code == 200
# Request on different slot — should NOT trigger clearing
res = server.make_request("POST", "/completion", data={
"prompt": "The quick brown fox",
"id_slot": 1,
"cache_prompt": True,
})
assert res.status_code == 200
assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log.drain()
|