server: rework prompt cache persistence into a dir with per-entry idle sync
Replace the single --prompt-cache-file with --prompt-cache-dir: each --cache-ram entry becomes one self-describing, fingerprint-checked file named by its content hash, written atomically (.tmp + fsync + rename). A background sweep (GC-thread pattern) mirrors the cache into the dir while the server runs: entries older than --prompt-cache-sync-idle (default 300s) are written without stopping serving, and files that no longer belong to the cache are pruned. Membership is re-checked under the cache mutex before each unlink, and any write error aborts the sweep (retry next pass, RAM stays the source of truth). To make off-thread snapshotting safe, cache entries are now refcounted and immutable after publish: two-phase alloc (make_state/insert_state stamps content_hash + t_finalized_ms), and the steal path erases from the list instead of clearing blobs in place. The dir is serialized with a flock lock file; a second instance disables persistence with an error instead of corrupting the mirror. The cache is also flushed when entering sleep mode, which otherwise drops the RAM cache on wake. Co-Authored-By: Hermes <hermes@nousresearch.com>
This commit is contained in:
hal9000
+16
-12
@@ -3666,21 +3666,25 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SLOT_SAVE_PATH"));
|
||||
add_opt(common_arg(
|
||||
{"--prompt-cache-file"}, "PATH",
|
||||
"file to persist the --cache-ram prompt cache across restarts (default: disabled)",
|
||||
{"--prompt-cache-dir"}, "PATH",
|
||||
"directory to persist the --cache-ram prompt cache across restarts, one file per entry (default: disabled)",
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.prompt_cache_file = value;
|
||||
// create the parent directory, so the save on shutdown does not fail
|
||||
const auto parent = std::filesystem::path(value).parent_path();
|
||||
if (!parent.empty()) {
|
||||
std::error_code ec;
|
||||
std::filesystem::create_directories(parent, ec);
|
||||
if (ec) {
|
||||
throw std::invalid_argument("failed to create the parent directory for the prompt cache file: " + ec.message());
|
||||
}
|
||||
params.prompt_cache_dir = value;
|
||||
// create the directory, so the sync does not fail
|
||||
std::error_code ec;
|
||||
std::filesystem::create_directories(value, ec);
|
||||
if (ec) {
|
||||
throw std::invalid_argument("failed to create the prompt cache directory: " + ec.message());
|
||||
}
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PROMPT_CACHE_FILE"));
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PROMPT_CACHE_DIR"));
|
||||
add_opt(common_arg(
|
||||
{"--prompt-cache-sync-idle"}, "N",
|
||||
string_format("minimum time in seconds a prompt cache entry must sit unchanged in RAM before the live sync writes it to --prompt-cache-dir (default: %d)", 300),
|
||||
[](common_params & params, const std::string & value) {
|
||||
params.prompt_cache_sync_idle = std::stoi(value);
|
||||
}
|
||||
).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_PROMPT_CACHE_SYNC_IDLE"));
|
||||
add_opt(common_arg(
|
||||
{"--media-path"}, "PATH",
|
||||
"directory for loading local media files; files can be accessed via file:// URLs using relative paths (default: disabled)",
|
||||
|
||||
+2
-1
@@ -705,7 +705,8 @@ struct common_params {
|
||||
bool log_json = false;
|
||||
|
||||
std::string slot_save_path;
|
||||
std::string prompt_cache_file; // file to persist the --cache-ram prompt cache across restarts
|
||||
std::string prompt_cache_dir; // directory to persist the --cache-ram prompt cache across restarts (one file per entry)
|
||||
int32_t prompt_cache_sync_idle = 300; // min seconds an entry must sit unchanged in RAM before the live sync writes it to prompt_cache_dir
|
||||
std::string media_path; // path to directory for loading media files
|
||||
|
||||
float slot_prompt_similarity = 0.1f;
|
||||
|
||||
@@ -17,6 +17,8 @@ add_library(${TARGET} STATIC
|
||||
server-decision.h
|
||||
server-context.cpp
|
||||
server-context.h
|
||||
server-pcache-sync.cpp
|
||||
server-pcache-sync.h
|
||||
server-stream.cpp
|
||||
server-stream.h
|
||||
server-tools.cpp
|
||||
|
||||
@@ -168,7 +168,8 @@ For the full list of features, please refer to [server's changelog](https://gith
|
||||
| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, 0 for decision models, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)<br/>(env: LLAMA_ARG_CACHE_RAM) |
|
||||
| `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)<br/>(env: LLAMA_ARG_KV_UNIFIED) |
|
||||
| `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)<br/>(env: LLAMA_ARG_CACHE_IDLE_SLOTS) |
|
||||
| `--prompt-cache-file PATH` | file to persist the `--cache-ram` prompt cache across restarts. On shutdown the cache is saved to PATH, on startup it is restored (only if the model + context fingerprint in the file header matches, otherwise the file is skipped). No-op when unset. (default: disabled)<br/>(env: LLAMA_ARG_PROMPT_CACHE_FILE) |
|
||||
| `--prompt-cache-dir PATH` | directory to persist the `--cache-ram` prompt cache across restarts, one file per entry, named after the entry's content hash. Entries idle for `--prompt-cache-sync-idle` seconds are synced to the directory in the background and the directory is kept a strict mirror of the cache (stale files are pruned). On startup the entries are restored (entry files that do not match the model, context and build fingerprint are skipped). Only one `llama-server` instance may use a directory at a time (`.lock` file). No-op when unset. (default: disabled)<br/>(env: LLAMA_ARG_PROMPT_CACHE_DIR) |
|
||||
| `--prompt-cache-sync-idle N` | minimum time in seconds a prompt cache entry must sit unchanged in RAM before the live sync writes it to `--prompt-cache-dir` (default: 300). Only used together with `--prompt-cache-dir`.<br/>(env: LLAMA_ARG_PROMPT_CACHE_SYNC_IDLE) |
|
||||
| `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)<br/>(env: LLAMA_ARG_CONTEXT_SHIFT) |
|
||||
| `-r, --reverse-prompt PROMPT` | halt generation at PROMPT, return control in interactive mode |
|
||||
| `-sp, --special` | special tokens output enabled (default: false) |
|
||||
|
||||
+104
-37
@@ -7,6 +7,7 @@
|
||||
#include "server-common.h"
|
||||
#include "server-decision.h"
|
||||
#include "server-http.h"
|
||||
#include "server-pcache-sync.h"
|
||||
#include "server-task.h"
|
||||
#include "server-queue.h"
|
||||
#include "server-schema.h"
|
||||
@@ -352,7 +353,9 @@ struct server_slot {
|
||||
SRV_TRC(" - saving prompt with length %d, total state size = %.3f MiB (draft: %.3f MiB)\n",
|
||||
(int) prompt.tokens.size(), cur_size / (1024.0 * 1024.0), cur_size_dft / (1024.0 * 1024.0));
|
||||
|
||||
auto * cur = prompt_cache.alloc(prompt, cur_size_tgt, cur_size_dft);
|
||||
// two-phase publish: the entry is created unpublished, the blobs are filled from the
|
||||
// context, then it is published to the cache (see server_prompt_cache::make_state())
|
||||
auto cur = prompt_cache.make_state(prompt, cur_size_tgt, cur_size_dft);
|
||||
if (cur == nullptr) {
|
||||
return false;
|
||||
}
|
||||
@@ -362,6 +365,8 @@ struct server_slot {
|
||||
llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE);
|
||||
}
|
||||
|
||||
prompt_cache.insert_state(std::move(cur));
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1094,6 +1099,11 @@ private:
|
||||
|
||||
std::unique_ptr<server_prompt_cache> prompt_cache;
|
||||
|
||||
// background mirror of the prompt cache to --prompt-cache-dir; (re-)created by load_model()
|
||||
// - on wake from sleeping, the old instance is destroyed before the new prompt_cache, so
|
||||
// the sync thread can never point at a freed cache
|
||||
std::unique_ptr<server_prompt_cache_sync> prompt_cache_sync;
|
||||
|
||||
server_metrics metrics;
|
||||
|
||||
// queued prompt stats - llama_decode() is async, so the timing is only valid after a sync
|
||||
@@ -1139,6 +1149,12 @@ private:
|
||||
// note: for sleeping == false, event is emitted by load_model()
|
||||
}
|
||||
SRV_INF("%s", "server is entering sleeping state\n");
|
||||
// flush the prompt cache to disk while the contexts are still alive and the queue
|
||||
// is quiesced (we run from an on_sleeping_state callback, before destroy()); the
|
||||
// sync thread is stopped first and the flush runs on this thread - the sync
|
||||
// instance itself is re-created by load_model() on wake
|
||||
prompt_cache_sync_stop();
|
||||
prompt_cache_persist("sleeping");
|
||||
destroy();
|
||||
} else {
|
||||
SRV_INF("%s", "server is exiting sleeping state\n");
|
||||
@@ -1149,6 +1165,64 @@ private:
|
||||
sleeping = new_state;
|
||||
}
|
||||
|
||||
// stop the persistence sync thread, keep the instance (flush() must still work after it -
|
||||
// see prompt_cache_persist)
|
||||
void prompt_cache_sync_stop() {
|
||||
if (prompt_cache_sync) {
|
||||
prompt_cache_sync->stop();
|
||||
}
|
||||
}
|
||||
|
||||
// capture the live slots into the prompt cache, then fully sync the cache to
|
||||
// --prompt-cache-dir: entry files that are missing get written, stale files are pruned
|
||||
// runs entirely on the calling (main) thread, before the prompt cache is torn down; used
|
||||
// at shutdown and when entering sleeping state, with the sync thread stopped
|
||||
void prompt_cache_persist(const char * reason) {
|
||||
if (prompt_cache == nullptr || params_base.prompt_cache_dir.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (prompt_cache_sync == nullptr || !prompt_cache_sync->owns_dir()) {
|
||||
// never started, or another instance holds the dir lock - do not touch the files
|
||||
SRV_TRC("prompt cache dir '%s' is not persisted by this instance, skipping save (%s)\n",
|
||||
params_base.prompt_cache_dir.c_str(), reason);
|
||||
return;
|
||||
}
|
||||
|
||||
auto & prompt_cache_ref = *prompt_cache;
|
||||
|
||||
// capture the live slots as well: the prompt cache only holds contexts that were saved
|
||||
// when a slot was about to be reused or cleared, so an active conversation is not in it
|
||||
// the size limit is lifted for the capture, so contexts larger than --cache-ram are persisted too
|
||||
// note: skipped while sleeping - destroy() has already freed the llama contexts that
|
||||
// prompt_save() needs to serialize the KV (use-after-free); nothing can have
|
||||
// changed since sleep, because handle_sleeping_state() captured and flushed
|
||||
// before calling destroy()
|
||||
if (!sleeping) {
|
||||
const size_t limit_size_old = prompt_cache_ref.limit_size;
|
||||
prompt_cache_ref.limit_size = 0;
|
||||
|
||||
for (auto & slot : slots) {
|
||||
if (slot.prompt.tokens.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (slot.prompt_save(prompt_cache_ref)) {
|
||||
SRV_INF("prompt cache save: captured live slot %d (%zu tokens)\n", slot.id, slot.prompt.tokens.size());
|
||||
}
|
||||
}
|
||||
|
||||
prompt_cache_ref.limit_size = limit_size_old;
|
||||
}
|
||||
|
||||
// full dir sync on this thread, it must not spawn threads (see server_prompt_cache_sync::flush())
|
||||
prompt_cache_sync->flush();
|
||||
|
||||
SRV_INF("saved %zu prompt cache entries (%.3f MiB, %zu tokens) to '%s' (%s)\n",
|
||||
prompt_cache_ref.states.size(), prompt_cache_ref.size() / (1024.0 * 1024.0),
|
||||
prompt_cache_ref.n_tokens(), params_base.prompt_cache_dir.c_str(), reason);
|
||||
}
|
||||
|
||||
struct load_progress_data {
|
||||
server_context_impl * ctx;
|
||||
std::string stage;
|
||||
@@ -1580,32 +1654,45 @@ private:
|
||||
}
|
||||
SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n");
|
||||
|
||||
// stop the persistence sync thread before re-creating the cache below - on wake
|
||||
// from sleeping, a previous instance may still reference the old prompt_cache
|
||||
prompt_cache_sync_stop();
|
||||
prompt_cache_sync.reset();
|
||||
|
||||
prompt_cache = std::make_unique<server_prompt_cache>(params_base.cache_ram_mib, n_ctx);
|
||||
|
||||
// fingerprint of the loaded model, written to / checked against the prompt cache file
|
||||
// fingerprint of the loaded model, written to / checked against the prompt cache entry files
|
||||
prompt_cache->fp_n_ctx = n_ctx;
|
||||
prompt_cache->fp_n_embd = llama_model_n_embd_inp(model_tgt);
|
||||
prompt_cache->fp_n_layer = llama_model_n_layer(model_tgt);
|
||||
prompt_cache->fp_n_head = llama_model_n_head(model_tgt);
|
||||
prompt_cache->fp_n_vocab = (uint32_t) llama_vocab_n_tokens(vocab);
|
||||
prompt_cache->fp_model_path = params_base.model.path;
|
||||
prompt_cache->fp_build = std::string(llama_build_info());
|
||||
|
||||
// optionally restore the prompt cache from disk (see --prompt-cache-file)
|
||||
if (!params_base.prompt_cache_file.empty()) {
|
||||
const std::string & cache_file = params_base.prompt_cache_file;
|
||||
// optionally restore the prompt cache from disk (see --prompt-cache-dir)
|
||||
if (!params_base.prompt_cache_dir.empty()) {
|
||||
const std::string & cache_dir = params_base.prompt_cache_dir;
|
||||
std::error_code ec;
|
||||
if (std::filesystem::exists(std::filesystem::u8path(cache_file), ec) && !ec) {
|
||||
SRV_INF("loading prompt cache from '%s'\n", cache_file.c_str());
|
||||
if (std::filesystem::is_directory(std::filesystem::u8path(cache_dir), ec) && !ec) {
|
||||
SRV_INF("loading prompt cache from '%s'\n", cache_dir.c_str());
|
||||
|
||||
const size_t n_loaded = prompt_cache->load_from_file(cache_file, ctx_dft != nullptr, mctx != nullptr);
|
||||
const size_t n_loaded = prompt_cache->load_from_dir(cache_dir, ctx_dft != nullptr, mctx != nullptr);
|
||||
if (n_loaded > 0) {
|
||||
SRV_INF("restored %zu prompt cache entries (%.3f MiB, %zu tokens) from '%s'\n",
|
||||
prompt_cache->states.size(), prompt_cache->size() / (1024.0 * 1024.0),
|
||||
prompt_cache->n_tokens(), cache_file.c_str());
|
||||
prompt_cache->n_tokens(), cache_dir.c_str());
|
||||
}
|
||||
} else {
|
||||
SRV_TRC("prompt cache file '%s' not found, starting with an empty cache\n", cache_file.c_str());
|
||||
SRV_TRC("prompt cache dir '%s' not found, starting with an empty cache\n", cache_dir.c_str());
|
||||
}
|
||||
|
||||
// start the live sync thread (see --prompt-cache-sync-idle); if another
|
||||
// instance holds the dir lock, persistence is disabled but the server keeps
|
||||
// running on the RAM cache
|
||||
prompt_cache_sync = std::make_unique<server_prompt_cache_sync>(
|
||||
prompt_cache.get(), cache_dir, (int64_t) params_base.prompt_cache_sync_idle * 1000);
|
||||
prompt_cache_sync->start();
|
||||
}
|
||||
} else {
|
||||
SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n");
|
||||
@@ -5071,34 +5158,14 @@ void server_context::terminate() {
|
||||
impl->queue_tasks.terminate();
|
||||
}
|
||||
|
||||
void server_context::prompt_cache_sync_stop() {
|
||||
impl->prompt_cache_sync_stop();
|
||||
}
|
||||
|
||||
void server_context::prompt_cache_save() {
|
||||
if (impl->prompt_cache == nullptr || impl->params_base.prompt_cache_file.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
auto & prompt_cache = *impl->prompt_cache;
|
||||
|
||||
// capture the live slots as well: the prompt cache only holds contexts that were saved
|
||||
// when a slot was about to be reused or cleared, so an active conversation is not in it
|
||||
// the size limit is lifted for the capture, so contexts larger than --cache-ram are persisted too
|
||||
{
|
||||
const size_t limit_size_old = prompt_cache.limit_size;
|
||||
prompt_cache.limit_size = 0;
|
||||
|
||||
for (auto & slot : impl->slots) {
|
||||
if (slot.prompt.tokens.empty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (slot.prompt_save(prompt_cache)) {
|
||||
SRV_INF("prompt cache save: captured live slot %d (%zu tokens)\n", slot.id, slot.prompt.tokens.size());
|
||||
}
|
||||
}
|
||||
|
||||
prompt_cache.limit_size = limit_size_old;
|
||||
}
|
||||
|
||||
prompt_cache.save_to_file(impl->params_base.prompt_cache_file);
|
||||
// capture the live slots and do the full dir sync on this thread, with the sync thread
|
||||
// stopped (see server.cpp clean_up)
|
||||
impl->prompt_cache_persist("shutdown");
|
||||
}
|
||||
|
||||
llama_context * server_context::get_llama_context() const {
|
||||
|
||||
@@ -101,12 +101,19 @@ struct server_context {
|
||||
// terminate main loop (will unblock start_loop)
|
||||
void terminate();
|
||||
|
||||
// persist the prompt cache to params.prompt_cache_file: the live slots are captured as well,
|
||||
// so the current contexts are saved even though they are not (yet) in the RAM prompt cache
|
||||
// no-op when the flag is unset or there is nothing to save
|
||||
// not thread-safe, must be called from the main thread before terminate()
|
||||
// persist the prompt cache to params.prompt_cache_dir: the live slots are captured as
|
||||
// well, so the current contexts are saved even though they are not (yet) in the RAM prompt
|
||||
// cache; entry files that are missing get written, stale ones are pruned
|
||||
// no-op when the flag is unset or another instance holds the dir lock
|
||||
// not thread-safe, must be called from the main thread after prompt_cache_sync_stop() and
|
||||
// before terminate()
|
||||
void prompt_cache_save();
|
||||
|
||||
// stop the prompt cache persistence thread (see --prompt-cache-dir), if it is running
|
||||
// no-op when the feature is disabled; the final prompt_cache_save() must run afterwards,
|
||||
// on the main thread
|
||||
void prompt_cache_sync_stop();
|
||||
|
||||
// get the underlaying llama_context, can return nullptr if sleeping
|
||||
// not thread-safe, should only be used from the main thread
|
||||
llama_context * get_llama_context() const;
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
#include "server-pcache-sync.h"
|
||||
|
||||
#include "server-common.h"
|
||||
#include "server-task.h"
|
||||
|
||||
#include <chrono>
|
||||
#include <filesystem>
|
||||
#include <set>
|
||||
#include <vector>
|
||||
|
||||
#if defined(_WIN32)
|
||||
# include <io.h>
|
||||
#else
|
||||
# include <fcntl.h>
|
||||
# include <sys/file.h>
|
||||
# include <unistd.h>
|
||||
#endif
|
||||
|
||||
// how often the sweep runs; independent from --prompt-cache-sync-idle, which only decides
|
||||
// how long an entry must have been unchanged before a sweep may write it
|
||||
static constexpr int PCACHE_SYNC_SWEEP_INTERVAL_S = 30;
|
||||
|
||||
server_prompt_cache_sync::server_prompt_cache_sync(server_prompt_cache * cache, std::string dir_, int64_t sync_idle_ms)
|
||||
: cache(cache), dir(std::move(dir_)), sync_idle_ms(sync_idle_ms) {
|
||||
}
|
||||
|
||||
server_prompt_cache_sync::~server_prompt_cache_sync() {
|
||||
stop();
|
||||
|
||||
// remove the lock file when this instance owned the dir, so a clean shutdown leaves no
|
||||
// stray files (a killed instance cannot clean up - the flock itself is released by the
|
||||
// kernel, the next instance simply re-creates the file)
|
||||
if (owns_dir_) {
|
||||
std::error_code ec;
|
||||
std::filesystem::remove(std::filesystem::u8path(dir) / ".lock", ec);
|
||||
if (ec) {
|
||||
SRV_TRC("failed to remove the prompt cache dir lock file: %s\n", ec.message().c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache_sync::start() {
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(wake_mu);
|
||||
if (running) {
|
||||
// already running - the guard for the wake-from-sleep path, load_model() re-creates
|
||||
// the instance before starting it, so this only protects against double starts
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
std::error_code ec;
|
||||
std::filesystem::create_directories(std::filesystem::u8path(dir), ec);
|
||||
if (ec) {
|
||||
SRV_ERR("failed to create the prompt cache dir '%s': %s - persistence disabled\n",
|
||||
dir.c_str(), ec.message().c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
// serialize the whole directory with a lock file: one llama-server instance per dir
|
||||
// note: a separate .lock file, so that the sweep never mistakes it for an entry
|
||||
const std::filesystem::path lock_path = std::filesystem::u8path(dir) / ".lock";
|
||||
|
||||
#if defined(_WIN32)
|
||||
// no flock equivalent implemented (same guard as the fsync in save_entry()) - best effort,
|
||||
// still open the file so the layout stays the same
|
||||
int fd = _open(lock_path.string().c_str(), O_CREAT | O_RDWR, 0644);
|
||||
bool locked = fd >= 0;
|
||||
#else
|
||||
int fd = ::open(lock_path.string().c_str(), O_CREAT | O_RDWR, 0644);
|
||||
bool locked = fd >= 0 && flock(fd, LOCK_EX | LOCK_NB) == 0;
|
||||
#endif
|
||||
|
||||
if (!locked) {
|
||||
if (fd >= 0) {
|
||||
#if defined(_WIN32)
|
||||
_close(fd);
|
||||
#else
|
||||
::close(fd);
|
||||
#endif
|
||||
}
|
||||
// do not kill the server - it keeps running on the RAM cache alone
|
||||
SRV_ERR("%s", "another llama-server instance holds the prompt cache dir - persistence disabled\n");
|
||||
return;
|
||||
}
|
||||
|
||||
lock_fd = fd;
|
||||
owns_dir_ = true;
|
||||
|
||||
// seed the mirror set from the freshly loaded cache: files restored by load_from_dir()
|
||||
// belong to the mirror and must be prunable when their entries are later consumed or
|
||||
// evicted, even if no sweep ran before that (a sweep's own contains() check could miss
|
||||
// an entry that was stolen within the first interval)
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(cache->mu);
|
||||
for (const auto & state : cache->states) {
|
||||
if (!state->content_hash.empty()) {
|
||||
synced.insert(state->content_hash);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(wake_mu);
|
||||
running = true;
|
||||
}
|
||||
|
||||
SRV_INF("prompt cache persistence enabled, dir = '%s', sync idle = %lld s\n",
|
||||
dir.c_str(), (long long) (sync_idle_ms / 1000));
|
||||
|
||||
thread = std::thread([this] { loop(); });
|
||||
}
|
||||
|
||||
void server_prompt_cache_sync::stop() {
|
||||
bool was_running;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(wake_mu);
|
||||
was_running = running;
|
||||
running = false;
|
||||
}
|
||||
if (was_running) {
|
||||
wake_cv.notify_all();
|
||||
if (thread.joinable()) {
|
||||
thread.join();
|
||||
}
|
||||
}
|
||||
|
||||
// release the dir lock; note: owns_dir_ is kept - flush() must still work after stop(),
|
||||
// it is the final save at shutdown, and the dir stays ours until the process exits
|
||||
if (lock_fd >= 0) {
|
||||
#if !defined(_WIN32)
|
||||
flock(lock_fd, LOCK_UN);
|
||||
::close(lock_fd);
|
||||
#else
|
||||
_close(lock_fd);
|
||||
#endif
|
||||
lock_fd = -1;
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache_sync::loop() {
|
||||
// same shape as the stream session GC (see server-stream.cpp)
|
||||
while (true) {
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(wake_mu);
|
||||
wake_cv.wait_for(lock,
|
||||
std::chrono::seconds(PCACHE_SYNC_SWEEP_INTERVAL_S),
|
||||
[this] { return !running; });
|
||||
if (!running) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
sweep(/* idle_only = */ true);
|
||||
}
|
||||
}
|
||||
|
||||
void server_prompt_cache_sync::flush() {
|
||||
if (!owns_dir_) {
|
||||
SRV_TRC("prompt cache dir '%s' is not owned by this instance, skipping flush\n", dir.c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
sweep(/* idle_only = */ false);
|
||||
}
|
||||
|
||||
void server_prompt_cache_sync::sweep(bool idle_only) {
|
||||
// (a) one listing of the dir: which entries are already on disk
|
||||
std::set<std::string> on_disk;
|
||||
|
||||
std::error_code ec;
|
||||
for (auto it = std::filesystem::directory_iterator(std::filesystem::u8path(dir), ec),
|
||||
end = std::filesystem::directory_iterator();
|
||||
it != end && !ec; it.increment(ec)) {
|
||||
if (!it->is_regular_file()) {
|
||||
continue;
|
||||
}
|
||||
const std::filesystem::path & path = it->path();
|
||||
if (path.extension() == ".bin" && !path.stem().empty()) {
|
||||
on_disk.insert(path.stem().string());
|
||||
}
|
||||
}
|
||||
if (ec) {
|
||||
SRV_ERR("prompt cache dir '%s': %s\n", dir.c_str(), ec.message().c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
// (b) snapshot the entries to write under mu - published entries are immutable, so their
|
||||
// blobs can be serialized outside the lock
|
||||
const int64_t t_now = ggml_time_ms();
|
||||
|
||||
std::vector<std::shared_ptr<server_prompt_cache_state>> to_write;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(cache->mu);
|
||||
|
||||
for (const auto & state : cache->states) {
|
||||
if (state->content_hash.empty() || on_disk.count(state->content_hash)) {
|
||||
continue;
|
||||
}
|
||||
if (idle_only && state->t_finalized_ms + sync_idle_ms >= t_now) {
|
||||
continue;
|
||||
}
|
||||
to_write.push_back(state);
|
||||
}
|
||||
}
|
||||
|
||||
// (c) write outside the lock; on any failure abort the sweep - a broken dir is a reason
|
||||
// not to prune, the write is retried on the next pass
|
||||
size_t n_written = 0;
|
||||
for (const auto & state : to_write) {
|
||||
if (!cache->save_entry(dir, *state)) {
|
||||
SRV_ERR("%s", "prompt cache sync: aborting the sweep after a write error, nothing pruned\n");
|
||||
return;
|
||||
}
|
||||
synced.insert(state->content_hash);
|
||||
n_written++;
|
||||
}
|
||||
|
||||
if (n_written > 0) {
|
||||
SRV_INF("prompt cache sync: wrote %zu entr%s to '%s'\n",
|
||||
n_written, n_written == 1 ? "y" : "ies", dir.c_str());
|
||||
}
|
||||
|
||||
// (d) prune files that no longer belong to the cache
|
||||
// membership is re-checked under mu right before each unlink: the cache may have been
|
||||
// re-created in the meantime (model reload on wake), and the listing above is stale
|
||||
// only files this instance has seen as part of the mirror (loaded, or written by an
|
||||
// earlier sweep) may be deleted - a file that never loaded (fingerprint mismatch,
|
||||
// corrupt, foreign) is left alone, so e.g. one run with the wrong -c value cannot wipe
|
||||
// a good mirror
|
||||
size_t n_unknown = 0;
|
||||
for (const auto & hash : on_disk) {
|
||||
if (cache->contains(hash)) {
|
||||
// part of the mirror - remember it, so that a later eviction can prune it
|
||||
synced.insert(hash);
|
||||
continue;
|
||||
}
|
||||
if (!synced.count(hash)) {
|
||||
n_unknown++;
|
||||
continue;
|
||||
}
|
||||
|
||||
const std::filesystem::path path = std::filesystem::u8path(dir) / (hash + ".bin");
|
||||
|
||||
std::error_code ec_rm;
|
||||
std::filesystem::remove(path, ec_rm);
|
||||
if (ec_rm) {
|
||||
SRV_WRN("prompt cache sync: failed to remove '%s': %s\n", path.string().c_str(), ec_rm.message().c_str());
|
||||
} else {
|
||||
SRV_TRC("prompt cache sync: pruned '%s'\n", path.string().c_str());
|
||||
}
|
||||
synced.erase(hash);
|
||||
}
|
||||
if (n_unknown > 0) {
|
||||
SRV_TRC("prompt cache sync: left %zu file(s) in '%s' untouched - not part of this instance's mirror\n",
|
||||
n_unknown, dir.c_str());
|
||||
}
|
||||
|
||||
// note: '.lock' and non-.bin files are never deleted, only the sweep's own .tmp files are
|
||||
// (handled inside save_entry())
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
#pragma once
|
||||
|
||||
#include <condition_variable>
|
||||
#include <cstdint>
|
||||
#include <mutex>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <thread>
|
||||
|
||||
struct server_prompt_cache;
|
||||
|
||||
// background mirror of the in-RAM prompt cache to a directory (see --prompt-cache-dir)
|
||||
//
|
||||
// one file per entry, named after the entry's content hash; the directory strictly mirrors
|
||||
// the cache: entries that have been sitting unchanged for --prompt-cache-sync-idle seconds
|
||||
// are written, files that no longer belong to the cache are pruned
|
||||
//
|
||||
// the thread only reads cache byte blobs and the filesystem - it never touches the
|
||||
// llama_context or the llama_model, so it is safe while requests are in flight
|
||||
//
|
||||
// the directory is serialized with a '.lock' file (flock): if another llama-server instance
|
||||
// holds it, persistence is disabled for this instance and the thread does not run, but the
|
||||
// server itself keeps working on the RAM cache
|
||||
struct server_prompt_cache_sync {
|
||||
server_prompt_cache_sync(server_prompt_cache * cache, std::string dir, int64_t sync_idle_ms);
|
||||
~server_prompt_cache_sync();
|
||||
|
||||
// create the dir, acquire the dir lock and start the thread
|
||||
// no-op if the thread is already running (e.g. wake from sleeping)
|
||||
// on lock failure: warning, no thread, flush() becomes a no-op
|
||||
void start();
|
||||
|
||||
// stop + join the thread, release the dir lock; safe to call when not running
|
||||
// note: the instance stays flushable afterwards - the final save at shutdown runs on the
|
||||
// main thread (see flush()), while the process still owns the dir until it exits
|
||||
void stop();
|
||||
|
||||
// one full sync pass on the calling thread: write every cache entry that has no file yet
|
||||
// (not just the idle ones), then prune stale files
|
||||
// never spawns threads; intended for shutdown and for entering sleeping state, with the
|
||||
// background thread stopped, so that it cannot race with the sweep
|
||||
// no-op when the dir lock was never acquired (another instance owns the dir)
|
||||
void flush();
|
||||
|
||||
// did this instance acquire the dir lock at least once?
|
||||
bool owns_dir() const {
|
||||
return owns_dir_;
|
||||
}
|
||||
|
||||
private:
|
||||
void loop();
|
||||
|
||||
// write the pending entries + prune; with idle_only, only entries older than sync_idle_ms
|
||||
void sweep(bool idle_only);
|
||||
|
||||
server_prompt_cache * cache;
|
||||
const std::string dir;
|
||||
const int64_t sync_idle_ms;
|
||||
|
||||
bool running = false;
|
||||
std::mutex wake_mu;
|
||||
std::condition_variable wake_cv;
|
||||
std::thread thread;
|
||||
|
||||
int lock_fd = -1;
|
||||
bool owns_dir_ = false;
|
||||
|
||||
// hashes this instance has seen as part of the mirror: loaded into the cache from the
|
||||
// dir, or written by a previous sweep. prune only deletes files in this set - a file
|
||||
// that never loaded (fingerprint mismatch, corrupt, foreign) is left alone, so one run
|
||||
// with the wrong config cannot wipe a good mirror
|
||||
std::set<std::string> synced;
|
||||
};
|
||||
+544
-325
File diff suppressed because it is too large
Load Diff
@@ -7,6 +7,8 @@
|
||||
#include <unordered_set>
|
||||
#include <list>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
|
||||
// TODO: prevent including the whole server-common.h as we only use server_tokens
|
||||
#include "server-common.h"
|
||||
@@ -682,6 +684,16 @@ struct server_prompt_cache_state {
|
||||
server_prompt prompt;
|
||||
server_prompt_data data;
|
||||
|
||||
// deterministic hash over the full entry payload (tokens + main + drft + checkpoints),
|
||||
// stamped by insert_state() once the blobs are fully filled; used as the entry's filename
|
||||
// in the persistent cache dir (see server_prompt_cache::save_entry())
|
||||
// empty = the payload could not be hashed, the entry will not be persisted
|
||||
std::string content_hash;
|
||||
|
||||
// time the entry was published to the cache (ggml_time_ms), used by the persistence sync
|
||||
// thread to skip entries that have not been idle long enough (see --prompt-cache-sync-idle)
|
||||
int64_t t_finalized_ms = 0;
|
||||
|
||||
size_t size() const {
|
||||
size_t res = data.size();
|
||||
|
||||
@@ -699,7 +711,17 @@ struct server_prompt_cache {
|
||||
this->limit_tokens = limit_tokens;
|
||||
}
|
||||
|
||||
std::list<server_prompt_cache_state> states;
|
||||
// entries are held as shared_ptr so that the persistence sync thread can keep an entry's
|
||||
// bytes alive while it writes them to disk, even if the entry is evicted or consumed
|
||||
// (stolen by load()) in the meantime
|
||||
std::list<std::shared_ptr<server_prompt_cache_state>> states;
|
||||
|
||||
// guards all surgery on and iteration of `states` (make_state, insert_state, load, update,
|
||||
// save_entry, contains) - these operations already run mutually exclusively (main loop or
|
||||
// the HTTP thread doing yield_to_queue work), the mutex's only job is to fence the
|
||||
// persistence sync thread while it snapshots entries; no file I/O and no llama_* calls
|
||||
// happen while the lock is held, and it is uncontended by design
|
||||
mutable std::mutex mu;
|
||||
|
||||
// in bytes, 0 = no limit
|
||||
size_t limit_size = 0;
|
||||
@@ -707,33 +729,58 @@ struct server_prompt_cache {
|
||||
// in tokens, 0 = no limit
|
||||
size_t limit_tokens = 0;
|
||||
|
||||
// note: these iterate `states` without locking; when the persistence sync thread may be
|
||||
// running, call them while holding `mu`
|
||||
size_t size() const;
|
||||
|
||||
size_t n_tokens() const;
|
||||
|
||||
server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
|
||||
// two-phase publish: make_state() runs the duplicate / contained / limit checks and returns
|
||||
// a new entry, NOT yet in the list, with allocated but not yet filled blobs; the caller
|
||||
// fills data.main / data.drft and then publishes it with insert_state(), which stamps
|
||||
// content_hash + t_finalized_ms and appends to `states` under mu
|
||||
// an entry is never visible to load() or to the persistence sync thread before it is fully
|
||||
// filled, so a published entry's payload is immutable for its whole lifetime
|
||||
// make_state() returns nullptr if the prompt is already in the cache or does not fit
|
||||
std::shared_ptr<server_prompt_cache_state> make_state(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
|
||||
|
||||
void insert_state(std::shared_ptr<server_prompt_cache_state> state);
|
||||
|
||||
bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot);
|
||||
|
||||
void update();
|
||||
|
||||
// fingerprint of the loaded model, written to and checked against the on-disk cache file (see save_to_file())
|
||||
// same as update(), for callers already holding mu
|
||||
void update_locked();
|
||||
|
||||
// is there a live entry with this content hash? (used by the persistence sync pruning)
|
||||
bool contains(const std::string & content_hash) const;
|
||||
|
||||
// fingerprint of the loaded model, written to and checked against each on-disk entry file
|
||||
// (see save_entry() / load_from_dir())
|
||||
int32_t fp_n_ctx = 0;
|
||||
int32_t fp_n_embd = 0;
|
||||
int32_t fp_n_layer = 0;
|
||||
int32_t fp_n_head = 0;
|
||||
uint32_t fp_n_vocab = 0;
|
||||
std::string fp_model_path;
|
||||
std::string fp_build; // build string (llama_build_info()), so entries from another build are skipped
|
||||
|
||||
// save all states to a file (atomic: write to a .tmp file, fsync, rename over the target)
|
||||
// returns false on I/O errors
|
||||
bool save_to_file(const std::string & path) const;
|
||||
// write one entry as '<dir>/<content_hash>.bin' (atomic: write a .tmp file, fsync, rename)
|
||||
// an existing file is skipped: the name is the payload's hash, same hash = same bytes
|
||||
// the payload is serialized under mu, the file I/O happens outside of it
|
||||
// returns false on I/O or serialization errors
|
||||
bool save_entry(const std::string & dir, const server_prompt_cache_state & state) const;
|
||||
|
||||
// load states from a file, returns the number of entries loaded
|
||||
// a fingerprint mismatch (model or context) skips the whole file
|
||||
// build the full file image (header + payload) for one entry; the caller must hold mu
|
||||
std::vector<uint8_t> serialize_entry_locked(const server_prompt_cache_state & state) const;
|
||||
|
||||
// load all '<hash>.bin' entry files from dir, returns the number of entries loaded
|
||||
// each file is validated on its own (magic / version / size / fingerprint), invalid files
|
||||
// are skipped with a warning naming the file, the others still load
|
||||
// if !has_draft, the draft states of the loaded entries are dropped
|
||||
// if !has_mtmd, entries with media tokens are skipped
|
||||
size_t load_from_file(const std::string & path, bool has_draft = true, bool has_mtmd = false);
|
||||
size_t load_from_dir(const std::string & dir, bool has_draft = true, bool has_mtmd = false);
|
||||
};
|
||||
|
||||
// used exclusively by router mode
|
||||
|
||||
@@ -489,7 +489,11 @@ int llama_server(common_params & params, int argc, char ** argv, server_child &
|
||||
// stop the session GC first, it finalizes live sessions and wakes pending readers
|
||||
server_stream_session_manager_stop();
|
||||
ctx_http.stop();
|
||||
// persist the prompt cache before terminating, if --prompt-cache-file is set
|
||||
// stop the prompt cache sync thread first - the final save below runs on this thread
|
||||
ctx_server.prompt_cache_sync_stop();
|
||||
// persist the prompt cache before terminating, if --prompt-cache-dir is set
|
||||
// note: the sync thread is started by load_model(), which runs after ctx_http.start()
|
||||
// succeeds below, so it cannot exist when clean_up runs on the HTTP-failure path
|
||||
ctx_server.prompt_cache_save();
|
||||
ctx_server.terminate();
|
||||
mcp_mgr.shutdown();
|
||||
|
||||
Reference in New Issue
Block a user