|
| 1 | +import os |
| 2 | +import tempfile |
| 3 | +import pytest |
| 4 | +from utils import * |
| 5 | + |
| 6 | +server = ServerPreset.tinyllama2() |
| 7 | + |
| 8 | + |
| 9 | +class LogReader: |
| 10 | + def __init__(self, path): |
| 11 | + self.path = path |
| 12 | + self.pos = 0 |
| 13 | + |
| 14 | + def drain(self): |
| 15 | + with open(self.path, errors="ignore") as f: |
| 16 | + f.seek(self.pos) |
| 17 | + content = f.read() |
| 18 | + self.pos = f.tell() |
| 19 | + return content |
| 20 | + |
| 21 | + |
| 22 | +@pytest.fixture(autouse=True) |
| 23 | +def create_server(): |
| 24 | + global server |
| 25 | + server = ServerPreset.tinyllama2() |
| 26 | + # single slot: a conversation that is not currently loaded lives ONLY |
| 27 | + # in the prompt cache, which is what exposes the consume-on-load bug. |
| 28 | + server.n_slots = 1 |
| 29 | + server.n_predict = 4 |
| 30 | + server.temperature = 0.0 |
| 31 | + server.server_slots = True |
| 32 | + server.cache_ram = 100 |
| 33 | + # force every slot reuse through the prompt cache (save + load) instead |
| 34 | + # of the in-place LCP-similarity reuse, so load() is actually exercised. |
| 35 | + server.slot_prompt_similarity = 0 |
| 36 | + # isolate the cache save/load to get_available_slot (the idle-slot |
| 37 | + # clearing path would add another save and muddy the scenario). |
| 38 | + server.no_cache_idle_slots = True |
| 39 | + server.debug = True |
| 40 | + fd, server.log_path = tempfile.mkstemp(suffix='.log') |
| 41 | + os.close(fd) |
| 42 | + yield |
| 43 | + |
| 44 | + |
| 45 | +# A and C share a long common prefix (so a request for C matches A's cached |
| 46 | +# entry, f_keep >= 0.25), then diverge. B is unrelated to both. |
| 47 | +COMMON_AC = ( |
| 48 | + "Once upon a time in a quiet village by the sea there lived an old " |
| 49 | + "fisherman who every morning rowed his small wooden boat out past the " |
| 50 | + "harbour wall to cast his nets beneath the pale light of the rising sun." |
| 51 | +) |
| 52 | +CONV_A = COMMON_AC + ( |
| 53 | + " On this particular day he caught a silver fish that spoke to him and " |
| 54 | + "promised three wishes in exchange for its freedom and a safe return home." |
| 55 | +) |
| 56 | +CONV_C = COMMON_AC + ( |
| 57 | + " But the storm clouds gathered quickly that afternoon and the waves grew " |
| 58 | + "tall and angry as the wind tore the sails and scattered the frightened gulls." |
| 59 | +) |
| 60 | +CONV_B = ( |
| 61 | + "In a bustling city far inland a young clockmaker tinkered late into the " |
| 62 | + "night with brass gears and tiny springs trying to build a machine that " |
| 63 | + "could measure not the hours but the quiet weight of a person's memories." |
| 64 | +) |
| 65 | +# A continuation of A (strict superset). Used for A's return so that |
| 66 | +# n_past < task tokens and we avoid the identical-prompt path that, on |
| 67 | +# SWA / hybrid / recurrent models, cannot partially remove the final |
| 68 | +# token and would reset regardless of caching. |
| 69 | +CONV_A_CONT = CONV_A + ( |
| 70 | + " The fisherman closed his eyes and made his first wish very carefully." |
| 71 | +) |
| 72 | + |
| 73 | + |
| 74 | +def _total_prompt_tokens(res): |
| 75 | + t = res.body["timings"] |
| 76 | + return t["prompt_n"] + t["cache_n"] |
| 77 | + |
| 78 | + |
| 79 | +# A prompt-cache entry must survive being matched by a DIFFERENT conversation. |
| 80 | +# Regression test for load() consuming (erasing) the matched entry: with one |
| 81 | +# slot and three conversations, conversation A lives only in the cache while |
| 82 | +# conversation C — which shares A's prefix — is loaded. C's load must not |
| 83 | +# destroy A's entry, otherwise A pays a full re-prefill when it returns. |
| 84 | +def test_cache_entry_survives_cross_conversation_load(): |
| 85 | + global server |
| 86 | + server.start() |
| 87 | + log = LogReader(server.log_path) |
| 88 | + |
| 89 | + # 1) Conversation A, cold. Capture its full token length. |
| 90 | + res_a1 = server.make_request("POST", "/completion", data={ |
| 91 | + "prompt": CONV_A, |
| 92 | + "cache_prompt": True, |
| 93 | + }) |
| 94 | + assert res_a1.status_code == 200 |
| 95 | + assert res_a1.body["timings"]["cache_n"] == 0 # nothing cached yet |
| 96 | + n_tokens_a = _total_prompt_tokens(res_a1) |
| 97 | + |
| 98 | + # 2) Conversation B (unrelated). Selecting the slot saves A into the |
| 99 | + # cache; B does not match A, so A is parked in the cache untouched. |
| 100 | + res_b = server.make_request("POST", "/completion", data={ |
| 101 | + "prompt": CONV_B, |
| 102 | + "cache_prompt": True, |
| 103 | + }) |
| 104 | + assert res_b.status_code == 200 |
| 105 | + assert "updating prompt cache" in log.drain() |
| 106 | + |
| 107 | + # 3) Conversation C, which shares A's long prefix. Selecting the slot |
| 108 | + # saves B; loading C matches A's cached entry (f_keep >= 0.25). The |
| 109 | + # buggy behaviour erased A here; the fix keeps it. |
| 110 | + res_c = server.make_request("POST", "/completion", data={ |
| 111 | + "prompt": CONV_C, |
| 112 | + "cache_prompt": True, |
| 113 | + }) |
| 114 | + assert res_c.status_code == 200 |
| 115 | + assert res_c.body["timings"]["cache_n"] > 0 # C reused A's shared prefix |
| 116 | + |
| 117 | + # 4) Conversation A returns (as a strict superset, so n_past < task |
| 118 | + # tokens). It was only in the cache. With the fix its entry survived |
| 119 | + # step 3, so all of A is reused and only the new continuation is |
| 120 | + # processed. Without the fix A's entry was consumed in step 3 and |
| 121 | + # only the prefix A shares with C can be reused. |
| 122 | + res_a2 = server.make_request("POST", "/completion", data={ |
| 123 | + "prompt": CONV_A_CONT, |
| 124 | + "cache_prompt": True, |
| 125 | + }) |
| 126 | + assert res_a2.status_code == 200 |
| 127 | + cache_n_a2 = res_a2.body["timings"]["cache_n"] |
| 128 | + |
| 129 | + # The full original A prompt is reused from cache. Under the bug this |
| 130 | + # would only be the A/C shared prefix, which is well below n_tokens_a. |
| 131 | + assert cache_n_a2 >= n_tokens_a - 2 |
0 commit comments