mirror of
https://github.com/Niko1221/Strata.git
synced 2026-11-08 08:12:16 -10:00
Add disk lifecycle parity gate and draft KV read-back checks
Prepare private-model tests for eviction-only writes, restart, admission denial, changed tokenizer identity and corrupted files. Require known answers, output/main-state parity, bounded retention and matching draft fingerprints. Verify the offline gate rejects incomplete or inconsistent evidence, including under optimized Python. Add opt-in byte-for-byte draft read-back after shared restore, including resident ring pages. Remove the GPU test refill that could hide a broken restore. Use bounded diagnostic read buffers and fail state hashing on transfer errors. Host transfer sanitizers pass 1188 checks and the five disk-gate verifier tests pass. Updated runtime/GPU test translation units pass syntax checking. Full-model and real-GPU execution remain pending an exclusive test window.
This commit is contained in:
1 parent
e7eb674431
commit
4d73bb393c
9 files changed
+462
-8
No files matched your search
@@ -125,3 +125,9 @@ allowance; these limits cover vector storage, not allocator or process RSS.
|
||||
|
||||
The integration is implemented but has only host tests and C++ syntax checks so
|
||||
far. Full-model restart, output/state parity and pressure validation remain open.
|
||||
|
||||
For model validation, `STRATA_SNAPSHOT_VERIFY=1` records which draft-prefill path
|
||||
ran and compares restored draft KV bytes with the saved image immediately after
|
||||
restore. The read-back checks authoritative storage and resident ring pages with
|
||||
64 KiB of workspace, then emits a fingerprint only on success. A mismatch or
|
||||
read-back failure stops the test engine. This diagnostic is off by default.
|
||||
@@ -41,6 +41,21 @@ key export. This is not a linked engine build. The frontend suite passes 67 test
|
||||
paths. Full-model NVMe promotion, restart, output parity and state hashes are still
|
||||
unverified; no running benchmark server or GPU was used for these checks.
|
||||
|
||||
The draft read-back diagnostic passes 1,188 host transfer checks under ASan/UBSan,
|
||||
including injected corruption and failed copies. The host backend wraps CUDA
|
||||
copies/synchronization; unused ring/stream GPU functions abort if reached. The
|
||||
real GPU fixture now checks the ring produced by restore directly, without an
|
||||
extra refill that could mask a bug; that updated GPU fixture has not run yet.
|
||||
State fingerprint reads now use fixed-size chunks and reject transfer failures.
|
||||
|
||||
`tools/conversation_cache_disk.py` prepares six sequential private engines for
|
||||
baseline, eviction, restart, admission denial, changed tokenizer identity and
|
||||
corrupted files. It requires known answers, token/main-state parity, matching
|
||||
draft read-back fingerprints across restart, and the requested draft-prefill
|
||||
path. Its five offline verifier tests pass, including 28 rejected evidence
|
||||
mutations and optimized Python. This establishes the verifier's checks, not
|
||||
full-model success. All model executions remain pending the exclusive test window.
|
||||
|
||||
```sh
|
||||
cmake -S . -B build-conversation-host -DSTRATA_ENABLE_CUDA=OFF -DSTRATA_ENABLE_HIP=OFF -DSTRATA_NATIVE_EXPERTS=OFF -DSTRATA_BUILD_TESTS=OFF -DSTRATA_BUILD_CONVERSATION_TESTS=ON -DSTRATA_ENABLE_CONVERSATION_DISK=ON
|
||||
cmake --build build-conversation-host --target conversation_file_test conversation_store_test conversation_cache_test conversation_memory_test -j 1
|
||||
@@ -102,6 +117,19 @@ at least 30 cycles and three lengths crossing resident KV and approaching the
|
||||
configured context limit. HTTP smoke requires an exclusive idle test endpoint.
|
||||
Run model/GPU gates only in an exclusive test window.
|
||||
|
||||
The disk gate is dry-run by default. Run it separately for batched draft prefill
|
||||
and the draft ring; each run owns a new output directory and modifies only its
|
||||
copied tokenizer and generated cache files:
|
||||
|
||||
```sh
|
||||
python tools/conversation_cache_disk.py --config CONFIG --engine ENGINE --output NEW_DIRECTORY --draft-path batched
|
||||
python tools/conversation_cache_disk.py --config CONFIG --engine ENGINE --output ANOTHER_NEW_DIRECTORY --draft-path ring
|
||||
```
|
||||
|
||||
Add `--run` in the exclusive window. Run INT8 and K8V4 configurations where
|
||||
supported. The gate records actual paths/modes and refuses to count an absent
|
||||
batched pass or ring restore as coverage.
|
||||
|
||||
## Outstanding evidence
|
||||
|
||||
- Repeat affected build/model gates on 0.1.26, including batched draft prompt KV
|
||||
|
||||
@@ -20,6 +20,11 @@ bool conversation_kv_validate(const ConversationKv& image, const QsaState& state
|
||||
int64_t upto, bool include_index, std::string& error);
|
||||
bool conversation_kv_restore(const ConversationKv& image, const QsaState& state, const ModelGeometry& g,
|
||||
int64_t upto, bool include_index, std::string& error);
|
||||
// Diagnostic read-back after a synchronized restore. Uses 64 KiB of stack
|
||||
// workspace, compares authoritative bytes and resident draft-ring pages, and
|
||||
// fingerprints the authoritative payload only. Never changes model state.
|
||||
bool conversation_kv_verify(const ConversationKv& image, const QsaState& state, const ModelGeometry& g,
|
||||
int64_t upto, bool include_index, uint64_t& fingerprint, std::string& error);
|
||||
|
||||
struct ConversationStateSizes {
|
||||
size_t gdn = 0, ple = 0, tail = 0, dead = 0, block_pos = 0;
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <array>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
|
||||
namespace strata::core {
|
||||
@@ -65,8 +66,8 @@ bool layout(const QsaState& st, const ModelGeometry& g, int64_t upto, bool index
|
||||
return true;
|
||||
}
|
||||
|
||||
std::array<void*, 5> pools(const QsaState& st) {
|
||||
const bool host = st.kv_mode != 0;
|
||||
std::array<void*, 5> pools(const QsaState& st, bool resident = false) {
|
||||
const bool host = st.kv_mode != 0 && !resident;
|
||||
if (st.kv_hybrid) return {st.k_q, st.v_q4, st.k_scale, nullptr, st.idx_pooled};
|
||||
if (st.kv_q4)
|
||||
return {host ? st.host.k_q4 : st.k_q4, host ? st.host.v_q4 : st.v_q4, nullptr, nullptr, st.idx_pooled};
|
||||
@@ -188,4 +189,42 @@ bool conversation_kv_restore(const ConversationKv& image, const QsaState& st, co
|
||||
return false;
|
||||
}
|
||||
|
||||
bool conversation_kv_verify(const ConversationKv& image, const QsaState& st, const ModelGeometry& g,
|
||||
int64_t upto, bool index, uint64_t& fingerprint, std::string& error) {
|
||||
if (!conversation_kv_validate(image, st, g, upto, index, error)) return false;
|
||||
std::array<uint8_t, 65536> buffer;
|
||||
uint64_t hash = 1469598103934665603ull;
|
||||
auto compare = [&](const void* device, const uint8_t* expected, size_t bytes, bool include_hash) {
|
||||
for (size_t offset = 0; offset < bytes;) {
|
||||
const size_t n = std::min(buffer.size(), bytes - offset);
|
||||
if (!transfer(buffer.data(), static_cast<const uint8_t*>(device) + offset, n, error)) return false;
|
||||
if (std::memcmp(buffer.data(), expected + offset, n) != 0) {
|
||||
error = "conversation snapshot: restored K/V bytes differ";
|
||||
return false;
|
||||
}
|
||||
if (include_hash) for (size_t i = 0; i < n; ++i) { hash ^= buffer[i]; hash *= 1099511628211ull; }
|
||||
offset += n;
|
||||
}
|
||||
return true;
|
||||
};
|
||||
const std::array<const std::vector<uint8_t>*, 5> saved = {&image.k, &image.v, &image.k_scale, &image.v_scale, &image.pooled};
|
||||
const auto authoritative = pools(st);
|
||||
for (size_t i = 0; i < saved.size(); ++i)
|
||||
if (!compare(authoritative[i], saved[i]->data(), saved[i]->size(), true)) return false;
|
||||
if (st.kv_mode == 2 && image.cells > 0) {
|
||||
const auto resident = pools(st, true);
|
||||
const int64_t end = image.cells / image.page_size;
|
||||
const int64_t begin = std::max<int64_t>(0, end - st.n_slots);
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
if (saved[i]->empty()) continue;
|
||||
const size_t page_bytes = saved[i]->size() / size_t(end);
|
||||
for (int64_t page = begin; page < end; ++page)
|
||||
if (!compare(static_cast<const uint8_t*>(resident[i]) + size_t(page % st.n_slots) * page_bytes,
|
||||
saved[i]->data() + size_t(page) * page_bytes, page_bytes, false)) return false;
|
||||
}
|
||||
}
|
||||
fingerprint = hash;
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace strata::core
|
||||
@@ -197,6 +197,8 @@ int main() {
|
||||
cuda_check(cudaDeviceSynchronize());
|
||||
check(conversation_kv_save(restored,f.state,f.g,upto,index,err),"read back restored A");
|
||||
check(equal(a,restored),"A/B/A byte-exact K/V and indexer state");
|
||||
uint64_t fingerprint = 0;
|
||||
check(conversation_kv_verify(a,f.state,f.g,upto,index,fingerprint,err),"verify restored authoritative and resident bytes without rewriting them");
|
||||
if (mode==1) {
|
||||
std::vector<int32_t> table((size_t)f.state.n_pages);
|
||||
cuda_check(cudaMemcpy(table.data(),f.state.page_table,table.size()*4,cudaMemcpyDeviceToHost));
|
||||
@@ -205,8 +207,6 @@ int main() {
|
||||
if (mode==2 && upto>0) {
|
||||
const auto s=qsa_real_shapes();
|
||||
const int64_t b1=(upto+s.page_size-1)/s.page_size,b0=std::max<int64_t>(0,b1-f.state.n_slots);
|
||||
kv_ring_restore(qsa_attn_pools(f.state),f.state.host,fmt,b0,b1,f.state.n_slots,s,nullptr);
|
||||
cuda_check(cudaDeviceSynchronize());
|
||||
const size_t block=kv_block_bytes(s,fmt)/2;
|
||||
// Codes and scales are separate in INT8; check code bytes here.
|
||||
const size_t code_block=fmt==kKvInt8 ? s.page_size*s.n_head_kv*s.head_dim : block;
|
||||
|
||||
@@ -171,6 +171,13 @@ void fixture(int format, int experts, bool zero_qsa, bool ple) {
|
||||
"injected host transfer backend completes a valid restore");
|
||||
check(gdn==image.live.gdn && history==image.live.ple && ss.ple_prev[0]==8 && ss.ple_prev[1]==9,
|
||||
"successful restore publishes correct running state and PLE window");
|
||||
uint64_t fingerprint = 0;
|
||||
check(conversation_kv_verify(image.kv.back(),draft.st,g,9,false,fingerprint,error), "read-back verifies complete draft payload");
|
||||
draft.data[0][0] ^= 1;
|
||||
check(!conversation_kv_verify(image.kv.back(),draft.st,g,9,false,fingerprint,error), "read-back detects corrupted draft byte");
|
||||
draft.data[0][0] ^= 1;
|
||||
fail_copy = copy_calls + 1;
|
||||
check(!conversation_kv_verify(image.kv.back(),draft.st,g,9,false,fingerprint,error), "read-back transfer failure is not a successful fingerprint");
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3462,6 +3462,9 @@ int main(int argc, char** argv) {
|
||||
// E-9: batched through the prompt path when it can (one GPU: a layer split's drafter is on the last stage)
|
||||
const bool batched = !multi_gpu && sp.draft_kv(mtp, R_rows, nxt.data(), T, p0, e);
|
||||
if (!e.empty() || (!batched && !mtp.prefill(R_rows, nxt.data(), T, p0, e))) return false;
|
||||
if (std::getenv("STRATA_SNAPSHOT_VERIFY") != nullptr)
|
||||
std::fprintf(stderr, "strata serve: DRAFT_PREFILL path=%s mode=%d cells=%lld\n",
|
||||
batched ? "batched" : "token", mtp.kv_state().kv_mode, (long long) T);
|
||||
// progress for the server window: PP <position reached> <prompt tokens> <ms> <fresh tokens/s>
|
||||
const int64_t done = p0 + T;
|
||||
const double ms = std::chrono::duration<double, std::milli>(Clock::now() - pp_t0).count();
|
||||
@@ -4013,6 +4016,17 @@ int main(int argc, char** argv) {
|
||||
std::printf("ERR restoring parked conversation: %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
if (std::getenv("STRATA_SNAPSHOT_VERIFY") != nullptr) {
|
||||
uint64_t draft_hash = 0;
|
||||
if (!strata::core::conversation_kv_verify(incoming->kv.back(), mtp.kv_state(), g,
|
||||
int64_t(incoming->live.ids.size()), false, draft_hash, err)) {
|
||||
std::printf("ERR verifying restored draft KV: %s\n", err.c_str());
|
||||
return 1;
|
||||
}
|
||||
std::fprintf(stderr, "strata serve: SNAPSHOT_VERIFY draft=%016llx cells=%lld mode=%d source=%s\n",
|
||||
(unsigned long long) draft_hash, (long long) incoming->kv.back().cells,
|
||||
mtp.kv_state().kv_mode, incoming_from_disk ? "disk" : "ram");
|
||||
}
|
||||
live = std::move(incoming->live.ids);
|
||||
live_imgs = std::move(incoming->live.imgs);
|
||||
checks = std::move(incoming->checkpoints);
|
||||
@@ -4459,17 +4473,30 @@ int main(int argc, char** argv) {
|
||||
if (state_hash && live_ok) {
|
||||
// DEBUG: a fingerprint of every part of the session over the positions it holds ([0, L)), and
|
||||
// separately of what lies past them in the last KV page (stale cells, fine unless something reads them)
|
||||
cudaDeviceSynchronize();
|
||||
if (cudaDeviceSynchronize() != cudaSuccess) {
|
||||
std::printf("ERR synchronizing state fingerprint\n");
|
||||
return 1;
|
||||
}
|
||||
const int64_t L = (int64_t) live.size();
|
||||
const strata::kernels::QsaShapes qs = [&] {
|
||||
strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes();
|
||||
s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_dim = g.idx_key_dim;
|
||||
return s;
|
||||
}();
|
||||
bool hash_ok = true;
|
||||
std::array<uint8_t, 65536> hash_buffer;
|
||||
auto hash_dev = [&](const void* p, size_t bytes, uint64_t h) {
|
||||
std::vector<uint8_t> b(bytes);
|
||||
if (bytes) cudaMemcpy(b.data(), p, bytes, cudaMemcpyDefault); // VRAM or a streamed host copy
|
||||
return fnv1a(b.data(), b.size(), h);
|
||||
for (size_t offset = 0; hash_ok && offset < bytes;) {
|
||||
const size_t n = std::min(hash_buffer.size(), bytes - offset);
|
||||
// VRAM or a streamed host copy, with fixed diagnostic workspace.
|
||||
if (cudaMemcpy(hash_buffer.data(), static_cast<const uint8_t*>(p) + offset, n, cudaMemcpyDefault) != cudaSuccess) {
|
||||
hash_ok = false;
|
||||
break;
|
||||
}
|
||||
h = fnv1a(hash_buffer.data(), n, h);
|
||||
offset += n;
|
||||
}
|
||||
return h;
|
||||
};
|
||||
// the cells [c0, c1) of one int8 K or V pool ([page][kv_head][page_size][head_dim]), bytes per value `w`
|
||||
auto hash_cells = [&](const void* pool, int64_t per_cell, int64_t c0, int64_t c1, uint64_t h) {
|
||||
@@ -4536,6 +4563,10 @@ int main(int argc, char** argv) {
|
||||
const int64_t mL = std::min<int64_t>(L, ms.max_cells);
|
||||
for (const auto& [pool, w] : kv_arrays(ms))
|
||||
if (pool != nullptr) h_mtp = hash_cells(pool, w, 0, mL, h_mtp);
|
||||
if (!hash_ok) {
|
||||
std::printf("ERR reading state fingerprint\n");
|
||||
return 1;
|
||||
}
|
||||
std::fprintf(stderr, "strata serve: STATE_HASH L=%lld gdn=%016llx ple=%016llx tail=%016llx pooled=%016llx "
|
||||
"kv=%016llx mtp=%016llx stale=%016llx dead=%016llx ple_prev=%d,%d\n", (long long) L,
|
||||
(unsigned long long) h_gdn, (unsigned long long) h_ple, (unsigned long long) h_tail,
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Private-engine NVMe restart, corruption, identity and admission gate.
|
||||
|
||||
Dry-run by default. Six sequential model loads require an exclusive GPU window.
|
||||
The test only mutates snapshots and tokenizer copies in a new private directory.
|
||||
"""
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import threading
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path[:0] = [str(ROOT), str(ROOT / 'tools')]
|
||||
from conversation_cache_parity import STATE_KEYS, engine_args, load_tokenizer, require, state_hashes
|
||||
from conversation_cache_soak import answer_text
|
||||
from serve.frontend import ChatTemplate
|
||||
from serve.server import StrataEngine, child_env
|
||||
|
||||
|
||||
def evidence(text):
|
||||
return {
|
||||
'verified': [dict(draft=h, cells=int(c), mode=int(m), source=s) for h, c, m, s in re.findall(
|
||||
r'SNAPSHOT_VERIFY draft=([0-9a-f]{16}) cells=(\d+) mode=(\d+) source=(ram|disk)', text)],
|
||||
'draft_prefill': [dict(path=p, mode=int(m), cells=int(c)) for p, m, c in re.findall(
|
||||
r'DRAFT_PREFILL path=(batched|token) mode=(\d+) cells=(\d+)', text)],
|
||||
'spills': text.count('disk cache: spilled '),
|
||||
'promotion_skips': text.count('disk cache: promotion skipped '),
|
||||
'integrity_failures': text.count('snapshot integrity check failed'),
|
||||
'parks': [dict(parked=int(n), bytes=int(b)) for n, b in re.findall(
|
||||
r'conversation cache: parked \d+ tokens .*?parked=(\d+) bytes=(\d+)', text)],
|
||||
}
|
||||
|
||||
|
||||
def verify(results):
|
||||
phases = results['phases']
|
||||
require(set(phases) == {'baseline', 'producer', 'restart', 'admission', 'foreign', 'corrupt'}, 'incomplete disk lifecycle')
|
||||
baseline, producer = phases['baseline'], phases['producer']
|
||||
require([r['name'] for r in baseline['records']] == ['A', 'warm', 'B', 'cold'], 'incomplete baseline controls')
|
||||
require([r['name'] for r in producer['records']] == ['A', 'B', 'C', 'return'], 'incomplete eviction sequence')
|
||||
warm, cold = baseline['records'][1], baseline['records'][3]
|
||||
require(cold['reused'] == 0, 'cold control unexpectedly reused active state')
|
||||
require(len(baseline['records'][0]['ids']) == len(producer['records'][0]['ids']) == 1, 'initial A must consume exactly its prompt')
|
||||
require(producer['records'][0]['ids'] == baseline['records'][0]['ids'], 'initial A differs')
|
||||
require(producer['disk_counts'][0:2] == [0, 0] and producer['disk_counts'][2] > 0,
|
||||
'disk writes must start on RAM eviction, not on continuing turns or ordinary parking')
|
||||
require(producer['disk_counts'][-1] == producer['files_after_close'] > 0, 'shutdown unexpectedly changed disk retention')
|
||||
require(producer['evidence']['spills'] > 0, 'no disk spill observed')
|
||||
require(baseline['info']['conversation_cache_mib'] == baseline['info']['conversation_disk_mib'] == 0,
|
||||
'baseline cache configuration differs')
|
||||
for name, phase in phases.items():
|
||||
for key in ('context', 'kv', 'kv_resident', 'expert_slots', 'spec', 'mtp_max', 'lookup', 'cvec', 'pcie_frac'):
|
||||
require(phase['info'][key] == baseline['info'][key], f'{name}: inference setting differs: {key}')
|
||||
if name != 'baseline':
|
||||
require(phase['info']['conversation_cache_mib'] == results['cache_mib'] and
|
||||
phase['info']['conversation_cache_slots'] == 1 and
|
||||
phase['info']['conversation_disk_mib'] == results['disk_mib'] and
|
||||
phase['info']['conversation_disk_slots'] == 8, f'{name}: cache configuration differs')
|
||||
for record in phase['records']:
|
||||
require(record['ids'] and record['finish'] in ('length', 'stop'), f'{name}: request did not complete')
|
||||
require(set(STATE_KEYS) <= record['state'].keys(), f'{name}: missing main-state fingerprint')
|
||||
require(all(p['parked'] <= 1 and p['bytes'] <= results['cache_mib'] * 1024 * 1024
|
||||
for p in phase['evidence']['parks']), f'{name}: RAM retention exceeded its budget')
|
||||
require(len(phase['disk_bytes']) == len(phase['records']) and
|
||||
all(0 <= n <= results['disk_mib'] * 1024 * 1024 for n in phase['disk_bytes']),
|
||||
f'{name}: missing or over-budget disk occupancy')
|
||||
require(len(phase['disk_counts']) == len(phase['records']) and all(0 <= n <= 8 for n in phase['disk_counts']),
|
||||
f'{name}: missing or over-budget disk entry count')
|
||||
mode = 0 if results['draft_path'] == 'batched' else 2
|
||||
path = 'batched' if mode == 0 else 'token'
|
||||
require(any(p['path'] == path and p['mode'] == mode and p['cells'] > 0
|
||||
for p in producer['evidence']['draft_prefill']), 'requested draft prefill path never ran')
|
||||
require(phases['admission']['info']['conversation_cache_min_free_mib'] == results['denial_floor_mib'],
|
||||
'admission phase did not use the requested floor')
|
||||
require(results['foreign_asset_before'] != results['foreign_asset_after'], 'tokenizer identity was not changed')
|
||||
require(results['corrupted_files'] > 0, 'no persisted files were corrupted')
|
||||
fingerprints = []
|
||||
for name in ('producer', 'restart', 'admission', 'foreign', 'corrupt'):
|
||||
phase = phases[name]
|
||||
record = phase['records'][-1]
|
||||
require(record['name'] == 'return' and (name == 'producer' or len(phase['records']) == 1), 'invalid consumer sequence')
|
||||
reference = warm if name in ('producer', 'restart') else cold
|
||||
require(answer_text(reference['text']) == results['expected'].casefold(), 'baseline known answer is wrong')
|
||||
require(answer_text(record['text']) == results['expected'].casefold(), f'{name}: known answer is wrong')
|
||||
require(record['ids'] == reference['ids'], f'{name}: continuation token parity differs')
|
||||
require(record['state'] == reference['state'], f'{name}: main-model state parity differs')
|
||||
verified = [v for v in phase['evidence']['verified'] if v['source'] == 'disk']
|
||||
if name in ('producer', 'restart'):
|
||||
require(record['reused'] >= results['prompt_tokens'], f'{name}: full disk prefix not restored')
|
||||
require(len(verified) == 1 and verified[0]['mode'] == mode and verified[0]['cells'] >= results['prompt_tokens'],
|
||||
f'{name}: missing draft read-back proof for the requested layout')
|
||||
fingerprints.append(verified[0]['draft'])
|
||||
else:
|
||||
require(record['reused'] == 0 and not verified, f'{name}: incompatible/unadmitted file reused')
|
||||
require(fingerprints[0] == fingerprints[1], 'persisted draft payload changed across restart')
|
||||
require(phases['admission']['evidence']['promotion_skips'] > 0, 'no disk staging admission rejection observed')
|
||||
require(phases['corrupt']['evidence']['integrity_failures'] > 0, 'no integrity rejection observed')
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__)
|
||||
ap.add_argument('--config', type=Path, required=True)
|
||||
ap.add_argument('--engine', type=Path, required=True)
|
||||
ap.add_argument('--output', type=Path, required=True)
|
||||
ap.add_argument('--cache-mib', type=int, default=8192)
|
||||
ap.add_argument('--disk-mib', type=int, default=16384)
|
||||
ap.add_argument('--paragraphs', type=int, default=128)
|
||||
ap.add_argument('--draft-path', choices=('batched', 'ring'), required=True)
|
||||
ap.add_argument('--resident-cells', type=int, default=32768)
|
||||
ap.add_argument('--min-free-mib', type=int, default=2560)
|
||||
ap.add_argument('--run', action='store_true')
|
||||
a = ap.parse_args()
|
||||
limit = (2**63 - 1) // (1024 * 1024)
|
||||
if not 0 < a.cache_mib <= limit or not 0 < a.disk_mib <= limit or a.paragraphs < 1 or a.resident_cells < 4:
|
||||
ap.error('positive in-range cache/disk budgets, paragraphs and resident-cells >= 4 are required')
|
||||
if not 0 <= a.min_free_mib <= limit:
|
||||
ap.error('min-free-mib must be a nonnegative engine-range integer')
|
||||
if not a.run:
|
||||
print(f'Dry run: {a.draft_path} draft KV; six sequential engines; restart, admission, tokenizer identity and corruption.')
|
||||
print('No model loaded. --run requires an exclusive GPU/model window.')
|
||||
return
|
||||
cfg = json.loads(a.config.read_text())
|
||||
a.output.mkdir(mode=0o700, parents=False, exist_ok=False)
|
||||
output = a.output.resolve()
|
||||
tokenizer = output / 'tokenizer'
|
||||
tokenizer.mkdir(mode=0o700)
|
||||
for name in ('vocab.json', 'merges.txt', 'token_type.json', 'chat_template.jinja'):
|
||||
source = Path(cfg['tokenizer']) / name
|
||||
if name == 'chat_template.jinja' and not source.exists():
|
||||
source = ROOT / 'serve/chat_template.jinja'
|
||||
shutil.copyfile(source, tokenizer / name)
|
||||
tok = load_tokenizer(tokenizer)
|
||||
tpl = ChatTemplate(tokenizer / 'chat_template.jinja')
|
||||
expected = 'AZURE-314159'
|
||||
def encode(text):
|
||||
return tok.encode(text, parse_special=True)
|
||||
def prompt(label, code):
|
||||
content = f'Conversation {label}. Remember the exact code {code}.\n'
|
||||
content += '\n'.join(f'Record {i}: blue square, green triangle, red circle.' for i in range(a.paragraphs))
|
||||
content += f'\nThe exact code to remember is {code}. Reply OK.'
|
||||
return encode(tpl.render([{'role': 'user', 'content': content}], enable_thinking=False))
|
||||
A, B, C = prompt('A', expected), prompt('B', 'BRONZE-271828'), prompt('C', 'CORAL-161803')
|
||||
suffix = encode('<|im_end|>\n<|im_start|>user\nWhat exact code did I ask you to remember? '
|
||||
'Reply with only that code.<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n')
|
||||
disk = output / 'disk'
|
||||
env = child_env(cfg)
|
||||
env.update(STRATA_STATE_HASH='1', STRATA_SNAPSHOT_VERIFY='1', STRATA_MTP_BATCH='1')
|
||||
results = {'phases': {}, 'expected': expected, 'prompt_tokens': len(A), 'cache_mib': a.cache_mib,
|
||||
'disk_mib': a.disk_mib, 'draft_path': a.draft_path, 'denial_floor_mib': limit}
|
||||
continuation = None
|
||||
def files():
|
||||
return sorted((disk / 'strata-conversations-v1').glob('*.snap'))
|
||||
def save():
|
||||
(output / 'results.json').write_text(json.dumps(results, indent=2) + '\n')
|
||||
for phase in ('baseline', 'producer', 'restart', 'admission', 'foreign', 'corrupt'):
|
||||
vocab = tokenizer / 'vocab.json'
|
||||
if phase == 'foreign':
|
||||
results['foreign_asset_before'] = hashlib.sha256(vocab.read_bytes()).hexdigest()
|
||||
with vocab.open('ab') as stream:
|
||||
stream.write(b'\n') # valid JSON with identical tokenizer semantics
|
||||
results['foreign_asset_after'] = hashlib.sha256(vocab.read_bytes()).hexdigest()
|
||||
if phase == 'corrupt':
|
||||
with vocab.open('r+b') as stream:
|
||||
stream.seek(-1, 2)
|
||||
stream.truncate() # undo only this test's appended newline
|
||||
damaged = files()
|
||||
require(damaged, 'producer left no snapshots to corrupt')
|
||||
for path in damaged:
|
||||
with path.open('r+b') as stream:
|
||||
stream.seek(-1, 2)
|
||||
byte = stream.read(1)
|
||||
stream.seek(-1, 2)
|
||||
stream.write(bytes([byte[0] ^ 1]))
|
||||
results['corrupted_files'] = len(damaged)
|
||||
args = engine_args(cfg, 0 if phase == 'baseline' else a.cache_mib, 1)
|
||||
args += ['--conversation-cache-slots', '1', '--conversation-cache-disk', str(disk),
|
||||
'--conversation-cache-disk-mib', '0' if phase == 'baseline' else str(a.disk_mib),
|
||||
'--conversation-cache-disk-slots', '8', '--conversation-cache-tokenizer', str(tokenizer),
|
||||
'--conversation-cache-template', str(tokenizer / 'chat_template.jinja'),
|
||||
'--conversation-cache-min-free-mib', str(limit if phase == 'admission' else a.min_free_mib),
|
||||
'--kv-resident', str(0 if a.draft_path == 'batched' else a.resident_cells)]
|
||||
log = output / f'{phase}.log'
|
||||
engine = StrataEngine(str(a.engine.resolve()), args, cwd=cfg.get('cwd'), log=str(log), env=env)
|
||||
entry = {'info': dict(engine.info), 'records': [], 'disk_counts': [], 'disk_bytes': []}
|
||||
results['phases'][phase] = entry
|
||||
def generate(ids, count, name):
|
||||
tokens = [t for t in engine.generate(ids, count, {'temperature': 0}, threading.Event()) if t is not None]
|
||||
entry['records'].append({'name': name, 'ids': tokens, 'text': tok.decode(tokens), **engine.last})
|
||||
entry['disk_counts'].append(len(files()))
|
||||
entry['disk_bytes'].append(sum(path.stat().st_size for path in files()))
|
||||
return tokens
|
||||
try:
|
||||
if phase == 'baseline':
|
||||
head = generate(A, 1, 'A')
|
||||
continuation = A + head + suffix
|
||||
generate(continuation, 24, 'warm')
|
||||
generate(B, 1, 'B')
|
||||
generate(continuation, 24, 'cold')
|
||||
elif phase == 'producer':
|
||||
generate(A, 1, 'A')
|
||||
generate(B, 1, 'B')
|
||||
generate(C, 1, 'C')
|
||||
generate(continuation, 24, 'return')
|
||||
else:
|
||||
generate(continuation, 24, 'return')
|
||||
finally:
|
||||
engine.close()
|
||||
engine.proc.wait(timeout=30) # close() may have killed a slow child; wait before the next model load
|
||||
engine.log.close()
|
||||
save()
|
||||
text = log.read_text()
|
||||
hashes = state_hashes(text)
|
||||
require(len(hashes) == len(entry['records']), f'{phase}: missing state hashes')
|
||||
for record, state in zip(entry['records'], hashes):
|
||||
record['state'] = state
|
||||
entry['evidence'] = evidence(text)
|
||||
entry['files_after_close'] = len(files())
|
||||
save()
|
||||
verify(results)
|
||||
print('PASS: eviction-only persistence, restart, draft read-back, known answer, main-state parity, admission and foreign/corrupt fallback')
|
||||
print(f'Results: {output / "results.json"}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,111 @@
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from tools.conversation_cache_disk import STATE_KEYS, evidence, verify
|
||||
|
||||
|
||||
def fixture(path='batched'):
|
||||
mode = 0 if path == 'batched' else 2
|
||||
info = dict(context=4096, kv='int8', kv_resident=0 if mode == 0 else 1024, expert_slots=1000,
|
||||
spec=2, mtp_max=1, lookup=0, cvec=0, pcie_frac=.55, conversation_cache_mib=8192,
|
||||
conversation_cache_slots=1, conversation_disk_mib=16384, conversation_disk_slots=8, conversation_cache_min_free_mib=2560)
|
||||
state = {k: '1234' for k in STATE_KEYS}
|
||||
def record(name, reused=0, initial=False):
|
||||
return dict(name=name, ids=[1] if initial else [31, 41], text='AZURE-314159', finish='stop', reused=reused, state=dict(state))
|
||||
def proof(disk=False):
|
||||
return dict(verified=[dict(draft='0123456789abcdef', cells=100, mode=mode, source='disk')] if disk else [],
|
||||
draft_prefill=[dict(path='batched' if mode == 0 else 'token', mode=mode, cells=99)],
|
||||
spills=1 if disk else 0, promotion_skips=0, integrity_failures=0,
|
||||
parks=[dict(parked=1, bytes=1234)] if disk else [])
|
||||
phases = {
|
||||
'baseline': dict(info={**info, 'conversation_cache_mib': 0, 'conversation_disk_mib': 0},
|
||||
records=[record('A', initial=True), record('warm', 100), record('B', initial=True), record('cold')], evidence=proof()),
|
||||
'producer': dict(info=info.copy(), records=[record('A', initial=True), record('B', initial=True), record('C', initial=True), record('return', 100)],
|
||||
evidence=proof(True), disk_counts=[0, 0, 1, 2], files_after_close=2),
|
||||
}
|
||||
for name in ('restart', 'admission', 'foreign', 'corrupt'):
|
||||
phases[name] = dict(info=info.copy(), records=[record('return', 100 if name == 'restart' else 0)], evidence=proof(name == 'restart'))
|
||||
phases['admission']['info']['conversation_cache_min_free_mib'] = 999999
|
||||
phases['admission']['evidence']['promotion_skips'] = 1
|
||||
phases['corrupt']['evidence']['integrity_failures'] = 1
|
||||
for phase in phases.values():
|
||||
phase['disk_bytes'] = [2048] * len(phase['records'])
|
||||
phase.setdefault('disk_counts', [0] * len(phase['records']))
|
||||
return dict(phases=phases, expected='AZURE-314159', prompt_tokens=100, cache_mib=8192, disk_mib=16384,
|
||||
draft_path=path, denial_floor_mib=999999, foreign_asset_before='aaa', foreign_asset_after='bbb', corrupted_files=2)
|
||||
|
||||
|
||||
class DiskGate(unittest.TestCase):
|
||||
def test_complete_batched_and_ring_evidence(self):
|
||||
verify(fixture())
|
||||
verify(fixture('ring'))
|
||||
|
||||
def test_missing_or_wrong_evidence_is_rejected(self):
|
||||
mutations = {
|
||||
'missing phase': lambda d: d['phases'].pop('restart'),
|
||||
'wrong answer': lambda d: d['phases']['restart']['records'][0].update(text='BRONZE'),
|
||||
'no disk reuse': lambda d: d['phases']['restart']['records'][0].update(reused=0),
|
||||
'changed output': lambda d: d['phases']['restart']['records'][0].update(ids=[99]),
|
||||
'changed state': lambda d: d['phases']['restart']['records'][0]['state'].update(kv='bad'),
|
||||
'missing state': lambda d: d['phases']['restart']['records'][0]['state'].pop('ple'),
|
||||
'incomplete request': lambda d: d['phases']['restart']['records'][0].update(finish='error'),
|
||||
'fake disk hit': lambda d: d['phases']['restart']['evidence'].update(verified=[]),
|
||||
'RAM instead of disk': lambda d: d['phases']['restart']['evidence']['verified'][0].update(source='ram'),
|
||||
'wrong draft payload': lambda d: d['phases']['restart']['evidence']['verified'][0].update(draft='f' * 16),
|
||||
'wrong resident mode': lambda d: d['phases']['restart']['evidence']['verified'][0].update(mode=2),
|
||||
'no batched prefill': lambda d: d['phases']['producer']['evidence']['draft_prefill'][0].update(path='token'),
|
||||
'per-turn write': lambda d: d['phases']['producer'].update(disk_counts=[1, 1, 2, 2]),
|
||||
'no eviction write': lambda d: d['phases']['producer'].update(disk_counts=[0, 0, 0, 2]),
|
||||
'shutdown write': lambda d: d['phases']['producer'].update(files_after_close=3),
|
||||
'wrong engine residency': lambda d: d['phases']['restart']['info'].update(expert_slots=999),
|
||||
'wrong cache quota': lambda d: d['phases']['restart']['info'].update(conversation_disk_mib=0),
|
||||
'RAM budget exceeded': lambda d: d['phases']['producer']['evidence']['parks'][0].update(bytes=8192 * 1024 * 1024 + 1),
|
||||
'disk budget exceeded': lambda d: d['phases']['restart'].update(disk_bytes=[16384 * 1024 * 1024 + 1]),
|
||||
'disk entry limit exceeded': lambda d: d['phases']['restart'].update(disk_counts=[9]),
|
||||
'missing disk occupancy': lambda d: d['phases']['restart'].update(disk_bytes=[]),
|
||||
'foreign reused': lambda d: d['phases']['foreign']['records'][0].update(reused=100),
|
||||
'no identity change': lambda d: d.update(foreign_asset_after=d['foreign_asset_before']),
|
||||
'no corrupted file': lambda d: d.update(corrupted_files=0),
|
||||
'no integrity rejection': lambda d: d['phases']['corrupt']['evidence'].update(integrity_failures=0),
|
||||
'no admission rejection': lambda d: d['phases']['admission']['evidence'].update(promotion_skips=0),
|
||||
'wrong admission floor': lambda d: d['phases']['admission']['info'].update(conversation_cache_min_free_mib=2560),
|
||||
'warm cold-control': lambda d: d['phases']['baseline']['records'][3].update(reused=100),
|
||||
}
|
||||
for label, mutate in mutations.items():
|
||||
with self.subTest(label=label):
|
||||
data = fixture()
|
||||
mutate(data)
|
||||
with self.assertRaises(AssertionError):
|
||||
verify(data)
|
||||
|
||||
def test_log_evidence(self):
|
||||
parsed = evidence('''strata serve: DRAFT_PREFILL path=batched mode=0 cells=1024
|
||||
strata serve: SNAPSHOT_VERIFY draft=0123456789abcdef cells=2048 mode=0 source=disk
|
||||
strata serve: conversation cache: parked 2048 tokens in 1.0 ms; parked=1 bytes=8192 evictions=1
|
||||
strata serve: disk cache: spilled 2048 tokens
|
||||
strata serve: disk cache: promotion skipped (snapshot integrity check failed)
|
||||
''')
|
||||
self.assertEqual(parsed['verified'], [dict(draft='0123456789abcdef', cells=2048, mode=0, source='disk')])
|
||||
self.assertEqual(parsed['parks'], [dict(parked=1, bytes=8192)])
|
||||
self.assertEqual(parsed['spills'], 1)
|
||||
self.assertEqual(parsed['integrity_failures'], 1)
|
||||
self.assertEqual(parsed['promotion_skips'], 1)
|
||||
|
||||
def test_dry_run_does_not_open_config_or_engine(self):
|
||||
script = Path(__file__).with_name('conversation_cache_disk.py')
|
||||
result = subprocess.run([sys.executable, str(script), '--config', '/does-not-exist', '--engine', '/also-missing',
|
||||
'--output', '/not-created', '--draft-path', 'batched'], capture_output=True, text=True)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertIn('No model loaded', result.stdout)
|
||||
|
||||
def test_verifier_survives_optimized_python(self):
|
||||
command = 'from tools.test_conversation_cache_disk import fixture; from tools.conversation_cache_disk import verify; d=fixture(); d["phases"]["restart"]["records"][0]["text"]="wrong"; verify(d)'
|
||||
result = subprocess.run([sys.executable, '-O', '-c', command], capture_output=True, text=True)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn('known answer is wrong', result.stderr)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user