diff --git a/common/arg.cpp b/common/arg.cpp index 7ce7eaad16ab..77fdd9ca4644 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -313,7 +313,7 @@ const std::vector kv_cache_types = { GGML_TYPE_Q5_1, }; -static ggml_type kv_cache_type_from_str(const std::string & s) { +ggml_type common_kv_cache_type_from_str(const std::string & s) { for (const auto & type : kv_cache_types) { if (ggml_type_name(type) == s) { return type; @@ -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), @@ -2432,7 +2461,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.cache_type_k) ), [](common_params & params, const std::string & value) { - params.cache_type_k = kv_cache_type_from_str(value); + params.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_K")); add_opt(common_arg( @@ -2445,7 +2474,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.cache_type_v) ), [](common_params & params, const std::string & value) { - params.cache_type_v = kv_cache_type_from_str(value); + params.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_V")); add_opt(common_arg( @@ -4056,7 +4085,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.speculative.draft.cache_type_k) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_k = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K")); add_opt(common_arg( @@ -4069,7 +4098,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ggml_type_name(params.speculative.draft.cache_type_v) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_v = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V")); add_opt(common_arg( diff --git a/common/arg.h b/common/arg.h index 421bc295fc21..c82700bfc848 100644 --- a/common/arg.h +++ b/common/arg.h @@ -127,6 +127,9 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e // load all backends and print the list of available (non-CPU) devices to stdout void common_print_available_devices(); +// a KV cache type by its ggml name (the --cache-type-k/-v values); throws std::runtime_error otherwise +ggml_type common_kv_cache_type_from_str(const std::string & s); + // parse input arguments from CLI into a map bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map & out_map); diff --git a/common/common.cpp b/common/common.cpp index bda497246329..97e35d63df6d 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1271,6 +1271,35 @@ static void common_init_sampler_from_model( get_float(llama_model_meta_key_str(LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA), sparams.mirostat_eta, common_params_sampling_config::COMMON_PARAMS_SAMPLING_CONFIG_MIROSTAT_ETA); } +// the per-context state a context takes from `params`: control vectors and the adapter set +static bool common_context_attach(llama_context * lctx, llama_model * model, common_params & params) { + if (!params.control_vectors.empty()) { + if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; + if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); + + const auto cvec = common_control_vector_load(params.control_vectors); + if (cvec.n_embd == -1) { + return false; + } + + int err = llama_set_adapter_cvec( + lctx, + cvec.data.data(), + cvec.data.size(), + cvec.n_embd, + params.control_vector_layer_start, + params.control_vector_layer_end); + if (err) { + return false; + } + } + + if (!params.lora_init_without_apply) { + common_set_adapter_lora(lctx, params.lora_adapters); + } + return true; +} + struct common_init_result::impl { impl() = default; ~impl() = default; @@ -1286,6 +1315,9 @@ struct common_init_result::impl { std::vector samplers; std::vector samplers_seq_config; + + // the backend samplers (one per sequence) and a context on `model`; replaces any context held + llama_context * make_context(common_params & params, llama_context_params cparams); }; common_init_result::common_init_result(common_params & params, bool model_only) : @@ -1387,31 +1419,65 @@ common_init_result::common_init_result(common_params & params, bool model_only) params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end()); } + llama_context * lctx = pimpl->make_context(params, cparams); + if (lctx == NULL) { + return; + } + + set_process_priority(params.cpuparams.priority); + + pimpl->threadpools.init(lctx, params); +} + +llama_context * common_init_result::impl::make_context(common_params & params, llama_context_params cparams) { + llama_model * mdl = model.get(); + // init the backend samplers as part of the context creation - pimpl->samplers.resize(cparams.n_seq_max); - pimpl->samplers_seq_config.resize(cparams.n_seq_max); + samplers.clear(); + samplers.resize(cparams.n_seq_max); + samplers_seq_config.resize(cparams.n_seq_max); for (int i = 0; i < (int) cparams.n_seq_max; ++i) { - pimpl->samplers[i].reset(common_sampler_init(model, params.sampling)); - pimpl->samplers_seq_config[i] = { i, common_sampler_get(pimpl->samplers[i].get()) }; + samplers[i].reset(common_sampler_init(mdl, params.sampling)); + samplers_seq_config[i] = { i, common_sampler_get(samplers[i].get()) }; } if (params.sampling.backend_sampling) { - cparams.samplers = pimpl->samplers_seq_config.data(); - cparams.n_samplers = pimpl->samplers_seq_config.size(); + cparams.samplers = samplers_seq_config.data(); + cparams.n_samplers = samplers_seq_config.size(); } - llama_context * lctx = llama_init_from_model(model, cparams); + llama_context * lctx = llama_init_from_model(mdl, cparams); if (lctx == NULL) { COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); - return; + return nullptr; } - pimpl->context.reset(lctx); + context.reset(lctx); + return lctx; +} - set_process_priority(params.cpuparams.priority); +void common_init_result::free_context() { + pimpl->context.reset(); +} - pimpl->threadpools.init(lctx, params); +llama_context * common_init_result::reinit_context(common_params & params) { + GGML_ASSERT(pimpl->model); + + // free first: the old and the new context never hold memory at the same time + pimpl->context.reset(); + + llama_context * lctx = pimpl->make_context(params, common_context_params_to_llama(params)); + if (lctx == nullptr) { + return nullptr; + } + pimpl->threadpools.attach(lctx); + + if (!common_context_attach(lctx, pimpl->model.get(), params)) { + pimpl->context.reset(); + return nullptr; + } + return lctx; } llama_model * common_init_result::model() { @@ -1465,25 +1531,8 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode params.ctx_shift = false; } - if (!params.control_vectors.empty()) { - if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; - if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_model_n_layer(model); - - const auto cvec = common_control_vector_load(params.control_vectors); - if (cvec.n_embd == -1) { - return res; - } - - int err = llama_set_adapter_cvec( - lctx, - cvec.data.data(), - cvec.data.size(), - cvec.n_embd, - params.control_vector_layer_start, - params.control_vector_layer_end); - if (err) { - return res; - } + if (!common_context_attach(lctx, model, params)) { + return res; } if (llama_pooling_type(lctx) == LLAMA_POOLING_TYPE_RANK) { @@ -1510,10 +1559,6 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode } } - if (!params.lora_init_without_apply) { - common_set_adapter_lora(lctx, params.lora_adapters); - } - if (params.warmup) { COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n"); @@ -1829,7 +1874,13 @@ void common_threadpools::init(llama_context * ctx, const common_params & params) return; } - llama_attach_threadpool(ctx, threadpool, threadpool_batch); + attach(ctx); +} + +void common_threadpools::attach(llama_context * ctx) const { + if (threadpool != nullptr) { + llama_attach_threadpool(ctx, threadpool, threadpool_batch); + } } // diff --git a/common/common.h b/common/common.h index 29b171b7dda0..aebff5f398ab 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"; @@ -929,6 +931,15 @@ struct common_init_result { std::vector & lora(); + // Continuum fork: frees the context and builds one from `params` (n_ctx, n_parallel, KV types) + // on the SAME model: its weights, mmap and adapters stay loaded. Returns nullptr when the new + // context cannot be created; the result then holds no context and the caller may call it + // again with the old geometry. + llama_context * reinit_context(common_params & params); + + // Continuum fork: frees the context only; the model and adapters stay (reinit_context builds the next) + void free_context(); + private: struct impl; std::unique_ptr pimpl; @@ -965,6 +976,10 @@ struct common_threadpools { void init(llama_context * ctx, const common_params & params); + // Continuum fork: hands the pools init() made to another context of the same process (a + // context rebuilt with a new geometry keeps its CPU threads) + void attach(llama_context * ctx) const; + private: ggml_threadpool * threadpool = nullptr; ggml_threadpool * threadpool_batch = nullptr; diff --git a/common/speculative.cpp b/common/speculative.cpp index 4b67bbb88599..2f594f7c7e0b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2375,17 +2375,8 @@ struct common_speculative_init_result::impl { llama_context_ptr context; }; -common_speculative_init_result::common_speculative_init_result( - common_params & params, - llama_model * model_tgt, - llama_context * ctx_tgt) : - pimpl(new impl{}) { - const bool has_draft = params.speculative.has_dft(); - const bool spec_mtp = std::find(params.speculative.types.begin(), - params.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - - auto mparams = common_model_params_to_llama(params); +// the draft context's parameters: as many tokens per sequence as the target context, which it follows +static llama_context_params common_speculative_cparams(common_params & params, llama_context * ctx_tgt, bool spec_mtp) { auto cparams = common_context_params_to_llama(params); if (spec_mtp) { @@ -2399,6 +2390,24 @@ common_speculative_init_result::common_speculative_init_result( // the extra memory for small models is likely negligible? cparams.n_rs_seq = 0; cparams.ctx_other = ctx_tgt; + return cparams; +} + +static bool common_speculative_is_mtp(const common_params & params) { + return std::find(params.speculative.types.begin(), params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); +} + +common_speculative_init_result::common_speculative_init_result( + common_params & params, + llama_model * model_tgt, + llama_context * ctx_tgt) : + pimpl(new impl{}) { + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = common_speculative_is_mtp(params); + + auto mparams = common_model_params_to_llama(params); + auto cparams = common_speculative_cparams(params, ctx_tgt, spec_mtp); std::string model_path; if (has_draft) { @@ -2435,6 +2444,28 @@ common_speculative_init_result::common_speculative_init_result( } } +void common_speculative_init_result::free_context() { + pimpl->context.reset(); +} + +llama_context * common_speculative_init_result::reinit_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + pimpl->context.reset(); + + const bool spec_mtp = common_speculative_is_mtp(params); + llama_model * model = pimpl->model ? pimpl->model.get() : (spec_mtp ? model_tgt : nullptr); + if (model == nullptr) { + return nullptr; + } + + llama_context * ctx_dft = llama_init_from_model(model, common_speculative_cparams(params, ctx_tgt, spec_mtp)); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create the draft context\n", __func__); + return nullptr; + } + pimpl->context.reset(ctx_dft); + return ctx_dft; +} + common_speculative_init_result::~common_speculative_init_result() = default; llama_model * common_speculative_init_result::model() { diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de59..7f1dede52a02 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -96,6 +96,15 @@ struct common_speculative_init_result { llama_model * model(); llama_context * context(); + // Continuum fork: frees the draft context and builds one for `ctx_tgt` (a rebuilt target + // context) on the same draft model, or on `model_tgt` for MTP. The draft model stays loaded. + // Returns nullptr when the new context cannot be created. Free the draft context (this call + // does it first) before the target it points at. + llama_context * reinit_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + + // Continuum fork: frees the draft context only (it points at the target context, so it goes first) + void free_context(); + private: struct impl; std::unique_ptr pimpl; diff --git a/src/llama-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/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index fc118c44814c..362b25e28255 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2069,11 +2069,14 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama uint32_t n_stream_cur; io.read(&n_stream_cur, sizeof(n_stream_cur)); - if (n_stream_cur != n_stream) { + // Continuum fork: one sequence's cells all land in that sequence's stream, so its state + // reads into a cache split into any number of streams (a server whose slot count changed, + // or a unified cache from a split one). The whole cache (seq_id == -1) keeps its layout. + if (n_stream_cur != n_stream && seq_id == -1) { throw std::runtime_error("n_stream mismatch"); } - for (uint32_t s = 0; s < n_stream; ++s) { + for (uint32_t s = 0; s < n_stream_cur; ++s) { uint32_t cell_count; io.read(&cell_count, sizeof(cell_count)); diff --git a/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..4980750c4fdf 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-reconfigure.cpp + server-reconfigure.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..ff2faf994e24 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -482,6 +482,15 @@ 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_reconfigure = 0; // Continuum fork: /reconfigure runs (server-reconfigure.h) + uint64_t n_reconfigure_failed = 0; + uint64_t n_page_in = 0; + uint64_t n_page_failed = 0; // page_out, page_in or spill refused or failed + uint64_t n_ckpt_evicted_budget = 0; // checkpoints dropped to keep within --ctx-checkpoints-bytes + uint64_t n_ckpt_skipped_budget = 0; // checkpoints not taken: one alone exceeds the budget + void init() { t_start = ggml_time_us(); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 784d5426eeba..bfff1e46033a 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-reconfigure.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 @@ -121,6 +109,17 @@ struct server_batch { } } + // Continuum fork: a reconfigure to more slots than n_batch needs a larger batch + void reinit(int32_t n_tokens_alloc, int32_t n_embd) { + if (batch.pos != nullptr) { + clear(); + llama_batch_free(batch); + batch.pos = nullptr; + } + tokens.clear(); + init(n_tokens_alloc, n_embd); + } + void init(int32_t n_tokens_alloc, int32_t n_embd) { this->n_tokens_alloc = n_tokens_alloc; this->n_embd = n_embd; @@ -450,6 +449,18 @@ struct server_slot { return state != SLOT_STATE_IDLE; } + // Continuum fork (server-reconfigure.h): between two update_slots passes every token the slot + // holds is in its KV and nothing waits on the context's outputs, so its state can leave the + // context and come back. A prompt whose last token is decoded but not yet sampled, or a draft + // decoded but not yet verified, still reads the outputs. + bool at_boundary() const { + return state != SLOT_STATE_DONE_PROMPT && i_batch == -1 && spec_i_batch.empty(); + } + + // Continuum fork: the task registered this slot's sampler with the context (backend sampling); + // a rebuilt context registers it again + bool backend_sampler = false; + bool can_speculate() const { return !!spec; } @@ -875,6 +886,10 @@ struct server_context_impl { return metrics; } + // what the routes report about the loaded model and context (server_context::get_meta); + // read on the task thread, or before the loop starts + server_context_meta get_meta(); + void reset_metrics_bucket() { metrics.reset_bucket(); } @@ -938,6 +953,18 @@ 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; + + // Continuum fork: POST /reconfigure (server-reconfigure.h). The request waits here until every + // slot is at a boundary; admission stays closed meanwhile and the requests that arrive queue. + std::unique_ptr reconfigure_task; + server_reconfigure_request reconfigure_req; + int64_t t_reconfigure_posted_us = 0; + size_t n_reconfigure_deferred = 0; + std::string reconfigure_waiting_on; + int n_reconfigures = 0; + server_metrics metrics; // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync @@ -960,6 +987,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(); @@ -1275,32 +1308,7 @@ struct server_context_impl { } for (int i = 0; i < params_base.n_parallel; i++) { - server_slot & slot = slots[i]; - - slot.id = i; - slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft; - slot.mem.init(ctx_tgt, ctx_dft); - slot.spec = spec.get(); - slot.n_ctx = n_ctx_slot; - - slot.mctx = mctx; - slot.prompt.tokens.has_mtmd = mctx != nullptr; - - SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); - - slot.callback_on_release = [this](int id_slot) { - queue_tasks.pop_deferred_task(id_slot); - }; - - slot.callback_on_reset = [this](const server_slot & slot) { - // flush the generated token stats before reset() - if (slot.stats.n_gen > 0) { - metrics_on_prediction(slot); - } - }; - - slot.reset(); + init_slot(slots[i], i, n_ctx_slot); } { @@ -1352,7 +1360,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 { @@ -1387,6 +1400,38 @@ struct server_context_impl { return true; } + // points a slot at the current contexts and speculative state (load_model, reconfigure) + void bind_slot(server_slot & slot, int n_ctx_slot) { + slot.ctx_tgt = ctx_tgt; + slot.ctx_dft = ctx_dft; + slot.mem.init(ctx_tgt, ctx_dft); + slot.spec = spec.get(); + slot.n_ctx = n_ctx_slot; + } + + void init_slot(server_slot & slot, int id, int n_ctx_slot) { + slot.id = id; + bind_slot(slot, n_ctx_slot); + + slot.mctx = mctx; + slot.prompt.tokens.has_mtmd = mctx != nullptr; + + SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); + + slot.callback_on_release = [this](int id_slot) { + queue_tasks.pop_deferred_task(id_slot); + }; + + slot.callback_on_reset = [this](const server_slot & slot) { + // flush the generated token stats before reset() + if (slot.stats.n_gen > 0) { + metrics_on_prediction(slot); + } + }; + + slot.reset(); + } + // unlike load_model(), this is only called once during initialization bool init() { GGML_ASSERT(ctx_tgt != nullptr); @@ -1399,7 +1444,10 @@ struct server_context_impl { return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { + // a reconfigure runs between two passes, never inside one (server-reconfigure.h) + reconfigure_if_ready(); update_slots(); + reconfigure_if_ready(); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1784,6 +1832,7 @@ struct server_context_impl { } else { llama_set_sampler(ctx_tgt, slot.id, nullptr); } + slot.backend_sampler = use_backend_sampling; SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str()); @@ -2472,7 +2521,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 +2578,270 @@ 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; + } + + // + // Continuum fork: reconfigure without reloading the model (server-reconfigure.h). The mechanism + // lives there; this is the server's side of server_reconfigure_host. + // + + server_geometry geometry() const { + server_geometry g; + g.n_ctx = llama_n_ctx(ctx_tgt); + g.n_parallel = (int32_t) slots.size(); + g.type_k = params_base.cache_type_k; + g.type_v = params_base.cache_type_v; + return g; + } + + struct reconfigure_host final : server_reconfigure_host { + server_context_impl & s; + const bool with_spec; // the speculative state is rebuilt when the server had one + + explicit reconfigure_host(server_context_impl & s) : s(s), with_spec(s.spec != nullptr) {} + + server_geometry geometry() const override { return s.geometry(); } + int32_t n_ctx_seq() const override { return s.slots.empty() ? 0 : s.slots.front().n_ctx; } + + llama_context * target() const override { return s.ctx_tgt; } + llama_context * draft() const override { return s.ctx_dft; } + common_speculative * speculative() const override { return s.spec.get(); } + + server_prompt & slot_prompt(int id) override { return s.slots[id].prompt; } + std::vector & slot_lora(int id) override { return s.slots[id].lora; } + bool slot_processing(int id) const override { return s.slots[id].is_processing(); } + void slot_fail(int id, const std::string & why) override { + server_slot & slot = s.slots[id]; + s.send_error(slot, why, ERROR_TYPE_SERVER); + slot.prompt_clear(); + slot.release(); + } + + void free_contexts() override { + s.spec.reset(); // holds both contexts + if (s.spec_init) { + s.spec_init->free_context(); // the draft points at the target: it goes first + } + s.ctx_dft = nullptr; + s.llama_init->free_context(); + s.ctx_tgt = nullptr; + } + + bool build_contexts(const server_geometry & geo) override { + common_params p = s.params_base; + p.n_ctx = geo.n_ctx; + p.n_parallel = geo.n_parallel; + p.cache_type_k = geo.type_k; + p.cache_type_v = geo.type_v; + const auto output_limits = server_output_limits(p); + p.n_outputs_max = output_limits.total; + p.n_outputs_max_per_seq = output_limits.per_seq; + + s.ctx_tgt = s.llama_init->reinit_context(p); + if (s.ctx_tgt == nullptr) { + return false; + } + if (s.spec_init) { + common_params p_dft = common_base_params_to_speculative(p); + s.ctx_dft = s.spec_init->reinit_context(p_dft, s.model_tgt, s.ctx_tgt); + if (s.ctx_dft == nullptr) { + s.llama_init->free_context(); + s.ctx_tgt = nullptr; + return false; + } + p.speculative.draft.ctx_tgt = s.ctx_tgt; + p.speculative.draft.ctx_dft = s.ctx_dft; + } + if (with_spec) { + s.spec.reset(common_speculative_init(p.speculative, p.n_parallel)); + } + s.params_base = p; // a wake from sleep reloads this geometry + return true; + } + + void set_slots(int n) override { + const int n_ctx_slot = std::min(llama_n_ctx_seq(s.ctx_tgt), llama_model_n_ctx_train(s.model_tgt)); + while ((int) s.slots.size() > n) { + GGML_ASSERT(!s.slots.back().is_processing()); + s.slots.pop_back(); + } + for (auto & slot : s.slots) { + s.bind_slot(slot, n_ctx_slot); + if (slot.backend_sampler && slot.is_processing()) { + llama_set_sampler(s.ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); + } + } + while ((int) s.slots.size() < n) { + s.slots.emplace_back(); + s.init_slot(s.slots.back(), (int) s.slots.size() - 1, n_ctx_slot); + } + const int32_t n_alloc = std::max(llama_n_batch(s.ctx_tgt), (int32_t) s.slots.size()); + if (n_alloc > s.batch.n_tokens_alloc) { + s.batch.reinit(n_alloc, llama_model_n_embd_inp(s.model_tgt)); + } + s.learn_lora_id = -1; // the optimizer state lived in the freed context + if (s.prompt_cache) { + s.prompt_cache->limit_tokens = llama_n_ctx(s.ctx_tgt); + } + } + + server_page_store & pages() override { return s.page_store; } + + size_t drop_prompt_cache() override { + if (!s.prompt_cache) { + return 0; + } + const size_t n = s.prompt_cache->size(); + s.prompt_cache->states.clear(); + return n; + } + }; + + // runs a pending reconfigure once every slot it needs is at a boundary; called between passes + void reconfigure_if_ready() { + if (!reconfigure_task) { + return; + } + const server_task & task = *reconfigure_task; + + std::vector view; + view.reserve(slots.size()); + for (const auto & slot : slots) { + view.push_back({ slot.id, slot.is_processing(), slot.at_boundary(), (int32_t) slot.prompt.n_tokens() }); + } + + try { + auto plan = server_reconfigure_decide(reconfigure_req, geometry(), view, params_base.kv_unified, + llama_model_n_ctx_train(model_tgt), + string_format("reconfigure%d", n_reconfigures + 1)); + if (plan.wait) { + if (plan.wait_reason != reconfigure_waiting_on) { + SRV_INF("reconfigure waits: %s\n", plan.wait_reason.c_str()); + reconfigure_waiting_on = plan.wait_reason; + } + return; + } + auto res = std::make_unique(); + res->id = task.id; + res->receipt.t_drain_ms = (ggml_time_us() - t_reconfigure_posted_us) / 1000.0; + res->receipt.n_tasks_deferred = n_reconfigure_deferred; + reconfigure_host host(*this); + server_reconfigure_run(host, reconfigure_req, plan, res->receipt, string_format("reconfigure%d", n_reconfigures + 1)); + n_reconfigures++; + metrics.n_reconfigure++; + res->meta = std::make_shared(get_meta()); + queue_results.send(std::move(res)); + } catch (const server_page_error & err) { + metrics.n_reconfigure_failed++; + send_error(task, err.what(), err.type); + } + + // admission opens: the requests that queued meanwhile go back to the queue + reconfigure_task.reset(); + reconfigure_waiting_on.clear(); + for (size_t i = 0; i < slots.size(); ++i) { + queue_tasks.pop_deferred_task((int) i); + } + } + // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { // while yielding, an encode / decode is running and only reading the server state is safe @@ -2526,6 +2866,13 @@ struct server_context_impl { const int id_task = task.id; + if (reconfigure_task) { + // Continuum fork: admission is closed while a reconfigure waits (server-reconfigure.h) + n_reconfigure_deferred++; + queue_tasks.defer(std::move(task)); + break; + } + server_slot * slot = get_available_slot(task); // @@ -2645,6 +2992,11 @@ struct server_context_impl { res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); res->metrics = metrics; + res->page_store_bytes = page_store.n_bytes(); + res->page_store_pages = page_store.n_pages(); + for (const server_slot & slot : slots) { + res->ctx_checkpoint_bytes += server_checkpoints_bytes(slot.prompt.checkpoints); + } if (task.metrics_reset_bucket) { metrics.reset_bucket(); @@ -2729,39 +3081,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 +3120,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 +3144,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 +3208,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 +3327,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 +3395,51 @@ struct server_context_impl { res->n_erased = n_erased; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_RECONFIGURE: + { + // Continuum fork: runs between two passes once the slots allow it (server-reconfigure.h) + if (reconfigure_task) { + send_error(task, "a reconfigure is already pending", ERROR_TYPE_UNAVAILABLE); + break; + } + try { + reconfigure_req = server_reconfigure_request::parse(task.reconfigure, geometry()); + } catch (const server_page_error & err) { + metrics.n_reconfigure_failed++; + send_error(task, err.what(), err.type); + break; + } + t_reconfigure_posted_us = ggml_time_us(); + n_reconfigure_deferred = 0; + reconfigure_task = std::make_unique(std::move(task)); + reconfigure_if_ready(); + } break; + case SERVER_TASK_TYPE_PAGE_OUT: + case SERVER_TASK_TYPE_PAGE_IN: + case SERVER_TASK_TYPE_PAGE_SPILL: + case SERVER_TASK_TYPE_PAGE_DROP: + { + // Continuum fork: named host-RAM pages (server-pages.h). page_out and page_in + // mutate the slot's KV, so they defer while it is processing, as restore does. + const bool on_slot = task.type == SERVER_TASK_TYPE_PAGE_OUT || task.type == SERVER_TASK_TYPE_PAGE_IN; + server_slot * slot = on_slot ? get_slot_by_id(task.slot_action.id_slot) : nullptr; + if (on_slot && slot == nullptr) { + send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); + break; + } + if (on_slot && slot->is_processing()) { + SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); + queue_tasks.defer(std::move(task)); + break; + } + try { + auto res = page_task(task, slot); + queue_results.send(std::move(res)); + } catch (const server_page_error & err) { + metrics.n_page_failed += task.type == SERVER_TASK_TYPE_PAGE_DROP ? 0 : 1; + send_error(task, err.what(), err.type); + } + } break; case SERVER_TASK_TYPE_GET_LORA: { // TODO @ngxson : make lora_adapters a dedicated member of server_context @@ -3195,6 +3599,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. @@ -4737,47 +5142,52 @@ server_response_reader server_context::get_response_reader() { } server_context_meta server_context::get_meta() const { - auto bos_id = llama_vocab_bos(impl->vocab); - auto eos_id = llama_vocab_eos(impl->vocab); - auto bos_token_str = bos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(impl->ctx_tgt, bos_id, true) : ""; - auto eos_token_str = eos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(impl->ctx_tgt, eos_id, true) : ""; + return impl->get_meta(); +} + +server_context_meta server_context_impl::get_meta() { + auto bos_id = llama_vocab_bos(this->vocab); + auto eos_id = llama_vocab_eos(this->vocab); + auto bos_token_str = bos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(this->ctx_tgt, bos_id, true) : ""; + auto eos_token_str = eos_id != LLAMA_TOKEN_NULL ? common_token_to_piece(this->ctx_tgt, eos_id, true) : ""; - const char * ftype_name = llama_ftype_name(llama_model_ftype(impl->model_tgt)); + const char * ftype_name = llama_ftype_name(llama_model_ftype(this->model_tgt)); return server_context_meta { /* build_info */ std::string(llama_build_info()), - /* model_name */ impl->model_name, - /* model_aliases */ impl->model_aliases, - /* model_tags */ impl->model_tags, - /* model_path */ impl->params_base.model.path, - /* has_mtmd */ impl->mctx != nullptr, - /* has_inp_image */ impl->chat_params.allow_image, - /* has_inp_audio */ impl->chat_params.allow_audio, - /* has_inp_video */ impl->chat_params.allow_video, - /* json_ui_settings */ impl->json_ui_settings, - /* slot_n_ctx */ impl->get_slot_n_ctx(), - /* pooling_type */ llama_pooling_type(impl->ctx_tgt), - - /* chat_params */ impl->chat_params, - /* chat_template_caps */ common_chat_templates_get_caps(impl->chat_params.tmpls.get()), + /* model_name */ this->model_name, + /* model_aliases */ this->model_aliases, + /* model_tags */ this->model_tags, + /* model_path */ this->params_base.model.path, + /* has_mtmd */ this->mctx != nullptr, + /* has_inp_image */ this->chat_params.allow_image, + /* has_inp_audio */ this->chat_params.allow_audio, + /* has_inp_video */ this->chat_params.allow_video, + /* json_ui_settings */ this->json_ui_settings, + /* slot_n_ctx */ this->get_slot_n_ctx(), + /* pooling_type */ llama_pooling_type(this->ctx_tgt), + /* n_slots */ (int32_t) this->slots.size(), + + /* chat_params */ this->chat_params, + /* chat_template_caps */ common_chat_templates_get_caps(this->chat_params.tmpls.get()), /* bos_token_str */ bos_token_str, /* eos_token_str */ eos_token_str, - /* fim_pre_token */ llama_vocab_fim_pre(impl->vocab), - /* fim_sub_token */ llama_vocab_fim_suf(impl->vocab), - /* fim_mid_token */ llama_vocab_fim_mid(impl->vocab), - /* fim_pad_token */ llama_vocab_fim_pad(impl->vocab), - /* fim_rep_token */ llama_vocab_fim_rep(impl->vocab), - /* fim_sep_token */ llama_vocab_fim_sep(impl->vocab), - - /* logit_bias_eog */ impl->params_base.sampling.logit_bias_eog, - - /* model_vocab_type */ llama_vocab_type(impl->vocab), - /* model_vocab_n_tokens */ llama_vocab_n_tokens(impl->vocab), - /* model_n_ctx_train */ llama_model_n_ctx_train(impl->model_tgt), - /* model_n_embd_inp */ llama_model_n_embd(impl->model_tgt), - /* model_n_params */ llama_model_n_params(impl->model_tgt), - /* model_size */ llama_model_size(impl->model_tgt), + /* fim_pre_token */ llama_vocab_fim_pre(this->vocab), + /* fim_sub_token */ llama_vocab_fim_suf(this->vocab), + /* fim_mid_token */ llama_vocab_fim_mid(this->vocab), + /* fim_pad_token */ llama_vocab_fim_pad(this->vocab), + /* fim_rep_token */ llama_vocab_fim_rep(this->vocab), + /* fim_sep_token */ llama_vocab_fim_sep(this->vocab), + + /* logit_bias_eog */ this->params_base.sampling.logit_bias_eog, + + /* model_vocab_type */ llama_vocab_type(this->vocab), + /* model_vocab_n_tokens */ llama_vocab_n_tokens(this->vocab), + /* model_n_ctx_train */ llama_model_n_ctx_train(this->model_tgt), + /* model_n_embd_inp */ llama_model_n_embd(this->model_tgt), + /* model_n_params */ llama_model_n_params(this->model_tgt), + /* model_size */ llama_model_size(this->model_tgt), /* model_ftype */ ftype_name, /* model_weight_buffers */ [&] { // MEASURED placement: where the weights were actually allocated. @@ -4786,19 +5196,19 @@ server_context_meta server_context::get_meta() const { // so a lane that asked for the GPU and fell back to the CPU prints // numbers identical to one that succeeded. json bufs = json::array(); - const size_t n = llama_model_n_weight_backends(impl->model_tgt); + const size_t n = llama_model_n_weight_backends(this->model_tgt); for (size_t i = 0; i < n; i++) { - const char * backend = llama_model_weight_backend_name(impl->model_tgt, i); + const char * backend = llama_model_weight_backend_name(this->model_tgt, i); bufs.push_back(json { { "backend", backend ? backend : "" }, - { "size_bytes", llama_model_weight_backend_bytes(impl->model_tgt, i) }, + { "size_bytes", llama_model_weight_backend_bytes(this->model_tgt, i) }, }); } return bufs; }(), /* memory_breakdown */ [&] { json out = json::array(); - for (const auto & [buft, mb] : llama_get_memory_breakdown(impl->ctx_tgt)) { + for (const auto & [buft, mb] : llama_get_memory_breakdown(this->ctx_tgt)) { out.push_back(json { { "buffer_type", ggml_backend_buft_name(buft) }, { "model_bytes", mb.model }, @@ -4809,9 +5219,9 @@ server_context_meta server_context::get_meta() const { return out; }(), /* context_geometry */ json { - { "allocated_tokens", llama_n_ctx(impl->ctx_tgt) }, - { "max_tokens_per_sequence", llama_n_ctx_seq(impl->ctx_tgt) }, - { "sequence_capacity", llama_n_seq_max(impl->ctx_tgt) }, + { "allocated_tokens", llama_n_ctx(this->ctx_tgt) }, + { "max_tokens_per_sequence", llama_n_ctx_seq(this->ctx_tgt) }, + { "sequence_capacity", llama_n_seq_max(this->ctx_tgt) }, }, }; } @@ -4915,7 +5325,7 @@ std::unique_ptr server_routes::handle_completions_impl( task.params = server_schema::eval_llama_cmpl_schema( ctx_server.vocab, params, - meta->logit_bias_eog, + meta_now()->logit_bias_eog, data); task.params.message_spans = task.tokens.find_message_spans(delimiters); @@ -4926,7 +5336,7 @@ std::unique_ptr server_routes::handle_completions_impl( // OAI-compat task.params.res_type = res_type; task.params.oaicompat_cmpl_id = completion_id; - task.params.oaicompat_model = meta->model_name; + task.params.oaicompat_model = meta_now()->model_name; // prepare child tasks if (task.params.n_cmpl > 1) { @@ -5211,8 +5621,9 @@ static json get_res_props(const server_context_meta & meta, const common_params json props = { { "default_generation_settings", default_generation_settings_for_props }, - { "total_slots", params.n_parallel }, + { "total_slots", meta.n_slots }, { "slot_save_size_version", 1 }, + { "slot_page_version", 1 }, // Continuum fork: /slots page_out, page_in, spill, page_drop { "expert_container_custody_version", 1 }, { "model_alias", meta.model_name }, { "model_ftype", meta.model_ftype }, @@ -5249,7 +5660,7 @@ static json get_res_props(const server_context_meta & meta, const common_params } json server_routes::get_model_info() const { - return get_res_model_info(*meta); + return get_res_model_info(*meta_now()); } void server_routes::init_routes() { @@ -5388,7 +5799,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 +5818,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"); } @@ -5426,7 +5842,7 @@ void server_routes::init_routes() { std::unique_lock lock(mutex_cache); res->ok(cached_props); } else { - res->ok(get_res_props(*meta, params, false)); + res->ok(get_res_props(*meta_now(), params, false)); } return res; }; @@ -5507,7 +5923,7 @@ void server_routes::init_routes() { data.at("input_extra"), params.n_batch, params.n_predict, - meta->slot_n_ctx, + meta_now()->slot_n_ctx, params.spm_infill, tokenized_prompts[0].get_tokens() // TODO: this could maybe be multimodal. ); @@ -5551,7 +5967,7 @@ void server_routes::init_routes() { json body = json::parse(req.body); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5610,7 +6026,7 @@ void server_routes::init_routes() { SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5627,7 +6043,7 @@ void server_routes::init_routes() { this->post_transcriptions_oai = [this](const server_http_req & req) { auto res = create_response(); - if (!meta->has_mtmd || !meta->chat_params.allow_audio) { + if (!meta_now()->has_mtmd || !meta_now()->chat_params.allow_audio) { res->error(format_error_response("The current model does not support audio input.", ERROR_TYPE_NOT_SUPPORTED)); return res; } @@ -5635,14 +6051,14 @@ void server_routes::init_routes() { std::vector files; json body = convert_transcriptions_to_chatcmpl( json::parse(req.body), - meta->chat_params.tmpls.get(), + meta_now()->chat_params.tmpls.get(), req.files, files); SRV_DBG("%s\n", "Request converted: OpenAI Transcriptions -> OpenAI Chat Completions"); SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5660,7 +6076,7 @@ void server_routes::init_routes() { SRV_DBG("converted request: %s\n", body.dump().c_str()); json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); return handle_completions_impl( req, @@ -5681,7 +6097,7 @@ void server_routes::init_routes() { json body = json::parse(req.body); json data = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); res->ok({{ "prompt", std::move(data.at("prompt")) }}); return res; @@ -5694,7 +6110,7 @@ void server_routes::init_routes() { std::unique_lock lock(mutex_cache); res->ok(cached_models); } else { - res->ok(get_res_models(*meta)); + res->ok(get_res_models(*meta_now())); } return res; }; @@ -5832,7 +6248,7 @@ void server_routes::init_routes() { // write JSON response json root = format_response_rerank( body, - meta->model_name, + meta_now()->model_name, responses, is_tei_format, documents, @@ -5965,6 +6381,42 @@ void server_routes::init_routes() { // Duty rest returns a deferred outcome: the caller retains replayable input context, reply or // labels, and adapter parameters, then retries after retry_after_ms. A slot's prompt may have // changed by then. The slot-release queue cannot wake a cooldown on an already idle slot. + // Continuum fork: POST /reconfigure {"n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force"} + // changes the context geometry without reloading the model; the answer is the receipt + // (server-reconfigure.h). It waits while slots drain, so a caller gives it a long timeout. + this->post_reconfigure = [this](const server_http_req & req) { + auto res = create_response(); + json body = json::object(); + if (!req.body.empty()) { + try { + body = json::parse(req.body); + } catch (const std::exception &) { + res->error(format_error_response("The body must be JSON", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + } + { + server_task task(SERVER_TASK_TYPE_RECONFIGURE); + task.id = res->rd.get_new_id(); + task.reconfigure = std::move(body); + res->rd.post_task(std::move(task)); + } + auto result = res->rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + auto * done = dynamic_cast(result.get()); + GGML_ASSERT(done != nullptr); + set_meta(done->meta); // /props and the slot bounds now report the new geometry + res->ok(done->to_json()); + return res; + }; + this->post_learn = [this](const server_http_req & req) { auto res = create_response(); if (params.train_dir.empty()) { @@ -5997,7 +6449,7 @@ void server_routes::init_routes() { served["messages"].push_back(messages[i]); } std::vector files; - const json parsed = oaicompat_chat_params_parse(served, meta->chat_params, files); + const json parsed = oaicompat_chat_params_parse(served, meta_now()->chat_params, files); if (!files.empty()) { res->error(format_error_response("a conversation with media cannot be learned yet: only text prompts", ERROR_TYPE_INVALID_REQUEST)); return res; @@ -6012,7 +6464,7 @@ void server_routes::init_routes() { framed["tools"] = tools; framed["add_generation_prompt"] = generation; std::vector media; - const json rendered = oaicompat_chat_params_parse(framed, meta->chat_params, media, + const json rendered = oaicompat_chat_params_parse(framed, meta_now()->chat_params, media, generation ? server_chat_render::serving : server_chat_render::closed_transcript); if (!media.empty()) { throw std::invalid_argument("learning media is not supported"); @@ -6227,7 +6679,7 @@ std::unique_ptr server_routes::handle_slots_restore(const valid = (i == 8 || i == 13 || i == 18 || i == 23) ? id[i] == '-' : ((id[i] >= '0' && id[i] <= '9') || (id[i] >= 'a' && id[i] <= 'f')); } - if (!valid || id_slot < 0 || id_slot >= params.n_parallel) { + if (!valid || id_slot < 0 || id_slot >= meta_now()->n_slots) { res->error(format_error_response("Invalid restore operation identity or slot", ERROR_TYPE_INVALID_REQUEST)); return res; } @@ -6293,6 +6745,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 >= meta_now()->n_slots)) { + 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; @@ -6331,7 +6892,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons return res; } - if (res_type != TASK_RESPONSE_TYPE_NONE && meta->pooling_type == LLAMA_POOLING_TYPE_NONE) { + if (res_type != TASK_RESPONSE_TYPE_NONE && meta_now()->pooling_type == LLAMA_POOLING_TYPE_NONE) { res->error(format_error_response("Pooling type 'none' is not OAI compatible. Please use a different pooling type", ERROR_TYPE_INVALID_REQUEST)); return res; } @@ -6373,8 +6934,8 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons int embd_normalize = params.embd_normalize; if (body.count("embd_normalize") != 0) { embd_normalize = body.at("embd_normalize").get(); - if (meta->pooling_type == LLAMA_POOLING_TYPE_NONE) { - SRV_DBG("embd_normalize is not supported by pooling type %d, ignoring it\n", meta->pooling_type); + if (meta_now()->pooling_type == LLAMA_POOLING_TYPE_NONE) { + SRV_DBG("embd_normalize is not supported by pooling type %d, ignoring it\n", meta_now()->pooling_type); } } @@ -6416,7 +6977,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons // write JSON response json root = res_type == TASK_RESPONSE_TYPE_OAI_EMBD - ? format_embeddings_response_oaicompat(body, meta->model_name, responses, use_base64) + ? format_embeddings_response_oaicompat(body, meta_now()->model_name, responses, use_base64) : json(responses); res->ok(root); return res; @@ -6449,7 +7010,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l json body_parsed = oaicompat_chat_params_parse( body, - meta->chat_params, + meta_now()->chat_params, files); json prompt = body_parsed.at("prompt"); // SRV_DBG("prompt = %s\n", prompt.dump().c_str()); @@ -6478,8 +7039,8 @@ void server_routes::update_cached_responses(bool is_sleeping) { std::unique_lock lock(mutex_cache); if (is_sleeping) { - cached_models = get_res_models(*meta); - cached_props = get_res_props(*meta, params, true); + cached_models = get_res_models(*meta_now()); + cached_props = get_res_props(*meta_now(), params, true); cached_metrics = ctx_server.get_metrics(); should_reset_buckets = false; diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 6deb3a13cc06..0b4489155613 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -26,6 +26,7 @@ struct server_context_meta { json json_ui_settings; int slot_n_ctx; enum llama_pooling_type pooling_type; + int32_t n_slots; // Continuum fork: changes with /reconfigure, so routes read it here, not from params // chat params server_chat_params & chat_params; @@ -136,7 +137,7 @@ struct server_routes { // note: this is not thread-safe and can only when ctx_http.is_ready is false void update_meta(const server_context & ctx_server) { - this->meta = std::make_unique(ctx_server.get_meta()); + set_meta(std::make_shared(ctx_server.get_meta())); } // handlers using lambda function, so that they can capture `this` without `std::bind` @@ -176,6 +177,7 @@ struct server_routes { server_http_context::handler_t post_train_measure; server_http_context::handler_t post_train_pause; server_http_context::handler_t post_train_resume; + server_http_context::handler_t post_reconfigure; // Continuum fork (server-reconfigure.h) // to be used in router mode json get_model_info() const; @@ -190,11 +192,15 @@ 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); - // using unique_ptr to allow late initialization of const - std::unique_ptr meta; + // late-initialized, and replaced by /reconfigure while other HTTP threads read it: every access + // goes through meta_now() / set_meta() (atomic shared_ptr operations) + std::shared_ptr meta; + std::shared_ptr meta_now() const { return std::atomic_load(&meta); } + void set_meta(std::shared_ptr m) { std::atomic_store(&meta, std::move(m)); } const common_params & params; server_context_impl & ctx_server; diff --git a/tools/server/server-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-reconfigure.cpp b/tools/server/server-reconfigure.cpp new file mode 100644 index 000000000000..686c5bdddcb3 --- /dev/null +++ b/tools/server/server-reconfigure.cpp @@ -0,0 +1,379 @@ +#include "server-reconfigure.h" + +#include "server-checkpoints.h" + +#include "arg.h" + +#include +#include + +json server_geometry::to_json() const { + return json { + { "n_ctx", n_ctx }, + { "n_parallel", n_parallel }, + { "cache_type_k", ggml_type_name(type_k) }, + { "cache_type_v", ggml_type_name(type_v) }, + }; +} + +int32_t server_geometry_n_ctx_seq(const server_geometry & geo, bool kv_unified, int32_t n_ctx_train) { + const int32_t n_ctx = GGML_PAD(geo.n_ctx, 256); + int32_t n_ctx_seq = kv_unified ? n_ctx : GGML_PAD(n_ctx / std::max(1, geo.n_parallel), 256); + if (n_ctx_train > 0) { + n_ctx_seq = std::min(n_ctx_seq, n_ctx_train); + } + return n_ctx_seq; +} + +static int32_t parse_positive(const json & body, const char * key, int32_t cur, int32_t max) { + if (!body.contains(key)) { + return cur; + } + const json & v = body.at(key); + if (!v.is_number_integer() || v.get() < 1 || v.get() > max) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format("%s must be an integer in [1, %d]", key, max)); + } + return v.get(); +} + +static ggml_type parse_kv_type(const json & body, const char * key, ggml_type cur) { + if (!body.contains(key)) { + return cur; + } + if (!body.at(key).is_string()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format("%s must be a KV cache type name", key)); + } + try { + return common_kv_cache_type_from_str(body.at(key).get()); + } catch (const std::runtime_error & e) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, e.what()); + } +} + +server_reconfigure_request server_reconfigure_request::parse(const json & body, const server_geometry & cur) { + if (!body.is_object()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "the body must be a JSON object"); + } + static const char * known[] = { "n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force" }; + for (const auto & item : body.items()) { + const std::string & key = item.key(); + if (std::find_if(std::begin(known), std::end(known), [&](const char * k) { return key == k; }) == std::end(known)) { + // a misspelled field silently keeping the old geometry is the bug this refuses + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "unknown field: " + key); + } + } + server_reconfigure_request req; + req.geo.n_ctx = parse_positive(body, "n_ctx", cur.n_ctx, INT32_MAX / 2); + req.geo.n_parallel = parse_positive(body, "n_parallel", cur.n_parallel, (int32_t) llama_max_parallel_sequences()); + req.geo.type_k = parse_kv_type (body, "cache_type_k", cur.type_k); + req.geo.type_v = parse_kv_type (body, "cache_type_v", cur.type_v); + if (body.contains("force")) { + if (!body.at("force").is_boolean()) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, "force must be a boolean"); + } + req.force = body.at("force").get(); + } + return req; +} + +server_reconfigure_plan server_reconfigure_decide( + const server_reconfigure_request & req, + const server_geometry & cur, + const std::vector & slots, + bool kv_unified, + int32_t n_ctx_train, + const std::string & name_prefix) { + server_reconfigure_plan plan; + plan.n_ctx_seq = server_geometry_n_ctx_seq(req.geo, kv_unified, n_ctx_train); + plan.cold_start = !req.geo.same_kv_format(cur); + + if (plan.cold_start && !req.force) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format( + "a KV type change (%s/%s -> %s/%s) cannot carry the residents' states over: pass \"force\": true to cold start", + ggml_type_name(cur.type_k), ggml_type_name(cur.type_v), ggml_type_name(req.geo.type_k), ggml_type_name(req.geo.type_v))); + } + + for (const auto & s : slots) { + if (!s.processing) { + continue; + } + if (s.n_tokens + 1 > plan.n_ctx_seq) { + throw server_page_error(ERROR_TYPE_INVALID_REQUEST, string_format( + "slot %d is serving a request holding %d tokens and the new geometry gives a sequence %d: " + "shrinking under a live request would truncate it", + s.id, s.n_tokens, plan.n_ctx_seq)); + } + if (!s.at_boundary) { + plan.wait = true; + plan.wait_reason = string_format("slot %d is between a decode and its sampling", s.id); + } else if (s.id >= req.geo.n_parallel) { + plan.wait = true; + plan.wait_reason = string_format("slot %d serves a request and has no place among %d slots", s.id, req.geo.n_parallel); + } else if (plan.cold_start) { + plan.wait = true; + plan.wait_reason = string_format("slot %d serves a request and a cold start has no state to resume it from", s.id); + } + } + if (plan.wait) { + return plan; + } + + // slots of the new geometry that nobody keeps: an occupied idle slot past n_parallel moves into one + std::vector taken(req.geo.n_parallel, false); + for (const auto & s : slots) { + if (s.n_tokens > 0 && s.id < req.geo.n_parallel && s.n_tokens <= plan.n_ctx_seq) { + taken[s.id] = true; + } + } + + for (const auto & s : slots) { + if (s.n_tokens == 0) { + continue; + } + server_reconfigure_move m; + m.id_from = s.id; + m.processing = s.processing; + m.n_tokens = s.n_tokens; + if (plan.cold_start) { + m.fate = SERVER_RECONFIGURE_DROPPED; + m.reason = "KV type change (cold start)"; + } else if (s.n_tokens > plan.n_ctx_seq) { + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.reason = string_format("%d tokens do not fit a sequence of %d", s.n_tokens, plan.n_ctx_seq); + } else if (s.id < req.geo.n_parallel) { + m.id_to = s.id; + } else { + const auto free = std::find(taken.begin(), taken.end(), false); + if (free == taken.end()) { + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.reason = string_format("no free slot among %d", req.geo.n_parallel); + } else { + *free = true; + m.id_to = (int) (free - taken.begin()); + } + } + if (m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE) { + m.page = string_format("%s.slot%d", name_prefix.c_str(), s.id); + } + plan.moves.push_back(std::move(m)); + } + return plan; +} + +static const char * fate_name(server_reconfigure_fate f) { + switch (f) { + case SERVER_RECONFIGURE_RESTORED: return "restored"; + case SERVER_RECONFIGURE_KEPT_IN_STORE: return "kept_in_store"; + case SERVER_RECONFIGURE_DROPPED: return "dropped"; + } + return "unknown"; +} + +json server_reconfigure_move::to_json() const { + json j = { + { "id_slot_before", id_from }, + { "id_slot_after", id_to }, + { "in_flight", processing }, + { "n_tokens", n_tokens }, + { "n_bytes", n_bytes }, + { "fate", fate_name(fate) }, + }; + if (!page.empty()) { + j["page"] = page; + } + if (!reason.empty()) { + j["reason"] = reason; + } + return j; +} + +json server_reconfigure_receipt::to_json() const { + json moved = json::array(); + int n_restored = 0, n_kept = 0, n_dropped = 0; + for (const auto & m : moves) { + moved.push_back(m.to_json()); + n_restored += m.fate == SERVER_RECONFIGURE_RESTORED; + n_kept += m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE; + n_dropped += m.fate == SERVER_RECONFIGURE_DROPPED; + } + json j = { + { "before", before.to_json() }, + { "after", after.to_json() }, + { "n_ctx_seq_before", n_ctx_seq_before }, + { "n_ctx_seq_after", n_ctx_seq_after }, + { "cold_start", cold_start }, + { "rolled_back", rolled_back }, + { "t_drain_ms", t_drain_ms }, + { "t_page_out_ms", t_page_out_ms }, + { "t_free_ms", t_free_ms }, + { "t_init_ms", t_init_ms }, + { "t_page_in_ms", t_page_in_ms }, + { "t_total_ms", t_total_ms }, + { "n_bytes", n_bytes }, + { "n_bytes_dropped", n_bytes_dropped }, + { "n_tasks_deferred", n_tasks_deferred }, + { "n_restored", n_restored }, + { "n_kept_in_store", n_kept }, + { "n_dropped", n_dropped }, + { "pages", moved }, + }; + if (rolled_back) { + j["rollback_reason"] = rollback_reason; + } + return j; +} + +void server_reconfigure_run( + server_reconfigure_host & host, + const server_reconfigure_request & req, + server_reconfigure_plan & plan, + server_reconfigure_receipt & rc, + const std::string & page_prefix) { + rc.before = host.geometry(); + rc.n_ctx_seq_before = host.n_ctx_seq(); + server_page_store & store = host.pages(); + + // the store must take every page that cannot be restored before anything moves, counted as it charges them + std::vector shapes; + shapes.reserve(plan.moves.size()); + { + size_t n_kept = 0; + for (const auto & m : plan.moves) { + shapes.push_back(server_page_measure(host.target(), host.draft(), host.speculative(), m.id_from)); + if (m.fate == SERVER_RECONFIGURE_KEPT_IN_STORE) { + n_kept += shapes.back().n_bytes(server_checkpoints_bytes(host.slot_prompt(m.id_from).checkpoints)); + } + } + if (store.limit_bytes > 0 && store.n_bytes() + n_kept > store.limit_bytes) { + throw server_page_error(ERROR_TYPE_UNAVAILABLE, string_format( + "%zu bytes of slot state cannot be restored under the new geometry and the page store (%zu of %llu bytes) cannot hold them", + n_kept, store.n_bytes(), (unsigned long long) store.limit_bytes)); + } + } + + // 1. page out: every occupied slot's state into host RAM. A cold start's states are captured + // too and dropped only once the new geometry stands, so a failed build rolls back whole. + const int64_t t0 = ggml_time_us(); + std::vector pages(plan.moves.size()); + for (size_t i = 0; i < plan.moves.size(); ++i) { + auto & m = plan.moves[i]; + pages[i] = server_page_capture(host.target(), host.draft(), m.id_from, std::move(shapes[i])); + server_prompt & prompt = host.slot_prompt(m.id_from); + const bool in_place = m.fate == SERVER_RECONFIGURE_RESTORED && m.id_to == m.id_from; + if (!in_place) { + // it leaves its slot: the page carries the tokens, checkpoints and adapters + pages[i]->prompt = std::move(prompt); + pages[i]->lora = host.slot_lora(m.id_from); + prompt.clear(); + } + m.n_bytes = pages[i]->n_bytes() + (in_place ? server_checkpoints_bytes(prompt.checkpoints) : 0); + rc.n_bytes += m.fate == SERVER_RECONFIGURE_DROPPED ? 0 : m.n_bytes; + } + const int64_t t1 = ggml_time_us(); + + // 2. free both contexts; the model, its mmap and the adapters stay + host.free_contexts(); + const int64_t t2 = ggml_time_us(); + + // 3. build the new geometry, or the old one again + if (!host.build_contexts(req.geo)) { + rc.rolled_back = true; + rc.rollback_reason = "the new geometry could not be built (see the log); the old one was rebuilt"; + SRV_ERR("reconfigure to n_ctx = %d, n_parallel = %d failed; rebuilding the old geometry\n", req.geo.n_ctx, req.geo.n_parallel); + if (!host.build_contexts(rc.before)) { + GGML_ABORT("reconfigure: neither the new nor the old context geometry could be built"); + } + } + const int64_t t3 = ggml_time_us(); + + // 4. the slots of the geometry that stands + host.set_slots(rc.rolled_back ? rc.before.n_parallel : req.geo.n_parallel); + + // 5. page in: every state back, or into the store with a typed reason; never truncated + for (size_t i = 0; i < plan.moves.size(); ++i) { + auto & m = plan.moves[i]; + if (rc.rolled_back) { + m.fate = SERVER_RECONFIGURE_RESTORED; // the old geometry holds every state it held + m.id_to = m.id_from; + m.reason.clear(); + m.page.clear(); + } + if (m.fate == SERVER_RECONFIGURE_DROPPED) { + rc.n_bytes_dropped += m.n_bytes; + pages[i].reset(); + continue; + } + // a page that left its slot carries the slot's tokens; one restored in place does not + const bool carries_prompt = pages[i]->prompt.tokens.size() > 0; + if (m.fate == SERVER_RECONFIGURE_RESTORED) { + try { + server_page_restore(*pages[i], host.target(), host.draft(), host.speculative(), m.id_to); + if (carries_prompt) { + host.slot_prompt(m.id_to) = std::move(pages[i]->prompt); + host.slot_lora(m.id_to) = pages[i]->lora; + } + continue; + } catch (const server_page_error & err) { + m.reason = err.what(); + if (host.slot_processing(m.id_to)) { + // its KV is gone and its request cannot continue: it ends with a typed error + // instead of going on with a context it no longer has + host.slot_fail(m.id_to, std::string("the reconfigure could not restore this request's state: ") + err.what()); + m.fate = SERVER_RECONFIGURE_DROPPED; + continue; + } + if (!carries_prompt) { + pages[i]->prompt = std::move(host.slot_prompt(m.id_to)); + pages[i]->lora = host.slot_lora(m.id_to); + } + host.slot_prompt(m.id_to).clear(); + m.fate = SERVER_RECONFIGURE_KEPT_IN_STORE; + m.page = string_format("%s.slot%d", page_prefix.c_str(), m.id_from); + } + } + m.id_to = -1; + try { + store.put(m.page, pages[i]); + SRV_WRN("reconfigure: slot %d's %d tokens stay in the page store as '%s': %s\n", + m.id_from, m.n_tokens, m.page.c_str(), m.reason.c_str()); + } catch (const server_page_error & err) { + // admitted above unless its restore failed here: a state the store cannot take is dropped, typed + m.reason = (m.reason.empty() ? "" : m.reason + "; ") + "the page store refused it: " + err.what(); + SRV_ERR("reconfigure: slot %d's %d tokens are dropped: %s\n", m.id_from, m.n_tokens, m.reason.c_str()); + m.fate = SERVER_RECONFIGURE_DROPPED; + m.page.clear(); + rc.n_bytes_dropped += m.n_bytes; + } + } + + rc.cold_start = plan.cold_start && !rc.rolled_back; + if (rc.cold_start) { + // every state in the old KV format: the store's pages and the prompt cache go with them + const size_t n_resident = rc.n_bytes_dropped; + const size_t n_store = store.n_bytes(); + const size_t n_pages = store.n_pages(); + const size_t n_cache = host.drop_prompt_cache(); + store.clear(); + rc.n_bytes_dropped += n_store + n_cache; + SRV_WRN("reconfigure: KV type change %s/%s -> %s/%s is a cold start: dropped %zu resident bytes, %zu page-store bytes (%zu pages), %zu prompt-cache bytes\n", + ggml_type_name(rc.before.type_k), ggml_type_name(rc.before.type_v), + ggml_type_name(req.geo.type_k), ggml_type_name(req.geo.type_v), + n_resident, n_store, n_pages, n_cache); + } + const int64_t t4 = ggml_time_us(); + + rc.after = host.geometry(); + rc.n_ctx_seq_after = host.n_ctx_seq(); + rc.moves = plan.moves; + rc.t_page_out_ms = (t1 - t0) / 1000.0; + rc.t_free_ms = (t2 - t1) / 1000.0; + rc.t_init_ms = (t3 - t2) / 1000.0; + rc.t_page_in_ms = (t4 - t3) / 1000.0; + rc.t_total_ms = (t4 - t0) / 1000.0; + + SRV_INF("reconfigure: n_ctx %d -> %d, n_parallel %d -> %d, %s/%s; pause %.1f ms (page out %.1f, free %.1f, init %.1f, page in %.1f) for %zu bytes; drain %.1f ms%s\n", + rc.before.n_ctx, rc.after.n_ctx, rc.before.n_parallel, rc.after.n_parallel, + ggml_type_name(rc.after.type_k), ggml_type_name(rc.after.type_v), + rc.t_total_ms, rc.t_page_out_ms, rc.t_free_ms, rc.t_init_ms, rc.t_page_in_ms, rc.n_bytes, rc.t_drain_ms, + rc.rolled_back ? " (ROLLED BACK)" : ""); +} diff --git a/tools/server/server-reconfigure.h b/tools/server/server-reconfigure.h new file mode 100644 index 000000000000..45e3e225b7f1 --- /dev/null +++ b/tools/server/server-reconfigure.h @@ -0,0 +1,182 @@ +#pragma once + +// Continuum fork: a running server changes its context geometry without reloading the model. +// +// POST /reconfigure {"n_ctx", "n_parallel", "cache_type_k", "cache_type_v", "force"} is a pause, not +// a relaunch: admission stops, the slots reach a ubatch boundary, every occupied slot's sequence state +// (target KV, draft KV, speculative tail, checkpoints) moves to host RAM (server-pages.h), the target +// and draft contexts are freed and rebuilt on the SAME llama_model (weights and mmap stay), the slots +// are rebuilt for the new n_parallel, the states go back, and serving resumes. A request in flight +// keeps its slot and continues where it paused; nothing is truncated and nothing is lost. +// Design: continuum docs/architecture/SERVING-ENGINE.md ยง3 ("a relaunch is a pause"), slice S2. +// +// What the core sees: the receipt (server_reconfigure_receipt::to_json) and /props context_geometry. + +#include "server-common.h" +#include "server-pages.h" +#include "server-task.h" + +#include "common.h" +#include "llama.h" + +#include +#include +#include +#include + +struct server_context_meta; + +// The context geometry a reconfigure may change. Everything else (the model, adapters, batch +// sizes, speculative types, checkpoint budget) stays as launched. +struct server_geometry { + int32_t n_ctx = 0; // KV cells of the whole context (llama_n_ctx) + int32_t n_parallel = 0; // slots, one sequence each (llama_n_seq_max) + ggml_type type_k = GGML_TYPE_F16; + ggml_type type_v = GGML_TYPE_F16; + + // a sequence state is written in its KV types: across a type change no page can be carried + bool same_kv_format(const server_geometry & other) const { + return type_k == other.type_k && type_v == other.type_v; + } + + json to_json() const; +}; + +// The tokens one sequence may hold under `geo`: libllama's rule (llama-context.cpp: n_ctx padded to +// 256; a split cache gives each sequence n_ctx / n_seq_max, padded) and the server's cap at the +// model's training context. Used before the new context exists, so a refusal costs nothing. +int32_t server_geometry_n_ctx_seq(const server_geometry & geo, bool kv_unified, int32_t n_ctx_train); + +struct server_reconfigure_request { + server_geometry geo; + // a KV type change is an intentional cold start: refused unless `force`, then every resident + // state, page and prompt-cache entry is dropped (logged loudly and counted in the receipt) + bool force = false; + + // Absent fields keep `cur`. Throws server_page_error (ERROR_TYPE_INVALID_REQUEST). + static server_reconfigure_request parse(const json & body, const server_geometry & cur); +}; + +// One slot as the reconfigure sees it on the task thread. +struct server_reconfigure_slot { + int id = -1; + bool processing = false; // a request is in flight + bool at_boundary = true; // no sampling or verification is half done (see server-context.cpp) + int32_t n_tokens = 0; // tokens its KV holds +}; + +// Where an occupied slot's state goes. +enum server_reconfigure_fate { + SERVER_RECONFIGURE_RESTORED, // into a slot of the new geometry + SERVER_RECONFIGURE_KEPT_IN_STORE, // a named RAM page (/slots/{id}?action=page_in brings it back) + SERVER_RECONFIGURE_DROPPED, // cold start (KV type change with force) +}; + +struct server_reconfigure_move { + int id_from = -1; + int id_to = -1; // -1 unless restored + bool processing = false; + int32_t n_tokens = 0; + size_t n_bytes = 0; + + server_reconfigure_fate fate = SERVER_RECONFIGURE_RESTORED; + std::string page; // the store name when kept + std::string reason; // why it was not restored + + json to_json() const; +}; + +// The decision before anything is touched. +struct server_reconfigure_plan { + bool wait = false; // a slot must first finish (it sits past the new n_parallel) or reach a boundary + std::string wait_reason; + + bool cold_start = false; // KV type change with force + int32_t n_ctx_seq = 0; // predicted tokens per sequence after + std::vector moves; // one per occupied slot, fate decided +}; + +// Decides the reconfigure. Throws server_page_error for a refusal: a KV type change without force, +// an in-flight request whose tokens do not fit the new window (shrinking under a live act would +// truncate it), or a geometry libllama would reject. `name_prefix` names pages that stay in the store. +server_reconfigure_plan server_reconfigure_decide( + const server_reconfigure_request & req, + const server_geometry & cur, + const std::vector & slots, + bool kv_unified, + int32_t n_ctx_train, + const std::string & name_prefix); + +struct server_reconfigure_receipt { + server_geometry before; + server_geometry after; + int32_t n_ctx_seq_before = 0; + int32_t n_ctx_seq_after = 0; + + bool cold_start = false; + bool rolled_back = false; // the new geometry could not be built; the old one was rebuilt and every state restored + std::string rollback_reason; + + double t_drain_ms = 0.0; // request posted -> every slot at a boundary (and the ones past n_parallel done) + double t_page_out_ms = 0.0; + double t_free_ms = 0.0; + double t_init_ms = 0.0; + double t_page_in_ms = 0.0; + double t_total_ms = 0.0; // the pause serving saw: page out through page in + + size_t n_bytes = 0; // paged out and back + size_t n_bytes_dropped = 0; // cold start: resident states, pages and prompt-cache entries dropped + size_t n_tasks_deferred = 0; // requests that arrived during the pause, queued and resumed after + + std::vector moves; + + json to_json() const; +}; + +// What a reconfigure needs from the server and nothing more; server-context.cpp implements it over +// its private slots and contexts. Slot ids are dense: 0 .. geometry().n_parallel - 1. +struct server_reconfigure_host { + virtual ~server_reconfigure_host() = default; + + virtual server_geometry geometry() const = 0; + virtual int32_t n_ctx_seq() const = 0; // tokens one slot holds now + + virtual llama_context * target() const = 0; + virtual llama_context * draft() const = 0; + virtual common_speculative * speculative() const = 0; + + virtual server_prompt & slot_prompt(int id_slot) = 0; + virtual std::vector & slot_lora(int id_slot) = 0; + virtual bool slot_processing(int id_slot) const = 0; + // ends the slot's request with a typed error and empties the slot + virtual void slot_fail(int id_slot, const std::string & why) = 0; + + // frees the draft and target contexts and the speculative state; the model(s) stay loaded + virtual void free_contexts() = 0; + // builds both contexts for `geo` on the loaded model(s); false leaves no context + virtual bool build_contexts(const server_geometry & geo) = 0; + // `n` slots bound to the current contexts; slots past `n` are idle and go + virtual void set_slots(int n) = 0; + + virtual server_page_store & pages() = 0; + virtual size_t drop_prompt_cache() = 0; // the bytes it held +}; + +// Runs a decided plan: page out, free, build (or roll back to the old geometry), page in. Throws +// server_page_error only before anything is touched. `page_prefix` names pages that stay in the store. +void server_reconfigure_run( + server_reconfigure_host & host, + const server_reconfigure_request & req, + server_reconfigure_plan & plan, + server_reconfigure_receipt & rc, + const std::string & page_prefix); + +struct server_task_result_reconfigure : server_task_result { + server_reconfigure_receipt receipt; + // the routes' view of the server after the reconfigure, built on the task thread + std::shared_ptr meta; + + virtual json to_json() override { + return receipt.to_json(); + } +}; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 746c3a1f98c8..b5a3e319bd82 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1591,6 +1591,34 @@ 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 + }, { + "reconfigure_total", + "Context geometry changes made without reloading the model (/reconfigure)", + (double) metrics.n_reconfigure + }, { + "reconfigure_failed_total", + "/reconfigure requests refused", + (double) metrics.n_reconfigure_failed + }, { + "ctx_checkpoints_evicted_budget_total", + "Context checkpoints dropped to keep within --ctx-checkpoints-bytes", + (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 +1643,18 @@ std::string server_task_result_metrics::to_metrics() { "n_busy_slots_per_decode", "Average number of busy slots per llama_decode() call", (double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0) + }, { + "page_store_bytes", + "Bytes held by the host-RAM page store", + (double) page_store_bytes + }, { + "page_store_pages", + "Pages held by the host-RAM page store", + (double) page_store_pages + }, { + "ctx_checkpoint_bytes", + "Bytes held by every slot's context checkpoints", + (double) ctx_checkpoint_bytes }, }; @@ -1675,6 +1715,34 @@ json server_task_result_slot_save_load::to_json() { }; } +// +// server_task_result_page +// +json server_task_result_page::to_json() { + json res = { + { "id_slot", id_slot }, + { "action", action }, + { "page", page }, + { "gen", gen }, + { "n_tokens", n_tokens }, + { "n_bytes", n_bytes }, + { "n_bytes_draft", n_bytes_draft }, + { "n_bytes_spec", n_bytes_spec }, + { "n_checkpoints", n_checkpoints }, + { "n_bytes_checkpoints", n_bytes_checkpoints }, + { "store_bytes", store_bytes }, + { "store_pages", store_pages }, + { "timings", { + { "ms", t_ms } + }}, + }; + if (!filename.empty()) { + res["filename"] = filename; + res["n_written"] = n_written; + } + return res; +} + // // server_task_result_slot_erase // diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 87a344e174cf..16b31d0f66d5 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -56,6 +56,11 @@ 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_RECONFIGURE, // Continuum fork: new context geometry, same model (server-reconfigure.h) SERVER_TASK_TYPE_GET_LORA, SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, @@ -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; @@ -257,6 +266,10 @@ struct server_task { // used by SERVER_TASK_TYPE_LEARN learn_request learn; + // used by SERVER_TASK_TYPE_RECONFIGURE: the request body, parsed on the task thread against the + // geometry it then has (server_reconfigure_request::parse) + json reconfigure; + server_task() = default; server_task(server_task_type type) : type(type) {} @@ -578,6 +591,9 @@ struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) int n_processing_slots = 0; int n_tasks_deferred = 0; + size_t page_store_bytes = 0; // Continuum fork: the host-RAM page store + size_t page_store_pages = 0; + size_t ctx_checkpoint_bytes = 0; // every slot's context checkpoints server_metrics metrics; @@ -614,6 +630,33 @@ struct server_task_result_slot_save_load : server_task_result { virtual json to_json() override; }; +struct server_page; + +// Continuum fork: the outcome of a PAGE_* task, with the store's totals after it +struct server_task_result_page : server_task_result { + std::string action; + std::string page; + uint64_t gen = 0; + + size_t n_tokens = 0; + size_t n_bytes = 0; // the page: target + draft state + speculative state + checkpoints + size_t n_bytes_draft = 0; + size_t n_bytes_spec = 0; // the speculative tail state + size_t n_checkpoints = 0; + size_t n_bytes_checkpoints = 0; + double t_ms = 0.0; + + size_t store_bytes = 0; + size_t store_pages = 0; + + // PAGE_SPILL: the page the HTTP thread writes; shared, so the task thread never waits on disk + std::shared_ptr spilled; + std::string filename; + size_t n_written = 0; + + virtual json to_json() override; +}; + struct server_task_result_slot_erase : server_task_result { size_t n_erased; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index f449dc00f11e..9eec965e8283 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -220,6 +220,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; routes.post_learn = models_routes->proxy_post; + routes.post_reconfigure = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -275,6 +276,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); ctx_http.post("/learn", ex_wrapper(routes.post_learn)); + ctx_http.post("/reconfigure", ex_wrapper(routes.post_reconfigure)); // Continuum fork: new geometry, same model ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); ctx_http.post("/train/cancel", ex_wrapper(routes.post_train_cancel)); diff --git a/tools/server/tests/unit/test_reconfigure.py b/tools/server/tests/unit/test_reconfigure.py new file mode 100644 index 000000000000..d3be16d51ebc --- /dev/null +++ b/tools/server/tests/unit/test_reconfigure.py @@ -0,0 +1,281 @@ +import os +import platform +import re +import subprocess +import threading +import time + +import pytest +from utils import * + +# Continuum fork: POST /reconfigure changes the context geometry without reloading the model +# (server-reconfigure.h; continuum docs/architecture/SERVING-ENGINE.md, slice S2). The residents' +# states are paged to host RAM, the contexts are rebuilt on the loaded model, and the states go back. + +server = ServerPreset.tinyllama2() + +HYBRID = os.environ.get("LLAMA_SERVER_TEST_HYBRID_GGUF") + +MIND_A = "Once upon a time, a little fox lived at the edge of a dark forest." +MIND_B = "The old lighthouse keeper climbed the stairs every night to light the lamp." +FOLLOW = " One day," + + +def use_model(kind): + global server + server = ServerPreset.tinyllama2() + server.temperature = 0.0 + server.server_metrics = True + server.server_slots = True + server.slot_save_path = "./tmp" + server.n_slots = 2 + server.n_ctx = 1024 + if kind == "hybrid": + if not HYBRID: + pytest.skip("set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF") + server.model_hf_repo = None + server.model_file = HYBRID + server.n_batch = 512 + server.n_ctx = 4096 + + +def complete(prompt, id_slot, n_predict=16, ignore_eos=False): + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, "id_slot": id_slot, "cache_prompt": True, "n_predict": n_predict, + "ignore_eos": ignore_eos, + }, timeout=600) + assert res.status_code == 200, res.body + return res.body + + +def reconfigure(body): + return server.make_request("POST", "/reconfigure", data=body, timeout=600) + + +def props(): + return server.make_request("GET", "/props").body + + +def metric(name): + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + for line in res.body.splitlines(): + if line.startswith(f"llamacpp:{name} "): + return float(line.split()[1]) + raise AssertionError(f"metric {name} missing") + + +def erase(id_slot): + assert server.make_request("POST", f"/slots/{id_slot}?action=erase").status_code == 200 + + +@pytest.mark.parametrize("kind", ["attention", "hybrid"]) +def test_reconfigure_keeps_both_minds_warm_and_identical(kind): + # what this catches: a reconfigure that reloads the model, restores the KV without the slot's + # tokens (the next act re-prefills from 0), restores it corrupt (the output differs from a mind + # that was never reconfigured), or leaves /props reporting the old geometry. + use_model(kind) + server.start() + n_ctx0 = props()["context_geometry"]["allocated_tokens"] + + complete(MIND_A, 0) + complete(MIND_B, 1) + + up = reconfigure({"n_ctx": n_ctx0 * 2, "n_parallel": 3}) + assert up.status_code == 200, up.body + r = up.body + assert r["before"]["n_parallel"] == 2 and r["after"]["n_parallel"] == 3 + assert r["after"]["n_ctx"] >= n_ctx0 * 2 + assert r["n_restored"] == 2 and r["n_kept_in_store"] == 0 and r["n_dropped"] == 0 + assert not r["rolled_back"] and not r["cold_start"] + assert r["n_bytes"] > 0 and all(p["fate"] == "restored" for p in r["pages"]) + for key in ("t_drain_ms", "t_page_out_ms", "t_free_ms", "t_init_ms", "t_page_in_ms", "t_total_ms"): + assert r[key] >= 0 + p = props() + assert p["total_slots"] == 3 + assert p["context_geometry"]["sequence_capacity"] == 3 + assert p["context_geometry"]["allocated_tokens"] == r["after"]["n_ctx"] + assert metric("reconfigure_total") == 1 + + a = complete(MIND_A + FOLLOW, 0) + b = complete(MIND_B + FOLLOW, 1) + c = complete("A third mind arrives.", 2) # the new seat serves + assert c["timings"]["prompt_n"] == c["tokens_evaluated"] + + # back down: slot 2's state has no seat among 2 occupied slots, so it stays a named page + down = reconfigure({"n_ctx": n_ctx0, "n_parallel": 2}) + assert down.status_code == 200, down.body + d = down.body + assert d["after"]["n_parallel"] == 2 and d["n_restored"] == 2 and d["n_kept_in_store"] == 1 + kept = [p for p in d["pages"] if p["fate"] == "kept_in_store"][0] + assert kept["id_slot_before"] == 2 and kept["page"] and kept["reason"] + assert props()["total_slots"] == 2 + a2 = complete(MIND_A + FOLLOW + a["content"] + FOLLOW, 0) + + # the reference: the same acts on slots that were never reconfigured + erase(0) + erase(1) + complete(MIND_A, 0) + complete(MIND_B, 1) + ref_a = complete(MIND_A + FOLLOW, 0) + ref_b = complete(MIND_B + FOLLOW, 1) + ref_a2 = complete(MIND_A + FOLLOW + ref_a["content"] + FOLLOW, 0) + + for got, ref in ((a, ref_a), (b, ref_b), (a2, ref_a2)): + assert got["content"] == ref["content"] + assert got["timings"]["prompt_n"] == ref["timings"]["prompt_n"] + assert got["timings"]["prompt_n"] < got["tokens_evaluated"] # the prefix survived + + # the kept page comes back into a slot like any page + restored = server.make_request("POST", f"/slots/1?action=page_in&page={kept['page']}") + assert restored.status_code == 200, restored.body + assert restored.body["n_tokens"] == kept["n_tokens"] + + +def long_act_model(): + # a request must outlast an HTTP round trip: the 0.8B hybrid decodes ~100x slower than stories260K + use_model("hybrid" if HYBRID else "attention") + server.n_ctx = 8192 + server.start() + return 1500 if not HYBRID else 300 + + +def wait_processing(id_slot): + while not server.make_request("GET", "/slots").body[id_slot]["is_processing"]: + time.sleep(0.005) + + +def test_an_in_flight_request_pauses_and_finishes_identically(): + # what this catches: a reconfigure that drops or corrupts a request mid-generation (its KV, + # sampler or speculative state lost with the freed context), or one that blocks until it ends. + n_predict = long_act_model() + prompt = "Tell me a long story about a dragon." + reference = complete(prompt, 0, n_predict=n_predict, ignore_eos=True) + erase(0) + + out = {} + + def generate(): + out["res"] = complete(prompt, 0, n_predict=n_predict, ignore_eos=True) + + t = threading.Thread(target=generate) + t.start() + wait_processing(0) + # grow, then shrink back, both with the request in flight + caught = 0 + for n_ctx in (16384, 8192): + r = reconfigure({"n_ctx": n_ctx, "n_parallel": 2}) + assert r.status_code == 200, r.body + caught += sum(1 for p in r.body["pages"] if p["in_flight"]) + t.join() + assert caught >= 1, "the generation ended before a reconfigure caught it in flight" + assert out["res"]["content"] == reference["content"] + assert out["res"]["tokens_predicted"] == reference["tokens_predicted"] + + +def test_a_slot_past_the_new_count_drains_and_requests_queue(): + # what this catches: a shrink that kills the request on a removed slot, or a request that + # arrives during the pause and is lost instead of queued. + n_predict = long_act_model() + out = {} + + def generate(): + out["res"] = complete("Tell me a long story about a castle.", 1, n_predict=n_predict, ignore_eos=True) + + def shrink(): + out["shrink"] = reconfigure({"n_parallel": 1}) + + t = threading.Thread(target=generate) + t.start() + wait_processing(1) + s = threading.Thread(target=shrink) + s.start() + time.sleep(0.1) # the shrink is pending: admission is closed + late = complete("A request during the pause.", -1) # queued, then served after + t.join() + s.join() + assert out["res"]["tokens_predicted"] == n_predict # the drained request finished whole + assert out["shrink"].status_code == 200, out["shrink"].body + assert out["shrink"].body["after"]["n_parallel"] == 1 + assert out["shrink"].body["t_drain_ms"] > 0 + assert out["shrink"].body["n_tasks_deferred"] >= 1 + assert late["tokens_predicted"] > 0 + + +def test_refusals_are_typed_and_change_nothing(): + # what this catches: a misspelled field that silently keeps the old geometry, a KV type change + # that carries pages of the wrong format, and a shrink that truncates a mind's state. + use_model("attention") + server.start() + first = complete(MIND_A, 0) + assert reconfigure({"n_slots": 3}).status_code == 400 + assert reconfigure({"n_parallel": 0}).status_code == 400 + assert reconfigure({"cache_type_k": "f32"}).status_code == 400 # needs force + assert reconfigure({"cache_type_k": "nope", "force": True}).status_code == 400 + again = complete(MIND_A, 0) + assert again["timings"]["prompt_n"] == 1 # nothing moved + assert again["content"] == first["content"] + assert metric("reconfigure_failed_total") >= 3 + + # a window smaller than the mind: her state stays whole in the store, never truncated + long_prompt = "The cat sat on the mat and looked at the moon. " * 16 + n_long = len(server.make_request("POST", "/tokenize", data={"content": long_prompt}).body["tokens"]) + assert 280 < n_long < 480, n_long # past one 256-token sequence, inside a 512 slot + complete(long_prompt, 0) + small = reconfigure({"n_ctx": 256, "n_parallel": 1}) + assert small.status_code == 200, small.body + kept = small.body["pages"][0] + assert kept["fate"] == "kept_in_store" and "do not fit" in kept["reason"] + assert reconfigure({"n_ctx": 1024, "n_parallel": 2}).status_code == 200 + back = server.make_request("POST", f"/slots/0?action=page_in&page={kept['page']}") + assert back.status_code == 200 and back.body["n_tokens"] == kept["n_tokens"] + + # a geometry libllama cannot build rolls back with every mind intact (a quantized V cache + # needs flash attention, which this model's head size does not get) + fail = reconfigure({"cache_type_k": "q8_0", "cache_type_v": "q8_0", "force": True}) + assert fail.status_code == 200, fail.body + assert fail.body["rolled_back"] and not fail.body["cold_start"] and fail.body["n_restored"] == 1 + assert fail.body["after"]["cache_type_v"] == "f16" + warm = complete(long_prompt, 0) + assert warm["timings"]["prompt_n"] == 1 + + # with force, a KV type change is an intentional cold start, counted (f32: this model's + # head size is not a multiple of a quantized block) + cold = reconfigure({"cache_type_k": "f32", "force": True}) + assert cold.status_code == 200, cold.body + assert cold.body["cold_start"] and cold.body["n_dropped"] == 1 and cold.body["n_bytes_dropped"] > 0 + assert cold.body["after"]["cache_type_k"] == "f32" + fresh = complete(MIND_A, 0) + assert fresh["timings"]["prompt_n"] == fresh["tokens_evaluated"] + + +CYCLES = int(os.environ.get("LLAMA_SERVER_TEST_RECONFIGURE_CYCLES", "0")) + + +def footprint_bytes(pid): + # macOS: phys_footprint, what the OS charges the process (RSS undercounts GPU-mapped memory) + out = subprocess.run(["footprint", "-p", str(pid)], capture_output=True, text=True).stdout + m = re.search(r"phys_footprint:\s*([\d.]+)\s*([KMG]B)", out) + assert m, out + return float(m.group(1)) * {"KB": 2**10, "MB": 2**20, "GB": 2**30}[m.group(2)] + + +@pytest.mark.skipif(not CYCLES or platform.system() != "Darwin", reason="set LLAMA_SERVER_TEST_RECONFIGURE_CYCLES (macOS footprint)") +def test_footprint_is_flat_over_many_cycles(): + # what this catches: a reconfigure that leaks a context, a draft context, its compute buffers + # or the paged copies each cycle (the plan line 3b gate: flat over 100 resize cycles). + use_model("hybrid" if HYBRID else "attention") + server.start() + n_ctx0 = props()["context_geometry"]["allocated_tokens"] + complete(MIND_A, 0) + complete(MIND_B, 1) + samples = [] + for i in range(CYCLES): + assert reconfigure({"n_ctx": n_ctx0 * 2, "n_parallel": 3}).status_code == 200 + assert reconfigure({"n_ctx": n_ctx0, "n_parallel": 2}).status_code == 200 + samples.append(footprint_bytes(server.process.pid)) + print("footprint MB per cycle:", [round(s / 2**20, 1) for s in samples]) + # after warm-up, the last cycle's footprint is within 2% of the tenth's + assert samples[-1] <= samples[min(9, len(samples) - 1)] * 1.02 + again = complete(MIND_A + FOLLOW, 0) + assert again["timings"]["prompt_n"] < again["tokens_evaluated"] diff --git a/tools/server/tests/unit/test_slot_page.py b/tools/server/tests/unit/test_slot_page.py new file mode 100644 index 000000000000..069aff494fa7 --- /dev/null +++ b/tools/server/tests/unit/test_slot_page.py @@ -0,0 +1,251 @@ +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from utils import * + +# Continuum fork: a slot's sequence state paged to a named host-RAM page and back +# (server-pages.h; continuum docs/architecture/SERVING-ENGINE.md, slice S1). + +server = ServerPreset.tinyllama2() + +PROMPT = "What is the capital of France?" +FOLLOW = "What is the capital of France? And what is the capital of Germany?" + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.slot_save_path = "./tmp" + server.temperature = 0.0 + server.server_metrics = True + + +def complete(prompt, id_slot): + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, "id_slot": id_slot, "cache_prompt": True, "n_predict": 16, + }) + assert res.status_code == 200 + return res.body + + +def page(id_slot, action, name, body=None): + return server.make_request("POST", f"/slots/{id_slot}?action={action}&page={name}", data=body or {}) + + +def metric(name): + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + for line in res.body.splitlines(): + if line.startswith(f"llamacpp:{name} "): + return float(line.split()[1]) + raise AssertionError(f"metric {name} missing") + + +def test_page_out_in_is_identical_to_never_paging(): + # what this catches: a page_in that restores the KV but not the slot's tokens (the next + # prompt re-prefills from 0), or restores them into a corrupt state (the output differs from + # a slot that was never paged). The page goes to ANOTHER slot, and the source slot is reused + # in between, so the page cannot be served by KV that happened to stay in place. + global server + server.start() + assert server.make_request("GET", "/props").body["slot_page_version"] == 1 + + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200, out.body + n_tokens = out.body["n_tokens"] + assert n_tokens == first["timings"]["prompt_n"] + first["timings"]["predicted_n"] - 1 + assert out.body["n_bytes"] > 0 and out.body["store_pages"] == 1 + assert out.body["store_bytes"] == out.body["n_bytes"] + assert metric("page_store_bytes") == out.body["n_bytes"] + + # the page left the slot: its KV is free and its cached tokens are gone + scrub = complete("Once upon a time", 1) + assert scrub["timings"]["prompt_n"] == scrub["tokens_evaluated"] + + restored = page(0, "page_in", "mind-a") + assert restored.status_code == 200, restored.body + assert restored.body["n_tokens"] == n_tokens + assert restored.body["store_pages"] == 0 and restored.body["store_bytes"] == 0 + paged = complete(FOLLOW, 0) + + # the reference: the same two requests on a slot that never paged + assert server.make_request("POST", "/slots/1?action=erase").status_code == 200 + complete(PROMPT, 1) + reference = complete(FOLLOW, 1) + + assert paged["timings"]["prompt_n"] == reference["timings"]["prompt_n"] + assert paged["timings"]["prompt_n"] < paged["tokens_evaluated"] # the prefix came from the page + assert paged["content"] == reference["content"] + assert metric("page_out_total") == 1 and metric("page_in_total") == 1 + + +def test_page_errors_are_typed_and_leave_state_intact(): + # what this catches: a missing or consumed page answered as a silent cold start, or a + # failed page_out that frees the slot anyway. + global server + server.start() + assert page(0, "page_in", "nobody").status_code == 404 + assert page(0, "page_out", "empty").status_code == 400 # nothing in the slot + assert page(0, "page_out", "bad%01name").status_code == 400 + complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 200 + assert page(1, "page_in", "mind-a").status_code == 200 + assert page(0, "page_in", "mind-a").status_code == 404 # a page_in consumes the page + assert metric("page_failed_total") >= 3 + + +def test_page_ram_ceiling_refuses_and_keeps_the_slot(): + # what this catches: a page store that grows past its ceiling on unified memory, or a + # refused page_out that still drops the slot's sequence. + global server + server.page_ram_bytes = 1 + server.start() + first = complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 503 + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 # the slot still holds the sequence + assert again["content"] == first["content"] + + +def test_spill_writes_the_slot_save_format(): + # what this catches: a spilled page that today's restore cannot read, a spill that blocks + # on a page replaced meanwhile, or a drop that leaves the RAM bytes counted. + global server + server.start() + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200 + spilled = page(1, "spill", "mind-a", {"filename": "mind-a.bin", "drop": True}) + assert spilled.status_code == 200, spilled.body + files = [Path(server.slot_save_path) / ("mind-a.bin" + s) for s in ("", ".ckpt", ".lora", ".dft")] + assert all(f.exists() for f in files) + assert spilled.body["n_written"] == sum(f.stat().st_size for f in files) + assert spilled.body["store_pages"] == 0 and metric("page_store_bytes") == 0 + + res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "mind-a.bin"}) + assert res.status_code == 200 + assert res.body["n_restored"] == out.body["n_tokens"] + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 + assert again["content"] == first["content"] + + # a spill names a filename like save does + assert page(1, "spill", "mind-a", {"filename": "../x.bin"}).status_code == 400 + + +def test_spill_publishes_whole_and_restore_cleans_a_crashed_one(): + # what this catches: a spill that leaves a stage beside the page dir, or a restore that reads + # a generation a crashed publication left uncommitted (the committed page must be what restores, + # and the leftovers must go). + global server + server.start() + first = complete(PROMPT, 1) + out = page(1, "page_out", "mind-a") + assert out.status_code == 200 + assert page(1, "spill", "mind-a", {"filename": "mind-a.bin", "drop": True}).status_code == 200 + staging = Path(server.slot_save_path) / "..page-staging" + assert not staging.exists() or not any(staging.iterdir()) + + # a server that died between writing its files and committing them: its pid is gone + dead = subprocess.Popen([sys.executable, "-c", "pass"]) + dead.wait() + stage = staging / f"{dead.pid}-1" + stage.mkdir(parents=True) + for suffix in ("", ".ckpt", ".lora", ".dft"): + (stage / ("mind-a.bin" + suffix)).write_bytes(b"torn") + + res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "mind-a.bin"}) + assert res.status_code == 200, res.body + assert res.body["n_restored"] == out.body["n_tokens"] + assert not stage.exists() + again = complete(PROMPT, 0) + assert again["timings"]["prompt_n"] == 1 + assert again["content"] == first["content"] + + +HYBRID = os.environ.get("LLAMA_SERVER_TEST_HYBRID_GGUF") + + +@pytest.mark.skipif(not HYBRID, reason="set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF") +def test_checkpoints_stay_within_the_byte_budget_and_page_with_the_slot(): + # what this catches: checkpoints counted instead of measured (a hybrid's checkpoints are the + # hidden RAM cost), and a page that drops a slot's checkpoints (the next divergence would + # re-prefill from the head). + global server + server.model_hf_repo = None + server.model_file = HYBRID + server.n_ctx = 4096 + server.n_batch = 512 + server.n_slots = 2 + server.start() + long_prompt = " ".join(f"line {i} of the story continues." for i in range(120)) + complete(long_prompt, 0) + per_slot = metric("ctx_checkpoint_bytes") + assert per_slot > 0 + out = page(0, "page_out", "mind-a") + assert out.status_code == 200 and out.body["n_checkpoints"] > 0 + assert metric("ctx_checkpoint_bytes") == 0 # they left with the page + restored = page(0, "page_in", "mind-a") + assert restored.body["n_checkpoints"] == out.body["n_checkpoints"] + server.stop() + + budget = int(per_slot) + int(per_slot) // 2 # room for one slot's set and half another's + server.ctx_checkpoints_bytes = budget + server.start() + complete(long_prompt, 0) + complete("A different tale: " + long_prompt, 1) + assert metric("ctx_checkpoint_bytes") <= budget + assert metric("ctx_checkpoints_evicted_budget_total") > 0 + + +@pytest.mark.skipif(not HYBRID, reason="set LLAMA_SERVER_TEST_HYBRID_GGUF to a small hybrid (recurrent) GGUF with an MTP head") +def test_page_with_draft_state_fits_its_ceiling_exactly(): + # what this catches: a page_out admitted on a count that leaves out part of what the page holds + # (draft or speculative state), so the store passes --page-ram-bytes after a successful admission. + # At a ceiling of exactly the page's bytes it is stored; one byte less and it is refused, the slot kept. + global server + + def mtp_server(ceiling): + global server + server = ServerPreset.tinyllama2() + server.slot_save_path = "./tmp" + server.temperature = 0.0 + server.server_metrics = True + server.model_hf_repo = None + server.model_file = HYBRID + server.n_ctx = 4096 + server.n_batch = 512 + server.spec_type = "draft-mtp" + server.page_ram_bytes = ceiling + server.start() + + mtp_server(None) + complete(PROMPT, 0) + out = page(0, "page_out", "mind-a") + assert out.status_code == 200, out.body + n_bytes = out.body["n_bytes"] + assert out.body["n_bytes_draft"] > 0 # the draft (MTP) state rides in the page + assert n_bytes == out.body["store_bytes"] + assert metric("page_store_bytes") == pytest.approx(n_bytes, rel=1e-5) # the gauge prints 6 digits + server.stop() + + mtp_server(n_bytes - 1) + first = complete(PROMPT, 0) + assert page(0, "page_out", "mind-a").status_code == 503 + again = complete(PROMPT, 0) + # the refused page_out kept the slot's sequence: a hybrid resumes from its kept checkpoint + # (a cleared slot would evaluate the whole prompt) + assert again["timings"]["prompt_n"] < again["tokens_evaluated"] + assert again["content"] == first["content"] + server.stop() + + mtp_server(n_bytes) + complete(PROMPT, 0) + out = page(0, "page_out", "mind-a") + assert out.status_code == 200, out.body + assert out.body["n_bytes"] == n_bytes == out.body["store_bytes"] diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c591..7bc7215acf29 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -121,6 +121,8 @@ class ServerProcess: mcp_servers_config: str | None = None mcp_servers_json: str | None = None cors_origins: str | None = None + ctx_checkpoints_bytes: int | None = None + page_ram_bytes: int | None = None # session variables process: subprocess.Popen | None = None @@ -285,6 +287,10 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--mcp-servers-json", self.mcp_servers_json]) if self.backend_sampling: server_args.append("--backend_sampling") + if self.ctx_checkpoints_bytes is not None: + server_args.extend(["--ctx-checkpoints-bytes", self.ctx_checkpoints_bytes]) + if self.page_ram_bytes is not None: + server_args.extend(["--page-ram-bytes", self.page_ram_bytes]) if self.gcp_compat: env["AIP_MODE"] = "PREDICTION"