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/include/llama.h b/include/llama.h index 277fc9d881cb..b2144b206cbd 100644 --- a/include/llama.h +++ b/include/llama.h @@ -796,6 +796,20 @@ extern "C" { // NOTE: loaded adapters that are not manually freed will be freed when the associated model is deleted LLAMA_API void llama_adapter_lora_free(struct llama_adapter_lora * adapter); + // Continuum fork: page an adapter's tensors between their device buffers and host memory, so a + // loaded adapter no context applies costs host RAM instead of device memory, and comes back + // with a copy instead of a read from disk. The adapter stays loaded and keeps its identity; + // a paged-out adapter must be paged in before a context applies it. The caller serializes + // these with every use of the adapter (synchronize the contexts that ran it first). + // Return the bytes moved; 0 when it is already there; -1 when an allocation failed (the + // adapter is then left exactly as it was). + LLAMA_API int64_t llama_adapter_lora_page_out(struct llama_adapter_lora * adapter); + LLAMA_API int64_t llama_adapter_lora_page_in (struct llama_adapter_lora * adapter); + LLAMA_API bool llama_adapter_lora_is_paged_out(const struct llama_adapter_lora * adapter); + // The bytes the adapter's tensors take where they are now: its device buffers when resident, + // host memory when paged out. + LLAMA_API uint64_t llama_adapter_lora_n_bytes(const struct llama_adapter_lora * adapter); + // Get the invocation tokens if the current lora is an alora LLAMA_API uint64_t llama_adapter_get_alora_n_invocation_tokens(const struct llama_adapter_lora * adapter); LLAMA_API const llama_token * llama_adapter_get_alora_invocation_tokens (const struct llama_adapter_lora * adapter); diff --git a/src/llama-adapter.cpp b/src/llama-adapter.cpp index a6ac5596353b..47bd02a1a07e 100644 --- a/src/llama-adapter.cpp +++ b/src/llama-adapter.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -479,6 +480,97 @@ int32_t llama_adapter_meta_val_str_by_index(const llama_adapter_lora * adapter, return snprintf(buf, buf_size, "%s", it->second.c_str()); } +// Continuum fork: adapter paging between device buffers and host memory (llama.h). + +static void lora_for_each_tensor(const llama_adapter_lora & adapter, const std::function & fn) { + for (size_t i = 0; i < adapter.ctxs.size(); ++i) { + for (ggml_tensor * t = ggml_get_first_tensor(adapter.ctxs[i].get()); t; t = ggml_get_next_tensor(adapter.ctxs[i].get(), t)) { + fn(i, t); + } + } +} + +int64_t llama_adapter_lora_page_out(llama_adapter_lora * adapter) { + if (adapter->paged) { + return 0; + } + auto page = std::make_unique(); + page->bufts.assign(adapter->ctxs.size(), nullptr); + size_t total = 0; + lora_for_each_tensor(*adapter, [&](size_t i, ggml_tensor * t) { + if (page->bufts[i] == nullptr && t->buffer != nullptr) { + page->bufts[i] = ggml_backend_buffer_get_type(t->buffer); + } + total += ggml_nbytes(t); + }); + try { + page->host.resize(total); + } catch (const std::bad_alloc &) { + LLAMA_LOG_ERROR("%s: no host memory for %zu bytes of LoRA tensors\n", __func__, total); + return -1; + } + size_t off = 0; + lora_for_each_tensor(*adapter, [&](size_t, ggml_tensor * t) { + ggml_backend_tensor_get(t, page->host.data() + off, 0, ggml_nbytes(t)); + off += ggml_nbytes(t); + }); + adapter->bufs.clear(); // frees the device buffers + lora_for_each_tensor(*adapter, [](size_t, ggml_tensor * t) { + t->data = nullptr; + t->buffer = nullptr; + }); + adapter->paged = std::move(page); + return (int64_t) total; +} + +int64_t llama_adapter_lora_page_in(llama_adapter_lora * adapter) { + if (!adapter->paged) { + return 0; + } + const llama_adapter_lora_host_page & page = *adapter->paged; + std::vector bufs; + for (size_t i = 0; i < adapter->ctxs.size(); ++i) { + if (page.bufts[i] == nullptr) { + continue; // a context with no tensors + } + ggml_backend_buffer_ptr buf { ggml_backend_alloc_ctx_tensors_from_buft(adapter->ctxs[i].get(), page.bufts[i]) }; + if (!buf) { + LLAMA_LOG_ERROR("%s: failed to re-allocate a %s LoRA buffer\n", __func__, ggml_backend_buft_name(page.bufts[i])); + bufs.clear(); + lora_for_each_tensor(*adapter, [](size_t, ggml_tensor * t) { + t->data = nullptr; + t->buffer = nullptr; + }); + return -1; + } + bufs.emplace_back(std::move(buf)); + } + size_t off = 0; + lora_for_each_tensor(*adapter, [&](size_t, ggml_tensor * t) { + ggml_backend_tensor_set(t, page.host.data() + off, 0, ggml_nbytes(t)); + off += ggml_nbytes(t); + }); + adapter->bufs = std::move(bufs); + const int64_t moved = (int64_t) page.host.size(); + adapter->paged.reset(); + return moved; +} + +bool llama_adapter_lora_is_paged_out(const llama_adapter_lora * adapter) { + return adapter->paged != nullptr; +} + +uint64_t llama_adapter_lora_n_bytes(const llama_adapter_lora * adapter) { + if (adapter->paged) { + return adapter->paged->host.size(); + } + uint64_t n = 0; + for (const auto & buf : adapter->bufs) { + n += ggml_backend_buffer_get_size(buf.get()); + } + return n; +} + void llama_adapter_lora_free(llama_adapter_lora * adapter) { if (adapter == nullptr) { return; diff --git a/src/llama-adapter.h b/src/llama-adapter.h index 9102ff66ab0f..663de5f5f1ef 100644 --- a/src/llama-adapter.h +++ b/src/llama-adapter.h @@ -60,6 +60,15 @@ struct llama_adapter_lora_weight { llama_adapter_lora_weight(ggml_tensor * a, ggml_tensor * b) : a(a), b(b) {} }; +// Continuum fork: an adapter's tensors held in host memory while its device buffers are freed +// (llama_adapter_lora_page_out). `bufts` is aligned with `ctxs`: the buffer type each context's +// tensors are re-allocated from on page-in; `host` is every tensor's bytes, contexts in order and +// tensors in context order. +struct llama_adapter_lora_host_page { + std::vector bufts; + std::vector host; +}; + struct llama_adapter_lora { llama_model * model = nullptr; @@ -79,6 +88,9 @@ struct llama_adapter_lora { // activated lora (aLoRA) std::vector alora_invocation_tokens; + // non-null while paged out to host memory: the tensors then have no device data + std::unique_ptr paged; + explicit llama_adapter_lora(llama_model * model) : model(model) {} ~llama_adapter_lora() = default; 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..61a391fc803f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -172,6 +172,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) + # Continuum fork: the RAM page store's ceiling and the page files' whole-generation publication + llama_build_and_test(test-server-pages.cpp) + target_include_directories(test-server-pages PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) + target_link_libraries(test-server-pages PRIVATE server-context) # TODO: disabled on loongarch64 because the ggml-ci node lacks Python 3.8 if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "loongarch64") llama_build_and_test(test-json-schema-to-grammar.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) diff --git a/tests/test-server-pages.cpp b/tests/test-server-pages.cpp new file mode 100644 index 000000000000..b02df7f34758 --- /dev/null +++ b/tests/test-server-pages.cpp @@ -0,0 +1,268 @@ +// what this catches (Continuum fork, server-pages.h + server-page-files.h, no model needed): +// - a page admitted on a count that leaves out part of what it holds (its speculative state), so the +// store passes its ceiling after a successful page_out; +// - a page out of the store but still held by a spill that stops counting, so the live RAM passes the ceiling; +// - a refused page_out that has already moved the slot's state out (the incumbent is lost); +// - a publication that removes or overwrites the committed files before its own are whole: a failed +// sidecar write, or a crash between files, that loses the last durable page. + +#include "server-page-files.h" +#include "server-pages.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +static long current_pid() { return (long) _getpid(); } +#else +# include +static long current_pid() { return (long) getpid(); } +#endif + +#define CHECK(cond) do { if (!(cond)) { fprintf(stderr, "FAILED %s:%d: %s\n", __FILE__, __LINE__, #cond); exit(1); } } while (0) + +namespace fs = std::filesystem; + +// a page holding `n` bytes in each part, the speculative tail state included +static server_page_ptr make_page(size_t n) { + auto page = std::make_shared(); + page->data.main.assign(n, 1); + page->data.drft.assign(n, 2); + page->spec.assign(n, 3); + return page; +} + +// a slot's prompt: tokens and one checkpoint of `n` bytes +static server_prompt make_prompt(size_t n) { + server_prompt prompt; + prompt.tokens = server_tokens(llama_tokens{ 1, 2, 3 }, false); + auto & c = prompt.checkpoints.emplace_back(); + c.data_tgt.assign(n, 4); + return prompt; +} + +static error_type refusal(const std::function & fn) { + try { + fn(); + } catch (const server_page_error & err) { + return err.type; + } + return ERROR_TYPE_SERVER; // stands for "not refused"; every refusal checked here is UNAVAILABLE +} + +static void test_admission_counts_what_the_store_charges() { + const size_t n = 64; + const size_t bytes = server_page_bytes(2 * n, n, n); // main + draft, speculative, one checkpoint + { + server_page_store store; + store.limit_bytes = bytes; // fits exactly + auto page = make_page(n); + auto from = make_prompt(n); + store.admit("a", bytes); + store.put("a", page, &from); + CHECK(page->n_bytes() == bytes); + CHECK(store.n_bytes() == bytes); + CHECK(from.checkpoints.empty() && page->prompt.checkpoints.size() == 1); + } + { + server_page_store store; + store.limit_bytes = bytes - 1; // one byte short: the speculative state alone tips it + CHECK(refusal([&] { store.admit("a", bytes); }) == ERROR_TYPE_UNAVAILABLE); + // a put is checked on the same count, so a short pre-count cannot get a page past the ceiling + auto page = make_page(n); + auto from = make_prompt(n); + CHECK(refusal([&] { store.put("a", page, &from); }) == ERROR_TYPE_UNAVAILABLE); + // the refusal leaves the incumbent: the slot's prompt is where it was, and nothing is stored + CHECK(from.tokens.size() == 3 && from.checkpoints.size() == 1 && from.checkpoints.front().data_tgt.size() == n); + CHECK(page->prompt.tokens.size() == 0); + CHECK(store.n_bytes() == 0 && store.n_pages() == 0 && store.find("a") == nullptr); + } +} + +static void test_spill_held_pages_stay_counted() { + const size_t n = 64; + const size_t bytes = server_page_bytes(2 * n, n, n); + server_page_store store; + store.limit_bytes = 2 * bytes - 1; // room for one page, not two + + auto from = make_prompt(n); + store.put("a", make_page(n), &from); + server_page_ptr spilling = store.find("a"); // a spill's reference on its HTTP thread + + // replaced under its name while spilling: the old page is still live, so the new one does not fit + auto from2 = make_prompt(n); + CHECK(refusal([&] { store.put("a", make_page(n), &from2); }) == ERROR_TYPE_UNAVAILABLE); + CHECK(from2.checkpoints.size() == 1); + + // left the store (page_in, page_drop) while spilling: still counted + CHECK(store.erase("a")); + CHECK(store.n_pages() == 0 && store.n_bytes() == bytes); + CHECK(refusal([&] { store.put("b", make_page(n), &from2); }) == ERROR_TYPE_UNAVAILABLE); + + // the spill finishes: its bytes are released and the page fits + spilling.reset(); + CHECK(store.n_bytes() == 0); + store.put("b", make_page(n), &from2); + CHECK(store.n_bytes() == bytes); + + // a replacement the store solely owns frees the page it replaces, so it fits + auto from3 = make_prompt(n); + store.put("b", make_page(n), &from3); + CHECK(store.n_bytes() == bytes && store.n_pages() == 1); + + // a cleared store (the model unloaded) still counts a page a spill holds + spilling = store.find("b"); + store.clear(); + CHECK(store.n_bytes() == bytes); + spilling.reset(); + CHECK(store.n_bytes() == 0); +} + +static std::string read_file(const fs::path & p) { + std::ifstream in(p, std::ios::binary); + return in ? std::string(std::istreambuf_iterator(in), std::istreambuf_iterator()) : std::string(""); +} + +static void write_file(const fs::path & p, const std::string & bytes) { + std::ofstream out(p, std::ios::binary | std::ios::trunc); + out << bytes; +} + +static server_page_file text_file(const std::string & suffix, const std::string & text) { + return server_page_sidecar(suffix, [text](const server_page_put & put) { put(text.data(), text.size()); }); +} + +static const char * const SUFFIXES[] = { "", ".ckpt", ".lora", ".dft" }; + +static void check_generation(const fs::path & file, const std::string & tag) { + for (const char * s : SUFFIXES) { + CHECK(read_file(file.string() + s) == tag + s); + } +} + +// no stage and no commit record is left beside the page +static void check_no_leftovers(const fs::path & dir) { + const fs::path staging = dir / "..page-staging"; + if (fs::exists(staging)) { + CHECK(fs::is_empty(staging)); + } +} + +static void publish_generation(const fs::path & file, const std::string & tag) { + std::vector files; + for (const char * s : SUFFIXES) { + files.push_back(text_file(s, tag + s)); + } + server_page_publish(file.string(), files); +} + +static void test_a_failed_publication_keeps_the_committed_generation(const fs::path & dir) { + const fs::path file = dir / "mind-a.bin"; + publish_generation(file, "one"); + check_generation(file, "one"); + + // the .lora sidecar fails after the state and .ckpt of the new generation were written + const auto failing = server_page_file{ ".lora", [](const std::string &) -> size_t { return 0; } }; + bool refused = false; + try { + server_page_publish(file.string(), { text_file("", "two"), text_file(".ckpt", "two.ckpt"), failing, text_file(".dft", "two.dft") }); + } catch (const server_page_error &) { + refused = true; + } + CHECK(refused); + { server_page_reader reader(file.string()); check_generation(file, "one"); } + check_no_leftovers(dir); + + // a writer that throws is the same failure + refused = false; + try { + server_page_publish(file.string(), { text_file("", "two"), + server_page_sidecar(".ckpt", [](const server_page_put &) { throw std::runtime_error("disk full"); }) }); + } catch (const server_page_error &) { + refused = true; + } + CHECK(refused); + check_generation(file, "one"); + check_no_leftovers(dir); + + // a generation without a sidecar removes the older one: a stale .dft is never read beside new state + server_page_publish(file.string(), { text_file("", "three"), text_file(".ckpt", "three.ckpt"), text_file(".lora", "three.lora") }); + CHECK(read_file(file) == "three" && read_file(file.string() + ".dft") == ""); + check_no_leftovers(dir); +} + +static void test_an_interrupted_publication_restores_whole(const fs::path & dir) { + const fs::path file = dir / "mind-b.bin"; + const fs::path staging = dir / "..page-staging"; + publish_generation(file, "one"); + + // a crash after staging, before the commit record: this process's pid, a stage no publication owns + const fs::path stage = staging / (std::to_string(current_pid()) + "-999999"); + fs::create_directories(stage); + write_file(stage / "mind-b.bin", "two"); + write_file(stage / "mind-b.bin.ckpt", "two.ckpt"); + { server_page_reader reader(file.string()); check_generation(file, "one"); } + check_no_leftovers(dir); + + // a crash after the commit record, with one file moved: the reader completes the generation, + // and the sidecar the generation lacks (.dft) is removed + fs::create_directories(stage); + for (const char * s : { "", ".ckpt", ".lora" }) { + write_file(stage / (std::string("mind-b.bin") + s), std::string("two") + s); + } + fs::rename(stage / "mind-b.bin", file); + write_file(staging / "mind-b.bin.commit", stage.filename().string() + "\nf\nf.ckpt\nf.lora\n"); + { + server_page_reader reader(file.string()); + CHECK(read_file(file) == "two" && read_file(file.string() + ".ckpt") == "two.ckpt" && read_file(file.string() + ".lora") == "two.lora"); + CHECK(read_file(file.string() + ".dft") == ""); + } + check_no_leftovers(dir); +} + +static void test_spill_writes_a_page_restore_reads(const fs::path & dir) { + // a sequence state as llama_state_seq_get_data writes it: [io magic][seq id][state] + auto page = std::make_shared(); + const uint32_t io_magic = 0xaf143cd8; + page->data.main.resize(sizeof(uint32_t) + sizeof(llama_seq_id) + 16, 7); + memcpy(page->data.main.data(), &io_magic, sizeof(io_magic)); + page->data.drft.assign(8, 5); + page->spec.assign(4, 6); + page->prompt = make_prompt(32); + + const fs::path file = dir / "mind-c.bin"; + const size_t n = server_page_spill(*page, file.string()); + size_t on_disk = 0; + for (const char * s : SUFFIXES) { + on_disk += fs::file_size(file.string() + s); + } + CHECK(n == on_disk); + std::vector dft, spec; + std::string why; + CHECK(server_page_read_draft(file.string() + ".dft", dft, spec, why)); + CHECK(dft == page->data.drft && spec == page->spec); + check_no_leftovers(dir); +} + +int main() { + test_admission_counts_what_the_store_charges(); + test_spill_held_pages_stay_counted(); + + const fs::path dir = fs::temp_directory_path() / ("test-server-pages-" + std::to_string(current_pid())); + fs::remove_all(dir); + fs::create_directories(dir); + test_a_failed_publication_keeps_the_committed_generation(dir); + test_an_interrupted_publication_restores_whole(dir); + test_spill_writes_a_page_restore_reads(dir); + fs::remove_all(dir); + + printf("test-server-pages: OK\n"); + return 0; +} diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 141d44126731..3687d269ff24 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -16,6 +16,14 @@ 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-genes.cpp + server-genes.h + server-checkpoints.cpp + server-checkpoints.h server-stream.cpp server-stream.h server-tools.cpp diff --git a/tools/server/server-checkpoints.cpp b/tools/server/server-checkpoints.cpp new file mode 100644 index 000000000000..efcce5eb409c --- /dev/null +++ b/tools/server/server-checkpoints.cpp @@ -0,0 +1,42 @@ +#include "server-checkpoints.h" + +#include + +size_t server_checkpoints_bytes(const server_checkpoint_list & list) { + size_t n = 0; + for (const auto & c : list) { + n += c.size(); + } + return n; +} + +bool server_checkpoints_fit( + const std::vector & lists, + size_t own, + uint64_t budget, + size_t incoming, + const std::function & on_evict) { + if (incoming > budget) { + return false; + } + for (;;) { + uint64_t total = incoming; + size_t victim = own; + size_t victim_bytes = 0; + for (size_t i = 0; i < lists.size(); ++i) { + const size_t n = server_checkpoints_bytes(*lists[i]); + total += n; + if (n > victim_bytes || (n == victim_bytes && i == own)) { + victim = i; + victim_bytes = n; + } + } + if (total <= budget || victim_bytes == 0) { + return true; + } + auto & list = *lists[victim]; + auto it = list.size() >= 2 ? std::next(list.begin()) : list.begin(); + on_evict(victim, *it); + list.erase(it); + } +} diff --git a/tools/server/server-checkpoints.h b/tools/server/server-checkpoints.h new file mode 100644 index 000000000000..1ab7750087a6 --- /dev/null +++ b/tools/server/server-checkpoints.h @@ -0,0 +1,28 @@ +#pragma once + +// Continuum fork: context checkpoints kept within one byte budget across all slots +// (--ctx-checkpoints-bytes), in place of a per-slot count. On a hybrid 27B a checkpoint is +// ~420 MB of host RAM, which on unified memory is the GPU's RAM, so the count is the wrong unit. + +#include "common.h" + +#include +#include +#include +#include + +using server_checkpoint_list = std::list; + +size_t server_checkpoints_bytes(const server_checkpoint_list & list); + +// Makes room for a new checkpoint of `incoming` bytes. While the lists' total plus `incoming` +// exceeds `budget`, the list holding the most bytes (the creator's on a tie) drops its oldest +// checkpoint after its first, so each slot keeps its head (the start of its first user turn) and +// its newest longest, as a saved page's subset does. `on_evict(list index, checkpoint)` runs before +// each drop. Returns false, dropping nothing, when the new checkpoint alone exceeds the budget. +bool server_checkpoints_fit( + const std::vector & lists, + size_t own, + uint64_t budget, + size_t incoming, + const std::function & on_evict); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index c55d5b79be80..0d58c9b3384c 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -482,6 +482,18 @@ 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 + // Continuum fork: genes (adapters) paged between the device and host RAM (server-genes.h) + uint64_t n_gene_page_out = 0; + uint64_t n_gene_page_in = 0; + uint64_t n_gene_page_in_demand = 0; // a task applied a host-resident adapter: paged in on its path + uint64_t n_gene_page_failed = 0; + void init() { t_start = ggml_time_us(); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 784d5426eeba..1c96206a2b43 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2,6 +2,10 @@ #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-genes.h" #include "server-chat.h" #include "server-common.h" #include "server-http.h" @@ -31,22 +35,6 @@ #include #include -// Continuum fork: an adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it -// cannot be read), stored in a saved slot's `.lora` sidecar so a restore refuses a page computed -// under different weights at the same path. -static std::pair lora_file_revision(const std::string & path) { - std::error_code ec; - const auto size = std::filesystem::file_size(path, ec); - if (ec) { - return {0, 0}; - } - const auto mtime = std::filesystem::last_write_time(path, ec); - if (ec) { - return {0, 0}; - } - return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; -} - // fix problem with std::min and std::max #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN @@ -938,6 +926,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 +951,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 +1349,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 { @@ -1694,6 +1696,21 @@ struct server_context_impl { { auto task_loras = task.params.lora.empty() ? params_base.lora_adapters : construct_lora_list(task.params.lora); + // A gene this task applies that the core paged out to host RAM comes back before the + // slot computes with it, never a decode on freed memory; each one is a counted miss + // (the core pages genes in ahead of the turns that need them). + try { + for (const auto & m : server_gene_page_in_applied(task_loras)) { + metrics.n_gene_page_in++; + metrics.n_gene_page_in_demand++; + SLT_WRN(slot, "LoRA adapter %d paged in on demand: %llu bytes in %.2f ms\n", + m.id, (unsigned long long) m.n_moved, m.t_ms); + } + } catch (const server_page_error & e) { + metrics.n_gene_page_failed++; + send_error(task, e.what(), e.type); + return false; + } if (!are_lora_equal(task_loras, slot.lora)) { // if lora has changed, check to see if the cache should be cleared if (lora_should_clear_cache(slot.lora, task_loras)) { @@ -2101,6 +2118,15 @@ struct server_context_impl { if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { return refused("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); } + if (llama_adapter_lora_is_paged_out(loras[L.lora_id].ptr)) { + // the gene it trains is in host RAM: back on the device before the optimizer binds it + if (llama_adapter_lora_page_in(loras[L.lora_id].ptr) < 0) { + metrics.n_gene_page_failed++; + return refused("adapter " + std::to_string(L.lora_id) + " is paged out and could not be paged back in"); + } + metrics.n_gene_page_in++; + metrics.n_gene_page_in_demand++; + } const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label if (n0 < 0 || L.n_prefix > slot.prompt.n_tokens() || L.reply.empty()) { return refused("slot " + std::to_string(slot.id) + " holds " + std::to_string(slot.prompt.n_tokens()) @@ -2472,7 +2498,34 @@ struct server_context_impl { ++it; } - while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { + // Continuum fork: a byte budget across all slots replaces the per-slot count when set + if (params_base.ctx_checkpoints_bytes > 0) { + const size_t incoming = + llama_state_seq_get_size_ext(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) + + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) : 0); + std::vector lists; + size_t own = 0; + for (auto & s : slots) { + if (&s == &slot) { + own = lists.size(); + } + lists.push_back(&s.prompt.checkpoints); + } + const bool fits = server_checkpoints_fit(lists, own, params_base.ctx_checkpoints_bytes, incoming, + [&](size_t i, const common_prompt_checkpoint & c) { + metrics.n_ckpt_evicted_budget++; + SLT_WRN(slots[i], "checkpoint budget: erasing context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB) for slot %d's new one\n", + c.pos_min, c.pos_max, c.n_tokens, (float) c.size() / 1024 / 1024, slot.id); + }); + if (!fits) { + metrics.n_ckpt_skipped_budget++; + SLT_WRN(slot, "checkpoint budget: a checkpoint of %.3f MiB alone exceeds --ctx-checkpoints-bytes %llu, not taken\n", + (float) incoming / 1024 / 1024, (unsigned long long) params_base.ctx_checkpoints_bytes); + return; + } + } + + while (params_base.ctx_checkpoints_bytes == 0 && slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { // make room for the new checkpoint, if needed const auto & cur = slot.prompt.checkpoints.front(); @@ -2502,6 +2555,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 +2798,12 @@ 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); + } + std::tie(res->gene_device_bytes, res->gene_host_bytes) = server_gene_bytes(params_base.lora_adapters); if (task.metrics_reset_bucket) { metrics.reset_bucket(); @@ -2729,39 +2888,26 @@ struct server_context_impl { ++i; } } - auto write_checkpoints = [&](auto && put) { - const uint32_t magic = 0x434B5054u; // 'CKPT' - const uint32_t version = 1; - const uint32_t n_ckpt = (uint32_t) kept_checkpoints.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); - for (const auto * kept : kept_checkpoints) { - const auto & c = *kept; - const int64_t n_tokens = c.n_tokens; - const int32_t pos_min = c.pos_min; - const int32_t pos_max = c.pos_max; - const uint64_t sz_tgt = c.data_tgt.size(); - const uint64_t sz_dft = c.data_dft.size(); - const uint64_t sz_spec = c.data_spec.size(); - put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); - put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); - put(c.data_tgt.data(), c.data_tgt.size()); - put(c.data_dft.data(), c.data_dft.size()); - put(c.data_spec.data(), c.data_spec.size()); - } - }; - auto write_adapters = [&](auto && put) { - const uint32_t magic = 0x4C4F5241u; // 'LORA' - const uint32_t version = 2; - const uint32_t n_lora = (uint32_t) slot->lora.size(); - put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); - for (const auto & l : slot->lora) { - const uint32_t n_path = (uint32_t) l.path.size(); - const auto rev = lora_file_revision(l.path); - put(&n_path, sizeof(n_path)); - put(l.path.data(), n_path); - put(&l.scale, sizeof(l.scale)); - put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); + // Continuum fork: the draft (MTP) state rides in a `.dft` sidecar, so a restored + // page drafts warm; a server without a draft context writes none. + // The quote counts the draft bytes without copying them. + std::vector draft_state; + std::vector draft_spec; + const size_t n_draft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (ctx_dft) { + common_speculative_get_state(spec.get(), slot->id, draft_spec); + } + if (ctx_dft && !task.slot_action.size_only) { + draft_state.resize(n_draft); + if (llama_state_seq_get_data_ext(ctx_dft, draft_state.data(), n_draft, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) != n_draft) { + send_error(task, "Unable to read the draft state", ERROR_TYPE_SERVER); + break; } + } + auto write_checkpoints = [&](const server_page_put & put) { server_page_write_checkpoints(kept_checkpoints, put); }; + auto write_adapters = [&](const server_page_put & put) { server_page_write_adapters(slot->lora, put); }; + auto write_draft = [&](const server_page_put & put) { + server_page_write_draft(draft_state.empty() ? nullptr : draft_state.data(), n_draft, draft_spec, put); }; if (task.slot_action.size_only || task.slot_action.has_byte_ceiling) { size_t projected_bytes = llama_state_seq_file_size( @@ -2781,6 +2927,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 +2951,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(); @@ -2868,6 +3015,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 +3134,23 @@ struct server_context_impl { SLT_WRN(*slot, "no usable checkpoint sidecar (%s): the restore is cold — the next divergence re-prefills from the head\n", ckpt_path.c_str()); } } + // Continuum fork: the draft (MTP) state from the `.dft` sidecar, so drafting + // resumes warm. A page saved before the sidecar existed drafts cold, as it did. + if (ctx_dft) { + std::vector draft_state; + std::vector draft_spec; + std::string why; + const std::string dft_path = filepath + ".dft"; + if (!server_page_read_draft(dft_path, draft_state, draft_spec, why)) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, -1, -1); + SLT_WRN(*slot, "draft sidecar %s (%s): the draft starts cold\n", why.c_str(), dft_path.c_str()); + } else if (llama_state_seq_set_data_ext(ctx_dft, draft_state.data(), draft_state.size(), slot->id, LLAMA_STATE_SEQ_FLAGS_NONE) != draft_state.size()) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, -1, -1); + SLT_WRN(*slot, "draft sidecar did not fit the draft context (%s): the draft starts cold\n", dft_path.c_str()); + } else { + common_speculative_set_state(spec.get(), slot->id, draft_spec); + } + } if (progress && progress->cancel_requested.load(std::memory_order_relaxed)) { throw std::runtime_error("Restore operation cancelled"); } @@ -3036,6 +3202,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 @@ -3062,6 +3254,8 @@ struct server_context_impl { lora, alora_invocation_string, alora_invocation_tokens, + !llama_adapter_lora_is_paged_out(lora.ptr), + llama_adapter_lora_n_bytes(lora.ptr), }); } queue_results.send(std::move(res)); @@ -3146,6 +3340,42 @@ struct server_context_impl { res->outcome = learn_on_slot(*slot, L, t0); queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_PAGE_LORA: + { + const int32_t id = task.page_lora; + auto applied_by_processing_slot = [&](int32_t lora_id) -> int { + for (const auto & slot : slots) { + if (slot.is_processing() && (size_t) lora_id < slot.lora.size() && + slot.lora[lora_id].ptr == params_base.lora_adapters[lora_id].ptr && + slot.lora[lora_id].scale != 0.0f) { + return slot.id; + } + } + return -1; + }; + try { + if (!task.page_lora_to_device) { + // the last batch may still be reading it, and the context may still + // hold it: settle both before its device buffers are freed. Only + // ctx_tgt ever receives adapters (common_set_adapter_lora in update_slots). + llama_synchronize(ctx_tgt); + llama_set_adapters_lora(ctx_tgt, nullptr, 0, nullptr); + } + auto res = std::make_unique(); + res->id = task.id; + res->move = server_gene_page(params_base.lora_adapters, id, task.page_lora_to_device, + learn_lora_id, applied_by_processing_slot); + if (res->move.n_moved > 0) { + (task.page_lora_to_device ? metrics.n_gene_page_in : metrics.n_gene_page_out)++; + } + SRV_INF("lora adapter id=%d %s: %llu bytes in %.2f ms\n", id, res->move.on_device ? "on device" : "in host RAM", + (unsigned long long) res->move.n_moved, res->move.t_ms); + queue_results.send(std::move(res)); + } catch (const server_page_error & e) { + metrics.n_gene_page_failed++; + send_error(task, e.what(), e.type); + } + } break; case SERVER_TASK_TYPE_UNLOAD_LORA: { auto & loras = params_base.lora_adapters; @@ -3195,6 +3425,7 @@ struct server_context_impl { if (prompt_cache) { prompt_cache->evict_lora(ptr); } + page_store.evict_lora(ptr); // the context may still hold it from the last batch; clear before freeing. // Only ctx_tgt ever receives adapters (common_set_adapter_lora in // update_slots); the draft/MTP context never does. @@ -5213,6 +5444,8 @@ 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 + { "lora_page_version", 1 }, // Continuum fork: /lora-adapters/page_out, page_in (server-genes.h) { "expert_container_custody_version", 1 }, { "model_alias", meta.model_name }, { "model_ftype", meta.model_ftype }, @@ -5388,7 +5621,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 +5640,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"); } @@ -6091,6 +6329,38 @@ void server_routes::init_routes() { return res; }; + // POST /lora-adapters/page_out and /page_in {"id": N}: a loaded gene's tensors to host RAM or back + // to the device (server-genes.h). Idempotent: a gene already there answers n_moved 0. + auto page_lora = [this](const server_http_req & req, bool to_device) { + auto res = create_response(); + const json body = json::parse(req.body); + if (!body.contains("id") || !body.at("id").is_number_integer()) { + res->error(format_error_response("\"id\" (an integer adapter id) is required", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + auto & rd = res->rd; + { + server_task task(SERVER_TASK_TYPE_PAGE_LORA); + task.id = rd.get_new_id(); + task.page_lora = body.at("id").get(); + task.page_lora_to_device = to_device; + rd.post_task(std::move(task)); + } + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + res->ok(result->to_json()); + return res; + }; + this->post_lora_adapters_page_out = [page_lora](const server_http_req & req) { return page_lora(req, false); }; + this->post_lora_adapters_page_in = [page_lora](const server_http_req & req) { return page_lora(req, true); }; + auto train_control = [this](const server_http_req & req, bool pause) { auto res = create_response(); const json body = json::parse(req.body); @@ -6293,6 +6563,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..f6d43bcff611 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -169,6 +169,8 @@ struct server_routes { server_http_context::handler_t post_lora_adapters; server_http_context::handler_t post_lora_adapters_load; server_http_context::handler_t post_lora_adapters_unload; + server_http_context::handler_t post_lora_adapters_page_out; // Continuum fork: server-genes.h + server_http_context::handler_t post_lora_adapters_page_in; server_http_context::handler_t post_learn; server_http_context::handler_t get_train; server_http_context::handler_t post_train; @@ -190,6 +192,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-genes.cpp b/tools/server/server-genes.cpp new file mode 100644 index 000000000000..5d36f0f9f85d --- /dev/null +++ b/tools/server/server-genes.cpp @@ -0,0 +1,78 @@ +#include "server-genes.h" + +#include "ggml.h" + +json server_gene_move::to_json() const { + return json { + { "success", true }, + { "id", id }, + { "path", path }, + { "resident", on_device ? "device" : "host" }, + { "n_bytes", n_bytes }, + { "n_moved", n_moved }, + { "t_ms", t_ms }, + }; +} + +static server_gene_move gene_move(const common_adapter_lora_info & lora, int32_t id, bool to_device) { + const int64_t t0 = ggml_time_us(); + const int64_t moved = to_device ? llama_adapter_lora_page_in(lora.ptr) : llama_adapter_lora_page_out(lora.ptr); + if (moved < 0) { + throw server_page_error(ERROR_TYPE_SERVER, "LoRA adapter " + std::to_string(id) + " could not be paged " + + (to_device ? "in: the device refused its buffers" : "out: no host memory for its tensors")); + } + server_gene_move m; + m.id = id; + m.path = lora.path; + m.on_device = !llama_adapter_lora_is_paged_out(lora.ptr); + m.n_bytes = llama_adapter_lora_n_bytes(lora.ptr); + m.n_moved = (uint64_t) moved; + m.t_ms = (ggml_time_us() - t0) / 1000.0; + return m; +} + +server_gene_move server_gene_page(const std::vector & loras, int32_t id, bool to_device, + int32_t learn_lora_id, const std::function & applied_by_processing_slot) { + if (id < 0 || (size_t) id >= loras.size() || loras[id].ptr == nullptr) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "no loaded LoRA adapter with id " + std::to_string(id)); + } + if (!to_device) { + if (id == learn_lora_id) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "LoRA adapter " + std::to_string(id) + + " is bound to the learning optimizer; it stays on the device"); + } + if (loras[id].scale != 0.0f) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "LoRA adapter " + std::to_string(id) + + " is active by default (scale " + std::to_string(loras[id].scale) + "); set its scale to 0 first"); + } + const int slot = applied_by_processing_slot(id); + if (slot >= 0) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, "LoRA adapter " + std::to_string(id) + " is in use by slot " + + std::to_string(slot) + "; retry when it finishes"); + } + } + return gene_move(loras[id], id, to_device); +} + +std::vector server_gene_page_in_applied(const std::vector & applied) { + std::vector moved; + for (size_t i = 0; i < applied.size(); ++i) { + const auto & lora = applied[i]; + if (lora.ptr != nullptr && lora.scale != 0.0f && llama_adapter_lora_is_paged_out(lora.ptr)) { + moved.push_back(gene_move(lora, (int32_t) i, true)); + } + } + return moved; +} + +std::pair server_gene_bytes(const std::vector & loras) { + uint64_t device = 0; + uint64_t host = 0; + for (const auto & lora : loras) { + if (lora.ptr == nullptr) { + continue; + } + (llama_adapter_lora_is_paged_out(lora.ptr) ? host : device) += llama_adapter_lora_n_bytes(lora.ptr); + } + return { device, host }; +} diff --git a/tools/server/server-genes.h b/tools/server/server-genes.h new file mode 100644 index 000000000000..5d615f9b1e26 --- /dev/null +++ b/tools/server/server-genes.h @@ -0,0 +1,57 @@ +#pragma once + +// Continuum fork: genes (LoRA adapters) paged between device memory and host RAM. +// +// A mind runs only the genes her turn is near (continuum genome::selection); the rest of the +// genes the node has loaded need not hold device memory. The core's working-set owner decides, +// by temperature and within one memory budget, which genes stay on the device; these verbs are +// the mechanism: `page_out` copies an adapter's tensors to host RAM and frees its device buffers, +// `page_in` re-allocates them and copies back, `unload` (POST /lora-adapters/unload) drops it. +// An adapter keeps its id and identity throughout, so a slot's adapter list, its KV and a saved +// page computed under it stay valid. A task that applies a host-resident adapter pages it in +// first and counts the miss, so the engine never computes on freed memory. +// Design: continuum docs/architecture/SERVING-ENGINE.md (the `genes` module, §3 one budget). + +#include "server-pages.h" + +#include "common.h" + +#include +#include +#include +#include +#include + +// Where an adapter's tensors are, after a move. +struct server_gene_move { + int32_t id = -1; + std::string path; + bool on_device = true; + uint64_t n_bytes = 0; // the adapter's bytes where it is now + uint64_t n_moved = 0; // bytes this call copied; 0 when it was already there + double t_ms = 0.0; + + json to_json() const; +}; + +// The result of SERVER_TASK_TYPE_PAGE_LORA. +struct server_task_result_gene_move : server_task_result { + server_gene_move move; + virtual json to_json() override { return move.to_json(); } +}; + +// Moves adapter `id` of `loras` to the device (`to_device`) or to host RAM. Refuses with +// server_page_error: an unknown or retired id (INVALID_REQUEST); to host, the adapter the learning +// optimizer holds or one active by default (INVALID_REQUEST), or one a processing slot applies +// (UNAVAILABLE: retry when it finishes; `applied_by_processing_slot` returns that slot's id or -1); +// a failed allocation (SERVER), the adapter then left where it was. The caller synchronizes the +// contexts that ran the adapter before paging it out. +server_gene_move server_gene_page(const std::vector & loras, int32_t id, bool to_device, + int32_t learn_lora_id, const std::function & applied_by_processing_slot); + +// Pages in every adapter `applied` runs (scale != 0) that is in host RAM: a task is about to run +// it. Returns what moved; empty on the fast path. Throws server_page_error on a failed allocation. +std::vector server_gene_page_in_applied(const std::vector & applied); + +// (device bytes, host bytes) of every loaded adapter. +std::pair server_gene_bytes(const std::vector & loras); diff --git a/tools/server/server-page-files.cpp b/tools/server/server-page-files.cpp new file mode 100644 index 000000000000..5a136639ae6d --- /dev/null +++ b/tools/server/server-page-files.cpp @@ -0,0 +1,444 @@ +#include "server-page-files.h" + +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# ifndef NOMINMAX +# define NOMINMAX +# endif +# include +# include +# include +# include +#else +# include +# include +# include +# include +#endif + +namespace fs = std::filesystem; + +static constexpr uint32_t PAGE_CKPT_MAGIC = 0x434B5054u; // 'CKPT' +static constexpr uint32_t PAGE_LORA_MAGIC = 0x4C4F5241u; // 'LORA' +static constexpr uint32_t PAGE_DRFT_MAGIC = 0x44524654u; // 'DRFT' + +// +// page files +// + +std::pair lora_file_revision(const std::string & path) { + std::error_code ec; + const auto size = std::filesystem::file_size(path, ec); + if (ec) { + return {0, 0}; + } + const auto mtime = std::filesystem::last_write_time(path, ec); + if (ec) { + return {0, 0}; + } + return {(uint64_t) size, (int64_t) std::chrono::duration_cast(mtime.time_since_epoch()).count()}; +} + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put) { + const uint32_t magic = PAGE_CKPT_MAGIC; + const uint32_t version = 1; + const uint32_t n_ckpt = (uint32_t) kept.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_ckpt, sizeof(n_ckpt)); + for (const auto * c : kept) { + const int64_t n_tokens = c->n_tokens; + const int32_t pos_min = c->pos_min; + const int32_t pos_max = c->pos_max; + const uint64_t sz_tgt = c->data_tgt.size(); + const uint64_t sz_dft = c->data_dft.size(); + const uint64_t sz_spec = c->data_spec.size(); + put(&n_tokens, sizeof(n_tokens)); put(&pos_min, sizeof(pos_min)); put(&pos_max, sizeof(pos_max)); + put(&sz_tgt, sizeof(sz_tgt)); put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(c->data_tgt.data(), c->data_tgt.size()); + put(c->data_dft.data(), c->data_dft.size()); + put(c->data_spec.data(), c->data_spec.size()); + } +} + +void server_page_write_adapters(const std::vector & lora, const server_page_put & put) { + const uint32_t magic = PAGE_LORA_MAGIC; + const uint32_t version = 2; + const uint32_t n_lora = (uint32_t) lora.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); put(&n_lora, sizeof(n_lora)); + for (const auto & l : lora) { + const uint32_t n_path = (uint32_t) l.path.size(); + const auto rev = lora_file_revision(l.path); + put(&n_path, sizeof(n_path)); + put(l.path.data(), n_path); + put(&l.scale, sizeof(l.scale)); + put(&rev.first, sizeof(rev.first)); put(&rev.second, sizeof(rev.second)); + } +} + +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put) { + const uint32_t magic = PAGE_DRFT_MAGIC; + const uint32_t version = 1; + const uint64_t sz_dft = n_dft; + const uint64_t sz_spec = spec.size(); + put(&magic, sizeof(magic)); put(&version, sizeof(version)); + put(&sz_dft, sizeof(sz_dft)); put(&sz_spec, sizeof(sz_spec)); + put(dft, n_dft); + put(spec.data(), spec.size()); +} + +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why) { + std::ifstream in(path, std::ios::binary | std::ios::ate); + if (!in) { + why = "absent"; + return false; + } + const auto total = (uint64_t) in.tellg(); + in.seekg(0); + uint32_t magic = 0, version = 0; + uint64_t sz_dft = 0, sz_spec = 0; + in.read(reinterpret_cast(&magic), sizeof(magic)); + in.read(reinterpret_cast(&version), sizeof(version)); + in.read(reinterpret_cast(&sz_dft), sizeof(sz_dft)); + in.read(reinterpret_cast(&sz_spec), sizeof(sz_spec)); + const uint64_t header = 2 * sizeof(uint32_t) + 2 * sizeof(uint64_t); + if (!in || magic != PAGE_DRFT_MAGIC || version != 1 || sz_dft > total || sz_spec > total || header + sz_dft + sz_spec != total) { + why = "malformed"; + return false; + } + dft.resize(sz_dft); + spec.resize(sz_spec); + in.read(reinterpret_cast(dft.data()), sz_dft); + in.read(reinterpret_cast(spec.data()), sz_spec); + if (!in) { + why = "truncated"; + return false; + } + return true; +} + +// +// generation publication +// + +// Every file a page may have. A generation names the ones it carries; the others are removed at its commit. +static const char * const PAGE_SUFFIXES[] = { "", ".ckpt", ".lora", ".dft" }; + +// Beside the page dir; ".." in it is a name fs_validate_filename refuses, so no page can take it. +static const char * const STAGING_DIR = "..page-staging"; + +// Publications in this process: the commit boundary every reader shares, and the stages being written. +struct page_publications { + std::shared_mutex commit; // a commit (and a reader's repair) exclusive; a reader's read shared + std::mutex stages_mutex; + std::set stages; + uint64_t next = 1; +}; + +static page_publications & publications() { + static page_publications pubs; + return pubs; +} + +static long current_pid() { +#if defined(_WIN32) + return (long) _getpid(); +#else + return (long) getpid(); +#endif +} + +static bool process_alive(long pid) { +#if defined(_WIN32) + HANDLE h = OpenProcess(SYNCHRONIZE, FALSE, (DWORD) pid); + if (h == nullptr) { + return false; + } + const bool alive = WaitForSingleObject(h, 0) == WAIT_TIMEOUT; + CloseHandle(h); + return alive; +#else + return kill((pid_t) pid, 0) == 0 || errno == EPERM; +#endif +} + +// A stage is "-": live while its publication runs in this process, or while its process runs. +static bool stage_live(const std::string & id) { + const auto dash = id.find('-'); + long pid = -1; + try { + pid = std::stol(id.substr(0, dash)); + } catch (const std::exception &) { + return false; + } + if (pid <= 0) { + return false; + } + if (pid != current_pid()) { + return process_alive(pid); + } + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + return pubs.stages.count(id) > 0; +} + +// Makes the file's contents durable before anything that depends on them (the commit record). +static bool sync_file(const fs::path & path) { +#if defined(_WIN32) + const int fd = _wopen(path.c_str(), _O_RDWR | _O_BINARY); + if (fd < 0) { + return false; + } + const bool ok = _commit(fd) == 0; + _close(fd); + return ok; +#else + const int fd = open(path.c_str(), O_RDONLY); + if (fd < 0) { + return false; + } +# if defined(__APPLE__) + // macOS fsync stops at the drive's cache; only F_FULLFSYNC orders the data before the record + bool ok = fcntl(fd, F_FULLFSYNC) == 0 || fsync(fd) == 0; +# else + bool ok = fsync(fd) == 0; +# endif + ok = close(fd) == 0 && ok; + return ok; +#endif +} + +// Makes a dir's entries (a rename into it) durable. Windows has no handle for this; NTFS journals them. +static void sync_dir(const fs::path & dir) { +#if !defined(_WIN32) + const int fd = open(dir.c_str(), O_RDONLY); + if (fd >= 0) { + fsync(fd); + close(fd); + } +#else + (void) dir; +#endif +} + +static fs::path staging_dir(const fs::path & file) { + return file.parent_path() / STAGING_DIR; +} + +// Its presence is the commit: the stage it names holds the generation, whole and synced. +static fs::path commit_record(const fs::path & file) { + return staging_dir(file) / (file.filename().string() + ".commit"); +} + +// "\n" then "f\n" per file of the generation. +static bool read_commit_record(const fs::path & record, std::string & id, std::vector & suffixes) { + std::ifstream in(record); + if (!in || !std::getline(in, id) || id.empty()) { + return false; + } + for (std::string line; std::getline(in, line);) { + if (line.empty() || line[0] != 'f') { + return false; + } + suffixes.push_back(line.substr(1)); + } + return !suffixes.empty(); +} + +// Writes the record inside the stage, syncs it, then renames it into place: the rename is the commit. +static void write_commit_record(const fs::path & file, const fs::path & stage, const std::string & id, const std::vector & suffixes) { + const fs::path tmp = stage / "..commit"; // ".." is in no page name + { + std::ofstream out(tmp, std::ios::trunc); + out << id << '\n'; + for (const auto & s : suffixes) { + out << 'f' << s << '\n'; + } + out.flush(); + if (!out) { + throw server_page_error(ERROR_TYPE_SERVER, "the page's commit record could not be written: " + tmp.string()); + } + } + std::error_code ec; + if (!sync_file(tmp) || (fs::rename(tmp, commit_record(file), ec), ec)) { + throw server_page_error(ERROR_TYPE_SERVER, "the page's commit record could not be placed: " + commit_record(file).string()); + } + sync_dir(staging_dir(file)); +} + +// Moves a committed generation over the page's files, removes the files it lacks, then its stage and +// record. Idempotent: a file a crash already moved is not in the stage. Throws, keeping the record, so +// the next reader finishes it. +static void roll_forward(const fs::path & file, const fs::path & stage, const std::vector & suffixes) { + const std::string name = file.filename().string(); + std::error_code ec; + for (const auto & s : suffixes) { + const fs::path staged = stage / (name + s); + if (fs::exists(staged, ec) && (fs::rename(staged, file.string() + s, ec), ec)) { + throw server_page_error(ERROR_TYPE_SERVER, "a committed page file could not be moved into place: " + staged.string() + ": " + ec.message()); + } + } + for (const char * s : PAGE_SUFFIXES) { + if (std::find(suffixes.begin(), suffixes.end(), s) == suffixes.end() && !fs::remove(file.string() + s, ec) && ec) { + throw server_page_error(ERROR_TYPE_SERVER, "an older page file could not be removed: " + file.string() + s + ": " + ec.message()); + } + } + sync_dir(file.parent_path()); + fs::remove_all(stage, ec); + fs::remove(commit_record(file), ec); + sync_dir(staging_dir(file)); +} + +// Removes every stage no commit record names and no live publication owns: what a crash before a +// commit left. Called under the exclusive commit lock. +static void remove_orphan_stages(const fs::path & staging) { + std::error_code ec; + if (!fs::is_directory(staging, ec)) { + return; + } + std::set committed; + for (const auto & e : fs::directory_iterator(staging, ec)) { + std::string id; + std::vector suffixes; + if (e.path().extension() == ".commit" && read_commit_record(e.path(), id, suffixes)) { + committed.insert(id); + } + } + std::vector orphans; + for (const auto & e : fs::directory_iterator(staging, ec)) { + const std::string id = e.path().filename().string(); + if (e.is_directory(ec) && committed.count(id) == 0 && !stage_live(id)) { + orphans.push_back(e.path()); + } + } + for (const auto & p : orphans) { + SRV_WRN("removing a page generation a publication left uncommitted: %s\n", p.string().c_str()); + fs::remove_all(p, ec); + } +} + +// A publication's stage, live while it exists: uncommitted, it is removed with its owner; committed, it +// stays until its roll forward removes it. +struct stage_owner { + std::string id; + fs::path stage; + bool committed = false; + + explicit stage_owner(const fs::path & staging) { + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + id = std::to_string(current_pid()) + "-" + std::to_string(pubs.next++); + stage = staging / id; + pubs.stages.insert(id); + } + + ~stage_owner() { + std::error_code ec; + if (!committed) { + fs::remove_all(stage, ec); + } + auto & pubs = publications(); + std::lock_guard lock(pubs.stages_mutex); + pubs.stages.erase(id); + } + + stage_owner(const stage_owner &) = delete; + stage_owner & operator=(const stage_owner &) = delete; +}; + +server_page_file server_page_sidecar(std::string suffix, std::function body) { + return { std::move(suffix), [body = std::move(body)](const std::string & path) -> size_t { + std::ofstream out(path, std::ios::binary | std::ios::trunc); + size_t n = 0; + body([&](const void * ptr, size_t len) { + out.write(reinterpret_cast(ptr), len); + n += len; + }); + out.flush(); + return out ? n : 0; + } }; +} + +size_t server_page_publish(const std::string & filepath, const std::vector & files) { + const fs::path file(filepath); + auto & pubs = publications(); + + stage_owner owner(staging_dir(file)); + const fs::path & stage = owner.stage; + + std::error_code ec; + fs::create_directories(stage, ec); + if (ec) { + throw server_page_error(ERROR_TYPE_SERVER, "the page could not be staged: " + stage.string() + ": " + ec.message()); + } + size_t n_written = 0; + std::vector suffixes; + for (const auto & f : files) { + GGML_ASSERT(std::find(std::begin(PAGE_SUFFIXES), std::end(PAGE_SUFFIXES), f.suffix) != std::end(PAGE_SUFFIXES)); + const fs::path staged = stage / (file.filename().string() + f.suffix); + size_t n = 0; + try { + n = f.write(staged.string()); + } catch (const std::exception & err) { + throw server_page_error(ERROR_TYPE_SERVER, "a page file could not be written: " + staged.string() + ": " + err.what()); + } + if (n == 0 || !sync_file(staged)) { + throw server_page_error(ERROR_TYPE_SERVER, "a page file could not be written: " + staged.string()); + } + n_written += n; + suffixes.push_back(f.suffix); + } + sync_dir(stage); + + std::unique_lock lock(pubs.commit); + write_commit_record(file, stage, owner.id, suffixes); + owner.committed = true; + roll_forward(file, stage, suffixes); + return n_written; +} + +size_t server_page_spill(const server_page & page, const std::string & filepath) { + std::vector packed; + try { + packed = page.prompt.tokens.serialize(); + } catch (const std::exception & err) { + throw server_page_error(ERROR_TYPE_NOT_SUPPORTED, err.what()); + } + std::vector all; + for (const auto & c : page.prompt.checkpoints) { + all.push_back(&c); + } + return server_page_publish(filepath, { + { "", [&](const std::string & path) { + return llama_state_seq_save_file_from_data(path.c_str(), page.data.main.data(), page.data.main.size(), + reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); + } }, + server_page_sidecar(".ckpt", [&](const server_page_put & put) { server_page_write_checkpoints(all, put); }), + server_page_sidecar(".lora", [&](const server_page_put & put) { server_page_write_adapters(page.lora, put); }), + server_page_sidecar(".dft", [&](const server_page_put & put) { server_page_write_draft(page.data.drft.data(), page.data.drft.size(), page.spec, put); }), + }); +} + +server_page_reader::server_page_reader(const std::string & filepath) { + const fs::path file(filepath); + auto & pubs = publications(); + { + std::unique_lock repair(pubs.commit); + const fs::path record = commit_record(file); + std::string id; + std::vector suffixes; + if (read_commit_record(record, id, suffixes)) { + SRV_WRN("completing the publication of %s that stopped after its commit\n", filepath.c_str()); + roll_forward(file, staging_dir(file) / id, suffixes); + } + remove_orphan_stages(staging_dir(file)); + } + lock = std::shared_lock(pubs.commit); +} diff --git a/tools/server/server-page-files.h b/tools/server/server-page-files.h new file mode 100644 index 000000000000..ea12dfd9ce1c --- /dev/null +++ b/tools/server/server-page-files.h @@ -0,0 +1,67 @@ +#pragma once + +// Continuum fork: a page on disk, `` (sequence state) and its sidecars `.ckpt` (context +// checkpoints), `.lora` (the adapters the state was computed under) and `.dft` (draft state), +// shared by save, spill and restore. +// +// The files of a page are one generation, published whole: a writer stages every file beside the +// page dir and syncs them, writes a commit record, then moves them over the committed files. A +// failure before the record leaves the committed generation as it was; a crash after it is +// completed by the next reader. Restore reads only through a server_page_reader, so it sees one +// committed generation, never a mix. Two servers publishing the SAME name into one dir at once are +// not ordered against each other: a page is seated in one server at a time. + +#include "server-pages.h" + +#include +#include +#include +#include +#include +#include + +using server_page_put = std::function; + +// An adapter file's revision (size, mtime in ns since the epoch; (0, 0) when it cannot be read), +// stored in the `.lora` sidecar so a restore refuses a page computed under different weights. +std::pair lora_file_revision(const std::string & path); + +void server_page_write_checkpoints(const std::vector & kept, const server_page_put & put); // .ckpt +void server_page_write_adapters(const std::vector & lora, const server_page_put & put); // .lora +// `dft` may be null when `put` only counts bytes +void server_page_write_draft(const uint8_t * dft, size_t n_dft, const std::vector & spec, const server_page_put & put); // .dft + +// Reads a `.dft` sidecar. Returns false with `why` set when it is absent or malformed. +bool server_page_read_draft(const std::string & path, std::vector & dft, std::vector & spec, std::string & why); + +// One file of a generation: its suffix ("" is the sequence state; one of "", ".ckpt", ".lora", +// ".dft") and the writer that fills the staged `path`, returning the bytes written (0 = failed). +struct server_page_file { + std::string suffix; + std::function write; +}; + +// A sidecar written by one of the serializers above. +server_page_file server_page_sidecar(std::string suffix, std::function body); + +// Publishes `files` as the generation of the page at `filepath`, whole or not at all; a page file +// the generation lacks is removed at the commit. Throws server_page_error; on any failure before +// the commit record the committed generation is untouched and nothing staged is left. Returns the +// bytes written. +size_t server_page_publish(const std::string & filepath, const std::vector & files); + +// Writes `page` to `filepath` (+ .ckpt, .lora, .dft) with every checkpoint, as one generation. +// Touches no context, so it runs off the task thread. Returns the bytes written. +size_t server_page_spill(const server_page & page, const std::string & filepath); + +// Holds the committed generation at `filepath` still while it is read. Constructing one completes a +// publication a crash cut short after its commit record and removes staged files no live publication +// owns (a crash before the commit), then keeps publication off until it is destroyed. Throws +// server_page_error when an interrupted commit cannot be completed. +class server_page_reader { +public: + explicit server_page_reader(const std::string & filepath); + +private: + std::shared_lock lock; +}; diff --git a/tools/server/server-pages.cpp b/tools/server/server-pages.cpp 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..09b741a9f75d --- /dev/null +++ b/tools/server/server-pages.h @@ -0,0 +1,126 @@ +#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 { + 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 + + server_page_charge charge; + + 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..f65020f81abb 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1591,6 +1591,42 @@ 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 + }, { + "lora_page_out_total", + "LoRA adapters paged out to host RAM", + (double) metrics.n_gene_page_out + }, { + "lora_page_in_total", + "LoRA adapters paged back onto the device (on demand included)", + (double) metrics.n_gene_page_in + }, { + "lora_page_in_on_demand_total", + "LoRA adapters a task applied while in host RAM, paged in on its path", + (double) metrics.n_gene_page_in_demand + }, { + "lora_page_failed_total", + "LoRA page_out or page_in refused or failed", + (double) metrics.n_gene_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 +1651,26 @@ 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 + }, { + "lora_device_bytes", + "Bytes of loaded LoRA adapters on the device", + (double) gene_device_bytes + }, { + "lora_host_bytes", + "Bytes of loaded LoRA adapters paged out to host RAM", + (double) gene_host_bytes }, }; @@ -1675,6 +1731,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 // @@ -1699,6 +1783,8 @@ json server_task_result_get_lora::to_json() { {"scale", lora.info.scale}, {"task_name", lora.info.task_name}, {"prompt_prefix", lora.info.prompt_prefix}, + {"resident", lora.on_device ? "device" : "host"}, + {"n_bytes", lora.n_bytes}, }; if (!lora.alora_invocation_tokens.empty()) { entry["alora_invocation_string"] = lora.alora_invocation_string; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 87a344e174cf..2b795976e777 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -56,10 +56,15 @@ 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, SERVER_TASK_TYPE_UNLOAD_LORA, + SERVER_TASK_TYPE_PAGE_LORA, // Continuum fork: an adapter's tensors to host RAM or back (server-genes.h) SERVER_TASK_TYPE_LEARN, }; @@ -237,6 +242,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; @@ -254,6 +263,10 @@ struct server_task { // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire int32_t unload_lora = -1; + // used by SERVER_TASK_TYPE_PAGE_LORA: the adapter id, and where it goes + int32_t page_lora = -1; + bool page_lora_to_device = true; + // used by SERVER_TASK_TYPE_LEARN learn_request learn; @@ -578,6 +591,11 @@ 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 + uint64_t gene_device_bytes = 0; // Continuum fork: loaded adapters on the device + uint64_t gene_host_bytes = 0; // and paged out to host RAM server_metrics metrics; @@ -614,6 +632,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; @@ -639,6 +684,8 @@ struct server_task_result_get_lora : server_task_result { common_adapter_lora_info info; std::string alora_invocation_string; llama_tokens alora_invocation_tokens; + bool on_device = true; // Continuum fork: false while paged out to host RAM + uint64_t n_bytes = 0; // its bytes where it is }; std::vector loras; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index f449dc00f11e..95edbacc7547 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -219,6 +219,8 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters = models_routes->proxy_post; routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; + routes.post_lora_adapters_page_out = models_routes->proxy_post; + routes.post_lora_adapters_page_in = models_routes->proxy_post; routes.post_learn = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -274,6 +276,8 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters", ex_wrapper(routes.post_lora_adapters)); ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); + ctx_http.post("/lora-adapters/page_out", ex_wrapper(routes.post_lora_adapters_page_out)); + ctx_http.post("/lora-adapters/page_in", ex_wrapper(routes.post_lora_adapters_page_in)); ctx_http.post("/learn", ex_wrapper(routes.post_learn)); ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); diff --git a/tools/server/tests/unit/test_lora_page.py b/tools/server/tests/unit/test_lora_page.py new file mode 100644 index 000000000000..4a3adb36faad --- /dev/null +++ b/tools/server/tests/unit/test_lora_page.py @@ -0,0 +1,146 @@ +import shutil +from pathlib import Path + +import pytest +from utils import * + +# Continuum fork: a loaded gene (LoRA adapter) paged between the device and host RAM, and dropped +# (server-genes.h; continuum docs/architecture/SERVING-ENGINE.md, the `genes` module). + +LORA_FILE_URL = "https://huggingface.co/ggml-org/stories15M_MOE/resolve/main/moe_shakespeare15M.gguf" +TRAIN_DIR = Path("./tmp/lora-page-train") +GENE = "shakespeare.gguf" +PROMPT = "Look in thy glass" + +server = ServerPreset.stories15m_moe() + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.stories15m_moe() + server.server_metrics = True + TRAIN_DIR.mkdir(parents=True, exist_ok=True) + shutil.copyfile(download_file(LORA_FILE_URL), TRAIN_DIR / GENE) + server.train_dir = str(TRAIN_DIR) + + +def load_gene(): + res = server.make_request("POST", "/lora-adapters/load", data={"name": GENE}) + assert res.status_code == 200, res.body + return res.body["id"] + + +def page(action, gene_id): + return server.make_request("POST", f"/lora-adapters/{action}", data={"id": gene_id}) + + +def complete(gene_id, scale): + res = server.make_request("POST", "/completion", data={ + "prompt": PROMPT, "lora": [{"id": gene_id, "scale": scale}], + "seed": 42, "temperature": 0.0, "n_predict": 24, "cache_prompt": False, + }) + assert res.status_code == 200, res.body + return res.body["content"] + + +def listed(gene_id): + res = server.make_request("GET", "/lora-adapters") + assert res.status_code == 200 + return next(entry for entry in res.body if entry["id"] == gene_id) + + +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 re-binds the tensors to the wrong bytes or the wrong buffer + # type (the gene's output changes), a page_out that does not actually free its device bytes, or + # sizes and timings that are not reported where the core's working-set owner reads them. + global server + server.start() + assert server.make_request("GET", "/props").body["lora_page_version"] == 1 + gene = load_gene() + reference = complete(gene, 1.0) + assert reference != complete(gene, 0.0), "the gene must change the output for this test to see it" + on_device = listed(gene) + assert on_device["resident"] == "device" and on_device["n_bytes"] > 0 + + out = page("page_out", gene) + assert out.status_code == 200, out.body + assert out.body["resident"] == "host" and out.body["n_moved"] == out.body["n_bytes"] > 0 + assert out.body["t_ms"] >= 0 + assert listed(gene)["resident"] == "host" + assert metric("lora_host_bytes") == pytest.approx(out.body["n_bytes"], rel=1e-5) # /metrics prints 6 digits + assert metric("lora_device_bytes") == 0 + assert metric("lora_page_out_total") == 1 + + back = page("page_in", gene) + assert back.status_code == 200, back.body + assert back.body["resident"] == "device" and back.body["n_moved"] == out.body["n_bytes"] + assert metric("lora_host_bytes") == 0 + assert complete(gene, 1.0) == reference + assert metric("lora_page_in_on_demand_total") == 0, "an explicit page_in is not a miss" + + +def test_a_task_applying_a_paged_out_gene_pages_it_in_first(): + # what this catches: a decode computing with an adapter whose device buffers were freed (a + # crash, or garbage) when the core has not paged it back in ahead of the turn. The task pages it + # in on its own path, produces the same output, and the miss is counted for the core to see. + global server + server.start() + gene = load_gene() + reference = complete(gene, 1.0) + assert page("page_out", gene).status_code == 200 + assert complete(gene, 1.0) == reference + assert listed(gene)["resident"] == "device" + assert metric("lora_page_in_on_demand_total") == 1 + # a task that does not apply it leaves it where it is + assert page("page_out", gene).status_code == 200 + complete(gene, 0.0) + assert listed(gene)["resident"] == "host" + assert metric("lora_page_in_on_demand_total") == 1 + + +def test_page_refusals_are_typed_and_repeats_are_idempotent(): + # what this catches: paging out an adapter every default request runs (it would page back in + # on every request and thrash), silently accepting an unknown id, and a retried page_out + # copying twice. + global server + server.start() + gene = load_gene() + assert page("page_out", 99).status_code == 400 + assert page("page_in", 99).status_code == 400 + + assert server.make_request("POST", "/lora-adapters", data=[{"id": gene, "scale": 1.0}]).status_code == 200 + refused = page("page_out", gene) + assert refused.status_code == 400 and "active by default" in refused.body["error"]["message"] + assert server.make_request("POST", "/lora-adapters", data=[{"id": gene, "scale": 0.0}]).status_code == 200 + + first = page("page_out", gene) + again = page("page_out", gene) + assert first.status_code == again.status_code == 200 + assert again.body["n_moved"] == 0 and again.body["n_bytes"] == first.body["n_bytes"] + assert metric("lora_page_out_total") == 1 + assert page("page_in", gene).body["n_moved"] > 0 + assert page("page_in", gene).body["n_moved"] == 0 + + +def test_unload_drops_a_paged_out_gene(): + # what this catches: a gene dropped while in host RAM leaking its host copy (the host gauge + # stays up) or staying addressable. + global server + server.start() + gene = load_gene() + assert page("page_out", gene).status_code == 200 + dropped = server.make_request("POST", "/lora-adapters/unload", data={"id": gene}) + assert dropped.status_code == 200 and dropped.body["retired"] is True + assert all(entry["id"] != gene for entry in server.make_request("GET", "/lora-adapters").body) + assert metric("lora_host_bytes") == 0 + assert page("page_in", gene).status_code == 400 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..d36f288fe25e 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -121,6 +121,9 @@ 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 + train_dir: str | None = None # session variables process: subprocess.Popen | None = None @@ -285,6 +288,12 @@ 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.train_dir is not None: + server_args.extend(["--train-dir", self.train_dir]) if self.gcp_compat: env["AIP_MODE"] = "PREDICTION"