From 971962df4606db3b594c9a552218998787fb6a4c Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Fri, 23 Oct 2026 22:55:22 -0500 Subject: [PATCH 1/3] server: a slot pages to named host RAM and back; checkpoints keep a byte budget; a page carries its draft KV A mind switch costs a memcpy instead of a 66-100 s disk save or a 70-160 s re-prefill (continuum SERVING-ENGINE.md, slice S1, plan line 3b). - /slots/{id}?action=page_out&page=: the slot's target KV, draft (MTP) KV, speculative tail state and context checkpoints move into a server-owned RAM page store, and the slot is freed. page_in puts them into any slot with its tokens, so the next request's prefix matches; the page leaves the store. spill writes a page in the slot-save format (.bin + .ckpt + .lora + new .dft) on the HTTP thread, off the engine thread; "drop": true then frees that page's generation only. page_drop frees a page. page_out/page_in defer while the slot processes, as restore does. Errors are typed (404 unknown page, 503 over --page-ram-bytes, 400 empty slot); a failed page_out leaves the slot intact, a failed page_in leaves the page in the store. /props slot_page_version 1. - save writes the draft state to `.dft` and counts it in its quote; restore reads it, so MTP is warm after a restore (no `.dft`: the draft seq is cleared and the cold start logged). - --ctx-checkpoints-bytes: one budget across all slots in place of the per-slot count; the slot holding the most bytes drops its oldest checkpoint after its first. Every eviction logs and counts. - /metrics: page_store_bytes/pages, ctx_checkpoint_bytes, page_{out,in,failed}_total, ctx_checkpoints_{evicted,skipped}_budget_total. - llama_state_seq_save_file_from_data (llama-ext.h): libllama stays the only owner of the sequence-file format; the sidecar writers move from the save arm into server-pages.cpp, shared by save and spill. Measured (Qwen3.5 0.8B hybrid, M5 Metal, 19,794 tokens, 2 checkpoints, median of 5): page_out 10.4 ms, page_in 7.8 ms, 304 MB; save 134 ms, restore 37 ms (page-cache warm), 304 MB on disk; spill 101 ms. The paged slot's next completion processed 6 prompt tokens and its output matched a slot that never paged. Tests: tools/server/tests/unit/test_slot_page.py (5, hybrid case via LLAMA_SERVER_TEST_HYBRID_GGUF) + test_slot_save.py (11): 16 passed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- common/arg.cpp | 29 ++ common/common.h | 2 + src/llama-context.cpp | 30 ++ src/llama-ext.h | 16 + tools/server/CMakeLists.txt | 4 + tools/server/server-checkpoints.cpp | 42 +++ tools/server/server-checkpoints.h | 28 ++ tools/server/server-common.h | 7 + tools/server/server-context.cpp | 391 +++++++++++++++++++--- tools/server/server-context.h | 1 + tools/server/server-pages.cpp | 301 +++++++++++++++++ tools/server/server-pages.h | 111 ++++++ tools/server/server-task.cpp | 59 ++++ tools/server/server-task.h | 37 ++ tools/server/tests/unit/test_slot_page.py | 171 ++++++++++ tools/server/tests/utils.py | 6 + 16 files changed, 1184 insertions(+), 51 deletions(-) create mode 100644 tools/server/server-checkpoints.cpp create mode 100644 tools/server/server-checkpoints.h create mode 100644 tools/server/server-pages.cpp create mode 100644 tools/server/server-pages.h create mode 100644 tools/server/tests/unit/test_slot_page.py diff --git a/common/arg.cpp b/common/arg.cpp index 7ce7eaad16ab..04f677503828 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -330,6 +330,18 @@ static std::string get_all_kv_cache_types() { return msg.str(); } +// an exact byte count: the caller passes a measured number, so no unit suffixes are guessed +static uint64_t parse_byte_count(const std::string & value, const char * flag) { + if (value.empty() || value.size() > 20 || value.find_first_not_of("0123456789") != std::string::npos) { + throw std::invalid_argument(std::string(flag) + " takes an unsigned byte count, got: " + value); + } + try { + return std::stoull(value); + } catch (const std::exception &) { + throw std::invalid_argument(std::string(flag) + " is out of range: " + value); + } +} + static bool parse_bool_value(const std::string & value) { if (is_truthy(value)) { return true; @@ -1691,6 +1703,23 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.n_ctx_checkpoints = value; } ).set_env("LLAMA_ARG_CTX_CHECKPOINTS").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--ctx-checkpoints-bytes"}, "BYTES", + "byte budget for the context checkpoints of all slots together, in place of the per-slot count " + "(default: 0 = use --ctx-checkpoints); when a new checkpoint would exceed it, the slot holding the " + "most checkpoint bytes drops its oldest checkpoint after its first, so each slot keeps its head and its newest", + [](common_params & params, const std::string & value) { + params.ctx_checkpoints_bytes = parse_byte_count(value, "--ctx-checkpoints-bytes"); + } + ).set_env("LLAMA_ARG_CTX_CHECKPOINTS_BYTES").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--page-ram-bytes"}, "BYTES", + "ceiling of the host-RAM page store that /slots page_out fills (default: 0 = no ceiling); " + "a page_out that would exceed it is refused, so the caller spills a page to disk first", + [](common_params & params, const std::string & value) { + params.page_ram_bytes = parse_byte_count(value, "--page-ram-bytes"); + } + ).set_env("LLAMA_ARG_PAGE_RAM_BYTES").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-cms", "--checkpoint-min-step"}, "N", string_format("minimum spacing between context checkpoints in tokens (default: %d, 0 = no minimum)", params.checkpoint_min_step), diff --git a/common/common.h b/common/common.h index 29b171b7dda0..9a2e6824fdac 100644 --- a/common/common.h +++ b/common/common.h @@ -621,6 +621,8 @@ struct common_params { bool cache_idle_slots = true; // save and clear idle slots upon starting a new task int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints + uint64_t ctx_checkpoints_bytes = 0; // checkpoint byte budget across all slots; 0 = the count above governs + uint64_t page_ram_bytes = 0; // ceiling of the named host-RAM page store; 0 = no ceiling int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. std::string hostname = "127.0.0.1"; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 654d386b63f3..86818e5bfe4c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3995,6 +3995,36 @@ size_t llama_state_seq_save_file(llama_context * ctx, const char * filepath, lla } } +size_t llama_state_seq_save_file_from_data(const char * filepath, const uint8_t * data, size_t size, const llama_token * tokens, size_t n_token_count) { + // get_data writes [io_magic][seq_id][state]; the file is [magic][version][count][tokens][state] + const size_t prefix = sizeof(io_magic) + sizeof(llama_seq_id); + if (data == nullptr || size <= prefix) { + LLAMA_LOG_ERROR("%s: no sequence state to write\n", __func__); + return 0; + } + uint32_t magic_read = 0; + std::memcpy(&magic_read, data, sizeof(magic_read)); + if (magic_read != io_magic) { + LLAMA_LOG_ERROR("%s: the bytes are not a sequence state (magic %08x)\n", __func__, magic_read); + return 0; + } + try { + llama_file file(filepath, "wb"); + const uint32_t magic = LLAMA_STATE_SEQ_MAGIC; + const uint32_t version = LLAMA_STATE_SEQ_VERSION; + const uint32_t count = (uint32_t) n_token_count; + file.write_raw(&magic, sizeof(magic)); + file.write_raw(&version, sizeof(version)); + file.write_raw(&count, sizeof(count)); + file.write_raw(tokens, sizeof(llama_token) * n_token_count); + file.write_raw(data + prefix, size - prefix); + return file.tell(); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error writing sequence state file: %s\n", __func__, err.what()); + return 0; + } +} + size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { return llama_state_seq_load_file_with_progress(ctx, filepath, dest_seq_id, tokens_out, n_token_capacity, n_token_count_out, nullptr, nullptr); } diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa8..9303d2cf0d0e 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -130,3 +130,19 @@ LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_mode // if out is nullptr, returns the number of tokens without writing to out // caller must allocate enough memory for out before calling LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out); + +// +// sequence state pages (Continuum fork) +// + +// Writes a sequence-state file (the llama_state_seq_save_file format) from bytes that +// llama_state_seq_get_data_ext returned with LLAMA_STATE_SEQ_FLAGS_NONE, without a context: +// a page held in host RAM spills to disk off the engine thread, and llama_state_seq_load_file +// reads it back. Returns the bytes written, or 0 when the bytes are not a sequence state or +// the write fails. +LLAMA_API size_t llama_state_seq_save_file_from_data( + const char * filepath, + const uint8_t * data, + size_t size, + const llama_token * tokens, + size_t n_token_count); diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 141d44126731..2e1c2ed087ca 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -16,6 +16,10 @@ add_library(${TARGET} STATIC server-context.cpp server-train.cpp server-context.h + server-pages.cpp + server-pages.h + server-checkpoints.cpp + server-checkpoints.h server-stream.cpp server-stream.h server-tools.cpp diff --git a/tools/server/server-checkpoints.cpp b/tools/server/server-checkpoints.cpp new file mode 100644 index 000000000000..efcce5eb409c --- /dev/null +++ b/tools/server/server-checkpoints.cpp @@ -0,0 +1,42 @@ +#include "server-checkpoints.h" + +#include + +size_t server_checkpoints_bytes(const server_checkpoint_list & list) { + size_t n = 0; + for (const auto & c : list) { + n += c.size(); + } + return n; +} + +bool server_checkpoints_fit( + const std::vector & lists, + size_t own, + uint64_t budget, + size_t incoming, + const std::function & on_evict) { + if (incoming > budget) { + return false; + } + for (;;) { + uint64_t total = incoming; + size_t victim = own; + size_t victim_bytes = 0; + for (size_t i = 0; i < lists.size(); ++i) { + const size_t n = server_checkpoints_bytes(*lists[i]); + total += n; + if (n > victim_bytes || (n == victim_bytes && i == own)) { + victim = i; + victim_bytes = n; + } + } + if (total <= budget || victim_bytes == 0) { + return true; + } + auto & list = *lists[victim]; + auto it = list.size() >= 2 ? std::next(list.begin()) : list.begin(); + on_evict(victim, *it); + list.erase(it); + } +} diff --git a/tools/server/server-checkpoints.h b/tools/server/server-checkpoints.h new file mode 100644 index 000000000000..1ab7750087a6 --- /dev/null +++ b/tools/server/server-checkpoints.h @@ -0,0 +1,28 @@ +#pragma once + +// Continuum fork: context checkpoints kept within one byte budget across all slots +// (--ctx-checkpoints-bytes), in place of a per-slot count. On a hybrid 27B a checkpoint is +// ~420 MB of host RAM, which on unified memory is the GPU's RAM, so the count is the wrong unit. + +#include "common.h" + +#include +#include +#include +#include + +using server_checkpoint_list = std::list; + +size_t server_checkpoints_bytes(const server_checkpoint_list & list); + +// Makes room for a new checkpoint of `incoming` bytes. While the lists' total plus `incoming` +// exceeds `budget`, the list holding the most bytes (the creator's on a tie) drops its oldest +// checkpoint after its first, so each slot keeps its head (the start of its first user turn) and +// its newest longest, as a saved page's subset does. `on_evict(list index, checkpoint)` runs before +// each drop. Returns false, dropping nothing, when the new checkpoint alone exceeds the budget. +bool server_checkpoints_fit( + const std::vector & lists, + size_t own, + uint64_t budget, + size_t incoming, + const std::function & on_evict); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index c55d5b79be80..81bf294e4bc1 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -482,6 +482,13 @@ struct server_metrics { uint64_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model std::vector n_accepted_per_pos; // Accepted tokens per draft position + // Continuum fork: RAM pages and the checkpoint byte budget + uint64_t n_page_out = 0; + uint64_t n_page_in = 0; + uint64_t n_page_failed = 0; // page_out, page_in or spill refused or failed + uint64_t n_ckpt_evicted_budget = 0; // checkpoints dropped to keep within --ctx-checkpoints-bytes + uint64_t n_ckpt_skipped_budget = 0; // checkpoints not taken: one alone exceeds the budget + void init() { t_start = ggml_time_us(); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 784d5426eeba..d121ebcb2f5c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2,6 +2,8 @@ #include "server-context.h" #include "../../src/llama-ext.h" #include "server-train.h" +#include "server-checkpoints.h" +#include "server-pages.h" #include "server-chat.h" #include "server-common.h" #include "server-http.h" @@ -31,22 +33,6 @@ #include #include -// Continuum fork: an adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it -// cannot be read), stored in a saved slot's `.lora` sidecar so a restore refuses a page computed -// under different weights at the same path. -static std::pair lora_file_revision(const std::string & path) { - std::error_code ec; - const auto size = std::filesystem::file_size(path, ec); - if (ec) { - return {0, 0}; - } - const auto mtime = std::filesystem::last_write_time(path, ec); - if (ec) { - return {0, 0}; - } - return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; -} - // fix problem with std::min and std::max #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN @@ -938,6 +924,9 @@ struct server_context_impl { std::unique_ptr prompt_cache; + // Continuum fork: named slot states in host RAM (/slots page_out, page_in, spill, page_drop) + server_page_store page_store; + server_metrics metrics; // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync @@ -960,6 +949,12 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { + // a page names adapters by pointer, and the model's adapters are freed with it + if (page_store.n_pages() > 0) { + SRV_WRN("dropping %zu RAM pages (%zu bytes) with the model: spill them before sleep to keep them\n", + page_store.n_pages(), page_store.n_bytes()); + page_store = server_page_store{}; + } spec.reset(); spec_init.reset(); @@ -1352,7 +1347,12 @@ struct server_context_impl { } SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); - if (params_base.n_ctx_checkpoints > 0) { + page_store.limit_bytes = params_base.page_ram_bytes; + + if (params_base.n_ctx_checkpoints > 0 && params_base.ctx_checkpoints_bytes > 0) { + SRV_INF("context checkpoints enabled, budget = %llu bytes across all slots, min spacing = %d\n", + (unsigned long long) params_base.ctx_checkpoints_bytes, params_base.checkpoint_min_step); + } else if (params_base.n_ctx_checkpoints > 0) { SRV_TRC("context checkpoints enabled, max = %d, min spacing = %d\n", params_base.n_ctx_checkpoints, params_base.checkpoint_min_step); } else { @@ -2472,7 +2472,34 @@ struct server_context_impl { ++it; } - while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { + // Continuum fork: a byte budget across all slots replaces the per-slot count when set + if (params_base.ctx_checkpoints_bytes > 0) { + const size_t incoming = + llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) + + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) : 0); + std::vector lists; + size_t own = 0; + for (auto & s : slots) { + if (&s == &slot) { + own = lists.size(); + } + lists.push_back(&s.prompt.checkpoints); + } + const bool fits = server_checkpoints_fit(lists, own, params_base.ctx_checkpoints_bytes, incoming, + [&](size_t i, const common_prompt_checkpoint & c) { + metrics.n_ckpt_evicted_budget++; + SLT_WRN(slots[i], "checkpoint budget: erasing context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB) for slot %d's new one\n", + c.pos_min, c.pos_max, c.n_tokens, (float) c.size() / 1024 / 1024, slot.id); + }); + if (!fits) { + metrics.n_ckpt_skipped_budget++; + SLT_WRN(slot, "checkpoint budget: a checkpoint of %.3f MiB alone exceeds --ctx-checkpoints-bytes %llu, not taken\n", + (float) incoming / 1024 / 1024, (unsigned long long) params_base.ctx_checkpoints_bytes); + return; + } + } + + while (params_base.ctx_checkpoints_bytes == 0 && slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { // make room for the new checkpoint, if needed const auto & cur = slot.prompt.checkpoints.front(); @@ -2502,6 +2529,103 @@ struct server_context_impl { cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); } + // Continuum fork: one PAGE_* task on the task thread. Throws server_page_error: a failed + // page_out leaves the slot and the store as they were; a failed page_in leaves the page in + // the store and the slot empty, never half-restored. + std::unique_ptr page_task(const server_task & task, server_slot * slot) { + const int64_t t_start = ggml_time_us(); + const std::string & name = task.slot_action.page; + + auto res = std::make_unique(); + res->id = task.id; + res->id_slot = task.slot_action.id_slot; + res->page = name; + + server_page_ptr page; + switch (task.type) { + case SERVER_TASK_TYPE_PAGE_OUT: + { + res->action = "page_out"; + if (slot->prompt.tokens.size() == 0) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the slot holds no sequence to page out"); + } + const auto [n_tgt, n_dft] = server_page_state_size(ctx_tgt, ctx_dft, slot->id); + page_store.admit(name, n_tgt + n_dft + server_checkpoints_bytes(slot->prompt.checkpoints)); + page = server_page_capture(ctx_tgt, ctx_dft, spec.get(), slot->id); + page->prompt = std::move(slot->prompt); + page->lora = slot->lora; + slot->prompt_clear(); + res->gen = page_store.put(name, page); + metrics.n_page_out++; + } break; + case SERVER_TASK_TYPE_PAGE_IN: + { + res->action = "page_in"; + page = page_store.find(name); + if (!page) { + throw server_page_error(ERROR_TYPE_NOT_FOUND, "no RAM page named '" + name + "'"); + } + if (page->prompt.tokens.size() > (size_t) slot->n_ctx) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the page does not fit the slot context"); + } + slot->prompt_clear(); + try { + server_page_restore(*page, ctx_tgt, ctx_dft, spec.get(), slot->id); + } catch (...) { + slot->prompt_clear(); + throw; + } + // the page leaves the store; its buffers are copied only while a spill still reads them + const bool sole = page.use_count() == 2; // the store's reference and this one + slot->prompt = sole ? std::move(page->prompt) : page->prompt.clone(); + slot->lora = page->lora; + res->gen = page->gen; + page_store.erase(name, page->gen); + metrics.n_page_in++; + } break; + case SERVER_TASK_TYPE_PAGE_SPILL: + { + res->action = "spill"; + page = page_store.find(name); + if (!page) { + throw server_page_error(ERROR_TYPE_NOT_FOUND, "no RAM page named '" + name + "'"); + } + res->gen = page->gen; + res->spilled = page; + } break; + case SERVER_TASK_TYPE_PAGE_DROP: + { + res->action = "page_drop"; + page = page_store.find(name); + if (!page) { + throw server_page_error(ERROR_TYPE_NOT_FOUND, "no RAM page named '" + name + "'"); + } + if (page_store.erase(name, task.slot_action.page_gen) == 0) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the page under this name was replaced; nothing dropped"); + } + res->gen = page->gen; + } break; + default: + GGML_ABORT("not a page task"); + } + + // a page_in moved its prompt into the slot; report what the slot now holds + const server_prompt & held = task.type == SERVER_TASK_TYPE_PAGE_IN ? slot->prompt : page->prompt; + res->n_tokens = held.tokens.size(); + res->n_checkpoints = held.checkpoints.size(); + res->n_bytes_checkpoints = server_checkpoints_bytes(held.checkpoints); + res->n_bytes_draft = page->data.drft.size(); + res->n_bytes = page->data.size() + page->spec.size() + res->n_bytes_checkpoints; + res->store_bytes = page_store.n_bytes(); + res->store_pages = page_store.n_pages(); + res->t_ms = (ggml_time_us() - t_start) / 1000.0; + + SRV_INF("%s '%s' slot %d: %zu tokens, %zu bytes (%zu draft, %zu in %zu checkpoints) in %.3f ms; store %zu bytes in %zu pages\n", + res->action.c_str(), name.c_str(), res->id_slot, res->n_tokens, res->n_bytes, res->n_bytes_draft, + res->n_bytes_checkpoints, res->n_checkpoints, res->t_ms, res->store_bytes, res->store_pages); + return res; + } + // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { // while yielding, an encode / decode is running and only reading the server state is safe @@ -2645,6 +2769,11 @@ struct server_context_impl { res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); res->metrics = metrics; + res->page_store_bytes = page_store.n_bytes(); + res->page_store_pages = page_store.n_pages(); + for (const server_slot & slot : slots) { + res->ctx_checkpoint_bytes += server_checkpoints_bytes(slot.prompt.checkpoints); + } if (task.metrics_reset_bucket) { metrics.reset_bucket(); @@ -2729,39 +2858,26 @@ struct server_context_impl { ++i; } } - auto write_checkpoints = [&](auto && put) { - const uint32_t magic = 0x434B5054u; // 'CKPT' - const uint32_t version = 1; - const uint32_t n_ckpt = (uint32_t) kept_checkpoints.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); - for (const auto * kept : kept_checkpoints) { - const auto & c = *kept; - const int64_t n_tokens = c.n_tokens; - const int32_t pos_min = c.pos_min; - const int32_t pos_max = c.pos_max; - const uint64_t sz_tgt = c.data_tgt.size(); - const uint64_t sz_dft = c.data_dft.size(); - const uint64_t sz_spec = c.data_spec.size(); - put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); - put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); - put(c.data_tgt.data(), c.data_tgt.size()); - put(c.data_dft.data(), c.data_dft.size()); - put(c.data_spec.data(), c.data_spec.size()); - } - }; - auto write_adapters = [&](auto && put) { - const uint32_t magic = 0x4C4F5241u; // 'LORA' - const uint32_t version = 2; - const uint32_t n_lora = (uint32_t) slot->lora.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); - for (const auto & l : slot->lora) { - const uint32_t n_path = (uint32_t) l.path.size(); - const auto rev = lora_file_revision(l.path); - put(&n_path, sizeof(n_path)); - put(l.path.data(), n_path); - put(&l.scale, sizeof(l.scale)); - put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); + // Continuum fork: the draft (MTP) state rides in a `.dft` sidecar, so a restored + // page drafts warm; a server without a draft context writes none. + // The quote counts the draft bytes without copying them. + std::vector draft_state; + std::vector draft_spec; + const size_t n_draft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (ctx_dft) { + common_speculative_get_state(spec.get(), slot->id, draft_spec); + } + if (ctx_dft && !task.slot_action.size_only) { + draft_state.resize(n_draft); + if (llama_state_seq_get_data_ext(ctx_dft, draft_state.data(), n_draft, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) != n_draft) { + send_error(task, "Unable to read the draft state", ERROR_TYPE_SERVER); + break; } + } + auto write_checkpoints = [&](const server_page_put & put) { server_page_write_checkpoints(kept_checkpoints, put); }; + auto write_adapters = [&](const server_page_put & put) { server_page_write_adapters(slot->lora, put); }; + auto write_draft = [&](const server_page_put & put) { + server_page_write_draft(draft_state.empty() ? nullptr : draft_state.data(), n_draft, draft_spec, put); }; if (task.slot_action.size_only || task.slot_action.has_byte_ceiling) { size_t projected_bytes = llama_state_seq_file_size( @@ -2781,6 +2897,9 @@ struct server_context_impl { }; write_checkpoints(count); write_adapters(count); + if (ctx_dft) { + write_draft(count); + } if (size_overflow) { send_error(task, "Slot file size overflow", ERROR_TYPE_SERVER); break; @@ -2824,6 +2943,17 @@ struct server_context_impl { SLT_WRN(*slot, "adapter sidecar could not be written: %s (this page will not restore)\n", filepath.c_str()); } } + if (ctx_dft) { + std::ofstream out(filepath + ".dft", std::ios::binary | std::ios::trunc); + write_draft([&](const void * ptr, size_t n) { out.write(reinterpret_cast(ptr), n); }); + if (!out) { + SLT_WRN(*slot, "draft sidecar could not be written: %s (the restore will draft cold)\n", filepath.c_str()); + } + } else { + // a draft state saved under this name by an earlier server is not this page's + std::error_code ec; + std::filesystem::remove(filepath + ".dft", ec); + } const int64_t t_end = ggml_time_us(); const double t_save_ms = (t_end - t_start) / 1000.0; @@ -2985,6 +3115,23 @@ struct server_context_impl { SLT_WRN(*slot, "no usable checkpoint sidecar (%s): the restore is cold — the next divergence re-prefills from the head\n", ckpt_path.c_str()); } } + // Continuum fork: the draft (MTP) state from the `.dft` sidecar, so drafting + // resumes warm. A page saved before the sidecar existed drafts cold, as it did. + if (ctx_dft) { + std::vector draft_state; + std::vector draft_spec; + std::string why; + const std::string dft_path = filepath + ".dft"; + if (!server_page_read_draft(dft_path, draft_state, draft_spec, why)) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, -1, -1); + SLT_WRN(*slot, "draft sidecar %s (%s): the draft starts cold\n", why.c_str(), dft_path.c_str()); + } else if (llama_state_seq_set_data_ext(ctx_dft, draft_state.data(), draft_state.size(), slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) != draft_state.size()) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, -1, -1); + SLT_WRN(*slot, "draft sidecar did not fit the draft context (%s): the draft starts cold\n", dft_path.c_str()); + } else { + common_speculative_set_state(spec.get(), slot->id, draft_spec); + } + } if (progress && progress->cancel_requested.load(std::memory_order_relaxed)) { throw std::runtime_error("Restore operation cancelled"); } @@ -3036,6 +3183,32 @@ struct server_context_impl { res->n_erased = n_erased; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_PAGE_OUT: + case SERVER_TASK_TYPE_PAGE_IN: + case SERVER_TASK_TYPE_PAGE_SPILL: + case SERVER_TASK_TYPE_PAGE_DROP: + { + // Continuum fork: named host-RAM pages (server-pages.h). page_out and page_in + // mutate the slot's KV, so they defer while it is processing, as restore does. + const bool on_slot = task.type == SERVER_TASK_TYPE_PAGE_OUT || task.type == SERVER_TASK_TYPE_PAGE_IN; + server_slot * slot = on_slot ? get_slot_by_id(task.slot_action.id_slot) : nullptr; + if (on_slot && slot == nullptr) { + send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); + break; + } + if (on_slot && slot->is_processing()) { + SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); + queue_tasks.defer(std::move(task)); + break; + } + try { + auto res = page_task(task, slot); + queue_results.send(std::move(res)); + } catch (const server_page_error & err) { + metrics.n_page_failed += task.type == SERVER_TASK_TYPE_PAGE_DROP ? 0 : 1; + send_error(task, err.what(), err.type); + } + } break; case SERVER_TASK_TYPE_GET_LORA: { // TODO @ngxson : make lora_adapters a dedicated member of server_context @@ -3195,6 +3368,7 @@ struct server_context_impl { if (prompt_cache) { prompt_cache->evict_lora(ptr); } + page_store.evict_lora(ptr); // the context may still hold it from the last batch; clear before freeing. // Only ctx_tgt ever receives adapters (common_set_adapter_lora in // update_slots); the draft/MTP context never does. @@ -5213,6 +5387,7 @@ static json get_res_props(const server_context_meta & meta, const common_params { "default_generation_settings", default_generation_settings_for_props }, { "total_slots", params.n_parallel }, { "slot_save_size_version", 1 }, + { "slot_page_version", 1 }, // Continuum fork: /slots page_out, page_in, spill, page_drop { "expert_container_custody_version", 1 }, { "model_alias", meta.model_name }, { "model_ftype", meta.model_ftype }, @@ -5388,7 +5563,9 @@ void server_routes::init_routes() { this->post_slots = [this](const server_http_req & req) { auto res = create_response(); - if (params.slot_save_path.empty()) { + const std::string action_name = req.get_param("action"); + const bool ram_only = action_name == "page_out" || action_name == "page_in" || action_name == "page_drop"; + if (params.slot_save_path.empty() && !ram_only) { res->error(format_error_response("This server does not support slots action. Start it with `--slot-save-path`", ERROR_TYPE_NOT_SUPPORTED)); return res; } @@ -5405,6 +5582,9 @@ void server_routes::init_routes() { std::string action = req.get_param("action"); + if (action == "page_out" || action == "page_in" || action == "spill" || action == "page_drop") { + return handle_slots_page(req, id_slot, action); + } if (action == "save" || action == "save-size") { return handle_slots_save(req, id_slot, action == "save-size"); } @@ -6293,6 +6473,115 @@ std::unique_ptr server_routes::handle_slots_restore(const return res; } +// Continuum fork: the RAM page verbs (server-pages.h). +// POST /slots/{id}?action=page_out&page= slot -> RAM page; the slot is freed +// POST /slots/{id}?action=page_in&page= RAM page -> slot; the page leaves the store +// POST /slots/{id}?action=spill&page= body {"filename": f, "drop": bool}: RAM page -> f (+.ckpt .lora .dft) +// POST /slots/{id}?action=page_drop&page= frees the page +// {id} is not used by spill and page_drop. The page name may also come as "page" in the body. +std::unique_ptr server_routes::handle_slots_page(const server_http_req & req, int id_slot, const std::string & action) { + auto res = create_response(); + json request_data = json::object(); + if (!req.body.empty()) { + try { + request_data = json::parse(req.body); + } catch (const std::exception &) { + res->error(format_error_response("The body must be JSON", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + } + std::string page = req.get_param("page"); + if (page.empty() && request_data.contains("page") && request_data.at("page").is_string()) { + page = request_data.at("page").get(); + } + if (!server_page_name_valid(page)) { + res->error(format_error_response("A page name is 1-255 bytes without control characters", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + const bool on_slot = action == "page_out" || action == "page_in"; + if (on_slot && (id_slot < 0 || id_slot >= params.n_parallel)) { + res->error(format_error_response("Invalid slot ID", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + std::string filename; + bool drop = false; + if (action == "spill") { + filename = request_data.value("filename", std::string()); + if (!fs_validate_filename(filename)) { + res->error(format_error_response("Invalid filename", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + if (request_data.contains("drop") && !request_data.at("drop").is_boolean()) { + res->error(format_error_response("drop must be a boolean", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + drop = request_data.value("drop", false); + } + + // a reader posts once, so the spill's drop gets its own + server_response_reader rd_drop(queue_tasks, queue_results, HTTP_POLLING_SECONDS); + auto run = [&](server_response_reader & rd, server_task_type type, uint64_t gen) -> server_task_result_ptr { + server_task task(type); + task.id = rd.get_new_id(); + // bound to its slot, so a deferred page_out/page_in is serviced when that slot releases + task.id_slot = on_slot ? id_slot : -1; + task.slot_action.id_slot = id_slot; + task.slot_action.page = page; + task.slot_action.page_gen = gen; + rd.post_task(std::move(task)); + return rd.next(req.should_stop); + }; + + const server_task_type type = + action == "page_out" ? SERVER_TASK_TYPE_PAGE_OUT : + action == "page_in" ? SERVER_TASK_TYPE_PAGE_IN : + action == "spill" ? SERVER_TASK_TYPE_PAGE_SPILL : SERVER_TASK_TYPE_PAGE_DROP; + auto result = run(res->rd, type, 0); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + auto * page_res = dynamic_cast(result.get()); + GGML_ASSERT(page_res != nullptr); + + if (action == "spill") { + // the disk write runs here, on bytes already in RAM, so the engine never waits on it + const int64_t t_start = ggml_time_us(); + try { + page_res->n_written = server_page_spill(*page_res->spilled, params.slot_save_path + filename); + } catch (const server_page_error & err) { + res->error(format_error_response(err.what(), err.type)); + return res; + } + page_res->filename = filename; + page_res->t_ms = (ggml_time_us() - t_start) / 1000.0; + const uint64_t gen = page_res->gen; + page_res->spilled.reset(); + if (drop) { + // drops this generation only: a page_out under the same name since then is kept + auto dropped = run(rd_drop, SERVER_TASK_TYPE_PAGE_DROP, gen); + if (!dropped) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (dropped->is_error()) { + res->error(dropped->to_json()); + return res; + } + auto * d = dynamic_cast(dropped.get()); + page_res->store_bytes = d->store_bytes; + page_res->store_pages = d->store_pages; + } + } + + res->ok(page_res->to_json()); + return res; +} + std::unique_ptr server_routes::handle_slots_erase(const server_http_req & req, int id_slot) { auto res = create_response(); auto & rd = res->rd; diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 6deb3a13cc06..38efbb264491 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -190,6 +190,7 @@ struct server_routes { std::unique_ptr handle_slots_save(const server_http_req & req, int id_slot, bool size_only); std::unique_ptr handle_slots_restore(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_erase(const server_http_req &, int id_slot); + std::unique_ptr handle_slots_page(const server_http_req & req, int id_slot, const std::string & action); std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); diff --git a/tools/server/server-pages.cpp b/tools/server/server-pages.cpp new file mode 100644 index 000000000000..9a8ff43474e0 --- /dev/null +++ b/tools/server/server-pages.cpp @@ -0,0 +1,301 @@ +#include "server-pages.h" + +#include "../../src/llama-ext.h" +#include "speculative.h" + +#include +#include +#include +#include + +static constexpr uint32_t PAGE_CKPT_MAGIC = 0x434B5054u; // 'CKPT' +static constexpr uint32_t PAGE_LORA_MAGIC = 0x4C4F5241u; // 'LORA' +static constexpr uint32_t PAGE_DRFT_MAGIC = 0x44524654u; // 'DRFT' + +size_t server_page::n_checkpoint_bytes() const { + size_t n = 0; + for (const auto & c : prompt.checkpoints) { + n += c.size(); + } + return n; +} + +size_t server_page::n_bytes() const { + return data.size() + spec.size() + n_checkpoint_bytes(); +} + +std::pair server_page_state_size(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq) { + return { + llama_state_seq_get_size_ext(ctx_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE), + ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) : 0, + }; +} + +server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq) { + const auto [n_tgt, n_dft] = server_page_state_size(ctx_tgt, ctx_dft, seq); + if (n_tgt == 0) { + throw server_page_error(ERROR_TYPE_SERVER, "the target sequence state could not be sized"); + } + auto page = std::make_shared(); + try { + page->data.main.resize(n_tgt); + page->data.drft.resize(n_dft); + } catch (const std::bad_alloc &) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, "host RAM for the page could not be allocated"); + } + if (llama_state_seq_get_data_ext(ctx_tgt, page->data.main.data(), n_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE) != n_tgt) { + throw server_page_error(ERROR_TYPE_SERVER, "the target sequence state read short"); + } + if (n_dft > 0 && llama_state_seq_get_data_ext(ctx_dft, page->data.drft.data(), n_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) != n_dft) { + throw server_page_error(ERROR_TYPE_SERVER, "the draft sequence state read short"); + } + if (spec != nullptr) { + common_speculative_get_state(spec, seq, page->spec); + } + return page; +} + +void server_page_restore(const server_page & page, llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq) { + const auto & tgt = page.data.main; + const auto & dft = page.data.drft; + if (llama_state_seq_set_data_ext(ctx_tgt, tgt.data(), tgt.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE) != tgt.size()) { + throw server_page_error(ERROR_TYPE_SERVER, "the target sequence state did not fit the context (KV cache full?)"); + } + if (!dft.empty()) { + if (ctx_dft == nullptr) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the page carries a draft state and this server has no draft context"); + } + if (llama_state_seq_set_data_ext(ctx_dft, dft.data(), dft.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE) != dft.size()) { + throw server_page_error(ERROR_TYPE_SERVER, "the draft sequence state did not fit the draft context"); + } + } else if (ctx_dft != nullptr) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the page carries no draft state and this server drafts: the draft would be cold"); + } + if (spec != nullptr && !page.spec.empty()) { + common_speculative_set_state(spec, seq, page.spec); + } +} + +// +// server_page_store +// + +void server_page_store::admit(const std::string & name, size_t n_bytes_new) const { + if (limit_bytes == 0) { + return; + } + size_t held = n_bytes(); + const auto it = pages.find(name); + if (it != pages.end()) { + held -= it->second->n_bytes(); + } + if (held + n_bytes_new > limit_bytes) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, string_format( + "the page store holds %zu bytes; a page of %zu bytes would exceed its ceiling of %llu bytes: spill or drop a page first", + held, n_bytes_new, (unsigned long long) limit_bytes)); + } +} + +uint64_t server_page_store::put(const std::string & name, server_page_ptr page) { + page->gen = next_gen++; + const auto it = pages.find(name); + if (it != pages.end()) { + SRV_WRN("page '%s' replaced (%zu bytes freed)\n", name.c_str(), it->second->n_bytes()); + } + pages[name] = std::move(page); + return pages[name]->gen; +} + +server_page_ptr server_page_store::find(const std::string & name) const { + const auto it = pages.find(name); + return it == pages.end() ? nullptr : it->second; +} + +size_t server_page_store::erase(const std::string & name, uint64_t gen) { + const auto it = pages.find(name); + if (it == pages.end() || (gen != 0 && it->second->gen != gen)) { + return 0; + } + const size_t n = it->second->n_bytes(); + pages.erase(it); + return n; +} + +void server_page_store::evict_lora(const llama_adapter_lora * adapter) { + for (auto it = pages.begin(); it != pages.end();) { + bool uses = false; + for (const auto & l : it->second->lora) { + uses = uses || (l.ptr == adapter && l.scale != 0.0f); + } + if (uses) { + SRV_WRN("page '%s' dropped: its adapter is being unloaded (%zu bytes freed)\n", it->first.c_str(), it->second->n_bytes()); + it = pages.erase(it); + } else { + ++it; + } + } +} + +size_t server_page_store::n_bytes() const { + size_t n = 0; + for (const auto & [name, page] : pages) { + n += page->n_bytes(); + } + return n; +} + +bool server_page_name_valid(const std::string & name) { + if (name.empty() || name.size() > 255) { + return false; + } + for (const unsigned char c : name) { + if (c < 0x20 || c == 0x7f) { + return false; + } + } + return true; +} + +// +// page files +// + +std::pair lora_file_revision(const std::string & path) { + std::error_code ec; + const auto size = std::filesystem::file_size(path, ec); + if (ec) { + return {0, 0}; + } + const auto mtime = std::filesystem::last_write_time(path, ec); + if (ec) { + return {0, 0}; + } + return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; +} + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put) { + const uint32_t magic = PAGE_CKPT_MAGIC; + const uint32_t version = 1; + const uint32_t n_ckpt = (uint32_t) kept.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); + for (const auto * c : kept) { + const int64_t n_tokens = c->n_tokens; + const int32_t pos_min = c->pos_min; + const int32_t pos_max = c->pos_max; + const uint64_t sz_tgt = c->data_tgt.size(); + const uint64_t sz_dft = c->data_dft.size(); + const uint64_t sz_spec = c->data_spec.size(); + put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); + put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(c->data_tgt.data(), c->data_tgt.size()); + put(c->data_dft.data(), c->data_dft.size()); + put(c->data_spec.data(), c->data_spec.size()); + } +} + +void server_page_write_adapters(const std::vector & lora, const server_page_put & put) { + const uint32_t magic = PAGE_LORA_MAGIC; + const uint32_t version = 2; + const uint32_t n_lora = (uint32_t) lora.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); + for (const auto & l : lora) { + const uint32_t n_path = (uint32_t) l.path.size(); + const auto rev = lora_file_revision(l.path); + put(&n_path, sizeof(n_path)); + put(l.path.data(), n_path); + put(&l.scale, sizeof(l.scale)); + put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); + } +} + +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put) { + const uint32_t magic = PAGE_DRFT_MAGIC; + const uint32_t version = 1; + const uint64_t sz_dft = n_dft; + const uint64_t sz_spec = spec.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); + put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(dft, n_dft); + put(spec.data(), spec.size()); +} + +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why) { + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) { + why = "absent"; + return false; + } + const auto total = (uint64_t) in.tellg(); + in.seekg(0); + uint32_t magic = 0, version = 0; + uint64_t sz_dft = 0, sz_spec = 0; + in.read(reinterpret_cast(&magic), sizeof(magic)); + in.read(reinterpret_cast(&version), sizeof(version)); + in.read(reinterpret_cast(&sz_dft), sizeof(sz_dft)); + in.read(reinterpret_cast(&sz_spec), sizeof(sz_spec)); + const uint64_t header = 2 * sizeof(uint32_t) + 2 * sizeof(uint64_t); + if (!in || magic != PAGE_DRFT_MAGIC || version != 1 || sz_dft > total || sz_spec > total || header + sz_dft + sz_spec != total) { + why = "malformed"; + return false; + } + dft.resize(sz_dft); + spec.resize(sz_spec); + in.read(reinterpret_cast(dft.data()), sz_dft); + in.read(reinterpret_cast(spec.data()), sz_spec); + if (!in) { + why = "truncated"; + return false; + } + return true; +} + +static bool write_sidecar(const std::string & path, const std::function & body, size_t & n_written) { + std::ofstream out(path, std::ios::binary | std::ios::trunc); + body([&](const void * ptr, size_t n) { + out.write(reinterpret_cast(ptr), n); + n_written += n; + }); + out.flush(); + return (bool) out; +} + +size_t server_page_spill(const server_page & page, const std::string & filepath) { + const std::string sidecars[] = { filepath + ".ckpt", filepath + ".lora", filepath + ".dft" }; + auto remove_all = [&]() { + std::error_code ec; + std::filesystem::remove(filepath, ec); + for (const auto & p : sidecars) { + std::filesystem::remove(p, ec); + } + }; + // no file of an older page under this name may survive beside the new state + remove_all(); + + std::vector packed; + try { + packed = page.prompt.tokens.serialize(); + } catch (const std::exception & err) { + throw server_page_error(ERROR_TYPE_NOT_SUPPORTED, err.what()); + } + size_t n_written = llama_state_seq_save_file_from_data(filepath.c_str(), + page.data.main.data(), page.data.main.size(), + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + if (n_written == 0) { + remove_all(); + throw server_page_error(ERROR_TYPE_SERVER, "the page's sequence state could not be written: " + filepath); + } + + std::vector all; + for (const auto & c : page.prompt.checkpoints) { + all.push_back(&c); + } + const bool ok = + write_sidecar(sidecars[0], [&](const server_page_put & put) { server_page_write_checkpoints(all, put); }, n_written) && + write_sidecar(sidecars[1], [&](const server_page_put & put) { server_page_write_adapters(page.lora, put); }, n_written) && + write_sidecar(sidecars[2], [&](const server_page_put & put) { server_page_write_draft(page.data.drft.data(), page.data.drft.size(), page.spec, put); }, n_written); + if (!ok) { + remove_all(); + throw server_page_error(ERROR_TYPE_SERVER, "a page sidecar could not be written beside " + filepath); + } + return n_written; +} diff --git a/tools/server/server-pages.h b/tools/server/server-pages.h new file mode 100644 index 000000000000..5ce42b33a6d5 --- /dev/null +++ b/tools/server/server-pages.h @@ -0,0 +1,111 @@ +#pragma once + +// Continuum fork: a slot's sequence state held as a named page in host RAM. +// +// A mind's switch costs a memcpy instead of a disk save (66-100 s at 27B) or a re-prefill +// (70-160 s): `page_out` moves a slot's target KV, draft (MTP) KV, speculative tail state +// and context checkpoints into the store and frees the slot; `page_in` puts them back into +// any slot with its tokens, so the next request's prefix matches. `spill` writes a page in +// the slot-save format (.bin + .ckpt + .lora + .dft) that `restore` reads. +// Design: continuum docs/architecture/SERVING-ENGINE.md, slice S1. + +#include "server-common.h" +#include "server-task.h" + +#include "common.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +struct common_speculative; + +// A page operation that failed, with the HTTP error class the caller sees. +struct server_page_error : std::runtime_error { + error_type type; + server_page_error(error_type type, const std::string & msg) : std::runtime_error(msg), type(type) {} +}; + +struct server_page { + uint64_t gen = 0; // identity within the store: a page replaced under the same name is another page + + server_prompt prompt; // tokens and context checkpoints + server_prompt_data data; // target (main) and draft (drft) sequence state + std::vector spec; // speculative tail state (common_speculative_get_state); may be empty + + std::vector lora; // the slot's adapter list the state was computed under + + size_t n_checkpoint_bytes() const; + size_t n_bytes() const; +}; + +using server_page_ptr = std::shared_ptr; + +// The sequence state bytes a page of `seq` will hold (target, draft), before it is taken. +std::pair server_page_state_size(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq); + +// Copies `seq`'s state out of both contexts (each get_data synchronizes its context first) and +// the speculative tail state; the caller moves in the prompt and adapters. Throws on a short read. +server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq); + +// Writes a page's state into `seq` of both contexts. The page is left intact (it may be spilling). +// Throws server_page_error; the caller clears the slot so a failure is never a silent half-state. +void server_page_restore(const server_page & page, llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq); + +// The named host-RAM page store. Owned by the server's task thread. +struct server_page_store { + uint64_t limit_bytes = 0; // --page-ram-bytes; 0 = no ceiling + + // Refuses (server_page_error, UNAVAILABLE) a page of `n_bytes` under `name` that would take the + // store past its ceiling; a page it replaces under the same name is not counted. + void admit(const std::string & name, size_t n_bytes) const; + + // Stores `page` under `name`, replacing any page there; returns its generation. + uint64_t put(const std::string & name, server_page_ptr page); + + server_page_ptr find(const std::string & name) const; + + // Removes the page under `name` when it is generation `gen` (0 = any); returns the bytes freed. + size_t erase(const std::string & name, uint64_t gen = 0); + + // Drops every page computed with `adapter` active: the adapter is about to be freed. + void evict_lora(const llama_adapter_lora * adapter); + + size_t n_bytes() const; + size_t n_pages() const { return pages.size(); } + +private: + std::map pages; + uint64_t next_gen = 1; +}; + +// A page name: a non-empty opaque key of at most 255 bytes without control characters. +bool server_page_name_valid(const std::string & name); + +// +// the slot page files: `` (sequence state) and its sidecars, shared by save, spill and restore +// + +using server_page_put = std::function; + +// An adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it cannot be read), +// stored in the `.lora` sidecar so a restore refuses a page computed under different weights. +std::pair lora_file_revision(const std::string & path); + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put); // .ckpt +void server_page_write_adapters(const std::vector & lora, const server_page_put & put); // .lora +// `dft` may be null when `put` only counts bytes +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put); // .dft + +// Reads a `.dft` sidecar. Returns false with `why` set when it is absent or malformed. +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why); + +// Writes `page` to `filepath` (+ .ckpt, .lora, .dft) with every checkpoint. Touches no context, so it +// runs off the task thread. All four files are written or none are left. Returns the bytes written. +size_t server_page_spill(const server_page & page, const std::string & filepath); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 746c3a1f98c8..32a1502a6e31 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1591,6 +1591,26 @@ std::string server_task_result_metrics::to_metrics() { "spec_decode_num_drafts_total", "Speculative: Total speculative decoding verification steps", (double) metrics.n_draft_verif_steps + }, { + "page_out_total", + "Slot states paged out to the host-RAM page store", + (double) metrics.n_page_out + }, { + "page_in_total", + "Pages restored from the host-RAM page store into a slot", + (double) metrics.n_page_in + }, { + "page_failed_total", + "page_out, page_in or spill operations refused or failed", + (double) metrics.n_page_failed + }, { + "ctx_checkpoints_evicted_budget_total", + "Context checkpoints dropped to keep within --ctx-checkpoints-bytes", + (double) metrics.n_ckpt_evicted_budget + }, { + "ctx_checkpoints_skipped_budget_total", + "Context checkpoints not taken because one alone exceeds --ctx-checkpoints-bytes", + (double) metrics.n_ckpt_skipped_budget }, }; @@ -1615,6 +1635,18 @@ std::string server_task_result_metrics::to_metrics() { "n_busy_slots_per_decode", "Average number of busy slots per llama_decode() call", (double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0) + }, { + "page_store_bytes", + "Bytes held by the host-RAM page store", + (double) page_store_bytes + }, { + "page_store_pages", + "Pages held by the host-RAM page store", + (double) page_store_pages + }, { + "ctx_checkpoint_bytes", + "Bytes held by every slot's context checkpoints", + (double) ctx_checkpoint_bytes }, }; @@ -1675,6 +1707,33 @@ json server_task_result_slot_save_load::to_json() { }; } +// +// server_task_result_page +// +json server_task_result_page::to_json() { + json res = { + { "id_slot", id_slot }, + { "action", action }, + { "page", page }, + { "gen", gen }, + { "n_tokens", n_tokens }, + { "n_bytes", n_bytes }, + { "n_bytes_draft", n_bytes_draft }, + { "n_checkpoints", n_checkpoints }, + { "n_bytes_checkpoints", n_bytes_checkpoints }, + { "store_bytes", store_bytes }, + { "store_pages", store_pages }, + { "timings", { + { "ms", t_ms } + }}, + }; + if (!filename.empty()) { + res["filename"] = filename; + res["n_written"] = n_written; + } + return res; +} + // // server_task_result_slot_erase // diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 87a344e174cf..ff2b8fad6846 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -56,6 +56,10 @@ enum server_task_type { SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE, + SERVER_TASK_TYPE_PAGE_OUT, // Continuum fork: slot -> named host-RAM page (server-pages.h) + SERVER_TASK_TYPE_PAGE_IN, // named page -> slot + SERVER_TASK_TYPE_PAGE_SPILL, // hands a page to the HTTP thread, which writes it to disk + SERVER_TASK_TYPE_PAGE_DROP, // frees a page SERVER_TASK_TYPE_GET_LORA, SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, @@ -237,6 +241,10 @@ struct server_task { // head of its first user turn, so the middle of the list is bytes nobody reads. int32_t keep_newest_checkpoints = -1; std::shared_ptr progress; + // Continuum fork: the RAM page a PAGE_* task names, and for PAGE_DROP the generation it + // drops (0 = whichever page holds the name), so a spill's drop never frees a newer page + std::string page; + uint64_t page_gen = 0; }; slot_action slot_action; @@ -578,6 +586,9 @@ struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) int n_processing_slots = 0; int n_tasks_deferred = 0; + size_t page_store_bytes = 0; // Continuum fork: the host-RAM page store + size_t page_store_pages = 0; + size_t ctx_checkpoint_bytes = 0; // every slot's context checkpoints server_metrics metrics; @@ -614,6 +625,32 @@ struct server_task_result_slot_save_load : server_task_result { virtual json to_json() override; }; +struct server_page; + +// Continuum fork: the outcome of a PAGE_* task, with the store's totals after it +struct server_task_result_page : server_task_result { + std::string action; + std::string page; + uint64_t gen = 0; + + size_t n_tokens = 0; + size_t n_bytes = 0; // the page: target + draft state + speculative state + checkpoints + size_t n_bytes_draft = 0; + size_t n_checkpoints = 0; + size_t n_bytes_checkpoints = 0; + double t_ms = 0.0; + + size_t store_bytes = 0; + size_t store_pages = 0; + + // PAGE_SPILL: the page the HTTP thread writes; shared, so the task thread never waits on disk + std::shared_ptr spilled; + std::string filename; + size_t n_written = 0; + + virtual json to_json() override; +}; + struct server_task_result_slot_erase : server_task_result { size_t n_erased; diff --git a/tools/server/tests/unit/test_slot_page.py b/tools/server/tests/unit/test_slot_page.py new file mode 100644 index 000000000000..29dfb20a83af --- /dev/null +++ b/tools/server/tests/unit/test_slot_page.py @@ -0,0 +1,171 @@ +import os +from pathlib import Path + +import pytest +from utils import * + +# Continuum fork: a slot's sequence state paged to a named host-RAM page and back +# (server-pages.h; continuum docs/architecture/SERVING-ENGINE.md, slice S1). + +server = ServerPreset.tinyllama2() + +PROMPT = "What is the capital of France?" +FOLLOW = "What is the capital of France? And what is the capital of Germany?" + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.slot_save_path = "./tmp" + server.temperature = 0.0 + server.server_metrics = True + + +def complete(prompt, id_slot): + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, "id_slot": id_slot, "cache_prompt": True, "n_predict": 16, + }) + assert res.status_code == 200 + return res.body + + +def page(id_slot, action, name, body=None): + return server.make_request("POST", f"/slots/{id_slot}?action={action}&page={name}", data=body or {}) + + +def metric(name): + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + for line in res.body.splitlines(): + if line.startswith(f"llamacpp:{name} "): + return float(line.split()[1]) + raise AssertionError(f"metric {name} missing") + + +def test_page_out_in_is_identical_to_never_paging(): + # what this catches: a page_in that restores the KV but not the slot's tokens (the next + # prompt re-prefills from 0), or restores them into a corrupt state (the output differs from + # a slot that was never paged). The page goes to ANOTHER slot, and the source slot is reused + # in between, so the page cannot be served by KV that happened to stay in place. + global server + server.start() + assert server.make_request("GET", "/props").body["slot_page_version"] == 1 + + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200, out.body + n_tokens = out.body["n_tokens"] + assert n_tokens == first["timings"]["prompt_n"] + first["timings"]["predicted_n"] - 1 + assert out.body["n_bytes"] > 0 and out.body["store_pages"] == 1 + assert out.body["store_bytes"] == out.body["n_bytes"] + assert metric("page_store_bytes") == out.body["n_bytes"] + + # the page left the slot: its KV is free and its cached tokens are gone + scrub = complete("Once upon a time", 1) + assert scrub["timings"]["prompt_n"] == scrub["tokens_evaluated"] + + restored = page(0, "page_in", "mind-a") + assert restored.status_code == 200, restored.body + assert restored.body["n_tokens"] == n_tokens + assert restored.body["store_pages"] == 0 and restored.body["store_bytes"] == 0 + paged = complete(FOLLOW, 0) + + # the reference: the same two requests on a slot that never paged + assert server.make_request("POST", "/slots/1?action=erase").status_code == 200 + complete(PROMPT, 1) + reference = complete(FOLLOW, 1) + + assert paged["timings"]["prompt_n"] == reference["timings"]["prompt_n"] + assert paged["timings"]["prompt_n"] < paged["tokens_evaluated"] # the prefix came from the page + assert paged["content"] == reference["content"] + assert metric("page_out_total") == 1 and metric("page_in_total") == 1 + + +def test_page_errors_are_typed_and_leave_state_intact(): + # what this catches: a missing or consumed page answered as a silent cold start, or a + # failed page_out that frees the slot anyway. + global server + server.start() + assert page(0, "page_in", "nobody").status_code == 404 + assert page(0, "page_out", "empty").status_code == 400 # nothing in the slot + assert page(0, "page_out", "bad%01name").status_code == 400 + complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 200 + assert page(1, "page_in", "mind-a").status_code == 200 + assert page(0, "page_in", "mind-a").status_code == 404 # a page_in consumes the page + assert metric("page_failed_total") >= 3 + + +def test_page_ram_ceiling_refuses_and_keeps_the_slot(): + # what this catches: a page store that grows past its ceiling on unified memory, or a + # refused page_out that still drops the slot's sequence. + global server + server.page_ram_bytes = 1 + server.start() + first = complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 503 + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 # the slot still holds the sequence + assert again["content"] == first["content"] + + +def test_spill_writes_the_slot_save_format(): + # what this catches: a spilled page that today's restore cannot read, a spill that blocks + # on a page replaced meanwhile, or a drop that leaves the RAM bytes counted. + global server + server.start() + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200 + spilled = page(1, "spill", "mind-a", {"filename": "mind-a.bin", "drop": True}) + assert spilled.status_code == 200, spilled.body + files = [Path(server.slot_save_path) / ("mind-a.bin" + s) for s in ("", ".ckpt", ".lora", ".dft")] + assert all(f.exists() for f in files) + assert spilled.body["n_written"] == sum(f.stat().st_size for f in files) + assert spilled.body["store_pages"] == 0 and metric("page_store_bytes") == 0 + + res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "mind-a.bin"}) + assert res.status_code == 200 + assert res.body["n_restored"] == out.body["n_tokens"] + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 + assert again["content"] == first["content"] + + # a spill names a filename like save does + assert page(1, "spill", "mind-a", {"filename": "../x.bin"}).status_code == 400 + + +HYBRID = os.environ.get("LLAMA_SERVER_TEST_HYBRID_GGUF") + + +@pytest.mark.skipif(not HYBRID, reason="set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF") +def test_checkpoints_stay_within_the_byte_budget_and_page_with_the_slot(): + # what this catches: checkpoints counted instead of measured (a hybrid's checkpoints are the + # hidden RAM cost), and a page that drops a slot's checkpoints (the next divergence would + # re-prefill from the head). + global server + server.model_hf_repo = None + server.model_file = HYBRID + server.n_ctx = 4096 + server.n_batch = 512 + server.n_slots = 2 + server.start() + long_prompt = " ".join(f"line {i} of the story continues." for i in range(120)) + complete(long_prompt, 0) + per_slot = metric("ctx_checkpoint_bytes") + assert per_slot > 0 + out = page(0, "page_out", "mind-a") + assert out.status_code == 200 and out.body["n_checkpoints"] > 0 + assert metric("ctx_checkpoint_bytes") == 0 # they left with the page + restored = page(0, "page_in", "mind-a") + assert restored.body["n_checkpoints"] == out.body["n_checkpoints"] + server.stop() + + budget = int(per_slot) + int(per_slot) // 2 # room for one slot's set and half another's + server.ctx_checkpoints_bytes = budget + server.start() + complete(long_prompt, 0) + complete("A different tale: " + long_prompt, 1) + assert metric("ctx_checkpoint_bytes") <= budget + assert metric("ctx_checkpoints_evicted_budget_total") > 0 diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c591..7bc7215acf29 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -121,6 +121,8 @@ class ServerProcess: mcp_servers_config: str | None = None mcp_servers_json: str | None = None cors_origins: str | None = None + ctx_checkpoints_bytes: int | None = None + page_ram_bytes: int | None = None # session variables process: subprocess.Popen | None = None @@ -285,6 +287,10 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--mcp-servers-json", self.mcp_servers_json]) if self.backend_sampling: server_args.append("--backend_sampling") + if self.ctx_checkpoints_bytes is not None: + server_args.extend(["--ctx-checkpoints-bytes", self.ctx_checkpoints_bytes]) + if self.page_ram_bytes is not None: + server_args.extend(["--page-ram-bytes", self.page_ram_bytes]) if self.gcp_compat: env["AIP_MODE"] = "PREDICTION" From e08b11ed1ca3fa8c217112d967c29bddab0f8042 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sat, 24 Oct 2026 04:37:49 -0500 Subject: [PATCH 2/3] server: page RAM ceiling counts every live page; page files publish whole generations Review fixes on the RAM page verbs (Codex on #70): - One count: server_page_bytes is what page_out admits (measured before any copy, speculative tail state included) and what the store charges; put re-checks the ceiling on exactly the page it stores, so no pre-count can get a page past it. - The ceiling counts live ownership: a stored page carries a charge released when its LAST owner drops it, so a page a spill still holds after page_in, page_drop, replacement or a model unload keeps counting. Admission subtracts a replaced page only when the store is its sole owner. - page_out moves the slot's prompt only into a stored page: put inserts first and moves the prompt back on any refusal, so a refused or failed page_out leaves the slot's sequence intact. - Page files (server-page-files.{h,cpp}, moved out of server-pages.cpp): save and spill publish `` + .ckpt/.lora/.dft as one generation. Every file is written to `..page-staging/-/` and synced, a commit record is renamed into place, then the files move over the committed ones and a sidecar the generation lacks is removed. A failure before the record leaves the committed generation untouched and nothing staged; restore reads through server_page_reader, which completes a publication cut short after its record, removes stages no live publication owns, and holds commits off while it reads. A save whose sidecar cannot be written now fails instead of leaving a mixed set. - /slots page results report n_bytes_spec. Tests: tests/test-server-pages.cpp (admission = charge with speculative state, spill-held pages stay counted, refusal keeps the incumbent, failed and interrupted publications keep the committed generation and leave nothing staged); test_slot_page.py adds a crashed-publication restore and an MTP page at an exact ceiling (hybrid GGUF). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tests/CMakeLists.txt | 4 + tests/test-server-pages.cpp | 268 +++++++++++++ tools/server/CMakeLists.txt | 2 + tools/server/server-context.cpp | 77 ++-- tools/server/server-page-files.cpp | 444 ++++++++++++++++++++++ tools/server/server-page-files.h | 67 ++++ tools/server/server-pages.cpp | 254 ++++--------- tools/server/server-pages.h | 89 +++-- tools/server/server-task.cpp | 1 + tools/server/server-task.h | 1 + tools/server/tests/unit/test_slot_page.py | 80 ++++ 11 files changed, 1018 insertions(+), 269 deletions(-) create mode 100644 tests/test-server-pages.cpp create mode 100644 tools/server/server-page-files.cpp create mode 100644 tools/server/server-page-files.h diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 498b70553bba..61a391fc803f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -172,6 +172,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) + # Continuum fork: the RAM page store's ceiling and the page files' whole-generation publication + llama_build_and_test(test-server-pages.cpp) + target_include_directories(test-server-pages PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) + target_link_libraries(test-server-pages PRIVATE server-context) # TODO: disabled on loongarch64 because the ggml-ci node lacks Python 3.8 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "loongarch64") llama_build_and_test(test-json-schema-to-grammar.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) diff --git a/tests/test-server-pages.cpp b/tests/test-server-pages.cpp new file mode 100644 index 000000000000..b02df7f34758 --- /dev/null +++ b/tests/test-server-pages.cpp @@ -0,0 +1,268 @@ +// what this catches (Continuum fork, server-pages.h + server-page-files.h, no model needed): +// - a page admitted on a count that leaves out part of what it holds (its speculative state), so the +// store passes its ceiling after a successful page_out; +// - a page out of the store but still held by a spill that stops counting, so the live RAM passes the ceiling; +// - a refused page_out that has already moved the slot's state out (the incumbent is lost); +// - a publication that removes or overwrites the committed files before its own are whole: a failed +// sidecar write, or a crash between files, that loses the last durable page. + +#include "server-page-files.h" +#include "server-pages.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +static long current_pid() { return (long) _getpid(); } +#else +# include +static long current_pid() { return (long) getpid(); } +#endif + +#define CHECK(cond) do { if (!(cond)) { fprintf(stderr, "FAILED %s:%d: %s\n", __FILE__, __LINE__, #cond); exit(1); } } while (0) + +namespace fs = std::filesystem; + +// a page holding `n` bytes in each part, the speculative tail state included +static server_page_ptr make_page(size_t n) { + auto page = std::make_shared(); + page->data.main.assign(n, 1); + page->data.drft.assign(n, 2); + page->spec.assign(n, 3); + return page; +} + +// a slot's prompt: tokens and one checkpoint of `n` bytes +static server_prompt make_prompt(size_t n) { + server_prompt prompt; + prompt.tokens = server_tokens(llama_tokens{ 1, 2, 3 }, false); + auto & c = prompt.checkpoints.emplace_back(); + c.data_tgt.assign(n, 4); + return prompt; +} + +static error_type refusal(const std::function & fn) { + try { + fn(); + } catch (const server_page_error & err) { + return err.type; + } + return ERROR_TYPE_SERVER; // stands for "not refused"; every refusal checked here is UNAVAILABLE +} + +static void test_admission_counts_what_the_store_charges() { + const size_t n = 64; + const size_t bytes = server_page_bytes(2 * n, n, n); // main + draft, speculative, one checkpoint + { + server_page_store store; + store.limit_bytes = bytes; // fits exactly + auto page = make_page(n); + auto from = make_prompt(n); + store.admit("a", bytes); + store.put("a", page, &from); + CHECK(page->n_bytes() == bytes); + CHECK(store.n_bytes() == bytes); + CHECK(from.checkpoints.empty() && page->prompt.checkpoints.size() == 1); + } + { + server_page_store store; + store.limit_bytes = bytes - 1; // one byte short: the speculative state alone tips it + CHECK(refusal([&] { store.admit("a", bytes); }) == ERROR_TYPE_UNAVAILABLE); + // a put is checked on the same count, so a short pre-count cannot get a page past the ceiling + auto page = make_page(n); + auto from = make_prompt(n); + CHECK(refusal([&] { store.put("a", page, &from); }) == ERROR_TYPE_UNAVAILABLE); + // the refusal leaves the incumbent: the slot's prompt is where it was, and nothing is stored + CHECK(from.tokens.size() == 3 && from.checkpoints.size() == 1 && from.checkpoints.front().data_tgt.size() == n); + CHECK(page->prompt.tokens.size() == 0); + CHECK(store.n_bytes() == 0 && store.n_pages() == 0 && store.find("a") == nullptr); + } +} + +static void test_spill_held_pages_stay_counted() { + const size_t n = 64; + const size_t bytes = server_page_bytes(2 * n, n, n); + server_page_store store; + store.limit_bytes = 2 * bytes - 1; // room for one page, not two + + auto from = make_prompt(n); + store.put("a", make_page(n), &from); + server_page_ptr spilling = store.find("a"); // a spill's reference on its HTTP thread + + // replaced under its name while spilling: the old page is still live, so the new one does not fit + auto from2 = make_prompt(n); + CHECK(refusal([&] { store.put("a", make_page(n), &from2); }) == ERROR_TYPE_UNAVAILABLE); + CHECK(from2.checkpoints.size() == 1); + + // left the store (page_in, page_drop) while spilling: still counted + CHECK(store.erase("a")); + CHECK(store.n_pages() == 0 && store.n_bytes() == bytes); + CHECK(refusal([&] { store.put("b", make_page(n), &from2); }) == ERROR_TYPE_UNAVAILABLE); + + // the spill finishes: its bytes are released and the page fits + spilling.reset(); + CHECK(store.n_bytes() == 0); + store.put("b", make_page(n), &from2); + CHECK(store.n_bytes() == bytes); + + // a replacement the store solely owns frees the page it replaces, so it fits + auto from3 = make_prompt(n); + store.put("b", make_page(n), &from3); + CHECK(store.n_bytes() == bytes && store.n_pages() == 1); + + // a cleared store (the model unloaded) still counts a page a spill holds + spilling = store.find("b"); + store.clear(); + CHECK(store.n_bytes() == bytes); + spilling.reset(); + CHECK(store.n_bytes() == 0); +} + +static std::string read_file(const fs::path & p) { + std::ifstream in(p, std::ios::binary); + return in ? std::string(std::istreambuf_iterator(in), std::istreambuf_iterator()) : std::string(""); +} + +static void write_file(const fs::path & p, const std::string & bytes) { + std::ofstream out(p, std::ios::binary | std::ios::trunc); + out << bytes; +} + +static server_page_file text_file(const std::string & suffix, const std::string & text) { + return server_page_sidecar(suffix, [text](const server_page_put & put) { put(text.data(), text.size()); }); +} + +static const char * const SUFFIXES[] = { "", ".ckpt", ".lora", ".dft" }; + +static void check_generation(const fs::path & file, const std::string & tag) { + for (const char * s : SUFFIXES) { + CHECK(read_file(file.string() + s) == tag + s); + } +} + +// no stage and no commit record is left beside the page +static void check_no_leftovers(const fs::path & dir) { + const fs::path staging = dir / "..page-staging"; + if (fs::exists(staging)) { + CHECK(fs::is_empty(staging)); + } +} + +static void publish_generation(const fs::path & file, const std::string & tag) { + std::vector files; + for (const char * s : SUFFIXES) { + files.push_back(text_file(s, tag + s)); + } + server_page_publish(file.string(), files); +} + +static void test_a_failed_publication_keeps_the_committed_generation(const fs::path & dir) { + const fs::path file = dir / "mind-a.bin"; + publish_generation(file, "one"); + check_generation(file, "one"); + + // the .lora sidecar fails after the state and .ckpt of the new generation were written + const auto failing = server_page_file{ ".lora", [](const std::string &) -> size_t { return 0; } }; + bool refused = false; + try { + server_page_publish(file.string(), { text_file("", "two"), text_file(".ckpt", "two.ckpt"), failing, text_file(".dft", "two.dft") }); + } catch (const server_page_error &) { + refused = true; + } + CHECK(refused); + { server_page_reader reader(file.string()); check_generation(file, "one"); } + check_no_leftovers(dir); + + // a writer that throws is the same failure + refused = false; + try { + server_page_publish(file.string(), { text_file("", "two"), + server_page_sidecar(".ckpt", [](const server_page_put &) { throw std::runtime_error("disk full"); }) }); + } catch (const server_page_error &) { + refused = true; + } + CHECK(refused); + check_generation(file, "one"); + check_no_leftovers(dir); + + // a generation without a sidecar removes the older one: a stale .dft is never read beside new state + server_page_publish(file.string(), { text_file("", "three"), text_file(".ckpt", "three.ckpt"), text_file(".lora", "three.lora") }); + CHECK(read_file(file) == "three" && read_file(file.string() + ".dft") == ""); + check_no_leftovers(dir); +} + +static void test_an_interrupted_publication_restores_whole(const fs::path & dir) { + const fs::path file = dir / "mind-b.bin"; + const fs::path staging = dir / "..page-staging"; + publish_generation(file, "one"); + + // a crash after staging, before the commit record: this process's pid, a stage no publication owns + const fs::path stage = staging / (std::to_string(current_pid()) + "-999999"); + fs::create_directories(stage); + write_file(stage / "mind-b.bin", "two"); + write_file(stage / "mind-b.bin.ckpt", "two.ckpt"); + { server_page_reader reader(file.string()); check_generation(file, "one"); } + check_no_leftovers(dir); + + // a crash after the commit record, with one file moved: the reader completes the generation, + // and the sidecar the generation lacks (.dft) is removed + fs::create_directories(stage); + for (const char * s : { "", ".ckpt", ".lora" }) { + write_file(stage / (std::string("mind-b.bin") + s), std::string("two") + s); + } + fs::rename(stage / "mind-b.bin", file); + write_file(staging / "mind-b.bin.commit", stage.filename().string() + "\nf\nf.ckpt\nf.lora\n"); + { + server_page_reader reader(file.string()); + CHECK(read_file(file) == "two" && read_file(file.string() + ".ckpt") == "two.ckpt" && read_file(file.string() + ".lora") == "two.lora"); + CHECK(read_file(file.string() + ".dft") == ""); + } + check_no_leftovers(dir); +} + +static void test_spill_writes_a_page_restore_reads(const fs::path & dir) { + // a sequence state as llama_state_seq_get_data writes it: [io magic][seq id][state] + auto page = std::make_shared(); + const uint32_t io_magic = 0xaf143cd8; + page->data.main.resize(sizeof(uint32_t) + sizeof(llama_seq_id) + 16, 7); + memcpy(page->data.main.data(), &io_magic, sizeof(io_magic)); + page->data.drft.assign(8, 5); + page->spec.assign(4, 6); + page->prompt = make_prompt(32); + + const fs::path file = dir / "mind-c.bin"; + const size_t n = server_page_spill(*page, file.string()); + size_t on_disk = 0; + for (const char * s : SUFFIXES) { + on_disk += fs::file_size(file.string() + s); + } + CHECK(n == on_disk); + std::vector dft, spec; + std::string why; + CHECK(server_page_read_draft(file.string() + ".dft", dft, spec, why)); + CHECK(dft == page->data.drft && spec == page->spec); + check_no_leftovers(dir); +} + +int main() { + test_admission_counts_what_the_store_charges(); + test_spill_held_pages_stay_counted(); + + const fs::path dir = fs::temp_directory_path() / ("test-server-pages-" + std::to_string(current_pid())); + fs::remove_all(dir); + fs::create_directories(dir); + test_a_failed_publication_keeps_the_committed_generation(dir); + test_an_interrupted_publication_restores_whole(dir); + test_spill_writes_a_page_restore_reads(dir); + fs::remove_all(dir); + + printf("test-server-pages: OK\n"); + return 0; +} diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 2e1c2ed087ca..e4a9783970a8 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -18,6 +18,8 @@ add_library(${TARGET} STATIC server-context.h server-pages.cpp server-pages.h + server-page-files.cpp + server-page-files.h server-checkpoints.cpp server-checkpoints.h server-stream.cpp diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d121ebcb2f5c..8846f3038323 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3,6 +3,7 @@ #include "../../src/llama-ext.h" #include "server-train.h" #include "server-checkpoints.h" +#include "server-page-files.h" #include "server-pages.h" #include "server-chat.h" #include "server-common.h" @@ -953,7 +954,7 @@ struct server_context_impl { if (page_store.n_pages() > 0) { SRV_WRN("dropping %zu RAM pages (%zu bytes) with the model: spill them before sleep to keep them\n", page_store.n_pages(), page_store.n_bytes()); - page_store = server_page_store{}; + page_store.clear(); } spec.reset(); spec_init.reset(); @@ -2549,13 +2550,14 @@ struct server_context_impl { if (slot->prompt.tokens.size() == 0) { throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the slot holds no sequence to page out"); } - const auto [n_tgt, n_dft] = server_page_state_size(ctx_tgt, ctx_dft, slot->id); - page_store.admit(name, n_tgt + n_dft + server_checkpoints_bytes(slot->prompt.checkpoints)); - page = server_page_capture(ctx_tgt, ctx_dft, spec.get(), slot->id); - page->prompt = std::move(slot->prompt); - page->lora = slot->lora; + // refused before anything is copied, on the count the store charges + auto shape = server_page_measure(ctx_tgt, ctx_dft, spec.get(), slot->id); + page_store.admit(name, shape.n_bytes(server_checkpoints_bytes(slot->prompt.checkpoints))); + page = server_page_capture(ctx_tgt, ctx_dft, slot->id, std::move(shape)); + page->lora = slot->lora; + // the slot's prompt moves only into a stored page; the slot is freed after + res->gen = page_store.put(name, page, &slot->prompt); slot->prompt_clear(); - res->gen = page_store.put(name, page); metrics.n_page_out++; } break; case SERVER_TASK_TYPE_PAGE_IN: @@ -2600,7 +2602,7 @@ struct server_context_impl { if (!page) { throw server_page_error(ERROR_TYPE_NOT_FOUND, "no RAM page named '" + name + "'"); } - if (page_store.erase(name, task.slot_action.page_gen) == 0) { + if (!page_store.erase(name, task.slot_action.page_gen)) { throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the page under this name was replaced; nothing dropped"); } res->gen = page->gen; @@ -2615,13 +2617,15 @@ struct server_context_impl { res->n_checkpoints = held.checkpoints.size(); res->n_bytes_checkpoints = server_checkpoints_bytes(held.checkpoints); res->n_bytes_draft = page->data.drft.size(); - res->n_bytes = page->data.size() + page->spec.size() + res->n_bytes_checkpoints; + res->n_bytes_spec = page->spec.size(); + res->n_bytes = page->charge.n; + page.reset(); // the store's count below includes this page only while another owner holds it res->store_bytes = page_store.n_bytes(); res->store_pages = page_store.n_pages(); res->t_ms = (ggml_time_us() - t_start) / 1000.0; - SRV_INF("%s '%s' slot %d: %zu tokens, %zu bytes (%zu draft, %zu in %zu checkpoints) in %.3f ms; store %zu bytes in %zu pages\n", - res->action.c_str(), name.c_str(), res->id_slot, res->n_tokens, res->n_bytes, res->n_bytes_draft, + SRV_INF("%s '%s' slot %d: %zu tokens, %zu bytes (%zu draft, %zu speculative, %zu in %zu checkpoints) in %.3f ms; store %zu bytes in %zu pages\n", + res->action.c_str(), name.c_str(), res->id_slot, res->n_tokens, res->n_bytes, res->n_bytes_draft, res->n_bytes_spec, res->n_bytes_checkpoints, res->n_checkpoints, res->t_ms, res->store_bytes, res->store_pages); return res; } @@ -2921,38 +2925,25 @@ struct server_context_impl { break; } } - const size_t nwrite = llama_state_seq_save_file( - ctx_tgt, filepath.c_str(), slot->id, - reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); - if (nwrite == 0) { - send_error(task, "Unable to save slot", ERROR_TYPE_SERVER); - break; - } - // Companion serialization is shared with the pre-write size quote. - { - std::ofstream out(filepath + ".ckpt", std::ios::binary | std::ios::trunc); - write_checkpoints([&](const void * ptr, size_t n) { out.write(reinterpret_cast(ptr), n); }); - if (!out) { - SLT_WRN(*slot, "checkpoint sidecar could not be written: %s (the restore will be cold)\n", filepath.c_str()); - } - } - { - std::ofstream out(filepath + ".lora", std::ios::binary | std::ios::trunc); - write_adapters([&](const void * ptr, size_t n) { out.write(reinterpret_cast(ptr), n); }); - if (!out) { - SLT_WRN(*slot, "adapter sidecar could not be written: %s (this page will not restore)\n", filepath.c_str()); - } - } + // one generation, published whole: a failed save leaves the page that was there + size_t nwrite = 0; + std::vector files = { + { "", [&](const std::string & path) { + nwrite = llama_state_seq_save_file(ctx_tgt, path.c_str(), slot->id, + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + return nwrite; + } }, + server_page_sidecar(".ckpt", write_checkpoints), + server_page_sidecar(".lora", write_adapters), + }; if (ctx_dft) { - std::ofstream out(filepath + ".dft", std::ios::binary | std::ios::trunc); - write_draft([&](const void * ptr, size_t n) { out.write(reinterpret_cast(ptr), n); }); - if (!out) { - SLT_WRN(*slot, "draft sidecar could not be written: %s (the restore will draft cold)\n", filepath.c_str()); - } - } else { - // a draft state saved under this name by an earlier server is not this page's - std::error_code ec; - std::filesystem::remove(filepath + ".dft", ec); + files.push_back(server_page_sidecar(".dft", write_draft)); // without one, an older `.dft` is removed + } + try { + server_page_publish(filepath, files); + } catch (const server_page_error & err) { + send_error(task, std::string("Unable to save slot: ") + err.what(), err.type); + break; } const int64_t t_end = ggml_time_us(); @@ -2998,6 +2989,8 @@ struct server_context_impl { size_t nread = 0; try { + // Continuum fork: only a committed generation is read, and it stays one while it is + const server_page_reader committed(filepath); // Continuum fork: the adapters the state was computed under (see SLOT_SAVE), read // FIRST so a page whose identity is unknown or foreign never touches the KV. A // page without them cannot be reused safely (its KV may belong to other adapters), diff --git a/tools/server/server-page-files.cpp b/tools/server/server-page-files.cpp new file mode 100644 index 000000000000..5a136639ae6d --- /dev/null +++ b/tools/server/server-page-files.cpp @@ -0,0 +1,444 @@ +#include "server-page-files.h" + +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +# include +# include +# include +#else +# include +# include +# include +# include +#endif + +namespace fs = std::filesystem; + +static constexpr uint32_t PAGE_CKPT_MAGIC = 0x434B5054u; // 'CKPT' +static constexpr uint32_t PAGE_LORA_MAGIC = 0x4C4F5241u; // 'LORA' +static constexpr uint32_t PAGE_DRFT_MAGIC = 0x44524654u; // 'DRFT' + +// +// page files +// + +std::pair lora_file_revision(const std::string & path) { + std::error_code ec; + const auto size = std::filesystem::file_size(path, ec); + if (ec) { + return {0, 0}; + } + const auto mtime = std::filesystem::last_write_time(path, ec); + if (ec) { + return {0, 0}; + } + return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; +} + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put) { + const uint32_t magic = PAGE_CKPT_MAGIC; + const uint32_t version = 1; + const uint32_t n_ckpt = (uint32_t) kept.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); + for (const auto * c : kept) { + const int64_t n_tokens = c->n_tokens; + const int32_t pos_min = c->pos_min; + const int32_t pos_max = c->pos_max; + const uint64_t sz_tgt = c->data_tgt.size(); + const uint64_t sz_dft = c->data_dft.size(); + const uint64_t sz_spec = c->data_spec.size(); + put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); + put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(c->data_tgt.data(), c->data_tgt.size()); + put(c->data_dft.data(), c->data_dft.size()); + put(c->data_spec.data(), c->data_spec.size()); + } +} + +void server_page_write_adapters(const std::vector & lora, const server_page_put & put) { + const uint32_t magic = PAGE_LORA_MAGIC; + const uint32_t version = 2; + const uint32_t n_lora = (uint32_t) lora.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); + for (const auto & l : lora) { + const uint32_t n_path = (uint32_t) l.path.size(); + const auto rev = lora_file_revision(l.path); + put(&n_path, sizeof(n_path)); + put(l.path.data(), n_path); + put(&l.scale, sizeof(l.scale)); + put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); + } +} + +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put) { + const uint32_t magic = PAGE_DRFT_MAGIC; + const uint32_t version = 1; + const uint64_t sz_dft = n_dft; + const uint64_t sz_spec = spec.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); + put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(dft, n_dft); + put(spec.data(), spec.size()); +} + +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why) { + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) { + why = "absent"; + return false; + } + const auto total = (uint64_t) in.tellg(); + in.seekg(0); + uint32_t magic = 0, version = 0; + uint64_t sz_dft = 0, sz_spec = 0; + in.read(reinterpret_cast(&magic), sizeof(magic)); + in.read(reinterpret_cast(&version), sizeof(version)); + in.read(reinterpret_cast(&sz_dft), sizeof(sz_dft)); + in.read(reinterpret_cast(&sz_spec), sizeof(sz_spec)); + const uint64_t header = 2 * sizeof(uint32_t) + 2 * sizeof(uint64_t); + if (!in || magic != PAGE_DRFT_MAGIC || version != 1 || sz_dft > total || sz_spec > total || header + sz_dft + sz_spec != total) { + why = "malformed"; + return false; + } + dft.resize(sz_dft); + spec.resize(sz_spec); + in.read(reinterpret_cast(dft.data()), sz_dft); + in.read(reinterpret_cast(spec.data()), sz_spec); + if (!in) { + why = "truncated"; + return false; + } + return true; +} + +// +// generation publication +// + +// Every file a page may have. A generation names the ones it carries; the others are removed at its commit. +static const char * const PAGE_SUFFIXES[] = { "", ".ckpt", ".lora", ".dft" }; + +// Beside the page dir; ".." in it is a name fs_validate_filename refuses, so no page can take it. +static const char * const STAGING_DIR = "..page-staging"; + +// Publications in this process: the commit boundary every reader shares, and the stages being written. +struct page_publications { + std::shared_mutex commit; // a commit (and a reader's repair) exclusive; a reader's read shared + std::mutex stages_mutex; + std::set stages; + uint64_t next = 1; +}; + +static page_publications & publications() { + static page_publications pubs; + return pubs; +} + +static long current_pid() { +#if defined(_WIN32) + return (long) _getpid(); +#else + return (long) getpid(); +#endif +} + +static bool process_alive(long pid) { +#if defined(_WIN32) + HANDLE h = OpenProcess(SYNCHRONIZE, FALSE, (DWORD) pid); + if (h == nullptr) { + return false; + } + const bool alive = WaitForSingleObject(h, 0) == WAIT_TIMEOUT; + CloseHandle(h); + return alive; +#else + return kill((pid_t) pid, 0) == 0 || errno == EPERM; +#endif +} + +// A stage is "-": live while its publication runs in this process, or while its process runs. +static bool stage_live(const std::string & id) { + const auto dash = id.find('-'); + long pid = -1; + try { + pid = std::stol(id.substr(0, dash)); + } catch (const std::exception &) { + return false; + } + if (pid <= 0) { + return false; + } + if (pid != current_pid()) { + return process_alive(pid); + } + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + return pubs.stages.count(id) > 0; +} + +// Makes the file's contents durable before anything that depends on them (the commit record). +static bool sync_file(const fs::path & path) { +#if defined(_WIN32) + const int fd = _wopen(path.c_str(), _O_RDWR | _O_BINARY); + if (fd < 0) { + return false; + } + const bool ok = _commit(fd) == 0; + _close(fd); + return ok; +#else + const int fd = open(path.c_str(), O_RDONLY); + if (fd < 0) { + return false; + } +# if defined(__APPLE__) + // macOS fsync stops at the drive's cache; only F_FULLFSYNC orders the data before the record + bool ok = fcntl(fd, F_FULLFSYNC) == 0 || fsync(fd) == 0; +# else + bool ok = fsync(fd) == 0; +# endif + ok = close(fd) == 0 && ok; + return ok; +#endif +} + +// Makes a dir's entries (a rename into it) durable. Windows has no handle for this; NTFS journals them. +static void sync_dir(const fs::path & dir) { +#if !defined(_WIN32) + const int fd = open(dir.c_str(), O_RDONLY); + if (fd >= 0) { + fsync(fd); + close(fd); + } +#else + (void) dir; +#endif +} + +static fs::path staging_dir(const fs::path & file) { + return file.parent_path() / STAGING_DIR; +} + +// Its presence is the commit: the stage it names holds the generation, whole and synced. +static fs::path commit_record(const fs::path & file) { + return staging_dir(file) / (file.filename().string() + ".commit"); +} + +// "\n" then "f\n" per file of the generation. +static bool read_commit_record(const fs::path & record, std::string & id, std::vector & suffixes) { + std::ifstream in(record); + if (!in || !std::getline(in, id) || id.empty()) { + return false; + } + for (std::string line; std::getline(in, line);) { + if (line.empty() || line[0] != 'f') { + return false; + } + suffixes.push_back(line.substr(1)); + } + return !suffixes.empty(); +} + +// Writes the record inside the stage, syncs it, then renames it into place: the rename is the commit. +static void write_commit_record(const fs::path & file, const fs::path & stage, const std::string & id, const std::vector & suffixes) { + const fs::path tmp = stage / "..commit"; // ".." is in no page name + { + std::ofstream out(tmp, std::ios::trunc); + out << id << '\n'; + for (const auto & s : suffixes) { + out << 'f' << s << '\n'; + } + out.flush(); + if (!out) { + throw server_page_error(ERROR_TYPE_SERVER, "the page's commit record could not be written: " + tmp.string()); + } + } + std::error_code ec; + if (!sync_file(tmp) || (fs::rename(tmp, commit_record(file), ec), ec)) { + throw server_page_error(ERROR_TYPE_SERVER, "the page's commit record could not be placed: " + commit_record(file).string()); + } + sync_dir(staging_dir(file)); +} + +// Moves a committed generation over the page's files, removes the files it lacks, then its stage and +// record. Idempotent: a file a crash already moved is not in the stage. Throws, keeping the record, so +// the next reader finishes it. +static void roll_forward(const fs::path & file, const fs::path & stage, const std::vector & suffixes) { + const std::string name = file.filename().string(); + std::error_code ec; + for (const auto & s : suffixes) { + const fs::path staged = stage / (name + s); + if (fs::exists(staged, ec) && (fs::rename(staged, file.string() + s, ec), ec)) { + throw server_page_error(ERROR_TYPE_SERVER, "a committed page file could not be moved into place: " + staged.string() + ": " + ec.message()); + } + } + for (const char * s : PAGE_SUFFIXES) { + if (std::find(suffixes.begin(), suffixes.end(), s) == suffixes.end() && !fs::remove(file.string() + s, ec) && ec) { + throw server_page_error(ERROR_TYPE_SERVER, "an older page file could not be removed: " + file.string() + s + ": " + ec.message()); + } + } + sync_dir(file.parent_path()); + fs::remove_all(stage, ec); + fs::remove(commit_record(file), ec); + sync_dir(staging_dir(file)); +} + +// Removes every stage no commit record names and no live publication owns: what a crash before a +// commit left. Called under the exclusive commit lock. +static void remove_orphan_stages(const fs::path & staging) { + std::error_code ec; + if (!fs::is_directory(staging, ec)) { + return; + } + std::set committed; + for (const auto & e : fs::directory_iterator(staging, ec)) { + std::string id; + std::vector suffixes; + if (e.path().extension() == ".commit" && read_commit_record(e.path(), id, suffixes)) { + committed.insert(id); + } + } + std::vector orphans; + for (const auto & e : fs::directory_iterator(staging, ec)) { + const std::string id = e.path().filename().string(); + if (e.is_directory(ec) && committed.count(id) == 0 && !stage_live(id)) { + orphans.push_back(e.path()); + } + } + for (const auto & p : orphans) { + SRV_WRN("removing a page generation a publication left uncommitted: %s\n", p.string().c_str()); + fs::remove_all(p, ec); + } +} + +// A publication's stage, live while it exists: uncommitted, it is removed with its owner; committed, it +// stays until its roll forward removes it. +struct stage_owner { + std::string id; + fs::path stage; + bool committed = false; + + explicit stage_owner(const fs::path & staging) { + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + id = std::to_string(current_pid()) + "-" + std::to_string(pubs.next++); + stage = staging / id; + pubs.stages.insert(id); + } + + ~stage_owner() { + std::error_code ec; + if (!committed) { + fs::remove_all(stage, ec); + } + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + pubs.stages.erase(id); + } + + stage_owner(const stage_owner &) = delete; + stage_owner & operator=(const stage_owner &) = delete; +}; + +server_page_file server_page_sidecar(std::string suffix, std::function body) { + return { std::move(suffix), [body = std::move(body)](const std::string & path) -> size_t { + std::ofstream out(path, std::ios::binary | std::ios::trunc); + size_t n = 0; + body([&](const void * ptr, size_t len) { + out.write(reinterpret_cast(ptr), len); + n += len; + }); + out.flush(); + return out ? n : 0; + } }; +} + +size_t server_page_publish(const std::string & filepath, const std::vector & files) { + const fs::path file(filepath); + auto & pubs = publications(); + + stage_owner owner(staging_dir(file)); + const fs::path & stage = owner.stage; + + std::error_code ec; + fs::create_directories(stage, ec); + if (ec) { + throw server_page_error(ERROR_TYPE_SERVER, "the page could not be staged: " + stage.string() + ": " + ec.message()); + } + size_t n_written = 0; + std::vector suffixes; + for (const auto & f : files) { + GGML_ASSERT(std::find(std::begin(PAGE_SUFFIXES), std::end(PAGE_SUFFIXES), f.suffix) != std::end(PAGE_SUFFIXES)); + const fs::path staged = stage / (file.filename().string() + f.suffix); + size_t n = 0; + try { + n = f.write(staged.string()); + } catch (const std::exception & err) { + throw server_page_error(ERROR_TYPE_SERVER, "a page file could not be written: " + staged.string() + ": " + err.what()); + } + if (n == 0 || !sync_file(staged)) { + throw server_page_error(ERROR_TYPE_SERVER, "a page file could not be written: " + staged.string()); + } + n_written += n; + suffixes.push_back(f.suffix); + } + sync_dir(stage); + + std::unique_lock lock(pubs.commit); + write_commit_record(file, stage, owner.id, suffixes); + owner.committed = true; + roll_forward(file, stage, suffixes); + return n_written; +} + +size_t server_page_spill(const server_page & page, const std::string & filepath) { + std::vector packed; + try { + packed = page.prompt.tokens.serialize(); + } catch (const std::exception & err) { + throw server_page_error(ERROR_TYPE_NOT_SUPPORTED, err.what()); + } + std::vector all; + for (const auto & c : page.prompt.checkpoints) { + all.push_back(&c); + } + return server_page_publish(filepath, { + { "", [&](const std::string & path) { + return llama_state_seq_save_file_from_data(path.c_str(), page.data.main.data(), page.data.main.size(), + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + } }, + server_page_sidecar(".ckpt", [&](const server_page_put & put) { server_page_write_checkpoints(all, put); }), + server_page_sidecar(".lora", [&](const server_page_put & put) { server_page_write_adapters(page.lora, put); }), + server_page_sidecar(".dft", [&](const server_page_put & put) { server_page_write_draft(page.data.drft.data(), page.data.drft.size(), page.spec, put); }), + }); +} + +server_page_reader::server_page_reader(const std::string & filepath) { + const fs::path file(filepath); + auto & pubs = publications(); + { + std::unique_lock repair(pubs.commit); + const fs::path record = commit_record(file); + std::string id; + std::vector suffixes; + if (read_commit_record(record, id, suffixes)) { + SRV_WRN("completing the publication of %s that stopped after its commit\n", filepath.c_str()); + roll_forward(file, staging_dir(file) / id, suffixes); + } + remove_orphan_stages(staging_dir(file)); + } + lock = std::shared_lock(pubs.commit); +} diff --git a/tools/server/server-page-files.h b/tools/server/server-page-files.h new file mode 100644 index 000000000000..ea12dfd9ce1c --- /dev/null +++ b/tools/server/server-page-files.h @@ -0,0 +1,67 @@ +#pragma once + +// Continuum fork: a page on disk, `` (sequence state) and its sidecars `.ckpt` (context +// checkpoints), `.lora` (the adapters the state was computed under) and `.dft` (draft state), +// shared by save, spill and restore. +// +// The files of a page are one generation, published whole: a writer stages every file beside the +// page dir and syncs them, writes a commit record, then moves them over the committed files. A +// failure before the record leaves the committed generation as it was; a crash after it is +// completed by the next reader. Restore reads only through a server_page_reader, so it sees one +// committed generation, never a mix. Two servers publishing the SAME name into one dir at once are +// not ordered against each other: a page is seated in one server at a time. + +#include "server-pages.h" + +#include +#include +#include +#include +#include +#include + +using server_page_put = std::function; + +// An adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it cannot be read), +// stored in the `.lora` sidecar so a restore refuses a page computed under different weights. +std::pair lora_file_revision(const std::string & path); + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put); // .ckpt +void server_page_write_adapters(const std::vector & lora, const server_page_put & put); // .lora +// `dft` may be null when `put` only counts bytes +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put); // .dft + +// Reads a `.dft` sidecar. Returns false with `why` set when it is absent or malformed. +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why); + +// One file of a generation: its suffix ("" is the sequence state; one of "", ".ckpt", ".lora", +// ".dft") and the writer that fills the staged `path`, returning the bytes written (0 = failed). +struct server_page_file { + std::string suffix; + std::function write; +}; + +// A sidecar written by one of the serializers above. +server_page_file server_page_sidecar(std::string suffix, std::function body); + +// Publishes `files` as the generation of the page at `filepath`, whole or not at all; a page file +// the generation lacks is removed at the commit. Throws server_page_error; on any failure before +// the commit record the committed generation is untouched and nothing staged is left. Returns the +// bytes written. +size_t server_page_publish(const std::string & filepath, const std::vector & files); + +// Writes `page` to `filepath` (+ .ckpt, .lora, .dft) with every checkpoint, as one generation. +// Touches no context, so it runs off the task thread. Returns the bytes written. +size_t server_page_spill(const server_page & page, const std::string & filepath); + +// Holds the committed generation at `filepath` still while it is read. Constructing one completes a +// publication a crash cut short after its commit record and removes staged files no live publication +// owns (a crash before the commit), then keeps publication off until it is destroyed. Throws +// server_page_error when an interrupted commit cannot be completed. +class server_page_reader { +public: + explicit server_page_reader(const std::string & filepath); + +private: + std::shared_lock lock; +}; diff --git a/tools/server/server-pages.cpp b/tools/server/server-pages.cpp index 9a8ff43474e0..2a5d98abcf27 100644 --- a/tools/server/server-pages.cpp +++ b/tools/server/server-pages.cpp @@ -1,16 +1,18 @@ #include "server-pages.h" -#include "../../src/llama-ext.h" #include "speculative.h" -#include -#include -#include -#include +#include -static constexpr uint32_t PAGE_CKPT_MAGIC = 0x434B5054u; // 'CKPT' -static constexpr uint32_t PAGE_LORA_MAGIC = 0x4C4F5241u; // 'LORA' -static constexpr uint32_t PAGE_DRFT_MAGIC = 0x44524654u; // 'DRFT' +size_t server_page_bytes(size_t n_state, size_t n_spec, size_t n_checkpoint_bytes) { + return n_state + n_spec + n_checkpoint_bytes; +} + +server_page_charge::~server_page_charge() { + if (held) { + held->fetch_sub(n); + } +} size_t server_page::n_checkpoint_bytes() const { size_t n = 0; @@ -21,23 +23,26 @@ size_t server_page::n_checkpoint_bytes() const { } size_t server_page::n_bytes() const { - return data.size() + spec.size() + n_checkpoint_bytes(); + return server_page_bytes(data.size(), spec.size(), n_checkpoint_bytes()); } -std::pair server_page_state_size(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq) { - return { - llama_state_seq_get_size_ext(ctx_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE), - ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) : 0, - }; +server_page_shape server_page_measure(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq) { + server_page_shape shape; + shape.n_tgt = llama_state_seq_get_size_ext(ctx_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE); + shape.n_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + common_speculative_get_state(spec, seq, shape.spec); + return shape; } -server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq) { - const auto [n_tgt, n_dft] = server_page_state_size(ctx_tgt, ctx_dft, seq); +server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, server_page_shape shape) { + const size_t n_tgt = shape.n_tgt; + const size_t n_dft = shape.n_dft; if (n_tgt == 0) { throw server_page_error(ERROR_TYPE_SERVER, "the target sequence state could not be sized"); } - auto page = std::make_shared(); + server_page_ptr page; try { + page = std::make_shared(); page->data.main.resize(n_tgt); page->data.drft.resize(n_dft); } catch (const std::bad_alloc &) { @@ -49,9 +54,7 @@ server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx if (n_dft > 0 && llama_state_seq_get_data_ext(ctx_dft, page->data.drft.data(), n_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) != n_dft) { throw server_page_error(ERROR_TYPE_SERVER, "the draft sequence state read short"); } - if (spec != nullptr) { - common_speculative_get_state(spec, seq, page->spec); - } + page->spec = std::move(shape.spec); return page; } @@ -84,26 +87,50 @@ void server_page_store::admit(const std::string & name, size_t n_bytes_new) cons if (limit_bytes == 0) { return; } - size_t held = n_bytes(); + size_t live = n_bytes(); const auto it = pages.find(name); - if (it != pages.end()) { - held -= it->second->n_bytes(); + // only this thread makes owners, so a sole owner seen here stays sole until the replacement frees it + if (it != pages.end() && it->second.use_count() == 1) { + live -= it->second->charge.n; } - if (held + n_bytes_new > limit_bytes) { + if (live + n_bytes_new > limit_bytes) { throw server_page_error(ERROR_TYPE_UNAVAILABLE, string_format( - "the page store holds %zu bytes; a page of %zu bytes would exceed its ceiling of %llu bytes: spill or drop a page first", - held, n_bytes_new, (unsigned long long) limit_bytes)); + "the live pages hold %zu bytes; a page of %zu bytes would exceed the ceiling of %llu bytes: spill or drop a page first", + live, n_bytes_new, (unsigned long long) limit_bytes)); } } -uint64_t server_page_store::put(const std::string & name, server_page_ptr page) { - page->gen = next_gen++; - const auto it = pages.find(name); - if (it != pages.end()) { - SRV_WRN("page '%s' replaced (%zu bytes freed)\n", name.c_str(), it->second->n_bytes()); +uint64_t server_page_store::put(const std::string & name, const server_page_ptr & page, server_prompt * from) { + GGML_ASSERT(page->charge.held == nullptr && "a page is stored once"); + const auto give_back = [&]() { + if (from != nullptr) { + *from = std::move(page->prompt); + } + }; + if (from != nullptr) { + page->prompt = std::move(*from); } - pages[name] = std::move(page); - return pages[name]->gen; + std::map::iterator it; + try { + admit(name, page->n_bytes()); + it = pages.try_emplace(name).first; + if (it->second) { + SRV_WRN("page '%s' replaced (%zu bytes, freed when no spill holds them)\n", name.c_str(), it->second->charge.n); + } + } catch (const std::bad_alloc &) { + give_back(); + throw server_page_error(ERROR_TYPE_UNAVAILABLE, "host RAM for the page's entry could not be allocated"); + } catch (...) { + give_back(); + throw; + } + // nothing below can fail: the page is stored whole or not at all + page->gen = next_gen++; + page->charge.n = page->n_bytes(); + page->charge.held = held; + held->fetch_add(page->charge.n); + it->second = page; + return page->gen; } server_page_ptr server_page_store::find(const std::string & name) const { @@ -111,14 +138,13 @@ server_page_ptr server_page_store::find(const std::string & name) const { return it == pages.end() ? nullptr : it->second; } -size_t server_page_store::erase(const std::string & name, uint64_t gen) { +bool server_page_store::erase(const std::string & name, uint64_t gen) { const auto it = pages.find(name); if (it == pages.end() || (gen != 0 && it->second->gen != gen)) { - return 0; + return false; } - const size_t n = it->second->n_bytes(); pages.erase(it); - return n; + return true; } void server_page_store::evict_lora(const llama_adapter_lora * adapter) { @@ -128,7 +154,7 @@ void server_page_store::evict_lora(const llama_adapter_lora * adapter) { uses = uses || (l.ptr == adapter && l.scale != 0.0f); } if (uses) { - SRV_WRN("page '%s' dropped: its adapter is being unloaded (%zu bytes freed)\n", it->first.c_str(), it->second->n_bytes()); + SRV_WRN("page '%s' dropped: its adapter is being unloaded (%zu bytes)\n", it->first.c_str(), it->second->charge.n); it = pages.erase(it); } else { ++it; @@ -136,14 +162,6 @@ void server_page_store::evict_lora(const llama_adapter_lora * adapter) { } } -size_t server_page_store::n_bytes() const { - size_t n = 0; - for (const auto & [name, page] : pages) { - n += page->n_bytes(); - } - return n; -} - bool server_page_name_valid(const std::string & name) { if (name.empty() || name.size() > 255) { return false; @@ -155,147 +173,3 @@ bool server_page_name_valid(const std::string & name) { } return true; } - -// -// page files -// - -std::pair lora_file_revision(const std::string & path) { - std::error_code ec; - const auto size = std::filesystem::file_size(path, ec); - if (ec) { - return {0, 0}; - } - const auto mtime = std::filesystem::last_write_time(path, ec); - if (ec) { - return {0, 0}; - } - return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; -} - -void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put) { - const uint32_t magic = PAGE_CKPT_MAGIC; - const uint32_t version = 1; - const uint32_t n_ckpt = (uint32_t) kept.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); - for (const auto * c : kept) { - const int64_t n_tokens = c->n_tokens; - const int32_t pos_min = c->pos_min; - const int32_t pos_max = c->pos_max; - const uint64_t sz_tgt = c->data_tgt.size(); - const uint64_t sz_dft = c->data_dft.size(); - const uint64_t sz_spec = c->data_spec.size(); - put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); - put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); - put(c->data_tgt.data(), c->data_tgt.size()); - put(c->data_dft.data(), c->data_dft.size()); - put(c->data_spec.data(), c->data_spec.size()); - } -} - -void server_page_write_adapters(const std::vector & lora, const server_page_put & put) { - const uint32_t magic = PAGE_LORA_MAGIC; - const uint32_t version = 2; - const uint32_t n_lora = (uint32_t) lora.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); - for (const auto & l : lora) { - const uint32_t n_path = (uint32_t) l.path.size(); - const auto rev = lora_file_revision(l.path); - put(&n_path, sizeof(n_path)); - put(l.path.data(), n_path); - put(&l.scale, sizeof(l.scale)); - put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); - } -} - -void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put) { - const uint32_t magic = PAGE_DRFT_MAGIC; - const uint32_t version = 1; - const uint64_t sz_dft = n_dft; - const uint64_t sz_spec = spec.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); - put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); - put(dft, n_dft); - put(spec.data(), spec.size()); -} - -bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why) { - std::ifstream in(path, std::ios::binary | std::ios::ate); - if (!in) { - why = "absent"; - return false; - } - const auto total = (uint64_t) in.tellg(); - in.seekg(0); - uint32_t magic = 0, version = 0; - uint64_t sz_dft = 0, sz_spec = 0; - in.read(reinterpret_cast(&magic), sizeof(magic)); - in.read(reinterpret_cast(&version), sizeof(version)); - in.read(reinterpret_cast(&sz_dft), sizeof(sz_dft)); - in.read(reinterpret_cast(&sz_spec), sizeof(sz_spec)); - const uint64_t header = 2 * sizeof(uint32_t) + 2 * sizeof(uint64_t); - if (!in || magic != PAGE_DRFT_MAGIC || version != 1 || sz_dft > total || sz_spec > total || header + sz_dft + sz_spec != total) { - why = "malformed"; - return false; - } - dft.resize(sz_dft); - spec.resize(sz_spec); - in.read(reinterpret_cast(dft.data()), sz_dft); - in.read(reinterpret_cast(spec.data()), sz_spec); - if (!in) { - why = "truncated"; - return false; - } - return true; -} - -static bool write_sidecar(const std::string & path, const std::function & body, size_t & n_written) { - std::ofstream out(path, std::ios::binary | std::ios::trunc); - body([&](const void * ptr, size_t n) { - out.write(reinterpret_cast(ptr), n); - n_written += n; - }); - out.flush(); - return (bool) out; -} - -size_t server_page_spill(const server_page & page, const std::string & filepath) { - const std::string sidecars[] = { filepath + ".ckpt", filepath + ".lora", filepath + ".dft" }; - auto remove_all = [&]() { - std::error_code ec; - std::filesystem::remove(filepath, ec); - for (const auto & p : sidecars) { - std::filesystem::remove(p, ec); - } - }; - // no file of an older page under this name may survive beside the new state - remove_all(); - - std::vector packed; - try { - packed = page.prompt.tokens.serialize(); - } catch (const std::exception & err) { - throw server_page_error(ERROR_TYPE_NOT_SUPPORTED, err.what()); - } - size_t n_written = llama_state_seq_save_file_from_data(filepath.c_str(), - page.data.main.data(), page.data.main.size(), - reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); - if (n_written == 0) { - remove_all(); - throw server_page_error(ERROR_TYPE_SERVER, "the page's sequence state could not be written: " + filepath); - } - - std::vector all; - for (const auto & c : page.prompt.checkpoints) { - all.push_back(&c); - } - const bool ok = - write_sidecar(sidecars[0], [&](const server_page_put & put) { server_page_write_checkpoints(all, put); }, n_written) && - write_sidecar(sidecars[1], [&](const server_page_put & put) { server_page_write_adapters(page.lora, put); }, n_written) && - write_sidecar(sidecars[2], [&](const server_page_put & put) { server_page_write_draft(page.data.drft.data(), page.data.drft.size(), page.spec, put); }, n_written); - if (!ok) { - remove_all(); - throw server_page_error(ERROR_TYPE_SERVER, "a page sidecar could not be written beside " + filepath); - } - return n_written; -} diff --git a/tools/server/server-pages.h b/tools/server/server-pages.h index 5ce42b33a6d5..09b741a9f75d 100644 --- a/tools/server/server-pages.h +++ b/tools/server/server-pages.h @@ -6,7 +6,7 @@ // (70-160 s): `page_out` moves a slot's target KV, draft (MTP) KV, speculative tail state // and context checkpoints into the store and frees the slot; `page_in` puts them back into // any slot with its tokens, so the next request's prefix matches. `spill` writes a page in -// the slot-save format (.bin + .ckpt + .lora + .dft) that `restore` reads. +// the slot-save format (server-page-files.h) that `restore` reads. // Design: continuum docs/architecture/SERVING-ENGINE.md, slice S1. #include "server-common.h" @@ -15,8 +15,8 @@ #include "common.h" #include "llama.h" +#include #include -#include #include #include #include @@ -32,6 +32,23 @@ struct server_page_error : std::runtime_error { server_page_error(error_type type, const std::string & msg) : std::runtime_error(msg), type(type) {} }; +// The host RAM a page holds: its target + draft sequence state, its speculative tail state and its +// context checkpoints. The one count: page_out admits a page with it before copying anything, and +// the store charges a stored page with it. +size_t server_page_bytes(size_t n_state, size_t n_spec, size_t n_checkpoint_bytes); + +// The bytes the store counts for one page, released when the page's LAST owner drops it: a spill +// writing on an HTTP thread keeps a page alive after it left the store, and its RAM is still held. +struct server_page_charge { + std::shared_ptr> held; // the store's live-bytes ledger; null until stored + size_t n = 0; + + server_page_charge() = default; + server_page_charge(const server_page_charge &) = delete; + server_page_charge & operator=(const server_page_charge &) = delete; + ~server_page_charge(); +}; + struct server_page { uint64_t gen = 0; // identity within the store: a page replaced under the same name is another page @@ -41,71 +58,69 @@ struct server_page { std::vector lora; // the slot's adapter list the state was computed under + server_page_charge charge; + size_t n_checkpoint_bytes() const; - size_t n_bytes() const; + size_t n_bytes() const; // server_page_bytes of what the page holds }; using server_page_ptr = std::shared_ptr; -// The sequence state bytes a page of `seq` will hold (target, draft), before it is taken. -std::pair server_page_state_size(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq); +// A page of `seq` sized before anything is copied: what admission counts, then what capture takes. +struct server_page_shape { + size_t n_tgt = 0; + size_t n_dft = 0; + std::vector spec; // the speculative tail state (common_speculative_get_state), small: taken now + + size_t n_bytes(size_t n_checkpoint_bytes) const { return server_page_bytes(n_tgt + n_dft, spec.size(), n_checkpoint_bytes); } +}; + +server_page_shape server_page_measure(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq); -// Copies `seq`'s state out of both contexts (each get_data synchronizes its context first) and -// the speculative tail state; the caller moves in the prompt and adapters. Throws on a short read. -server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq); +// Copies `seq`'s state out of both contexts (each get_data synchronizes its context first): exactly +// the bytes `shape` measured. The caller moves in the prompt and adapters. Throws on a short read. +server_page_ptr server_page_capture(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, server_page_shape shape); // Writes a page's state into `seq` of both contexts. The page is left intact (it may be spilling). // Throws server_page_error; the caller clears the slot so a failure is never a silent half-state. void server_page_restore(const server_page & page, llama_context * ctx_tgt, llama_context * ctx_dft, common_speculative * spec, llama_seq_id seq); -// The named host-RAM page store. Owned by the server's task thread. +// The named host-RAM page store. Owned by the server's task thread: only that thread finds pages and +// so makes new owners; a spill's thread only drops its own. The ceiling counts every live page, +// stored or still held by a spill. struct server_page_store { uint64_t limit_bytes = 0; // --page-ram-bytes; 0 = no ceiling // Refuses (server_page_error, UNAVAILABLE) a page of `n_bytes` under `name` that would take the - // store past its ceiling; a page it replaces under the same name is not counted. + // live pages past the ceiling. The page it replaces is not counted when the store is its only + // owner (the replacement frees it); a spill-held one stays live and is counted. void admit(const std::string & name, size_t n_bytes) const; - // Stores `page` under `name`, replacing any page there; returns its generation. - uint64_t put(const std::string & name, server_page_ptr page); + // Stores `page` under `name`, replacing any page there; returns its generation. The ceiling is + // checked on exactly what is stored. `from` (a slot's prompt), when given, moves into the page + // first; a refused or failed put moves it back and throws, leaving the caller's state as it was. + uint64_t put(const std::string & name, const server_page_ptr & page, server_prompt * from = nullptr); server_page_ptr find(const std::string & name) const; - // Removes the page under `name` when it is generation `gen` (0 = any); returns the bytes freed. - size_t erase(const std::string & name, uint64_t gen = 0); + // Removes the page under `name` when it is generation `gen` (0 = any); returns whether it did. + // Its bytes stay counted until its last owner drops it. + bool erase(const std::string & name, uint64_t gen = 0); // Drops every page computed with `adapter` active: the adapter is about to be freed. void evict_lora(const llama_adapter_lora * adapter); - size_t n_bytes() const; + // Drops every stored page; pages a spill still holds stay counted until it finishes. + void clear() { pages.clear(); } + + size_t n_bytes() const { return held->load(); } // every live page, stored or spilling size_t n_pages() const { return pages.size(); } private: std::map pages; + std::shared_ptr> held = std::make_shared>(0); uint64_t next_gen = 1; }; // A page name: a non-empty opaque key of at most 255 bytes without control characters. bool server_page_name_valid(const std::string & name); - -// -// the slot page files: `` (sequence state) and its sidecars, shared by save, spill and restore -// - -using server_page_put = std::function; - -// An adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it cannot be read), -// stored in the `.lora` sidecar so a restore refuses a page computed under different weights. -std::pair lora_file_revision(const std::string & path); - -void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put); // .ckpt -void server_page_write_adapters(const std::vector & lora, const server_page_put & put); // .lora -// `dft` may be null when `put` only counts bytes -void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put); // .dft - -// Reads a `.dft` sidecar. Returns false with `why` set when it is absent or malformed. -bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why); - -// Writes `page` to `filepath` (+ .ckpt, .lora, .dft) with every checkpoint. Touches no context, so it -// runs off the task thread. All four files are written or none are left. Returns the bytes written. -size_t server_page_spill(const server_page & page, const std::string & filepath); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 32a1502a6e31..12c1e5a008fe 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1719,6 +1719,7 @@ json server_task_result_page::to_json() { { "n_tokens", n_tokens }, { "n_bytes", n_bytes }, { "n_bytes_draft", n_bytes_draft }, + { "n_bytes_spec", n_bytes_spec }, { "n_checkpoints", n_checkpoints }, { "n_bytes_checkpoints", n_bytes_checkpoints }, { "store_bytes", store_bytes }, diff --git a/tools/server/server-task.h b/tools/server/server-task.h index ff2b8fad6846..d9f420121944 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -636,6 +636,7 @@ struct server_task_result_page : server_task_result { size_t n_tokens = 0; size_t n_bytes = 0; // the page: target + draft state + speculative state + checkpoints size_t n_bytes_draft = 0; + size_t n_bytes_spec = 0; // the speculative tail state size_t n_checkpoints = 0; size_t n_bytes_checkpoints = 0; double t_ms = 0.0; diff --git a/tools/server/tests/unit/test_slot_page.py b/tools/server/tests/unit/test_slot_page.py index 29dfb20a83af..069aff494fa7 100644 --- a/tools/server/tests/unit/test_slot_page.py +++ b/tools/server/tests/unit/test_slot_page.py @@ -1,4 +1,6 @@ import os +import subprocess +import sys from pathlib import Path import pytest @@ -136,6 +138,36 @@ def test_spill_writes_the_slot_save_format(): assert page(1, "spill", "mind-a", {"filename": "../x.bin"}).status_code == 400 +def test_spill_publishes_whole_and_restore_cleans_a_crashed_one(): + # what this catches: a spill that leaves a stage beside the page dir, or a restore that reads + # a generation a crashed publication left uncommitted (the committed page must be what restores, + # and the leftovers must go). + global server + server.start() + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200 + assert page(1, "spill", "mind-a", {"filename": "mind-a.bin", "drop": True}).status_code == 200 + staging = Path(server.slot_save_path) / "..page-staging" + assert not staging.exists() or not any(staging.iterdir()) + + # a server that died between writing its files and committing them: its pid is gone + dead = subprocess.Popen([sys.executable, "-c", "pass"]) + dead.wait() + stage = staging / f"{dead.pid}-1" + stage.mkdir(parents=True) + for suffix in ("", ".ckpt", ".lora", ".dft"): + (stage / ("mind-a.bin" + suffix)).write_bytes(b"torn") + + res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "mind-a.bin"}) + assert res.status_code == 200, res.body + assert res.body["n_restored"] == out.body["n_tokens"] + assert not stage.exists() + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 + assert again["content"] == first["content"] + + HYBRID = os.environ.get("LLAMA_SERVER_TEST_HYBRID_GGUF") @@ -169,3 +201,51 @@ def test_checkpoints_stay_within_the_byte_budget_and_page_with_the_slot(): complete("A different tale: " + long_prompt, 1) assert metric("ctx_checkpoint_bytes") <= budget assert metric("ctx_checkpoints_evicted_budget_total") > 0 + + +@pytest.mark.skipif(not HYBRID, reason="set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF with an MTP head") +def test_page_with_draft_state_fits_its_ceiling_exactly(): + # what this catches: a page_out admitted on a count that leaves out part of what the page holds + # (draft or speculative state), so the store passes --page-ram-bytes after a successful admission. + # At a ceiling of exactly the page's bytes it is stored; one byte less and it is refused, the slot kept. + global server + + def mtp_server(ceiling): + global server + server = ServerPreset.tinyllama2() + server.slot_save_path = "./tmp" + server.temperature = 0.0 + server.server_metrics = True + server.model_hf_repo = None + server.model_file = HYBRID + server.n_ctx = 4096 + server.n_batch = 512 + server.spec_type = "draft-mtp" + server.page_ram_bytes = ceiling + server.start() + + mtp_server(None) + complete(PROMPT, 0) + out = page(0, "page_out", "mind-a") + assert out.status_code == 200, out.body + n_bytes = out.body["n_bytes"] + assert out.body["n_bytes_draft"] > 0 # the draft (MTP) state rides in the page + assert n_bytes == out.body["store_bytes"] + assert metric("page_store_bytes") == pytest.approx(n_bytes, rel=1e-5) # the gauge prints 6 digits + server.stop() + + mtp_server(n_bytes - 1) + first = complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 503 + again = complete(PROMPT, 0) + # the refused page_out kept the slot's sequence: a hybrid resumes from its kept checkpoint + # (a cleared slot would evaluate the whole prompt) + assert again["timings"]["prompt_n"] < again["tokens_evaluated"] + assert again["content"] == first["content"] + server.stop() + + mtp_server(n_bytes) + complete(PROMPT, 0) + out = page(0, "page_out", "mind-a") + assert out.status_code == 200, out.body + assert out.body["n_bytes"] == n_bytes == out.body["store_bytes"] From 77cfa0be3a6ba331f885d566aeaf5fada2fbdfa2 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sat, 24 Oct 2026 01:52:03 -0500 Subject: [PATCH 3/3] server: /reconfigure changes the context geometry without reloading the model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A relaunch is a pause, not a reload (continuum SERVING-ENGINE.md §3, slice S2, plan line 3b). POST /reconfigure {"n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force"}: admission closes (requests that arrive queue), the slots reach a ubatch boundary, every occupied slot's target KV, draft KV, speculative tail and checkpoints move to host RAM (S1's page capture), the draft and target contexts are freed and rebuilt on the SAME llama_model (weights and mmap stay), the slots are rebuilt for the new n_parallel, the states go back, and serving resumes. - A request in flight keeps its slot and finishes as if nothing happened (test: identical output). A request on a slot past the new n_parallel drains first. A request whose tokens do not fit the new window refuses the reconfigure; it is never truncated. - An idle mind with no seat moves into a free slot, or stays a named RAM page (page_in brings it back) with a typed reason; one larger than the new window is likewise kept whole. - A KV type change is an intentional cold start: refused without "force"; with it, resident states, the page store and the prompt cache are dropped, logged loudly and counted. - A geometry libllama cannot build rolls back to the old one with every state restored. - Receipt: drain, page out, free, init, page in and total ms; bytes; per-slot fate. /props context_geometry and total_slots follow the change (routes read an atomic meta snapshot); /metrics reconfigure_total, reconfigure_failed_total. A wake from sleep reloads the new geometry. - libllama: one sequence's state reads into a cache split into a different number of streams (its cells all land in its own stream); the whole-cache read still requires the same layout. - common: common_init_result::reinit_context/free_context and the speculative init result's rebuild the context on the loaded model; the context, sampler and threadpool setup are shared with the first build rather than copied. common_kv_cache_type_from_str is public. - The mechanism is tools/server/server-reconfigure.{h,cpp} behind server_reconfigure_host; server-context.cpp implements the host over its slots and wires the task, the route and the admission gate. Slot init is one function shared by load_model and the reconfigure. Measured (M5 Metal, this build): - Qwen3.5 0.8B hybrid, lane geometry (174,592 ctx, q4_0 KV, 32 checkpoints), two minds of 25,226 and 17,883 tokens: pause 72-104 ms (page out 6-13, free 3-5, init 57-80, page in 6) for 271 MB; median 78 ms over 40 cycles. Each mind's next act prefilled 15 tokens. - Footprint: 100 grow/shrink cycles (2 -> 3 slots, n_ctx x2 and back) 395 MB -> 327 MB, never growing; 20 cycles at lane geometry 1,685 -> 1,617 MB. - Qwen3.8 27B with the M5 lane's flags (mmproj, MTP draft model, q4_0 KV, 2 x 87,296), same two minds: 2 -> 3 slots 826 ms, 3 -> 4 (4 x 59,392) 802 ms, back to 2 626 ms, 2.25 GB paged; init (context + MTP context allocation) is 0.5-0.6 s of it. Next acts prefilled 15 tokens (6.4 s and 7.2 s wall, against 248 s and 173 s to seat them). CONTENDED: a second 27B serving lane started on the node during the run, so these are upper bounds; the 20-cycle 27B pauses (0.54-4.6 s, median 0.70 s) ran beside it too. Tests: tools/server/tests/unit/test_reconfigure.py: 5 passed (attention and hybrid), plus the opt-in 100-cycle footprint validator (LLAMA_SERVER_TEST_RECONFIGURE_CYCLES); with test_slot_page, test_slot_save, test_speculative: 27 passed. Full unit suite: 395 passed, 185 failed, the same 185 as S1's binary (test_tool_call real models not cached offline, one test_completion). Rebased on #70's review fixes: the pre-check counts each kept page with server_page_measure (the store's own count, speculative state included); a put the store refuses after a failed restore drops that state with a typed reason instead of aborting the reconfigure; a cold start clears the store (a spill-held page stays counted until its spill ends). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- common/arg.cpp | 10 +- common/arg.h | 3 + common/common.cpp | 121 +++-- common/common.h | 13 + common/speculative.cpp | 53 ++- common/speculative.h | 9 + src/llama-kv-cache.cpp | 7 +- tools/server/CMakeLists.txt | 2 + tools/server/server-common.h | 2 + tools/server/server-context.cpp | 465 ++++++++++++++++---- tools/server/server-context.h | 11 +- tools/server/server-reconfigure.cpp | 379 ++++++++++++++++ tools/server/server-reconfigure.h | 182 ++++++++ tools/server/server-task.cpp | 8 + tools/server/server-task.h | 5 + tools/server/server.cpp | 2 + tools/server/tests/unit/test_reconfigure.py | 281 ++++++++++++ 17 files changed, 1404 insertions(+), 149 deletions(-) create mode 100644 tools/server/server-reconfigure.cpp create mode 100644 tools/server/server-reconfigure.h create mode 100644 tools/server/tests/unit/test_reconfigure.py diff --git a/common/arg.cpp b/common/arg.cpp index 04f677503828..77fdd9ca4644 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -313,7 +313,7 @@ const std::vector kv_cache_types = { GGML_TYPE_Q5_1, }; -static ggml_type kv_cache_type_from_str(const std::string & s) { +ggml_type common_kv_cache_type_from_str(const std::string & s) { for (const auto & type : kv_cache_types) { if (ggml_type_name(type) == s) { return type; @@ -2461,7 +2461,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.cache_type_k) ), [](common_params & params, const std::string & value) { - params.cache_type_k = kv_cache_type_from_str(value); + params.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_K")); add_opt(common_arg( @@ -2474,7 +2474,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.cache_type_v) ), [](common_params & params, const std::string & value) { - params.cache_type_v = kv_cache_type_from_str(value); + params.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_V")); add_opt(common_arg( @@ -4085,7 +4085,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.speculative.draft.cache_type_k) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_k = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K")); add_opt(common_arg( @@ -4098,7 +4098,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.speculative.draft.cache_type_v) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_v = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V")); add_opt(common_arg( diff --git a/common/arg.h b/common/arg.h index 421bc295fc21..c82700bfc848 100644 --- a/common/arg.h +++ b/common/arg.h @@ -127,6 +127,9 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e // load all backends and print the list of available (non-CPU) devices to stdout void common_print_available_devices(); +// a KV cache type by its ggml name (the --cache-type-k/-v values); throws std::runtime_error otherwise +ggml_type common_kv_cache_type_from_str(const std::string & s); + // parse input arguments from CLI into a map bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map); diff --git a/common/common.cpp b/common/common.cpp index bda497246329..97e35d63df6d 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1271,6 +1271,35 @@ static void common_init_sampler_from_model( get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA), sparams.mirostat_eta, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA); } +// the per-context state a context takes from `params`: control vectors and the adapter set +static bool common_context_attach(llama_context * lctx, llama_model * model, common_params & params) { + if (!params.control_vectors.empty()) { + if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; + if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); + + const auto cvec = common_control_vector_load(params.control_vectors); + if (cvec.n_embd == -1) { + return false; + } + + int err = llama_set_adapter_cvec( + lctx, + cvec.data.data(), + cvec.data.size(), + cvec.n_embd, + params.control_vector_layer_start, + params.control_vector_layer_end); + if (err) { + return false; + } + } + + if (!params.lora_init_without_apply) { + common_set_adapter_lora(lctx, params.lora_adapters); + } + return true; +} + struct common_init_result::impl { impl() = default; ~impl() = default; @@ -1286,6 +1315,9 @@ struct common_init_result::impl { std::vector samplers; std::vector samplers_seq_config; + + // the backend samplers (one per sequence) and a context on `model`; replaces any context held + llama_context * make_context(common_params & params, llama_context_params cparams); }; common_init_result::common_init_result(common_params & params, bool model_only) : @@ -1387,31 +1419,65 @@ common_init_result::common_init_result(common_params & params, bool model_only) params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end()); } + llama_context * lctx = pimpl->make_context(params, cparams); + if (lctx == NULL) { + return; + } + + set_process_priority(params.cpuparams.priority); + + pimpl->threadpools.init(lctx, params); +} + +llama_context * common_init_result::impl::make_context(common_params & params, llama_context_params cparams) { + llama_model * mdl = model.get(); + // init the backend samplers as part of the context creation - pimpl->samplers.resize(cparams.n_seq_max); - pimpl->samplers_seq_config.resize(cparams.n_seq_max); + samplers.clear(); + samplers.resize(cparams.n_seq_max); + samplers_seq_config.resize(cparams.n_seq_max); for (int i = 0; i < (int) cparams.n_seq_max; ++i) { - pimpl->samplers[i].reset(common_sampler_init(model, params.sampling)); - pimpl->samplers_seq_config[i] = { i, common_sampler_get(pimpl->samplers[i].get()) }; + samplers[i].reset(common_sampler_init(mdl, params.sampling)); + samplers_seq_config[i] = { i, common_sampler_get(samplers[i].get()) }; } if (params.sampling.backend_sampling) { - cparams.samplers = pimpl->samplers_seq_config.data(); - cparams.n_samplers = pimpl->samplers_seq_config.size(); + cparams.samplers = samplers_seq_config.data(); + cparams.n_samplers = samplers_seq_config.size(); } - llama_context * lctx = llama_init_from_model(model, cparams); + llama_context * lctx = llama_init_from_model(mdl, cparams); if (lctx == NULL) { COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); - return; + return nullptr; } - pimpl->context.reset(lctx); + context.reset(lctx); + return lctx; +} - set_process_priority(params.cpuparams.priority); +void common_init_result::free_context() { + pimpl->context.reset(); +} - pimpl->threadpools.init(lctx, params); +llama_context * common_init_result::reinit_context(common_params & params) { + GGML_ASSERT(pimpl->model); + + // free first: the old and the new context never hold memory at the same time + pimpl->context.reset(); + + llama_context * lctx = pimpl->make_context(params, common_context_params_to_llama(params)); + if (lctx == nullptr) { + return nullptr; + } + pimpl->threadpools.attach(lctx); + + if (!common_context_attach(lctx, pimpl->model.get(), params)) { + pimpl->context.reset(); + return nullptr; + } + return lctx; } llama_model * common_init_result::model() { @@ -1465,25 +1531,8 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode params.ctx_shift = false; } - if (!params.control_vectors.empty()) { - if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; - if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); - - const auto cvec = common_control_vector_load(params.control_vectors); - if (cvec.n_embd == -1) { - return res; - } - - int err = llama_set_adapter_cvec( - lctx, - cvec.data.data(), - cvec.data.size(), - cvec.n_embd, - params.control_vector_layer_start, - params.control_vector_layer_end); - if (err) { - return res; - } + if (!common_context_attach(lctx, model, params)) { + return res; } if (llama_pooling_type(lctx) == LLAMA_POOLING_TYPE_RANK) { @@ -1510,10 +1559,6 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode } } - if (!params.lora_init_without_apply) { - common_set_adapter_lora(lctx, params.lora_adapters); - } - if (params.warmup) { COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n"); @@ -1829,7 +1874,13 @@ void common_threadpools::init(llama_context * ctx, const common_params & params) return; } - llama_attach_threadpool(ctx, threadpool, threadpool_batch); + attach(ctx); +} + +void common_threadpools::attach(llama_context * ctx) const { + if (threadpool != nullptr) { + llama_attach_threadpool(ctx, threadpool, threadpool_batch); + } } // diff --git a/common/common.h b/common/common.h index 9a2e6824fdac..aebff5f398ab 100644 --- a/common/common.h +++ b/common/common.h @@ -931,6 +931,15 @@ struct common_init_result { std::vector & lora(); + // Continuum fork: frees the context and builds one from `params` (n_ctx, n_parallel, KV types) + // on the SAME model: its weights, mmap and adapters stay loaded. Returns nullptr when the new + // context cannot be created; the result then holds no context and the caller may call it + // again with the old geometry. + llama_context * reinit_context(common_params & params); + + // Continuum fork: frees the context only; the model and adapters stay (reinit_context builds the next) + void free_context(); + private: struct impl; std::unique_ptr pimpl; @@ -967,6 +976,10 @@ struct common_threadpools { void init(llama_context * ctx, const common_params & params); + // Continuum fork: hands the pools init() made to another context of the same process (a + // context rebuilt with a new geometry keeps its CPU threads) + void attach(llama_context * ctx) const; + private: ggml_threadpool * threadpool = nullptr; ggml_threadpool * threadpool_batch = nullptr; diff --git a/common/speculative.cpp b/common/speculative.cpp index 4b67bbb88599..2f594f7c7e0b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2375,17 +2375,8 @@ struct common_speculative_init_result::impl { llama_context_ptr context; }; -common_speculative_init_result::common_speculative_init_result( - common_params & params, - llama_model * model_tgt, - llama_context * ctx_tgt) : - pimpl(new impl{}) { - const bool has_draft = params.speculative.has_dft(); - const bool spec_mtp = std::find(params.speculative.types.begin(), - params.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - - auto mparams = common_model_params_to_llama(params); +// the draft context's parameters: as many tokens per sequence as the target context, which it follows +static llama_context_params common_speculative_cparams(common_params & params, llama_context * ctx_tgt, bool spec_mtp) { auto cparams = common_context_params_to_llama(params); if (spec_mtp) { @@ -2399,6 +2390,24 @@ common_speculative_init_result::common_speculative_init_result( // the extra memory for small models is likely negligible? cparams.n_rs_seq = 0; cparams.ctx_other = ctx_tgt; + return cparams; +} + +static bool common_speculative_is_mtp(const common_params & params) { + return std::find(params.speculative.types.begin(), params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); +} + +common_speculative_init_result::common_speculative_init_result( + common_params & params, + llama_model * model_tgt, + llama_context * ctx_tgt) : + pimpl(new impl{}) { + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = common_speculative_is_mtp(params); + + auto mparams = common_model_params_to_llama(params); + auto cparams = common_speculative_cparams(params, ctx_tgt, spec_mtp); std::string model_path; if (has_draft) { @@ -2435,6 +2444,28 @@ common_speculative_init_result::common_speculative_init_result( } } +void common_speculative_init_result::free_context() { + pimpl->context.reset(); +} + +llama_context * common_speculative_init_result::reinit_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + pimpl->context.reset(); + + const bool spec_mtp = common_speculative_is_mtp(params); + llama_model * model = pimpl->model ? pimpl->model.get() : (spec_mtp ? model_tgt : nullptr); + if (model == nullptr) { + return nullptr; + } + + llama_context * ctx_dft = llama_init_from_model(model, common_speculative_cparams(params, ctx_tgt, spec_mtp)); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create the draft context\n", __func__); + return nullptr; + } + pimpl->context.reset(ctx_dft); + return ctx_dft; +} + common_speculative_init_result::~common_speculative_init_result() = default; llama_model * common_speculative_init_result::model() { diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de59..7f1dede52a02 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -96,6 +96,15 @@ struct common_speculative_init_result { llama_model * model(); llama_context * context(); + // Continuum fork: frees the draft context and builds one for `ctx_tgt` (a rebuilt target + // context) on the same draft model, or on `model_tgt` for MTP. The draft model stays loaded. + // Returns nullptr when the new context cannot be created. Free the draft context (this call + // does it first) before the target it points at. + llama_context * reinit_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + + // Continuum fork: frees the draft context only (it points at the target context, so it goes first) + void free_context(); + private: struct impl; std::unique_ptr pimpl; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index fc118c44814c..362b25e28255 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2069,11 +2069,14 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama uint32_t n_stream_cur; io.read(&n_stream_cur, sizeof(n_stream_cur)); - if (n_stream_cur != n_stream) { + // Continuum fork: one sequence's cells all land in that sequence's stream, so its state + // reads into a cache split into any number of streams (a server whose slot count changed, + // or a unified cache from a split one). The whole cache (seq_id == -1) keeps its layout. + if (n_stream_cur != n_stream && seq_id == -1) { throw std::runtime_error("n_stream mismatch"); } - for (uint32_t s = 0; s < n_stream; ++s) { + for (uint32_t s = 0; s < n_stream_cur; ++s) { uint32_t cell_count; io.read(&cell_count, sizeof(cell_count)); diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index e4a9783970a8..4980750c4fdf 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -20,6 +20,8 @@ add_library(${TARGET} STATIC server-pages.h server-page-files.cpp server-page-files.h + server-reconfigure.cpp + server-reconfigure.h server-checkpoints.cpp server-checkpoints.h server-stream.cpp diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 81bf294e4bc1..ff2faf994e24 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -484,6 +484,8 @@ struct server_metrics { // Continuum fork: RAM pages and the checkpoint byte budget uint64_t n_page_out = 0; + uint64_t n_reconfigure = 0; // Continuum fork: /reconfigure runs (server-reconfigure.h) + uint64_t n_reconfigure_failed = 0; uint64_t n_page_in = 0; uint64_t n_page_failed = 0; // page_out, page_in or spill refused or failed uint64_t n_ckpt_evicted_budget = 0; // checkpoints dropped to keep within --ctx-checkpoints-bytes diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8846f3038323..bfff1e46033a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -5,6 +5,7 @@ #include "server-checkpoints.h" #include "server-page-files.h" #include "server-pages.h" +#include "server-reconfigure.h" #include "server-chat.h" #include "server-common.h" #include "server-http.h" @@ -108,6 +109,17 @@ struct server_batch { } } + // Continuum fork: a reconfigure to more slots than n_batch needs a larger batch + void reinit(int32_t n_tokens_alloc, int32_t n_embd) { + if (batch.pos != nullptr) { + clear(); + llama_batch_free(batch); + batch.pos = nullptr; + } + tokens.clear(); + init(n_tokens_alloc, n_embd); + } + void init(int32_t n_tokens_alloc, int32_t n_embd) { this->n_tokens_alloc = n_tokens_alloc; this->n_embd = n_embd; @@ -437,6 +449,18 @@ struct server_slot { return state != SLOT_STATE_IDLE; } + // Continuum fork (server-reconfigure.h): between two update_slots passes every token the slot + // holds is in its KV and nothing waits on the context's outputs, so its state can leave the + // context and come back. A prompt whose last token is decoded but not yet sampled, or a draft + // decoded but not yet verified, still reads the outputs. + bool at_boundary() const { + return state != SLOT_STATE_DONE_PROMPT && i_batch == -1 && spec_i_batch.empty(); + } + + // Continuum fork: the task registered this slot's sampler with the context (backend sampling); + // a rebuilt context registers it again + bool backend_sampler = false; + bool can_speculate() const { return !!spec; } @@ -862,6 +886,10 @@ struct server_context_impl { return metrics; } + // what the routes report about the loaded model and context (server_context::get_meta); + // read on the task thread, or before the loop starts + server_context_meta get_meta(); + void reset_metrics_bucket() { metrics.reset_bucket(); } @@ -928,6 +956,15 @@ struct server_context_impl { // Continuum fork: named slot states in host RAM (/slots page_out, page_in, spill, page_drop) server_page_store page_store; + // Continuum fork: POST /reconfigure (server-reconfigure.h). The request waits here until every + // slot is at a boundary; admission stays closed meanwhile and the requests that arrive queue. + std::unique_ptr reconfigure_task; + server_reconfigure_request reconfigure_req; + int64_t t_reconfigure_posted_us = 0; + size_t n_reconfigure_deferred = 0; + std::string reconfigure_waiting_on; + int n_reconfigures = 0; + server_metrics metrics; // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync @@ -1271,32 +1308,7 @@ struct server_context_impl { } for (int i = 0; i < params_base.n_parallel; i++) { - server_slot & slot = slots[i]; - - slot.id = i; - slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft; - slot.mem.init(ctx_tgt, ctx_dft); - slot.spec = spec.get(); - slot.n_ctx = n_ctx_slot; - - slot.mctx = mctx; - slot.prompt.tokens.has_mtmd = mctx != nullptr; - - SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); - - slot.callback_on_release = [this](int id_slot) { - queue_tasks.pop_deferred_task(id_slot); - }; - - slot.callback_on_reset = [this](const server_slot & slot) { - // flush the generated token stats before reset() - if (slot.stats.n_gen > 0) { - metrics_on_prediction(slot); - } - }; - - slot.reset(); + init_slot(slots[i], i, n_ctx_slot); } { @@ -1388,6 +1400,38 @@ struct server_context_impl { return true; } + // points a slot at the current contexts and speculative state (load_model, reconfigure) + void bind_slot(server_slot & slot, int n_ctx_slot) { + slot.ctx_tgt = ctx_tgt; + slot.ctx_dft = ctx_dft; + slot.mem.init(ctx_tgt, ctx_dft); + slot.spec = spec.get(); + slot.n_ctx = n_ctx_slot; + } + + void init_slot(server_slot & slot, int id, int n_ctx_slot) { + slot.id = id; + bind_slot(slot, n_ctx_slot); + + slot.mctx = mctx; + slot.prompt.tokens.has_mtmd = mctx != nullptr; + + SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); + + slot.callback_on_release = [this](int id_slot) { + queue_tasks.pop_deferred_task(id_slot); + }; + + slot.callback_on_reset = [this](const server_slot & slot) { + // flush the generated token stats before reset() + if (slot.stats.n_gen > 0) { + metrics_on_prediction(slot); + } + }; + + slot.reset(); + } + // unlike load_model(), this is only called once during initialization bool init() { GGML_ASSERT(ctx_tgt != nullptr); @@ -1400,7 +1444,10 @@ struct server_context_impl { return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { + // a reconfigure runs between two passes, never inside one (server-reconfigure.h) + reconfigure_if_ready(); update_slots(); + reconfigure_if_ready(); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1785,6 +1832,7 @@ struct server_context_impl { } else { llama_set_sampler(ctx_tgt, slot.id, nullptr); } + slot.backend_sampler = use_backend_sampling; SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str()); @@ -2630,6 +2678,170 @@ struct server_context_impl { return res; } + // + // Continuum fork: reconfigure without reloading the model (server-reconfigure.h). The mechanism + // lives there; this is the server's side of server_reconfigure_host. + // + + server_geometry geometry() const { + server_geometry g; + g.n_ctx = llama_n_ctx(ctx_tgt); + g.n_parallel = (int32_t) slots.size(); + g.type_k = params_base.cache_type_k; + g.type_v = params_base.cache_type_v; + return g; + } + + struct reconfigure_host final : server_reconfigure_host { + server_context_impl & s; + const bool with_spec; // the speculative state is rebuilt when the server had one + + explicit reconfigure_host(server_context_impl & s) : s(s), with_spec(s.spec != nullptr) {} + + server_geometry geometry() const override { return s.geometry(); } + int32_t n_ctx_seq() const override { return s.slots.empty() ? 0 : s.slots.front().n_ctx; } + + llama_context * target() const override { return s.ctx_tgt; } + llama_context * draft() const override { return s.ctx_dft; } + common_speculative * speculative() const override { return s.spec.get(); } + + server_prompt & slot_prompt(int id) override { return s.slots[id].prompt; } + std::vector & slot_lora(int id) override { return s.slots[id].lora; } + bool slot_processing(int id) const override { return s.slots[id].is_processing(); } + void slot_fail(int id, const std::string & why) override { + server_slot & slot = s.slots[id]; + s.send_error(slot, why, ERROR_TYPE_SERVER); + slot.prompt_clear(); + slot.release(); + } + + void free_contexts() override { + s.spec.reset(); // holds both contexts + if (s.spec_init) { + s.spec_init->free_context(); // the draft points at the target: it goes first + } + s.ctx_dft = nullptr; + s.llama_init->free_context(); + s.ctx_tgt = nullptr; + } + + bool build_contexts(const server_geometry & geo) override { + common_params p = s.params_base; + p.n_ctx = geo.n_ctx; + p.n_parallel = geo.n_parallel; + p.cache_type_k = geo.type_k; + p.cache_type_v = geo.type_v; + const auto output_limits = server_output_limits(p); + p.n_outputs_max = output_limits.total; + p.n_outputs_max_per_seq = output_limits.per_seq; + + s.ctx_tgt = s.llama_init->reinit_context(p); + if (s.ctx_tgt == nullptr) { + return false; + } + if (s.spec_init) { + common_params p_dft = common_base_params_to_speculative(p); + s.ctx_dft = s.spec_init->reinit_context(p_dft, s.model_tgt, s.ctx_tgt); + if (s.ctx_dft == nullptr) { + s.llama_init->free_context(); + s.ctx_tgt = nullptr; + return false; + } + p.speculative.draft.ctx_tgt = s.ctx_tgt; + p.speculative.draft.ctx_dft = s.ctx_dft; + } + if (with_spec) { + s.spec.reset(common_speculative_init(p.speculative, p.n_parallel)); + } + s.params_base = p; // a wake from sleep reloads this geometry + return true; + } + + void set_slots(int n) override { + const int n_ctx_slot = std::min(llama_n_ctx_seq(s.ctx_tgt), llama_model_n_ctx_train(s.model_tgt)); + while ((int) s.slots.size() > n) { + GGML_ASSERT(!s.slots.back().is_processing()); + s.slots.pop_back(); + } + for (auto & slot : s.slots) { + s.bind_slot(slot, n_ctx_slot); + if (slot.backend_sampler && slot.is_processing()) { + llama_set_sampler(s.ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); + } + } + while ((int) s.slots.size() < n) { + s.slots.emplace_back(); + s.init_slot(s.slots.back(), (int) s.slots.size() - 1, n_ctx_slot); + } + const int32_t n_alloc = std::max(llama_n_batch(s.ctx_tgt), (int32_t) s.slots.size()); + if (n_alloc > s.batch.n_tokens_alloc) { + s.batch.reinit(n_alloc, llama_model_n_embd_inp(s.model_tgt)); + } + s.learn_lora_id = -1; // the optimizer state lived in the freed context + if (s.prompt_cache) { + s.prompt_cache->limit_tokens = llama_n_ctx(s.ctx_tgt); + } + } + + server_page_store & pages() override { return s.page_store; } + + size_t drop_prompt_cache() override { + if (!s.prompt_cache) { + return 0; + } + const size_t n = s.prompt_cache->size(); + s.prompt_cache->states.clear(); + return n; + } + }; + + // runs a pending reconfigure once every slot it needs is at a boundary; called between passes + void reconfigure_if_ready() { + if (!reconfigure_task) { + return; + } + const server_task & task = *reconfigure_task; + + std::vector view; + view.reserve(slots.size()); + for (const auto & slot : slots) { + view.push_back({ slot.id, slot.is_processing(), slot.at_boundary(), (int32_t) slot.prompt.n_tokens() }); + } + + try { + auto plan = server_reconfigure_decide(reconfigure_req, geometry(), view, params_base.kv_unified, + llama_model_n_ctx_train(model_tgt), + string_format("reconfigure%d", n_reconfigures + 1)); + if (plan.wait) { + if (plan.wait_reason != reconfigure_waiting_on) { + SRV_INF("reconfigure waits: %s\n", plan.wait_reason.c_str()); + reconfigure_waiting_on = plan.wait_reason; + } + return; + } + auto res = std::make_unique(); + res->id = task.id; + res->receipt.t_drain_ms = (ggml_time_us() - t_reconfigure_posted_us) / 1000.0; + res->receipt.n_tasks_deferred = n_reconfigure_deferred; + reconfigure_host host(*this); + server_reconfigure_run(host, reconfigure_req, plan, res->receipt, string_format("reconfigure%d", n_reconfigures + 1)); + n_reconfigures++; + metrics.n_reconfigure++; + res->meta = std::make_shared(get_meta()); + queue_results.send(std::move(res)); + } catch (const server_page_error & err) { + metrics.n_reconfigure_failed++; + send_error(task, err.what(), err.type); + } + + // admission opens: the requests that queued meanwhile go back to the queue + reconfigure_task.reset(); + reconfigure_waiting_on.clear(); + for (size_t i = 0; i < slots.size(); ++i) { + queue_tasks.pop_deferred_task((int) i); + } + } + // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { // while yielding, an encode / decode is running and only reading the server state is safe @@ -2654,6 +2866,13 @@ struct server_context_impl { const int id_task = task.id; + if (reconfigure_task) { + // Continuum fork: admission is closed while a reconfigure waits (server-reconfigure.h) + n_reconfigure_deferred++; + queue_tasks.defer(std::move(task)); + break; + } + server_slot * slot = get_available_slot(task); // @@ -3176,6 +3395,25 @@ struct server_context_impl { res->n_erased = n_erased; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_RECONFIGURE: + { + // Continuum fork: runs between two passes once the slots allow it (server-reconfigure.h) + if (reconfigure_task) { + send_error(task, "a reconfigure is already pending", ERROR_TYPE_UNAVAILABLE); + break; + } + try { + reconfigure_req = server_reconfigure_request::parse(task.reconfigure, geometry()); + } catch (const server_page_error & err) { + metrics.n_reconfigure_failed++; + send_error(task, err.what(), err.type); + break; + } + t_reconfigure_posted_us = ggml_time_us(); + n_reconfigure_deferred = 0; + reconfigure_task = std::make_unique(std::move(task)); + reconfigure_if_ready(); + } break; case SERVER_TASK_TYPE_PAGE_OUT: case SERVER_TASK_TYPE_PAGE_IN: case SERVER_TASK_TYPE_PAGE_SPILL: @@ -4904,47 +5142,52 @@ server_response_reader server_context::get_response_reader() { } server_context_meta server_context::get_meta() const { - auto bos_id = llama_vocab_bos(impl->vocab); - auto eos_id = llama_vocab_eos(impl->vocab); - auto bos_token_str = bos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(impl->ctx_tgt, bos_id, true) : ""; - auto eos_token_str = eos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(impl->ctx_tgt, eos_id, true) : ""; + return impl->get_meta(); +} + +server_context_meta server_context_impl::get_meta() { + auto bos_id = llama_vocab_bos(this->vocab); + auto eos_id = llama_vocab_eos(this->vocab); + auto bos_token_str = bos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(this->ctx_tgt, bos_id, true) : ""; + auto eos_token_str = eos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(this->ctx_tgt, eos_id, true) : ""; - const char * ftype_name = llama_ftype_name(llama_model_ftype(impl->model_tgt)); + const char * ftype_name = llama_ftype_name(llama_model_ftype(this->model_tgt)); return server_context_meta { /* build_info */ std::string(llama_build_info()), - /* model_name */ impl->model_name, - /* model_aliases */ impl->model_aliases, - /* model_tags */ impl->model_tags, - /* model_path */ impl->params_base.model.path, - /* has_mtmd */ impl->mctx != nullptr, - /* has_inp_image */ impl->chat_params.allow_image, - /* has_inp_audio */ impl->chat_params.allow_audio, - /* has_inp_video */ impl->chat_params.allow_video, - /* json_ui_settings */ impl->json_ui_settings, - /* slot_n_ctx */ impl->get_slot_n_ctx(), - /* pooling_type */ llama_pooling_type(impl->ctx_tgt), - - /* chat_params */ impl->chat_params, - /* chat_template_caps */ common_chat_templates_get_caps(impl->chat_params.tmpls.get()), + /* model_name */ this->model_name, + /* model_aliases */ this->model_aliases, + /* model_tags */ this->model_tags, + /* model_path */ this->params_base.model.path, + /* has_mtmd */ this->mctx != nullptr, + /* has_inp_image */ this->chat_params.allow_image, + /* has_inp_audio */ this->chat_params.allow_audio, + /* has_inp_video */ this->chat_params.allow_video, + /* json_ui_settings */ this->json_ui_settings, + /* slot_n_ctx */ this->get_slot_n_ctx(), + /* pooling_type */ llama_pooling_type(this->ctx_tgt), + /* n_slots */ (int32_t) this->slots.size(), + + /* chat_params */ this->chat_params, + /* chat_template_caps */ common_chat_templates_get_caps(this->chat_params.tmpls.get()), /* bos_token_str */ bos_token_str, /* eos_token_str */ eos_token_str, - /* fim_pre_token */ llama_vocab_fim_pre(impl->vocab), - /* fim_sub_token */ llama_vocab_fim_suf(impl->vocab), - /* fim_mid_token */ llama_vocab_fim_mid(impl->vocab), - /* fim_pad_token */ llama_vocab_fim_pad(impl->vocab), - /* fim_rep_token */ llama_vocab_fim_rep(impl->vocab), - /* fim_sep_token */ llama_vocab_fim_sep(impl->vocab), - - /* logit_bias_eog */ impl->params_base.sampling.logit_bias_eog, - - /* model_vocab_type */ llama_vocab_type(impl->vocab), - /* model_vocab_n_tokens */ llama_vocab_n_tokens(impl->vocab), - /* model_n_ctx_train */ llama_model_n_ctx_train(impl->model_tgt), - /* model_n_embd_inp */ llama_model_n_embd(impl->model_tgt), - /* model_n_params */ llama_model_n_params(impl->model_tgt), - /* model_size */ llama_model_size(impl->model_tgt), + /* fim_pre_token */ llama_vocab_fim_pre(this->vocab), + /* fim_sub_token */ llama_vocab_fim_suf(this->vocab), + /* fim_mid_token */ llama_vocab_fim_mid(this->vocab), + /* fim_pad_token */ llama_vocab_fim_pad(this->vocab), + /* fim_rep_token */ llama_vocab_fim_rep(this->vocab), + /* fim_sep_token */ llama_vocab_fim_sep(this->vocab), + + /* logit_bias_eog */ this->params_base.sampling.logit_bias_eog, + + /* model_vocab_type */ llama_vocab_type(this->vocab), + /* model_vocab_n_tokens */ llama_vocab_n_tokens(this->vocab), + /* model_n_ctx_train */ llama_model_n_ctx_train(this->model_tgt), + /* model_n_embd_inp */ llama_model_n_embd(this->model_tgt), + /* model_n_params */ llama_model_n_params(this->model_tgt), + /* model_size */ llama_model_size(this->model_tgt), /* model_ftype */ ftype_name, /* model_weight_buffers */ [&] { // MEASURED placement: where the weights were actually allocated. @@ -4953,19 +5196,19 @@ server_context_meta server_context::get_meta() const { // so a lane that asked for the GPU and fell back to the CPU prints // numbers identical to one that succeeded. json bufs = json::array(); - const size_t n = llama_model_n_weight_backends(impl->model_tgt); + const size_t n = llama_model_n_weight_backends(this->model_tgt); for (size_t i = 0; i < n; i++) { - const char * backend = llama_model_weight_backend_name(impl->model_tgt, i); + const char * backend = llama_model_weight_backend_name(this->model_tgt, i); bufs.push_back(json { { "backend", backend ? backend : "" }, - { "size_bytes", llama_model_weight_backend_bytes(impl->model_tgt, i) }, + { "size_bytes", llama_model_weight_backend_bytes(this->model_tgt, i) }, }); } return bufs; }(), /* memory_breakdown */ [&] { json out = json::array(); - for (const auto & [buft, mb] : llama_get_memory_breakdown(impl->ctx_tgt)) { + for (const auto & [buft, mb] : llama_get_memory_breakdown(this->ctx_tgt)) { out.push_back(json { { "buffer_type", ggml_backend_buft_name(buft) }, { "model_bytes", mb.model }, @@ -4976,9 +5219,9 @@ server_context_meta server_context::get_meta() const { return out; }(), /* context_geometry */ json { - { "allocated_tokens", llama_n_ctx(impl->ctx_tgt) }, - { "max_tokens_per_sequence", llama_n_ctx_seq(impl->ctx_tgt) }, - { "sequence_capacity", llama_n_seq_max(impl->ctx_tgt) }, + { "allocated_tokens", llama_n_ctx(this->ctx_tgt) }, + { "max_tokens_per_sequence", llama_n_ctx_seq(this->ctx_tgt) }, + { "sequence_capacity", llama_n_seq_max(this->ctx_tgt) }, }, }; } @@ -5082,7 +5325,7 @@ std::unique_ptr server_routes::handle_completions_impl( task.params = server_schema::eval_llama_cmpl_schema( ctx_server.vocab, params, - meta->logit_bias_eog, + meta_now()->logit_bias_eog, data); task.params.message_spans = task.tokens.find_message_spans(delimiters); @@ -5093,7 +5336,7 @@ std::unique_ptr server_routes::handle_completions_impl( // OAI-compat task.params.res_type = res_type; task.params.oaicompat_cmpl_id = completion_id; - task.params.oaicompat_model = meta->model_name; + task.params.oaicompat_model = meta_now()->model_name; // prepare child tasks if (task.params.n_cmpl > 1) { @@ -5378,7 +5621,7 @@ static json get_res_props(const server_context_meta & meta, const common_params json props = { { "default_generation_settings", default_generation_settings_for_props }, - { "total_slots", params.n_parallel }, + { "total_slots", meta.n_slots }, { "slot_save_size_version", 1 }, { "slot_page_version", 1 }, // Continuum fork: /slots page_out, page_in, spill, page_drop { "expert_container_custody_version", 1 }, @@ -5417,7 +5660,7 @@ static json get_res_props(const server_context_meta & meta, const common_params } json server_routes::get_model_info() const { - return get_res_model_info(*meta); + return get_res_model_info(*meta_now()); } void server_routes::init_routes() { @@ -5599,7 +5842,7 @@ void server_routes::init_routes() { std::unique_lock lock(mutex_cache); res->ok(cached_props); } else { - res->ok(get_res_props(*meta, params, false)); + res->ok(get_res_props(*meta_now(), params, false)); } return res; }; @@ -5680,7 +5923,7 @@ void server_routes::init_routes() { data.at("input_extra"), params.n_batch, params.n_predict, - meta->slot_n_ctx, + meta_now()->slot_n_ctx, params.spm_infill, tokenized_prompts[0].get_tokens() // TODO: this could maybe be multimodal. ); @@ -5724,7 +5967,7 @@ void server_routes::init_routes() { json body = json::parse(req.body); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5783,7 +6026,7 @@ void server_routes::init_routes() { SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5800,7 +6043,7 @@ void server_routes::init_routes() { this->post_transcriptions_oai = [this](const server_http_req & req) { auto res = create_response(); - if (!meta->has_mtmd || !meta->chat_params.allow_audio) { + if (!meta_now()->has_mtmd || !meta_now()->chat_params.allow_audio) { res->error(format_error_response("The current model does not support audio input.", ERROR_TYPE_NOT_SUPPORTED)); return res; } @@ -5808,14 +6051,14 @@ void server_routes::init_routes() { std::vector files; json body = convert_transcriptions_to_chatcmpl( json::parse(req.body), - meta->chat_params.tmpls.get(), + meta_now()->chat_params.tmpls.get(), req.files, files); SRV_DBG("%s\n", "Request converted: OpenAI Transcriptions -> OpenAI Chat Completions"); SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5833,7 +6076,7 @@ void server_routes::init_routes() { SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5854,7 +6097,7 @@ void server_routes::init_routes() { json body = json::parse(req.body); json data = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); res->ok({{ "prompt", std::move(data.at("prompt")) }}); return res; @@ -5867,7 +6110,7 @@ void server_routes::init_routes() { std::unique_lock lock(mutex_cache); res->ok(cached_models); } else { - res->ok(get_res_models(*meta)); + res->ok(get_res_models(*meta_now())); } return res; }; @@ -6005,7 +6248,7 @@ void server_routes::init_routes() { // write JSON response json root = format_response_rerank( body, - meta->model_name, + meta_now()->model_name, responses, is_tei_format, documents, @@ -6138,6 +6381,42 @@ void server_routes::init_routes() { // Duty rest returns a deferred outcome: the caller retains replayable input context, reply or // labels, and adapter parameters, then retries after retry_after_ms. A slot's prompt may have // changed by then. The slot-release queue cannot wake a cooldown on an already idle slot. + // Continuum fork: POST /reconfigure {"n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force"} + // changes the context geometry without reloading the model; the answer is the receipt + // (server-reconfigure.h). It waits while slots drain, so a caller gives it a long timeout. + this->post_reconfigure = [this](const server_http_req & req) { + auto res = create_response(); + json body = json::object(); + if (!req.body.empty()) { + try { + body = json::parse(req.body); + } catch (const std::exception &) { + res->error(format_error_response("The body must be JSON", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + } + { + server_task task(SERVER_TASK_TYPE_RECONFIGURE); + task.id = res->rd.get_new_id(); + task.reconfigure = std::move(body); + res->rd.post_task(std::move(task)); + } + auto result = res->rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + auto * done = dynamic_cast(result.get()); + GGML_ASSERT(done != nullptr); + set_meta(done->meta); // /props and the slot bounds now report the new geometry + res->ok(done->to_json()); + return res; + }; + this->post_learn = [this](const server_http_req & req) { auto res = create_response(); if (params.train_dir.empty()) { @@ -6170,7 +6449,7 @@ void server_routes::init_routes() { served["messages"].push_back(messages[i]); } std::vector files; - const json parsed = oaicompat_chat_params_parse(served, meta->chat_params, files); + const json parsed = oaicompat_chat_params_parse(served, meta_now()->chat_params, files); if (!files.empty()) { res->error(format_error_response("a conversation with media cannot be learned yet: only text prompts", ERROR_TYPE_INVALID_REQUEST)); return res; @@ -6185,7 +6464,7 @@ void server_routes::init_routes() { framed["tools"] = tools; framed["add_generation_prompt"] = generation; std::vector media; - const json rendered = oaicompat_chat_params_parse(framed, meta->chat_params, media, + const json rendered = oaicompat_chat_params_parse(framed, meta_now()->chat_params, media, generation ? server_chat_render::serving : server_chat_render::closed_transcript); if (!media.empty()) { throw std::invalid_argument("learning media is not supported"); @@ -6400,7 +6679,7 @@ std::unique_ptr server_routes::handle_slots_restore(const valid = (i == 8 || i == 13 || i == 18 || i == 23) ? id[i] == '-' : ((id[i] >= '0' && id[i] <= '9') || (id[i] >= 'a' && id[i] <= 'f')); } - if (!valid || id_slot < 0 || id_slot >= params.n_parallel) { + if (!valid || id_slot < 0 || id_slot >= meta_now()->n_slots) { res->error(format_error_response("Invalid restore operation identity or slot", ERROR_TYPE_INVALID_REQUEST)); return res; } @@ -6492,7 +6771,7 @@ std::unique_ptr server_routes::handle_slots_page(const ser return res; } const bool on_slot = action == "page_out" || action == "page_in"; - if (on_slot && (id_slot < 0 || id_slot >= params.n_parallel)) { + if (on_slot && (id_slot < 0 || id_slot >= meta_now()->n_slots)) { res->error(format_error_response("Invalid slot ID", ERROR_TYPE_INVALID_REQUEST)); return res; } @@ -6613,7 +6892,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons return res; } - if (res_type != TASK_RESPONSE_TYPE_NONE && meta->pooling_type == LLAMA_POOLING_TYPE_NONE) { + if (res_type != TASK_RESPONSE_TYPE_NONE && meta_now()->pooling_type == LLAMA_POOLING_TYPE_NONE) { res->error(format_error_response("Pooling type 'none' is not OAI compatible. Please use a different pooling type", ERROR_TYPE_INVALID_REQUEST)); return res; } @@ -6655,8 +6934,8 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons int embd_normalize = params.embd_normalize; if (body.count("embd_normalize") != 0) { embd_normalize = body.at("embd_normalize").get(); - if (meta->pooling_type == LLAMA_POOLING_TYPE_NONE) { - SRV_DBG("embd_normalize is not supported by pooling type %d, ignoring it\n", meta->pooling_type); + if (meta_now()->pooling_type == LLAMA_POOLING_TYPE_NONE) { + SRV_DBG("embd_normalize is not supported by pooling type %d, ignoring it\n", meta_now()->pooling_type); } } @@ -6698,7 +6977,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons // write JSON response json root = res_type == TASK_RESPONSE_TYPE_OAI_EMBD - ? format_embeddings_response_oaicompat(body, meta->model_name, responses, use_base64) + ? format_embeddings_response_oaicompat(body, meta_now()->model_name, responses, use_base64) : json(responses); res->ok(root); return res; @@ -6731,7 +7010,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); json prompt = body_parsed.at("prompt"); // SRV_DBG("prompt = %s\n", prompt.dump().c_str()); @@ -6760,8 +7039,8 @@ void server_routes::update_cached_responses(bool is_sleeping) { std::unique_lock lock(mutex_cache); if (is_sleeping) { - cached_models = get_res_models(*meta); - cached_props = get_res_props(*meta, params, true); + cached_models = get_res_models(*meta_now()); + cached_props = get_res_props(*meta_now(), params, true); cached_metrics = ctx_server.get_metrics(); should_reset_buckets = false; diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 38efbb264491..0b4489155613 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -26,6 +26,7 @@ struct server_context_meta { json json_ui_settings; int slot_n_ctx; enum llama_pooling_type pooling_type; + int32_t n_slots; // Continuum fork: changes with /reconfigure, so routes read it here, not from params // chat params server_chat_params & chat_params; @@ -136,7 +137,7 @@ struct server_routes { // note: this is not thread-safe and can only when ctx_http.is_ready is false void update_meta(const server_context & ctx_server) { - this->meta = std::make_unique(ctx_server.get_meta()); + set_meta(std::make_shared(ctx_server.get_meta())); } // handlers using lambda function, so that they can capture `this` without `std::bind` @@ -176,6 +177,7 @@ struct server_routes { server_http_context::handler_t post_train_measure; server_http_context::handler_t post_train_pause; server_http_context::handler_t post_train_resume; + server_http_context::handler_t post_reconfigure; // Continuum fork (server-reconfigure.h) // to be used in router mode json get_model_info() const; @@ -194,8 +196,11 @@ struct server_routes { std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); - // using unique_ptr to allow late initialization of const - std::unique_ptr meta; + // late-initialized, and replaced by /reconfigure while other HTTP threads read it: every access + // goes through meta_now() / set_meta() (atomic shared_ptr operations) + std::shared_ptr meta; + std::shared_ptr meta_now() const { return std::atomic_load(&meta); } + void set_meta(std::shared_ptr m) { std::atomic_store(&meta, std::move(m)); } const common_params & params; server_context_impl & ctx_server; diff --git a/tools/server/server-reconfigure.cpp b/tools/server/server-reconfigure.cpp new file mode 100644 index 000000000000..686c5bdddcb3 --- /dev/null +++ b/tools/server/server-reconfigure.cpp @@ -0,0 +1,379 @@ +#include "server-reconfigure.h" + +#include "server-checkpoints.h" + +#include "arg.h" + +#include +#include + +json server_geometry::to_json() const { + return json { + { "n_ctx", n_ctx }, + { "n_parallel", n_parallel }, + { "cache_type_k", ggml_type_name(type_k) }, + { "cache_type_v", ggml_type_name(type_v) }, + }; +} + +int32_t server_geometry_n_ctx_seq(const server_geometry & geo, bool kv_unified, int32_t n_ctx_train) { + const int32_t n_ctx = GGML_PAD(geo.n_ctx, 256); + int32_t n_ctx_seq = kv_unified ? n_ctx : GGML_PAD(n_ctx / std::max(1, geo.n_parallel), 256); + if (n_ctx_train > 0) { + n_ctx_seq = std::min(n_ctx_seq, n_ctx_train); + } + return n_ctx_seq; +} + +static int32_t parse_positive(const json & body, const char * key, int32_t cur, int32_t max) { + if (!body.contains(key)) { + return cur; + } + const json & v = body.at(key); + if (!v.is_number_integer() || v.get() < 1 || v.get() > max) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format("%s must be an integer in [1, %d]", key, max)); + } + return v.get(); +} + +static ggml_type parse_kv_type(const json & body, const char * key, ggml_type cur) { + if (!body.contains(key)) { + return cur; + } + if (!body.at(key).is_string()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format("%s must be a KV cache type name", key)); + } + try { + return common_kv_cache_type_from_str(body.at(key).get()); + } catch (const std::runtime_error & e) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, e.what()); + } +} + +server_reconfigure_request server_reconfigure_request::parse(const json & body, const server_geometry & cur) { + if (!body.is_object()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the body must be a JSON object"); + } + static const char * known[] = { "n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force" }; + for (const auto & item : body.items()) { + const std::string & key = item.key(); + if (std::find_if(std::begin(known), std::end(known), [&](const char * k) { return key == k; }) == std::end(known)) { + // a misspelled field silently keeping the old geometry is the bug this refuses + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "unknown field: " + key); + } + } + server_reconfigure_request req; + req.geo.n_ctx = parse_positive(body, "n_ctx", cur.n_ctx, INT32_MAX / 2); + req.geo.n_parallel = parse_positive(body, "n_parallel", cur.n_parallel, (int32_t) llama_max_parallel_sequences()); + req.geo.type_k = parse_kv_type (body, "cache_type_k", cur.type_k); + req.geo.type_v = parse_kv_type (body, "cache_type_v", cur.type_v); + if (body.contains("force")) { + if (!body.at("force").is_boolean()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "force must be a boolean"); + } + req.force = body.at("force").get(); + } + return req; +} + +server_reconfigure_plan server_reconfigure_decide( + const server_reconfigure_request & req, + const server_geometry & cur, + const std::vector & slots, + bool kv_unified, + int32_t n_ctx_train, + const std::string & name_prefix) { + server_reconfigure_plan plan; + plan.n_ctx_seq = server_geometry_n_ctx_seq(req.geo, kv_unified, n_ctx_train); + plan.cold_start = !req.geo.same_kv_format(cur); + + if (plan.cold_start && !req.force) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format( + "a KV type change (%s/%s -> %s/%s) cannot carry the residents' states over: pass \"force\": true to cold start", + ggml_type_name(cur.type_k), ggml_type_name(cur.type_v), ggml_type_name(req.geo.type_k), ggml_type_name(req.geo.type_v))); + } + + for (const auto & s : slots) { + if (!s.processing) { + continue; + } + if (s.n_tokens + 1 > plan.n_ctx_seq) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format( + "slot %d is serving a request holding %d tokens and the new geometry gives a sequence %d: " + "shrinking under a live request would truncate it", + s.id, s.n_tokens, plan.n_ctx_seq)); + } + if (!s.at_boundary) { + plan.wait = true; + plan.wait_reason = string_format("slot %d is between a decode and its sampling", s.id); + } else if (s.id >= req.geo.n_parallel) { + plan.wait = true; + plan.wait_reason = string_format("slot %d serves a request and has no place among %d slots", s.id, req.geo.n_parallel); + } else if (plan.cold_start) { + plan.wait = true; + plan.wait_reason = string_format("slot %d serves a request and a cold start has no state to resume it from", s.id); + } + } + if (plan.wait) { + return plan; + } + + // slots of the new geometry that nobody keeps: an occupied idle slot past n_parallel moves into one + std::vector taken(req.geo.n_parallel, false); + for (const auto & s : slots) { + if (s.n_tokens > 0 && s.id < req.geo.n_parallel && s.n_tokens <= plan.n_ctx_seq) { + taken[s.id] = true; + } + } + + for (const auto & s : slots) { + if (s.n_tokens == 0) { + continue; + } + server_reconfigure_move m; + m.id_from = s.id; + m.processing = s.processing; + m.n_tokens = s.n_tokens; + if (plan.cold_start) { + m.fate = SERVER_RECONFIGURE_DROPPED; + m.reason = "KV type change (cold start)"; + } else if (s.n_tokens > plan.n_ctx_seq) { + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.reason = string_format("%d tokens do not fit a sequence of %d", s.n_tokens, plan.n_ctx_seq); + } else if (s.id < req.geo.n_parallel) { + m.id_to = s.id; + } else { + const auto free = std::find(taken.begin(), taken.end(), false); + if (free == taken.end()) { + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.reason = string_format("no free slot among %d", req.geo.n_parallel); + } else { + *free = true; + m.id_to = (int) (free - taken.begin()); + } + } + if (m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE) { + m.page = string_format("%s.slot%d", name_prefix.c_str(), s.id); + } + plan.moves.push_back(std::move(m)); + } + return plan; +} + +static const char * fate_name(server_reconfigure_fate f) { + switch (f) { + case SERVER_RECONFIGURE_RESTORED: return "restored"; + case SERVER_RECONFIGURE_KEPT_IN_STORE: return "kept_in_store"; + case SERVER_RECONFIGURE_DROPPED: return "dropped"; + } + return "unknown"; +} + +json server_reconfigure_move::to_json() const { + json j = { + { "id_slot_before", id_from }, + { "id_slot_after", id_to }, + { "in_flight", processing }, + { "n_tokens", n_tokens }, + { "n_bytes", n_bytes }, + { "fate", fate_name(fate) }, + }; + if (!page.empty()) { + j["page"] = page; + } + if (!reason.empty()) { + j["reason"] = reason; + } + return j; +} + +json server_reconfigure_receipt::to_json() const { + json moved = json::array(); + int n_restored = 0, n_kept = 0, n_dropped = 0; + for (const auto & m : moves) { + moved.push_back(m.to_json()); + n_restored += m.fate == SERVER_RECONFIGURE_RESTORED; + n_kept += m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE; + n_dropped += m.fate == SERVER_RECONFIGURE_DROPPED; + } + json j = { + { "before", before.to_json() }, + { "after", after.to_json() }, + { "n_ctx_seq_before", n_ctx_seq_before }, + { "n_ctx_seq_after", n_ctx_seq_after }, + { "cold_start", cold_start }, + { "rolled_back", rolled_back }, + { "t_drain_ms", t_drain_ms }, + { "t_page_out_ms", t_page_out_ms }, + { "t_free_ms", t_free_ms }, + { "t_init_ms", t_init_ms }, + { "t_page_in_ms", t_page_in_ms }, + { "t_total_ms", t_total_ms }, + { "n_bytes", n_bytes }, + { "n_bytes_dropped", n_bytes_dropped }, + { "n_tasks_deferred", n_tasks_deferred }, + { "n_restored", n_restored }, + { "n_kept_in_store", n_kept }, + { "n_dropped", n_dropped }, + { "pages", moved }, + }; + if (rolled_back) { + j["rollback_reason"] = rollback_reason; + } + return j; +} + +void server_reconfigure_run( + server_reconfigure_host & host, + const server_reconfigure_request & req, + server_reconfigure_plan & plan, + server_reconfigure_receipt & rc, + const std::string & page_prefix) { + rc.before = host.geometry(); + rc.n_ctx_seq_before = host.n_ctx_seq(); + server_page_store & store = host.pages(); + + // the store must take every page that cannot be restored before anything moves, counted as it charges them + std::vector shapes; + shapes.reserve(plan.moves.size()); + { + size_t n_kept = 0; + for (const auto & m : plan.moves) { + shapes.push_back(server_page_measure(host.target(), host.draft(), host.speculative(), m.id_from)); + if (m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE) { + n_kept += shapes.back().n_bytes(server_checkpoints_bytes(host.slot_prompt(m.id_from).checkpoints)); + } + } + if (store.limit_bytes > 0 && store.n_bytes() + n_kept > store.limit_bytes) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, string_format( + "%zu bytes of slot state cannot be restored under the new geometry and the page store (%zu of %llu bytes) cannot hold them", + n_kept, store.n_bytes(), (unsigned long long) store.limit_bytes)); + } + } + + // 1. page out: every occupied slot's state into host RAM. A cold start's states are captured + // too and dropped only once the new geometry stands, so a failed build rolls back whole. + const int64_t t0 = ggml_time_us(); + std::vector pages(plan.moves.size()); + for (size_t i = 0; i < plan.moves.size(); ++i) { + auto & m = plan.moves[i]; + pages[i] = server_page_capture(host.target(), host.draft(), m.id_from, std::move(shapes[i])); + server_prompt & prompt = host.slot_prompt(m.id_from); + const bool in_place = m.fate == SERVER_RECONFIGURE_RESTORED && m.id_to == m.id_from; + if (!in_place) { + // it leaves its slot: the page carries the tokens, checkpoints and adapters + pages[i]->prompt = std::move(prompt); + pages[i]->lora = host.slot_lora(m.id_from); + prompt.clear(); + } + m.n_bytes = pages[i]->n_bytes() + (in_place ? server_checkpoints_bytes(prompt.checkpoints) : 0); + rc.n_bytes += m.fate == SERVER_RECONFIGURE_DROPPED ? 0 : m.n_bytes; + } + const int64_t t1 = ggml_time_us(); + + // 2. free both contexts; the model, its mmap and the adapters stay + host.free_contexts(); + const int64_t t2 = ggml_time_us(); + + // 3. build the new geometry, or the old one again + if (!host.build_contexts(req.geo)) { + rc.rolled_back = true; + rc.rollback_reason = "the new geometry could not be built (see the log); the old one was rebuilt"; + SRV_ERR("reconfigure to n_ctx = %d, n_parallel = %d failed; rebuilding the old geometry\n", req.geo.n_ctx, req.geo.n_parallel); + if (!host.build_contexts(rc.before)) { + GGML_ABORT("reconfigure: neither the new nor the old context geometry could be built"); + } + } + const int64_t t3 = ggml_time_us(); + + // 4. the slots of the geometry that stands + host.set_slots(rc.rolled_back ? rc.before.n_parallel : req.geo.n_parallel); + + // 5. page in: every state back, or into the store with a typed reason; never truncated + for (size_t i = 0; i < plan.moves.size(); ++i) { + auto & m = plan.moves[i]; + if (rc.rolled_back) { + m.fate = SERVER_RECONFIGURE_RESTORED; // the old geometry holds every state it held + m.id_to = m.id_from; + m.reason.clear(); + m.page.clear(); + } + if (m.fate == SERVER_RECONFIGURE_DROPPED) { + rc.n_bytes_dropped += m.n_bytes; + pages[i].reset(); + continue; + } + // a page that left its slot carries the slot's tokens; one restored in place does not + const bool carries_prompt = pages[i]->prompt.tokens.size() > 0; + if (m.fate == SERVER_RECONFIGURE_RESTORED) { + try { + server_page_restore(*pages[i], host.target(), host.draft(), host.speculative(), m.id_to); + if (carries_prompt) { + host.slot_prompt(m.id_to) = std::move(pages[i]->prompt); + host.slot_lora(m.id_to) = pages[i]->lora; + } + continue; + } catch (const server_page_error & err) { + m.reason = err.what(); + if (host.slot_processing(m.id_to)) { + // its KV is gone and its request cannot continue: it ends with a typed error + // instead of going on with a context it no longer has + host.slot_fail(m.id_to, std::string("the reconfigure could not restore this request's state: ") + err.what()); + m.fate = SERVER_RECONFIGURE_DROPPED; + continue; + } + if (!carries_prompt) { + pages[i]->prompt = std::move(host.slot_prompt(m.id_to)); + pages[i]->lora = host.slot_lora(m.id_to); + } + host.slot_prompt(m.id_to).clear(); + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.page = string_format("%s.slot%d", page_prefix.c_str(), m.id_from); + } + } + m.id_to = -1; + try { + store.put(m.page, pages[i]); + SRV_WRN("reconfigure: slot %d's %d tokens stay in the page store as '%s': %s\n", + m.id_from, m.n_tokens, m.page.c_str(), m.reason.c_str()); + } catch (const server_page_error & err) { + // admitted above unless its restore failed here: a state the store cannot take is dropped, typed + m.reason = (m.reason.empty() ? "" : m.reason + "; ") + "the page store refused it: " + err.what(); + SRV_ERR("reconfigure: slot %d's %d tokens are dropped: %s\n", m.id_from, m.n_tokens, m.reason.c_str()); + m.fate = SERVER_RECONFIGURE_DROPPED; + m.page.clear(); + rc.n_bytes_dropped += m.n_bytes; + } + } + + rc.cold_start = plan.cold_start && !rc.rolled_back; + if (rc.cold_start) { + // every state in the old KV format: the store's pages and the prompt cache go with them + const size_t n_resident = rc.n_bytes_dropped; + const size_t n_store = store.n_bytes(); + const size_t n_pages = store.n_pages(); + const size_t n_cache = host.drop_prompt_cache(); + store.clear(); + rc.n_bytes_dropped += n_store + n_cache; + SRV_WRN("reconfigure: KV type change %s/%s -> %s/%s is a cold start: dropped %zu resident bytes, %zu page-store bytes (%zu pages), %zu prompt-cache bytes\n", + ggml_type_name(rc.before.type_k), ggml_type_name(rc.before.type_v), + ggml_type_name(req.geo.type_k), ggml_type_name(req.geo.type_v), + n_resident, n_store, n_pages, n_cache); + } + const int64_t t4 = ggml_time_us(); + + rc.after = host.geometry(); + rc.n_ctx_seq_after = host.n_ctx_seq(); + rc.moves = plan.moves; + rc.t_page_out_ms = (t1 - t0) / 1000.0; + rc.t_free_ms = (t2 - t1) / 1000.0; + rc.t_init_ms = (t3 - t2) / 1000.0; + rc.t_page_in_ms = (t4 - t3) / 1000.0; + rc.t_total_ms = (t4 - t0) / 1000.0; + + SRV_INF("reconfigure: n_ctx %d -> %d, n_parallel %d -> %d, %s/%s; pause %.1f ms (page out %.1f, free %.1f, init %.1f, page in %.1f) for %zu bytes; drain %.1f ms%s\n", + rc.before.n_ctx, rc.after.n_ctx, rc.before.n_parallel, rc.after.n_parallel, + ggml_type_name(rc.after.type_k), ggml_type_name(rc.after.type_v), + rc.t_total_ms, rc.t_page_out_ms, rc.t_free_ms, rc.t_init_ms, rc.t_page_in_ms, rc.n_bytes, rc.t_drain_ms, + rc.rolled_back ? " (ROLLED BACK)" : ""); +} diff --git a/tools/server/server-reconfigure.h b/tools/server/server-reconfigure.h new file mode 100644 index 000000000000..45e3e225b7f1 --- /dev/null +++ b/tools/server/server-reconfigure.h @@ -0,0 +1,182 @@ +#pragma once + +// Continuum fork: a running server changes its context geometry without reloading the model. +// +// POST /reconfigure {"n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force"} is a pause, not +// a relaunch: admission stops, the slots reach a ubatch boundary, every occupied slot's sequence state +// (target KV, draft KV, speculative tail, checkpoints) moves to host RAM (server-pages.h), the target +// and draft contexts are freed and rebuilt on the SAME llama_model (weights and mmap stay), the slots +// are rebuilt for the new n_parallel, the states go back, and serving resumes. A request in flight +// keeps its slot and continues where it paused; nothing is truncated and nothing is lost. +// Design: continuum docs/architecture/SERVING-ENGINE.md §3 ("a relaunch is a pause"), slice S2. +// +// What the core sees: the receipt (server_reconfigure_receipt::to_json) and /props context_geometry. + +#include "server-common.h" +#include "server-pages.h" +#include "server-task.h" + +#include "common.h" +#include "llama.h" + +#include +#include +#include +#include + +struct server_context_meta; + +// The context geometry a reconfigure may change. Everything else (the model, adapters, batch +// sizes, speculative types, checkpoint budget) stays as launched. +struct server_geometry { + int32_t n_ctx = 0; // KV cells of the whole context (llama_n_ctx) + int32_t n_parallel = 0; // slots, one sequence each (llama_n_seq_max) + ggml_type type_k = GGML_TYPE_F16; + ggml_type type_v = GGML_TYPE_F16; + + // a sequence state is written in its KV types: across a type change no page can be carried + bool same_kv_format(const server_geometry & other) const { + return type_k == other.type_k && type_v == other.type_v; + } + + json to_json() const; +}; + +// The tokens one sequence may hold under `geo`: libllama's rule (llama-context.cpp: n_ctx padded to +// 256; a split cache gives each sequence n_ctx / n_seq_max, padded) and the server's cap at the +// model's training context. Used before the new context exists, so a refusal costs nothing. +int32_t server_geometry_n_ctx_seq(const server_geometry & geo, bool kv_unified, int32_t n_ctx_train); + +struct server_reconfigure_request { + server_geometry geo; + // a KV type change is an intentional cold start: refused unless `force`, then every resident + // state, page and prompt-cache entry is dropped (logged loudly and counted in the receipt) + bool force = false; + + // Absent fields keep `cur`. Throws server_page_error (ERROR_TYPE_INVALID_REQUEST). + static server_reconfigure_request parse(const json & body, const server_geometry & cur); +}; + +// One slot as the reconfigure sees it on the task thread. +struct server_reconfigure_slot { + int id = -1; + bool processing = false; // a request is in flight + bool at_boundary = true; // no sampling or verification is half done (see server-context.cpp) + int32_t n_tokens = 0; // tokens its KV holds +}; + +// Where an occupied slot's state goes. +enum server_reconfigure_fate { + SERVER_RECONFIGURE_RESTORED, // into a slot of the new geometry + SERVER_RECONFIGURE_KEPT_IN_STORE, // a named RAM page (/slots/{id}?action=page_in brings it back) + SERVER_RECONFIGURE_DROPPED, // cold start (KV type change with force) +}; + +struct server_reconfigure_move { + int id_from = -1; + int id_to = -1; // -1 unless restored + bool processing = false; + int32_t n_tokens = 0; + size_t n_bytes = 0; + + server_reconfigure_fate fate = SERVER_RECONFIGURE_RESTORED; + std::string page; // the store name when kept + std::string reason; // why it was not restored + + json to_json() const; +}; + +// The decision before anything is touched. +struct server_reconfigure_plan { + bool wait = false; // a slot must first finish (it sits past the new n_parallel) or reach a boundary + std::string wait_reason; + + bool cold_start = false; // KV type change with force + int32_t n_ctx_seq = 0; // predicted tokens per sequence after + std::vector moves; // one per occupied slot, fate decided +}; + +// Decides the reconfigure. Throws server_page_error for a refusal: a KV type change without force, +// an in-flight request whose tokens do not fit the new window (shrinking under a live act would +// truncate it), or a geometry libllama would reject. `name_prefix` names pages that stay in the store. +server_reconfigure_plan server_reconfigure_decide( + const server_reconfigure_request & req, + const server_geometry & cur, + const std::vector & slots, + bool kv_unified, + int32_t n_ctx_train, + const std::string & name_prefix); + +struct server_reconfigure_receipt { + server_geometry before; + server_geometry after; + int32_t n_ctx_seq_before = 0; + int32_t n_ctx_seq_after = 0; + + bool cold_start = false; + bool rolled_back = false; // the new geometry could not be built; the old one was rebuilt and every state restored + std::string rollback_reason; + + double t_drain_ms = 0.0; // request posted -> every slot at a boundary (and the ones past n_parallel done) + double t_page_out_ms = 0.0; + double t_free_ms = 0.0; + double t_init_ms = 0.0; + double t_page_in_ms = 0.0; + double t_total_ms = 0.0; // the pause serving saw: page out through page in + + size_t n_bytes = 0; // paged out and back + size_t n_bytes_dropped = 0; // cold start: resident states, pages and prompt-cache entries dropped + size_t n_tasks_deferred = 0; // requests that arrived during the pause, queued and resumed after + + std::vector moves; + + json to_json() const; +}; + +// What a reconfigure needs from the server and nothing more; server-context.cpp implements it over +// its private slots and contexts. Slot ids are dense: 0 .. geometry().n_parallel - 1. +struct server_reconfigure_host { + virtual ~server_reconfigure_host() = default; + + virtual server_geometry geometry() const = 0; + virtual int32_t n_ctx_seq() const = 0; // tokens one slot holds now + + virtual llama_context * target() const = 0; + virtual llama_context * draft() const = 0; + virtual common_speculative * speculative() const = 0; + + virtual server_prompt & slot_prompt(int id_slot) = 0; + virtual std::vector & slot_lora(int id_slot) = 0; + virtual bool slot_processing(int id_slot) const = 0; + // ends the slot's request with a typed error and empties the slot + virtual void slot_fail(int id_slot, const std::string & why) = 0; + + // frees the draft and target contexts and the speculative state; the model(s) stay loaded + virtual void free_contexts() = 0; + // builds both contexts for `geo` on the loaded model(s); false leaves no context + virtual bool build_contexts(const server_geometry & geo) = 0; + // `n` slots bound to the current contexts; slots past `n` are idle and go + virtual void set_slots(int n) = 0; + + virtual server_page_store & pages() = 0; + virtual size_t drop_prompt_cache() = 0; // the bytes it held +}; + +// Runs a decided plan: page out, free, build (or roll back to the old geometry), page in. Throws +// server_page_error only before anything is touched. `page_prefix` names pages that stay in the store. +void server_reconfigure_run( + server_reconfigure_host & host, + const server_reconfigure_request & req, + server_reconfigure_plan & plan, + server_reconfigure_receipt & rc, + const std::string & page_prefix); + +struct server_task_result_reconfigure : server_task_result { + server_reconfigure_receipt receipt; + // the routes' view of the server after the reconfigure, built on the task thread + std::shared_ptr meta; + + virtual json to_json() override { + return receipt.to_json(); + } +}; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 12c1e5a008fe..b5a3e319bd82 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1603,6 +1603,14 @@ std::string server_task_result_metrics::to_metrics() { "page_failed_total", "page_out, page_in or spill operations refused or failed", (double) metrics.n_page_failed + }, { + "reconfigure_total", + "Context geometry changes made without reloading the model (/reconfigure)", + (double) metrics.n_reconfigure + }, { + "reconfigure_failed_total", + "/reconfigure requests refused", + (double) metrics.n_reconfigure_failed }, { "ctx_checkpoints_evicted_budget_total", "Context checkpoints dropped to keep within --ctx-checkpoints-bytes", diff --git a/tools/server/server-task.h b/tools/server/server-task.h index d9f420121944..16b31d0f66d5 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -60,6 +60,7 @@ enum server_task_type { SERVER_TASK_TYPE_PAGE_IN, // named page -> slot SERVER_TASK_TYPE_PAGE_SPILL, // hands a page to the HTTP thread, which writes it to disk SERVER_TASK_TYPE_PAGE_DROP, // frees a page + SERVER_TASK_TYPE_RECONFIGURE, // Continuum fork: new context geometry, same model (server-reconfigure.h) SERVER_TASK_TYPE_GET_LORA, SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, @@ -265,6 +266,10 @@ struct server_task { // used by SERVER_TASK_TYPE_LEARN learn_request learn; + // used by SERVER_TASK_TYPE_RECONFIGURE: the request body, parsed on the task thread against the + // geometry it then has (server_reconfigure_request::parse) + json reconfigure; + server_task() = default; server_task(server_task_type type) : type(type) {} diff --git a/tools/server/server.cpp b/tools/server/server.cpp index f449dc00f11e..9eec965e8283 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -220,6 +220,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; routes.post_learn = models_routes->proxy_post; + routes.post_reconfigure = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -275,6 +276,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); ctx_http.post("/learn", ex_wrapper(routes.post_learn)); + ctx_http.post("/reconfigure", ex_wrapper(routes.post_reconfigure)); // Continuum fork: new geometry, same model ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); ctx_http.post("/train/cancel", ex_wrapper(routes.post_train_cancel)); diff --git a/tools/server/tests/unit/test_reconfigure.py b/tools/server/tests/unit/test_reconfigure.py new file mode 100644 index 000000000000..d3be16d51ebc --- /dev/null +++ b/tools/server/tests/unit/test_reconfigure.py @@ -0,0 +1,281 @@ +import os +import platform +import re +import subprocess +import threading +import time + +import pytest +from utils import * + +# Continuum fork: POST /reconfigure changes the context geometry without reloading the model +# (server-reconfigure.h; continuum docs/architecture/SERVING-ENGINE.md, slice S2). The residents' +# states are paged to host RAM, the contexts are rebuilt on the loaded model, and the states go back. + +server = ServerPreset.tinyllama2() + +HYBRID = os.environ.get("LLAMA_SERVER_TEST_HYBRID_GGUF") + +MIND_A = "Once upon a time, a little fox lived at the edge of a dark forest." +MIND_B = "The old lighthouse keeper climbed the stairs every night to light the lamp." +FOLLOW = " One day," + + +def use_model(kind): + global server + server = ServerPreset.tinyllama2() + server.temperature = 0.0 + server.server_metrics = True + server.server_slots = True + server.slot_save_path = "./tmp" + server.n_slots = 2 + server.n_ctx = 1024 + if kind == "hybrid": + if not HYBRID: + pytest.skip("set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF") + server.model_hf_repo = None + server.model_file = HYBRID + server.n_batch = 512 + server.n_ctx = 4096 + + +def complete(prompt, id_slot, n_predict=16, ignore_eos=False): + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, "id_slot": id_slot, "cache_prompt": True, "n_predict": n_predict, + "ignore_eos": ignore_eos, + }, timeout=600) + assert res.status_code == 200, res.body + return res.body + + +def reconfigure(body): + return server.make_request("POST", "/reconfigure", data=body, timeout=600) + + +def props(): + return server.make_request("GET", "/props").body + + +def metric(name): + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + for line in res.body.splitlines(): + if line.startswith(f"llamacpp:{name} "): + return float(line.split()[1]) + raise AssertionError(f"metric {name} missing") + + +def erase(id_slot): + assert server.make_request("POST", f"/slots/{id_slot}?action=erase").status_code == 200 + + +@pytest.mark.parametrize("kind", ["attention", "hybrid"]) +def test_reconfigure_keeps_both_minds_warm_and_identical(kind): + # what this catches: a reconfigure that reloads the model, restores the KV without the slot's + # tokens (the next act re-prefills from 0), restores it corrupt (the output differs from a mind + # that was never reconfigured), or leaves /props reporting the old geometry. + use_model(kind) + server.start() + n_ctx0 = props()["context_geometry"]["allocated_tokens"] + + complete(MIND_A, 0) + complete(MIND_B, 1) + + up = reconfigure({"n_ctx": n_ctx0 * 2, "n_parallel": 3}) + assert up.status_code == 200, up.body + r = up.body + assert r["before"]["n_parallel"] == 2 and r["after"]["n_parallel"] == 3 + assert r["after"]["n_ctx"] >= n_ctx0 * 2 + assert r["n_restored"] == 2 and r["n_kept_in_store"] == 0 and r["n_dropped"] == 0 + assert not r["rolled_back"] and not r["cold_start"] + assert r["n_bytes"] > 0 and all(p["fate"] == "restored" for p in r["pages"]) + for key in ("t_drain_ms", "t_page_out_ms", "t_free_ms", "t_init_ms", "t_page_in_ms", "t_total_ms"): + assert r[key] >= 0 + p = props() + assert p["total_slots"] == 3 + assert p["context_geometry"]["sequence_capacity"] == 3 + assert p["context_geometry"]["allocated_tokens"] == r["after"]["n_ctx"] + assert metric("reconfigure_total") == 1 + + a = complete(MIND_A + FOLLOW, 0) + b = complete(MIND_B + FOLLOW, 1) + c = complete("A third mind arrives.", 2) # the new seat serves + assert c["timings"]["prompt_n"] == c["tokens_evaluated"] + + # back down: slot 2's state has no seat among 2 occupied slots, so it stays a named page + down = reconfigure({"n_ctx": n_ctx0, "n_parallel": 2}) + assert down.status_code == 200, down.body + d = down.body + assert d["after"]["n_parallel"] == 2 and d["n_restored"] == 2 and d["n_kept_in_store"] == 1 + kept = [p for p in d["pages"] if p["fate"] == "kept_in_store"][0] + assert kept["id_slot_before"] == 2 and kept["page"] and kept["reason"] + assert props()["total_slots"] == 2 + a2 = complete(MIND_A + FOLLOW + a["content"] + FOLLOW, 0) + + # the reference: the same acts on slots that were never reconfigured + erase(0) + erase(1) + complete(MIND_A, 0) + complete(MIND_B, 1) + ref_a = complete(MIND_A + FOLLOW, 0) + ref_b = complete(MIND_B + FOLLOW, 1) + ref_a2 = complete(MIND_A + FOLLOW + ref_a["content"] + FOLLOW, 0) + + for got, ref in ((a, ref_a), (b, ref_b), (a2, ref_a2)): + assert got["content"] == ref["content"] + assert got["timings"]["prompt_n"] == ref["timings"]["prompt_n"] + assert got["timings"]["prompt_n"] < got["tokens_evaluated"] # the prefix survived + + # the kept page comes back into a slot like any page + restored = server.make_request("POST", f"/slots/1?action=page_in&page={kept['page']}") + assert restored.status_code == 200, restored.body + assert restored.body["n_tokens"] == kept["n_tokens"] + + +def long_act_model(): + # a request must outlast an HTTP round trip: the 0.8B hybrid decodes ~100x slower than stories260K + use_model("hybrid" if HYBRID else "attention") + server.n_ctx = 8192 + server.start() + return 1500 if not HYBRID else 300 + + +def wait_processing(id_slot): + while not server.make_request("GET", "/slots").body[id_slot]["is_processing"]: + time.sleep(0.005) + + +def test_an_in_flight_request_pauses_and_finishes_identically(): + # what this catches: a reconfigure that drops or corrupts a request mid-generation (its KV, + # sampler or speculative state lost with the freed context), or one that blocks until it ends. + n_predict = long_act_model() + prompt = "Tell me a long story about a dragon." + reference = complete(prompt, 0, n_predict=n_predict, ignore_eos=True) + erase(0) + + out = {} + + def generate(): + out["res"] = complete(prompt, 0, n_predict=n_predict, ignore_eos=True) + + t = threading.Thread(target=generate) + t.start() + wait_processing(0) + # grow, then shrink back, both with the request in flight + caught = 0 + for n_ctx in (16384, 8192): + r = reconfigure({"n_ctx": n_ctx, "n_parallel": 2}) + assert r.status_code == 200, r.body + caught += sum(1 for p in r.body["pages"] if p["in_flight"]) + t.join() + assert caught >= 1, "the generation ended before a reconfigure caught it in flight" + assert out["res"]["content"] == reference["content"] + assert out["res"]["tokens_predicted"] == reference["tokens_predicted"] + + +def test_a_slot_past_the_new_count_drains_and_requests_queue(): + # what this catches: a shrink that kills the request on a removed slot, or a request that + # arrives during the pause and is lost instead of queued. + n_predict = long_act_model() + out = {} + + def generate(): + out["res"] = complete("Tell me a long story about a castle.", 1, n_predict=n_predict, ignore_eos=True) + + def shrink(): + out["shrink"] = reconfigure({"n_parallel": 1}) + + t = threading.Thread(target=generate) + t.start() + wait_processing(1) + s = threading.Thread(target=shrink) + s.start() + time.sleep(0.1) # the shrink is pending: admission is closed + late = complete("A request during the pause.", -1) # queued, then served after + t.join() + s.join() + assert out["res"]["tokens_predicted"] == n_predict # the drained request finished whole + assert out["shrink"].status_code == 200, out["shrink"].body + assert out["shrink"].body["after"]["n_parallel"] == 1 + assert out["shrink"].body["t_drain_ms"] > 0 + assert out["shrink"].body["n_tasks_deferred"] >= 1 + assert late["tokens_predicted"] > 0 + + +def test_refusals_are_typed_and_change_nothing(): + # what this catches: a misspelled field that silently keeps the old geometry, a KV type change + # that carries pages of the wrong format, and a shrink that truncates a mind's state. + use_model("attention") + server.start() + first = complete(MIND_A, 0) + assert reconfigure({"n_slots": 3}).status_code == 400 + assert reconfigure({"n_parallel": 0}).status_code == 400 + assert reconfigure({"cache_type_k": "f32"}).status_code == 400 # needs force + assert reconfigure({"cache_type_k": "nope", "force": True}).status_code == 400 + again = complete(MIND_A, 0) + assert again["timings"]["prompt_n"] == 1 # nothing moved + assert again["content"] == first["content"] + assert metric("reconfigure_failed_total") >= 3 + + # a window smaller than the mind: her state stays whole in the store, never truncated + long_prompt = "The cat sat on the mat and looked at the moon. " * 16 + n_long = len(server.make_request("POST", "/tokenize", data={"content": long_prompt}).body["tokens"]) + assert 280 < n_long < 480, n_long # past one 256-token sequence, inside a 512 slot + complete(long_prompt, 0) + small = reconfigure({"n_ctx": 256, "n_parallel": 1}) + assert small.status_code == 200, small.body + kept = small.body["pages"][0] + assert kept["fate"] == "kept_in_store" and "do not fit" in kept["reason"] + assert reconfigure({"n_ctx": 1024, "n_parallel": 2}).status_code == 200 + back = server.make_request("POST", f"/slots/0?action=page_in&page={kept['page']}") + assert back.status_code == 200 and back.body["n_tokens"] == kept["n_tokens"] + + # a geometry libllama cannot build rolls back with every mind intact (a quantized V cache + # needs flash attention, which this model's head size does not get) + fail = reconfigure({"cache_type_k": "q8_0", "cache_type_v": "q8_0", "force": True}) + assert fail.status_code == 200, fail.body + assert fail.body["rolled_back"] and not fail.body["cold_start"] and fail.body["n_restored"] == 1 + assert fail.body["after"]["cache_type_v"] == "f16" + warm = complete(long_prompt, 0) + assert warm["timings"]["prompt_n"] == 1 + + # with force, a KV type change is an intentional cold start, counted (f32: this model's + # head size is not a multiple of a quantized block) + cold = reconfigure({"cache_type_k": "f32", "force": True}) + assert cold.status_code == 200, cold.body + assert cold.body["cold_start"] and cold.body["n_dropped"] == 1 and cold.body["n_bytes_dropped"] > 0 + assert cold.body["after"]["cache_type_k"] == "f32" + fresh = complete(MIND_A, 0) + assert fresh["timings"]["prompt_n"] == fresh["tokens_evaluated"] + + +CYCLES = int(os.environ.get("LLAMA_SERVER_TEST_RECONFIGURE_CYCLES", "0")) + + +def footprint_bytes(pid): + # macOS: phys_footprint, what the OS charges the process (RSS undercounts GPU-mapped memory) + out = subprocess.run(["footprint", "-p", str(pid)], capture_output=True, text=True).stdout + m = re.search(r"phys_footprint:\s*([\d.]+)\s*([KMG]B)", out) + assert m, out + return float(m.group(1)) * {"KB": 2**10, "MB": 2**20, "GB": 2**30}[m.group(2)] + + +@pytest.mark.skipif(not CYCLES or platform.system() != "Darwin", reason="set LLAMA_SERVER_TEST_RECONFIGURE_CYCLES (macOS footprint)") +def test_footprint_is_flat_over_many_cycles(): + # what this catches: a reconfigure that leaks a context, a draft context, its compute buffers + # or the paged copies each cycle (the plan line 3b gate: flat over 100 resize cycles). + use_model("hybrid" if HYBRID else "attention") + server.start() + n_ctx0 = props()["context_geometry"]["allocated_tokens"] + complete(MIND_A, 0) + complete(MIND_B, 1) + samples = [] + for i in range(CYCLES): + assert reconfigure({"n_ctx": n_ctx0 * 2, "n_parallel": 3}).status_code == 200 + assert reconfigure({"n_ctx": n_ctx0, "n_parallel": 2}).status_code == 200 + samples.append(footprint_bytes(server.process.pid)) + print("footprint MB per cycle:", [round(s / 2**20, 1) for s in samples]) + # after warm-up, the last cycle's footprint is within 2% of the tenth's + assert samples[-1] <= samples[min(9, len(samples) - 1)] * 1.02 + again = complete(MIND_A + FOLLOW, 0) + assert again["timings"]["prompt_n"] < again["tokens_evaluated"]