diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index 530ace7cda89..f157fde20ad8 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -166,7 +166,13 @@ jobs: -DLLAMA_BUILD_BORINGSSL=ON ^ -DGGML_SCHED_NO_REALLOC=ON set /A NINJA_JOBS=%NUMBER_OF_PROCESSORS%-1 - cmake --build build --config Release -j %NINJA_JOBS% --target llama-server + cmake --build build --config Release -j %NINJA_JOBS% --target llama-server test-server-pages + + - name: Page publication tests + id: server_page_tests + shell: cmd + run: | + ctest --test-dir build -C Release -R test-server-pages --output-on-failure - name: Python setup id: setup_python 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/tests/CMakeLists.txt b/tests/CMakeLists.txt index 498b70553bba..2d16f336f82c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -156,6 +156,13 @@ llama_build_and_test(test-grad-leaf.cpp) # ggml-opt with per-step graphs (the llama training path) starts each optimizer period from zero accumulators llama_build_and_test(test-opt-dynamic-accum.cpp) +# Continuum fork: the RAM page store's ceiling and the page files' whole-generation publication. +# Built on every platform: it links only server-context, as llama-server does, and Windows is +# where its durable-move adapter differs most. +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) + if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-unicode.cpp) @@ -172,6 +179,9 @@ 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) + llama_build_and_test(test-page-checkpoints.cpp) + target_include_directories(test-page-checkpoints PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) + target_link_libraries(test-page-checkpoints 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-page-checkpoints.cpp b/tests/test-page-checkpoints.cpp new file mode 100644 index 000000000000..34ab955ce985 --- /dev/null +++ b/tests/test-page-checkpoints.cpp @@ -0,0 +1,114 @@ +// What a saved page keeps of a slot's context checkpoints (Continuum fork, select_page_checkpoints). +// +// what this catches: the M5 on 2026-10-25 (a hybrid model, 4 residents paging over 3 lanes) re- +// prefilled 20k+ tokens a turn and once fell to "full re-prefill: no checkpoint passed the search +// (raw lcp 9710)": a page kept the oldest checkpoint and the newest four, all near the end of a +// 35-40k prompt, so a restored slot had nothing under the point where its next prompt diverged +// (about 11k). The page must keep, within the same budget, checkpoints just under the slot's last +// raw common prefix. +// +// The byte budget's eviction (server_checkpoint_victim, server_checkpoints_fit) dropped each +// slot's oldest checkpoint after its first: exactly the ones under the last divergence. It now +// gives up what the page subset gives up first, so a slot keeps its head, the newest at or below +// its last raw common prefix and its newest longest. + +#include "server-checkpoints.h" +#include "server-common.h" + +#include +#include +#include + +#define CHECK(cond) do { if (!(cond)) { std::fprintf(stderr, "%s:%d: CHECK failed: %s\n", __FILE__, __LINE__, #cond); std::exit(1); } } while (0) + +int main() { + // A 35-40k prompt's checkpoints at user-message starts, about 1,576 apart. + std::vector pos; + for (int64_t p = 300; p < 38000; p += 1576) { + pos.push_back(p); + } + const size_t n = pos.size(); + + // keep < 0: every checkpoint, as before the subset existed. + CHECK(select_page_checkpoints(pos, -1, 11000).size() == n); + + // No divergence known (a fresh slot, no cache_prompt): the oldest and the newest four, as today. + { + auto k = select_page_checkpoints(pos, 4, -1); + CHECK(k.size() == 5); + CHECK(k[0] == 0); + for (size_t j = 1; j < 5; ++j) { + CHECK(k[j] == n - 5 + j); + } + } + + // Divergence at 11,000: the oldest, the two newest at or below 11,000, the newest two. + { + auto k = select_page_checkpoints(pos, 4, 11000); + CHECK(k.size() == 5); + CHECK(k[0] == 0); + CHECK(pos[k[1]] <= 11000 && pos[k[2]] <= 11000); + CHECK(pos[k[2]] + 1576 > 11000); // the nearest under the break is kept + CHECK(k[3] == n - 2 && k[4] == n - 1); // and the end of the prompt + } + + // The budget never grows: at most 1 + keep, whatever the divergence. + for (int32_t keep = 0; keep <= 6; ++keep) { + CHECK(select_page_checkpoints(pos, keep, 11000).size() <= (size_t) keep + 1); + } + + // A divergence under every checkpoint but the first keeps no extra "below" points. + { + auto k = select_page_checkpoints(pos, 4, 100); + CHECK(k.size() == 5 && k[0] == 0 && k[1] == n - 4); + } + + // keep 0: the oldest only. An empty list: nothing. + CHECK(select_page_checkpoints(pos, 0, 11000).size() == 1); + CHECK(select_page_checkpoints({}, 4, 11000).empty()); + + // Eviction gives up what the subset gives up first: never the head, never the two newest at + // or below the divergence, never the newest; with no divergence known, the oldest after the + // head, as before. + { + auto list_of = [&](size_t bytes) { + server_checkpoint_list list; + for (int64_t p : pos) { + auto & c = list.emplace_back(); + c.n_tokens = p; + c.data_tgt.resize(bytes); + } + return list; + }; + const auto list = list_of(1); + const size_t victim = server_checkpoint_victim(list, 11000); + const auto keep = select_page_checkpoints(pos, (int32_t) n - 2, 11000); + CHECK(keep.size() == n - 1); + for (size_t i : keep) { + CHECK(i != victim); + } + CHECK(victim != 0 && victim != n - 1); + CHECK(!(pos[victim] <= 11000 && pos[victim] + 2 * 1576 > 11000)); // not one of the two under the break + CHECK(server_checkpoint_victim(list, -1) == 1); // unknown divergence: oldest after the head + + // The budget across slots: the fuller list drops its victims until the new one fits, and + // the checkpoints just under its divergence survive. + auto a = list_of(10); + auto b = list_of(1); + std::vector lists = { &a, &b }; + const uint64_t budget = server_checkpoints_bytes(a) + server_checkpoints_bytes(b) - 25; + size_t evicted = 0; + CHECK(server_checkpoints_fit(lists, { 11000, -1 }, 1, budget, 10, [&](size_t i, const common_prompt_checkpoint &) { CHECK(i == 0); ++evicted; })); + CHECK(evicted == 4); + CHECK(a.front().n_tokens == pos[0] && a.back().n_tokens == pos[n - 1]); + size_t under = 0; + for (const auto & c : a) { + under += (c.n_tokens <= 11000 && c.n_tokens + 2 * 1576 > 11000) ? 1 : 0; + } + CHECK(under == 2); + CHECK(!server_checkpoints_fit(lists, { 11000, -1 }, 1, 5, 10, [&](size_t, const common_prompt_checkpoint &) { CHECK(false); })); + } + + std::printf("test-page-checkpoints: ok\n"); + return 0; +} diff --git a/tests/test-server-pages.cpp b/tests/test-server-pages.cpp new file mode 100644 index 000000000000..3ee1b144ce2f --- /dev/null +++ b/tests/test-server-pages.cpp @@ -0,0 +1,358 @@ +// 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; +// - a directory sync that fails (stage, commit record, destination) yet reports a publication or drops +// the commit record before the moved names are durable; +// - a repeated publication over existing page files (rename onto an existing name) that leaves the old one. + +#include "server-page-files.h" +#include "server-page-os.h" +#include "server-pages.h" + +#include +#include +#include +#include +#include +#include +#include +#include + + +using page_os::current_pid; + +#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); +} + +// `op` fails on every path ending in `suffix` (and, when given, containing `within`) while this lives +struct page_fault { + page_fault(server_page_fault op, std::string suffix, std::string within = "") { + server_page_files_inject_fault([=](server_page_fault o, const std::string & path) { + return o == op && path.find(within) != std::string::npos && path.size() >= suffix.size() && + path.compare(path.size() - suffix.size(), suffix.size(), suffix) == 0; + }); + } + ~page_fault() { server_page_files_inject_fault(nullptr); } +}; + +static bool publish_refused(const fs::path & file, const std::string & tag) { + try { + publish_generation(file, tag); + } catch (const server_page_error &) { + return true; + } + return false; +} + +static bool reader_refused(const fs::path & file) { + try { + server_page_reader reader(file.string()); + } catch (const server_page_error &) { + return true; + } + return false; +} + +[[noreturn]] static void disk_full(const server_page_put &) { + throw std::runtime_error("disk full"); +} + +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", 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); + + // a staged member's status cannot be read after the commit: never taken for already moved, so + // the record and stage keep custody until a reader can read it, then the generation is whole + const fs::path record = dir / "..page-staging" / "mind-a.bin.commit"; + { + page_fault fault(server_page_fault::status, "mind-a.bin.ckpt", "..page-staging"); + CHECK(publish_refused(file, "four")); + CHECK(fs::exists(record)); + CHECK(reader_refused(file) && fs::exists(record)); + } + { server_page_reader reader(file.string()); check_generation(file, "four"); } + 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_a_failed_dir_sync_never_reports_or_loses_a_generation(const fs::path & dir) { + const fs::path file = dir / "mind-d.bin"; + const fs::path staging = dir / "..page-staging"; + publish_generation(file, "one"); + + // the stage's sync fails (every sync fails; the stage's is the first reached): refused before + // the commit, the committed generation as it was + { + page_fault fault(server_page_fault::sync, ""); + CHECK(publish_refused(file, "two")); + } + { server_page_reader reader(file.string()); check_generation(file, "one"); } + check_no_leftovers(dir); + + // the commit record's dir sync fails: refused, but the record keeps custody; the next reader completes it + { + page_fault fault(server_page_fault::sync, "..page-staging"); + CHECK(publish_refused(file, "three")); + CHECK(fs::exists(staging / "mind-d.bin.commit")); + } + { server_page_reader reader(file.string()); check_generation(file, "three"); } + check_no_leftovers(dir); + + // the destination's sync fails after the files moved: refused, the record kept until a sync succeeds + { + page_fault fault(server_page_fault::sync, dir.filename().string()); + CHECK(publish_refused(file, "four")); + CHECK(fs::exists(staging / "mind-d.bin.commit")); + CHECK(reader_refused(file) && fs::exists(staging / "mind-d.bin.commit")); + } + { server_page_reader reader(file.string()); check_generation(file, "four"); } + check_no_leftovers(dir); +} + +// every publication after the first renames onto existing page files: each must replace them +static void test_repeated_publication_replaces_every_file(const fs::path & dir) { + const fs::path file = dir / "mind-e.bin"; + for (const char * tag : { "one", "two", "three", "four" }) { + publish_generation(file, tag); + { server_page_reader reader(file.string()); check_generation(file, tag); } + 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_a_failed_dir_sync_never_reports_or_loses_a_generation(dir); + test_repeated_publication_replaces_every_file(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 141d44126731..c526abdb3586 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -16,6 +16,13 @@ add_library(${TARGET} STATIC server-context.cpp server-train.cpp server-context.h + server-pages.cpp + server-pages.h + server-page-files.cpp + server-page-files.h + server-page-os.h + server-checkpoints.cpp + server-checkpoints.h server-stream.cpp server-stream.h server-tools.cpp @@ -30,6 +37,13 @@ if (BUILD_SHARED_LIBS) set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON) endif() +# the page files' OS adapter: each OS supplies only its own calls +if (WIN32) + target_sources(${TARGET} PRIVATE server-page-os-win32.cpp) +else() + target_sources(${TARGET} PRIVATE server-page-os-posix.cpp) +endif() + target_include_directories(${TARGET} PRIVATE ../mtmd) target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR}) target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT}) diff --git a/tools/server/server-checkpoints.cpp b/tools/server/server-checkpoints.cpp new file mode 100644 index 000000000000..5f72fcad5660 --- /dev/null +++ b/tools/server/server-checkpoints.cpp @@ -0,0 +1,67 @@ +#include "server-checkpoints.h" + +#include "server-common.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; +} + +size_t server_checkpoint_victim(const server_checkpoint_list & list, int32_t raw_lcp) { + if (list.size() < 2) { + return 0; + } + std::vector n_tokens; + n_tokens.reserve(list.size()); + for (const auto & c : list) { + n_tokens.push_back(c.n_tokens); + } + const auto kept = select_page_checkpoints(n_tokens, (int32_t) list.size() - 2, raw_lcp); + size_t k = 0; + for (size_t i = 0; i < list.size(); ++i) { + if (k < kept.size() && kept[k] == i) { + ++k; + } else { + return i; + } + } + return list.size() - 1; +} + +bool server_checkpoints_fit( + const std::vector & lists, + const std::vector & raw_lcps, + 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]; + const int32_t raw_lcp = victim < raw_lcps.size() ? raw_lcps[victim] : -1; + auto it = std::next(list.begin(), (std::ptrdiff_t) server_checkpoint_victim(list, raw_lcp)); + 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..981d3ba69da1 --- /dev/null +++ b/tools/server/server-checkpoints.h @@ -0,0 +1,34 @@ +#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); + +// The checkpoint a list gives up first: the one select_page_checkpoints leaves out when it keeps +// all but one, so a list keeps longest exactly what a saved page keeps (its head, the newest at or +// below the slot's last raw common prefix `raw_lcp`, then the newest), never its oldest first. +size_t server_checkpoint_victim(const server_checkpoint_list & list, int32_t raw_lcp); + +// 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 +// server_checkpoint_victim under its slot's `raw_lcps` entry. `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, + const std::vector & raw_lcps, + size_t own, + uint64_t budget, + size_t incoming, + const std::function & on_evict); diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 3269f86b2199..5dd7b097cd22 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -9,6 +9,7 @@ #include "server-common.h" +#include #include #include #include @@ -1844,3 +1845,41 @@ server_tokens format_prompt_rerank( return result; } + +std::vector select_page_checkpoints(const std::vector & n_tokens, int32_t keep, int32_t raw_lcp) { + const size_t n = n_tokens.size(); + std::vector kept; + if (keep < 0) { + for (size_t i = 0; i < n; ++i) { + kept.push_back(i); + } + return kept; + } + std::vector take(n, false); + if (n > 0) { + take[0] = true; + } + size_t budget = (size_t) keep; + if (raw_lcp >= 0) { + size_t below = std::min(2, budget / 2); + for (size_t i = n; i-- > 1 && below > 0;) { + if (n_tokens[i] <= raw_lcp && !take[i]) { + take[i] = true; + --below; + --budget; + } + } + } + for (size_t i = n; i-- > 1 && budget > 0;) { + if (!take[i]) { + take[i] = true; + --budget; + } + } + for (size_t i = 0; i < n; ++i) { + if (take[i]) { + kept.push_back(i); + } + } + return kept; +} diff --git a/tools/server/server-common.h b/tools/server/server-common.h index c55d5b79be80..f364464f4830 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(); } @@ -613,3 +620,12 @@ struct server_pipe { return true; } }; + +// The context checkpoints a saved page keeps (Continuum fork), as indices into the slot's +// checkpoint list (positions in engine tokens, oldest first), in list order. `keep < 0` keeps +// all. Otherwise the first checkpoint is always kept, and `keep` more: when the slot's last raw +// common prefix with a request is known (`raw_lcp >= 0`), up to min(2, keep / 2) of them are the +// newest checkpoints at or below it, so a restored page can resume from just under the point +// where its prompt last diverged; the rest are the newest (the end of the prompt). A page keeps +// at most 1 + keep checkpoints either way. +std::vector select_page_checkpoints(const std::vector & n_tokens, int32_t keep, int32_t raw_lcp); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 784d5426eeba..c8597aaae4bf 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2,6 +2,9 @@ #include "server-context.h" #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" #include "server-http.h" @@ -31,22 +34,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 @@ -245,6 +232,11 @@ struct server_slot { // generation props int32_t n_ctx = 0; // context size per slot + + // The raw common prefix (engine tokens) of this slot's prompt with its last request, before + // any checkpoint fallback, or -1 when unknown (no cache_prompt, a restored or erased slot). + // A saved page keeps checkpoints at or below it (select_page_checkpoints, Continuum fork). + int32_t resume_raw_lcp = -1; int32_t n_keep = 0; int32_t i_batch = -1; @@ -938,6 +930,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 +955,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.clear(); + } spec.reset(); spec_init.reset(); @@ -1352,7 +1353,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 +2478,36 @@ 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; + std::vector raw_lcps; + size_t own = 0; + for (auto & s : slots) { + if (&s == &slot) { + own = lists.size(); + } + lists.push_back(&s.prompt.checkpoints); + raw_lcps.push_back(s.resume_raw_lcp); + } + const bool fits = server_checkpoints_fit(lists, raw_lcps, 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 +2537,106 @@ 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"); + } + // 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(); + 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)) { + 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_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 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; + } + // 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 +2780,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(); @@ -2713,55 +2853,51 @@ struct server_context_impl { } GGML_ASSERT(packed.size() % sizeof(llama_token) == 0); - // The checkpoints this page carries: the oldest (the head of the first user - // turn) and the newest `keep_newest_checkpoints`, in list order, or all of them - // when the caller asked for no subset. The quote and the write share this list. + // The checkpoints this page carries (select_page_checkpoints): the oldest (the + // head of the first user turn), up to two just under where the slot's prompt last + // diverged from a request (its raw common prefix), and the newest, at most + // 1 + `keep_newest_checkpoints`, or all of them when the caller asked for no + // subset. A restored page then resumes from under its last break, not from 0 + // (the M5 2026-10-25: "full re-prefill: no checkpoint passed the search"). + // The quote and the write share this list. std::vector kept_checkpoints; { - const auto & all = slot->prompt.checkpoints; - const int32_t keep = task.slot_action.keep_newest_checkpoints; - const size_t n = all.size(); - size_t i = 0; + const auto & all = slot->prompt.checkpoints; + std::vector positions; for (const auto & c : all) { - if (keep < 0 || i == 0 || i + (size_t) keep >= n) { + positions.push_back(c.n_tokens); + } + const auto kept = select_page_checkpoints( + positions, task.slot_action.keep_newest_checkpoints, slot->resume_raw_lcp); + size_t i = 0, k = 0; + for (const auto & c : all) { + if (k < kept.size() && kept[k] == i) { kept_checkpoints.push_back(&c); + ++k; } ++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 +2917,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; @@ -2802,27 +2941,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()); - } + // 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) { + files.push_back(server_page_sidecar(".dft", write_draft)); // without one, an older `.dft` is removed } - { - 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()); - } + 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(); @@ -2852,6 +2989,8 @@ struct server_context_impl { queue_tasks.defer(std::move(task)); break; } + // Its prompt is replaced or dropped: the last divergence no longer applies. + slot->resume_raw_lcp = -1; auto progress = task.slot_action.progress; if (progress) { @@ -2868,6 +3007,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), @@ -2985,6 +3126,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"); } @@ -3024,6 +3182,8 @@ struct server_context_impl { queue_tasks.defer(std::move(task)); break; } + // Its prompt is replaced or dropped: the last divergence no longer applies. + slot->resume_raw_lcp = -1; // Erase token cache const size_t n_erased = slot->prompt.tokens.size(); @@ -3036,6 +3196,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 +3381,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. @@ -3672,6 +3859,7 @@ struct server_context_impl { int n_past = 0; // the raw common prefix with the slot's tokens, before any reuse adjustment (#56) int resume_raw_lcp = -1; + slot.resume_raw_lcp = -1; // unknown until this request computes it // empty prompt passed -> release the slot and send empty response if (input_tokens.empty()) { @@ -3733,6 +3921,7 @@ struct server_context_impl { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); resume_raw_lcp = n_past; + slot.resume_raw_lcp = n_past; // if there is an alora invoked, don't cache after the invocation start if (slot.alora_invocation_start > 0) { @@ -5213,6 +5402,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 +5578,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 +5597,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 +6488,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-page-files.cpp b/tools/server/server-page-files.cpp new file mode 100644 index 000000000000..e846a24e21bd --- /dev/null +++ b/tools/server/server-page-files.cpp @@ -0,0 +1,416 @@ +#include "server-page-files.h" +#include "server-page-os.h" + +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include + + +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; +} + +// 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 != page_os::current_pid()) { + return page_os::process_alive(pid); + } + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + return pubs.stages.count(id) > 0; +} + +static std::function & injected_fault() { + static std::function fault; + return fault; +} + +void server_page_files_inject_fault(std::function fault) { + injected_fault() = std::move(fault); +} + +static bool faulted(server_page_fault op, const fs::path & path) { + return injected_fault() && injected_fault()(op, path.string()); +} + +// Whether `path` exists, or throws: an I/O or permission error is never read as absence, since +// recovery takes an absent staged file for one already moved. +static bool present(const fs::path & path) { + std::error_code ec; + const auto st = fs::symlink_status(path, ec); + if (faulted(server_page_fault::status, path) || (ec && st.type() != fs::file_type::not_found)) { + throw server_page_error(ERROR_TYPE_SERVER, "a page file's status could not be read: " + path.string() + (ec ? ": " + ec.message() : "")); + } + return st.type() != fs::file_type::not_found; +} + +// Makes a dir's entries (a rename into it) durable, or throws: a rename not proven durable is never +// reported as published. +static void sync_dir(const fs::path & dir) { + if (faulted(server_page_fault::sync, dir) || !page_os::sync_dir(dir)) { + throw server_page_error(ERROR_TYPE_SERVER, "a page directory could not be synced: " + dir.string()); + } +} + +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()); + } + } + if (!page_os::sync_file(tmp) || page_os::place(tmp, commit_record(file))) { + throw server_page_error(ERROR_TYPE_SERVER, "the page's commit record could not be placed: " + commit_record(file).string()); + } +} + +// Moves a committed generation over the page's files, retires the files it lacks into its stage, then +// removes the 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(); + for (const auto & s : suffixes) { + const fs::path staged = stage / (name + s); + if (!present(staged)) { + continue; + } + if (const auto ec = page_os::place(staged, file.string() + s)) { + 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) { + const fs::path older = file.string() + s; + if (std::find(suffixes.begin(), suffixes.end(), s) != suffixes.end() || !present(older)) { + continue; + } + // a move, not a delete: on Windows only a move is made durable before the record goes + if (const auto ec = page_os::place(older, stage / (std::string("..retired") + s))) { + throw server_page_error(ERROR_TYPE_SERVER, "an older page file could not be removed: " + older.string() + ": " + ec.message()); + } + } + sync_dir(file.parent_path()); // throws, keeping the record: the moved names are not yet durable + std::error_code ec; + fs::remove_all(stage, ec); + fs::remove(commit_record(file), ec); + try { + sync_dir(staging_dir(file)); + } catch (const server_page_error & err) { + // the generation is durable; a record whose removal is lost names a stage already moved, + // and the next reader's roll_forward of it changes nothing + SRV_WRN("%s\n", err.what()); + } +} + +// 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(page_os::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); + const fs::path written = stage / ("..written" + f.suffix); // placed, so its name is durable too + size_t n = 0; + try { + n = f.write(written.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 || !page_os::sync_file(written) || page_os::place(written, 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; // the record names the stage: from here a failure leaves it to the next reader + sync_dir(staging_dir(file)); + 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..7e1d7d3a2110 --- /dev/null +++ b/tools/server/server-page-files.h @@ -0,0 +1,74 @@ +#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); + +// What a test seam can make fail: a directory's sync, or reading a file's status. +enum class server_page_fault { sync, status }; + +// Test seam: `fault(op, path)` returning true makes that operation on that path fail, as a full or +// failing disk would. Empty (the default) runs for real. +void server_page_files_inject_fault(std::function fault); + +// 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-page-os-posix.cpp b/tools/server/server-page-os-posix.cpp new file mode 100644 index 000000000000..0f5b61cec169 --- /dev/null +++ b/tools/server/server-page-os-posix.cpp @@ -0,0 +1,47 @@ +#include "server-page-os.h" + +#include +#include +#include +#include + +namespace page_os { + +long current_pid() { + return (long) getpid(); +} + +bool process_alive(long pid) { + return kill((pid_t) pid, 0) == 0 || errno == EPERM; +} + +static bool fsync_path(const std::filesystem::path & path, bool full) { + 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 what follows + bool ok = (full && fcntl(fd, F_FULLFSYNC) == 0) || fsync(fd) == 0; +#else + (void) full; + bool ok = fsync(fd) == 0; +#endif + return close(fd) == 0 && ok; +} + +bool sync_file(const std::filesystem::path & path) { + return fsync_path(path, true); +} + +bool sync_dir(const std::filesystem::path & dir) { + return fsync_path(dir, false); +} + +std::error_code place(const std::filesystem::path & from, const std::filesystem::path & to) { + std::error_code ec; + std::filesystem::rename(from, to, ec); + return ec; +} + +} diff --git a/tools/server/server-page-os-win32.cpp b/tools/server/server-page-os-win32.cpp new file mode 100644 index 000000000000..1622f30d170f --- /dev/null +++ b/tools/server/server-page-os-win32.cpp @@ -0,0 +1,49 @@ +#include "server-page-os.h" + +#ifndef NOMINMAX +# define NOMINMAX +#endif +#include +#include +#include +#include + +namespace page_os { + +long current_pid() { + return (long) _getpid(); +} + +bool process_alive(long pid) { + HANDLE h = OpenProcess(SYNCHRONIZE, FALSE, (DWORD) pid); + if (h == nullptr) { + return false; + } + const bool alive = WaitForSingleObject(h, 0) == WAIT_TIMEOUT; + CloseHandle(h); + return alive; +} + +bool sync_file(const std::filesystem::path & path) { + const int fd = _wopen(path.c_str(), _O_RDWR | _O_BINARY); + if (fd < 0) { + return false; + } + const bool ok = _commit(fd) == 0; + return _close(fd) == 0 && ok; +} + +// Windows has no directory sync; every entry was already made durable by place()'s write-through move. +bool sync_dir(const std::filesystem::path &) { + return true; +} + +// "MOVEFILE_WRITE_THROUGH: the function does not return until the file is actually moved on the disk." +std::error_code place(const std::filesystem::path & from, const std::filesystem::path & to) { + if (MoveFileExW(from.c_str(), to.c_str(), MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH)) { + return {}; + } + return std::error_code((int) GetLastError(), std::system_category()); +} + +} diff --git a/tools/server/server-page-os.h b/tools/server/server-page-os.h new file mode 100644 index 000000000000..bb9acdfaf54e --- /dev/null +++ b/tools/server/server-page-os.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +// The OS calls page publication needs. One signature each; each OS supplies only its own lines +// (server-page-os-posix.cpp, server-page-os-win32.cpp). The publication logic never branches on the OS. +namespace page_os { + +long current_pid(); + +// whether the process with this pid still runs +bool process_alive(long pid); + +// makes a file's contents durable before anything that depends on them +bool sync_file(const std::filesystem::path & path); + +// makes a dir's entries durable: every rename into it that place() made +bool sync_dir(const std::filesystem::path & dir); + +// moves `from` onto `to`, replacing it; durable once sync_dir of its dir returns +std::error_code place(const std::filesystem::path & from, const std::filesystem::path & to); + +} diff --git a/tools/server/server-pages.cpp b/tools/server/server-pages.cpp new file mode 100644 index 000000000000..2a5d98abcf27 --- /dev/null +++ b/tools/server/server-pages.cpp @@ -0,0 +1,175 @@ +#include "server-pages.h" + +#include "speculative.h" + +#include + +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; + for (const auto & c : prompt.checkpoints) { + n += c.size(); + } + return n; +} + +size_t server_page::n_bytes() const { + return server_page_bytes(data.size(), 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) { + 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, 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"); + } + 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 &) { + 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"); + } + page->spec = std::move(shape.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 live = n_bytes(); + const auto it = pages.find(name); + // 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 (live + n_bytes_new > limit_bytes) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, string_format( + "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, 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); + } + 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 { + const auto it = pages.find(name); + return it == pages.end() ? nullptr : it->second; +} + +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 false; + } + pages.erase(it); + return true; +} + +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)\n", it->first.c_str(), it->second->charge.n); + it = pages.erase(it); + } else { + ++it; + } + } +} + +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; +} diff --git a/tools/server/server-pages.h b/tools/server/server-pages.h new file mode 100644 index 000000000000..a94aabdad2bf --- /dev/null +++ b/tools/server/server-pages.h @@ -0,0 +1,127 @@ +#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 (server-page-files.h) 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) {} +}; + +// 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 { + // declared first so it is destroyed LAST: the bytes stay counted until every buffer below is freed + server_page_charge charge; + + 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; // server_page_bytes of what the page holds +}; + +using server_page_ptr = std::shared_ptr; + +// 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): 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: 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 + // 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. 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 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); + + // 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); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 746c3a1f98c8..12c1e5a008fe 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,34 @@ 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_bytes_spec", n_bytes_spec }, + { "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..d9f420121944 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,33 @@ 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_bytes_spec = 0; // the speculative tail state + 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..069aff494fa7 --- /dev/null +++ b/tools/server/tests/unit/test_slot_page.py @@ -0,0 +1,251 @@ +import os +import subprocess +import sys +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 + + +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") + + +@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 + + +@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"] 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"