From 5ee22b5740b3fb2fc33429ae09a9c4025864ae26 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sun, 11 Oct 2026 07:33:56 +0000 Subject: [PATCH 1/3] feat(server): persist evicted sessions to disk with a budget, TTL and validated restores Keep each conversation's newest snapshot and shared heads on disk when RAM evicts them, and flush resident sessions on graceful shutdown. Enable the cache by default with a hard LRU budget, a TTL, bounded asynchronous writes, checksums, exact token keys and model/config identity validation. Restore deeper disk states ahead of shallow RAM hits, load paged snapshots without blocking decoding, and replay tool-call text after restart. Validate Qwen snapshot layout and every disk tensor's type, alignment, dimensions and payload bounds before adoption or allocation; corrupt files become misses. Includes persistence, restart, eviction, corruption and identity regression tests. Compressed prompts do not receive raw disk keys. Co-Authored-By: Claude Opus 5.5 --- server/README.md | 3 +- server/src/common/concurrency/seq_engine.h | 4 + server/src/common/model_backend.h | 101 ++ server/src/internal.h | 3 +- .../qwen35/concurrency/qwen35_seq_engine.cpp | 3 +- .../qwen35/concurrency/qwen35_seq_engine.h | 1 + server/src/qwen35/qwen35_backend.cpp | 176 +- server/src/qwen35/qwen35_backend.h | 9 + server/src/qwen35/qwen35_target_graph.cpp | 9 +- server/src/server/disk_prefix_cache.cpp | 1539 ++++++++++++----- server/src/server/disk_prefix_cache.h | 249 ++- server/src/server/http_server.cpp | 359 +++- server/src/server/http_server.h | 33 +- server/src/server/prefix_cache.cpp | 47 +- server/src/server/prefix_cache.h | 39 +- server/src/server/scheduler.cpp | 169 +- server/src/server/server_main.cpp | 81 +- server/test/unit/test_server_unit.cpp | 1135 +++++++++++- 18 files changed, 3413 insertions(+), 547 deletions(-) diff --git a/server/README.md b/server/README.md index f264586abb..b495ac0d60 100644 --- a/server/README.md +++ b/server/README.md @@ -329,7 +329,8 @@ See the [current six-expert Strix Halo profile](https://www.lucebox.com/blog/dee | `--kv-pool-tokens ` | auto | Shared physical K/V capacity for concurrent serving, rounded to the backend page size. | | `--decode-kv-offload-mb ` | `auto` | Automatically size the RAM budget for active KV suspension; `N` sets a per-model cap in MiB, `0` disables. Applies to concurrent serving. | | `--kv-cache-dir ` | none | Enable persistent disk KV cache in this directory. | -| `--kv-cache-budget ` | `4096` | Disk KV-cache size cap. | +| `--kv-cache-budget ` | `4096` | Disk KV-cache size cap (hard cap, LRU eviction). | +| `--kv-cache-ttl ` | `72` | Delete disk snapshots unused this long; `0` = never. | | `--kv-cache-min-tokens ` | `512` | Minimum prefix length to persist. | | `--kv-cache-interval ` | `10240` | Continued-checkpoint interval. | | `--kv-cache-cold-max ` | `10240` | Cold-prefix limit for long prompts. | diff --git a/server/src/common/concurrency/seq_engine.h b/server/src/common/concurrency/seq_engine.h index f58de81d4b..4306fa33ff 100644 --- a/server/src/common/concurrency/seq_engine.h +++ b/server/src/common/concurrency/seq_engine.h @@ -208,6 +208,10 @@ class SeqEngine { // Optional prefix-checkpoint admission. Unsupported engines remain on // cold admission and never receive a plan from the scheduler. virtual bool supports_prefix_store() const { return false; } + // True when checkpoint id N lives in the backend's snapshot slot N-1 + // (ModelBackend::snapshot_detach/snapshot_adopt reach it), which the + // disk session cache needs. + virtual bool prefix_store_in_backend_snapshots() const { return false; } // Conservative resident-byte estimate for one checkpoint. Returning zero // means the engine cannot safely participate in a configured byte budget. virtual size_t estimate_prefix_store_bytes(int) const { return 0; } diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index 9bcc9df271..6dd9129e77 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -22,6 +22,7 @@ #include "ggml.h" #include "ggml-backend.h" +#include "ggml-cpu.h" #include "generation_types.h" #include "sampler.h" #include "image_prompt.h" @@ -30,6 +31,82 @@ namespace luce::common { +// A snapshot payload taken out of a backend slot: owns its tensor context and +// buffer and frees both when destroyed. Move-only. The disk cache writes it +// in the background after a RAM eviction (see DiskPrefixCache::save_async). +class OwnedSnapshot { +public: + OwnedSnapshot() = default; + OwnedSnapshot(ggml_context * ctx, ggml_backend_buffer_t buf, + int cur_pos, int32_t last_tok) + : ctx_(ctx), buf_(buf), cur_pos_(cur_pos), last_tok_(last_tok) {} + ~OwnedSnapshot() { reset(); } + OwnedSnapshot(const OwnedSnapshot &) = delete; + OwnedSnapshot & operator=(const OwnedSnapshot &) = delete; + OwnedSnapshot(OwnedSnapshot && o) noexcept { *this = std::move(o); } + OwnedSnapshot & operator=(OwnedSnapshot && o) noexcept { + if (this != &o) { + reset(); + ctx_ = o.ctx_; buf_ = o.buf_; + cur_pos_ = o.cur_pos_; last_tok_ = o.last_tok_; + o.ctx_ = nullptr; o.buf_ = nullptr; + } + return *this; + } + void reset() { + if (buf_) ggml_backend_buffer_free(buf_); + if (ctx_) ggml_free(ctx_); + ctx_ = nullptr; buf_ = nullptr; + } + // Hand ctx and buf to a new owner (e.g. ModelBackend::snapshot_adopt). + void release() { ctx_ = nullptr; buf_ = nullptr; } + bool empty() const { return ctx_ == nullptr; } + ggml_context * ctx() const { return ctx_; } + ggml_backend_buffer_t buf() const { return buf_; } + int cur_pos() const { return cur_pos_; } + int32_t last_tok() const { return last_tok_; } + uint64_t bytes() const { + uint64_t n = 0; + for (ggml_tensor * t = ctx_ ? ggml_get_first_tensor(ctx_) : nullptr; t; + t = ggml_get_next_tensor(ctx_, t)) n += ggml_nbytes(t); + return n; + } + + // Copy every tensor of `ctx` into a fresh host (CPU) buffer. + static OwnedSnapshot clone_to_host(ggml_context * ctx, int cur_pos, + int32_t last_tok) { + if (!ctx) return {}; + int n = 0; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; + t = ggml_get_next_tensor(ctx, t)) ++n; + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * (size_t)(n + 4) + 4096; + ip.no_alloc = true; + ggml_context * out = ggml_init(ip); + if (!out) return {}; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; + t = ggml_get_next_tensor(ctx, t)) { + ggml_tensor * c = ggml_new_tensor(out, t->type, GGML_MAX_DIMS, t->ne); + ggml_set_name(c, t->name); + } + ggml_backend_buffer_t buf = + ggml_backend_alloc_ctx_tensors_from_buft(out, ggml_backend_cpu_buffer_type()); + if (!buf) { ggml_free(out); return {}; } + for (ggml_tensor * t = ggml_get_first_tensor(ctx), + * c = ggml_get_first_tensor(out); + t && c; t = ggml_get_next_tensor(ctx, t), c = ggml_get_next_tensor(out, c)) { + ggml_backend_tensor_get(t, c->data, 0, ggml_nbytes(t)); + } + return OwnedSnapshot(out, buf, cur_pos, last_tok); + } + +private: + ggml_context * ctx_ = nullptr; + ggml_backend_buffer_t buf_ = nullptr; + int cur_pos_ = 0; + int32_t last_tok_ = -1; +}; + enum class ParkTarget { // NOTE: Empty preserves Qwen3's existing no-target park/unpark behavior. Empty, @@ -365,6 +442,30 @@ struct ModelBackend { return false; } + // Take a slot's snapshot out of the backend, leaving the slot empty, as a + // host-memory payload the caller owns. Used when a RAM eviction hands the + // snapshot to the disk writer. The default copies snapshot_ref() to a + // host buffer, then frees the slot; backends that keep snapshots in host + // memory hand over the buffer itself. Empty when the slot has none or + // the backend exports no SnapshotRef. + virtual OwnedSnapshot snapshot_detach(int slot) { + const SnapshotRef ref = snapshot_ref(slot); + OwnedSnapshot out = ref.ctx + ? OwnedSnapshot::clone_to_host(ref.ctx, ref.cur_pos, ref.last_tok) + : OwnedSnapshot{}; + snapshot_free(slot); + return out; + } + + // Concurrent (paged) checkpoints: the payload a capture replaced in + // `slot`, kept aside instead of freed so the scheduler can write its + // evicted entry to disk. Empty when nothing was displaced. Every capture + // the scheduler resolves claims it, so none is kept past that. + virtual OwnedSnapshot snapshot_take_displaced(int slot) { + (void)slot; + return {}; + } + // ── Compress (pflash) ──────────────────────────────────────────── // Backend owns the DrafterContext lifecycle and park/unpark policy. diff --git a/server/src/internal.h b/server/src/internal.h index 519d6ade2c..553b152dc0 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -695,7 +695,8 @@ bool replace_paged_target_cache( const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & destination); + PrefixSnapshot & destination, + PrefixSnapshot * displaced = nullptr); // Restore a copied paged snapshot into fresh destination pages and one // recurrent-state slab. `block_table` describes the destination sequence and diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp index c26af182e1..3dc711b5b3 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.cpp +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.cpp @@ -659,7 +659,8 @@ PrefixStoreEvent Qwen35SeqEngine::capture_prefix( PrefixSnapshot & snapshot = b_.prefix_snapshots_[checkpoint]; if (!replace_paged_target_cache( b_.cache_, slot, sequence.block_table, - (int)pool_.block_size(), ticket.checkpoint.tokens, snapshot)) { + (int)pool_.block_size(), ticket.checkpoint.tokens, snapshot, + &b_.displaced_snapshots_[checkpoint])) { event.elapsed_us = (uint64_t)std::chrono::duration_cast( std::chrono::steady_clock::now() - capture_started).count(); diff --git a/server/src/qwen35/concurrency/qwen35_seq_engine.h b/server/src/qwen35/concurrency/qwen35_seq_engine.h index 420f61631f..e3ce03dd09 100644 --- a/server/src/qwen35/concurrency/qwen35_seq_engine.h +++ b/server/src/qwen35/concurrency/qwen35_seq_engine.h @@ -73,6 +73,7 @@ class Qwen35SeqEngine final : public SeqEngine { int slot_count() const override { return slots_.slot_count(); } int max_context() const override { return slots_.max_context(); } bool supports_prefix_store() const override { return true; } + bool prefix_store_in_backend_snapshots() const override { return true; } size_t estimate_prefix_store_bytes(int tokens) const override; void discard_prefix_store(PrefixStoreRef checkpoint) override; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 51f3b90d17..13c0738a9f 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1161,11 +1161,126 @@ ModelBackend::SnapshotRef Qwen35Backend::snapshot_ref(int slot) const { return ref; } +// Hand a PrefixSnapshot's payload to an OwnedSnapshot and leave it empty. +// Only full states (dense, paged) are useful on disk; a thin KV-only one is +// freed instead. A payload outside host memory is copied to host first. +static OwnedSnapshot take_prefix_snapshot(PrefixSnapshot & snap) { + if (!snap.ctx || !snap.buf || + (snap.layout != PrefixSnapshot::Layout::dense && + snap.layout != PrefixSnapshot::Layout::paged)) { + free_prefix_snapshot(snap); + return {}; + } + OwnedSnapshot out; + if (ggml_backend_buffer_is_host(snap.buf)) { + out = OwnedSnapshot(snap.ctx, snap.buf, snap.cur_pos, snap.last_tok); + snap.ctx = nullptr; + snap.buf = nullptr; + } else { + out = OwnedSnapshot::clone_to_host(snap.ctx, snap.cur_pos, snap.last_tok); + } + free_prefix_snapshot(snap); + return out; +} + +OwnedSnapshot Qwen35Backend::snapshot_detach(int slot) { + if (slot < 0 || slot >= PREFIX_SLOTS) return {}; + // A deferred snapshot is still the live state: copy it out first. + if (live_snapshot_deferred_ && slot == live_snapshot_slot_) { + materialize_live_snapshot(); + } + if (slot == live_snapshot_slot_) { + live_snapshot_slot_ = -1; + live_snapshot_deferred_ = false; + } + return take_prefix_snapshot(prefix_snapshots_[slot]); +} + +OwnedSnapshot Qwen35Backend::snapshot_take_displaced(int slot) { + if (slot < 0 || slot >= PREFIX_SLOTS) return {}; + return take_prefix_snapshot(displaced_snapshots_[slot]); +} + +bool Qwen35Backend::adopt_paged_snapshot(int slot, ggml_context * ctx, + ggml_backend_buffer_t buf, + int cur_pos) { + snapshot_free(slot); + auto & snap = prefix_snapshots_[slot]; + snap.attn_k_snap.assign(cache_.attn_k.size(), nullptr); + snap.attn_v_snap.assign(cache_.attn_v.size(), nullptr); + snap.ssm_state_snap.assign(cache_.ssm_state.size(), nullptr); + snap.conv_state_snap.assign(cache_.conv_state.size(), nullptr); + snap.target_feat_snap = nullptr; + const char * mismatch = nullptr; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t && !mismatch; + t = ggml_get_next_tensor(ctx, t)) { + size_t idx = 0; + ggml_tensor ** dst = nullptr; + if (std::sscanf(t->name, "snap_cache_k_%zu", &idx) == 1 && + idx < snap.attn_k_snap.size()) { + dst = &snap.attn_k_snap[idx]; + } else if (std::sscanf(t->name, "snap_cache_v_%zu", &idx) == 1 && + idx < snap.attn_v_snap.size()) { + dst = &snap.attn_v_snap[idx]; + } else if (std::sscanf(t->name, "snap_ssm_state_%zu", &idx) == 1 && + idx < snap.ssm_state_snap.size()) { + dst = &snap.ssm_state_snap[idx]; + } else if (std::sscanf(t->name, "snap_conv_state_%zu", &idx) == 1 && + idx < snap.conv_state_snap.size()) { + dst = &snap.conv_state_snap[idx]; + } + if (!dst || *dst) mismatch = "unexpected or duplicate tensor"; + else *dst = t; + } + // Every live layer needs its tensor, and KV rows must equal cur_pos; + // restore_paged_target_cache() checks types and shapes against the + // cache and the page table before it copies anything. + for (size_t i = 0; i < cache_.attn_k.size() && !mismatch; ++i) { + const bool present = cache_.attn_k[i] != nullptr; + if ((snap.attn_k_snap[i] != nullptr) != present || + (snap.attn_v_snap[i] != nullptr) != present || + (present && (snap.attn_k_snap[i]->ne[1] != cur_pos || + snap.attn_v_snap[i]->ne[1] != cur_pos))) { + mismatch = "attention KV"; + } + } + for (size_t i = 0; i < cache_.ssm_state.size() && !mismatch; ++i) { + const bool present = cache_.ssm_state[i] != nullptr; + if ((snap.ssm_state_snap[i] != nullptr) != present || + (snap.conv_state_snap[i] != nullptr) != present) { + mismatch = "recurrent state"; + } + } + if (mismatch || cur_pos <= 0 || cur_pos > cache_.max_ctx) { + std::fprintf(stderr, + "[qwen35] paged snapshot adopt rejected slot=%d: %s\n", + slot, mismatch ? mismatch : "position"); + snap.attn_k_snap.clear(); snap.attn_v_snap.clear(); + snap.ssm_state_snap.clear(); snap.conv_state_snap.clear(); + return false; + } + snap.ctx = ctx; + snap.buf = buf; + snap.layout = PrefixSnapshot::Layout::paged; + snap.cur_pos = cur_pos; + snap.last_tok = -1; + snap.kv_k_type = cache_.kv_k_type; + snap.max_ctx = cache_.max_ctx; + snap.target_feat_cap = 0; + std::fprintf(stderr, "[qwen35] paged snapshot adopted slot=%d pos=%d\n", + slot, cur_pos); + return true; +} + bool Qwen35Backend::snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, int cur_pos, int32_t last_tok) { - if (cfg_.paged_attention) return false; if (slot < 0 || slot >= PREFIX_SLOTS) return false; + if (cfg_.paged_attention) { + // Paged restores scatter from host memory (paged_snapshot_buffer_type). + if (!ggml_backend_buffer_is_host(buf)) return false; + return adopt_paged_snapshot(slot, ctx, buf, cur_pos); + } snapshot_free(slot); auto & snap = prefix_snapshots_[slot]; @@ -1220,8 +1335,55 @@ bool Qwen35Backend::snapshot_adopt(int slot, ggml_context * ctx, return false; } + // Names alone do not prove compatibility: a snapshot taken with another + // KV type or model shape would be restored as garbage. Every tensor must + // match the live cache's type and per-row shape, and hold cur_pos rows. + const auto same_rows = [](const ggml_tensor * s, const ggml_tensor * c) { + return s && c && s->type == c->type && s->ne[0] == c->ne[0] && + s->ne[2] == c->ne[2] && s->ne[3] == c->ne[3]; + }; + const auto same_shape = [&](const ggml_tensor * s, const ggml_tensor * c) { + return same_rows(s, c) && s->ne[1] == c->ne[1]; + }; + const char * mismatch = nullptr; + for (int i = 0; i < n_full_attn && !mismatch; ++i) { + const ggml_tensor * ck = i < (int)cache_.attn_k.size() ? cache_.attn_k[i] : nullptr; + const ggml_tensor * cv = i < (int)cache_.attn_v.size() ? cache_.attn_v[i] : nullptr; + if (!same_rows(snap.attn_k_snap[i], ck) || + !same_rows(snap.attn_v_snap[i], cv) || + snap.attn_k_snap[i]->ne[1] != cur_pos || + snap.attn_v_snap[i]->ne[1] != cur_pos) { + mismatch = "attention KV"; + } + } + for (int i = 0; i < n_delta && !mismatch; ++i) { + const ggml_tensor * cs = i < (int)cache_.ssm_state.size() ? cache_.ssm_state[i] : nullptr; + const ggml_tensor * cc = i < (int)cache_.conv_state.size() ? cache_.conv_state[i] : nullptr; + if (!same_shape(snap.ssm_state_snap[i], cs) || + !same_shape(snap.conv_state_snap[i], cc)) { + mismatch = "recurrent state"; + } + } + if (!mismatch && cache_.target_feat && + (snap.target_feat_snap->type != cache_.target_feat->type || + snap.target_feat_snap->ne[0] != cache_.target_feat->ne[0] || + snap.target_feat_snap->ne[1] != + std::min(cur_pos, cache_.target_feat_cap))) { + mismatch = "target features"; + } + if (mismatch) { + std::fprintf(stderr, + "[qwen35] snapshot adopt rejected slot=%d: %s type/shape " + "differs from the live cache\n", slot, mismatch); + snap.attn_k_snap.clear(); snap.attn_v_snap.clear(); + snap.ssm_state_snap.clear(); snap.conv_state_snap.clear(); + snap.target_feat_snap = nullptr; + return false; + } + snap.ctx = ctx; snap.buf = buf; + snap.layout = PrefixSnapshot::Layout::dense; // restore_target_cache requires it snap.cur_pos = cur_pos; snap.last_tok = last_tok; snap.kv_k_type = cache_.kv_k_type; @@ -1459,6 +1621,7 @@ void Qwen35Backend::shutdown() { draft_feature_mirror_free(feature_mirror_); for (int i = 0; i < PREFIX_SLOTS; i++) { free_prefix_snapshot(prefix_snapshots_[i]); + free_prefix_snapshot(displaced_snapshots_[i]); } if (!target_parked_) free_target_weights(w_); vision_.reset(); // its buffers belong to target_backend_, freed below @@ -1729,7 +1892,16 @@ GenerateResult Qwen35Backend::restore_and_generate_impl(int slot, // mask-less FA span, so rows beyond the restored prefix must be ZERO, // not leftovers from the previous request. cudaMemset is ~0.2ms. if (cache_.base_buf) ggml_backend_buffer_clear(cache_.base_buf, 0); - restore_target_cache(prefix_snapshots_[slot], cache_); + // Decoding on a cache the restore did not fill yields garbage. + if (!restore_target_cache(prefix_snapshots_[slot], cache_)) { + std::fprintf(stderr, "[qwen35] restore slot=%d failed: %s\n", + slot, luce_last_error()); + cache_.cur_pos = 0; + result.fail(GenerateErrorCode::InvalidSnapshotSlot, + std::string("snapshot restore failed: ") + luce_last_error()); + out_io.emit(-1); + return result; + } } // Now generate from restored state diff --git a/server/src/qwen35/qwen35_backend.h b/server/src/qwen35/qwen35_backend.h index 8562524790..1387e6b1aa 100644 --- a/server/src/qwen35/qwen35_backend.h +++ b/server/src/qwen35/qwen35_backend.h @@ -155,6 +155,8 @@ class Qwen35Backend : public ModelBackend { bool snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, int cur_pos, int32_t last_tok = -1) override; + OwnedSnapshot snapshot_detach(int slot) override; + OwnedSnapshot snapshot_take_displaced(int slot) override; CompressResult compress(const CompressRequest & req) override; std::vector compress_batch( @@ -325,6 +327,13 @@ class Qwen35Backend : public ModelBackend { // ── Prefix cache (snapshots) ───────────────────────────────────── static constexpr int PREFIX_SLOTS = 64; PrefixSnapshot prefix_snapshots_[PREFIX_SLOTS]; + // Paged checkpoints a capture replaced (see snapshot_take_displaced()). + PrefixSnapshot displaced_snapshots_[PREFIX_SLOTS]; + // snapshot_adopt() for a paged (continuous-batching) checkpoint read + // from disk; restore_paged_target_cache() validates it against the page + // table before any copy. + bool adopt_paged_snapshot(int slot, ggml_context * ctx, + ggml_backend_buffer_t buf, int cur_pos); // ── Park state ─────────────────────────────────────────────────── bool target_parked_ = false; diff --git a/server/src/qwen35/qwen35_target_graph.cpp b/server/src/qwen35/qwen35_target_graph.cpp index 713799205c..afeb97c2d4 100644 --- a/server/src/qwen35/qwen35_target_graph.cpp +++ b/server/src/qwen35/qwen35_target_graph.cpp @@ -3407,7 +3407,8 @@ bool replace_paged_target_cache( const std::vector & block_table, int block_size, int token_count, - PrefixSnapshot & destination) { + PrefixSnapshot & destination, + PrefixSnapshot * displaced) { PrefixSnapshot candidate; if (!snapshot_paged_target_cache( cache, seq_slot, block_table, block_size, token_count, @@ -3417,6 +3418,12 @@ bool replace_paged_target_cache( } using std::swap; swap(candidate, destination); + // The replaced payload goes to `displaced` (its owner writes or frees + // it) when one is given, else it is freed here. + if (displaced && candidate.ctx) { + free_prefix_snapshot(*displaced); + swap(candidate, *displaced); + } free_prefix_snapshot(candidate); return true; } diff --git a/server/src/server/disk_prefix_cache.cpp b/server/src/server/disk_prefix_cache.cpp index 7d8e69bf4e..03174b753d 100644 --- a/server/src/server/disk_prefix_cache.cpp +++ b/server/src/server/disk_prefix_cache.cpp @@ -8,6 +8,8 @@ #include "ggml-cpu.h" #include +#include +#include #include #include #include @@ -17,6 +19,15 @@ #include #include #include +#include + +#if defined(_WIN32) +#include +#else +#include +#include +#include +#endif namespace luce::common { @@ -190,43 +201,467 @@ int disk_prefix_cache_auto_boundary( // Little-endian I/O helpers. static void write_u32(FILE * f, uint32_t v) { std::fwrite(&v, 4, 1, f); } static void write_u64(FILE * f, uint64_t v) { std::fwrite(&v, 8, 1, f); } -static void write_i64(FILE * f, int64_t v) { std::fwrite(&v, 8, 1, f); } -static void write_u16(FILE * f, uint16_t v) { std::fwrite(&v, 2, 1, f); } - static bool read_u32(FILE * f, uint32_t & out) { return std::fread(&out, 4, 1, f) == 1; } static bool read_u64(FILE * f, uint64_t & out) { return std::fread(&out, 8, 1, f) == 1; } static bool read_i64(FILE * f, int64_t & out) { return std::fread(&out, 8, 1, f) == 1; } static bool read_u16(FILE * f, uint16_t & out) { return std::fread(&out, 2, 1, f) == 1; } +// 64-bit safe absolute seek. +static bool seek_to(FILE * f, uint64_t offset) { +#if defined(_WIN32) + return _fseeki64(f, (__int64)offset, SEEK_SET) == 0; +#else + return fseeko(f, (off_t)offset, SEEK_SET) == 0; +#endif +} + +// ─── Checksum ─────────────────────────────────────────────────────────── +// Streaming XXH64 (seed 0): fast enough for multi-GB payloads, unlike SHA-1, +// and detects bit flips and torn writes in any part of the file. +namespace { +class Checksum64 { +public: + Checksum64() { + v_[0] = P1 + P2; + v_[1] = P2; + v_[2] = 0; + v_[3] = 0 - P1; + } + void update(const void * data, size_t n) { + const uint8_t * p = (const uint8_t *)data; + total_ += n; + if (buf_len_ > 0) { + const size_t take = std::min(n, sizeof(buf_) - buf_len_); + std::memcpy(buf_ + buf_len_, p, take); + buf_len_ += take; + p += take; + n -= take; + if (buf_len_ < sizeof(buf_)) return; + stripe(buf_); + buf_len_ = 0; + } + while (n >= 32) { + stripe(p); + p += 32; + n -= 32; + } + if (n > 0) { + std::memcpy(buf_, p, n); + buf_len_ = n; + } + } + uint64_t digest() const { + uint64_t h; + if (total_ >= 32) { + h = rotl(v_[0], 1) + rotl(v_[1], 7) + rotl(v_[2], 12) + rotl(v_[3], 18); + for (int i = 0; i < 4; ++i) h = merge(h, v_[i]); + } else { + h = P5; + } + h += total_; + const uint8_t * p = buf_; + size_t n = buf_len_; + while (n >= 8) { + h ^= round(0, read64(p)); + h = rotl(h, 27) * P1 + P4; + p += 8; + n -= 8; + } + if (n >= 4) { + uint32_t w; + std::memcpy(&w, p, 4); + h ^= (uint64_t)w * P1; + h = rotl(h, 23) * P2 + P3; + p += 4; + n -= 4; + } + while (n > 0) { + h ^= (*p) * P5; + h = rotl(h, 11) * P1; + ++p; + --n; + } + h ^= h >> 33; + h *= P2; + h ^= h >> 29; + h *= P3; + h ^= h >> 32; + return h; + } + +private: + static constexpr uint64_t P1 = 11400714785074694791ULL; + static constexpr uint64_t P2 = 14029467366897019727ULL; + static constexpr uint64_t P3 = 1609587929392839161ULL; + static constexpr uint64_t P4 = 9650029242287828579ULL; + static constexpr uint64_t P5 = 2870177450012600261ULL; + static uint64_t rotl(uint64_t x, int r) { return (x << r) | (x >> (64 - r)); } + static uint64_t read64(const uint8_t * p) { + uint64_t v; + std::memcpy(&v, p, 8); + return v; + } + static uint64_t round(uint64_t acc, uint64_t input) { + acc += input * P2; + acc = rotl(acc, 31); + return acc * P1; + } + static uint64_t merge(uint64_t acc, uint64_t v) { + acc ^= round(0, v); + return acc * P1 + P4; + } + void stripe(const uint8_t * p) { + v_[0] = round(v_[0], read64(p)); + v_[1] = round(v_[1], read64(p + 8)); + v_[2] = round(v_[2], read64(p + 16)); + v_[3] = round(v_[3], read64(p + 24)); + } + uint64_t v_[4]; + uint8_t buf_[32]; + size_t buf_len_ = 0; + uint64_t total_ = 0; +}; + +// Sequential file I/O that feeds every byte through the checksum. +struct HashedWriter { + FILE * f; + Checksum64 sum; + bool ok = true; + void write(const void * p, size_t n) { + if (!ok || n == 0) return; + ok = std::fwrite(p, 1, n, f) == n; + sum.update(p, n); + } +}; + +struct HashedReader { + FILE * f; + Checksum64 sum; + bool read(void * p, size_t n) { + if (n == 0) return true; + if (std::fread(p, 1, n, f) != n) return false; + sum.update(p, n); + return true; + } +}; +} // namespace + +bool disk_cache_digest_enabled() { + static const bool on = [] { + const char * v = std::getenv("LUCE_DISK_CACHE_DIGEST"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + return on; +} + +uint64_t snapshot_state_digest(const ModelBackend::SnapshotRef & ref, + int * n_tensors, uint64_t * n_bytes) { + Checksum64 sum; + int tensors = 0; + uint64_t bytes = 0; + const int64_t pos = ref.cur_pos; + const int64_t tok = ref.last_tok; + sum.update(&pos, sizeof(pos)); + sum.update(&tok, sizeof(tok)); + std::vector tmp; + if (ref.ctx) { + for (ggml_tensor * t = ggml_get_first_tensor(ref.ctx); t; + t = ggml_get_next_tensor(ref.ctx, t)) { + const size_t nb = ggml_nbytes(t); + const int32_t type = (int32_t)t->type; + sum.update(t->name, std::strlen(t->name)); + sum.update(&type, sizeof(type)); + sum.update(t->ne, sizeof(t->ne)); + if (t->buffer && nb > 0) { + if (ggml_backend_buffer_is_host(t->buffer)) { + sum.update(t->data, nb); + } else { + const size_t chunk = 64u << 20; + tmp.resize(std::min(nb, chunk)); + for (size_t off = 0; off < nb; off += chunk) { + const size_t n = std::min(chunk, nb - off); + ggml_backend_tensor_get(t, tmp.data(), off, n); + sum.update(tmp.data(), n); + } + } + } + ++tensors; + bytes += nb; + } + } + if (n_tensors) *n_tensors = tensors; + if (n_bytes) *n_bytes = bytes; + return sum.digest(); +} + +// ─── Header (de)serialization ─────────────────────────────────────────── + +static void serialize_header(const DiskCacheHeader & hdr, + uint8_t raw[DISK_CACHE_HEADER_SIZE]) { + std::memset(raw, 0, DISK_CACHE_HEADER_SIZE); + std::memcpy(raw + 0, hdr.magic, 4); + std::memcpy(raw + 4, &hdr.version, 4); + std::memcpy(raw + 8, hdr.layout_id, 16); + std::memcpy(raw + 24, &hdr.cur_pos, 4); + std::memcpy(raw + 28, &hdr.n_tensors, 4); + std::memcpy(raw + 32, &hdr.token_count, 4); + std::memcpy(raw + 36, hdr.token_hash, 16); + std::memcpy(raw + 52, &hdr.payload_bytes, 8); + std::memcpy(raw + 60, &hdr.created_at, 8); + std::memcpy(raw + DISK_CACHE_LAST_USED_OFFSET, &hdr.last_used, 8); + const uint32_t last_tok = (uint32_t)hdr.last_tok; + std::memcpy(raw + 76, &last_tok, 4); + std::memcpy(raw + 80, hdr.identity_salt, 16); +} + +static void parse_header(const uint8_t raw[DISK_CACHE_HEADER_SIZE], + DiskCacheHeader & hdr) { + std::memcpy(hdr.magic, raw + 0, 4); + std::memcpy(&hdr.version, raw + 4, 4); + std::memcpy(hdr.layout_id, raw + 8, 16); + std::memcpy(&hdr.cur_pos, raw + 24, 4); + std::memcpy(&hdr.n_tensors, raw + 28, 4); + std::memcpy(&hdr.token_count, raw + 32, 4); + std::memcpy(hdr.token_hash, raw + 36, 16); + std::memcpy(&hdr.payload_bytes, raw + 52, 8); + std::memcpy(&hdr.created_at, raw + 60, 8); + std::memcpy(&hdr.last_used, raw + DISK_CACHE_LAST_USED_OFFSET, 8); + uint32_t last_tok = 0; + std::memcpy(&last_tok, raw + 76, 4); + hdr.last_tok = (int32_t)last_tok; + std::memcpy(hdr.identity_salt, raw + 80, 16); +} + +// The checksum covers the header with last_used zeroed: a hit rewrites +// that field in place. +static void hash_header(Checksum64 & sum, const uint8_t raw[DISK_CACHE_HEADER_SIZE]) { + uint8_t copy[DISK_CACHE_HEADER_SIZE]; + std::memcpy(copy, raw, DISK_CACHE_HEADER_SIZE); + std::memset(copy + DISK_CACHE_LAST_USED_OFFSET, 0, 8); + sum.update(copy, DISK_CACHE_HEADER_SIZE); +} + +// Skip the tensor table and key tokens of a file positioned after its +// header, then read the auxiliary blob. +static bool read_aux_section(FILE * f, uint32_t n_tensors, uint32_t token_count, + std::string & aux) { + for (uint32_t i = 0; i < n_tensors; ++i) { + uint16_t name_len = 0; + if (!read_u16(f, name_len) || + std::fseek(f, (long)name_len + 4 + 32 + 8, SEEK_CUR) != 0) { + return false; + } + } + // Key tokens: up to 4 * 2^32 bytes, so step in bounded seeks. + uint64_t skip = 4ull * token_count; + while (skip > 0) { + const long step = (long)std::min(skip, 1u << 30); + if (std::fseek(f, step, SEEK_CUR) != 0) return false; + skip -= (uint64_t)step; + } + uint32_t len = 0; + if (!read_u32(f, len) || len > (64u << 20)) return false; + aux.resize(len); + return len == 0 || std::fread(aux.data(), 1, len, f) == len; +} + +// Flush a written file to stable storage. +static bool sync_file(FILE * f) { +#if defined(_WIN32) + return _commit(_fileno(f)) == 0; +#else + return ::fdatasync(fileno(f)) == 0; +#endif +} + +// Make a rename inside `dir` durable (POSIX; a no-op elsewhere). +static void sync_directory(const std::string & dir) { +#if !defined(_WIN32) + const int fd = ::open(dir.c_str(), O_RDONLY | O_DIRECTORY); + if (fd >= 0) { + (void)::fsync(fd); + ::close(fd); + } +#else + (void)dir; +#endif +} + +static double mib(uint64_t bytes) { return (double)bytes / (1024.0 * 1024.0); } + +// Check file metadata without constructing a GGML tensor: removed types have +// zero block sizes, and ggml_new_tensor/ggml_row_size assume valid dimensions. +static bool valid_tensor_metadata(uint32_t type, const int64_t ne[4], uint64_t nbytes) { + if (type >= GGML_TYPE_COUNT) return false; + const auto tensor_type = (ggml_type)type; + const int64_t block = ggml_blck_size(tensor_type); + const size_t width = ggml_type_size(tensor_type); + if (block <= 0 || width == 0 || ne[0] < 0 || ne[0] % block != 0) return false; + // ggml_row_size multiplies by width before dividing by the block size. + if ((uint64_t)ne[0] > SIZE_MAX / width) return false; + uint64_t bytes = width * (uint64_t)ne[0] / (uint64_t)block; + uint64_t elements = (uint64_t)ne[0]; + for (int d = 1; d < 4; ++d) { + if (ne[d] < 0) return false; + const uint64_t dim = (uint64_t)ne[d]; + if (dim && (bytes > SIZE_MAX / dim || elements > INT64_MAX / dim)) return false; + bytes *= dim; + elements *= dim; + } + return bytes == nbytes && bytes <= SIZE_MAX - GGML_MEM_ALIGN; +} + // ─── Construction ─────────────────────────────────────────────────────── DiskPrefixCache::DiskPrefixCache(const DiskCacheConfig & cfg, ModelBackend & backend) : config_(cfg), backend_(backend) {} +DiskPrefixCache::~DiskPrefixCache() { + if (writer_.joinable()) { + drain(); + { + std::lock_guard lock(mu_); + stop_writer_ = true; + } + writer_cv_.notify_all(); + writer_.join(); + } +#if !defined(_WIN32) + if (lock_fd_ >= 0) ::close(lock_fd_); +#endif +} + +uint64_t DiskPrefixCache::now() const { + return clock_ ? clock_() : now_unix(); +} + +static bool has_suffix(const std::string & s, const char * suffix) { + const size_t n = std::strlen(suffix); + return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; +} + +// Validate a snapshot file without reading tensor data: header, version, +// tensor table, payload sum and exact file size. A file cut short by a crash +// or a full disk, or written by an older format, fails here. Bit corruption +// inside the file is caught by the checksum when it is read (read_file). +static bool validate_snapshot_file(const std::string & path, + DiskCacheHeader & hdr, + uint64_t & file_size, + const char *& why) { + std::error_code ec; + file_size = (uint64_t) fs::file_size(path, ec); + if (ec) { why = "unreadable size"; return false; } + FILE * f = std::fopen(path.c_str(), "rb"); + if (!f) { why = "unreadable"; return false; } + bool ok = false; + do { + uint8_t raw[DISK_CACHE_HEADER_SIZE]; + if (std::fread(raw, 1, sizeof(raw), f) != sizeof(raw)) { + why = "short header"; break; + } + DiskCacheHeader h{}; + parse_header(raw, h); + if (std::memcmp(h.magic, "DKVC", 4) != 0) { why = "bad magic"; break; } + if (h.version != DISK_CACHE_VERSION) { why = "old or unknown version"; break; } + if (h.cur_pos == 0 || h.cur_pos > h.token_count) { + why = "snapshot past key"; break; + } + if (h.n_tensors == 0 || h.n_tensors > 4096) { why = "tensor count"; break; } + uint64_t table_bytes = 0; + uint64_t payload = 0; + bool table_ok = true; + for (uint32_t i = 0; i < h.n_tensors && table_ok; ++i) { + uint16_t name_len = 0; + uint32_t type = 0; + int64_t ne[4] = {}; + uint64_t nbytes = 0; + table_ok = read_u16(f, name_len) && name_len < GGML_MAX_NAME && + std::fseek(f, name_len, SEEK_CUR) == 0 && read_u32(f, type) && + type < GGML_TYPE_COUNT; + for (int d = 0; d < 4 && table_ok; ++d) table_ok = read_i64(f, ne[d]); + table_ok = table_ok && read_u64(f, nbytes) && + valid_tensor_metadata(type, ne, nbytes) && + payload <= file_size && nbytes <= file_size - payload; + payload += nbytes; + table_bytes += 2 + name_len + 4 + 32 + 8; + } + if (!table_ok) { why = "bad tensor table"; break; } + if (payload != h.payload_bytes) { why = "payload mismatch"; break; } + uint32_t aux_len = 0; + const uint64_t aux_at = DISK_CACHE_HEADER_SIZE + table_bytes + + 4ull * h.token_count; + if (aux_at + 4 > file_size || + !seek_to(f, aux_at) || !read_u32(f, aux_len)) { + why = "truncated metadata"; break; + } + const uint64_t expected = + aux_at + 4 + aux_len + payload + DISK_CACHE_CHECKSUM_SIZE; + if (file_size != expected) { why = "truncated or oversized"; break; } + hdr = h; + ok = true; + } while (false); + std::fclose(f); + return ok; +} + // ─── Initialization ───────────────────────────────────────────────────── bool DiskPrefixCache::init() { if (disabled()) return true; - if (!mkdir_p(config_.cache_dir)) { std::fprintf(stderr, "[disk-cache] failed to create dir: %s\n", config_.cache_dir.c_str()); return false; } - - // Try to learn layout from existing files (enables first-request disk hits). - try_learn_from_disk(); - - std::fprintf(stderr, "[disk-cache] initialized dir=%s budget=%.1f GB layout=%s\n", +#if !defined(_WIN32) + const std::string lock_path = config_.cache_dir + "/.lock"; + lock_fd_ = ::open(lock_path.c_str(), O_RDWR | O_CREAT | O_CLOEXEC, 0644); + if (lock_fd_ < 0 || ::flock(lock_fd_, LOCK_EX | LOCK_NB) != 0) { + std::fprintf(stderr, + "[disk-cache] disabled: %s is in use by another server " + "(pass --kv-cache-dir to give this one its own)\n", + config_.cache_dir.c_str()); + if (lock_fd_ >= 0) ::close(lock_fd_); + lock_fd_ = -1; + config_.cache_dir.clear(); + return true; + } +#endif + std::lock_guard lock(mu_); + scan_all(); + // Budget and TTL hold from the first moment, before any request. + expire_locked(); + make_room_locked(0); + // Restore into the layout this identity used most recently; learn_layout + // confirms it against the live model before anything is written. + const DiskEntry * best = nullptr; + for (const auto & e : entries_) { + if (e.own_salt && (!best || e.last_used > best->last_used)) best = &e; + } + if (best) { + layout_id_ = best->layout_id; + layout_known_ = true; + layout_from_disk_ = true; + layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); + } + std::fprintf(stderr, + "[disk-cache] initialized dir=%s budget=%.1f GB (configured %.1f GB) " + "ttl=%s files=%zu used=%.1f MB layout=%s\n", config_.cache_dir.c_str(), + (double)effective_budget_ / (1024.0 * 1024.0 * 1024.0), (double)config_.budget_bytes / (1024.0 * 1024.0 * 1024.0), + config_.ttl_seconds + ? (std::to_string(config_.ttl_seconds) + "s").c_str() + : "never", + entries_.size(), mib(total_bytes_), layout_known_ ? hex(layout_id_.data(), 16).c_str() : "pending"); + if (!writer_.joinable()) writer_ = std::thread([this] { writer_loop(); }); return true; } // ─── Layout fingerprint ───────────────────────────────────────────────── -void DiskPrefixCache::compute_layout_id(ggml_context * ctx) { +std::array DiskPrefixCache::layout_of(ggml_context * ctx) const { // Collect tensor metadata sorted by name for deterministic fingerprint. struct TInfo { std::string name; uint32_t type; int64_t ne[4]; }; std::vector tensors; @@ -243,11 +678,10 @@ void DiskPrefixCache::compute_layout_id(ggml_context * ctx) { std::sort(tensors.begin(), tensors.end(), [](const TInfo & a, const TInfo & b) { return a.name < b.name; }); - // Build a single buffer and hash it. // Prepend identity_salt_ so that config/model differences (model file, - // max_ctx, chat_template) rotate the layout_id independently of tensor - // structure. All-zero salt (the default) is back-compatible. + // KV types, max_ctx, chat_template) rotate the layout_id independently of + // tensor structure. All-zero salt (the default) is back-compatible. std::vector buf; buf.insert(buf.end(), identity_salt_.begin(), identity_salt_.end()); for (const auto & ti : tensors) { @@ -255,16 +689,37 @@ void DiskPrefixCache::compute_layout_id(ggml_context * ctx) { buf.insert(buf.end(), (uint8_t *)&ti.type, (uint8_t *)&ti.type + 4); buf.insert(buf.end(), (uint8_t *)ti.ne, (uint8_t *)ti.ne + 32); } - uint8_t digest[20]; sha1_hash(buf.data(), buf.size(), digest); - std::memcpy(layout_id_.data(), digest, 16); + std::array out{}; + std::memcpy(out.data(), digest, 16); + return out; +} + +void DiskPrefixCache::learn_layout_from(ggml_context * ctx) { + const auto id = layout_of(ctx); + std::lock_guard lock(mu_); + if (layout_known_ && !layout_from_disk_) return; // verified from the live model + if (layout_from_disk_ && id != layout_id_) { + // The files stay indexed (budget/TTL) but are never restored. + std::fprintf(stderr, "[disk-cache] layout mismatch: disk=%s model=%s — switching\n", + hex(layout_id_.data(), 16).c_str(), hex(id.data(), 16).c_str()); + } + layout_id_ = id; + layout_known_ = true; + layout_from_disk_ = false; + layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); + mkdir_p(layout_dir_); + std::fprintf(stderr, "[disk-cache] layout learned: %s\n", + hex(layout_id_.data(), 16).c_str()); } void DiskPrefixCache::learn_layout(int slot) { if (disabled()) return; - if (layout_known_ && !layout_from_disk_) return; // already verified from live model - + { + std::lock_guard lock(mu_); + if (layout_known_ && !layout_from_disk_) return; + } auto ref = backend_.snapshot_ref(slot); if (!ref.ctx) { if (backend_.snapshot_used(slot)) { @@ -275,263 +730,593 @@ void DiskPrefixCache::learn_layout(int slot) { } return; } - - std::array prev_id = layout_id_; - bool had_disk_layout = layout_from_disk_; - - compute_layout_id(ref.ctx); - - if (had_disk_layout && std::memcmp(prev_id.data(), layout_id_.data(), 16) != 0) { - // Model layout differs from what was learned from disk files. - std::fprintf(stderr, "[disk-cache] layout mismatch: disk=%s model=%s — switching\n", - hex(prev_id.data(), 16).c_str(), - hex(layout_id_.data(), 16).c_str()); - entries_.clear(); - total_bytes_ = 0; - } - - layout_known_ = true; - layout_from_disk_ = false; - layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); - mkdir_p(layout_dir_); - - std::fprintf(stderr, "[disk-cache] layout learned: %s\n", - hex(layout_id_.data(), 16).c_str()); - - // Scan for previously saved files matching this layout. - scan_directory(); + learn_layout_from(ref.ctx); } // ─── Directory scanning ───────────────────────────────────────────────── -void DiskPrefixCache::scan_directory() { +void DiskPrefixCache::scan_all() { entries_.clear(); total_bytes_ = 0; - - if (layout_dir_.empty()) return; - + size_t removed = 0; std::error_code ec; - for (const auto & de : fs::directory_iterator(layout_dir_, ec)) { - const std::string name = de.path().filename().string(); - size_t nlen = name.size(); - if (nlen < 36 || name.compare(nlen - 4, 4, ".dkv") != 0) continue; + for (const auto & de : fs::directory_iterator(config_.cache_dir, ec)) { + const std::string base = de.path().filename().string(); + std::error_code dec; + if (base.empty() || base[0] == '.' || !de.is_directory(dec)) continue; + bool kept_any = false; + std::error_code sec; + for (const auto & se : fs::directory_iterator(de.path(), sec)) { + const std::string name = se.path().filename().string(); + const std::string path = se.path().string(); + std::error_code rec; + if (has_suffix(name, ".dkv.tmp")) { + // A write the previous process never renamed into place. + fs::remove(path, rec); + ++removed; + continue; + } + if (!has_suffix(name, ".dkv")) continue; + DiskCacheHeader hdr{}; + uint64_t size = 0; + const char * why = ""; + bool ok = validate_snapshot_file(path, hdr, size, why); + if (ok && (hex(hdr.layout_id, 16) != base || + hex(hdr.token_hash, 16) + ".dkv" != name)) { + ok = false; + why = "name does not match header"; + } + if (!ok) { + std::fprintf(stderr, "[disk-cache] removing %s: %s\n", + path.c_str(), why); + fs::remove(path, rec); + ++removed; + continue; + } + DiskEntry entry; + entry.id = next_entry_id_++; + entry.path = path; + std::memcpy(entry.token_hash.data(), hdr.token_hash, 16); + std::memcpy(entry.layout_id.data(), hdr.layout_id, 16); + entry.token_count = hdr.token_count; + entry.cur_pos = hdr.cur_pos; + entry.last_used = hdr.last_used; + entry.file_size = size; + entry.own_salt = std::memcmp(hdr.identity_salt, + identity_salt_.data(), 16) == 0; + total_bytes_ += entry.file_size; + entries_.push_back(std::move(entry)); + kept_any = true; + } + std::error_code rec; + if (!kept_any) fs::remove(de.path(), rec); // only removes if empty + } + std::fprintf(stderr, "[disk-cache] scanned %zu files, %.1f MB (%zu removed)\n", + entries_.size(), mib(total_bytes_), removed); +} - std::string path = de.path().string(); - FILE * f = std::fopen(path.c_str(), "rb"); - if (!f) continue; +// ─── Index maintenance ────────────────────────────────────────────────── - DiskCacheHeader hdr{}; - if (!read_header(f, hdr)) { std::fclose(f); continue; } - std::fclose(f); +bool DiskPrefixCache::expired(const DiskEntry & e, uint64_t t) const { + return config_.ttl_seconds > 0 && t > e.last_used && + t - e.last_used > config_.ttl_seconds; +} - // Validate magic and layout. - if (std::memcmp(hdr.magic, "DKVC", 4) != 0) continue; - if (std::memcmp(hdr.layout_id, layout_id_.data(), 16) != 0) continue; +void DiskPrefixCache::release_pending_locked(DiskEntry & e) { + if (!e.pending) return; + pending_bytes_ -= std::min(pending_bytes_, e.file_size); + e.pending.reset(); +} - DiskEntry entry; - entry.path = path; - std::memcpy(entry.token_hash.data(), hdr.token_hash, 16); - entry.token_count = hdr.token_count; - entry.cur_pos = hdr.cur_pos; - entry.last_used = hdr.last_used; +void DiskPrefixCache::remove_entry(size_t idx, const char * why) { + DiskEntry & e = entries_[idx]; + const bool queued = e.pending != nullptr; + std::fprintf(stderr, "[disk-cache] %s: removing %s%s (%u tokens, %.1f MB)\n", + why, e.path.c_str(), queued ? " (not yet written)" : "", + e.token_count, mib(e.file_size)); + if (queued) { + // Not on disk yet: drop it from the queue; a write in progress + // finds its entry gone and deletes its temporary file. + queue_.erase(std::remove(queue_.begin(), queue_.end(), e.id), queue_.end()); + release_pending_locked(e); + ++stats_.write_cancelled; + } else { + std::error_code ec; + fs::remove(e.path, ec); + } + total_bytes_ -= std::min(total_bytes_, e.file_size); + entries_.erase(entries_.begin() + (long)idx); + done_cv_.notify_all(); +} - std::error_code fec; - auto fsz = fs::file_size(path, fec); - if (!fec) entry.file_size = (uint64_t)fsz; +void DiskPrefixCache::expire_locked() { + const uint64_t t = now(); + for (size_t i = entries_.size(); i-- > 0;) { + if (expired(entries_[i], t)) remove_entry(i, "ttl expired"); + } +} - total_bytes_ += entry.file_size; - entries_.push_back(std::move(entry)); +uint64_t DiskPrefixCache::budget_locked() { + uint64_t budget = config_.budget_bytes; + // Never plan past the filesystem: the cache's bytes already on disk + // plus the free space, minus a reserve for everything else. + uint64_t available = 0; + uint64_t reserve = config_.free_space_reserve_bytes; + bool known = false; + if (free_space_) { + available = free_space_(); + known = true; + } else { + std::error_code ec; + const auto space = fs::space(config_.cache_dir, ec); + if (!ec) { + available = (uint64_t)space.available; + reserve = std::max( + reserve, std::min((uint64_t)space.capacity / 20, 16ull << 30)); + known = true; + } + } + if (known) { + const uint64_t on_disk = total_bytes_ - std::min(total_bytes_, pending_bytes_); + const uint64_t room = available > reserve ? available - reserve : 0; + const uint64_t limit = on_disk + room; + if (limit < budget) budget = limit; } + const bool clamped = budget < config_.budget_bytes; + const bool was_clamped = effective_budget_ < config_.budget_bytes; + const uint64_t moved = budget > effective_budget_ + ? budget - effective_budget_ : effective_budget_ - budget; + if (effective_budget_ == 0 || clamped != was_clamped || + (clamped && moved >= (1ull << 30))) { + std::fprintf(stderr, + "[disk-cache] budget %.1f GB (configured %.1f GB, free space " + "%.1f GB)\n", + (double)budget / 1073741824.0, + (double)config_.budget_bytes / 1073741824.0, + (double)available / 1073741824.0); + } + effective_budget_ = budget; + return budget; +} + +uint64_t DiskPrefixCache::budget_bytes() const { + std::lock_guard lock(mu_); + return effective_budget_; +} - std::fprintf(stderr, "[disk-cache] scanned %zu files, %.1f MB\n", - entries_.size(), (double)total_bytes_ / (1024.0 * 1024.0)); +size_t DiskPrefixCache::total_bytes() const { + std::lock_guard lock(mu_); + return total_bytes_; } -// ─── Cold start: learn layout from existing files ─────────────────────── +size_t DiskPrefixCache::entry_count() const { + std::lock_guard lock(mu_); + return entries_.size(); +} -void DiskPrefixCache::try_learn_from_disk() { - // Scan cache_dir for subdirectories (each is a layout fingerprint). - std::error_code ec; - for (const auto & de : fs::directory_iterator(config_.cache_dir, ec)) { - const std::string base = de.path().filename().string(); - if (!base.empty() && base[0] == '.') continue; - std::error_code dec; - if (!de.is_directory(dec)) continue; - const std::string subdir = de.path().string(); +DiskPrefixCache::Stats DiskPrefixCache::stats() const { + std::lock_guard lock(mu_); + Stats out = stats_; + out.pending_bytes = pending_bytes_; + out.pending_files = 0; + for (const auto & e : entries_) out.pending_files += e.pending ? 1 : 0; + return out; +} - // Check if this subdir has any .dkv files. - std::error_code sec; - for (const auto & se : fs::directory_iterator(subdir, sec)) { - const std::string sname = se.path().filename().string(); - size_t nlen = sname.size(); - if (nlen < 4 || sname.compare(nlen - 4, 4, ".dkv") != 0) continue; +void DiskPrefixCache::make_room_locked(uint64_t incoming) { + const uint64_t budget = budget_locked(); + while (!entries_.empty() && (uint64_t)total_bytes_ + incoming > budget) { + size_t lru = 0; + for (size_t i = 1; i < entries_.size(); ++i) { + if (entries_[i].last_used < entries_[lru].last_used) lru = i; + } + remove_entry(lru, "budget (lru)"); + } +} - // Read the header to get the layout_id. - std::string fpath = se.path().string(); - FILE * f = std::fopen(fpath.c_str(), "rb"); - if (!f) continue; +// ─── Lookup ───────────────────────────────────────────────────────────── - DiskCacheHeader hdr{}; - if (read_header(f, hdr) && std::memcmp(hdr.magic, "DKVC", 4) == 0) { - std::memcpy(layout_id_.data(), hdr.layout_id, 16); - layout_known_ = true; - layout_from_disk_ = true; // unverified — must be confirmed by learn_layout() - layout_dir_ = subdir; +OwnedSnapshot DiskPrefixCache::load(const std::vector & prompt_ids) { + if (disabled()) return {}; + const PrefixHash hash = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); + std::shared_ptr pending; + std::string path; + uint64_t id = 0; + uint32_t cur_pos = 0; + { + std::lock_guard lock(mu_); + if (!layout_known_) return {}; + const int idx = find_entry(hash); + if (idx < 0) return {}; + if (expired(entries_[(size_t)idx], now())) { + remove_entry((size_t)idx, "ttl expired"); + return {}; + } + const DiskEntry & e = entries_[(size_t)idx]; + id = e.id; + pending = e.pending; + path = e.path; + cur_pos = e.cur_pos; + } + OwnedSnapshot out; + const auto t_read = std::chrono::steady_clock::now(); + if (pending) { + // Still queued for the writer: copy it from memory. The shared + // pointer keeps it alive even if the write finishes meanwhile. + if (pending->key != prompt_ids) return {}; + out = OwnedSnapshot::clone_to_host(pending->snap.ctx(), + pending->snap.cur_pos(), + pending->snap.last_tok()); + if (out.empty()) return {}; + } else { + out = read_file(path, prompt_ids); + } + std::lock_guard lock(mu_); + const int idx = find_entry_id(id); + if (out.empty()) { + ++stats_.load_failures; + // A corrupt, partial or mismatched file is never retried. + if (idx >= 0 && !entries_[(size_t)idx].pending) { + remove_entry((size_t)idx, "lookup failed"); + } + return {}; + } + std::fprintf(stderr, "[disk-cache] read %s (%u pos, %.1f MB, %.2fs)\n", + pending ? "pending write" : fs::path(path).filename().string().c_str(), + cur_pos, mib(out.bytes()), + std::chrono::duration( + std::chrono::steady_clock::now() - t_read).count()); + ++stats_.loads; + if (idx >= 0) { + // Persist last_used in the file so LRU order and the TTL survive a + // restart. Not fsynced: losing it in a crash only ages the entry. + DiskEntry & entry = entries_[(size_t)idx]; + entry.last_used = now(); + if (!entry.pending) { + if (FILE * f = std::fopen(entry.path.c_str(), "r+b")) { + if (std::fseek(f, DISK_CACHE_LAST_USED_OFFSET, SEEK_SET) == 0) { + write_u64(f, entry.last_used); + } std::fclose(f); - scan_directory(); - return; } - std::fclose(f); } } + return out; } -// ─── Lookup ───────────────────────────────────────────────────────────── +bool DiskPrefixCache::adopt(OwnedSnapshot snap, int slot) { + if (snap.empty()) return false; + { + // The first hit after startup confirms the layout read from disk. + const auto id = layout_of(snap.ctx()); + std::lock_guard lock(mu_); + if (layout_from_disk_) { + if (id != layout_id_) { + std::fprintf(stderr, + "[disk-cache] adopted layout mismatch: disk=%s model=%s\n", + hex(layout_id_.data(), 16).c_str(), + hex(id.data(), 16).c_str()); + layout_id_ = id; + layout_from_disk_ = false; + layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); + mkdir_p(layout_dir_); + return false; + } + } + } + if (!backend_.snapshot_adopt(slot, snap.ctx(), snap.buf(), snap.cur_pos(), + snap.last_tok())) { + return false; // `snap` still owns and frees the payload + } + snap.release(); + std::lock_guard lock(mu_); + layout_from_disk_ = false; + return true; +} bool DiskPrefixCache::lookup(const std::vector & prompt_ids, int slot) { - if (disabled() || !layout_known_) return false; + return adopt(load(prompt_ids), slot); +} - PrefixHash hash = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); +void DiskPrefixCache::for_each_aux( + const std::function & visit) { + if (disabled()) return; + std::vector> files; + { + std::lock_guard lock(mu_); + for (const auto & e : entries_) { + if (e.own_salt && !e.pending) files.push_back({e.last_used, e.path}); + } + } + std::sort(files.begin(), files.end()); + for (const auto & file : files) { + FILE * f = std::fopen(file.second.c_str(), "rb"); + if (!f) continue; + DiskCacheHeader hdr{}; + std::string aux; + const bool ok = read_header(f, hdr) && + read_aux_section(f, hdr.n_tensors, hdr.token_count, aux); + std::fclose(f); + if (ok && !aux.empty()) visit(aux); + } +} +DiskPrefixCache::PrefixProbe DiskPrefixCache::longest_prefix( + const std::vector & prompt_ids) { + PrefixProbe out; + if (disabled()) return out; std::lock_guard lock(mu_); - int idx = find_entry(hash); - if (idx < 0) return false; - - auto & entry = entries_[idx]; - if (!read_file(entry.path, slot)) { - // File is corrupt or incompatible — remove it. - std::fprintf(stderr, "[disk-cache] lookup failed, removing %s\n", - entry.path.c_str()); - std::remove(entry.path.c_str()); - total_bytes_ -= entry.file_size; - entries_.erase(entries_.begin() + idx); - return false; + if (!layout_known_) return out; + const uint64_t t = now(); + std::vector lengths; + for (const auto & e : entries_) { + if (e.layout_id != layout_id_ || expired(e, t)) continue; + if (e.token_count > prompt_ids.size()) continue; + if ((int)e.token_count < config_.min_tokens) continue; + lengths.push_back(e.token_count); } - - if (layout_from_disk_) { - const std::array disk_id = layout_id_; - auto ref = backend_.snapshot_ref(slot); - if (!ref.ctx) { - backend_.snapshot_free(slot); - return false; - } - compute_layout_id(ref.ctx); - if (std::memcmp(disk_id.data(), layout_id_.data(), 16) != 0) { - std::fprintf(stderr, - "[disk-cache] adopted layout mismatch: disk=%s model=%s\n", - hex(disk_id.data(), 16).c_str(), - hex(layout_id_.data(), 16).c_str()); - backend_.snapshot_free(slot); - entries_.clear(); - total_bytes_ = 0; - layout_from_disk_ = false; - layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); - mkdir_p(layout_dir_); - scan_directory(); - return false; + std::sort(lengths.begin(), lengths.end(), std::greater()); + lengths.erase(std::unique(lengths.begin(), lengths.end()), lengths.end()); + for (uint32_t len : lengths) { + const int idx = find_entry(hash_prefix(prompt_ids.data(), (int)len)); + if (idx >= 0) { + out.len = (int)len; + out.cur_pos = (int)entries_[(size_t)idx].cur_pos; + return out; } - layout_from_disk_ = false; - layout_dir_ = config_.cache_dir + "/" + hex(layout_id_.data(), 16); } + return out; +} - // Update last_used on disk. - entry.last_used = now_unix(); - entry.hits++; - // Optionally rewrite header timestamp (non-critical, skip for perf). - return true; +int DiskPrefixCache::prune_ancestors(const std::vector & tokens, + int keep_max) { + if (disabled()) return 0; + std::lock_guard lock(mu_); + if (!layout_known_) return 0; + // The session's older states on disk: strict prefixes of `tokens`. + std::vector ancestors; + bool head_kept = false; // one at or below keep_max: the shared head + size_t shallowest = SIZE_MAX; + for (size_t i = 0; i < entries_.size(); ++i) { + const DiskEntry & e = entries_[i]; + if (e.layout_id != layout_id_ || e.token_count >= tokens.size()) continue; + if (hash_prefix(tokens.data(), (int)e.token_count) != e.token_hash) continue; + if ((int)e.token_count <= keep_max) { + head_kept = true; + continue; + } + ancestors.push_back(i); + if (shallowest == SIZE_MAX || + e.token_count < entries_[shallowest].token_count) shallowest = i; + } + // Without a short head on disk, the shallowest ancestor stands in for + // it (a long system/tools head other sessions restore), as the RAM + // cache keeps its shallowest entry. + int removed = 0; + std::sort(ancestors.begin(), ancestors.end(), std::greater()); + for (const size_t i : ancestors) { + if (!head_kept && i == shallowest) continue; + remove_entry(i, "superseded by deeper session state"); + ++removed; + } + return removed; } // ─── Save ─────────────────────────────────────────────────────────────── -bool DiskPrefixCache::save(int slot, const std::vector & prompt_ids) { - if (disabled()) return false; - - // Learn layout on first save. - if (!layout_known_) { - learn_layout(slot); - if (!layout_known_) return false; - } - - // Check minimum token threshold. +bool DiskPrefixCache::save_async(OwnedSnapshot snap, + const std::vector & prompt_ids, + const std::string & aux) { + if (disabled() || snap.empty()) return false; if ((int)prompt_ids.size() < config_.min_tokens) return false; - - auto ref = backend_.snapshot_ref(slot); - if (!ref.ctx) return false; - // A file is keyed by the hash of `prompt_ids` only, and a hit restores // the whole snapshot. A snapshot that extends past its key would be // restored for any prompt sharing just the key prefix, with unverified // tokens in between, so such files are never written. - if (ref.cur_pos > (int)prompt_ids.size()) { + if (snap.cur_pos() > (int)prompt_ids.size() || snap.cur_pos() <= 0) { std::fprintf(stderr, - "[disk-cache] skip save: snapshot pos=%d exceeds key of " - "%zu tokens\n", ref.cur_pos, prompt_ids.size()); + "[disk-cache] skip save: snapshot pos=%d does not fit a key of " + "%zu tokens\n", snap.cur_pos(), prompt_ids.size()); return false; } - - PrefixHash hash = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); - - std::lock_guard lock(mu_); - - // Skip if already on disk. - if (find_entry(hash) >= 0) return true; - - // Pre-write budget check: estimate file size and reject if it would - // exceed budget even after evicting all evictable entries. - if (config_.budget_bytes > 0) { - uint64_t payload = 0; - uint32_t ntens = 0; - size_t table_est = 0; - for (ggml_tensor * t = ggml_get_first_tensor(ref.ctx); t; - t = ggml_get_next_tensor(ref.ctx, t)) { - ntens++; - payload += ggml_nbytes(t); - table_est += 2 + std::strlen(t->name) + 4 + 32 + 8; // name_len + name + type + ne[4] + nbytes - } - size_t est_size = DISK_CACHE_HEADER_SIZE + table_est + payload; - if (est_size > config_.budget_bytes) { - std::fprintf(stderr, "[disk-cache] skip save: estimated %.1f MB exceeds budget\n", - (double)est_size / (1024.0 * 1024.0)); - return false; - } + learn_layout_from(snap.ctx()); + + uint64_t payload = 0; + size_t table_bytes = 0; + for (ggml_tensor * t = ggml_get_first_tensor(snap.ctx()); t; + t = ggml_get_next_tensor(snap.ctx(), t)) { + payload += ggml_nbytes(t); + table_bytes += 2 + std::strlen(t->name) + 4 + 32 + 8; } - - std::string path = make_path(hash); - std::string tmp_path = path + ".tmp"; - - if (!write_file(tmp_path, ref, prompt_ids)) { - std::remove(tmp_path.c_str()); - return false; + const uint64_t est_size = DISK_CACHE_HEADER_SIZE + table_bytes + + 4ull * prompt_ids.size() + 4 + aux.size() + payload + + DISK_CACHE_CHECKSUM_SIZE; + const PrefixHash hash = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); + + std::unique_lock lock(mu_); + // Snapshots waiting for the writer hold host memory: bound it. + if (pending_bytes_ > 0 && pending_bytes_ + est_size > config_.max_pending_bytes) { + ++stats_.backpressure_waits; + const auto t0 = std::chrono::steady_clock::now(); + done_cv_.wait(lock, [&] { + return pending_bytes_ == 0 || + pending_bytes_ + est_size <= config_.max_pending_bytes; + }); + std::fprintf(stderr, + "[disk-cache] waited %.2fs for the writer (queued writes over " + "%.0f MB)\n", + std::chrono::duration( + std::chrono::steady_clock::now() - t0).count(), + mib(config_.max_pending_bytes)); } - - // Atomic rename. - if (std::rename(tmp_path.c_str(), path.c_str()) != 0) { - std::remove(tmp_path.c_str()); + expire_locked(); + // Already on disk (or queued): count it as a use. + const int existing = find_entry(hash); + if (existing >= 0) { + entries_[(size_t)existing].last_used = now(); + return true; + } + // Make room first, so the budget holds while the file is being written. + if (est_size > budget_locked()) { + std::fprintf(stderr, "[disk-cache] skip save: %.1f MB exceeds budget\n", + mib(est_size)); return false; } + make_room_locked(est_size); + + auto job = std::make_shared(); + job->snap = std::move(snap); + job->key = prompt_ids; + job->aux = aux; + job->path = make_path(hash); + job->created_at = now(); + job->layout_id = layout_id_; - // Update index. DiskEntry entry; - entry.path = path; - std::memcpy(entry.token_hash.data(), hash.data(), 16); + entry.id = next_entry_id_++; + entry.path = job->path; + entry.token_hash = hash; + entry.layout_id = layout_id_; entry.token_count = (uint32_t)prompt_ids.size(); - entry.cur_pos = (uint32_t)ref.cur_pos; - entry.last_used = now_unix(); - entry.created_at = entry.last_used; - std::error_code fec; - auto fsz = fs::file_size(path, fec); - if (!fec) entry.file_size = (uint64_t)fsz; - - total_bytes_ += entry.file_size; + entry.cur_pos = (uint32_t)job->snap.cur_pos(); + entry.last_used = job->created_at; + entry.own_salt = true; + entry.file_size = est_size; + entry.pending = job; + total_bytes_ += est_size; + pending_bytes_ += est_size; + queue_.push_back(entry.id); entries_.push_back(std::move(entry)); + lock.unlock(); + writer_cv_.notify_one(); + return true; +} - std::fprintf(stderr, "[disk-cache] saved %s (%u tokens, %d pos, %.1f MB)\n", - hex(hash.data(), 16).c_str(), - (uint32_t)prompt_ids.size(), ref.cur_pos, - (double)entries_.back().file_size / (1024.0 * 1024.0)); +bool DiskPrefixCache::save(int slot, const std::vector & prompt_ids, + const std::string & aux) { + if (disabled()) return false; + learn_layout(slot); + if (disabled()) return false; + if ((int)prompt_ids.size() < config_.min_tokens) return false; + const auto ref = backend_.snapshot_ref(slot); + if (!ref.ctx) return false; + if (ref.cur_pos > (int)prompt_ids.size()) { + std::fprintf(stderr, + "[disk-cache] skip save: snapshot pos=%d exceeds key of " + "%zu tokens\n", ref.cur_pos, prompt_ids.size()); + return false; + } + { + // Already on disk: no copy needed. + const PrefixHash hash = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); + std::lock_guard lock(mu_); + const int existing = find_entry(hash); + if (existing >= 0) { + entries_[(size_t)existing].last_used = now(); + return true; + } + } + return save_async(OwnedSnapshot::clone_to_host(ref.ctx, ref.cur_pos, ref.last_tok), + prompt_ids, aux); +} - enforce_budget(); - return true; +// ─── Writer ───────────────────────────────────────────────────────────── + +void DiskPrefixCache::writer_loop() { + std::unique_lock lock(mu_); + while (true) { + writer_cv_.wait(lock, [&] { return stop_writer_ || !queue_.empty(); }); + if (queue_.empty()) break; // stopping + const uint64_t id = queue_.front(); + queue_.pop_front(); + const int idx = find_entry_id(id); + if (idx < 0 || !entries_[(size_t)idx].pending) continue; + std::shared_ptr job = entries_[(size_t)idx].pending; + job->last_used = entries_[(size_t)idx].last_used; + writing_id_ = id; + abort_write_.store(false); + lock.unlock(); + + const std::string tmp_path = job->path + ".tmp"; + const std::string dir = fs::path(job->path).parent_path().string(); + mkdir_p(dir); + const auto t0 = std::chrono::steady_clock::now(); + bool ok = write_file(tmp_path, *job); + const int err = ok ? 0 : errno; + + lock.lock(); + const int at = find_entry_id(id); + if (ok && at >= 0 && !abort_write_.load()) { + ok = std::rename(tmp_path.c_str(), job->path.c_str()) == 0; + if (ok) sync_directory(dir); + } + if (!ok || at < 0 || abort_write_.load()) std::remove(tmp_path.c_str()); + if (at >= 0) { + DiskEntry & e = entries_[(size_t)at]; + if (ok && !abort_write_.load()) { + // A hit while the file was being written: record it too. + if (e.last_used != job->last_used) { + if (FILE * f = std::fopen(e.path.c_str(), "r+b")) { + if (std::fseek(f, DISK_CACHE_LAST_USED_OFFSET, SEEK_SET) == 0) { + write_u64(f, e.last_used); + } + std::fclose(f); + } + } + release_pending_locked(e); + ++stats_.writes; + stats_.write_bytes += e.file_size; + std::fprintf(stderr, + "[disk-cache] wrote %s (%u tokens, %u pos, %.1f MB, %.2fs)\n", + fs::path(e.path).filename().string().c_str(), + e.token_count, e.cur_pos, mib(e.file_size), + std::chrono::duration( + std::chrono::steady_clock::now() - t0).count()); + } else { + ++stats_.write_failures; + std::fprintf(stderr, "[disk-cache] write failed (%s): %s\n", + abort_write_.load() ? "time limit" : std::strerror(err), + e.path.c_str()); + remove_entry((size_t)at, "write failed"); + } + } + job.reset(); // frees the payload unless a lookup still copies it + writing_id_ = 0; + done_cv_.notify_all(); + } +} + +DiskPrefixCache::DrainResult DiskPrefixCache::drain(double max_seconds) { + DrainResult out; + if (!writer_.joinable()) return out; + const auto t0 = std::chrono::steady_clock::now(); + std::unique_lock lock(mu_); + const uint64_t files0 = stats_.writes; + const uint64_t bytes0 = stats_.write_bytes; + const auto idle = [&] { return queue_.empty() && writing_id_ == 0; }; + if (max_seconds < 0) { + done_cv_.wait(lock, idle); + } else if (!done_cv_.wait_until( + lock, t0 + std::chrono::duration_cast( + std::chrono::duration(max_seconds)), + idle)) { + out.cancelled = queue_.size() + (writing_id_ != 0 ? 1 : 0); + while (!queue_.empty()) { + const int idx = find_entry_id(queue_.front()); + if (idx >= 0) { + remove_entry((size_t)idx, "flush time limit"); + } else { + queue_.pop_front(); + } + } + abort_write_.store(true); + done_cv_.wait(lock, idle); + } + out.files = stats_.writes - files0; + out.bytes = stats_.write_bytes - bytes0; + out.seconds = std::chrono::duration( + std::chrono::steady_clock::now() - t0).count(); + return out; } // ─── Continued checkpoints ────────────────────────────────────────────── @@ -541,7 +1326,6 @@ bool DiskPrefixCache::maybe_store_continued(int slot, int cur_pos) { if (disabled()) return false; if (cur_pos < config_.min_tokens) return false; - const int interval = config_.continued_interval; if (interval <= 0) return false; @@ -567,8 +1351,7 @@ bool DiskPrefixCache::maybe_store_continued(int slot, int DiskPrefixCache::cold_prefix_boundary(const std::vector & prompt_ids, const std::vector & boundaries) { - if (disabled() || !layout_known_) return 0; - + if (disabled()) return 0; const int prompt_len = (int)prompt_ids.size(); if (prompt_len <= config_.cold_max_tokens) return 0; if (boundaries.empty()) return 0; @@ -585,61 +1368,11 @@ int DiskPrefixCache::cold_prefix_boundary(const std::vector & prompt_id // Check if we already have a disk entry covering this prefix. PrefixHash hash = hash_prefix(prompt_ids.data(), best); std::lock_guard lock(mu_); + if (!layout_known_) return 0; if (find_entry(hash) >= 0) return 0; // already cached - return best; } -// ─── Budget enforcement ───────────────────────────────────────────────── - -void DiskPrefixCache::enforce_budget() { - uint64_t now = now_unix(); - - // DS4-style eviction scoring: (effective_hits + 1) * tokens / file_size - // with exponential decay on hits (6-hour half-life). - auto score = [now](const DiskEntry & e) -> double { - // Protect recently-saved entries (< 60 seconds old). - if (e.created_at > 0 && now > 0 && (now - e.created_at) < 60) { - return 1e18; - } - // Decay hits: half-life 6h → decay_rate = ln(2) / (6*3600) ≈ 3.2e-5 - double age_s = (now > e.last_used) ? (double)(now - e.last_used) : 0.0; - double effective_hits = (double)e.hits * std::exp(-age_s * 3.2e-5); - double size_factor = (e.file_size > 0) ? (double)e.file_size : 1.0; - return (effective_hits + 1.0) * (double)e.token_count / size_factor; - }; - - while (total_bytes_ > config_.budget_bytes && !entries_.empty()) { - // Find entry with lowest eviction score. - auto it = std::min_element(entries_.begin(), entries_.end(), - [&score](const DiskEntry & a, const DiskEntry & b) { - return score(a) < score(b); - }); - - // Don't evict protected entries. - if (score(*it) >= 1e18) break; - - std::fprintf(stderr, "[disk-cache] evicting %s (%.1f MB, hits=%u, score=%.3f)\n", - hex(it->token_hash.data(), 16).c_str(), - (double)it->file_size / (1024.0 * 1024.0), - it->hits, score(*it)); - - std::remove(it->path.c_str()); - total_bytes_ -= it->file_size; - entries_.erase(it); - } -} - -// ─── Touch ────────────────────────────────────────────────────────────── - -void DiskPrefixCache::touch(const PrefixHash & hash) { - std::lock_guard lock(mu_); - int idx = find_entry(hash); - if (idx >= 0) { - entries_[idx].last_used = now_unix(); - } -} - // ─── Helpers ──────────────────────────────────────────────────────────── std::string DiskPrefixCache::make_path(const PrefixHash & hash) const { @@ -648,206 +1381,212 @@ std::string DiskPrefixCache::make_path(const PrefixHash & hash) const { int DiskPrefixCache::find_entry(const PrefixHash & hash) const { for (int i = 0; i < (int)entries_.size(); ++i) { - if (entries_[i].token_hash == hash) return i; + if (entries_[i].token_hash == hash && + entries_[i].layout_id == layout_id_) return i; + } + return -1; +} + +int DiskPrefixCache::find_entry_id(uint64_t id) const { + for (int i = 0; i < (int)entries_.size(); ++i) { + if (entries_[(size_t)i].id == id) return i; } return -1; } // ─── File I/O: Write ──────────────────────────────────────────────────── -bool DiskPrefixCache::write_file(const std::string & path, - const ModelBackend::SnapshotRef & ref, - const std::vector & prompt_ids) { +bool DiskPrefixCache::write_file(const std::string & path, const PendingWrite & job) { FILE * f = std::fopen(path.c_str(), "wb"); if (!f) return false; + ggml_context * ctx = job.snap.ctx(); - // Count tensors and compute payload size. uint32_t n_tensors = 0; uint64_t payload_bytes = 0; - for (ggml_tensor * t = ggml_get_first_tensor(ref.ctx); t; t = ggml_get_next_tensor(ref.ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { n_tensors++; payload_bytes += ggml_nbytes(t); } - // Write header. DiskCacheHeader hdr{}; std::memcpy(hdr.magic, "DKVC", 4); hdr.version = DISK_CACHE_VERSION; - std::memcpy(hdr.layout_id, layout_id_.data(), 16); - hdr.cur_pos = (uint32_t)ref.cur_pos; + std::memcpy(hdr.layout_id, job.layout_id.data(), 16); + hdr.cur_pos = (uint32_t)job.snap.cur_pos(); hdr.n_tensors = n_tensors; - hdr.token_count = (uint32_t)prompt_ids.size(); - PrefixHash ph = hash_prefix(prompt_ids.data(), (int)prompt_ids.size()); + hdr.token_count = (uint32_t)job.key.size(); + const PrefixHash ph = hash_prefix(job.key.data(), (int)job.key.size()); std::memcpy(hdr.token_hash, ph.data(), 16); hdr.payload_bytes = payload_bytes; - hdr.created_at = now_unix(); - hdr.last_used = hdr.created_at; - hdr.last_tok = ref.last_tok; - - if (!write_header(f, hdr)) { std::fclose(f); return false; } - - // Write tensor table. - for (ggml_tensor * t = ggml_get_first_tensor(ref.ctx); t; t = ggml_get_next_tensor(ref.ctx, t)) { - uint16_t name_len = (uint16_t)std::strlen(t->name); - write_u16(f, name_len); - std::fwrite(t->name, 1, name_len, f); - write_u32(f, (uint32_t)t->type); - for (int d = 0; d < 4; ++d) write_i64(f, t->ne[d]); - write_u64(f, ggml_nbytes(t)); - } - - // Write tensor data. - std::vector buf(4 * 1024 * 1024); // 4 MB transfer buffer - for (ggml_tensor * t = ggml_get_first_tensor(ref.ctx); t; t = ggml_get_next_tensor(ref.ctx, t)) { - size_t nbytes = ggml_nbytes(t); - size_t offset = 0; - while (offset < nbytes) { - size_t chunk = std::min(buf.size(), nbytes - offset); - ggml_backend_tensor_get(t, buf.data(), offset, chunk); - if (std::fwrite(buf.data(), 1, chunk, f) != chunk) { - std::fclose(f); - return false; - } - offset += chunk; + hdr.created_at = job.created_at; + hdr.last_used = job.last_used; + hdr.last_tok = job.snap.last_tok(); + std::memcpy(hdr.identity_salt, identity_salt_.data(), 16); + + uint8_t raw[DISK_CACHE_HEADER_SIZE]; + serialize_header(hdr, raw); + HashedWriter w{f, {}}; + w.ok = std::fwrite(raw, 1, sizeof(raw), f) == sizeof(raw); + hash_header(w.sum, raw); + + for (ggml_tensor * t = ggml_get_first_tensor(ctx); w.ok && t; t = ggml_get_next_tensor(ctx, t)) { + const uint16_t name_len = (uint16_t)std::strlen(t->name); + const uint32_t type = (uint32_t)t->type; + const uint64_t nbytes = ggml_nbytes(t); + w.write(&name_len, 2); + w.write(t->name, name_len); + w.write(&type, 4); + for (int d = 0; d < 4; ++d) w.write(&t->ne[d], 8); + w.write(&nbytes, 8); + } + // The exact key: a hit compares it with the prompt before restoring. + w.write(job.key.data(), job.key.size() * sizeof(int32_t)); + const uint32_t aux_len = (uint32_t)job.aux.size(); + w.write(&aux_len, 4); + w.write(job.aux.data(), job.aux.size()); + // Tensor data: the snapshot's bytes as stored, no conversion. The payload + // is host memory (OwnedSnapshot), written in bounded chunks so a drain + // time limit can stop it. + constexpr size_t kChunk = 8u << 20; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); w.ok && t; t = ggml_get_next_tensor(ctx, t)) { + const size_t nbytes = ggml_nbytes(t); + for (size_t offset = 0; w.ok && offset < nbytes; offset += kChunk) { + if (abort_write_.load(std::memory_order_relaxed)) w.ok = false; + w.write((const char *)t->data + offset, std::min(kChunk, nbytes - offset)); } } - - std::fclose(f); - return true; + const uint64_t digest = w.sum.digest(); + bool ok = w.ok && std::fwrite(&digest, 1, 8, f) == 8; + // Durable before the rename publishes it. + ok = ok && std::fflush(f) == 0 && sync_file(f); + ok = (std::fclose(f) == 0) && ok; + return ok; } // ─── File I/O: Read ───────────────────────────────────────────────────── -bool DiskPrefixCache::read_file(const std::string & path, int slot) { +OwnedSnapshot DiskPrefixCache::read_file(const std::string & path, + const std::vector & prompt_ids) { + std::array layout_id{}; + { + std::lock_guard lock(mu_); + layout_id = layout_id_; + } FILE * f = std::fopen(path.c_str(), "rb"); - if (!f) return false; - + if (!f) return {}; + const auto fail = [&](const char * why) { + if (why) std::fprintf(stderr, "[disk-cache] rejecting %s: %s\n", path.c_str(), why); + std::fclose(f); + return OwnedSnapshot{}; + }; + HashedReader r{f, {}}; + std::error_code ec; + const uint64_t file_size = fs::file_size(path, ec); + if (ec) return fail("unreadable size"); + uint8_t raw[DISK_CACHE_HEADER_SIZE]; + if (std::fread(raw, 1, sizeof(raw), f) != sizeof(raw)) return fail("short header"); + hash_header(r.sum, raw); DiskCacheHeader hdr{}; - if (!read_header(f, hdr)) { std::fclose(f); return false; } - - // Validate. - if (std::memcmp(hdr.magic, "DKVC", 4) != 0) { std::fclose(f); return false; } - if (hdr.version != DISK_CACHE_VERSION) { std::fclose(f); return false; } - if (std::memcmp(hdr.layout_id, layout_id_.data(), 16) != 0) { std::fclose(f); return false; } + parse_header(raw, hdr); + if (std::memcmp(hdr.magic, "DKVC", 4) != 0 || hdr.version != DISK_CACHE_VERSION) { + return fail("bad magic or version"); + } + if (std::memcmp(hdr.layout_id, layout_id.data(), 16) != 0) return fail("layout"); + if (std::memcmp(hdr.identity_salt, identity_salt_.data(), 16) != 0) { + return fail("written by another model/KV type/config"); + } // Never restore a snapshot that extends past the tokens its key covers - // (see save()). Older files written that way are dropped by the caller. - if (hdr.cur_pos > hdr.token_count) { - std::fprintf(stderr, - "[disk-cache] rejecting %s: snapshot pos=%u exceeds key of " - "%u tokens\n", path.c_str(), hdr.cur_pos, hdr.token_count); - std::fclose(f); - return false; + // (see save_async()). + if (hdr.cur_pos == 0 || hdr.cur_pos > hdr.token_count || hdr.cur_pos > INT_MAX) { + return fail("snapshot past its key"); } + if (hdr.token_count != prompt_ids.size()) return fail("key length"); + if (hdr.n_tensors == 0 || hdr.n_tensors > 4096) return fail("tensor count"); + if (hdr.payload_bytes > file_size) return fail("payload size"); - // Read tensor table. std::vector table; table.reserve(hdr.n_tensors); + uint64_t payload = 0; for (uint32_t i = 0; i < hdr.n_tensors; ++i) { DiskTensorEntry ent; - uint16_t name_len; - if (!read_u16(f, name_len)) { std::fclose(f); return false; } - if (name_len >= GGML_MAX_NAME) { std::fclose(f); return false; } + uint16_t name_len = 0; char name_buf[GGML_MAX_NAME] = {}; - if (std::fread(name_buf, 1, name_len, f) != name_len) { std::fclose(f); return false; } - ent.name = name_buf; - if (!read_u32(f, ent.type)) { std::fclose(f); return false; } + uint64_t nbytes = 0; + if (!r.read(&name_len, 2) || name_len >= GGML_MAX_NAME || + !r.read(name_buf, name_len) || !r.read(&ent.type, 4) || + ent.type >= GGML_TYPE_COUNT) { + return fail("bad tensor table"); + } for (int d = 0; d < 4; ++d) { - if (!read_i64(f, ent.ne[d])) { std::fclose(f); return false; } + if (!r.read(&ent.ne[d], 8) || ent.ne[d] < 0) return fail("bad tensor table"); + } + if (!r.read(&nbytes, 8) || !valid_tensor_metadata(ent.type, ent.ne, nbytes) || + payload > hdr.payload_bytes || nbytes > hdr.payload_bytes - payload) { + return fail("bad tensor metadata"); } - uint64_t nbytes_raw; - if (!read_u64(f, nbytes_raw)) { std::fclose(f); return false; } - ent.nbytes = (size_t)nbytes_raw; + ent.name = name_buf; + ent.nbytes = (size_t)nbytes; + payload += nbytes; table.push_back(std::move(ent)); } + if (payload != hdr.payload_bytes) return fail("payload size"); + // The stored key must be exactly this prompt (not just its hash): a + // state keyed on another token stream, e.g. a compressed view, can + // never be restored for it. + std::vector stored(prompt_ids.size()); + if (!r.read(stored.data(), stored.size() * sizeof(int32_t)) || stored != prompt_ids) { + return fail("stored tokens differ from the prompt"); + } + uint32_t aux_len = 0; + if (!r.read(&aux_len, 4) || aux_len > (64u << 20)) return fail("aux"); + std::string aux(aux_len, '\0'); + if (!r.read(aux.data(), aux_len)) return fail("aux"); - // Allocate ggml context + buffer. ggml_init_params ip{}; ip.mem_size = ggml_tensor_overhead() * (size_t)(hdr.n_tensors + 4) + 4096; ip.no_alloc = true; ggml_context * ctx = ggml_init(ip); - if (!ctx) { std::fclose(f); return false; } - - // Create tensors. + if (!ctx) return fail("ggml_init"); for (const auto & ent : table) { ggml_tensor * t = ggml_new_tensor(ctx, (ggml_type)ent.type, 4, ent.ne); - if (!t) { + if (!t || ggml_nbytes(t) != ent.nbytes) { ggml_free(ctx); - std::fclose(f); - return false; + return fail("tensor size"); } ggml_set_name(t, ent.name.c_str()); } - // Host buffer from the CPU buffer type: no backend object to keep alive; // the adopting backend frees the buffer in snapshot_free(). ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_cpu_buffer_type()); if (!buf) { ggml_free(ctx); - std::fclose(f); - return false; + return fail("buffer allocation"); } - - // Read tensor data. - std::vector read_buf(4 * 1024 * 1024); + OwnedSnapshot out(ctx, buf, (int)hdr.cur_pos, hdr.last_tok); for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { - size_t nbytes = ggml_nbytes(t); - size_t offset = 0; - while (offset < nbytes) { - size_t chunk = std::min(read_buf.size(), nbytes - offset); - if (std::fread(read_buf.data(), 1, chunk, f) != chunk) { - ggml_backend_buffer_free(buf); - ggml_free(ctx); - std::fclose(f); - return false; - } - ggml_backend_tensor_set(t, read_buf.data(), offset, chunk); - offset += chunk; - } + if (!r.read(t->data, ggml_nbytes(t))) return fail("truncated payload"); } + uint64_t stored_sum = 0; + if (std::fread(&stored_sum, 1, 8, f) != 8) return fail("missing checksum"); + if (std::fgetc(f) != EOF) return fail("trailing bytes"); + if (stored_sum != r.sum.digest()) return fail("checksum mismatch (corrupt file)"); std::fclose(f); - - // Hand off to backend. - if (!backend_.snapshot_adopt(slot, ctx, buf, (int)hdr.cur_pos, hdr.last_tok)) { - ggml_backend_buffer_free(buf); - ggml_free(ctx); - return false; - } - return true; + return out; } // ─── Header I/O ───────────────────────────────────────────────────────── bool DiskPrefixCache::write_header(FILE * f, const DiskCacheHeader & hdr) { - std::fwrite(hdr.magic, 1, 4, f); - write_u32(f, hdr.version); - std::fwrite(hdr.layout_id, 1, 16, f); - write_u32(f, hdr.cur_pos); - write_u32(f, hdr.n_tensors); - write_u32(f, hdr.token_count); - std::fwrite(hdr.token_hash, 1, 16, f); - write_u64(f, hdr.payload_bytes); - write_u64(f, hdr.created_at); - write_u64(f, hdr.last_used); - write_u32(f, (uint32_t)hdr.last_tok); // stored as u32, cast back on read - // Total: 4+4+16+4+4+4+16+8+8+8+4 = 80 bytes. No extra padding needed. - return !std::ferror(f); + uint8_t raw[DISK_CACHE_HEADER_SIZE]; + serialize_header(hdr, raw); + return std::fwrite(raw, 1, sizeof(raw), f) == sizeof(raw); } bool DiskPrefixCache::read_header(FILE * f, DiskCacheHeader & hdr) { - if (std::fread(hdr.magic, 1, 4, f) != 4) return false; - if (!read_u32(f, hdr.version)) return false; - if (std::fread(hdr.layout_id, 1, 16, f) != 16) return false; - if (!read_u32(f, hdr.cur_pos)) return false; - if (!read_u32(f, hdr.n_tensors)) return false; - if (!read_u32(f, hdr.token_count)) return false; - if (std::fread(hdr.token_hash, 1, 16, f) != 16) return false; - if (!read_u64(f, hdr.payload_bytes)) return false; - if (!read_u64(f, hdr.created_at)) return false; - if (!read_u64(f, hdr.last_used)) return false; - uint32_t last_tok_raw; - if (!read_u32(f, last_tok_raw)) return false; - hdr.last_tok = (int32_t)last_tok_raw; + uint8_t raw[DISK_CACHE_HEADER_SIZE]; + if (std::fread(raw, 1, sizeof(raw), f) != sizeof(raw)) return false; + parse_header(raw, hdr); return true; } diff --git a/server/src/server/disk_prefix_cache.h b/server/src/server/disk_prefix_cache.h index d65e46fcd7..998ef3da46 100644 --- a/server/src/server/disk_prefix_cache.h +++ b/server/src/server/disk_prefix_cache.h @@ -3,13 +3,30 @@ // Complements the in-memory PrefixCache by serializing snapshot tensors to // files, enabling cache survival across restarts and overflow to disk. // -// File format: 80-byte header + tensor table + raw tensor data. -// Files are keyed by SHA-1 of prompt token IDs (same as in-memory cache). -// A layout fingerprint (SHA-1 of tensor names/types/shapes) prevents loading -// snapshots from incompatible models. +// File format (v5): 96-byte header + tensor table + the exact key token ids +// + an auxiliary blob (u32 length + bytes; the server stores the session's +// tool-call replay text there) + raw tensor data + an 8-byte checksum of +// everything before it (the header with last_used zeroed, since a hit +// rewrites that field in place). Tensor bytes are copied verbatim in the +// cache's own type (e.g. q4_0 KV stays q4_0): no requantization on save or +// restore. Files are keyed by SHA-1 of prompt token IDs (same as in-memory +// cache); a hit also compares the stored token ids with the prompt before +// adopting anything, and verifies the checksum over the whole file: a +// corrupt or partial file is a logged miss and is deleted. A layout +// fingerprint (SHA-1 of an identity salt plus the tensor names/types/shapes) +// and the identity salt itself prevent loading snapshots from another +// model, KV type or config. +// +// Writes are asynchronous: save_async() takes ownership of a snapshot that +// left RAM, indexes it at once (it counts against the budget and a lookup +// serves it from memory until it is on disk), and a writer thread writes it +// (temporary file, fdatasync, rename, directory sync). // // Directory structure: // //.dkv +// +// Every layout directory is indexed for the disk budget and the TTL; only +// the live layout's files are ever restored. #pragma once @@ -17,9 +34,15 @@ #include "common/model_backend.h" #include +#include +#include #include +#include +#include +#include #include #include +#include #include namespace luce::common { @@ -28,10 +51,20 @@ namespace luce::common { struct DiskCacheConfig { std::string cache_dir; // base directory (empty = disabled) - size_t budget_bytes = (size_t)4 * 1024 * 1024 * 1024; // 4 GB default + size_t budget_bytes = (size_t)128 * 1024 * 1024 * 1024; // 128 GiB default int min_tokens = 512; // only persist snapshots >= this many tokens int continued_interval = 10240; // save every N tokens during long sessions int cold_max_tokens = 10240; // create cold checkpoint for prompts longer than this + // Files unused for longer than this are deleted (at startup, before each + // save, and when a lookup meets one). 0 = never expire. + uint64_t ttl_seconds = 72ull * 3600; + // Snapshots handed to save_async() but not yet on disk hold host memory. + // Past this many bytes, save_async() waits for the writer (logged). + uint64_t max_pending_bytes = 4ull * 1024 * 1024 * 1024; + // The budget never exceeds what fits on the filesystem: the cache's own + // bytes plus the free space, minus this reserve or 5% of the disk (at + // most 16 GiB), whichever is larger. + uint64_t free_space_reserve_bytes = 1ull * 1024 * 1024 * 1024; }; enum class DiskPrefixCacheMode { @@ -52,6 +85,14 @@ struct DiskPrefixCachePolicy { const char * disk_prefix_cache_mode_name(DiskPrefixCacheMode mode); std::string disk_prefix_cache_policy_name(const DiskPrefixCachePolicy & policy); + +// Debug aid for the SSD-vs-RAM numerics check: XXH64 over every tensor of a +// snapshot (name, type, shape, bytes) plus cur_pos and last_tok. Enabled at +// restore time by LUCE_DISK_CACHE_DIGEST=1; costs a full device-to-host copy. +bool disk_cache_digest_enabled(); +uint64_t snapshot_state_digest(const ModelBackend::SnapshotRef & ref, + int * n_tensors = nullptr, + uint64_t * n_bytes = nullptr); bool parse_disk_prefix_cache_policy(const std::string & value, DiskPrefixCachePolicy & out); @@ -80,7 +121,7 @@ int disk_prefix_cache_auto_boundary( const std::vector & safe_boundaries, int min_tokens); -// ─── File header (80 bytes, little-endian) ────────────────────────────── +// ─── File header (96 bytes, little-endian) ────────────────────────────── struct DiskCacheHeader { char magic[4]; // "DKVC" @@ -94,15 +135,21 @@ struct DiskCacheHeader { uint64_t created_at; // unix seconds uint64_t last_used; // unix seconds (updated on hit) int32_t last_tok; // last prefill token (needed for decode seeding) - // Padding to 80 bytes. + uint8_t identity_salt[16]; // model/KV-type/config identity (v4) }; // NOTE: The header is serialized field-by-field, not as a raw struct, -// to avoid alignment/packing issues. The on-disk size is exactly 80 bytes. - -static constexpr size_t DISK_CACHE_HEADER_SIZE = 80; -// v3: ggml_type id 42 changed meaning (TQ3_0 -> Q2_0). Bump so cache files -// written before the Q2_0 type existed are rejected instead of misread. -static constexpr uint32_t DISK_CACHE_VERSION = 3; +// to avoid alignment/packing issues. The on-disk size is exactly 96 bytes. + +static constexpr size_t DISK_CACHE_HEADER_SIZE = 96; +// Byte offset of last_used, rewritten in place when a hit touches a file. +static constexpr long DISK_CACHE_LAST_USED_OFFSET = 68; +// v3: ggml_type id 42 changed meaning (TQ3_0 -> Q2_0). +// v4: identity salt in the header and the key token ids after the tensor +// table. +// v5: an 8-byte checksum trailer over the whole file. Older files are +// deleted at startup. +static constexpr uint32_t DISK_CACHE_VERSION = 5; +static constexpr size_t DISK_CACHE_CHECKSUM_SIZE = 8; // ─── Tensor table entry (on-disk) ────────────────────────────────────── @@ -118,7 +165,8 @@ struct DiskTensorEntry { class DiskPrefixCache { public: DiskPrefixCache(const DiskCacheConfig & cfg, ModelBackend & backend); - ~DiskPrefixCache() = default; + // Finishes the writes already queued (see drain()). + ~DiskPrefixCache(); DiskPrefixCache(const DiskPrefixCache &) = delete; DiskPrefixCache & operator=(const DiskPrefixCache &) = delete; @@ -127,17 +175,71 @@ class DiskPrefixCache { return config_.cache_dir.empty() || backend_snapshot_unsupported_; } - // Initialize: create directory, scan existing files, learn layout from - // first available snapshot. Returns false on fatal error. + // Initialize: create directory, take its lock (another process using it + // turns this cache off, logged), index every snapshot file, delete + // corrupt/partial/outdated files and leftover temporaries, clamp the + // budget to the free space, apply the TTL and the budget, pick the most + // recently used layout written with this identity salt (enables + // first-request disk hits), and start the writer thread. Returns false + // on fatal error. bool init(); - // Look up a prompt on disk. On hit, loads the snapshot into `out_slot` + // Look up a prompt on disk. On hit, loads the snapshot into `slot` // using backend.snapshot_adopt(). Returns true if loaded successfully. + // Equivalent to adopt(load(prompt_ids), slot). bool lookup(const std::vector & prompt_ids, int slot); - // Save the snapshot in `slot` to disk, keyed by prompt_ids. - // Returns true on success. - bool save(int slot, const std::vector & prompt_ids); + // Read the snapshot keyed by exactly `prompt_ids` into host memory: from + // the file (verifying key tokens and checksum; a bad file is deleted), + // or from memory while its write is still queued. Touches the backend + // not at all, so it may run on another thread. Empty on miss. + OwnedSnapshot load(const std::vector & prompt_ids); + + // Hand a loaded snapshot to the backend in `slot` (checks the layout + // against the live model on the first hit after startup). + bool adopt(OwnedSnapshot snap, int slot); + + // Take ownership of a snapshot that left RAM and write it in the + // background, keyed by `prompt_ids`, with an optional auxiliary blob. + // It is indexed at once: it counts against the budget (older files are + // evicted first) and lookups find it. Returns false when it is not + // written (disabled, below min_tokens, larger than the budget, a + // snapshot past its key); true when queued or already on disk. + bool save_async(OwnedSnapshot snap, const std::vector & prompt_ids, + const std::string & aux = std::string()); + + // Queue a copy of the snapshot in `slot` (save_async); the slot keeps + // its own. Returns true on success (or when the entry already exists). + bool save(int slot, const std::vector & prompt_ids, + const std::string & aux = std::string()); + + // Wait until every queued write is on disk, at most `max_seconds` + // (< 0: no limit). Writes still queued then are cancelled and the one + // in progress is abandoned (its temporary file is deleted). + struct DrainResult { + uint64_t files = 0; // written during the drain + uint64_t bytes = 0; + uint64_t cancelled = 0; // dropped by the time limit + double seconds = 0.0; + }; + DrainResult drain(double max_seconds = -1.0); + + // Visit the auxiliary blobs of the restorable files (this identity), + // least recently used first. Reads only the small metadata sections. + void for_each_aux(const std::function & visit); + + // Index-only probe (no file read): the deepest restorable entry whose + // key is a prefix of `prompt_ids`. len = key tokens, cur_pos = KV rows + // the snapshot holds (<= len). {0, 0} on miss. + struct PrefixProbe { int len = 0; int cur_pos = 0; }; + PrefixProbe longest_prefix(const std::vector & prompt_ids); + + // Delete live-layout entries whose key is a strict prefix of `tokens` + // and longer than `keep_max` tokens: a deeper state of the same session + // supersedes them. Entries up to keep_max (the shared system/tools head) + // stay for other sessions; when there is none, the shallowest ancestor + // stays in its place (a longer head). Returns the number of files removed. + int prune_ancestors(const std::vector & tokens, int keep_max); // Check if a continued checkpoint should be saved after generation. // `all_tokens` = prompt + generated tokens, `cur_pos` = final position. @@ -156,14 +258,25 @@ class DiskPrefixCache { int cold_prefix_boundary(const std::vector & prompt_ids, const std::vector & boundaries); - // Evict files until total disk usage is within budget. - void enforce_budget(); - - // Update last_used timestamp for a file (on cache hit). - void touch(const PrefixHash & hash); - - // Get total bytes used on disk. - size_t total_bytes() const { return total_bytes_; } + // Get total bytes used on disk (all layouts), queued writes included. + size_t total_bytes() const; + size_t entry_count() const; + int cold_max_tokens() const { return config_.cold_max_tokens; } + // The budget in force: the configured one, clamped to the free space. + uint64_t budget_bytes() const; + + struct Stats { + uint64_t writes = 0; // files written + uint64_t write_bytes = 0; + uint64_t write_failures = 0; + uint64_t write_cancelled = 0; // removed (budget, TTL, superseded) before written + uint64_t pending_files = 0; + uint64_t pending_bytes = 0; + uint64_t backpressure_waits = 0; + uint64_t loads = 0; // hits served (file or queued write) + uint64_t load_failures = 0; // corrupt, partial or mismatched files + }; + Stats stats() const; // Get the continued-interval setting. int continued_interval() const { return config_.continued_interval; } @@ -179,9 +292,19 @@ class DiskPrefixCache { identity_salt_ = salt; } + // Unix-seconds clock for created_at/last_used/TTL. Tests inject one. + void set_clock(std::function clock) { clock_ = std::move(clock); } + + // Free bytes on the cache's filesystem. Tests inject one. + void set_free_space_probe(std::function probe) { + free_space_ = std::move(probe); + } + private: DiskCacheConfig config_; ModelBackend & backend_; + std::function clock_; + std::function free_space_; // Some backends support in-memory snapshots but deliberately do not // implement SnapshotRef/adoption. Detect that once from a live slot and // stop scheduling disk work for the rest of the process. @@ -194,38 +317,78 @@ class DiskPrefixCache { // All-zeroes by default → backward-compatible behavior. std::array identity_salt_{}; - // Layout fingerprint (learned from first snapshot). + // Layout fingerprint (learned from first snapshot). Guarded by mu_. std::array layout_id_{}; bool layout_known_ = false; bool layout_from_disk_ = false; // true if learned from file, unverified std::string layout_dir_; // // - // In-memory index of on-disk files. + // A snapshot waiting for the writer. Shared with the index entry so a + // lookup can copy it while the writer is busy with it. + struct PendingWrite { + OwnedSnapshot snap; + std::vector key; + std::string aux; + std::string path; + uint64_t created_at = 0; + uint64_t last_used = 0; // the entry's, when the write started + std::array layout_id{}; + }; + + // In-memory index of on-disk files (every layout directory). struct DiskEntry { + uint64_t id = 0; std::string path; PrefixHash token_hash; + std::array layout_id{}; uint32_t token_count = 0; uint32_t cur_pos = 0; uint64_t file_size = 0; uint64_t last_used = 0; - uint32_t hits = 0; - uint64_t created_at = 0; // for eviction protection + bool own_salt = false; // written with identity_salt_ + std::shared_ptr pending; // non-null until on disk }; std::vector entries_; size_t total_bytes_ = 0; - std::mutex mu_; - - // Helpers. - void compute_layout_id(ggml_context * ctx); - void scan_directory(); - void try_learn_from_disk(); + uint64_t next_entry_id_ = 1; + uint64_t effective_budget_ = 0; + mutable std::mutex mu_; + + // Writer thread and its queue (entry ids, oldest first). + std::deque queue_; + uint64_t writing_id_ = 0; // entry being written, 0 = none + std::atomic abort_write_{false}; // drain() time limit hit + bool stop_writer_ = false; + uint64_t pending_bytes_ = 0; + std::condition_variable writer_cv_; // work for the writer + std::condition_variable done_cv_; // a write finished + std::thread writer_; + Stats stats_; + // Exclusive lock on /.lock: one process (and one model) per + // directory, since startup deletes temporaries and the budget evicts + // files of every layout in it. + int lock_fd_ = -1; + + // Helpers (the _locked ones need mu_). + uint64_t now() const; + std::array layout_of(ggml_context * ctx) const; + void learn_layout_from(ggml_context * ctx); + void scan_all(); std::string make_path(const PrefixHash & hash) const; int find_entry(const PrefixHash & hash) const; - - bool write_file(const std::string & path, - const ModelBackend::SnapshotRef & ref, - const std::vector & prompt_ids); - bool read_file(const std::string & path, int slot); + int find_entry_id(uint64_t id) const; + bool expired(const DiskEntry & e, uint64_t t) const; + void remove_entry(size_t idx, const char * why); + void expire_locked(); + uint64_t budget_locked(); + // Evict least-recently-used files until `incoming` more bytes fit. + void make_room_locked(uint64_t incoming); + void writer_loop(); + void release_pending_locked(DiskEntry & e); + + bool write_file(const std::string & path, const PendingWrite & job); + OwnedSnapshot read_file(const std::string & path, + const std::vector & prompt_ids); // Header I/O. static bool write_header(FILE * f, const DiskCacheHeader & hdr); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index e71d34f108..b75ad6c157 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -984,6 +984,9 @@ json build_props_body(const ServerConfig & config, {"disk_bytes", pcfs.disk_bytes}, {"lifetime_hits", pcfs.lifetime_hits}, {"disk_policy", disk_prefix_cache_policy_name(config.disk_cache_policy)}, + {"disk_enabled", !config.disk_cache_dir.empty()}, + {"disk_budget_mb", config.disk_cache_budget_mb}, + {"disk_ttl_hours", config.disk_cache_ttl_hours}, }}, {"tool_replay", { {"max_entries", tms.max_entries}, @@ -1264,7 +1267,58 @@ std::vector normalize_chat_messages( // max_ctx, sha1(chat_template_src), and the effective K-rotation basis // (LUCE_KV_ROTATE resolves against the K-cache type; a cache written // rotated must never be adopted by an un-rotated session or vice versa). +// Also covered: the KV cache types (a q8_0 snapshot must never be adopted by +// a q4_0 session), the draft model (it decides the captured target-feature +// layers), and a content fingerprint of the model file (SHA-1 of its first +// and last MiB plus its size; hashing the whole multi-GB file at every +// startup would cost tens of seconds, and a re-quantized or replaced file +// changes its header and tail). // Returns all-zeroes if model_path is empty (back-compat / disk disabled). +static void append_file_identity(std::vector & buf, + const std::string & path, bool content) { + uint32_t plen = (uint32_t)path.size(); + buf.insert(buf.end(), (uint8_t *)&plen, (uint8_t *)&plen + 4); + buf.insert(buf.end(), path.begin(), path.end()); + if (path.empty()) return; + struct stat st{}; + int64_t file_size = 0, file_mtime = 0; + if (::stat(path.c_str(), &st) == 0) { + file_size = (int64_t)st.st_size; + file_mtime = (int64_t)st.st_mtime; + } + buf.insert(buf.end(), (uint8_t *)&file_size, (uint8_t *)&file_size + 8); + buf.insert(buf.end(), (uint8_t *)&file_mtime, (uint8_t *)&file_mtime + 8); + if (!content) return; + uint8_t digest[20] = {}; + std::ifstream in(path, std::ios::binary); + if (in) { + constexpr int64_t kChunk = 1 << 20; + std::vector data; + auto read_at = [&](int64_t off, int64_t len) { + if (len <= 0) return; + in.clear(); + in.seekg((std::streamoff)off); + const size_t start = data.size(); + data.resize(start + (size_t)len); + in.read((char *)data.data() + start, (std::streamsize)len); + data.resize(start + (size_t)std::max(0, in.gcount())); + }; + read_at(0, std::min(kChunk, file_size)); + if (file_size > kChunk) { + const int64_t tail = std::max(kChunk, file_size - kChunk); + read_at(tail, file_size - tail); + } + sha1_hash(data.data(), data.size(), digest); + } + buf.insert(buf.end(), digest, digest + 20); +} + +static uint64_t disk_cache_ttl_seconds(double hours) { + if (!(hours > 0.0)) return 0; // 0 (or invalid) = never expire + const double secs = hours * 3600.0; + return secs < 1.0 ? 1 : (uint64_t)std::llround(secs); +} + static std::array compute_disk_cache_salt(const ServerConfig & cfg) { std::array salt{}; if (cfg.model_path.empty()) return salt; @@ -1297,6 +1351,14 @@ static std::array compute_disk_cache_salt(const ServerConfig & cfg) buf.insert(buf.end(), tmpl_digest, tmpl_digest + 20); uint8_t k_rotated = dflash_kv_k_rotation_enabled(cfg.kv_cache_k) ? 1 : 0; buf.push_back(k_rotated); + // Snapshot bytes are stored in the cache's own type. + for (const std::string * kv : {&cfg.kv_cache_k, &cfg.kv_cache_v}) { + uint32_t n = (uint32_t)kv->size(); + buf.insert(buf.end(), (uint8_t *)&n, (uint8_t *)&n + 4); + buf.insert(buf.end(), kv->begin(), kv->end()); + } + append_file_identity(buf, cfg.model_path, /*content=*/true); + append_file_identity(buf, cfg.draft_path, /*content=*/false); uint8_t digest[20]; sha1_hash(buf.data(), buf.size(), digest); @@ -1456,7 +1518,8 @@ HttpServer::HttpServer(luce::engine::LuceEngine & engine, config.disk_cache_budget_mb * (size_t)(1024 * 1024), config.disk_cache_min_tokens, config.disk_cache_continued_interval, - config.disk_cache_cold_max_tokens}, backend_) + config.disk_cache_cold_max_tokens, + disk_cache_ttl_seconds(config.disk_cache_ttl_hours)}, backend_) { config_.image_input_enabled = backend_.supports_images() && config_.pflash_upstream_base.empty() && @@ -1481,6 +1544,7 @@ HttpServer::HttpServer(luce::engine::LuceEngine & engine, disk_cache_.set_identity_salt(compute_disk_cache_salt(config)); } disk_cache_.init(); + restore_replay_blobs(); status_html_path_ = resolve_status_html(); // PPP env overrides (operator-facing; no CLI flags required). @@ -1694,20 +1758,12 @@ void HttpServer::shutdown() { queue_tail_ = nullptr; } - // Shutdown save: persist all tracked snapshot slots to disk. - // Safe to access slot_tokens_ without locking — worker is joined. - if (!disk_cache_.disabled() && !slot_tokens_.empty()) { - std::fprintf(stderr, "[disk-cache] shutdown: saving %zu tracked slots\n", - slot_tokens_.size()); - backend_.snapshot_flush_deferred(); - for (auto & [slot, tokens] : slot_tokens_) { - if (backend_.snapshot_used(slot)) { - disk_cache_.learn_layout(slot); - disk_cache_.save(slot, tokens); - } - } - slot_tokens_.clear(); - } + // Graceful shutdown is one last eviction: write each conversation's + // newest snapshot and the heads. Safe without locking — the worker is + // joined. A crash or kill -9 loses what was only in RAM. + flush_sessions_to_disk("shutdown"); + slot_tokens_.clear(); + slot_aux_.clear(); } bool HttpServer::start_worker() { @@ -1921,19 +1977,11 @@ int HttpServer::run(const std::vector & models) { clients_cv_.wait(lk, [this]() { return active_clients_.load() == 0; }); } - // Persist disk cache (worker joined — no race on slot_tokens_). - if (!disk_cache_.disabled() && !slot_tokens_.empty()) { - std::fprintf(stderr, "[disk-cache] shutdown: saving %zu tracked slots\n", - slot_tokens_.size()); - backend_.snapshot_flush_deferred(); - for (auto & [slot, tokens] : slot_tokens_) { - if (backend_.snapshot_used(slot)) { - disk_cache_.learn_layout(slot); - disk_cache_.save(slot, tokens); - } - } - slot_tokens_.clear(); - } + // Graceful shutdown is one last eviction (worker joined — no race on + // slot_tokens_). + flush_sessions_to_disk("shutdown"); + slot_tokens_.clear(); + slot_aux_.clear(); #if defined(_WIN32) // Status-event sockets remain owned until shutdown(). Keep Winsock alive @@ -3688,6 +3736,7 @@ std::vector HttpServer::request_restore_points(const std::vector & HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( const ParsedRequest & req, PreparedPrompt & prepared, GenerateRequest & generate_request) { + request_session_slots_.clear(); if (req.images) return {}; auto & effective_prompt = prepared.tokens; // Tool-heavy requests prefer the reusable system/tool boundary under eviction. @@ -3842,6 +3891,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( } std::vector lookup_lengths; + DiskPrefixCache::PrefixProbe disk_probe; if (cache.disk_policy.mode == DiskPrefixCacheMode::Full && !effective_prompt.empty()) { // Probes are index lookups; only a hit reads a file. A restart then @@ -3851,6 +3901,17 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( (int) effective_prompt.size(), find_all_boundaries(effective_prompt, prefix_cache_.chat_markers()), config_.disk_cache_min_tokens); + // Session states (agent turns with their generated text) are keyed + // past the prompt's chat boundaries; the index knows their lengths. + const auto probe = disk_cache_.longest_prefix(effective_prompt); + disk_probe = probe; + if (probe.len > 0 && + std::find(lookup_lengths.begin(), lookup_lengths.end(), probe.len) == + lookup_lengths.end()) { + lookup_lengths.insert(lookup_lengths.begin(), probe.len); + std::sort(lookup_lengths.begin(), lookup_lengths.end(), + std::greater()); + } } else if (selected_boundary > 0) { lookup_lengths.push_back(selected_boundary); } else if (cache.disk_policy.mode == DiskPrefixCacheMode::Auto) { @@ -3863,29 +3924,42 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( } } - if (!cache.using_restore && !disk_cache_.disabled()) { + // A RAM hit does not end the search: a session whose deep snapshot was + // evicted to disk often still matches a shallow resident entry (the + // shared system/tools head). The deeper disk state wins; the index + // answers without a read, and a failed read keeps the RAM hit. + bool disk_lookup = !cache.using_restore; + const int ram_prefix_len = cache.using_restore ? cache.prefix_len : 0; + if (cache.using_restore && !disk_cache_.disabled() && + !prepared.compressed && disk_probe.len > 0 && disk_probe.cur_pos > cache.prefix_len && + disk_probe.len < (int) effective_prompt.size()) { + lookup_lengths.assign(1, disk_probe.len); + disk_lookup = true; + } + if (disk_lookup && !disk_cache_.disabled()) { for (int lookup_length : lookup_lengths) { std::vector prefix_tokens( effective_prompt.begin(), effective_prompt.begin() + lookup_length); if (!disk_cache_.lookup(prefix_tokens, kDiskStagingSlot)) continue; - cache.cache_slot = kDiskStagingSlot; - cache.prefix_len = backend_.snapshot_cur_pos(kDiskStagingSlot); - if (cache.prefix_len <= 0 || - cache.prefix_len > (int) effective_prompt.size()) { + const int hit_pos = backend_.snapshot_cur_pos(kDiskStagingSlot); + if (hit_pos <= 0 || hit_pos > (int) effective_prompt.size()) { std::fprintf(stderr, "[disk-cache] ignoring invalid hit pos=%d prompt=%zu\n", - cache.prefix_len, effective_prompt.size()); + hit_pos, effective_prompt.size()); backend_.snapshot_free(kDiskStagingSlot); continue; } + cache.cache_slot = kDiskStagingSlot; + cache.prefix_len = hit_pos; cache.using_restore = true; cache.disk_hit = true; std::fprintf(stderr, - "[disk-cache] hit policy=%s len=%d slot=%d pos=%d\n", + "[disk-cache] hit policy=%s len=%d slot=%d pos=%d (RAM had %d)\n", disk_prefix_cache_policy_name(cache.disk_policy).c_str(), - lookup_length, kDiskStagingSlot, cache.prefix_len); + lookup_length, kDiskStagingSlot, cache.prefix_len, + ram_prefix_len); break; } } @@ -3974,7 +4048,12 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( } } + // A compressed prompt is not raw history: never persist a prefix of it. + // With RAM prefix slots the inline snapshot holds the head and reaches + // disk when it is evicted, so this extra prefill-and-write is only the + // fallback without them. if (!cache.using_restore && !disk_cache_.disabled() && + prefix_cache_.disabled() && !prepared.compressed && cache.disk_policy.mode == DiskPrefixCacheMode::Full) { const auto boundaries = find_all_boundaries( effective_prompt, prefix_cache_.chat_markers()); @@ -4098,7 +4177,7 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( } cache.snap_prepared = cache.snap_slot >= 0; if (cache.snap_prepared) { - forget_inline_slot_metadata(cache.snap_slot); + evict_inline_slot(cache.snap_slot); backend_.snapshot_free(cache.snap_slot); generate_request.snap_slot = cache.snap_slot; generate_request.snap_pos = cache.snap_cut; @@ -4107,6 +4186,19 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache( backend_.snapshot_free(cache.full_snap_slot); } + if (cache.using_restore && disk_cache_digest_enabled()) { + backend_.snapshot_flush_deferred(); + int n_tensors = 0; + uint64_t n_bytes = 0; + const uint64_t digest = snapshot_state_digest( + backend_.snapshot_ref(cache.cache_slot), &n_tensors, &n_bytes); + std::fprintf(stderr, + "[disk-cache] digest %s source=%s pos=%d tensors=%d bytes=%llu " + "xxh64=%016llx\n", + req.response_id.c_str(), cache.disk_hit ? "ssd" : "ram", + cache.prefix_len, n_tensors, (unsigned long long)n_bytes, + (unsigned long long)digest); + } std::fprintf(stderr, "[server] chat CACHE %s restore=%s slot=%d prefix_len=%d " "effective_prompt=%zu pflash=%s disk_policy=%s disk_hit=%s " @@ -4193,17 +4285,16 @@ void HttpServer::finalize_generation_cache( effective_prompt, backend_.snapshot_bytes_estimate(saved_position)); trim_snapshots_after_commit(cache.snap_slot); - // The shutdown/disk key is exact: it must not claim rows the - // snapshot lacks. - slot_tokens_[cache.snap_slot] = std::vector( - effective_prompt.begin(), - effective_prompt.begin() + saved_position); - if (!disk_cache_.disabled()) { - disk_cache_.learn_layout(cache.snap_slot); - if (cache.disk_policy.mode == DiskPrefixCacheMode::Full) { - disk_cache_.save( - cache.snap_slot, slot_tokens_[cache.snap_slot]); - } + // The disk key is exact: it must not claim rows the snapshot + // lacks. A compressed view is never written to disk, so it + // gets no key. The write happens only if this snapshot is + // evicted from RAM (evict_inline_slot) or at shutdown. + if (!prepared.compressed && + cache.disk_policy.mode == DiskPrefixCacheMode::Full) { + slot_tokens_[cache.snap_slot] = std::vector( + effective_prompt.begin(), + effective_prompt.begin() + saved_position); + note_session_slot(cache.snap_slot); } } else { backend_.snapshot_free(cache.snap_slot); @@ -4219,7 +4310,10 @@ void HttpServer::finalize_generation_cache( // Long conversations get a continued checkpoint after crossing the // configured interval so a future turn can restore prompt plus output. - if (!disk_cache_.disabled() && + // With RAM prefix slots a session's state is written when it is evicted + // from RAM instead; this checkpoint is the fallback without them. + if (!disk_cache_.disabled() && prefix_cache_.disabled() && + !prepared.compressed && cache.disk_policy.mode == DiskPrefixCacheMode::Full && result.ok() && generation_produced_output) { const int final_position = @@ -4258,6 +4352,149 @@ void HttpServer::forget_inline_slot_metadata(int slot) { if (slot < 0) return; agent_turn_cache_slots_.erase(slot); slot_tokens_.erase(slot); + slot_aux_.erase(slot); +} + +bool HttpServer::persist_evicted_slot(int slot, bool persist, const char * reason) { + if (slot < 0 || !persist || disk_cache_.disabled()) return false; + const auto key = slot_tokens_.find(slot); + // No key: a compressed view, a non-`full` disk scope, or no snapshot. + if (key == slot_tokens_.end() || !backend_.snapshot_used(slot)) return false; + const auto aux = slot_aux_.find(slot); + OwnedSnapshot snap = backend_.snapshot_detach(slot); + if (snap.empty()) return false; + const int pos = snap.cur_pos(); + const bool queued = disk_cache_.save_async( + std::move(snap), key->second, + aux == slot_aux_.end() ? std::string() : aux->second); + const int pruned = queued + ? disk_cache_.prune_ancestors(key->second, config_.disk_cache_cold_max_tokens) + : 0; + std::fprintf(stderr, + "[disk-cache] %s slot=%d tokens=%zu pos=%d -> %s (pruned %d older)\n", + reason, slot, key->second.size(), pos, + queued ? "write queued" : "not written", pruned); + return queued; +} + +void HttpServer::evict_inline_slot(int slot) { + if (slot >= 0) { + persist_evicted_slot(slot, prefix_cache_.persist_on_evict(slot), "ram eviction"); + } + forget_inline_slot_metadata(slot); +} + +void HttpServer::note_session_slot(int slot) { + if (slot >= 0) request_session_slots_.push_back(slot); +} + +std::string HttpServer::session_replay_blob(const ParsedRequest & req, + const SseEmitter * emitter) { + json entries = json::array(); + if (tool_memory_.disabled()) return std::string(); + const auto messages = + normalize_chat_messages(req.messages, req.format, tool_memory_, true); + for (const auto & m : messages) { + if (!m.tool_calls_replayed || m.content.empty()) continue; + json ids = json::array(); + for (const auto & call : m.tool_calls) ids.push_back(call.id); + entries.push_back({{"ids", ids}, {"raw", m.content}}); + } + // This response's own calls, remembered by remember_agent_turn(). + std::vector ids; + if (emitter) { + for (const auto & call : emitter->tool_calls()) ids.push_back(call.id); + } + if (!ids.empty()) { + const std::string raw = tool_memory_.lookup(ids); + if (!raw.empty()) entries.push_back({{"ids", ids}, {"raw", raw}}); + } + if (entries.empty()) return std::string(); + return json{{"tool_memory", entries}}.dump(); +} + +void HttpServer::restore_replay_blobs() { + size_t restored = 0; + disk_cache_.for_each_aux([&](const std::string & blob) { + const json doc = json::parse(blob, nullptr, /*allow_exceptions=*/false); + if (!doc.is_object() || !doc.contains("tool_memory") || + !doc["tool_memory"].is_array()) { + return; + } + for (const auto & entry : doc["tool_memory"]) { + if (!entry.is_object() || !entry.contains("ids") || + !entry["ids"].is_array() || !entry.contains("raw") || + !entry["raw"].is_string()) { + continue; + } + std::vector ids; + for (const auto & id : entry["ids"]) { + if (id.is_string()) ids.push_back(id.get()); + } + tool_memory_.remember(ids, entry["raw"].get()); + ++restored; + } + }); + if (restored > 0) { + std::fprintf(stderr, + "[disk-cache] restored %zu tool-call replay entries\n", restored); + } +} + +void HttpServer::record_session_aux(const ParsedRequest & req, + const SseEmitter & emitter) { + std::vector slots; + slots.swap(request_session_slots_); + if (disk_cache_.disabled()) return; + std::string blob; + bool computed = false; + for (const int slot : slots) { + if (!slot_tokens_.count(slot)) continue; + if (!computed) { + blob = session_replay_blob(req, &emitter); + computed = true; + } + slot_aux_[slot] = blob; + } +} + +void HttpServer::flush_sessions_to_disk(const char * reason) { + if (disk_cache_.disabled()) return; + const auto t0 = std::chrono::steady_clock::now(); + const double limit = config_.disk_cache_shutdown_flush_s; + const auto elapsed = [&] { + return std::chrono::duration( + std::chrono::steady_clock::now() - t0).count(); + }; + size_t queued = 0, backups = 0, skipped = 0; + const auto before = disk_cache_.stats(); + // Graceful shutdown evicts everything: each conversation's newest + // snapshot and the heads, most recently used first, so the time limit + // drops the least valuable ones. Queuing can wait for the writer + // (bounded queued bytes), so the limit is checked between snapshots. + for (const auto & entry : prefix_cache_.entries_for_shutdown()) { + if (!entry.persist) { + ++backups; + continue; + } + if (limit >= 0 && elapsed() >= limit) { + ++skipped; + continue; + } + if (persist_evicted_slot(entry.slot, true, reason)) ++queued; + forget_inline_slot_metadata(entry.slot); + } + const auto drained = disk_cache_.drain( + limit < 0 ? -1.0 : (std::max)(0.0, limit - elapsed())); + std::fprintf(stderr, + "[disk-cache] %s flush: %zu snapshots queued (%zu edit backups not " + "written), %llu files %.1f MB written in %.2fs (limit %.0fs; %llu " + "writes and %zu snapshots dropped by it)\n", + reason, queued, backups, + (unsigned long long)(disk_cache_.stats().writes - before.writes), + (double)(disk_cache_.stats().write_bytes - before.write_bytes) / (1024.0 * 1024.0), + elapsed(), limit, + (unsigned long long)drained.cancelled, skipped); } void HttpServer::trim_snapshots_after_commit(int slot) { @@ -4268,7 +4505,13 @@ void HttpServer::trim_snapshots_after_commit(int slot) { } }; release(prefix_cache_.prune_superseded_ancestors(slot)); - release(prefix_cache_.enforce_resident_budget(slot)); + std::vector persist; + const auto evicted = prefix_cache_.enforce_resident_budget(slot, &persist); + for (size_t i = 0; i < evicted.size(); ++i) { + persist_evicted_slot(evicted[i], persist[i], "ram eviction"); + forget_inline_slot_metadata(evicted[i]); + backend_.snapshot_free(evicted[i]); + } } void HttpServer::publish_memory_report() { @@ -4392,7 +4635,7 @@ bool HttpServer::save_generated_turn( return false; } const int slot = reservation.slot(); - forget_inline_slot_metadata(slot); + evict_inline_slot(slot); backend_.snapshot_free(slot); // The next request usually continues this state, so the backend may // keep it live instead of copying it out. @@ -4405,12 +4648,16 @@ bool HttpServer::save_generated_turn( } reservation.commit_at( canonical, saved_pos, backend_.snapshot_bytes_estimate(saved_pos)); - for (const int evicted : prefix_cache_.enforce_resident_budget(slot)) { - forget_inline_slot_metadata(evicted); - backend_.snapshot_free(evicted); + std::vector persist; + const auto evicted = prefix_cache_.enforce_resident_budget(slot, &persist); + for (size_t i = 0; i < evicted.size(); ++i) { + persist_evicted_slot(evicted[i], persist[i], "ram eviction"); + forget_inline_slot_metadata(evicted[i]); + backend_.snapshot_free(evicted[i]); } slot_tokens_[slot] = std::vector( canonical.begin(), canonical.begin() + saved_pos); + note_session_slot(slot); agent_turn_cache_slots_.insert(slot); std::fprintf(stderr, "[agent-turn-cache] saved generated turn slot=%d prefix=%d " @@ -4535,7 +4782,7 @@ void HttpServer::remember_agent_turn( reservation.cancel(); return; } - forget_inline_slot_metadata(slot); + evict_inline_slot(slot); backend_.snapshot_free(slot); GenerateRequest replay; @@ -4557,6 +4804,7 @@ void HttpServer::remember_agent_turn( trim_snapshots_after_commit(slot); canonical_tokens.resize((size_t) saved_pos); slot_tokens_[slot] = std::move(canonical_tokens); + note_session_slot(slot); agent_turn_cache_slots_.insert(slot); std::fprintf(stderr, "[agent-turn-cache] saved slot=%d prefix=%d source=%d:%d " @@ -5130,6 +5378,7 @@ void HttpServer::process_job(ServerJob * job) { req, prepared, cache, result, emitter, completion_tokens, visible_output_seen, client_disconnected, /*replay_cache=*/true); + record_session_aux(req, emitter); if (client_disconnected) { std::fprintf(stderr, "[server] client disconnected — generation aborted " diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 0e9d9a029b..6d2f852302 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -254,10 +254,14 @@ struct ServerConfig { // Disk prefix cache std::string disk_cache_dir; // empty = disabled - size_t disk_cache_budget_mb = 4096; // max disk usage in MB + size_t disk_cache_budget_mb = 131072; // max disk usage in MB (128 GB) int disk_cache_min_tokens = 512; // only persist >= this many tokens int disk_cache_continued_interval = 10240; // continued checkpoint every N tokens int disk_cache_cold_max_tokens = 10240; // cold prefix for prompts longer than this + double disk_cache_ttl_hours = 72.0; // delete files unused this long (0 = never) + // Graceful shutdown writes the RAM snapshots to disk (one last + // eviction) for at most this long; < 0 = no limit. + double disk_cache_shutdown_flush_s = 120.0; DiskPrefixCachePolicy disk_cache_policy; // Optional Jinja chat template (overrides the hardcoded ChatFormat::QWEN3 @@ -569,6 +573,28 @@ class HttpServer { bool visible_output_seen, bool client_disconnected, bool replay_cache); void forget_inline_slot_metadata(int slot); + // Free a RAM snapshot slot for reuse, writing it to disk first when + // PrefixCache::persist_on_evict() says so (see persist_evicted_slot). + void evict_inline_slot(int slot); + // Disk persistence (docs/PREFIX_CACHE.md): a snapshot is written only + // when it leaves RAM, never after every turn. Hands the payload of + // `slot` to the disk writer when `persist` is set and the slot has a disk + // key (slot_tokens_); the slot is left empty. Returns true when queued. + bool persist_evicted_slot(int slot, bool persist, const char * reason); + // Slots committed during the current request get this request's + // tool-call replay text (slot_aux_) once its response is known. + void note_session_slot(int slot); + void record_session_aux(const ParsedRequest & req, const SseEmitter & emitter); + // The tool-call replay text a session's prompt was rendered with + // (ToolMemory), stored with its disk snapshot so a restart renders the + // same tokens; restored into ToolMemory at startup. + std::string session_replay_blob(const ParsedRequest & req, + const SseEmitter * emitter); + void restore_replay_blobs(); + // Graceful shutdown: evict every RAM snapshot to disk (newest of each + // conversation and the heads; not edit backups), bounded by + // disk_cache_shutdown_flush_s. + void flush_sessions_to_disk(const char * reason); // After a commit in `slot`: free the snapshots it supersedes, then evict // until committed snapshots fit the resident budget again. void trim_snapshots_after_commit(int slot); @@ -706,9 +732,12 @@ class HttpServer { std::string status_html_path_; std::string resolve_status_html(); - // Track prompt tokens for each snapshot slot (for shutdown save). + // Disk key of each RAM snapshot slot (written on eviction/shutdown). std::unordered_map> slot_tokens_; std::unordered_set agent_turn_cache_slots_; + // Replay blob stored with each slot's disk snapshot. + std::unordered_map slot_aux_; + std::vector request_session_slots_; // Memory held by the live cache and saved snapshots, rebuilt on the // generation thread after each request and read by /props (see diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index c2ee5bc628..e40136394b 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -901,6 +901,7 @@ void PrefixCache::replace_inline_entry( if (entries_[(size_t)i].slot == slot) { std::fprintf(stderr, "[pc] dropping stale entry for reused slot=%d\n", slot); + record_eviction(i); erase_inline_entry(i); } } @@ -932,7 +933,10 @@ bool PrefixCache::commit_inline_reservation( } if (reservation.has_victim_) { const int victim = find_entry(reservation.victim_); - if (victim >= 0) erase_inline_entry(victim); + if (victim >= 0) { + record_eviction(victim); + erase_inline_entry(victim); + } } replace_inline_entry( reservation.slot_, committed_cut, prompt_ids, @@ -1044,6 +1048,7 @@ std::vector PrefixCache::prune_superseded_ancestors( if (previous >= 0 && (max_resident_bytes_ == 0 || resident_bytes_ - others <= max_resident_bytes_)) { superseded.erase(std::find(superseded.begin(), superseded.end(), previous)); + entries_[(size_t)previous].edit_backup = true; } size_t freed = 0; // Erase from the back so earlier indices stay valid. @@ -1063,7 +1068,8 @@ std::vector PrefixCache::prune_superseded_ancestors( return pruned; } -std::vector PrefixCache::enforce_resident_budget(int keep_slot) { +std::vector PrefixCache::enforce_resident_budget( + int keep_slot, std::vector * persist) { std::vector evicted; if (disabled_ || max_resident_bytes_ == 0) return evicted; const size_t before = resident_bytes_; @@ -1074,7 +1080,10 @@ std::vector PrefixCache::enforce_resident_budget(int keep_slot) { // ancestor; never the new entry or a protected pin. const int victim = pick_evict_victim(keep); if (victim < 0) break; - evicted.push_back(entries_[(size_t)victim].slot); + const int slot = entries_[(size_t)victim].slot; + evicted.push_back(slot); + if (persist) persist->push_back(persist_on_evict(slot)); + record_eviction(victim); erase_inline_entry(victim); } if (!evicted.empty()) { @@ -1086,6 +1095,38 @@ std::vector PrefixCache::enforce_resident_budget(int keep_slot) { return evicted; } +bool PrefixCache::persist_on_evict(int slot) const { + const int idx = find_slot_entry(slot); + if (idx < 0) return false; + const auto & entry = entries_[(size_t)idx]; + if (!entry.edit_backup || entry.protect) return true; + for (int i = 0; i < (int)entries_.size(); ++i) { + if (i != idx && is_strict_prefix(entry.ids, entries_[(size_t)i].ids)) { + return false; // its conversation's deeper snapshot is resident + } + } + return true; +} + +std::vector PrefixCache::entry_tokens(int slot) const { + const int idx = find_slot_entry(slot); + return idx < 0 ? std::vector{} : entries_[(size_t)idx].ids; +} + +void PrefixCache::record_eviction(int idx) { + if (!record_evictions_ || idx < 0 || idx >= (int)entries_.size()) return; + const auto & entry = entries_[(size_t)idx]; + evicted_log_.push_back({entry.slot, entry.ids, persist_on_evict(entry.slot)}); +} + +std::vector PrefixCache::entries_for_shutdown() const { + std::vector out; + for (auto it = entries_.rbegin(); it != entries_.rend(); ++it) { + out.push_back({it->slot, it->ids, persist_on_evict(it->slot)}); + } + return out; +} + static void update_atomic_max(std::atomic & value, uint64_t candidate) { uint64_t current = value.load(std::memory_order_relaxed); diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index 73ce911aa9..e80b1d8d7d 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -266,7 +266,37 @@ class PrefixCache { // Evicts committed entries, never `keep_slot` or protected pins, until the // resident bytes fit the budget again. Returns the evicted slots; the // caller frees their snapshot payloads. - std::vector enforce_resident_budget(int keep_slot); + // `persist` (optional) gets persist_on_evict() of each evicted slot, + // decided before it was erased. + std::vector enforce_resident_budget( + int keep_slot, std::vector * persist = nullptr); + + // ── Disk persistence of evicted entries ───────────────────────── + // Whether the entry in `slot`, about to leave RAM, should be written to + // disk: every conversation's newest snapshot and the shared heads are; + // a previous turn kept only as an edit backup (see + // prune_superseded_ancestors) is not while a deeper snapshot of its + // conversation is resident. False when `slot` holds no entry. + bool persist_on_evict(int slot) const; + // Token ids an entry in `slot` covers (empty when none). + std::vector entry_tokens(int slot) const; + struct Evicted { + int slot = -1; + std::vector ids; + bool persist = false; + }; + // Opt-in log of entries a commit replaced (its reservation's victim) or + // the resident budget evicted, for callers whose payload outlives the + // metadata (the concurrent scheduler writes them to disk). + void set_record_evictions(bool on) { record_evictions_ = on; } + std::vector take_evicted() { + std::vector out; + out.swap(evicted_log_); + return out; + } + // Every committed entry, most recently used first, with its + // persist_on_evict() flag: graceful shutdown evicts them all. + std::vector entries_for_shutdown() const; // Declares that the caller prunes after every commit, so the budget check // counts the entries a capture supersedes as freed when it lands. @@ -346,6 +376,9 @@ class PrefixCache { std::vector ids; // prefix tokens [0, target_cut) for prefix-aware eviction bool protect = false; // sticky tools-boundary pin size_t resident_bytes = 0; + // Kept behind a deeper snapshot of its conversation only so an + // edit of the last message can resume (prune_superseded_ancestors). + bool edit_backup = false; }; std::vector entries_; int next_slot_ = 0; @@ -364,6 +397,10 @@ class PrefixCache { uint64_t next_inline_reservation_ = 1; size_t max_resident_bytes_ = 0; bool prunes_superseded_ = false; + bool record_evictions_ = false; + std::vector evicted_log_; + // Log entries_[idx] into evicted_log_ (when recording) before an erase. + void record_eviction(int idx); size_t resident_bytes_ = 0; // Full-cache state diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 0d228e306a..543f144dd1 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include @@ -74,6 +76,7 @@ struct SchedSlot { enum class AdmissionDisposition { Admitted, // job owns an engine slot and its first token is on the wire Deferred, // engine had no room; caller keeps the job and retries it first + Loading, // its disk snapshot is being read; retried first, without backoff Retired, // job finished or failed during admission; no slot was taken }; @@ -101,6 +104,43 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { ServerJob * deferred = nullptr; std::chrono::steady_clock::time_point deferred_retry_at{}; + // Disk session cache. Checkpoints that leave RAM (a capture replacing + // its reservation's victim) are written by the disk writer thread; a + // prompt whose deepest state is on disk waits at the head of the line + // while a reader thread loads it, so live streams keep decoding. Only + // engines whose checkpoints live in the backend's snapshot slots + // (Qwen paged) take part. + const bool disk_enabled = !disk_cache_.disabled() && + engine.supports_prefix_store() && !prefix_cache_.disabled() && + engine.prefix_store_in_backend_snapshots(); + if (!disk_cache_.disabled() && !disk_enabled) { + std::fprintf(stderr, + "[disk-cache] off for concurrent serving: this engine keeps no " + "checkpoints in backend snapshot slots\n"); + } + prefix_cache_.set_record_evictions(disk_enabled); + struct DiskLoad { + ServerJob * job = nullptr; + std::vector key; + struct Result { + std::atomic done{false}; + OwnedSnapshot snap; + }; + std::shared_ptr result; + std::thread reader; + }; + std::optional disk_load; + std::vector finished_readers; + // Staging checkpoint for a disk restore: PrefixCache never uses it. + constexpr int kDiskStagingIndex = PrefixCache::MAX_CACHE_SLOTS; + auto drop_disk_load = [&]() { + if (!disk_load) return; + if (disk_load->reader.joinable()) { + finished_readers.push_back(std::move(disk_load->reader)); + } + disk_load.reset(); + }; + // Degenerate-run guard shared with do_ar_decode: explicit env override, // else 32 when the min-tokens floor is active, else off. static const int repeat_guard = [] { @@ -450,6 +490,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { PrefixCaptureTxn prepared_capture; PrefixCache::InlineReservation capture_reservation; int restore_policy_slot = -1; + bool disk_restore = false; // Tokens alone do not identify an image: image requests never touch // the prefix cache. const bool prefix_supported = !req.images && @@ -464,6 +505,65 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefix_plan.restore = { (uint64_t)hit.first + 1, hit.second}; } + // A deeper state on disk (the index answers without a read): + // load it off this thread, then restore from the staging slot. + const auto probe = disk_enabled + ? disk_cache_.longest_prefix(req.prompt_tokens) + : DiskPrefixCache::PrefixProbe{}; + if (disk_load && disk_load->job != job) drop_disk_load(); + if (probe.cur_pos > prefix_plan.restore.tokens && + probe.cur_pos == probe.len && + probe.cur_pos < (int)req.prompt_tokens.size()) { + if (!disk_load) { + disk_load.emplace(); + disk_load->job = job; + disk_load->key.assign(req.prompt_tokens.begin(), + req.prompt_tokens.begin() + probe.len); + disk_load->result = std::make_shared(); + disk_load->reader = std::thread( + [this, key = disk_load->key, out = disk_load->result]() { + out->snap = disk_cache_.load(key); + out->done.store(true, std::memory_order_release); + }); + std::fprintf(stderr, + "[disk-cache] %s: reading %d-token state from disk " + "(RAM has %d)\n", req.response_id.c_str(), probe.len, + prefix_plan.restore.tokens); + } + if (!disk_load->result->done.load(std::memory_order_acquire)) { + return AdmissionDisposition::Loading; + } + disk_load->reader.join(); + OwnedSnapshot snap = std::move(disk_load->result->snap); + disk_load.reset(); + const int pos = snap.cur_pos(); + if (!snap.empty() && + disk_cache_.adopt(std::move(snap), kDiskStagingIndex)) { + disk_restore = true; + restore_policy_slot = -1; + prefix_plan.restore = {(uint64_t)kDiskStagingIndex + 1, pos}; + } else { + std::fprintf(stderr, + "[disk-cache] %s: disk state unusable, %s\n", + req.response_id.c_str(), + prefix_plan.restore.valid() ? "restoring from RAM" + : "prefilling"); + } + } + + if (prefix_plan.restore.valid() && disk_cache_digest_enabled()) { + int n_tensors = 0; + uint64_t n_bytes = 0; + const uint64_t digest = snapshot_state_digest( + backend_.snapshot_ref((int)prefix_plan.restore.id - 1), + &n_tensors, &n_bytes); + std::fprintf(stderr, + "[disk-cache] digest %s source=%s pos=%d tensors=%d " + "bytes=%llu xxh64=%016llx\n", + req.response_id.c_str(), disk_restore ? "ssd" : "ram", + prefix_plan.restore.tokens, n_tensors, + (unsigned long long)n_bytes, (unsigned long long)digest); + } // The restore source is never the capture destination: the new // checkpoint lands in a different slot so the restore point can @@ -539,7 +639,25 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefix_cache_.invalidate_inline_snap(restore_policy_slot); } } - if (prefix.invalidated.valid()) { + if (disk_restore) { + // The engine copied the staging checkpoint into the sequence's + // pages during admission (or rejected it); it is not kept. + backend_.snapshot_free(kDiskStagingIndex); + if (prefix.restored.valid() && + prefix.restored == requested_plan.restore) { + std::fprintf(stderr, + "[disk-cache] hit %s: restored %d tokens from disk\n", + req.response_id.c_str(), prefix.restored.tokens); + } + } + if (disk_restore && prefix.invalidated.valid() && + prefix.invalidated == requested_plan.restore) { + // A disk state that does not fit this cache layout: the engine + // already admitted the prompt cold. + std::fprintf(stderr, + "[disk-cache] %s: disk state rejected by the engine, prefilling\n", + req.response_id.c_str()); + } else if (prefix.invalidated.valid()) { if (prefix.invalidated != requested_plan.restore || restore_policy_slot < 0) { prefix_protocol_error = @@ -761,6 +879,12 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { continue; } const AdmissionDisposition outcome = admit_job(job); + if (outcome == AdmissionDisposition::Loading) { + deferred = job; + deferred_retry_at = std::chrono::steady_clock::now() + + std::chrono::milliseconds(2); + break; // keep decoding while the disk read runs + } if (outcome == AdmissionDisposition::Deferred) { deferred = job; deferred_retry_at = std::chrono::steady_clock::now() + @@ -1050,14 +1174,51 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { using Resolution = PrefixCaptureTxn::Resolution; const Resolution resolution = s.cache_capture.resolve( out.prefix_store, s.job->req.prompt_tokens); + const int captured = + (int)out.prefix_store.ticket.checkpoint.id - 1; + // The checkpoint this capture replaced in its slot (the + // reservation's victim) left RAM: write it to disk unless it + // was only an edit backup. Claimed on every outcome. + OwnedSnapshot displaced = disk_enabled + ? backend_.snapshot_take_displaced(captured) : OwnedSnapshot{}; + for (auto & evicted : prefix_cache_.take_evicted()) { + if (evicted.slot != captured || displaced.empty()) continue; + if (evicted.persist) { + const auto aux = slot_aux_.find(captured); + const int pos = displaced.cur_pos(); + const bool queued = disk_cache_.save_async( + std::move(displaced), evicted.ids, + aux == slot_aux_.end() ? std::string() : aux->second); + if (queued) { + disk_cache_.prune_ancestors( + evicted.ids, config_.disk_cache_cold_max_tokens); + } + std::fprintf(stderr, + "[disk-cache] ram eviction checkpoint=%d tokens=%zu " + "pos=%d -> %s\n", captured + 1, evicted.ids.size(), + pos, queued ? "write queued" : "not written"); + } + displaced.reset(); + } if (resolution == Resolution::saved) { + // The slot holds the new checkpoint now. + forget_inline_slot_metadata(captured); + if (disk_enabled) { + const int tokens = out.prefix_store.ticket.checkpoint.tokens; + const auto & prompt = s.job->req.prompt_tokens; + slot_tokens_[captured] = std::vector( + prompt.begin(), prompt.begin() + tokens); + slot_aux_[captured] = + session_replay_blob(s.job->req, nullptr); + } // A conversation's checkpoints are prefixes of each other; // restores use only the deepest, so the ones it supersedes // would hold memory the other conversations need. std::vector tokens; const auto pruned = prefix_cache_.prune_superseded_ancestors( - (int)out.prefix_store.ticket.checkpoint.id - 1, &tokens); + captured, &tokens); for (size_t i = 0; i < pruned.size(); ++i) { + forget_inline_slot_metadata(pruned[i]); engine.discard_prefix_store({(uint64_t)pruned[i] + 1, tokens[i]}); } } else if (resolution == Resolution::failed) { @@ -1139,6 +1300,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } } + drop_disk_load(); + for (auto & reader : finished_readers) reader.join(); + finished_readers.clear(); + // Shutdown: unblock every parked client thread. for (int i = 0; i < n_slots; i++) { if (slots[(size_t)i].job) { diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 1d76bc53b6..8a335e29de 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -311,8 +311,15 @@ static void print_usage(const char * prog) { " --prefill-upstream-model Model name on forwarded requests.\n" "\n" "Disk KV cache:\n" - " --kv-cache-dir Directory for ondisk KV cache (enables feature)\n" - " --kv-cache-budget Max disk usage in MB (default: 4096)\n" + " --kv-cache-dir |off Directory for the on-disk KV session cache\n" + " (default: $XDG_CACHE_HOME/lucebox/kv-cache, else\n" + " ~/.cache/lucebox/kv-cache; off disables it)\n" + " --no-kv-cache Same as --kv-cache-dir off\n" + " --kv-cache-budget Max disk usage in MB, hard cap with LRU eviction,\n" + " clamped to the free space (default: 131072 = 128 GB)\n" + " --kv-cache-ttl Delete snapshots unused for this long (default: 72, 0 = never)\n" + " --kv-cache-shutdown-flush Max seconds a graceful shutdown spends writing\n" + " RAM snapshots to disk (default: 120, -1 = no limit)\n" " --kv-cache-min-tokens Min tokens to persist (default: 512)\n" " --kv-cache-interval Continued checkpoint every N tokens (default: 10240)\n" " --kv-cache-cold-max Cold prefix for prompts longer than N tokens (default: 10240)\n" @@ -417,6 +424,16 @@ struct ModelOptions { }; +// $XDG_CACHE_HOME/lucebox/kv-cache, else ~/.cache/lucebox/kv-cache; empty +// (disk cache off) when neither variable is set. +static std::string default_kv_cache_dir() { + const char * xdg = std::getenv("XDG_CACHE_HOME"); + if (xdg && xdg[0] == '/') return std::string(xdg) + "/lucebox/kv-cache"; + const char * home = std::getenv("HOME"); + if (home && home[0] == '/') return std::string(home) + "/.cache/lucebox/kv-cache"; + return std::string(); +} + static int parse_model_options(int argc, char ** argv, ModelOptions & model, bool load_balancing, bool first_model) { if (argc < 2 || argv[1][0] == '-') { @@ -438,6 +455,8 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, bool target_device_seen = false; bool target_devices_seen = false; bargs.model_path = argv[1]; + // The disk session cache is on by default (--kv-cache-dir off disables). + sconfig.disk_cache_dir = default_kv_cache_dir(); for (int i = 2; i < argc; i++) { const std::string option = argv[i]; @@ -981,9 +1000,31 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, std::fprintf(stderr, "[server] loaded chat template from %s (%ld bytes)\n", path, n); } } else if (std::strcmp(argv[i], "--kv-cache-dir") == 0 && i + 1 < argc) { - sconfig.disk_cache_dir = argv[++i]; + const std::string dir = argv[++i]; + sconfig.disk_cache_dir = + (dir == "off" || dir == "none" || dir.empty()) ? std::string() : dir; + } else if (std::strcmp(argv[i], "--no-kv-cache") == 0) { + sconfig.disk_cache_dir.clear(); + } else if (std::strcmp(argv[i], "--kv-cache-shutdown-flush") == 0 && i + 1 < argc) { + char * end = nullptr; + const double secs = std::strtod(argv[++i], &end); + if (!end || *end != '\0') { + std::fprintf(stderr, "--kv-cache-shutdown-flush expects seconds, got %s\n", + argv[i]); + return 1; + } + sconfig.disk_cache_shutdown_flush_s = secs; } else if (std::strcmp(argv[i], "--kv-cache-budget") == 0 && i + 1 < argc) { sconfig.disk_cache_budget_mb = (size_t)std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--kv-cache-ttl") == 0 && i + 1 < argc) { + char * end = nullptr; + const double hours = std::strtod(argv[++i], &end); + if (!end || *end != '\0' || !(hours >= 0.0)) { + std::fprintf(stderr, "--kv-cache-ttl expects hours >= 0 (0 = never), got %s\n", + argv[i]); + return 1; + } + sconfig.disk_cache_ttl_hours = hours; } else if (std::strcmp(argv[i], "--kv-cache-min-tokens") == 0 && i + 1 < argc) { sconfig.disk_cache_min_tokens = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--kv-cache-interval") == 0 && i + 1 < argc) { @@ -1473,30 +1514,32 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod // Continuous Qwen serving supports copied in-memory prefix checkpoints: // restore allocates fresh pages and scatters logical K/V through the new - // sequence's block table. Exact-prefill and disk snapshots still use the - // classic single-sequence format and remain disabled in paged mode. + // sequence's block table. Checkpoints that leave RAM go to the disk + // session cache (the scheduler writes and reads them). Exact-prefill + // snapshots still use the classic single-sequence format and remain + // disabled in paged mode; single-sequence paged serving keeps no + // snapshots at all, so it has no disk cache either. // This rewrites ServerConfig rather than rejecting the launch, which is // why it lives here and not in the gate. if (backend_cache.paged_attention) { - if (backend_execution.max_concurrency > 1) { - if (sconfig.prefix_cache_cap > 0) { - std::fprintf(stderr, - "[server] concurrent paged serving enables copied in-memory " - "prefix checkpoints; full-prefill and disk caches remain disabled\n"); - } else { - std::fprintf(stderr, - "[server] concurrent paged serving: prefix checkpoints are " - "disabled; full-prefill and disk caches remain disabled\n"); - } + if (backend_execution.max_concurrency > 1 && sconfig.prefix_cache_cap > 0) { + std::fprintf(stderr, + "[server] concurrent paged serving enables copied in-memory " + "prefix checkpoints and %s; the full-prefill cache stays disabled\n", + sconfig.disk_cache_dir.empty() ? "no disk cache (off)" + : "the disk session cache"); } else { std::fprintf(stderr, - "[server] single-sequence --paged-attention still disables " - "prefix snapshots\n"); + "[server] %s: prefix checkpoints and the disk cache are " + "disabled\n", + backend_execution.max_concurrency > 1 + ? "concurrent paged serving without prefix slots" + : "single-sequence --paged-attention"); sconfig.prefix_cache_cap = 0; + sconfig.disk_cache_dir.clear(); + sconfig.disk_cache_policy.mode = DiskPrefixCacheMode::Off; } sconfig.prefill_cache_cap = 0; - sconfig.disk_cache_dir.clear(); - sconfig.disk_cache_policy.mode = DiskPrefixCacheMode::Off; } // Every concurrent engine takes its prefix checkpoints from the scheduler: // paged engines copy pages, qwen4exp copies full-cache snapshots. Both use diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index e1aa19ee81..39e890330f 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -163,6 +163,26 @@ struct SchedulerTestHarness { return server.slot_tokens_.at(slot); } + static void set_slot_tokens(HttpServer & server, int slot, + std::vector tokens) { + server.slot_tokens_[slot] = std::move(tokens); + } + + static DiskPrefixCache & disk_cache(HttpServer & server) { + return server.disk_cache_; + } + + static void flush_sessions(HttpServer & server) { + server.flush_sessions_to_disk("shutdown"); + } + + // What prepare_generation_cache does with a reservation's slot before + // the new snapshot overwrites it. + static void evict_inline_slot(HttpServer & server, int slot) { + server.evict_inline_slot(slot); + server.backend_.snapshot_free(slot); + } + static void publish_memory_report(HttpServer & server) { server.publish_memory_report(); } @@ -8900,10 +8920,11 @@ static void rm_rf(const std::string & path) { TEST_CASE(ServerUnitFixture, test_disk_cache_config_defaults) { DiskCacheConfig cfg; TEST_ASSERT(cfg.cache_dir.empty()); - TEST_ASSERT(cfg.budget_bytes == (size_t)4 * 1024 * 1024 * 1024); + TEST_ASSERT(cfg.budget_bytes == (size_t)128 * 1024 * 1024 * 1024); TEST_ASSERT(cfg.min_tokens == 512); TEST_ASSERT(cfg.continued_interval == 10240); TEST_ASSERT(cfg.cold_max_tokens == 10240); + TEST_ASSERT(cfg.ttl_seconds == 72ull * 3600); } TEST_CASE(ServerUnitFixture, test_disk_cache_policy_parse) { @@ -9056,12 +9077,14 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_init_creates_directory) { } TEST_CASE(ServerUnitFixture, test_disk_cache_header_size) { - // The header should be exactly 80 bytes. - TEST_ASSERT(DISK_CACHE_HEADER_SIZE == 80); + // The header should be exactly 96 bytes. + TEST_ASSERT(DISK_CACHE_HEADER_SIZE == 96); + TEST_ASSERT(DISK_CACHE_LAST_USED_OFFSET == 68); // Bumped to 2 when the K-rotation default changed: a cache written by an // older binary stores K in the rotated basis, and the layout id does not // cover that, so the version is what rejects it. - TEST_ASSERT(DISK_CACHE_VERSION == 3); + // v4 added the identity salt and the stored key tokens. + TEST_ASSERT(DISK_CACHE_VERSION == 5); } TEST_CASE(ServerUnitFixture, test_disk_cache_header_round_trip) { @@ -9082,6 +9105,7 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_header_round_trip) { hdr.created_at = 1700000000; hdr.last_used = 1700000100; hdr.last_tok = 151643; + std::memset(hdr.identity_salt, 0xEF, 16); // Use DiskPrefixCache's static write/read_header (they are private, so // we test indirectly through file I/O matching the on-disk format). @@ -9100,6 +9124,7 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_header_round_trip) { u64 = hdr.created_at; std::fwrite(&u64, 8, 1, f); u64 = hdr.last_used; std::fwrite(&u64, 8, 1, f); int32_t i32 = hdr.last_tok; std::fwrite(&i32, 4, 1, f); + std::fwrite(hdr.identity_salt, 16, 1, f); std::fclose(f); // Verify file size is DISK_CACHE_HEADER_SIZE. @@ -9123,6 +9148,8 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_header_round_trip) { std::fread(&ru64, 8, 1, f); TEST_ASSERT(ru64 == 1700000000); // created_at std::fread(&ru64, 8, 1, f); TEST_ASSERT(ru64 == 1700000100); // last_used int32_t ri32; std::fread(&ri32, 4, 1, f); TEST_ASSERT(ri32 == 151643); // last_tok + uint8_t salt[16]; std::fread(salt, 16, 1, f); + TEST_ASSERT(salt[0] == 0xEF && salt[15] == 0xEF); std::fclose(f); remove_test_path(path); @@ -9339,22 +9366,22 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_rejects_snapshot_past_key) { std::string dir = test_tmp_path("luce_test_past_key").string(); rm_rf(dir); DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; - DiskPrefixCache cache(cfg, backend); - TEST_ASSERT(cache.init()); - cache.learn_layout(0); - std::vector short_key; for (int i = 0; i < MockBackendWithLayout::kMaxPos - 4; ++i) short_key.push_back(i + 1); std::vector full_key; for (int i = 0; i < MockBackendWithLayout::kMaxPos; ++i) full_key.push_back(i + 1); - - TEST_ASSERT(!cache.save(0, short_key)); - TEST_ASSERT(cache.total_bytes() == 0); - TEST_ASSERT(cache.save(0, full_key)); - TEST_ASSERT(cache.total_bytes() > 0); - TEST_ASSERT(!cache.lookup(short_key, 1)); - TEST_ASSERT(cache.lookup(full_key, 1)); - TEST_ASSERT(backend.adopted_cur_pos_ == MockBackendWithLayout::kMaxPos); + { + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(!cache.save(0, short_key)); + TEST_ASSERT(cache.total_bytes() == 0); + TEST_ASSERT(cache.save(0, full_key)); + TEST_ASSERT(cache.total_bytes() > 0); + TEST_ASSERT(!cache.lookup(short_key, 1)); + TEST_ASSERT(cache.lookup(full_key, 1)); + TEST_ASSERT(backend.adopted_cur_pos_ == MockBackendWithLayout::kMaxPos); + } // the writer finishes and the directory lock is released // Forge the pre-fix shape on disk: key shorter than the snapshot. The // scan keys entries by the header's token hash, so only the header @@ -9537,6 +9564,1082 @@ TEST_CASE(ServerUnitFixture, test_disk_identity_salt_zero_is_backcompat) { rm_rf(dir2); } +// ─── Session persistence: TTL, LRU budget, restart, validation ────────── + +namespace { + +std::vector disk_test_key(int n, int32_t base = 1) { + std::vector key; + for (int i = 0; i < n; ++i) key.push_back(base + i); + return key; +} + +std::vector dkv_files(const std::string & dir) { + std::vector out; + std::error_code ec; + for (auto & e : fs::recursive_directory_iterator(dir, ec)) { + if (e.is_regular_file() && e.path().extension() == ".dkv") { + out.push_back(e.path().string()); + } + } + std::sort(out.begin(), out.end()); + return out; +} + +// Size of one snapshot file the mock produces for a key of `n` tokens. +uint64_t one_file_size(MockBackendWithAdopt & backend, const std::string & dir, int n) { + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + uint64_t size = 0; + { + DiskPrefixCache cache(cfg, backend); + cache.init(); + cache.learn_layout(0); + cache.save(0, disk_test_key(n)); + size = cache.total_bytes(); + } + rm_rf(dir); + return size; +} + +} // namespace + +TEST_CASE(ServerUnitFixture, test_disk_cache_ttl_expiry) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_ttl").string(); + rm_rf(dir); + uint64_t clock = 1000; + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + cfg.ttl_seconds = 100; + const auto a = disk_test_key(40, 1); + const auto b = disk_test_key(40, 500); + const auto c = disk_test_key(40, 900); + { + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, a)); + clock = 1090; + TEST_ASSERT(cache.lookup(a, 1)); // used at 1090 + clock = 1180; + TEST_ASSERT(cache.lookup(a, 1)); // 90 s idle: still alive + clock = 1281; + TEST_ASSERT(!cache.lookup(a, 1)); // 101 s idle: expired + deleted + TEST_ASSERT(cache.entry_count() == 0); + TEST_ASSERT(dkv_files(dir).empty()); + + // Insertion sweep: saving b deletes an expired c. + TEST_ASSERT(cache.save(0, c)); + clock = 1400; + TEST_ASSERT(cache.save(0, b)); + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(!cache.lookup(c, 1)); + } + { + // Startup sweep: b expires while the server is down. + clock = 1600; + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 0); + TEST_ASSERT(dkv_files(dir).empty()); + } + { + // ttl 0 = never expire. + DiskCacheConfig never = cfg; never.ttl_seconds = 0; + DiskPrefixCache cache(never, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, a)); + clock += 10ull * 365 * 24 * 3600; + TEST_ASSERT(cache.lookup(a, 1)); + } + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_budget_lru_survives_restart) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_lru").string(); + const uint64_t one = one_file_size(backend, dir, 40); + TEST_ASSERT(one > 0); + uint64_t clock = 1; + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + cfg.ttl_seconds = 0; + cfg.budget_bytes = (size_t)(one * 2 + one / 2); // room for two files + const auto a = disk_test_key(40, 1); + const auto b = disk_test_key(40, 500); + const auto c = disk_test_key(40, 900); + const auto d = disk_test_key(40, 1300); + { + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + clock = 1; TEST_ASSERT(cache.save(0, a)); + clock = 2; TEST_ASSERT(cache.save(0, b)); + clock = 3; TEST_ASSERT(cache.lookup(a, 1)); // a is now newer than b + } + { + // last_used came from the file: b (2) is the LRU, not a (created 1). + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 2); + cache.learn_layout(0); + clock = 4; TEST_ASSERT(cache.save(0, c)); + TEST_ASSERT(cache.total_bytes() <= cfg.budget_bytes); + TEST_ASSERT(cache.entry_count() == 2); + TEST_ASSERT(!cache.lookup(b, 1)); + clock = 5; TEST_ASSERT(cache.lookup(a, 1)); + clock = 6; TEST_ASSERT(cache.save(0, d)); // evicts c (4) + TEST_ASSERT(cache.total_bytes() <= cfg.budget_bytes); + TEST_ASSERT(!cache.lookup(c, 1)); + TEST_ASSERT(cache.lookup(a, 1)); + TEST_ASSERT(cache.lookup(d, 1)); + } + { + // A smaller budget is enforced at startup, before any request. + DiskCacheConfig small = cfg; small.budget_bytes = (size_t)(one + one / 2); + DiskPrefixCache cache(small, backend); + cache.set_clock([&]() { return clock; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(cache.total_bytes() <= small.budget_bytes); + } + { + // A snapshot larger than the whole budget is never written. + DiskCacheConfig tiny = cfg; tiny.budget_bytes = (size_t)(one / 2); + DiskPrefixCache t(tiny, backend); + TEST_ASSERT(t.init()); + t.learn_layout(0); + TEST_ASSERT(!t.disabled()); + TEST_ASSERT(!t.save(0, disk_test_key(40, 7000))); + TEST_ASSERT(t.total_bytes() <= tiny.budget_bytes); + } + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_restart_restores_and_rejects_other_identity) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_restart").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + std::array q4{}; q4[0] = 4; + std::array q8{}; q8[0] = 8; + const auto key = disk_test_key(40); + { + DiskPrefixCache cache(cfg, backend); + cache.set_identity_salt(q4); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, key, "replay-text")); + } + { + // Same identity after a restart: the index is rebuilt from disk and + // the first lookup hits before any live snapshot exists. + DiskPrefixCache cache(cfg, backend); + cache.set_identity_salt(q4); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 1); + std::vector aux; + cache.for_each_aux([&](const std::string & blob) { aux.push_back(blob); }); + TEST_ASSERT(aux.size() == 1 && aux[0] == "replay-text"); + TEST_ASSERT(cache.longest_prefix(key).len == 40); + backend.adopted_cur_pos_ = 0; + TEST_ASSERT(cache.lookup(key, 1)); + TEST_ASSERT(backend.adopted_cur_pos_ == MockBackendWithLayout::kMaxPos); + } + { + // Another KV type/model identity: never restored, still indexed so + // the budget and the TTL cover it. + DiskPrefixCache cache(cfg, backend); + cache.set_identity_salt(q8); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 1); + size_t other_aux = 0; + cache.for_each_aux([&](const std::string &) { ++other_aux; }); + TEST_ASSERT(other_aux == 0); + TEST_ASSERT(cache.longest_prefix(key).len == 0); + TEST_ASSERT(!cache.lookup(key, 1)); + cache.learn_layout(0); + TEST_ASSERT(cache.longest_prefix(key).len == 0); + TEST_ASSERT(!cache.lookup(key, 1)); + TEST_ASSERT(dkv_files(dir).size() == 1); + } + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_init_deletes_corrupt_files) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_corrupt").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + const auto good = disk_test_key(40, 1); + const auto cut = disk_test_key(40, 300); + const auto old = disk_test_key(40, 600); + std::string layout_dir; + { + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, good)); + TEST_ASSERT(cache.save(0, cut)); + TEST_ASSERT(cache.save(0, old)); + } + auto files = dkv_files(dir); + TEST_ASSERT(files.size() == 3); + layout_dir = fs::path(files[0]).parent_path().string(); + auto path_of = [&](const std::vector & k) { + const PrefixHash h = hash_prefix(k.data(), (int)k.size()); + static const char * hx = "0123456789abcdef"; + std::string name; + for (uint8_t byte : h) { name += hx[byte >> 4]; name += hx[byte & 15]; } + return layout_dir + "/" + name + ".dkv"; + }; + // Truncated (crash or full disk mid-write). + fs::resize_file(path_of(cut), fs::file_size(path_of(cut)) - 7); + // Older format version. + if (FILE * f = std::fopen(path_of(old).c_str(), "r+b")) { + const uint32_t v3 = 3; + std::fseek(f, 4, SEEK_SET); + std::fwrite(&v3, 4, 1, f); + std::fclose(f); + } + // Garbage and a temp file left by an interrupted write. + { FILE * f = std::fopen((layout_dir + "/" + std::string(32, 'a') + ".dkv").c_str(), "wb"); + std::fputs("not a snapshot", f); std::fclose(f); } + { FILE * f = std::fopen((layout_dir + "/" + std::string(32, 'b') + ".dkv.tmp").c_str(), "wb"); + std::fputs("partial", f); std::fclose(f); } + { + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 1); + files = dkv_files(dir); + TEST_ASSERT(files.size() == 1 && files[0] == path_of(good)); + TEST_ASSERT(cache.lookup(good, 1)); + } + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_verifies_stored_tokens) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_tokens").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + const auto key = disk_test_key(40); + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, key)); + cache.drain(); + const auto files = dkv_files(dir); + TEST_ASSERT(files.size() == 1); + // The stored key sits right before the (empty) aux blob, the payload + // and the 8-byte checksum. + uint64_t payload = 0; + FILE * f = std::fopen(files[0].c_str(), "r+b"); + TEST_ASSERT(f != nullptr); + std::fseek(f, 52, SEEK_SET); + TEST_ASSERT(std::fread(&payload, 8, 1, f) == 1); + const long tokens_at = (long)(fs::file_size(files[0]) - payload - 8 - 4 - 4 * key.size()); + const int32_t other = 999999; + std::fseek(f, tokens_at + 4 * 10, SEEK_SET); + std::fwrite(&other, 4, 1, f); + std::fclose(f); + // Same hash in the index, different stored tokens: rejected and deleted. + TEST_ASSERT(!cache.lookup(key, 1)); + TEST_ASSERT(dkv_files(dir).empty()); + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_longest_prefix_and_prune_ancestors) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_prefix_probe").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + const auto turn3 = disk_test_key(60); + const std::vector turn1(turn3.begin(), turn3.begin() + 40); + const std::vector turn2(turn3.begin(), turn3.begin() + 48); + TEST_ASSERT(cache.save(0, turn1)); + TEST_ASSERT(cache.save(0, turn2)); + + // Deepest restorable prefix, without reading any file. + auto probe = cache.longest_prefix(turn3); + TEST_ASSERT(probe.len == 48 && probe.cur_pos == MockBackendWithLayout::kMaxPos); + auto diverged = turn3; diverged[45] = -7; + TEST_ASSERT(cache.longest_prefix(diverged).len == 40); + TEST_ASSERT(cache.longest_prefix(disk_test_key(60, 5000)).len == 0); + TEST_ASSERT(cache.longest_prefix(turn1).len == 40); + + // A deeper state of the session supersedes turn2, but the shared head + // (<= keep_max) stays for other sessions. + TEST_ASSERT(cache.save(0, turn3)); + TEST_ASSERT(cache.prune_ancestors(turn3, 40) == 1); + TEST_ASSERT(cache.entry_count() == 2); + TEST_ASSERT(cache.longest_prefix(diverged).len == 40); + // Without a head at or below keep_max, the shallowest ancestor stays + // in its place (a long system/tools head). + TEST_ASSERT(cache.prune_ancestors(turn3, 0) == 0); + TEST_ASSERT(cache.entry_count() == 2); + TEST_ASSERT(cache.longest_prefix(diverged).len == 40); + TEST_ASSERT(cache.longest_prefix(turn3).len == 60); + const auto turn4 = disk_test_key(70); + TEST_ASSERT(cache.save(0, turn4)); + TEST_ASSERT(cache.prune_ancestors(turn4, 0) == 1); // turn3; turn1 stays + TEST_ASSERT(cache.longest_prefix(diverged).len == 40); + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_cache_write_failure_is_not_fatal) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_write_fail").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, disk_test_key(40))); + cache.drain(); + const std::string layout_dir = fs::path(dkv_files(dir)[0]).parent_path().string(); + // Unwritable directory (stands in for a full disk): the write, in the + // background, fails cleanly, leaves no temp file and drops the entry. + fs::permissions(layout_dir, fs::perms::owner_read | fs::perms::owner_exec); + const bool queued = cache.save(0, disk_test_key(40, 800)); + cache.drain(); + fs::permissions(layout_dir, fs::perms::owner_all); + TEST_ASSERT(queued); + if (::geteuid() != 0) { + TEST_ASSERT(cache.stats().write_failures == 1); + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(dkv_files(dir).size() == 1); + TEST_ASSERT(!cache.lookup(disk_test_key(40, 800), 1)); + } + rm_rf(dir); +} + +// ─── SSD session cache: eviction writes, edit backups, shutdown, paged ── + +namespace { + +// Snapshot payloads in real host ggml buffers, one per slot, so the server +// can detach them for the disk writer and adopt them back. A payload of +// `pos` positions holds one F32 tensor [4, pos] filled with `pos`. +struct DiskPayloadBackend : MockBackend { + struct Payload { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + int pos = 0; + }; + std::map slots; + std::map displaced; + std::vector freed; + std::vector adopted; + int live_position = 0; // what snapshot_save() captures + + static Payload make(int pos) { + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * 4 + 1024; + ip.no_alloc = true; + Payload p; + p.ctx = ggml_init(ip); + ggml_tensor * t = ggml_new_tensor_2d(p.ctx, GGML_TYPE_F32, 4, pos); + ggml_set_name(t, "snap_k_0"); + p.buf = ggml_backend_alloc_ctx_tensors_from_buft(p.ctx, ggml_backend_cpu_buffer_type()); + std::vector fill((size_t)ggml_nelements(t), (float)pos); + ggml_backend_tensor_set(t, fill.data(), 0, ggml_nbytes(t)); + p.pos = pos; + return p; + } + static void release(Payload & p) { + if (p.buf) ggml_backend_buffer_free(p.buf); + if (p.ctx) ggml_free(p.ctx); + p = Payload{}; + } + void put(int slot, int pos) { + drop(slot); + slots[slot] = make(pos); + } + void drop(int slot) { + auto it = slots.find(slot); + if (it == slots.end()) return; + release(it->second); + slots.erase(it); + } + ~DiskPayloadBackend() override { + for (auto & [slot, p] : slots) release(p); + for (auto & [slot, p] : displaced) release(p); + } + void snapshot_free(int slot) override { + if (slots.count(slot)) freed.push_back(slot); + drop(slot); + } + bool snapshot_save(int slot) override { + if (live_position <= 0) return false; + put(slot, live_position); + return true; + } + bool snapshot_used(int slot) const override { return slots.count(slot) != 0; } + int snapshot_cur_pos(int slot) const override { + auto it = slots.find(slot); + return it == slots.end() ? 0 : it->second.pos; + } + SnapshotRef snapshot_ref(int slot) const override { + SnapshotRef ref; + auto it = slots.find(slot); + if (it == slots.end()) return ref; + ref.ctx = it->second.ctx; + ref.buf = it->second.buf; + ref.cur_pos = it->second.pos; + ref.last_tok = 7; + return ref; + } + bool snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, + int cur_pos, int32_t) override { + drop(slot); + slots[slot] = {ctx, buf, cur_pos}; + adopted.push_back(slot); + return true; + } + OwnedSnapshot snapshot_take_displaced(int slot) override { + auto it = displaced.find(slot); + if (it == displaced.end()) return {}; + OwnedSnapshot out(it->second.ctx, it->second.buf, it->second.pos, -1); + displaced.erase(it); + return out; + } + size_t snapshot_bytes_estimate(int tokens) const override { + return (size_t)tokens * 100 + 1000; + } +}; + +ServerConfig disk_session_config(const std::string & dir) { + ServerConfig config; + config.prefix_cache_cap = 8; + config.prefix_cache_max_bytes = 0; + config.disk_cache_dir = dir; + config.disk_cache_min_tokens = 1; + config.disk_cache_cold_max_tokens = 2; // the 2-token head is "short" + return config; +} + +// Commit `prompt` as an inline snapshot at its end, the way a request does. +int commit_turn(HttpServer & server, DiskPayloadBackend & backend, + const std::vector & prompt, int restore_slot = -1) { + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + auto reservation = cache.reserve_inline_snap( + prompt, restore_slot >= 0 ? backend.snapshot_cur_pos(restore_slot) : 0, + /*prefer_tools_boundary=*/false, /*forced_cut=*/(int)prompt.size(), + restore_slot, [&](int tokens) { return backend.snapshot_bytes_estimate(tokens); }); + if (!reservation.active()) return -1; + const int slot = reservation.slot(); + SchedulerTestHarness::evict_inline_slot(server, slot); + backend.put(slot, (int)prompt.size()); + SchedulerTestHarness::finalize_inline_snapshot( + server, prompt, std::move(reservation), slot, (int)prompt.size(), + restore_slot); + return slot; +} + +} // namespace + +// Persist-on-evict policy: every conversation's newest snapshot and the +// shared heads are written; a previous turn kept only as an edit backup is +// not, while its conversation's deeper snapshot is resident. +TEST_CASE(ServerUnitFixture, test_prefix_cache_persist_on_evict_skips_edit_backups) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + PrefixCache cache(8, tokenizer, 0); + const std::vector head = {1, 100}; + const std::vector turn1 = {1, 100, 3, 101}; + const std::vector turn2 = {1, 100, 3, 101, 4, 102}; + const std::vector turn3 = {1, 100, 3, 101, 4, 102, 3, 103}; + cache.confirm_inline_snap(0, (int)head.size(), head); + cache.confirm_inline_snap(1, (int)turn1.size(), turn1); + cache.confirm_inline_snap(2, (int)turn2.size(), turn2); + TEST_ASSERT(cache.persist_on_evict(1)); // nothing marked yet + cache.confirm_inline_snap(3, (int)turn3.size(), turn3); + const auto pruned = cache.prune_superseded_ancestors(3); + TEST_ASSERT(pruned == std::vector({1})); + TEST_ASSERT(cache.persist_on_evict(0)); // head + TEST_ASSERT(!cache.persist_on_evict(2)); // edit backup of turn3 + TEST_ASSERT(cache.persist_on_evict(3)); // newest + TEST_ASSERT(!cache.persist_on_evict(5)); // no entry + const auto all = cache.entries_for_shutdown(); + TEST_ASSERT(all.size() == 3 && all[0].slot == 3); // most recent first + for (const auto & e : all) TEST_ASSERT(e.persist == (e.slot != 2)); + + // A protected pin (the system/tools head of a tool request) is always + // written, even with a deeper snapshot after it. + PrefixCache pinned(8, tokenizer, 0); + pinned.confirm_inline_snap(0, (int)turn1.size(), turn1, /*protect=*/true); + pinned.confirm_inline_snap(1, (int)turn2.size(), turn2); + pinned.confirm_inline_snap(2, (int)turn3.size(), turn3); + pinned.prune_superseded_ancestors(2); + TEST_ASSERT(pinned.persist_on_evict(0)); + TEST_ASSERT(pinned.persist_on_evict(2)); + + // Once the deeper snapshot is gone, the kept turn is the newest again. + cache.invalidate_inline_snap(3); + TEST_ASSERT(cache.persist_on_evict(2)); + + // Commit victims are logged (with their flag) only when recording. + cache.confirm_inline_snap(3, (int)turn3.size(), turn3); + cache.prune_superseded_ancestors(3); + TEST_ASSERT(!cache.persist_on_evict(2)); + PrefixCache small(2, tokenizer, 0); + small.set_record_evictions(true); + small.confirm_inline_snap(0, (int)turn1.size(), turn1); + small.confirm_inline_snap(1, (int)head.size(), {1, 200}); + const std::vector other = {1, 300, 3, 301}; + auto reservation = small.reserve_inline_snap(other, 0, false, (int)other.size()); + TEST_ASSERT(reservation.active()); + const int victim_slot = reservation.slot(); + TEST_ASSERT(reservation.commit(other)); + const auto evicted = small.take_evicted(); + TEST_ASSERT(evicted.size() == 1 && evicted[0].persist && + evicted[0].slot == victim_slot && + (evicted[0].ids == turn1 || evicted[0].ids == std::vector({1, 200}))); + TEST_ASSERT(small.take_evicted().empty()); + unlink(path.c_str()); +} + +// Single-sequence path: a snapshot the resident budget evicts is written to +// disk in the background and a later request restores it from there. +TEST_CASE(ServerUnitFixture, test_disk_session_ram_eviction_writes_and_restores) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_evict").string(); + rm_rf(dir); + const std::vector a = {1, 100, 3, 101}; + const std::vector b = {1, 200, 3, 201}; + const std::vector c = {1, 300, 3, 301}; + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config = disk_session_config(dir); + config.prefix_cache_max_bytes = 3000; // two 4-token snapshots (1400 B) + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(!disk.disabled()); + const int slot_a = commit_turn(server, backend, a); + TEST_ASSERT(slot_a >= 0 && commit_turn(server, backend, b) >= 0); + TEST_ASSERT(disk.entry_count() == 0); // nothing written per turn + TEST_ASSERT(commit_turn(server, backend, c) >= 0); // evicts a + // c reuses a's slot: a was handed to the disk writer first. + TEST_ASSERT(SchedulerTestHarness::prefix_cache(server).lookup(a).first != slot_a); + TEST_ASSERT(disk.entry_count() == 1); + TEST_ASSERT(disk.longest_prefix(a).len == (int)a.size()); + // The next turn of conversation a restores from disk (served from + // memory if the write is still queued), into the staging slot. + std::vector next = a; + next.insert(next.end(), {4, 102, 3, 103}); + auto prepared = SchedulerTestHarness::prepare_cache(server, next); + TEST_ASSERT(prepared.restore_slot == ModelBackend::kMaxSlots - 1); + TEST_ASSERT(prepared.prefix_len == (int)a.size()); + disk.drain(); + TEST_ASSERT(disk.stats().writes == 1 && disk.stats().pending_files == 0); + } + rm_rf(dir); + unlink(path.c_str()); +} + +// Tool path: an agent turn saved with its generated text (the live state +// after generation) is written when the resident budget evicts it, keyed +// by the tokens the next request renders. +TEST_CASE(ServerUnitFixture, test_disk_session_agent_turn_eviction_writes_live_state) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_agent").string(); + rm_rf(dir); + const std::vector prompt = {1, 100, 3, 101}; + const std::vector generated = {200, 201, 202}; + const std::vector canonical = {1, 100, 3, 101, 200, 201, 202, 7}; + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config = disk_session_config(dir); + config.prefix_cache_max_bytes = 3500; + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + backend.live_position = 6; // the last sampled token is not decoded + TEST_ASSERT(SchedulerTestHarness::save_generated_turn( + server, prompt, generated, canonical)); + TEST_ASSERT(commit_turn(server, backend, {1, 200, 3, 201}) >= 0); + TEST_ASSERT(disk.entry_count() == 0); + TEST_ASSERT(commit_turn(server, backend, {1, 300, 3, 301}) >= 0); + std::vector next = canonical; + next.insert(next.end(), {8, 9}); + const auto probe = disk.longest_prefix(next); + TEST_ASSERT(probe.len == 6 && probe.cur_pos == 6); + disk.drain(); + TEST_ASSERT(disk.stats().writes == 1); + } + rm_rf(dir); + unlink(path.c_str()); +} + +// Graceful shutdown is one last eviction: the head and each conversation's +// newest snapshot are written, the edit backup and pruned turns are not, and +// a restarted server restores the session from disk. +TEST_CASE(ServerUnitFixture, test_disk_session_shutdown_flush_then_restart_restores) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_flush").string(); + rm_rf(dir); + const std::vector head = {1, 100}; + const std::vector turn1 = {1, 100, 3, 101}; + const std::vector turn2 = {1, 100, 3, 101, 4, 102}; + const std::vector turn3 = {1, 100, 3, 101, 4, 102, 3, 103}; + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + HttpServer server(engine, tokenizer, disk_session_config(dir)); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + // The head, then three turns, each restoring the previous one. + int prev = -1; + for (const auto * turn : {&head, &turn1, &turn2, &turn3}) { + if (turn == &head) { + backend.put(0, (int)head.size()); + cache.confirm_inline_snap(0, (int)head.size(), head); + SchedulerTestHarness::set_slot_tokens(server, 0, head); + prev = 0; + continue; + } + prev = commit_turn(server, backend, *turn, prev); + TEST_ASSERT(prev >= 0); + } + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(disk.entry_count() == 0); + SchedulerTestHarness::flush_sessions(server); + TEST_ASSERT(disk.stats().pending_files == 0); + TEST_ASSERT(disk.entry_count() == 2); // head + turn3 + TEST_ASSERT(disk.longest_prefix(turn3).len == (int)turn3.size()); + TEST_ASSERT(disk.longest_prefix(turn2).len == (int)head.size()); + TEST_ASSERT(dkv_files(dir).size() == 2); + } + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + HttpServer server(engine, tokenizer, disk_session_config(dir)); + std::vector next = turn3; + next.insert(next.end(), {4, 104}); + const auto prepared = SchedulerTestHarness::prepare_cache(server, next); + TEST_ASSERT(prepared.restore_slot == ModelBackend::kMaxSlots - 1); + TEST_ASSERT(prepared.prefix_len == (int)turn3.size()); + TEST_ASSERT(backend.adopted == std::vector({ModelBackend::kMaxSlots - 1})); + // Another conversation on the same head restores the head. + const std::vector other = {1, 100, 3, 555, 4, 556}; + TEST_ASSERT(SchedulerTestHarness::prepare_cache(server, other).prefix_len == + (int)head.size()); + } + rm_rf(dir); + unlink(path.c_str()); +} + +// A session evicted to disk still matches the shared head in RAM; its +// next request restores the deeper disk state, not the head. When that read +// fails, the RAM head is still used. +TEST_CASE(ServerUnitFixture, test_disk_session_deeper_than_ram_hit_restores_from_disk) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_deeper").string(); + rm_rf(dir); + const std::vector head = {1, 100}; + const std::vector turn1 = {1, 100, 3, 101}; + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + HttpServer server(engine, tokenizer, disk_session_config(dir)); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + backend.put(0, (int)head.size()); + cache.confirm_inline_snap(0, (int)head.size(), head); + SchedulerTestHarness::set_slot_tokens(server, 0, head); + const int slot = commit_turn(server, backend, turn1, 0); + TEST_ASSERT(slot > 0); + SchedulerTestHarness::evict_inline_slot(server, slot); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + disk.drain(); + TEST_ASSERT(disk.stats().writes == 1); + std::vector next = turn1; + next.insert(next.end(), {4, 104}); + const auto prepared = SchedulerTestHarness::prepare_cache(server, next); + TEST_ASSERT(prepared.restore_slot == ModelBackend::kMaxSlots - 1); + TEST_ASSERT(prepared.prefix_len == (int)turn1.size()); + backend.snapshot_free(ModelBackend::kMaxSlots - 1); + // The file is gone: the index still names it, the read fails and + // the request restores the resident head. + for (const auto & f : dkv_files(dir)) fs::remove(f); + const auto fallback = SchedulerTestHarness::prepare_cache(server, next); + TEST_ASSERT(fallback.restore_slot == 0); + TEST_ASSERT(fallback.prefix_len == (int)head.size()); + } + rm_rf(dir); + unlink(path.c_str()); +} + +// Corruption fails closed: a flipped payload byte or a cut file is a logged +// miss, the file is deleted, and nothing is adopted. +TEST_CASE(ServerUnitFixture, test_disk_cache_checksum_rejects_corrupt_payload) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_checksum").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + cache.learn_layout(0); + const auto a = disk_test_key(40, 1); + const auto b = disk_test_key(40, 500); + TEST_ASSERT(cache.save(0, a)); + TEST_ASSERT(cache.save(0, b)); + cache.drain(); + const auto path_of = [&](const std::vector & k) { + for (const auto & f : dkv_files(dir)) { + const PrefixHash h = hash_prefix(k.data(), (int)k.size()); + static const char * hx = "0123456789abcdef"; + std::string name; + for (uint8_t byte : h) { name += hx[byte >> 4]; name += hx[byte & 15]; } + if (fs::path(f).filename().string() == name + ".dkv") return f; + } + return std::string(); + }; + const std::string fa = path_of(a), fb = path_of(b); + TEST_ASSERT(!fa.empty() && !fb.empty()); + // One bit in the middle of the tensor payload. + { + FILE * f = std::fopen(fa.c_str(), "r+b"); + TEST_ASSERT(f != nullptr); + const long at = (long)fs::file_size(fa) - 8 - 100; + std::fseek(f, at, SEEK_SET); + const int byte = std::fgetc(f); + std::fseek(f, at, SEEK_SET); + std::fputc(byte ^ 0x10, f); + std::fclose(f); + } + const size_t adopted_before = backend.adopted_.size(); + TEST_ASSERT(!cache.lookup(a, 1)); + TEST_ASSERT(!fs::exists(fa)); + TEST_ASSERT(cache.stats().load_failures == 1); + // Cut short after the index was built (no startup scan in between). + fs::resize_file(fb, fs::file_size(fb) - 3); + TEST_ASSERT(!cache.lookup(b, 1)); + TEST_ASSERT(!fs::exists(fb)); + TEST_ASSERT(backend.adopted_.size() == adopted_before); + TEST_ASSERT(cache.entry_count() == 0); + rm_rf(dir); +} + +// Tensor metadata is untrusted too: reject it before GGML can assert, +// divide by zero on a removed type, or allocate from overflowing dimensions. +TEST_CASE(ServerUnitFixture, test_disk_cache_rejects_corrupt_tensor_metadata) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_tensor_metadata").string(); + rm_rf(dir); + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + const auto key = disk_test_key(40); + for (int corruption = 0; corruption < 4; ++corruption) { + for (bool restart : {false, true}) { + rm_rf(dir); + auto cache = std::make_unique(cfg, backend); + TEST_ASSERT(cache->init()); + TEST_ASSERT(cache->save(0, key)); + cache->drain(); + const auto files = dkv_files(dir); + TEST_ASSERT(files.size() == 1); + FILE * f = std::fopen(files[0].c_str(), "r+b"); + TEST_ASSERT(f != nullptr); + std::fseek(f, DISK_CACHE_HEADER_SIZE, SEEK_SET); + uint16_t name_len = 0; + TEST_ASSERT(std::fread(&name_len, sizeof(name_len), 1, f) == 1); + const long type_at = DISK_CACHE_HEADER_SIZE + 2 + name_len; + std::fseek(f, type_at, SEEK_SET); + const uint32_t type = corruption == 0 ? 4 : GGML_TYPE_Q4_0; + TEST_ASSERT(std::fwrite(&type, sizeof(type), 1, f) == 1); + if (corruption != 0) { + const int64_t ne = corruption == 1 ? 1 : + (corruption == 2 ? INT64_MAX - 31 : 32); + TEST_ASSERT(std::fwrite(&ne, sizeof(ne), 1, f) == 1); + if (corruption == 3) { + const int64_t huge = INT64_MAX; + TEST_ASSERT(std::fwrite(&huge, sizeof(huge), 1, f) == 1); + } + } + std::fclose(f); + if (restart) { + cache.reset(); + cache = std::make_unique(cfg, backend); + TEST_ASSERT(cache->init()); + } + TEST_ASSERT(!cache->lookup(key, 1)); + TEST_ASSERT(!fs::exists(files[0])); + TEST_ASSERT(cache->entry_count() == 0); + TEST_ASSERT(backend.adopted_.empty()); + } + } + rm_rf(dir); +} + +// Writes are asynchronous: a queued snapshot counts against the budget and +// is served from memory at once; the payload handed over is owned by the +// cache; the budget is clamped to the free space. +TEST_CASE(ServerUnitFixture, test_disk_cache_save_async_queue_and_free_space) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_async").string(); + rm_rf(dir); + const auto key = disk_test_key(40, 1); + const auto ref = backend.snapshot_ref(0); + { + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.save_async( + OwnedSnapshot::clone_to_host(ref.ctx, ref.cur_pos, ref.last_tok), key, "aux")); + TEST_ASSERT(cache.total_bytes() > 0); + TEST_ASSERT(cache.lookup(key, 1)); // from memory or from the file + TEST_ASSERT(backend.adopted_cur_pos_ == ref.cur_pos); + const auto drained = cache.drain(); + TEST_ASSERT(drained.cancelled == 0); + TEST_ASSERT(cache.stats().writes == 1 && cache.stats().pending_bytes == 0); + TEST_ASSERT(dkv_files(dir).size() == 1); + // An empty payload or one past its key is refused. + TEST_ASSERT(!cache.save_async(OwnedSnapshot{}, disk_test_key(40, 99))); + TEST_ASSERT(!cache.save_async( + OwnedSnapshot::clone_to_host(ref.ctx, ref.cur_pos, ref.last_tok), + disk_test_key(ref.cur_pos - 1, 99))); + } + { + // 1 MB free and no reserve: the configured 128 GB budget shrinks to + // the bytes already used plus the free space. + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + cfg.free_space_reserve_bytes = 0; + DiskPrefixCache cache(cfg, backend); + cache.set_free_space_probe([] { return (uint64_t)1 << 20; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.budget_bytes() == cache.total_bytes() + (1u << 20)); + cache.set_free_space_probe([] { return (uint64_t)0; }); + cache.learn_layout(0); + TEST_ASSERT(cache.save(0, disk_test_key(40, 700))); // evicts the old one + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(!cache.lookup(key, 1)); + } + // A second server on the same directory leaves it alone. + { + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache first(cfg, backend); + TEST_ASSERT(first.init() && !first.disabled()); + DiskPrefixCache second(cfg, backend); + TEST_ASSERT(second.init()); + TEST_ASSERT(second.disabled()); + } + rm_rf(dir); +} + +#if !defined(_WIN32) +namespace { + +// A concurrent engine whose checkpoints live in the backend's snapshot +// slots (as Qwen's paged checkpoints do): a capture replaces the slot's +// payload and leaves the old one displaced; restoring checkpoint 64 reads +// the disk staging slot. +class DiskCheckpointEngine final : public SeqEngine { +public: + explicit DiskCheckpointEngine(DiskPayloadBackend & backend) + : backend_(backend), slots_(2) {} + int slot_count() const override { return (int)slots_.size(); } + int max_context() const override { return 64; } + bool supports_prefix_store() const override { return true; } + bool prefix_store_in_backend_snapshots() const override { return true; } + size_t estimate_prefix_store_bytes(int) const override { return 256; } + bool token_is_eos(int32_t token) const override { return token == 2; } + StepPlanLimits step_plan_limits(int) const override { + return {2, 64, 128, 64}; + } + AdmitResult admit(uint64_t id, const std::vector & prompt, + const SamplerCfg & sampler) override { + return admit_with_prefix(id, prompt, sampler, PrefixStorePlan{}); + } + AdmitResult admit_with_prefix(uint64_t, const std::vector & prompt, + const SamplerCfg &, + const PrefixStorePlan & plan) override { + AdmitResult result; + int chosen = -1; + for (int i = 0; i < (int)slots_.size(); ++i) { + if (!slots_[(size_t)i].active) { chosen = i; break; } + } + if (prompt.empty() || chosen < 0) { + result.status = AdmitResult::Status::busy; + return result; + } + Slot & slot = slots_[(size_t)chosen]; + slot = Slot{true, true, {}}; + result.status = AdmitResult::Status::admitted; + result.slot = chosen; + if (plan.restore.valid()) { + result.prefix_store.restore_attempted = true; + const int index = (int)plan.restore.id - 1; + if (backend_.snapshot_cur_pos(index) == plan.restore.tokens) { + result.prefix_store.restored = plan.restore; + restores.push_back(plan.restore); + } else { + result.prefix_store.invalidated = plan.restore; + } + } + if (!result.prefix_store.invalidated.valid() && plan.capture.valid()) { + slot.capture = plan.capture; + result.prefix_store.capture = plan.capture; + } + return result; + } + StepResult step(const StepPlan & plan) override { + StepResult result; + for (const StepInput & input : plan.decode) { + result.decode.push_back({input.slot, 2, false, {}}); + } + for (const PrefillSlice & slice : plan.prefills) { + Slot & slot = slots_[(size_t)slice.slot]; + slot.prefilling = false; + PrefillOutput output; + output.slot = slice.slot; + output.status = PrefillOutput::Status::completed; + output.token = 2; + if (slot.capture.valid()) { + const int index = (int)slot.capture.checkpoint.id - 1; + auto old = backend_.slots.find(index); + if (old != backend_.slots.end()) { + DiskPayloadBackend::release(backend_.displaced[index]); + backend_.displaced[index] = old->second; + backend_.slots.erase(old); + } + backend_.slots[index] = + DiskPayloadBackend::make(slot.capture.checkpoint.tokens); + output.prefix_store.status = PrefixStoreEvent::Status::saved; + output.prefix_store.ticket = slot.capture; + output.prefix_store.bytes = 256; + slot.capture = {}; + } + result.prefills.push_back(std::move(output)); + } + return result; + } + void retire(int slot) override { slots_[(size_t)slot] = Slot{}; } + void discard_prefix_store(PrefixStoreRef checkpoint) override { + backend_.drop((int)checkpoint.id - 1); + } + std::vector restores; + +private: + struct Slot { + bool active = false; + bool prefilling = false; + PrefixCaptureTicket capture; + }; + DiskPayloadBackend & backend_; + std::vector slots_; +}; + +struct DiskCheckpointBackend : DiskPayloadBackend { + DiskCheckpointEngine engine{*this}; + SeqEngine * seq_engine() override { return &engine; } +}; + +} // namespace + +// Concurrent (paged) serving: a checkpoint evicted by another capture is +// written to disk, the conversation's next request is restored from it +// through the staging checkpoint, and shutdown writes the rest. +TEST_CASE(ServerUnitFixture, test_disk_session_concurrent_eviction_and_restore) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_paged").string(); + rm_rf(dir); + + auto backend_owner = std::make_unique(); + DiskCheckpointBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config = disk_session_config(dir); + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; // a third conversation evicts the first + config.concurrent_prefix_cache_max_bytes = 4096; + config.concurrent_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(!disk.disabled()); + + std::vector> jobs; + std::vector fds; + const auto make_job = [&](std::vector prompt, const char * id) { + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + fds.push_back(sockets[0]); + fds.push_back(sockets[1]); + auto job = std::make_unique(); + job->fd = sockets[0]; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.prompt_tokens = std::move(prompt); + job->req.max_output = 1; + job->req.stream = false; + job->req.model = "scheduler-test"; + job->req.response_id = id; + jobs.push_back(std::move(job)); + return jobs.back().get(); + }; + const auto run_one = [&](std::vector prompt, const char * id) { + ServerJob * job = make_job(std::move(prompt), id); + SchedulerTestHarness::enqueue(server, job); + std::unique_lock lock(job->mu); + return job->cv.wait_for(lock, std::chrono::seconds(5), [&] { return job->done; }); + }; + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + const std::vector a = {1, 100, 3, 101, 4}; + const std::vector b = {1, 200, 3, 201, 4}; + const std::vector c = {1, 300, 3, 301, 4}; + const bool abc = run_one(a, "a") && run_one(b, "b") && run_one(c, "c"); + disk.drain(); + const auto written = disk.stats().writes; + const auto probe_a = disk.longest_prefix({1, 100, 3, 101, 4, 9}); + // Conversation a again: RAM lost it, disk has it. + const bool again = run_one({1, 100, 3, 101, 4, 9, 2, 3, 102, 4}, "a2"); + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + + TEST_ASSERT(abc && again); + TEST_ASSERT(written == 1); + TEST_ASSERT(probe_a.len == 3 && probe_a.cur_pos == 3); + bool restored_from_disk = false; + for (const auto & r : backend.engine.restores) { + restored_from_disk |= r == PrefixStoreRef{(uint64_t)ModelBackend::kMaxSlots, 3}; + } + TEST_ASSERT(restored_from_disk); + TEST_ASSERT(!backend.snapshot_used(ModelBackend::kMaxSlots - 1)); // staging freed + TEST_ASSERT(backend.displaced.empty()); // every displaced payload claimed + + // Shutdown writes what is still in RAM. + SchedulerTestHarness::flush_sessions(server); + TEST_ASSERT(disk.stats().pending_files == 0); + TEST_ASSERT(disk.longest_prefix({1, 300, 3, 301, 4, 9}).len == 3); + rm_rf(dir); + unlink(path.c_str()); +} +#endif + TEST_CASE(ServerUnitFixture, test_backend_ipc_rejects_file_work_dir) { const std::string file_path = test_tmp_path( "luce_test_backend_ipc_work_dir_file").string(); From 5f744cd83ce41d545041d84932f3eb5c5b8e5a2b Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sun, 11 Oct 2026 07:57:50 +0000 Subject: [PATCH 2/3] docs(server): document the disk session cache, its defaults and eviction-only writes Co-Authored-By: Claude Opus 5.5 --- docs/specs/props-endpoint.md | 13 +++- server/README.md | 5 +- server/docs/PREFIX_CACHE.md | 137 +++++++++++++++++++++++++++++++++-- 3 files changed, 147 insertions(+), 8 deletions(-) diff --git a/docs/specs/props-endpoint.md b/docs/specs/props-endpoint.md index afdad6bde7..160cce1766 100644 --- a/docs/specs/props-endpoint.md +++ b/docs/specs/props-endpoint.md @@ -237,7 +237,10 @@ banner. "disk_bytes": 0, "enabled": false, "in_use": 0, - "lifetime_hits": 0 + "lifetime_hits": 0, + "disk_enabled": false, + "disk_budget_mb": 131072, + "disk_ttl_hours": 72 } ``` @@ -245,6 +248,14 @@ The full-prompt cache (disk-backed). When `enabled = false`, `capacity`, `in_use`, `lifetime_hits`, and `disk_bytes` are all zero and ignored. +`disk_enabled` reports whether the separate SSD session cache initialized +successfully. It can be false even when RAM caching is enabled, for example +when the cache directory cannot be created or is locked by another server. +`disk_budget_mb` is the configured disk cap in MiB (1,048,576 bytes per unit), +and `disk_ttl_hours` is the idle expiry in hours (`0` means never expire). +These configured values remain present when `disk_enabled` is false; they +do not report live disk usage or the cap after the free-space clamp. + Counters use atomic loads (`std::memory_order_relaxed`); the snapshot is tear-free per field but the set of fields is **not internally consistent** — e.g. `in_use` and `lifetime_hits` may diff --git a/server/README.md b/server/README.md index b495ac0d60..4c1185fa37 100644 --- a/server/README.md +++ b/server/README.md @@ -328,9 +328,10 @@ See the [current six-expert Strix Halo profile](https://www.lucebox.com/blog/dee | `--admission-coalesce-ms ` | `20` | Idle-to-busy batching window from 0 through 1000 ms. | | `--kv-pool-tokens ` | auto | Shared physical K/V capacity for concurrent serving, rounded to the backend page size. | | `--decode-kv-offload-mb ` | `auto` | Automatically size the RAM budget for active KV suspension; `N` sets a per-model cap in MiB, `0` disables. Applies to concurrent serving. | -| `--kv-cache-dir ` | none | Enable persistent disk KV cache in this directory. | -| `--kv-cache-budget ` | `4096` | Disk KV-cache size cap (hard cap, LRU eviction). | +| `--kv-cache-dir \|off` | `~/.cache/lucebox/kv-cache` | Disk session cache directory (`$XDG_CACHE_HOME/lucebox/kv-cache` when set); `off` or `--no-kv-cache` disables it. Snapshots are written when they leave RAM and at graceful shutdown; a crash loses what was only in RAM. See [docs/PREFIX_CACHE.md](docs/PREFIX_CACHE.md#disk-session-cache). | +| `--kv-cache-budget ` | `131072` | Disk KV-cache size cap (hard cap, LRU eviction, clamped to the free space). | | `--kv-cache-ttl ` | `72` | Delete disk snapshots unused this long; `0` = never. | +| `--kv-cache-shutdown-flush ` | `120` | Longest a graceful shutdown spends writing RAM snapshots to disk; `-1` = no limit. | | `--kv-cache-min-tokens ` | `512` | Minimum prefix length to persist. | | `--kv-cache-interval ` | `10240` | Continued-checkpoint interval. | | `--kv-cache-cold-max ` | `10240` | Cold-prefix limit for long prompts. | diff --git a/server/docs/PREFIX_CACHE.md b/server/docs/PREFIX_CACHE.md index 26c423e601..09fce815c2 100644 --- a/server/docs/PREFIX_CACHE.md +++ b/server/docs/PREFIX_CACHE.md @@ -90,10 +90,136 @@ when the selected boundary advances beyond the restored prefix. The disk prefix cache is a separate persistence/overflow layer for token-keyed snapshots. Today it is integrated with the inline/effective-prompt path; exact prefill snapshots are kept in RAM unless a dedicated raw-prompt disk path is -added. With the default `full` policy a lookup probes the whole prompt and -then every chat boundary, deepest first, so a restart recovers the inline -snapshots the previous process persisted (probes are index lookups; only a -hit reads a file). +added. With the default `full` policy a lookup probes the whole prompt, every +chat boundary and the deepest stored entry that prefixes the prompt, deepest +first, so a restart recovers the sessions the previous process persisted +(probes are index lookups; only a hit reads a file). See +[Disk Session Cache](#disk-session-cache). + +## Disk Session Cache + +The disk session cache keeps agent and chat sessions on an SSD, so a session +whose RAM snapshots are gone (another session evicted them, the server +restarted, the user switched model) restores its KV and prefills only the new +suffix. It is on by default. + +| Flag | Default | Meaning | +|------|---------|---------| +| `--kv-cache-dir \|off` | `$XDG_CACHE_HOME/lucebox/kv-cache`, else `~/.cache/lucebox/kv-cache` | Directory of the cache; `off` (or `--no-kv-cache`) disables it. | +| `--kv-cache-budget ` | `131072` (128 GB) | Hard cap on the bytes of all snapshot files, clamped to the free space. | +| `--kv-cache-ttl ` | `72` | Delete files unused for this long; `0` = never. | +| `--kv-cache-shutdown-flush ` | `120` | Longest a graceful shutdown spends writing RAM snapshots; `-1` = no limit. | +| `--kv-cache-cold-max ` | `10240` | Longest head kept when a deeper state of its session is written; cold-head checkpoints without RAM slots. | +| `--kv-cache-interval ` | `10240` | Continued checkpoint pacing (fallback without RAM slots). | + +**When something is written.** Only when a snapshot leaves RAM, never after a +turn: + +- a RAM eviction: a new snapshot takes the slot of an older one (its + reservation's victim) or the resident budget evicts one after a commit; +- a graceful shutdown (SIGTERM/SIGINT): one last eviction of everything still + in RAM, most recently used first, for at most `--kv-cache-shutdown-flush` + seconds. The log line `[disk-cache] shutdown flush:` gives the sessions, + files, bytes and seconds, and what the limit dropped. + +A crash, `kill -9` or power loss loses what was only in RAM; the states written +before it still restore. An active session therefore writes nothing while it +stays resident: a 90-120K-token Qwen3.8-27B session with q4_0 KV is 1.5-2.2 GB +per snapshot. + +**What is written.** Since #835 every request snapshots the end of its last +message, a new turn's snapshot keeps the previous turn's one while the budget +holds it (an edit of the last message resumes there), and the turns before it +are pruned. On eviction (and at shutdown) the disk gets: + +- each conversation's newest snapshot: the inline snapshot of the last + request, or the agent-turn state with the generated reply (tool path, + `--agent-turn-cache`), on the single-sequence and the concurrent (paged) + paths alike; +- the system/tools heads (protected pins and the shallowest entries). + +It does not get the kept previous turn (an *edit backup*) while a deeper +snapshot of its conversation is resident (`PrefixCache::persist_on_evict`): +it only serves an edit of the last message, which the deeper snapshot's file +cannot, and writing it would double the bytes per conversation. Once the +deeper snapshot is gone, the kept turn is that conversation's newest and is +written like any other. Pruned (superseded) turns are not written either. +When a session's state is written, its shallower files on disk are deleted, +except a head at or below `--kv-cache-cold-max` tokens or, without one, the +shallowest (a longer head). + +Compressed prompts (PFlash, FlowKV) are not raw history and never get a disk +key; a request whose disk scope is not `full` gives its snapshots no key +either. + +**How it is written.** Eviction hands the snapshot's buffer to the disk writer +without a copy (Qwen keeps snapshots in host memory; other backends copy to +host first) and the request continues at once. The snapshot is indexed +immediately, counts against the budget, and a lookup serves it from memory +until its file exists. A writer thread writes `.tmp`, `fdatasync`s it, +renames it into place and syncs the directory. Queued writes are bounded +(4 GiB of payload); past that the evicting request waits, logged as +`[disk-cache] waited`. + +**Restore.** With the default `full` policy a lookup probes the whole prompt, +every chat boundary and the deepest stored key that prefixes the prompt +(index lookups; only a hit reads a file). A disk state deeper than the best +RAM hit wins over it: a session evicted to disk usually still matches the +shared head in RAM. If that read fails, the RAM hit is used. The single-sequence path reads the +file on the request that restores it (the only request running). The +concurrent scheduler reads it on a reader thread while the request waits at +the head of the queue and the live streams keep decoding, then restores it +into the new sequence's pages through a staging checkpoint. A hit's state is +bit-identical to the RAM snapshot it was written from (raw bytes in the +cache's own type; q4_0 stays q4_0). + +A session's prompt is rendered with the exact text the model generated for +its earlier tool calls (ToolMemory, which a new process does not have). Each +file stores that replay text and startup loads it back into ToolMemory. + +**Budget.** Every layout directory under the cache directory is indexed, also +files of another model or KV type, so the cap covers the whole directory. The +budget never exceeds the cache's bytes plus the free space minus a reserve +(the larger of 1 GiB and 5% of the disk, at most 16 GiB). Before a write the +least recently used files are deleted until the new file fits. A hit rewrites +the file's `last_used` field in place, so LRU order survives restarts. + +**TTL.** A file unused for longer than `--kv-cache-ttl` hours is deleted at +startup, before each write, and when a lookup meets it. + +**Safety.** + +- Format v5 ends with an XXH64 checksum of the whole file (the header with + `last_used` zeroed). A hit verifies it while reading: a corrupt or partial + file is a logged miss, it is deleted and the prompt is prefilled. Startup + deletes leftover `.tmp` files, files of older formats and truncated files. +- A file is restored only when its identity matches: model path, size, mtime + and a content fingerprint (SHA-1 of the first and last MiB), KV cache types, + K rotation, `--max-ctx`, chat template and draft model, plus the snapshot's + tensor names, types and shapes. Qwen checks every adopted tensor against the + live cache (the paged restore against the page table too). +- Each file stores the exact key token ids; a hit compares them with the + prompt. +- One server per directory: startup takes an exclusive lock on + `/.lock`; a second server on the same directory runs without a disk + cache and logs it. + +**Fallbacks without RAM slots.** With `--prefix-cache-slots 0` there is +nothing to evict, so the older per-request checkpoints remain: the cold head +(a prompt longer than `--kv-cache-cold-max` first prefills to the last chat +boundary at or below it) and continued checkpoints. The `fixed`/`auto` disk +scopes (`--disk-prefix-cache`) also keep their scoped-prefix write. All of +them go through the same asynchronous writer. + +**Not covered.** Single-sequence `--paged-attention` keeps no snapshots, and +the concurrent `qwen4exp` engine keeps its checkpoints outside the backend's +snapshot slots: both run without a disk cache (logged at startup). + +**Follow-up: cache-aware admission.** Compression admission that counts the +cached prefix as warm (on the PFlash branches, not on `main`) consults RAM +only. `DiskPrefixCache::longest_prefix(prompt)` returns the deepest restorable +key length and KV position without reading a file, so admission can take the +max of the RAM and disk prefixes. ## Snapshot Memory Management @@ -230,7 +356,7 @@ behind, keyed by the prompt plus the generated tokens the next request renders identically (see [Agent Turn Cache](API.md#agent-turn-cache)). Qwen defers that snapshot: while the next request continues the conversation it is never copied to system RAM, and it is copied out only before other work would overwrite the -live state, or when the disk cache persists it at shutdown. Restoring the +live state, or when it is evicted to disk. Restoring the snapshot the live state already holds copies nothing back either. Coding agents grow one conversation turn by turn, and each turn's snapshot is a @@ -284,6 +410,7 @@ cache staging, preventing disk loads from overwriting a live in-memory entry. | File | Role | |------|------| | `server/prefix_cache.{h,cpp}` | LRU logic, boundary detection, hashing | +| `server/disk_prefix_cache.{h,cpp}` | Disk snapshots: file format, index, budget, TTL | | `common/snapshot_backend.h` | Platform-aware snapshot backend selection | | `common/model_backend.h` | `snapshot_save/free/used/cur_pos` interface | | `qwen35/qwen35_target_graph.cpp` | `snapshot_target_cache()`, `restore_target_cache()` | From d6774ae4602d6987f5caa95709427336e0e9eaa6 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 12 Oct 2026 09:03:07 +0000 Subject: [PATCH 3/3] fix(server): harden disk cache lifecycle, scope and capture budgets Allow large captures to evict several eligible entries transactionally. Respect request scopes, refresh duplicate-save TTL metadata, disable failed cache initialization, and reap asynchronous readers during serving. Use an exclusive Windows directory lock and native cache defaults. Cover dense Qwen disk adoption, failed restores, shared capture reservations, and compressed-state exclusion with model-free regressions. Co-Authored-By: Claude Opus 5.5 --- docs/specs/props-endpoint.md | 3 +- server/docs/PREFIX_CACHE.md | 34 +- server/src/server/disk_prefix_cache.cpp | 62 ++- server/src/server/disk_prefix_cache.h | 11 +- server/src/server/http_server.cpp | 28 +- server/src/server/http_server.h | 5 + server/src/server/prefix_cache.cpp | 178 +++++- server/src/server/prefix_cache.h | 24 +- server/src/server/scheduler.cpp | 77 ++- server/src/server/server_main.cpp | 12 +- server/test/unit/test_server_unit.cpp | 710 +++++++++++++++++++++++- 11 files changed, 1058 insertions(+), 86 deletions(-) diff --git a/docs/specs/props-endpoint.md b/docs/specs/props-endpoint.md index 160cce1766..c71a8c5f5c 100644 --- a/docs/specs/props-endpoint.md +++ b/docs/specs/props-endpoint.md @@ -398,7 +398,8 @@ non-strictly-consistent semantics as `full_cache` (§4.7). `capacity = 0` means the cache is disabled. `max_resident_bytes` and `resident_bytes` cover committed copied checkpoints used by concurrent paged serving; a maximum of `0` means unlimited. `budget_skips` counts captures -declined because no single eligible LRU entry could make enough room. +declined because the eligible LRU entries, alone or together, could not make +enough room (or the checkpoint is larger than the whole budget). The capture and restore counters expose synchronous time spent copying checkpoint state on the scheduler thread. `*_attempts` include successful and diff --git a/server/docs/PREFIX_CACHE.md b/server/docs/PREFIX_CACHE.md index 09fce815c2..198f98c8d9 100644 --- a/server/docs/PREFIX_CACHE.md +++ b/server/docs/PREFIX_CACHE.md @@ -116,7 +116,9 @@ suffix. It is on by default. turn: - a RAM eviction: a new snapshot takes the slot of an older one (its - reservation's victim) or the resident budget evicts one after a commit; + reservation's victim), the resident budget evicts several to make room for + a large one (see *Several evictions for one capture*), or it evicts one + after a commit; - a graceful shutdown (SIGTERM/SIGINT): one last eviction of everything still in RAM, most recently used first, for at most `--kv-cache-shutdown-flush` seconds. The log line `[disk-cache] shutdown flush:` gives the sessions, @@ -366,7 +368,9 @@ shared system/tools head) and protected tool pins. A conversation therefore holds about two snapshots instead of one per turn. Chains that are abandoned, for example after a client compacts its context, are bounded by `--prefix-cache-max-mib`: when a new snapshot would not fit, the cache replaces -the oldest leaf first and skips the capture if nothing can make room. Snapshots +the oldest leaf first. When no single entry makes enough room, it evicts several, +least recently used first (see below), and skips the capture only if nothing +can make room or the snapshot is larger than the whole limit. Snapshots the new one supersedes (including the one it restored from) count as freed, because they are pruned as soon as it is committed; a failed capture keeps them. After each commit the cache also evicts, oldest leaf first, until the committed @@ -379,11 +383,27 @@ request runs, process memory can exceed the limit by up to one snapshot. Concurrent paged serving measures the exact backend allocation required for each checkpoint before copying it. The cache keeps committed checkpoints under `--concurrent-prefix-cache-max-mib`: when necessary it replaces one eligible least-recently-used -entry, and if no single eligible entry can make enough room it skips the new -checkpoint without disturbing the committed cache. The configured limit covers -resident committed checkpoint buffers. During an atomic replacement, the new -buffer and the selected victim can coexist briefly, so transient process memory -can exceed the limit by up to one checkpoint. +entry, and when no single eligible entry can make enough room it evicts several +(below). The configured limit covers resident committed checkpoint buffers. +During an atomic replacement, the new buffer and the entries it replaces can +coexist briefly, so transient process memory can exceed the limit by up to one +checkpoint. + +**Several evictions for one capture.** Both paths use the same rule. A capture +that fits the limit at all, but that no single entry can make room for, plans +several evictions, least recently used first: edit backups (not written to +disk while their deeper snapshot is resident), then leaves. The plan never +takes a protected pin, a head that other resident entries extend, the +capture's own ancestors or restore source, or an entry another capture in +flight writes, evicts or counts as freed. It is all or nothing: it exists only +if the capture fits once those entries are gone, otherwise the capture is +skipped and nothing is touched. The entries leave RAM only when the capture +commits (the single-sequence path writes and frees them just before the +commit), so a capture that fails or is cancelled evicts none of them; each one +goes to the disk cache like any other eviction. Without this a long session +that arrives after many short ones (subagent threads) was never captured, so +it never reached the disk either. A checkpoint larger than the whole limit is +still skipped (`budget_skips`). By default the limit is sized when the scheduler starts, from the batch engine's own estimate of one checkpoint at `--max-ctx`: every slot keeps its diff --git a/server/src/server/disk_prefix_cache.cpp b/server/src/server/disk_prefix_cache.cpp index 03174b753d..ec9969eb7e 100644 --- a/server/src/server/disk_prefix_cache.cpp +++ b/server/src/server/disk_prefix_cache.cpp @@ -22,6 +22,10 @@ #include #if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include #include #else #include @@ -526,7 +530,9 @@ DiskPrefixCache::~DiskPrefixCache() { writer_cv_.notify_all(); writer_.join(); } -#if !defined(_WIN32) +#if defined(_WIN32) + if (lock_handle_) CloseHandle(lock_handle_); +#else if (lock_fd_ >= 0) ::close(lock_fd_); #endif } @@ -609,24 +615,37 @@ static bool validate_snapshot_file(const std::string & path, bool DiskPrefixCache::init() { if (disabled()) return true; if (!mkdir_p(config_.cache_dir)) { - std::fprintf(stderr, "[disk-cache] failed to create dir: %s\n", + std::fprintf(stderr, "[disk-cache] disabled: failed to create dir: %s\n", config_.cache_dir.c_str()); + // No writer was started. Leaving the cache enabled would queue + // payloads forever and block serving once backpressure is reached. + config_.cache_dir.clear(); return false; } -#if !defined(_WIN32) const std::string lock_path = config_.cache_dir + "/.lock"; +#if defined(_WIN32) + // A share-mode-zero handle excludes every other server before scanning + // temporaries or enforcing the budget, just like flock on POSIX. + lock_handle_ = CreateFileW(fs::path(lock_path).c_str(), GENERIC_READ | GENERIC_WRITE, + 0, nullptr, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); + const bool locked = lock_handle_ != INVALID_HANDLE_VALUE; + if (!locked) lock_handle_ = nullptr; +#else lock_fd_ = ::open(lock_path.c_str(), O_RDWR | O_CREAT | O_CLOEXEC, 0644); - if (lock_fd_ < 0 || ::flock(lock_fd_, LOCK_EX | LOCK_NB) != 0) { + const bool locked = lock_fd_ >= 0 && ::flock(lock_fd_, LOCK_EX | LOCK_NB) == 0; +#endif + if (!locked) { std::fprintf(stderr, - "[disk-cache] disabled: %s is in use by another server " + "[disk-cache] disabled: cannot lock %s (in use or not writable) " "(pass --kv-cache-dir to give this one its own)\n", config_.cache_dir.c_str()); +#if !defined(_WIN32) if (lock_fd_ >= 0) ::close(lock_fd_); lock_fd_ = -1; +#endif config_.cache_dir.clear(); return true; } -#endif std::lock_guard lock(mu_); scan_all(); // Budget and TTL hold from the first moment, before any request. @@ -808,6 +827,16 @@ void DiskPrefixCache::release_pending_locked(DiskEntry & e) { e.pending.reset(); } +void DiskPrefixCache::touch_locked(DiskEntry & e) { + e.last_used = now(); + if (e.pending) return; // the writer publishes the latest timestamp + // Not fsynced: a crash may age an entry, but clean restarts retain uses. + if (FILE * f = std::fopen(e.path.c_str(), "r+b")) { + if (std::fseek(f, DISK_CACHE_LAST_USED_OFFSET, SEEK_SET) == 0) write_u64(f, e.last_used); + std::fclose(f); + } +} + void DiskPrefixCache::remove_entry(size_t idx, const char * why) { DiskEntry & e = entries_[idx]; const bool queued = e.pending != nullptr; @@ -970,16 +999,7 @@ OwnedSnapshot DiskPrefixCache::load(const std::vector & prompt_ids) { if (idx >= 0) { // Persist last_used in the file so LRU order and the TTL survive a // restart. Not fsynced: losing it in a crash only ages the entry. - DiskEntry & entry = entries_[(size_t)idx]; - entry.last_used = now(); - if (!entry.pending) { - if (FILE * f = std::fopen(entry.path.c_str(), "r+b")) { - if (std::fseek(f, DISK_CACHE_LAST_USED_OFFSET, SEEK_SET) == 0) { - write_u64(f, entry.last_used); - } - std::fclose(f); - } - } + touch_locked(entries_[(size_t)idx]); } return out; } @@ -1042,7 +1062,8 @@ void DiskPrefixCache::for_each_aux( } DiskPrefixCache::PrefixProbe DiskPrefixCache::longest_prefix( - const std::vector & prompt_ids) { + const std::vector & prompt_ids, + const std::vector * allowed_lengths) { PrefixProbe out; if (disabled()) return out; std::lock_guard lock(mu_); @@ -1053,6 +1074,9 @@ DiskPrefixCache::PrefixProbe DiskPrefixCache::longest_prefix( if (e.layout_id != layout_id_ || expired(e, t)) continue; if (e.token_count > prompt_ids.size()) continue; if ((int)e.token_count < config_.min_tokens) continue; + if (allowed_lengths && + std::find(allowed_lengths->begin(), allowed_lengths->end(), + (int)e.token_count) == allowed_lengths->end()) continue; lengths.push_back(e.token_count); } std::sort(lengths.begin(), lengths.end(), std::greater()); @@ -1153,7 +1177,7 @@ bool DiskPrefixCache::save_async(OwnedSnapshot snap, // Already on disk (or queued): count it as a use. const int existing = find_entry(hash); if (existing >= 0) { - entries_[(size_t)existing].last_used = now(); + touch_locked(entries_[(size_t)existing]); return true; } // Make room first, so the budget holds while the file is being written. @@ -1212,7 +1236,7 @@ bool DiskPrefixCache::save(int slot, const std::vector & prompt_ids, std::lock_guard lock(mu_); const int existing = find_entry(hash); if (existing >= 0) { - entries_[(size_t)existing].last_used = now(); + touch_locked(entries_[(size_t)existing]); return true; } } diff --git a/server/src/server/disk_prefix_cache.h b/server/src/server/disk_prefix_cache.h index 998ef3da46..09a8af9339 100644 --- a/server/src/server/disk_prefix_cache.h +++ b/server/src/server/disk_prefix_cache.h @@ -181,7 +181,7 @@ class DiskPrefixCache { // budget to the free space, apply the TTL and the budget, pick the most // recently used layout written with this identity salt (enables // first-request disk hits), and start the writer thread. Returns false - // on fatal error. + // when initialization fails; the disk cache disables itself in that case. bool init(); // Look up a prompt on disk. On hit, loads the snapshot into `slot` @@ -232,7 +232,9 @@ class DiskPrefixCache { // key is a prefix of `prompt_ids`. len = key tokens, cur_pos = KV rows // the snapshot holds (<= len). {0, 0} on miss. struct PrefixProbe { int len = 0; int cur_pos = 0; }; - PrefixProbe longest_prefix(const std::vector & prompt_ids); + // With allowed_lengths, only those exact scoped boundaries qualify. + PrefixProbe longest_prefix(const std::vector & prompt_ids, + const std::vector * allowed_lengths = nullptr); // Delete live-layout entries whose key is a strict prefix of `tokens` // and longer than `keep_max` tokens: a deeper state of the same session @@ -367,7 +369,11 @@ class DiskPrefixCache { // Exclusive lock on /.lock: one process (and one model) per // directory, since startup deletes temporaries and the budget evicts // files of every layout in it. +#if defined(_WIN32) + void * lock_handle_ = nullptr; +#else int lock_fd_ = -1; +#endif // Helpers (the _locked ones need mu_). uint64_t now() const; @@ -385,6 +391,7 @@ class DiskPrefixCache { void make_room_locked(uint64_t incoming); void writer_loop(); void release_pending_locked(DiskEntry & e); + void touch_locked(DiskEntry & e); bool write_file(const std::string & path, const PendingWrite & job); OwnedSnapshot read_file(const std::string & path, diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index b75ad6c157..9648507f59 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1544,6 +1544,9 @@ HttpServer::HttpServer(luce::engine::LuceEngine & engine, disk_cache_.set_identity_salt(compute_disk_cache_salt(config)); } disk_cache_.init(); + // Report the effective state after initialization failure or a directory + // lock conflict; serving continues with the ordinary RAM cache. + if (disk_cache_.disabled()) config_.disk_cache_dir.clear(); restore_replay_blobs(); status_html_path_ = resolve_status_html(); @@ -4281,6 +4284,7 @@ void HttpServer::finalize_generation_cache( // decision sees the requested cut as covered. Keying at the // saved position re-targets the same cut on every request // and the restore point never moves past it. + evict_extra_victims(cache.snap_reservation); cache.snap_reservation.commit( effective_prompt, backend_.snapshot_bytes_estimate(saved_position)); @@ -4384,6 +4388,14 @@ void HttpServer::evict_inline_slot(int slot) { forget_inline_slot_metadata(slot); } +void HttpServer::evict_extra_victims( + const PrefixCache::InlineReservation & reservation) { + for (const int slot : reservation.extra_victim_slots()) { + evict_inline_slot(slot); + backend_.snapshot_free(slot); + } +} + void HttpServer::note_session_slot(int slot) { if (slot >= 0) request_session_slots_.push_back(slot); } @@ -4646,6 +4658,7 @@ bool HttpServer::save_generated_turn( reservation.abort(); return false; } + evict_extra_victims(reservation); reservation.commit_at( canonical, saved_pos, backend_.snapshot_bytes_estimate(saved_pos)); std::vector persist; @@ -4655,9 +4668,11 @@ bool HttpServer::save_generated_turn( forget_inline_slot_metadata(evicted[i]); backend_.snapshot_free(evicted[i]); } - slot_tokens_[slot] = std::vector( - canonical.begin(), canonical.begin() + saved_pos); - note_session_slot(slot); + if (cache.disk_policy.mode == DiskPrefixCacheMode::Full) { + slot_tokens_[slot] = std::vector( + canonical.begin(), canonical.begin() + saved_pos); + note_session_slot(slot); + } agent_turn_cache_slots_.insert(slot); std::fprintf(stderr, "[agent-turn-cache] saved generated turn slot=%d prefix=%d " @@ -4798,13 +4813,16 @@ void HttpServer::remember_agent_turn( const int saved_pos = replay_result.ok() && backend_.snapshot_used(slot) ? backend_.snapshot_cur_pos(slot) : 0; if (saved_pos > source_pos && saved_pos <= canonical_end) { + evict_extra_victims(reservation); reservation.commit_at( canonical_tokens, saved_pos, backend_.snapshot_bytes_estimate(saved_pos)); trim_snapshots_after_commit(slot); canonical_tokens.resize((size_t) saved_pos); - slot_tokens_[slot] = std::move(canonical_tokens); - note_session_slot(slot); + if (cache.disk_policy.mode == DiskPrefixCacheMode::Full) { + slot_tokens_[slot] = std::move(canonical_tokens); + note_session_slot(slot); + } agent_turn_cache_slots_.insert(slot); std::fprintf(stderr, "[agent-turn-cache] saved slot=%d prefix=%d source=%d:%d " diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 6d2f852302..19b6b18f25 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -576,6 +576,11 @@ class HttpServer { // Free a RAM snapshot slot for reuse, writing it to disk first when // PrefixCache::persist_on_evict() says so (see persist_evicted_slot). void evict_inline_slot(int slot); + // Just before `reservation` commits: write (when persist_on_evict says + // so) and free the further entries its resident-budget plan evicts + // (PrefixCache::InlineReservation::extra_victim_slots). The commit then + // drops their metadata. A capture that fails never gets here. + void evict_extra_victims(const PrefixCache::InlineReservation & reservation); // Disk persistence (docs/PREFIX_CACHE.md): a snapshot is written only // when it leaves RAM, never after every turn. Hands the payload of // `slot` to the disk writer when `persist` is set and the slot has a disk diff --git a/server/src/server/prefix_cache.cpp b/server/src/server/prefix_cache.cpp index e40136394b..874c032ce9 100644 --- a/server/src/server/prefix_cache.cpp +++ b/server/src/server/prefix_cache.cpp @@ -541,6 +541,20 @@ bool PrefixCache::InlineReservation::active() const { return cache_ && cache_->inline_reservation_active(id_); } +std::vector PrefixCache::InlineReservation::extra_victim_slots() const { + std::vector out; + const ActiveReservation * active = + cache_ ? cache_->find_active_reservation(id_) : nullptr; + if (!active) return out; + for (size_t i = 0; i < active->extra_victims.size(); ++i) { + const int idx = cache_->find_entry(active->extra_victims[i]); + if (idx >= 0 && cache_->entries_[(size_t)idx].slot == active->extra_slots[i]) { + out.push_back(active->extra_slots[i]); + } + } + return out; +} + bool PrefixCache::InlineReservation::commit( const std::vector & prompt_ids, size_t resident_bytes, bool protect) { @@ -598,11 +612,16 @@ void PrefixCache::InlineReservation::take(InlineReservation && other) { } bool PrefixCache::inline_reservation_active(uint64_t id) const { - if (id == 0) return false; + return find_active_reservation(id) != nullptr; +} + +const PrefixCache::ActiveReservation * +PrefixCache::find_active_reservation(uint64_t id) const { + if (id == 0) return nullptr; for (const auto & active : active_reservations_) { - if (active.id == id) return true; + if (active.id == id) return &active; } - return false; + return nullptr; } void PrefixCache::release_inline_reservation(uint64_t id) { @@ -617,6 +636,18 @@ void PrefixCache::release_inline_reservation(uint64_t id) { bool PrefixCache::slot_reserved(int slot) const { for (const auto & active : active_reservations_) { if (active.slot == slot) return true; + for (const int extra : active.extra_slots) { + if (extra == slot) return true; + } + } + return false; +} + +bool PrefixCache::slot_credited(int slot) const { + for (const auto & active : active_reservations_) { + for (const int credited : active.credited_slots) { + if (credited == slot) return true; + } } return false; } @@ -783,6 +814,9 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( } size_t estimated_bytes = 0; + std::vector extra_victims; + std::vector extra_slots; + std::vector credited_slots; if (max_resident_bytes_ > 0) { estimated_bytes = estimate_bytes ? estimate_bytes(target_cut) : 0; // When the caller prunes after every commit, the entries this @@ -841,26 +875,97 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( // A turn kept behind a conversation's newest snapshot first (see // select_inline_evict_victim), then leaves, then the rest; // unprotected entries before protected ones. - int victim = -1; - for (int pass = 0; pass < 5 && victim < 0; ++pass) { - const bool require_middle = pass == 0; - const bool require_leaf = pass == 1 || pass == 2; - const bool allow_protected = pass == 2 || pass == 4; - for (int i = 0; i < (int)entries_.size(); ++i) { - if (restore_source_slot >= 0 && - entries_[(size_t)i].slot == restore_source_slot) - continue; - if (slot_reserved(entries_[(size_t)i].slot)) continue; - if (!fits(i)) continue; - if (require_middle && !is_middle(i)) continue; - if (require_leaf && !is_leaf(i)) continue; - if (!allow_protected && entries_[(size_t)i].protect) - continue; - victim = i; - break; + const auto pick_single = [&](int first_pass, int last_pass) { + for (int pass = first_pass; pass <= last_pass; ++pass) { + const bool require_middle = pass == 0; + const bool require_leaf = pass == 1 || pass == 2; + const bool allow_protected = pass == 2 || pass == 4; + for (int i = 0; i < (int)entries_.size(); ++i) { + if (restore_source_slot >= 0 && + entries_[(size_t)i].slot == restore_source_slot) + continue; + if (slot_reserved(entries_[(size_t)i].slot)) continue; + if (!fits(i)) continue; + if (require_middle && !is_middle(i)) continue; + if (require_leaf && !is_leaf(i)) continue; + if (!allow_protected && entries_[(size_t)i].protect) + continue; + return i; + } } - } - if (victim >= 0) { + return -1; + }; + // No single unprotected turn or leaf makes room: plan several, + // least recently used first, before falling back to a protected + // pin or a head. Each one goes to disk as it leaves RAM, so a + // long session that arrives after many short ones is still + // captured. Upstream skipped this capture instead (one victim + // per reservation) so a capture never disturbs the committed + // cache until it lands; the plan keeps that: it is all or + // nothing, decided before anything is evicted, and its entries + // are evicted only when the capture commits. A capture that + // fails or is cancelled evicts none of them, and one larger + // than the whole budget plans none. What a plan never takes: + // protected pins; heads (entries that other resident entries + // extend, unless kept only as an edit backup); this capture's + // own ancestors and restore source, which it builds on; and + // entries a capture in flight writes, evicts or already counts + // as freed, so two captures never count the same bytes. + const auto plan_several = [&]() { + std::vector plan; + if (estimated_bytes == 0 || + estimated_bytes > max_resident_bytes_) return plan; + size_t freed = reclaimed_bytes; + if (planned_victim >= 0 && !reclaimed[(size_t)planned_victim]) { + freed += entries_[(size_t)planned_victim].resident_bytes; + } + const size_t held = resident_bytes_ + pending_bytes; + const auto room = [&]() { + const size_t after_free = freed <= held ? held - freed : 0; + return after_free <= max_resident_bytes_ && + estimated_bytes <= max_resident_bytes_ - after_free; + }; + // Edit backups first (they are not written to disk while + // their deeper snapshot is resident), then leaves. + for (int pass = 0; pass < 2; ++pass) { + for (int i = 0; i < (int)entries_.size(); ++i) { + const auto & entry = entries_[(size_t)i]; + if (i == planned_victim || reclaimed[(size_t)i] || + entry.protect) continue; + if (restore_source_slot >= 0 && + entry.slot == restore_source_slot) continue; + if (slot_reserved(entry.slot) || + slot_credited(entry.slot)) continue; + if (is_strict_prefix(entry.ids, prompt_ids.data(), + (size_t)target_cut)) continue; + const bool leaf = is_leaf(i); + if (pass == 0 ? !(entry.edit_backup && !leaf) : !leaf) + continue; + plan.push_back(i); + freed += entry.resident_bytes; + if (room()) return plan; + } + } + plan.clear(); + return plan; + }; + int victim = pick_single(0, 1); + std::vector several; + if (victim < 0) several = plan_several(); + if (victim < 0 && several.empty()) victim = pick_single(2, 4); + if (!several.empty()) { + size_t bytes = 0; + for (const int i : several) { + extra_victims.push_back(entries_[(size_t)i].hash); + extra_slots.push_back(entries_[(size_t)i].slot); + bytes += entries_[(size_t)i].resident_bytes; + } + std::fprintf(stderr, + "[pc] resident budget plans %zu evictions (%zu bytes) for " + "capture slot=%d estimate=%zu resident=%zu budget=%zu\n", + several.size(), bytes, slot, estimated_bytes, + resident_bytes_, max_resident_bytes_); + } else if (victim >= 0) { victim_key = entries_[(size_t)victim].hash; has_victim = true; const int redirected_slot = entries_[(size_t)victim].slot; @@ -883,12 +988,19 @@ PrefixCache::InlineReservation PrefixCache::reserve_inline_snap( return {}; } } + for (size_t i = 0; i < reclaimed.size(); ++i) { + if (reclaimed[i]) credited_slots.push_back(entries_[i].slot); + } } uint64_t id = next_inline_reservation_++; if (id == 0) id = next_inline_reservation_++; + if (has_victim) credited_slots.push_back(slot); + credited_slots.insert(credited_slots.end(), extra_slots.begin(), extra_slots.end()); active_reservations_.push_back( - {id, slot, key, victim_key, has_victim, estimated_bytes}); + {id, slot, key, victim_key, has_victim, estimated_bytes, + std::move(extra_victims), std::move(extra_slots), + std::move(credited_slots)}); return InlineReservation( this, id, slot, target_cut, victim_key, has_victim, protect); } @@ -931,6 +1043,26 @@ bool PrefixCache::commit_inline_reservation( reservation.abort(); return false; } + // The further entries its plan evicts, before the slot's own victim: + // an edit backup leaves while its deeper snapshot is still resident, so + // it is not written to disk (persist_on_evict). + if (const ActiveReservation * active = find_active_reservation(reservation.id_)) { + size_t evicted = 0, freed = 0; + for (size_t i = 0; i < active->extra_victims.size(); ++i) { + const int idx = find_entry(active->extra_victims[i]); + if (idx < 0 || entries_[(size_t)idx].slot != active->extra_slots[i]) continue; + ++evicted; + freed += entries_[(size_t)idx].resident_bytes; + record_eviction(idx); + erase_inline_entry(idx); + } + if (evicted > 0) { + std::fprintf(stderr, + "[pc] resident budget evicted %zu snapshot(s), %zu MiB, for " + "capture slot=%d\n", + evicted, freed / (1024 * 1024), reservation.slot_); + } + } if (reservation.has_victim_) { const int victim = find_entry(reservation.victim_); if (victim >= 0) { diff --git a/server/src/server/prefix_cache.h b/server/src/server/prefix_cache.h index e80b1d8d7d..266215e0ce 100644 --- a/server/src/server/prefix_cache.h +++ b/server/src/server/prefix_cache.h @@ -179,6 +179,12 @@ class PrefixCache { bool active() const; int slot() const { return slot_; } int target_cut() const { return target_cut_; } + // Committed entries besides the slot's own victim that this capture + // evicts when it commits (see reserve_inline_snap), in eviction + // order. Only those still resident are listed. A caller that owns + // their payloads (the single-sequence path) writes and frees them + // just before the commit; the commit drops their metadata. + std::vector extra_victim_slots() const; // Commit after the engine saved the payload, cancel before the target // slot was touched, or abort after a failed write invalidated it. @@ -234,6 +240,13 @@ class PrefixCache { // With a `restore_point_spacing` (ModelBackend::restore_point_spacing), a // cut other than the end of the last message moves down to the deepest // spaced restore point, and nothing is reserved when that adds no prefix. + // Resident budget: when no single entry makes room for a capture that + // fits the budget at all, the reservation plans several, least recently + // used first (edit backups, then leaves); see the definition for what is + // never planned. All or nothing: a plan exists only when the capture + // fits once they are gone, else the capture is skipped and nothing is + // touched. They are evicted when it commits, never when it fails or is + // cancelled. InlineReservation reserve_inline_snap( const std::vector & prompt_ids, int restored_prefix_len = 0, @@ -392,6 +405,12 @@ class PrefixCache { PrefixHash victim; bool has_victim; size_t bytes; + // Further entries this capture evicts when it commits, in order. + std::vector extra_victims; + std::vector extra_slots; + // Slots of every entry it counts as freed (its superseded ancestors + // and its victims): no other capture counts them again. + std::vector credited_slots; }; std::vector active_reservations_; uint64_t next_inline_reservation_ = 1; @@ -459,8 +478,11 @@ class PrefixCache { bool record_hit); bool inline_reservation_active(uint64_t id) const; void release_inline_reservation(uint64_t id); - // A capture in flight writes this slot. + // A capture in flight writes this slot or evicts its entry. bool slot_reserved(int slot) const; + // A capture in flight counts this slot's entry as freed. + bool slot_credited(int slot) const; + const ActiveReservation * find_active_reservation(uint64_t id) const; bool commit_inline_reservation(InlineReservation & reservation, const std::vector & prompt_ids, int committed_cut, size_t resident_bytes, diff --git a/server/src/server/scheduler.cpp b/server/src/server/scheduler.cpp index 543f144dd1..e1f509c5db 100644 --- a/server/src/server/scheduler.cpp +++ b/server/src/server/scheduler.cpp @@ -118,7 +118,10 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { "[disk-cache] off for concurrent serving: this engine keeps no " "checkpoints in backend snapshot slots\n"); } - prefix_cache_.set_record_evictions(disk_enabled); + // Recorded with the disk off too: a capture whose resident-budget plan + // evicts further checkpoints (PrefixCache::reserve_inline_snap) frees + // their payloads here. + prefix_cache_.set_record_evictions(true); struct DiskLoad { ServerJob * job = nullptr; std::vector key; @@ -130,16 +133,26 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { std::thread reader; }; std::optional disk_load; - std::vector finished_readers; + std::vector finished_readers; // Staging checkpoint for a disk restore: PrefixCache never uses it. constexpr int kDiskStagingIndex = PrefixCache::MAX_CACHE_SLOTS; auto drop_disk_load = [&]() { if (!disk_load) return; if (disk_load->reader.joinable()) { - finished_readers.push_back(std::move(disk_load->reader)); + finished_readers.push_back(std::move(*disk_load)); } disk_load.reset(); }; + auto reap_readers = [&]() { + for (auto it = finished_readers.begin(); it != finished_readers.end();) { + if (it->result->done.load(std::memory_order_acquire)) { + it->reader.join(); + it = finished_readers.erase(it); + } else { + ++it; + } + } + }; // Degenerate-run guard shared with do_ar_decode: explicit env override, // else 32 when the min-tokens floor is active, else off. @@ -507,8 +520,19 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } // A deeper state on disk (the index answers without a read): // load it off this thread, then restore from the staging slot. - const auto probe = disk_enabled - ? disk_cache_.longest_prefix(req.prompt_tokens) + std::vector scoped_lengths; + const auto mode = req.disk_cache_policy.mode; + if (mode == DiskPrefixCacheMode::Fixed) { + scoped_lengths.push_back(disk_prefix_cache_fixed_boundary( + req.disk_cache_policy, (int)req.prompt_tokens.size(), + config_.disk_cache_min_tokens)); + } else if (mode == DiskPrefixCacheMode::Auto) { + scoped_lengths = find_all_boundaries( + req.prompt_tokens, prefix_cache_.chat_markers()); + } + const auto probe = disk_enabled && mode != DiskPrefixCacheMode::Off + ? disk_cache_.longest_prefix(req.prompt_tokens, + mode == DiskPrefixCacheMode::Full ? nullptr : &scoped_lengths) : DiskPrefixCache::PrefixProbe{}; if (disk_load && disk_load->job != job) drop_disk_load(); if (probe.cur_pos > prefix_plan.restore.tokens && @@ -549,6 +573,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { prefix_plan.restore.valid() ? "restoring from RAM" : "prefilling"); } + } else if (disk_load && disk_load->job == job) { + drop_disk_load(); } if (prefix_plan.restore.valid() && disk_cache_digest_enabled()) { @@ -835,6 +861,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { std::chrono::steady_clock::time_point decode_only_until{}; while (true) { + reap_readers(); // Phase 1 — Admission: deferred job first (FIFO), then the queue. // Blocking dequeue only when idle; between decode steps only a poll. // Engines reject atomically when their current slot, staging, or @@ -859,7 +886,8 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { // scratch memory back and park in a blocking dequeue. publish_live_count(); backend_.release_scratch(); - job = dequeue(); + job = finished_readers.empty() ? dequeue() + : dequeue_for(std::chrono::milliseconds(5)); woke_from_idle = job != nullptr; } else { // A fresh idle transition may spend its bounded batching @@ -875,6 +903,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } if (!job) break; // queue empty: go decode what is already live if (job->client_disconnected.load(std::memory_order_acquire)) { + if (disk_load && disk_load->job == job) drop_disk_load(); finish_job(job); continue; } @@ -1182,8 +1211,36 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { OwnedSnapshot displaced = disk_enabled ? backend_.snapshot_take_displaced(captured) : OwnedSnapshot{}; for (auto & evicted : prefix_cache_.take_evicted()) { - if (evicted.slot != captured || displaced.empty()) continue; - if (evicted.persist) { + if (evicted.slot != captured) { + // A further checkpoint the capture's resident-budget + // plan evicted as it committed: its payload is still + // in its own slot. Written like the displaced one, + // then freed. + bool queued = false; + if (disk_enabled && evicted.persist && slot_tokens_.count(evicted.slot)) { + const auto aux = slot_aux_.find(evicted.slot); + const std::string blob = + aux == slot_aux_.end() ? std::string() : aux->second; + OwnedSnapshot snap = backend_.snapshot_detach(evicted.slot); + const int pos = snap.empty() ? 0 : snap.cur_pos(); + queued = !snap.empty() && + disk_cache_.save_async(std::move(snap), evicted.ids, blob); + if (queued) { + disk_cache_.prune_ancestors( + evicted.ids, config_.disk_cache_cold_max_tokens); + } + std::fprintf(stderr, + "[disk-cache] ram eviction checkpoint=%d tokens=%zu " + "pos=%d -> %s\n", evicted.slot + 1, evicted.ids.size(), + pos, queued ? "write queued" : "not written"); + } + engine.discard_prefix_store( + {(uint64_t)evicted.slot + 1, (int)evicted.ids.size()}); + forget_inline_slot_metadata(evicted.slot); + continue; + } + if (displaced.empty()) continue; + if (evicted.persist && slot_tokens_.count(captured)) { const auto aux = slot_aux_.find(captured); const int pos = displaced.cur_pos(); const bool queued = disk_cache_.save_async( @@ -1203,7 +1260,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { if (resolution == Resolution::saved) { // The slot holds the new checkpoint now. forget_inline_slot_metadata(captured); - if (disk_enabled) { + if (disk_enabled && s.job->req.disk_cache_policy.mode == DiskPrefixCacheMode::Full) { const int tokens = out.prefix_store.ticket.checkpoint.tokens; const auto & prompt = s.job->req.prompt_tokens; slot_tokens_[captured] = std::vector( @@ -1301,7 +1358,7 @@ void HttpServer::scheduler_loop(SeqEngine & engine) { } drop_disk_load(); - for (auto & reader : finished_readers) reader.join(); + for (auto & reader : finished_readers) reader.reader.join(); finished_readers.clear(); // Shutdown: unblock every parked client thread. diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 8a335e29de..69e1d86448 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -314,6 +314,7 @@ static void print_usage(const char * prog) { " --kv-cache-dir |off Directory for the on-disk KV session cache\n" " (default: $XDG_CACHE_HOME/lucebox/kv-cache, else\n" " ~/.cache/lucebox/kv-cache; off disables it)\n" + " On Windows: %LOCALAPPDATA%/lucebox/kv-cache\n" " --no-kv-cache Same as --kv-cache-dir off\n" " --kv-cache-budget Max disk usage in MB, hard cap with LRU eviction,\n" " clamped to the free space (default: 131072 = 128 GB)\n" @@ -424,9 +425,16 @@ struct ModelOptions { }; -// $XDG_CACHE_HOME/lucebox/kv-cache, else ~/.cache/lucebox/kv-cache; empty -// (disk cache off) when neither variable is set. +// Native cache directory, or empty when no usable home/cache path is set. static std::string default_kv_cache_dir() { +#if defined(_WIN32) + for (const char * variable : {"LOCALAPPDATA", "USERPROFILE"}) { + const char * value = std::getenv(variable); + if (value && *value && std::filesystem::path(value).is_absolute()) { + return (std::filesystem::path(value) / "lucebox" / "kv-cache").string(); + } + } +#endif const char * xdg = std::getenv("XDG_CACHE_HOME"); if (xdg && xdg[0] == '/') return std::string(xdg) + "/lucebox/kv-cache"; const char * home = std::getenv("HOME"); diff --git a/server/test/unit/test_server_unit.cpp b/server/test/unit/test_server_unit.cpp index 39e890330f..8bc71e6560 100644 --- a/server/test/unit/test_server_unit.cpp +++ b/server/test/unit/test_server_unit.cpp @@ -137,11 +137,13 @@ struct SchedulerTestHarness { static void finalize_inline_snapshot( HttpServer & server, const std::vector & prompt, PrefixCache::InlineReservation reservation, - int slot, int requested_cut, int restore_slot = -1) { + int slot, int requested_cut, int restore_slot = -1, + bool compressed = false) { ParsedRequest req; req.prompt_tokens = prompt; HttpServer::PreparedPrompt prepared; prepared.tokens = prompt; + prepared.compressed = compressed; HttpServer::GenerationCacheState cache; cache.snap_reservation = std::move(reservation); cache.snap_slot = slot; @@ -171,6 +173,9 @@ struct SchedulerTestHarness { static DiskPrefixCache & disk_cache(HttpServer & server) { return server.disk_cache_; } + static json props(const HttpServer & server) { + return build_props_body(server.config_, server.prefix_cache_, server.tool_memory_); + } static void flush_sessions(HttpServer & server) { server.flush_sessions_to_disk("shutdown"); @@ -204,8 +209,10 @@ struct SchedulerTestHarness { HttpServer & server, const std::vector & prompt, const std::vector & generated, const std::vector & canonical, - int own_snapshot_slot = -1) { + int own_snapshot_slot = -1, + DiskPrefixCacheMode disk_scope = DiskPrefixCacheMode::Full) { HttpServer::GenerationCacheState cache; + cache.disk_policy.mode = disk_scope; cache.snap_prepared = own_snapshot_slot >= 0; cache.snap_slot = own_snapshot_slot; return server.save_generated_turn(prompt, generated, canonical, cache); @@ -3786,6 +3793,196 @@ TEST_CASE(ServerUnitFixture, test_prefix_cache_resident_budget_and_stall_stats) unlink(path.c_str()); } +// Resident budget: a capture that fits the budget but that no single entry +// makes room for evicts several, least recently used first, when it commits. +TEST_CASE(ServerUnitFixture, test_prefix_cache_budget_evicts_several_for_large_capture) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const auto bytes = [](size_t n) { + return [n](int) { return n; }; + }; + const std::vector> small = { + {1, 100, 3, 101}, {1, 200, 3, 201}, {1, 300, 3, 301}, {1, 400, 3, 401}}; + const std::vector big = {1, 900, 3, 901}; + const auto fill = [&](PrefixCache & cache, std::vector & slots) { + for (const auto & ids : small) { + auto r = cache.reserve_inline_snap(ids, 0, false, (int)ids.size(), -1, bytes(200)); + TEST_ASSERT(r.active()); + slots.push_back(r.slot()); + TEST_ASSERT(r.commit(ids, 200)); + } + TEST_ASSERT(cache.stats().resident_bytes == 800); + }; + + { + PrefixCache cache(8, tokenizer, /*max_resident_bytes=*/1000); + cache.set_prunes_superseded(true); + cache.set_record_evictions(true); + std::vector slots; + fill(cache, slots); + TEST_ASSERT(cache.take_evicted().empty()); + // 700 bytes: 500 must go, and every entry holds 200. + auto r = cache.reserve_inline_snap(big, 0, false, (int)big.size(), -1, bytes(700)); + TEST_ASSERT(r.active()); + TEST_ASSERT(r.extra_victim_slots() == + std::vector({slots[0], slots[1], slots[2]})); + // Nothing leaves RAM before the capture lands. + TEST_ASSERT(cache.stats().in_use == 4 && cache.stats().resident_bytes == 800); + TEST_ASSERT(cache.take_evicted().empty()); + TEST_ASSERT(r.commit(big, 700)); + const auto evicted = cache.take_evicted(); + TEST_ASSERT(evicted.size() == 3); + for (size_t i = 0; i < evicted.size(); ++i) { + TEST_ASSERT(evicted[i].slot == slots[i]); + TEST_ASSERT(evicted[i].ids == small[i]); + TEST_ASSERT(evicted[i].persist); // each one goes to disk + } + const auto stats = cache.stats(); + TEST_ASSERT(stats.in_use == 2 && stats.resident_bytes == 900); + TEST_ASSERT(stats.budget_skips == 0); + TEST_ASSERT(cache.lookup_candidate(small[3], 4).second == 4); + std::vector next = big; + next.insert(next.end(), {4, 902}); + TEST_ASSERT(cache.lookup_candidate(next, (int)next.size()).second == (int)big.size()); + } + { + // A capture that fails or is cancelled evicts none of them. + PrefixCache cache(8, tokenizer, 1000); + cache.set_record_evictions(true); + std::vector slots; + fill(cache, slots); + auto r = cache.reserve_inline_snap(big, 0, false, (int)big.size(), -1, bytes(700)); + TEST_ASSERT(r.active() && r.extra_victim_slots().size() == 3); + r.cancel(); + TEST_ASSERT(cache.stats().in_use == 4 && cache.stats().resident_bytes == 800); + r = cache.reserve_inline_snap(big, 0, false, (int)big.size(), -1, bytes(700)); + TEST_ASSERT(r.active()); + r.abort(); + TEST_ASSERT(cache.stats().in_use == 4 && cache.stats().resident_bytes == 800); + TEST_ASSERT(cache.take_evicted().empty()); + } + { + // Larger than the whole budget: skipped, nothing evicted. + PrefixCache cache(8, tokenizer, 1000); + cache.set_record_evictions(true); + std::vector slots; + fill(cache, slots); + auto r = cache.reserve_inline_snap(big, 0, false, (int)big.size(), -1, bytes(1001)); + TEST_ASSERT(!r.active()); + TEST_ASSERT(cache.stats().budget_skips == 1); + TEST_ASSERT(cache.stats().in_use == 4 && cache.stats().resident_bytes == 800); + TEST_ASSERT(cache.take_evicted().empty()); + // Exactly the budget: everything else goes. + r = cache.reserve_inline_snap(big, 0, false, (int)big.size(), -1, bytes(1000)); + TEST_ASSERT(r.active() && r.extra_victim_slots() == slots); + TEST_ASSERT(r.commit(big, 1000)); + TEST_ASSERT(cache.take_evicted().size() == 4); + TEST_ASSERT(cache.stats().in_use == 1 && cache.stats().resident_bytes == 1000); + } + unlink(path.c_str()); +} + +// What a several-entry plan never takes: protected pins, heads that other +// entries extend, the capture's own ancestors and restore source, and the +// entries another capture in flight evicts. +TEST_CASE(ServerUnitFixture, test_prefix_cache_budget_plan_keeps_protected_and_in_use) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + PrefixCache cache(16, tokenizer, /*max_resident_bytes=*/1200); + cache.set_record_evictions(true); + const std::vector s1 = {1, 110, 3, 111}; + const std::vector s2 = {1, 120, 3, 121}; + const std::vector pin = {1, 130}; + const std::vector head = {1, 140}; + const std::vector leaf = {1, 140, 3, 141}; + const std::vector ancestor = {1, 150}; + const std::vector source = {1, 150, 3, 151}; + const std::vector t1 = {1, 160, 3, 161}; + const std::vector t2 = {1, 170, 3, 171}; + // Oldest first. Every entry holds 100 bytes. + int slot = 0; + for (const auto * ids : {&s1, &s2, &pin, &head, &leaf, &ancestor, &source, &t1, &t2}) { + cache.confirm_inline_snap(slot++, (int)ids->size(), *ids, + /*protect=*/ids == &pin, /*resident_bytes=*/100); + } + TEST_ASSERT(cache.stats().resident_bytes == 900); + + // In flight: a 500-byte capture that needs two of them gone (s1, s2). + const std::vector other = {1, 180, 3, 181}; + auto first = cache.reserve_inline_snap( + other, 0, false, (int)other.size(), -1, [](int) { return (size_t)500; }); + TEST_ASSERT(first.active()); + TEST_ASSERT(first.extra_victim_slots() == std::vector({0, 1})); + + // The next turn of `source`, restored from it: 300 bytes must go. The + // oldest entries are s1/s2 (the first capture's), then the pin, the head + // and the ancestor; only leaf, t1 and t2 are eligible. + std::vector turn = source; + turn.insert(turn.end(), {4, 152}); + auto second = cache.reserve_inline_snap( + turn, (int)source.size(), false, (int)turn.size(), + /*restore_source_slot=*/6, [](int) { return (size_t)100; }); + TEST_ASSERT(second.active()); + TEST_ASSERT(second.extra_victim_slots() == std::vector({4, 7, 8})); + // A third capture needs 401 bytes gone: only `source` (100, not its + // ancestor here) is left to plan, so it is skipped. No protected pin, + // head or ancestor is taken instead. + auto third = cache.reserve_inline_snap( + {1, 190, 3, 191}, 0, false, 4, -1, [](int) { return (size_t)101; }); + TEST_ASSERT(!third.active()); + + TEST_ASSERT(second.commit(turn, 100)); + auto evicted = cache.take_evicted(); + TEST_ASSERT(evicted.size() == 3 && evicted[0].ids == leaf && + evicted[1].ids == t1 && evicted[2].ids == t2); + TEST_ASSERT(first.commit(other, 500)); + evicted = cache.take_evicted(); + TEST_ASSERT(evicted.size() == 2 && evicted[0].ids == s1 && evicted[1].ids == s2); + for (const auto * kept : {&pin, &head, &ancestor, &source}) { + TEST_ASSERT(cache.lookup_candidate(*kept, (int)kept->size()).second == + (int)kept->size()); + } + TEST_ASSERT(cache.stats().resident_bytes == 1000); + unlink(path.c_str()); +} + +// An edit backup goes before an older leaf, and is not written to disk while +// its conversation's deeper snapshot is resident. +TEST_CASE(ServerUnitFixture, test_prefix_cache_budget_plan_evicts_edit_backups_first) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + PrefixCache cache(16, tokenizer, /*max_resident_bytes=*/1000); + cache.set_prunes_superseded(true); + cache.set_record_evictions(true); + const std::vector y = {1, 200, 3, 201}; + const std::vector z = {1, 300, 3, 301}; + const std::vector head = {1, 100}; + const std::vector turn1 = {1, 100, 3, 101}; + const std::vector turn2 = {1, 100, 3, 101, 4, 102}; + cache.confirm_inline_snap(0, 4, y, false, 100); + cache.confirm_inline_snap(1, 4, z, false, 100); + cache.confirm_inline_snap(2, 2, head, false, 100); + cache.confirm_inline_snap(3, 4, turn1, false, 100); + cache.confirm_inline_snap(4, 6, turn2, false, 100); + TEST_ASSERT(cache.prune_superseded_ancestors(4).empty()); // turn1 kept + TEST_ASSERT(!cache.persist_on_evict(3)); // as an edit backup + + const std::vector big = {1, 400, 3, 401}; + auto r = cache.reserve_inline_snap( + big, 0, false, 4, -1, [](int) { return (size_t)700; }); + TEST_ASSERT(r.active()); + TEST_ASSERT(r.extra_victim_slots() == std::vector({3, 0})); + TEST_ASSERT(r.commit(big, 700)); + const auto evicted = cache.take_evicted(); + TEST_ASSERT(evicted.size() == 2); + TEST_ASSERT(evicted[0].ids == turn1 && !evicted[0].persist); + TEST_ASSERT(evicted[1].ids == y && evicted[1].persist); + unlink(path.c_str()); +} + TEST_CASE(ServerUnitFixture, test_canonical_turn_matches_replay_checkpoint) { TEST_ASSERT(http_detail::canonical_turn_matches_checkpoint( {1, 2, 3}, {1, 2, 9, 4}, 2)); @@ -4862,6 +5059,44 @@ TEST_CASE(ServerUnitFixture, test_budget_counts_superseded_snapshots_as_freed) { unlink(path.c_str()); } +TEST_CASE(ServerUnitFixture, test_parallel_reservations_do_not_double_credit_shared_ancestor) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::vector head = {1, 100}; + const std::vector ancestor = {1, 100, 3, 101}; + const std::vector a = {1, 100, 3, 101, 4, 102}; + const std::vector b = {1, 100, 3, 101, 4, 103}; + for (size_t budget : {300u, 320u}) { + for (bool reverse : {false, true}) { + PrefixCache cache(8, tokenizer, budget); + cache.set_prunes_superseded(true); + cache.confirm_inline_snap(0, 2, head, true, 20); + cache.confirm_inline_snap(1, 4, ancestor, false, 150); + auto first = cache.reserve_inline_snap(a, 2, false, 6, 0, + [](int) { return (size_t)150; }); + TEST_ASSERT(first.active()); + auto second = cache.reserve_inline_snap(b, 2, false, 6, 0, + [](int) { return (size_t)150; }); + // 20-byte head plus two 150-byte captures needs exactly 320; + // shared reclaimed bytes must never make it fit under 300. + TEST_ASSERT(second.active() == (budget == 320)); + const auto commit = [&](PrefixCache::InlineReservation & r, + const std::vector & ids) { + const int slot = r.slot(); + TEST_ASSERT(r.commit(ids, 150)); + cache.prune_superseded_ancestors(slot); + TEST_ASSERT(cache.enforce_resident_budget(slot).empty()); + TEST_ASSERT(cache.stats().resident_bytes <= budget); + }; + if (budget == 320 && reverse) commit(second, b); + commit(first, a); + if (budget == 320 && !reverse) commit(second, b); + } + } + unlink(path.c_str()); +} + // (i) A capture can commit shorter than it reserved (agent-turn replay saves // at the backend's position), so an ancestor credited as freed may stay // resident. Enforcing the budget after the commit restores the invariant. @@ -9358,6 +9593,36 @@ struct MockBackendWithAdopt : MockBackendWithLayout { } }; +TEST_CASE(ServerUnitFixture, test_disk_cache_failed_init_disables_pending_writes) { + MockBackendWithAdopt backend; + const std::string file = test_tmp_path("luce_test_disk_init_blocker").string(); + rm_rf(file); + FILE * blocker = std::fopen(file.c_str(), "wb"); + TEST_ASSERT(blocker != nullptr); + std::fclose(blocker); + for (const std::string & dir : {file, file + "/cache"}) { + DiskCacheConfig cfg; + cfg.cache_dir = dir; + cfg.min_tokens = 1; + cfg.max_pending_bytes = 1; // a second queued write would block without a writer + DiskPrefixCache cache(cfg, backend); + TEST_ASSERT(!cache.init()); + TEST_ASSERT(cache.disabled()); + const auto ref = backend.snapshot_ref(0); + std::vector key(40, 1); + for (int turn = 0; turn < 2; ++turn) { + ++key.back(); + TEST_ASSERT(!cache.save(0, key)); + TEST_ASSERT(!cache.save_async(OwnedSnapshot::clone_to_host( + ref.ctx, ref.cur_pos, ref.last_tok), key)); + } + TEST_ASSERT(cache.entry_count() == 0); + TEST_ASSERT(cache.stats().pending_bytes == 0); + TEST_ASSERT(cache.stats().backpressure_waits == 0); + } + rm_rf(file); +} + TEST_CASE(ServerUnitFixture, test_disk_cache_rejects_snapshot_past_key) { // A snapshot of kMaxPos positions may only be filed under a key that // covers at least kMaxPos tokens; shorter keys are refused on save and, @@ -9721,6 +9986,38 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_budget_lru_survives_restart) { rm_rf(dir); } +TEST_CASE(ServerUnitFixture, test_disk_cache_duplicate_save_refreshes_ttl_after_restart) { + MockBackendWithAdopt backend; + const std::string dir = test_tmp_path("luce_test_disk_duplicate_touch").string(); + const auto key = disk_test_key(40); + for (bool async : {false, true}) { + rm_rf(dir); + uint64_t now = 100000; + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; cfg.ttl_seconds = 100; + { + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&] { return now; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.save(0, key)); + cache.drain(); + now += 90; + const auto ref = backend.snapshot_ref(0); + TEST_ASSERT(async ? cache.save_async(OwnedSnapshot::clone_to_host( + ref.ctx, ref.cur_pos, ref.last_tok), key) : cache.save(0, key)); + TEST_ASSERT(cache.stats().writes == 1); + } + now += 60; // expired from creation, alive from the duplicate save + { + DiskPrefixCache cache(cfg, backend); + cache.set_clock([&] { return now; }); + TEST_ASSERT(cache.init()); + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(cache.lookup(key, 1)); + } + } + rm_rf(dir); +} + TEST_CASE(ServerUnitFixture, test_disk_cache_restart_restores_and_rejects_other_identity) { MockBackendWithAdopt backend; const std::string dir = test_tmp_path("luce_test_disk_restart").string(); @@ -9906,19 +10203,22 @@ TEST_CASE(ServerUnitFixture, test_disk_cache_write_failure_is_not_fatal) { TEST_ASSERT(cache.save(0, disk_test_key(40))); cache.drain(); const std::string layout_dir = fs::path(dkv_files(dir)[0]).parent_path().string(); - // Unwritable directory (stands in for a full disk): the write, in the - // background, fails cleanly, leaves no temp file and drops the entry. - fs::permissions(layout_dir, fs::perms::owner_read | fs::perms::owner_exec); + // Block the directory with a file: a deterministic write failure on + // Windows and POSIX, including runs as root. Restore the original file. + const std::string backup_dir = layout_dir + ".backup"; + fs::rename(layout_dir, backup_dir); + FILE * blocker = std::fopen(layout_dir.c_str(), "wb"); + TEST_ASSERT(blocker != nullptr); + std::fclose(blocker); const bool queued = cache.save(0, disk_test_key(40, 800)); cache.drain(); - fs::permissions(layout_dir, fs::perms::owner_all); + fs::remove(layout_dir); + fs::rename(backup_dir, layout_dir); TEST_ASSERT(queued); - if (::geteuid() != 0) { - TEST_ASSERT(cache.stats().write_failures == 1); - TEST_ASSERT(cache.entry_count() == 1); - TEST_ASSERT(dkv_files(dir).size() == 1); - TEST_ASSERT(!cache.lookup(disk_test_key(40, 800), 1)); - } + TEST_ASSERT(cache.stats().write_failures == 1); + TEST_ASSERT(cache.entry_count() == 1); + TEST_ASSERT(dkv_files(dir).size() == 1); + TEST_ASSERT(!cache.lookup(disk_test_key(40, 800), 1)); rm_rf(dir); } @@ -10029,7 +10329,8 @@ ServerConfig disk_session_config(const std::string & dir) { // Commit `prompt` as an inline snapshot at its end, the way a request does. int commit_turn(HttpServer & server, DiskPayloadBackend & backend, - const std::vector & prompt, int restore_slot = -1) { + const std::vector & prompt, int restore_slot = -1, + bool compressed = false) { PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); auto reservation = cache.reserve_inline_snap( prompt, restore_slot >= 0 ? backend.snapshot_cur_pos(restore_slot) : 0, @@ -10041,12 +10342,77 @@ int commit_turn(HttpServer & server, DiskPayloadBackend & backend, backend.put(slot, (int)prompt.size()); SchedulerTestHarness::finalize_inline_snapshot( server, prompt, std::move(reservation), slot, (int)prompt.size(), - restore_slot); + restore_slot, compressed); return slot; } } // namespace +TEST_CASE(ServerUnitFixture, test_disk_session_compressed_snapshot_keeps_ram_reuse_without_disk_key) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_compressed").string(); + rm_rf(dir); + { + auto owner = std::make_unique(); + auto & backend = *owner; + LuceEngine engine(std::move(owner)); + ServerConfig config = disk_session_config(dir); + config.prefix_cache_cap = 1; + HttpServer server(engine, tokenizer, config); + const std::vector view = {1, 100, 3, 101, 4}; + const int slot = commit_turn(server, backend, view, -1, true); + TEST_ASSERT(slot >= 0); + TEST_ASSERT(SchedulerTestHarness::prefix_cache(server).lookup_candidate(view, 5).second == 5); + SchedulerTestHarness::flush_sessions(server); + auto & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(disk.stats().writes == 0); + // Evict the compressed state: raw sessions still persist normally. + const std::vector raw = {1, 200, 3, 201, 4}; + TEST_ASSERT(commit_turn(server, backend, raw) >= 0); + disk.drain(); + TEST_ASSERT(disk.stats().writes == 0); + SchedulerTestHarness::flush_sessions(server); + TEST_ASSERT(disk.stats().writes == 1); + TEST_ASSERT(disk.longest_prefix(view).len == 0); + TEST_ASSERT(disk.longest_prefix(raw).len == 5); + } + rm_rf(dir); + unlink(path.c_str()); +} + +TEST_CASE(ServerUnitFixture, test_disk_session_failed_init_keeps_ram_cache_working) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + { + auto backend_owner = std::make_unique(); + auto & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + // A file cannot contain a cache directory: deterministic on every OS. + ServerConfig config = disk_session_config(path + "/cache"); + config.prefix_cache_max_bytes = 3000; + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(disk.disabled()); + TEST_ASSERT(!SchedulerTestHarness::props(server)["full_cache"].value("disk_enabled", true)); + const std::vector a = {1, 100, 3, 101}; + const std::vector b = {1, 200, 3, 201}; + const std::vector c = {1, 300, 3, 301}; + TEST_ASSERT(commit_turn(server, backend, a) >= 0); + const int slot_b = commit_turn(server, backend, b); + TEST_ASSERT(slot_b >= 0); + TEST_ASSERT(commit_turn(server, backend, c) >= 0); // evict a without disk work + const auto restored = SchedulerTestHarness::prepare_cache(server, {1, 200, 3, 201, 4, 102}); + TEST_ASSERT(restored.restore_slot == slot_b && restored.prefix_len == 4); + TEST_ASSERT(disk.entry_count() == 0 && disk.stats().pending_bytes == 0); + SchedulerTestHarness::flush_sessions(server); + TEST_ASSERT(disk.stats().writes == 0); + } + unlink(path.c_str()); +} + // Persist-on-evict policy: every conversation's newest snapshot and the // shared heads are written; a previous turn kept only as an edit backup is // not, while its conversation's deeper snapshot is resident. @@ -10151,6 +10517,62 @@ TEST_CASE(ServerUnitFixture, test_disk_session_ram_eviction_writes_and_restores) unlink(path.c_str()); } +// Single-sequence path: a long session that arrives after several short +// ones is captured; the short ones it pushes out of RAM go to disk first. +TEST_CASE(ServerUnitFixture, test_disk_session_large_capture_evicts_several_to_disk) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_several").string(); + rm_rf(dir); + const std::vector> small = { + {1, 100, 3, 101}, {1, 200, 3, 201}, {1, 300, 3, 301}, {1, 400, 3, 401}}; + std::vector big = {1, 900, 3}; + big.insert(big.end(), 40, 901); // 43 tokens: 5300 bytes, 1400 per short one + { + auto backend_owner = std::make_unique(); + DiskPayloadBackend & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config = disk_session_config(dir); + config.prefix_cache_max_bytes = 7000; + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + std::vector slots; + for (const auto & ids : small) { + slots.push_back(commit_turn(server, backend, ids)); + TEST_ASSERT(slots.back() >= 0); + } + TEST_ASSERT(disk.entry_count() == 0); + const int big_slot = commit_turn(server, backend, big); + TEST_ASSERT(big_slot >= 0); + TEST_ASSERT(cache.stats().budget_skips == 0); + TEST_ASSERT(cache.stats().resident_bytes == 1400 + 5300); + // The three oldest left RAM, written then freed; the newest stays. + for (int i = 0; i < 3; ++i) { + TEST_ASSERT(std::count(backend.freed.begin(), backend.freed.end(), slots[i]) >= 1); + TEST_ASSERT(!backend.snapshot_used(slots[i])); + TEST_ASSERT(cache.lookup_candidate(small[i], 4).first < 0); + TEST_ASSERT(disk.longest_prefix(small[i]).len == 4); + } + TEST_ASSERT(cache.lookup_candidate(small[3], 4).first == slots[3]); + TEST_ASSERT(disk.longest_prefix(small[3]).len == 0); + std::vector next = big; + next.insert(next.end(), {4, 902}); + TEST_ASSERT(cache.lookup_candidate(next, (int)next.size()).second == (int)big.size()); + disk.drain(); + TEST_ASSERT(disk.stats().writes == 3 && disk.stats().pending_files == 0); + // A short session restores from disk afterwards. + std::vector again = small[0]; + again.insert(again.end(), {4, 102}); + const auto prepared = SchedulerTestHarness::prepare_cache(server, again); + TEST_ASSERT(prepared.restore_slot == ModelBackend::kMaxSlots - 1); + TEST_ASSERT(prepared.prefix_len == 4); + } + rm_rf(dir); + unlink(path.c_str()); +} + // Tool path: an agent turn saved with its generated text (the live state // after generation) is written when the resident budget evicts it, keyed // by the tokens the next request renders. @@ -10188,6 +10610,94 @@ TEST_CASE(ServerUnitFixture, test_disk_session_agent_turn_eviction_writes_live_s unlink(path.c_str()); } +TEST_CASE(ServerUnitFixture, test_qwen_disk_adopt_restores_dense_state_and_rejects_failed_restore) { + struct CpuSnapshotBackend : Qwen35Backend { + CpuSnapshotBackend() : Qwen35Backend(Qwen35Config{}) { + target_weights().n_layer = 0; + target_weights().full_attention_interval = 1; + auto & c = target_cache(); + c.max_ctx = c.target_feat_cap = 16; + ggml_init_params ip{}; + ip.mem_size = 4 * ggml_tensor_overhead() + 1024; + ip.no_alloc = true; + c.base_ctx = ggml_init(ip); + c.target_feat = ggml_new_tensor_2d(c.base_ctx, GGML_TYPE_F32, 2, 16); + c.base_buf = ggml_backend_alloc_ctx_tensors_from_buft(c.base_ctx, ggml_backend_cpu_buffer_type()); + ggml_backend_buffer_clear(c.base_buf, 0); + } + bool restore(int slot) { return restore_target_cache_from_snapshot(slot); } + TargetCache & cache() { return target_cache(); } + } backend; + const std::string dir = test_tmp_path("luce_test_qwen_disk_adopt").string(); + rm_rf(dir); + { + DiskCacheConfig cfg; cfg.cache_dir = dir; cfg.min_tokens = 1; + DiskPrefixCache disk(cfg, backend); + TEST_ASSERT(disk.init()); + ggml_init_params ip{}; + ip.mem_size = 4 * ggml_tensor_overhead() + 1024; + ip.no_alloc = true; + auto * ctx = ggml_init(ip); + auto * t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2, 3); + ggml_set_name(t, "snap_target_feat"); + auto buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_cpu_buffer_type()); + const std::vector values = {1, 2, 3, 4, 5, 6}; + ggml_backend_tensor_set(t, values.data(), 0, values.size() * sizeof(float)); + const std::vector key = {1, 2, 3}; + TEST_ASSERT(disk.save_async(OwnedSnapshot(ctx, buf, 3, 3), key)); + disk.drain(); + TEST_ASSERT(disk.lookup(key, 0)); // actual Qwen35Backend::snapshot_adopt + TEST_ASSERT(backend.restore(0)); // requires the adopted layout to be dense + TEST_ASSERT(backend.cache().cur_pos == 3); + std::vector restored(values.size()); + ggml_backend_tensor_get(backend.cache().target_feat, restored.data(), 0, + restored.size() * sizeof(float)); + TEST_ASSERT(restored == values); + // A restore failure must return before any prefill or decode on cleared KV. + ++backend.cache().max_ctx; + GenerateRequest req; + req.prompt = key; + req.n_gen = 1; + DaemonIO io; + const auto failed = backend.restore_and_generate_impl(0, req, io); + TEST_ASSERT(!failed.ok() && failed.tokens.empty()); + TEST_ASSERT(failed.error && failed.error->code == GenerateErrorCode::InvalidSnapshotSlot); + TEST_ASSERT(backend.cache().cur_pos == 0); + } + rm_rf(dir); +} + +TEST_CASE(ServerUnitFixture, test_disk_agent_turn_respects_request_scope) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_agent_disk_scope").string(); + const std::vector prompt = {1, 100, 3, 101}; + const std::vector generated = {200, 201, 202}; + const std::vector canonical = {1, 100, 3, 101, 200, 201, 202, 7}; + for (auto scope : {DiskPrefixCacheMode::Off, DiskPrefixCacheMode::Auto, + DiskPrefixCacheMode::Fixed, DiskPrefixCacheMode::Full}) { + rm_rf(dir); + { + auto backend_owner = std::make_unique(); + auto & backend = *backend_owner; + LuceEngine engine(std::move(backend_owner)); + HttpServer server(engine, tokenizer, disk_session_config(dir)); + backend.live_position = 6; + TEST_ASSERT(SchedulerTestHarness::save_generated_turn( + server, prompt, generated, canonical, -1, scope)); + // All scopes keep the RAM checkpoint, but only full keys it for disk. + const auto hit = SchedulerTestHarness::prefix_cache(server).lookup_candidate(canonical, 8); + TEST_ASSERT(hit.first >= 0 && hit.second == 6); + SchedulerTestHarness::flush_sessions(server); + auto & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(disk.stats().writes == (scope == DiskPrefixCacheMode::Full ? 1 : 0)); + } + } + rm_rf(dir); + unlink(path.c_str()); +} + // Graceful shutdown is one last eviction: the head and each conversation's // newest snapshot are written, the edit backup and pruned turns are not, and // a restarted server restores the session from disk. @@ -10463,7 +10973,10 @@ class DiskCheckpointEngine final : public SeqEngine { int max_context() const override { return 64; } bool supports_prefix_store() const override { return true; } bool prefix_store_in_backend_snapshots() const override { return true; } - size_t estimate_prefix_store_bytes(int) const override { return 256; } + size_t estimate_prefix_store_bytes(int tokens) const override { + return bytes_per_token ? (size_t)tokens * bytes_per_token : 256; + } + size_t bytes_per_token = 0; // 0: every checkpoint is 256 bytes bool token_is_eos(int32_t token) const override { return token == 2; } StepPlanLimits step_plan_limits(int) const override { return {2, 64, 128, 64}; @@ -10528,7 +11041,8 @@ class DiskCheckpointEngine final : public SeqEngine { DiskPayloadBackend::make(slot.capture.checkpoint.tokens); output.prefix_store.status = PrefixStoreEvent::Status::saved; output.prefix_store.ticket = slot.capture; - output.prefix_store.bytes = 256; + output.prefix_store.bytes = + estimate_prefix_store_bytes(slot.capture.checkpoint.tokens); slot.capture = {}; } result.prefills.push_back(std::move(output)); @@ -10638,6 +11152,170 @@ TEST_CASE(ServerUnitFixture, test_disk_session_concurrent_eviction_and_restore) rm_rf(dir); unlink(path.c_str()); } + +TEST_CASE(ServerUnitFixture, test_disk_session_concurrent_respects_request_scope) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_paged_scope").string(); + for (const auto mode : {DiskPrefixCacheMode::Off, DiskPrefixCacheMode::Fixed, + DiskPrefixCacheMode::Auto}) { + rm_rf(dir); + auto owner = std::make_unique(); + auto & backend = *owner; + LuceEngine engine(std::move(owner)); + ServerConfig config = disk_session_config(dir); + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 2; + config.concurrent_prefix_cache_max_bytes = 4096; + config.concurrent_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + auto & disk = SchedulerTestHarness::disk_cache(server); + const std::vector a = {1, 100, 3, 101, 4}; + backend.slots[62] = DiskPayloadBackend::make(3); + disk.learn_layout(62); + TEST_ASSERT(disk.save(62, {1, 100, 3})); + disk.drain(); + backend.drop(62); + + std::vector> jobs; + std::vector fds; + const auto run_one = [&](std::vector prompt) { + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + fds.insert(fds.end(), {sockets[0], sockets[1]}); + auto job = std::make_unique(); + job->fd = sockets[0]; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.prompt_tokens = std::move(prompt); + job->req.max_output = 1; + job->req.model = "scheduler-test"; + job->req.disk_cache_policy.mode = mode; + job->req.disk_cache_policy.fixed_tokens = 2; + ServerJob * raw = job.get(); + jobs.push_back(std::move(job)); + SchedulerTestHarness::enqueue(server, raw); + std::unique_lock lock(raw->mu); + return raw->cv.wait_for(lock, std::chrono::seconds(5), [&] { return raw->done; }); + }; + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + const std::vector c = {1, 300, 3, 301, 4}; + const bool ok = run_one(a) && run_one({1, 200, 3, 201, 4}) && + run_one(c) && run_one(c); + SchedulerTestHarness::stop(server); + scheduler.join(); + for (int fd : fds) close(fd); + SchedulerTestHarness::flush_sessions(server); + + TEST_ASSERT(ok); + bool disk_restored = false, ram_restored = false; + for (const auto & r : backend.engine.restores) { + disk_restored |= r.id == (uint64_t)ModelBackend::kMaxSlots; + ram_restored |= r.id != (uint64_t)ModelBackend::kMaxSlots; + } + TEST_ASSERT(disk_restored == (mode == DiskPrefixCacheMode::Auto)); + TEST_ASSERT(ram_restored); // scoped disk policy still allows RAM reuse + TEST_ASSERT(disk.stats().writes == 1); // only the seeded Full entry + TEST_ASSERT(disk.entry_count() == 1); + TEST_ASSERT(backend.displaced.empty()); + } + rm_rf(dir); + unlink(path.c_str()); +} + +// Concurrent (paged) serving does the same: a long session's checkpoint +// that no single resident one makes room for evicts several when it lands, +// writes them to disk (with the disk on) and frees their payloads. +TEST_CASE(ServerUnitFixture, test_disk_session_concurrent_large_capture_evicts_several) { + const std::string path = write_deepseek_marker_tokenizer_fixture(); + Tokenizer tokenizer; + TEST_ASSERT(tokenizer.load_from_gguf(path.c_str())); + const std::string dir = test_tmp_path("luce_test_session_paged_several").string(); + for (const bool disk_on : {true, false}) { + rm_rf(dir); + auto backend_owner = std::make_unique(); + DiskCheckpointBackend & backend = *backend_owner; + backend.engine.bytes_per_token = 100; + LuceEngine engine(std::move(backend_owner)); + ServerConfig config = disk_session_config(disk_on ? dir : std::string()); + config.arch = "qwen35"; + config.max_ctx = 64; + config.prefix_cache_cap = 8; + config.concurrent_prefix_cache_max_bytes = 4500; + config.concurrent_prefix_cache = true; + config.admission_coalesce_ms = 0; + HttpServer server(engine, tokenizer, config); + DiskPrefixCache & disk = SchedulerTestHarness::disk_cache(server); + TEST_ASSERT(disk.disabled() == !disk_on); + + std::vector> jobs; + std::vector fds; + const auto run_one = [&](std::vector prompt, const char * id) { + int sockets[2] = {-1, -1}; + TEST_ASSERT(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + fds.push_back(sockets[0]); + fds.push_back(sockets[1]); + auto job = std::make_unique(); + job->fd = sockets[0]; + job->req.format = ApiFormat::OPENAI_CHAT; + job->req.prompt_tokens = std::move(prompt); + job->req.max_output = 1; + job->req.stream = false; + job->req.model = "scheduler-test"; + job->req.response_id = id; + jobs.push_back(std::move(job)); + ServerJob * raw = jobs.back().get(); + SchedulerTestHarness::enqueue(server, raw); + std::unique_lock lock(raw->mu); + return raw->cv.wait_for(lock, std::chrono::seconds(5), [&] { return raw->done; }); + }; + std::thread scheduler([&] { SchedulerTestHarness::run(server, backend.engine); }); + // Four short conversations save their 3-token heads (300 bytes each), + // then a long one saves its 42-token head (4200 bytes): 900 must go. + const std::vector> small = { + {1, 100, 3, 101, 4}, {1, 200, 3, 201, 4}, + {1, 300, 3, 301, 4}, {1, 400, 3, 401, 4}}; + std::vector big = {1}; + big.insert(big.end(), 40, 900); + big.insert(big.end(), {3, 901, 4}); + bool ok = true; + for (size_t i = 0; i < small.size(); ++i) { + ok = ok && run_one(small[i], "s"); + } + const auto before = SchedulerTestHarness::prefix_cache(server).stats(); + ok = ok && run_one(big, "big"); + disk.drain(); + SchedulerTestHarness::stop(server); + scheduler.join(); + for (const int fd : fds) close(fd); + + TEST_ASSERT(ok); + TEST_ASSERT(before.in_use == 4 && before.resident_bytes == 1200); + PrefixCache & cache = SchedulerTestHarness::prefix_cache(server); + const auto after = cache.stats(); + TEST_ASSERT(after.budget_skips == 0); + TEST_ASSERT(after.in_use == 2 && after.resident_bytes == 300 + 4200); + std::vector next = big; + next.push_back(9); + TEST_ASSERT(cache.lookup_candidate(next, (int)next.size()).second == 42); + TEST_ASSERT(cache.lookup_candidate(small[3], 5).second == 3); + // Only the newest short head and the long one keep a payload. + TEST_ASSERT(backend.slots.size() == 2); + TEST_ASSERT(backend.displaced.empty()); + if (disk_on) { + TEST_ASSERT(disk.stats().writes == 3); + for (int i = 0; i < 4; ++i) { + std::vector probe = small[(size_t)i]; + probe.push_back(9); + TEST_ASSERT(disk.longest_prefix(probe).len == (i < 3 ? 3 : 0)); + } + } + } + rm_rf(dir); + unlink(path.c_str()); +} #endif TEST_CASE(ServerUnitFixture, test_backend_ipc_rejects_file_work_dir) {