diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index 6d74e8bfd3af..bfbe23131677 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -121,11 +121,23 @@ jobs: - name: Python setup id: setup_python + # also after a failed step above, so the reconfigure tests below still run (see their note) + if: ${{ !cancelled() }} uses: actions/setup-python@v6 with: python-version: '3.11' pip-install: -r tools/server/tests/requirements.txt + # Continuum fork: the live reconfigure's validator (SERVING-ENGINE.md, S2), run on its own while + # pre-existing residency failures (the dream step above, unit/test_lora.py) stop the suite under -x + # before it; drop this step once the suite reaches test_reconfigure.py again + - name: Reconfigure tests + if: ${{ !cancelled() }} + run: | + cd tools/server/tests + ./tests.sh -v -x unit/test_reconfigure.py + LLAMA_ARG_BACKEND_SAMPLING=1 ./tests.sh -v -x unit/test_reconfigure.py + - name: Tests id: server_integration_tests run: | @@ -205,6 +217,14 @@ jobs: python-version: '3.11' pip-install: -r tools/server/tests/requirements.txt + # Continuum fork: the live reconfigure's validator, on its own (see the ubuntu job's note) + - name: Reconfigure tests + shell: bash + run: | + cd tools/server/tests + export PYTHONIOENCODING=":replace" + ./tests.sh -v -x unit/test_reconfigure.py + - name: Tests id: server_integration_tests shell: bash diff --git a/common/common.cpp b/common/common.cpp index 09fe1b77e3d7..6a092406d780 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1278,11 +1278,13 @@ struct common_init_result::impl { // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top common_threadpools threadpools; + bool threadpools_made = false; // init() runs once; a context created later is attached to them llama_model_ptr model; llama_context_ptr context; std::vector lora; + bool adapters_failed = false; // an adapter did not load: no context is created on this model std::vector samplers; std::vector samplers_seq_config; @@ -1291,9 +1293,10 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { auto mparams = common_model_params_to_llama(params); - auto cparams = common_context_params_to_llama(params); if (params.fit_params) { + auto cparams = common_context_params_to_llama(params); + COM_TRC("%s", "fitting params to device memory ...\n"); COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n"); @@ -1326,6 +1329,10 @@ common_init_result::common_init_result(common_params & params, bool model_only) params.fit_params_min_ctx, has_draft || spec_mtp ? &extra : nullptr, params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + + // Continuum fork: the context size fit chose becomes the params' own (fit changes only the + // context's n_ctx, common/fit.cpp), so init_context() creates the context it fitted + params.n_ctx = cparams.n_ctx; } llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); @@ -1352,6 +1359,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) lora.reset(llama_adapter_lora_init(model, la.path.c_str())); if (lora == nullptr) { COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str()); + pimpl->adapters_failed = true; return; } @@ -1387,8 +1395,23 @@ common_init_result::common_init_result(common_params & params, bool model_only) params.sampling.logit_bias.end(), params.sampling.logit_bias_eog.begin(), params.sampling.logit_bias_eog.end()); } +} + +bool common_init_result::init_context(common_params & params, bool warmup) { + GGML_ASSERT(pimpl->model && "init_context needs a loaded model"); + GGML_ASSERT(!pimpl->context && "free_context() before creating another context"); + + if (pimpl->adapters_failed) { + COM_ERR("%s", "an adapter failed to load, no context is created on this model\n"); + return false; + } + + llama_model * model = pimpl->model.get(); + auto cparams = common_context_params_to_llama(params); // init the backend samplers as part of the context creation + pimpl->samplers.clear(); + pimpl->samplers_seq_config.clear(); pimpl->samplers.resize(cparams.n_seq_max); pimpl->samplers_seq_config.resize(cparams.n_seq_max); @@ -1405,59 +1428,21 @@ common_init_result::common_init_result(common_params & params, bool model_only) llama_context * lctx = llama_init_from_model(model, cparams); if (lctx == NULL) { COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); - return; + return false; } pimpl->context.reset(lctx); set_process_priority(params.cpuparams.priority); - pimpl->threadpools.init(lctx, params); -} - -llama_model * common_init_result::model() { - return pimpl->model.get(); -} - -llama_context * common_init_result::context() { - return pimpl->context.get(); -} - -common_sampler * common_init_result::sampler(llama_seq_id seq_id) { - if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) { - return nullptr; - } - return pimpl->samplers[seq_id].get(); -} - -void common_init_result::reset_samplers() { - for (int i = 0; i < (int) pimpl->samplers.size(); ++i) { - llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get())); - } -} - -std::vector & common_init_result::lora() { - return pimpl->lora; -} - -common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { - common_init_result_ptr res(new common_init_result(params, model_only)); - - llama_model * model = res->model(); - if (model == NULL) { - COM_ERR("failed to load model '%s'\n", params.model.path.c_str()); - return res; - } - - if (model_only) { - return res; + if (pimpl->threadpools_made) { + pimpl->threadpools.attach(lctx); + } else { + pimpl->threadpools.init(lctx, params); + pimpl->threadpools_made = true; } - llama_context * lctx = res->context(); - if (lctx == NULL) { - COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); - return res; - } + // the context's own setup below; a failure in it leaves the context created, as it always has const llama_vocab * vocab = llama_model_get_vocab(model); @@ -1472,7 +1457,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode const auto cvec = common_control_vector_load(params.control_vectors); if (cvec.n_embd == -1) { - return res; + return true; } int err = llama_set_adapter_cvec( @@ -1483,7 +1468,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode params.control_vector_layer_start, params.control_vector_layer_end); if (err) { - return res; + return true; } } @@ -1507,7 +1492,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode } if (!ok) { - return res; + return true; } } @@ -1515,7 +1500,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode common_set_adapter_lora(lctx, params.lora_adapters); } - if (params.warmup) { + if (warmup) { COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n"); std::vector tmp; @@ -1550,7 +1535,67 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode llama_perf_context_reset(lctx); // reset samplers to reset RNG state after warmup to the seeded state - res->reset_samplers(); + reset_samplers(); + } + + return true; +} + +void common_init_result::free_context() { + pimpl->context.reset(); + // the backend sampler configs were the freed context's + pimpl->samplers.clear(); + pimpl->samplers_seq_config.clear(); +} + +llama_model * common_init_result::model() { + return pimpl->model.get(); +} + +llama_context * common_init_result::context() { + return pimpl->context.get(); +} + +common_sampler * common_init_result::sampler(llama_seq_id seq_id) { + if (seq_id < 0 || seq_id >= (int) pimpl->samplers.size()) { + return nullptr; + } + return pimpl->samplers[seq_id].get(); +} + +void common_init_result::reset_samplers() { + for (int i = 0; i < (int) pimpl->samplers.size(); ++i) { + llama_sampler_reset(common_sampler_get(pimpl->samplers[i].get())); + } +} + +std::vector & common_init_result::lora() { + return pimpl->lora; +} + +common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { + common_init_result_ptr res(new common_init_result(params, model_only)); + + if (res->model() == NULL) { + COM_ERR("failed to load model '%s'\n", params.model.path.c_str()); + return res; + } + + if (model_only) { + return res; + } + + // on failure the context stays null, which is what callers check + res->init_context(params, params.warmup); + + return res; +} + +common_init_result_ptr common_load_from_params(common_params & params) { + common_init_result_ptr res(new common_init_result(params, false)); + + if (res->model() == NULL) { + COM_ERR("failed to load model '%s'\n", params.model.path.c_str()); } return res; @@ -1833,6 +1878,13 @@ void common_threadpools::init(llama_context * ctx, const common_params & params) llama_attach_threadpool(ctx, threadpool, threadpool_batch); } +void common_threadpools::attach(llama_context * ctx) const { + if (threadpool == nullptr) { + return; // init() made none (no CPU backend, or the pools failed): the context keeps its own + } + llama_attach_threadpool(ctx, threadpool, threadpool_batch); +} + // // Batch utils // diff --git a/common/common.h b/common/common.h index f57bafb83295..83ca18d0c4e1 100644 --- a/common/common.h +++ b/common/common.h @@ -921,9 +921,18 @@ struct common_sampler; // note: defines the model, context, samplers, ets. lifetimes struct common_init_result { + // loads the model and, unless model_only, its adapters; the context is init_context()'s common_init_result(common_params & params, bool model_only = false); ~common_init_result(); + // Continuum fork: the context is its own lifetime, apart from the model's (SERVING-ENGINE.md, + // S2). init_context() creates it on the loaded model with params' geometry and runs its setup + // (control vectors, adapters, optionally the warmup); free_context() frees it and what is bound + // to it. The model, its mmap, its adapters and the CPU threadpools stay, so a server changes its + // geometry without loading the weights again. Returns false if no context could be created. + bool init_context(common_params & params, bool warmup); + void free_context(); + llama_model * model(); llama_context * context(); @@ -939,8 +948,13 @@ struct common_init_result { using common_init_result_ptr = std::unique_ptr; +// the model and its context (model_only: the model alone) common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false); +// Continuum fork: the model and its adapters with no context yet; the caller creates it with +// init_context(), and may free and create it again with another geometry +common_init_result_ptr common_load_from_params(common_params & params); + struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); @@ -968,6 +982,9 @@ struct common_threadpools { void init(llama_context * ctx, const common_params & params); + // Continuum fork: the pools already made, attached to a context created after them + 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..08d1c402bd64 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2375,17 +2375,54 @@ struct common_speculative_init_result::impl { llama_context_ptr context; }; +common_speculative_init_result::common_speculative_init_result(common_params & params) : + pimpl(new impl{}) { + if (!params.speculative.has_dft()) { + return; // an MTP context runs on the target model's weights: no draft model to load + } + + auto mparams = common_model_params_to_llama(params); + + LOG_INF("%s: loading draft model '%s'\n", __func__, params.speculative.draft.mparams.path.c_str()); + + llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); + if (model_dft == NULL) { + LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, params.speculative.draft.mparams.path.c_str()); + return; + } + + pimpl->model.reset(model_dft); +} + common_speculative_init_result::common_speculative_init_result( common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) : - pimpl(new impl{}) { + common_speculative_init_result(params) { + if (params.speculative.has_dft() && !pimpl->model) { + return; + } + init_context(params, model_tgt, ctx_tgt); +} + +bool common_speculative_init_result::init_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + GGML_ASSERT(!pimpl->context && "free_context() before creating another draft context"); + 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); + if (!has_draft && !spec_mtp) { + return true; // no draft context asked for + } + + llama_model * model = has_draft ? pimpl->model.get() : model_tgt; + if (model == nullptr) { + LOG_ERR("%s: no draft model is loaded\n", __func__); + return false; + } + auto cparams = common_context_params_to_llama(params); if (spec_mtp) { @@ -2400,39 +2437,23 @@ common_speculative_init_result::common_speculative_init_result( cparams.n_rs_seq = 0; cparams.ctx_other = ctx_tgt; - std::string model_path; - if (has_draft) { - model_path = params.speculative.draft.mparams.path; - LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str()); - - llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); - if (model_dft == NULL) { - LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); - return; - } - - pimpl->model.reset(model_dft); - - llama_context * ctx_dft = llama_init_from_model(model_dft, cparams); - if (ctx_dft == nullptr) { - LOG_ERR("%s: failed to create MTP context\n", __func__); - return; - } + if (!has_draft) { + LOG_INF("%s: creating MTP draft context against the target model '%s'\n", __func__, params.model.path.c_str()); + } - pimpl->context.reset(ctx_dft); - } else if (spec_mtp) { - model_path = params.model.path; + llama_context * ctx_dft = llama_init_from_model(model, cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return false; + } - LOG_INF("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str()); + pimpl->context.reset(ctx_dft); - llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams); - if (ctx_dft == nullptr) { - LOG_ERR("%s: failed to create MTP context\n", __func__); - return; - } + return true; +} - pimpl->context.reset(ctx_dft); - } +void common_speculative_init_result::free_context() { + pimpl->context.reset(); } common_speculative_init_result::~common_speculative_init_result() = default; diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de59..ffd190648e84 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -91,8 +91,17 @@ typedef std::unique_ptr common_s struct common_speculative_init_result { common_speculative_init_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + // Continuum fork: the draft model alone (loaded when params name one); the draft/MTP context is + // init_context()'s, so it can be freed and created again with another geometry while the draft + // model, and the target model an MTP context runs on, stay loaded (SERVING-ENGINE.md, S2) + explicit common_speculative_init_result(common_params & params); ~common_speculative_init_result(); + // creates the draft/MTP context for ctx_tgt's geometry; false if params ask for one and it could + // not be created + bool init_context(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + void free_context(); + llama_model * model(); llama_context * context(); diff --git a/examples/training/finetune-lora.cpp b/examples/training/finetune-lora.cpp index 1f97b2bcde03..11927a9899bc 100644 --- a/examples/training/finetune-lora.cpp +++ b/examples/training/finetune-lora.cpp @@ -136,6 +136,11 @@ int main(int argc, char ** argv) { /*recompute =*/false, // keep every layer's intermediates, as before /*walk_exact =*/false, // the per-chunk walk, as before /*walk_horizon =*/0, + // no caller-owned allocation authority: the example allocates as it always has + /*static_admission =*/nullptr, + /*static_admission_data =*/nullptr, + /*workspace_admission =*/nullptr, + /*workspace_admission_data =*/nullptr, }; llama_opt_init(ctx, model, lopt_params); diff --git a/examples/training/finetune.cpp b/examples/training/finetune.cpp index c19e5f82dd19..f15a373b6489 100644 --- a/examples/training/finetune.cpp +++ b/examples/training/finetune.cpp @@ -77,6 +77,11 @@ int main(int argc, char ** argv) { /*recompute =*/false, // keep every layer's intermediates, as before /*walk_exact =*/false, // the per-chunk walk, as before /*walk_horizon =*/0, + // no caller-owned allocation authority: the example allocates as it always has + /*static_admission =*/nullptr, + /*static_admission_data =*/nullptr, + /*workspace_admission =*/nullptr, + /*workspace_admission_data =*/nullptr, }; llama_opt_init(ctx, model, lopt_params); diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index 67fb03964d7c..83f06b13e272 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -14,6 +14,11 @@ add_library(${TARGET} STATIC server-common.cpp server-common.h server-context.cpp + server-geometry.h + server-reconfigure.cpp + server-reconfigure.h + server-slot-page.cpp + server-slot-page.h server-train.cpp server-context.h server-stream.cpp diff --git a/tools/server/server-context-learn-checkpoint.inc b/tools/server/server-context-learn-checkpoint.inc index d7bfa4e3f480..2dc0096fa2b4 100644 --- a/tools/server/server-context-learn-checkpoint.inc +++ b/tools/server/server-context-learn-checkpoint.inc @@ -42,8 +42,10 @@ gguf_get_val_i64(metadata.get(), iteration) <= 0 || gguf_get_val_i64(metadata.get(), iteration) - 1 != steps || std::string(gguf_get_val_str(metadata.get(), bound)) != binding) { return refuse("optimizer generation mismatch"); } - std::unique_ptr weights(fopen((path + "/adapter.gguf").c_str(), "rb"), fclose); - std::unique_ptr optimizer(fopen((path + "/optimizer.gguf").c_str(), "rb"), fclose); + // a deleter type, not decltype(&fclose): fclose's attributes are dropped from a template argument (-Wignored-attributes) + struct file_closer { void operator()(FILE * f) const { fclose(f); } }; + std::unique_ptr weights(fopen((path + "/adapter.gguf").c_str(), "rb")); + std::unique_ptr optimizer(fopen((path + "/optimizer.gguf").c_str(), "rb")); if (!weights || !optimizer) { return refuse("checkpoint components unreadable"); } const auto started = ggml_time_us(); auto held = learn_ledger.hold_restoration(request.grant, started); diff --git a/tools/server/server-context-learn.inc b/tools/server/server-context-learn.inc index 2b2a0aa08099..23adc3529068 100644 --- a/tools/server/server-context-learn.inc +++ b/tools/server/server-context-learn.inc @@ -392,8 +392,17 @@ const bool begun = llama_opt_learn_begin(ctx_tgt, slot.id, n0, toks.data(), labels.data(), (int32_t) toks.size(), served_tokens.data() + n0, (int32_t) served_tokens.size() - n0, result); common_set_adapter_lora(ctx_tgt, slot.lora); - learning = std::unique_ptr(new pending_learn{ &slot, cp, std::move(set), result, - toks[0], L.n_prefix, (int32_t) toks.size() - 1, loras[L.lora_id].ptr, t0 }); + // value-initialized, then named field by field: every member a learn does not set keeps its default + learning = std::make_unique(); + learning->slot = &slot; + learning->checkpoint = cp; + learning->adapters = std::move(set); + learning->result = result; + learning->last_prompt = toks[0]; + learning->prefix = L.n_prefix; + learning->n_labels = (int32_t) toks.size() - 1; + learning->shadow = loras[L.lora_id].ptr; + learning->started = t0; learning->began = begun; learning->grant = L.grant; learning->lora_id = L.lora_id; @@ -582,3 +591,41 @@ slot.learned = learn_on_slot(slot, L, ggml_time_us()); } + + // Continuum fork: what of learning cannot cross a live reconfigure (server-reconfigure.h), owned here + // beside the learn duties. A learn in progress holds its graph and accumulators in the serving context + // until its restoration ends. Each shadow's optimizer (its moments) lives in the context as well, and + // the fork moves it only through llama_opt_shadow_state_write and, on the other side, + // llama_opt_prepare_restore + llama_opt_state_restore, whose preparation needs the shadow's retained + // labelled experience on a fresh context: a reconfigure holds no such experience, so the moments + // cannot be paged with the slots' state and the rebuild is refused while any exist. + std::optional learning_blocks_reconfigure() const { + if (learning) { + const int32_t id = learning->slot->id; + return server_reconfigure_refused { SERVER_RECONFIGURE_LEARN_IN_PROGRESS, + "a learn is in progress on slot " + std::to_string(id) + "; retry when its restoration finishes", id }; + } + struct held_optimizers { + size_t bytes = 0; + std::set shadows; + } held; + llama_opt_visit_memory(ctx_tgt, [](const llama_adapter_lora * shadow, const char *, const char *, const char *, + size_t bytes, bool, void * user_data) { + auto & h = *static_cast(user_data); + h.bytes += bytes; + if (bytes > 0) { + h.shadows.insert(shadow); + } + }, &held); + if (held.bytes > 0) { + return server_reconfigure_refused { SERVER_RECONFIGURE_LEARNING_BOUND, + "the serving context holds the optimizer of " + std::to_string(held.shadows.size()) + " shadow(s), " + + std::to_string(held.bytes >> 20) + " MiB; their moments do not cross a rebuild (snapshot and unload them first)", -1 }; + } + if (!learning_workspace_residency.empty()) { + return server_reconfigure_refused { SERVER_RECONFIGURE_LEARNING_BOUND, + "the training workspace still replaces the serving scheduler (residency " + learning_workspace_residency + + "); it is reclaimed when the last shadow is released", -1 }; + } + return std::nullopt; + } diff --git a/tools/server/server-context-reconfigure.inc b/tools/server/server-context-reconfigure.inc new file mode 100644 index 000000000000..7254cbfcc406 --- /dev/null +++ b/tools/server/server-context-reconfigure.inc @@ -0,0 +1,193 @@ +// Members of server_context_impl (included in its body): what a live reconfigure needs of the server +// (server-reconfigure.h). The logic is server-reconfigure.cpp's; this binds it to the server's context +// lifetimes (destroy_contexts / build_contexts) and its slots. + // LIVE RECONFIGURE (server-reconfigure.h): the weights loaded since start (a reconfigure never adds + // one), the parked slots' host state (their pages are reconfigure_state's), and the geometry as + // /props reports it, republished by the main loop after every load and reconfigure + int64_t n_weight_loads = 0; + server_reconfigure_state reconfigure_state; + std::map parked_slots; + mutable std::mutex geometry_report_mutex; + json geometry_report = json::object(); + std::atomic n_seq_live{0}; + std::atomic n_ctx_seq_live{0}; + + bool kv_pool_unified() const { + // measured from the context, not the flag: one pool when a sequence may hold all of it + return llama_n_seq_max(ctx_tgt) <= 1 || llama_n_ctx_seq(ctx_tgt) >= llama_n_ctx(ctx_tgt); + } + + void publish_geometry() { + json parked = json::array(); + for (const auto & [id, p] : reconfigure_state.parked) { + parked.push_back({ { "slot", id }, { "tokens_held", p.tokens_held }, { "busy", p.busy } }); + } + json report = { + { "n_ctx", llama_n_ctx(ctx_tgt) }, + { "n_seq", (int32_t) slots.size() }, + { "n_ctx_seq", get_slot_n_ctx() }, + { "kv_unified", kv_pool_unified() }, + { "n_weight_loads", n_weight_loads }, + { "n_applied", reconfigure_state.n_applied }, + { "n_refused", reconfigure_state.n_refused }, + { "parked", parked }, + // the same keys /props' load-time context_geometry has, read from the live context + { "context_geometry", { + { "allocated_tokens", llama_n_ctx(ctx_tgt) }, + { "max_tokens_per_sequence", llama_n_ctx_seq(ctx_tgt) }, + { "sequence_capacity", llama_n_seq_max(ctx_tgt) }, + } }, + { "last", reconfigure_state.last ? server_reconfigure_to_json(*reconfigure_state.last) : json(nullptr) }, + { "cost", [&] { + json cost = server_reconfigure_cost_json(reconfigure_state); + // where the paged state lives now: the contexts' own buffers, per buffer type (the backend's + // name), so the price is read per backend + json buffers = json::array(); + for (const auto & [buft, mb] : llama_get_memory_breakdown(ctx_tgt)) { + if (mb.context > 0) { + buffers.push_back({ { "buffer_type", ggml_backend_buft_name(buft) }, { "context_bytes", mb.context } }); + } + } + cost["context_buffers"] = buffers; + return cost; + }() }, + }; + n_seq_live.store((int32_t) slots.size()); + n_ctx_seq_live.store(get_slot_n_ctx()); + std::lock_guard lock(geometry_report_mutex); + geometry_report = std::move(report); + } + + bool slot_id_is_parked(int id) const { + return parked_slots.find(id) != parked_slots.end(); + } + + // a slot leaves service whole and a fresh, empty slot takes its id + void park_slot(int id) { + GGML_ASSERT(parked_slots.find(id) == parked_slots.end()); + parked_slots.emplace(id, std::move(slots[id])); + slots[id] = server_slot(); + bind_slot(slots[id]); + init_slot(slots[id], id); + } + + // the parked slot takes its id back (the reconfigure pages its KV in) + void unpark_slot(int id) { + auto it = parked_slots.find(id); + GGML_ASSERT(it != parked_slots.end()); + slots[id] = std::move(it->second); + parked_slots.erase(it); + bind_slot(slots[id]); + } + + // a parked slot's work ends without its state ever returning: its client is told why. Its sampler + // was attached to a context since freed, so it is dropped, never released (release resets it there). + void drop_parked_slot(int id, const std::string & why) { + auto it = parked_slots.find(id); + if (it == parked_slots.end()) { + return; + } + if (it->second.task && !why.empty()) { + send_error(it->second, why, ERROR_TYPE_UNAVAILABLE); + } + parked_slots.erase(it); + reconfigure_state.parked.erase(id); + } + + void drop_parked(const std::string & why) { + while (!parked_slots.empty()) { + drop_parked_slot(parked_slots.begin()->first, why); + } + reconfigure_state.parked.clear(); + } + + // after a reconfigure placed the states: each live sampler on the new context, drafting begun again on + // the rebuilt speculative state, and the queued work offered the free slots (a grow adds some) + void resume_after_reconfigure() { + for (auto & slot : slots) { + if (slot.task && slot.smpl) { + if (server_slot::samples_on_backend(slot.task->params)) { + // a backend chain is initialized once, for the context it was first attached to + // (llama_sampler_chain_backend_init): on the rebuilt context the slot samples with a + // fresh chain that carries the same state (common_sampler_copy, as a rollback does) + try { + common_params_sampling sampling = slot.task->params.sampling; // init takes them mutable; the task's are const + common_sampler_ptr fresh(common_sampler_init(model_tgt, sampling)); + common_sampler_copy(slot.smpl.get(), fresh.get()); + slot.smpl = std::move(fresh); + } catch (const std::exception & e) { + send_error(slot, std::string("the sampler could not be made again after a reconfigure: ") + e.what()); + slot.release(); + continue; + } + } + slot.attach_sampler(slot.task->params); + } + if (slot.state == SLOT_STATE_GENERATING && slot.can_speculate()) { + common_speculative_begin(spec.get(), slot.id, slot.prompt.tokens.get_text_tokens()); + } + } + for (auto & slot : slots) { + if (!slot.is_processing()) { + queue_tasks.pop_deferred_task(slot.id); + } + } + publish_busy_slots(); + } + + // what a live reconfigure needs of this server: its contexts' lifetime and its slots + struct reconfigure_engine final : server_reconfigure_engine { + server_context_impl & s; + explicit reconfigure_engine(server_context_impl & owner) : s(owner) {} + + server_geometry live() const override { + // the geometry as asked for (fit's choice when the launch named none), so asking again is a no-op + const int32_t n_ctx = s.params_base.n_ctx > 0 ? s.params_base.n_ctx : (int32_t) llama_n_ctx(s.ctx_tgt); + return { n_ctx, (int32_t) s.slots.size() }; + } + int32_t n_ctx_seq() const override { return s.get_slot_n_ctx(); } + int32_t n_cells() const override { return (int32_t) llama_n_ctx(s.ctx_tgt); } + bool kv_unified() const override { return s.kv_pool_unified(); } + + std::optional blocked() const override { + if (auto learn = s.learning_blocks_reconfigure()) { + return learn; + } + for (const auto & slot : s.slots) { + if (slot.state == SLOT_STATE_WAIT_OTHER) { + return server_reconfigure_refused { SERVER_RECONFIGURE_FAN_OUT, + "a parent prompt is being copied into its child slots (n_cmpl > 1)", slot.id }; + } + } + return std::nullopt; + } + + int32_t tokens_held(int32_t id) const override { + llama_memory_t mem = llama_get_memory(s.ctx_tgt); + if (mem == nullptr) { + return 0; // a context with no memory module (an embedding model) holds no state between tasks + } + const llama_pos pos_max = llama_memory_seq_pos_max(mem, id); + if (pos_max < 0) { + return 0; + } + return std::max(s.slots[id].prompt.n_tokens(), pos_max + 1); + } + bool busy(int32_t id) const override { return s.slots[id].is_processing(); } + + bool page_out(int32_t id, server_slot_page & page) override { return s.slots[id].page_out(page); } + bool page_in (int32_t id, const server_slot_page & page) override { return s.slots[id].page_in(page); } + + void park (int32_t id) override { s.park_slot(id); } + void unpark(int32_t id) override { s.unpark_slot(id); } + + void destroy_contexts() override { s.destroy_contexts(); } + bool build_contexts(const server_geometry & g) override { + s.params_base.n_ctx = g.n_ctx; + s.params_base.n_parallel = g.n_seq; + // no warmup: the weights are already resident, and the warmup would only clear the new memory + return s.build_contexts(false); + } + + void resume() override { s.resume_after_reconfigure(); } + }; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c51277e13599..dd4a36c7b2cb 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,6 +14,8 @@ #include "server-prefill.h" #include "server-learn-ledger.h" #include "server-learn-residency.h" +#include "server-reconfigure.h" +#include "server-slot-page.h" #include "server-test-faults.h" #include "build-info.h" @@ -39,6 +41,7 @@ #include #include #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 @@ -132,6 +135,11 @@ struct server_batch { } void init(int32_t n_tokens_alloc, int32_t n_embd) { + // Continuum fork: every context build sizes the batch again (wake from sleep, a reconfigure) + if (batch.pos != nullptr) { + clear(); + llama_batch_free(batch); + } this->n_tokens_alloc = n_tokens_alloc; this->n_embd = n_embd; batch = llama_batch_init(n_tokens_alloc, 0, 1); @@ -314,27 +322,44 @@ struct server_slot { return false; } - const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); - const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; - - const size_t cur_size = cur_size_tgt + cur_size_dft; + const server_slot_page_size size = server_slot_page_measure(ctx_tgt, ctx_dft, id); SRV_TRC(" - saving prompt with length %d, total state size = %.3f MiB (draft: %.3f MiB)\n", - (int) prompt.tokens.size(), cur_size / (1024.0 * 1024.0), cur_size_dft / (1024.0 * 1024.0)); + (int) prompt.tokens.size(), size.total() / (1024.0 * 1024.0), size.drft / (1024.0 * 1024.0)); - auto * cur = prompt_cache.alloc(prompt, lora_active(lora), cur_size_tgt, cur_size_dft); + auto * cur = prompt_cache.alloc(prompt, lora_active(lora), size.main, size.drft); if (cur == nullptr) { return false; } - llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); - if (ctx_dft) { - llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE); + if (!server_slot_page_read(ctx_tgt, ctx_dft, id, cur->data)) { + prompt_cache.states.pop_back(); // alloc appended it; a short read is not this prompt's state + return false; } return true; } + // Continuum fork: page out / page in for a live reconfigure (server-reconfigure.h), through the + // one mover (server-slot-page.h). Only the KV moves: the prompt, its checkpoints and the sampler + // are host memory already and stay with the slot. + bool page_out(server_slot_page & page) const { + const server_slot_page_size size = server_slot_page_measure(ctx_tgt, ctx_dft, id); + try { + page.data.main.resize(size.main); + page.data.drft.resize(size.drft); + } catch (const std::bad_alloc & e) { + SLT_ERR(*this, "no host memory for a %.3f MiB page: %s\n", size.total() / (1024.0 * 1024.0), e.what()); + return false; + } + page.lora = lora_active(lora); + return server_slot_page_read(ctx_tgt, ctx_dft, id, page.data); + } + + bool page_in(const server_slot_page & page) const { + return server_slot_page_write(ctx_tgt, ctx_dft, id, page.data); + } + // `task_lora` is the adapter list the incoming task will run under: only a state // computed under the same active set is restored, and the slot then records it. bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens, @@ -449,6 +474,24 @@ struct server_slot { (ggml_time_us() - t_start) / 1000.0, n_text, (int) prompt.tokens.size()); } + // whether a task's sampler runs on the backend, attached to the context + static bool samples_on_backend(const task_params & params) { + const bool need_pre_sample_logits = params.sampling.n_probs > 0 && !params.post_sampling_probs; + + bool use_backend_sampling = params.sampling.backend_sampling; + + // TODO: getting pre sampling logits is not yet supported with backend sampling + use_backend_sampling &= !need_pre_sample_logits; + + return use_backend_sampling; + } + + // the backend sampler is bound to the context: a launch attaches it, a reconfigure attaches it again + void attach_sampler(const task_params & params) const { + // TODO: tmp until backend sampling is fully implemented + llama_set_sampler(ctx_tgt, id, samples_on_backend(params) ? common_sampler_get(smpl.get()) : nullptr); + } + bool need_embd() const { GGML_ASSERT(task); return task->need_embd(); @@ -916,6 +959,19 @@ struct server_context_impl { return lora_report; } + // the live geometry and the reconfigure counters as the main loop last published them + // (publish_geometry): /props reads them here, never through the task queue + json geometry_report_snapshot() const { + std::lock_guard lock(geometry_report_mutex); + return geometry_report; + } + int32_t live_n_seq() const { + return n_seq_live.load(); + } + int32_t live_n_ctx_seq() const { + return n_ctx_seq_live.load(); + } + mtmd_context * mctx = nullptr; const llama_vocab * vocab = nullptr; @@ -1110,18 +1166,44 @@ struct server_context_impl { lora_report = std::move(report); } +#include "server-context-reconfigure.inc" + int64_t t_last_load_progress_ms = 0; + // Continuum fork: the weights and the contexts are two lifetimes (SERVING-ENGINE.md, S2). + // destroy = destroy_contexts + free_weights; load_model = load_weights + build_contexts; a live + // reconfigure (server-reconfigure.h) is page out, destroy_contexts, build_contexts, page in, so the + // model and its mmap are never freed for a new geometry. void destroy() { + destroy_contexts(); + free_weights(); + // a slot's prompt describes KV that is gone with its context; the next load starts the slots empty + slots.clear(); + } + + // the target and draft contexts and the speculative state; the slots keep their host state + void destroy_contexts() { spec.reset(); - spec_init.reset(); - ctx_dft = nullptr; + if (spec_init) { + spec_init->free_context(); + } + ctx_dft = nullptr; + params_base.speculative.draft.ctx_tgt = nullptr; + params_base.speculative.draft.ctx_dft = nullptr; + + if (llama_init) { + llama_init->free_context(); + } + ctx_tgt = nullptr; + } + + // the target model with its adapters, the draft model, the multimodal projector + void free_weights() { + spec_init.reset(); model_dft = nullptr; llama_init.reset(); - - ctx_tgt = nullptr; model_tgt = nullptr; mtmd_free(mctx); @@ -1136,6 +1218,8 @@ struct server_context_impl { // note: for sleeping == false, event is emitted by load_model() } SRV_INF("%s", "server is entering sleeping state\n"); + // sleep frees the weights a parked slot's sampler was made on: its work ends, by name + drop_parked("the server went to sleep while this slot was paged out by a reconfigure"); destroy(); } else { SRV_INF("%s", "server is exiting sleeping state\n"); @@ -1188,9 +1272,6 @@ struct server_context_impl { const bool is_resume = sleeping; params_base = params; - const auto output_limits = server_output_limits(params_base); - params_base.n_outputs_max = output_limits.total; - params_base.n_outputs_max_per_seq = output_limits.per_seq; const bool has_mmproj = !params.mmproj.path.empty(); const bool has_draft = params.speculative.has_dft(); @@ -1219,6 +1300,116 @@ struct server_context_impl { SRV_INF("loading model '%s'\n", params.model.get_name().c_str()); SRV_TRC("local path '%s'\n", params.model.path.c_str()); + if (!load_weights(load_progress_text, load_progress_mmproj, load_progress_spec)) { + return false; + } + + if (!build_contexts(params_base.warmup)) { + return false; + } + + { + const char * LLAMA_TRACE = getenv("LLAMA_TRACE"); + trace = LLAMA_TRACE ? atoi(LLAMA_TRACE) : 0; + + if (trace) { + SRV_WRN("LLAMA_TRACE = %d\n", trace); + } + } + + { + const char * LLAMA_SERVER_SLOTS_DEBUG = getenv("LLAMA_SERVER_SLOTS_DEBUG"); + slots_debug = LLAMA_SERVER_SLOTS_DEBUG ? atoi(LLAMA_SERVER_SLOTS_DEBUG) : 0; + + if (slots_debug) { + SRV_WRN("LLAMA_SERVER_SLOTS_DEBUG = %d\n", slots_debug); + } + } + + { + const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF"); + slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0; + + if (slots_n_diff) { + SRV_WRN("LLAMA_SERVER_SLOTS_N_DIFF = %d\n", slots_n_diff); + } + } + + if constexpr (server_test_faults::enabled) { + test_faults = server_test_faults::from_env(); + if (test_faults.any()) { + SRV_WRN("test faults: prompt checkpoints = %d, speculative checkpoints = %d, prompt fills = %d, step delay = %d ms\n", + test_faults.prompt_checkpoints, test_faults.spec_checkpoints, test_faults.prompt_fills, test_faults.step_delay_ms); + } + } + + if (params_base.cache_ram_mib != 0) { + if (params_base.cache_ram_mib < 0) { + SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); + } else { + SRV_TRC("prompt cache is enabled, size limit: %d MiB\n", params_base.cache_ram_mib); + } + SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n"); + + prompt_cache = std::make_unique(params_base.cache_ram_mib, n_ctx); + } else { + SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); + } + SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); + + if (params_base.n_ctx_checkpoints > 0) { + if (params_base.checkpoint_min_step < 0) { + SRV_TRC("context checkpoints enabled, max = %d, min spacing = each prompt over %d\n", + params_base.n_ctx_checkpoints, params_base.n_ctx_checkpoints); + } else { + SRV_TRC("context checkpoints enabled, max = %d, min spacing = %d\n", + params_base.n_ctx_checkpoints, params_base.checkpoint_min_step); + } + } else { + SRV_TRC("%s", "context checkpoints disabled\n"); + } + + if (!params_base.model_alias.empty()) { + // backward compat: use first alias as model name + model_name = *params_base.model_alias.begin(); + } else if (!params_base.model.get_name().empty()) { + model_name = params_base.model.get_name(); + } else { + // fallback: derive model name from file name + auto model_path = std::filesystem::path(params_base.model.path); + model_name = model_path.filename().string(); + } + + model_aliases = params_base.model_alias; + model_tags = params_base.model_tags; + + // propagate new defaults back to caller + params = params_base; + + publish_lora_report(); + publish_geometry(); + + if (!is_resume) { + return init(); + } + + if (callback_state) { + callback_state(SERVER_STATE_READY, {}); + } + + return true; + } + + // THE WEIGHTS, which survive a reconfigure: the target model with its adapters, the draft model, + // the multimodal projector. No context is created here (build_contexts). + bool load_weights(load_progress_data & load_progress_text, load_progress_data & load_progress_mmproj, load_progress_data & load_progress_spec) { + const bool has_mmproj = !params_base.mmproj.path.empty(); + const bool has_draft = params_base.speculative.has_dft(); + const bool spec_mtp = std::find(params_base.speculative.types.begin(), + params_base.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); + const bool has_spec = has_draft || spec_mtp; + std::string & mmproj_path = params_base.mmproj.path; mtmd_context_params mparams = mtmd_context_params_default(); if (has_mmproj) { @@ -1273,25 +1464,19 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - llama_init = common_init_from_params(params_base); + llama_init = common_load_from_params(params_base); model_tgt = llama_init->model(); - ctx_tgt = llama_init->context(); if (model_tgt == nullptr) { SRV_ERR("failed to load model, '%s'\n", params_base.model.path.c_str()); return false; } - if (ctx_tgt == nullptr) { - SRV_ERR("failed to create_context with model '%s'\n", params_base.model.path.c_str()); - return false; - } + n_weight_loads++; vocab = llama_model_get_vocab(model_tgt); - n_ctx = llama_n_ctx(ctx_tgt); - add_bos_token = llama_vocab_get_add_bos(vocab); if (has_spec) { @@ -1306,22 +1491,13 @@ struct server_context_impl { params_dft.load_progress_callback = load_progress_callback; params_dft.load_progress_callback_user_data = &load_progress_spec; - spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); + spec_init = std::make_unique(params_dft); model_dft = spec_init->model(); - ctx_dft = spec_init->context(); if (has_draft && model_dft == nullptr) { SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); return false; } - - if (ctx_dft == nullptr) { - SRV_ERR("%s", "failed to create MTP context\n"); - return false; - } - - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft; } load_progress_callback(1.0f, &load_progress_spec); @@ -1332,7 +1508,7 @@ struct server_context_impl { callback_state(SERVER_STATE_LOADING, {{"stage", "mmproj_model"}}); } - if (!is_resume) { + if (!sleeping) { mtmd_helper_log_set(common_log_default_callback, nullptr); } @@ -1354,18 +1530,6 @@ struct server_context_impl { } } - if (!llama_memory_can_shift(llama_get_memory(ctx_tgt))) { - if (params_base.ctx_shift) { - params_base.ctx_shift = false; - SRV_WRN("%s\n", "ctx_shift is not supported by this context, it will be disabled"); - } - - if (params_base.n_cache_reuse) { - params_base.n_cache_reuse = 0; - SRV_WRN("%s\n", "cache_reuse is not supported by this context, it will be disabled"); - } - } - if (llama_model_n_swa(model_tgt) == 0) { if (params_base.swa_full) { params_base.swa_full = false; @@ -1378,6 +1542,59 @@ struct server_context_impl { // Necessary similarity of prompt for slot selection slot_prompt_similarity = params_base.slot_prompt_similarity; + return true; + } + + // THE CONTEXTS, which a reconfigure rebuilds on the same weights: the target and draft contexts, + // the speculative state, the slots and the batch, for params_base's n_ctx and n_parallel. On + // failure nothing of them is left, so the caller may build again (the previous geometry). + bool build_contexts(bool warmup) { + if (!build_contexts_impl(warmup)) { + destroy_contexts(); + return false; + } + return true; + } + + bool build_contexts_impl(bool warmup) { + const auto output_limits = server_output_limits(params_base); + params_base.n_outputs_max = output_limits.total; + params_base.n_outputs_max_per_seq = output_limits.per_seq; + + if (!llama_init->init_context(params_base, warmup)) { + SRV_ERR("failed to create_context with model '%s'\n", params_base.model.path.c_str()); + return false; + } + + ctx_tgt = llama_init->context(); + + n_ctx = llama_n_ctx(ctx_tgt); + + if (spec_init) { + common_params params_dft = common_base_params_to_speculative(params_base); + + if (!spec_init->init_context(params_dft, model_tgt, ctx_tgt) || spec_init->context() == nullptr) { + SRV_ERR("%s", "failed to create MTP context\n"); + return false; + } + ctx_dft = spec_init->context(); + + params_base.speculative.draft.ctx_tgt = ctx_tgt; + params_base.speculative.draft.ctx_dft = ctx_dft; + } + + if (!llama_memory_can_shift(llama_get_memory(ctx_tgt))) { + if (params_base.ctx_shift) { + params_base.ctx_shift = false; + SRV_WRN("%s\n", "ctx_shift is not supported by this context, it will be disabled"); + } + + if (params_base.n_cache_reuse) { + params_base.n_cache_reuse = 0; + SRV_WRN("%s\n", "cache_reuse is not supported by this context, it will be disabled"); + } + } + const int n_ctx_train = llama_model_n_ctx_train(model_tgt); int n_ctx_slot = llama_n_ctx_seq(ctx_tgt); @@ -1386,8 +1603,6 @@ struct server_context_impl { n_ctx_slot = n_ctx_train; } - slots.clear(); - ctx_tgt_seq_rm_type = common_context_can_seq_rm(ctx_tgt); if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_NO) { SRV_WRN("%s", "speculative decoding not supported by this context\n"); @@ -1401,11 +1616,6 @@ struct server_context_impl { SRV_INF("initializing, n_slots = %d, n_ctx_slot = %d, kv_unified = '%s'\n", params_base.n_parallel, n_ctx_slot, params_base.kv_unified ? "true" : "false"); - // initialize slots - for (int i = 0; i < params_base.n_parallel; i++) { - slots.emplace_back(); - } - // try speculative decoding if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { try { @@ -1427,69 +1637,7 @@ struct server_context_impl { model_dft = nullptr; } - 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(); - } - - { - const char * LLAMA_TRACE = getenv("LLAMA_TRACE"); - trace = LLAMA_TRACE ? atoi(LLAMA_TRACE) : 0; - - if (trace) { - SRV_WRN("LLAMA_TRACE = %d\n", trace); - } - } - - { - const char * LLAMA_SERVER_SLOTS_DEBUG = getenv("LLAMA_SERVER_SLOTS_DEBUG"); - slots_debug = LLAMA_SERVER_SLOTS_DEBUG ? atoi(LLAMA_SERVER_SLOTS_DEBUG) : 0; - - if (slots_debug) { - SRV_WRN("LLAMA_SERVER_SLOTS_DEBUG = %d\n", slots_debug); - } - } - - { - const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF"); - slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0; - - if (slots_n_diff) { - SRV_WRN("LLAMA_SERVER_SLOTS_N_DIFF = %d\n", slots_n_diff); - } - } - - if constexpr (server_test_faults::enabled) { - test_faults = server_test_faults::from_env(); - if (test_faults.any()) { - SRV_WRN("test faults: prompt checkpoints = %d, speculative checkpoints = %d, prompt fills = %d\n", - test_faults.prompt_checkpoints, test_faults.spec_checkpoints, test_faults.prompt_fills); - } - } + bind_slots(n_ctx_slot); // the update_slots() logic will always submit a maximum of n_batch or n_parallel tokens // note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used) @@ -1499,60 +1647,62 @@ struct server_context_impl { batch.init(std::max(n_batch, params_base.n_parallel), n_embd); } - if (params_base.cache_ram_mib != 0) { - if (params_base.cache_ram_mib < 0) { - SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); - } else { - SRV_TRC("prompt cache is enabled, size limit: %d MiB\n", params_base.cache_ram_mib); - } - SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n"); - - prompt_cache = std::make_unique(params_base.cache_ram_mib, n_ctx); - } else { - SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); + if (prompt_cache) { + prompt_cache->limit_tokens = n_ctx; // the cache's token limit follows the context (server_prompt_cache) } - SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); - if (params_base.n_ctx_checkpoints > 0) { - if (params_base.checkpoint_min_step < 0) { - SRV_TRC("context checkpoints enabled, max = %d, min spacing = each prompt over %d\n", - params_base.n_ctx_checkpoints, params_base.n_ctx_checkpoints); - } else { - SRV_TRC("context checkpoints enabled, max = %d, min spacing = %d\n", - params_base.n_ctx_checkpoints, params_base.checkpoint_min_step); + return true; + } + + // the slots on the contexts just built: their count follows params_base.n_parallel (the highest + // ids go, a new id comes in empty) and every slot is bound to the new contexts. A slot's host state + // (prompt, checkpoints, task, sampler) is untouched: its KV is paged in by the caller. + int32_t n_ctx_slot_live = 0; + + void bind_slots(int32_t n_ctx_slot) { + n_ctx_slot_live = n_ctx_slot; + + const size_t n_before = slots.size(); + slots.resize(params_base.n_parallel); + + for (int i = 0; i < params_base.n_parallel; i++) { + server_slot & slot = slots[i]; + bind_slot(slot); + if ((size_t) i >= n_before) { + init_slot(slot, i); } - } else { - SRV_TRC("%s", "context checkpoints disabled\n"); } + } - if (!params_base.model_alias.empty()) { - // backward compat: use first alias as model name - model_name = *params_base.model_alias.begin(); - } else if (!params_base.model.get_name().empty()) { - model_name = params_base.model.get_name(); - } else { - // fallback: derive model name from file name - auto model_path = std::filesystem::path(params_base.model.path); - model_name = model_path.filename().string(); - } + void bind_slot(server_slot & slot) const { + 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_live; + } - model_aliases = params_base.model_alias; - model_tags = params_base.model_tags; + // a new, empty slot (bound first: its reset detaches the context's sampler for its sequence) + void init_slot(server_slot & slot, int id) { + slot.id = id; - // propagate new defaults back to caller - params = params_base; + slot.mctx = mctx; + slot.prompt.tokens.has_mtmd = mctx != nullptr; - publish_lora_report(); + SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); - if (!is_resume) { - return init(); - } + slot.callback_on_release = [this](int id_slot) { + queue_tasks.pop_deferred_task(id_slot); + }; - if (callback_state) { - callback_state(SERVER_STATE_READY, {}); - } + 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); + } + }; - return true; + slot.reset(); } // unlike load_model(), this is only called once during initialization @@ -1665,8 +1815,12 @@ struct server_context_impl { } server_slot * get_slot_by_id(int id_slot) { - // note: allow id_slot to be out of bounds (wrap around) - id_slot = id_slot % slots.size(); + // Continuum fork: an id addresses exactly its slot. The upstream wrap-around (id % n_slots) sent a + // request for a slot a reconfigure parked to another mind's live slot; an id out of range is now + // no slot, and the caller refuses it (or, for a parked id, waits for it: slot_id_is_parked) + if (id_slot < 0 || id_slot >= (int) slots.size()) { + return nullptr; + } for (server_slot & slot : slots) { if (slot.id == id_slot) { @@ -1701,6 +1855,9 @@ struct server_context_impl { ret = get_slot_by_id(task.id_slot); if (ret) { SLT_INF(*ret, "selected slot by id (%d)\n", task.id_slot); + } else { + // a parked slot (server-reconfigure.h): the task waits for its own slot, never another's + return nullptr; } } @@ -1939,19 +2096,7 @@ struct server_context_impl { return false; } - const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs; - - bool use_backend_sampling = task.params.sampling.backend_sampling; - - // TODO: getting pre sampling logits is not yet supported with backend sampling - use_backend_sampling &= !need_pre_sample_logits; - - // TODO: tmp until backend sampling is fully implemented - if (use_backend_sampling) { - llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); - } else { - llama_set_sampler(ctx_tgt, slot.id, nullptr); - } + slot.attach_sampler(task.params); 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()); @@ -2599,6 +2744,13 @@ struct server_context_impl { const int id_task = task.id; + // Continuum fork: a slot id that is neither live nor parked by a reconfigure names no slot + if (task.id_slot != -1 && get_slot_by_id(task.id_slot) == nullptr && !slot_id_is_parked(task.id_slot)) { + send_error(task, "Invalid slot ID " + std::to_string(task.id_slot) + ": the server has " + + std::to_string(slots.size()) + " slots", ERROR_TYPE_INVALID_REQUEST); + break; + } + server_slot * slot = get_available_slot(task); // @@ -2664,6 +2816,14 @@ struct server_context_impl { break; } } + // or drop it where a reconfigure parked it (its client is gone: nobody to tell) + for (auto & [id, slot] : parked_slots) { + if (slot.task && slot.task->id == task.id_target) { + drop_parked_slot(id, ""); + publish_geometry(); + break; + } + } } break; case SERVER_TASK_TYPE_CONTROL: { @@ -2780,6 +2940,19 @@ struct server_context_impl { res->n_cached = std::min(slot->prompt.tokens.get_common_prefix(task.tokens), (size_t) std::max(held, 0)); queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_RECONFIGURE: + { + // Continuum fork: grow or shrink the contexts live (server-reconfigure.h). In the task + // phase no llama_decode is in flight and every sequence's KV is at an inter-batch + // boundary: the safe point SLOT_SAVE argues for below. Generations in progress are + // paged out and continue; new work waits in the queue behind this task. + reconfigure_engine engine(*this); + auto res = std::make_unique(); + res->id = task.id; + res->outcome = server_reconfigure(engine, reconfigure_state, task.reconfigure); + publish_geometry(); + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_SLOT_SAVE: { const int id_slot = task.slot_action.id_slot; @@ -2858,25 +3031,9 @@ struct server_context_impl { kept_checkpoints.size(), n, kept_at.c_str(), view_end, end_from); } } + // the page format's one owner (server-slot-page.h) 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()); - } + server_slot_page_put_checkpoints(kept_checkpoints, put); }; auto write_adapters = [&](auto && put) { const uint32_t magic = 0x4C4F5241u; // 'LORA' @@ -3088,36 +3245,18 @@ struct server_context_impl { } } std::list loaded; - bool ok = (bool) in; auto get = [&](void * ptr, size_t n) { if (progress && progress->cancel_requested.load(std::memory_order_relaxed)) { throw std::runtime_error("Restore operation cancelled"); } in.read(reinterpret_cast(ptr), n); - ok = ok && (bool) in; if (progress && in.gcount() > 0) progress->sidecar_bytes.fetch_add(uint64_t(in.gcount()), std::memory_order_relaxed); + return (bool) in; }; - if (ok) { - uint32_t magic = 0, version = 0, n_ckpt = 0; - get(&magic, sizeof(magic)); get(&version, sizeof(version)); get(&n_ckpt, sizeof(n_ckpt)); - ok = ok && magic == 0x434B5054u && version == 1 && n_ckpt <= 4096; - for (uint32_t i = 0; ok && i < n_ckpt; ++i) { - int64_t n_tokens = 0; int32_t pos_min = 0, pos_max = 0; - uint64_t sz_tgt = 0, sz_dft = 0, sz_spec = 0; - get(&n_tokens, sizeof(n_tokens)); get(&pos_min, sizeof(pos_min)); get(&pos_max, sizeof(pos_max)); - get(&sz_tgt, sizeof(sz_tgt)); get(&sz_dft, sizeof(sz_dft)); get(&sz_spec, sizeof(sz_spec)); - // a checkpoint is at most one sequence state; refuse absurd sizes - ok = ok && sz_tgt < (1ull << 36) && sz_dft < (1ull << 36) && sz_spec < (1ull << 30); - if (!ok) break; - auto & c = loaded.emplace_back(); - c.update_pos(n_tokens, pos_min, pos_max); - c.id_task = -1; - c.data_tgt.resize(sz_tgt); get(c.data_tgt.data(), sz_tgt); - c.data_dft.resize(sz_dft); get(c.data_dft.data(), sz_dft); - c.data_spec.resize(sz_spec); get(c.data_spec.data(), sz_spec); - if (progress && ok) progress->checkpoints.fetch_add(1, std::memory_order_relaxed); - } - } + // the page format's one owner (server-slot-page.h) + const bool ok = in && server_slot_page_get_checkpoints(get, loaded, [&]() { + if (progress) progress->checkpoints.fetch_add(1, std::memory_order_relaxed); + }); if (ok) { slot->prompt.checkpoints = std::move(loaded); SLT_INF(*slot, "restored %zu context checkpoints from the sidecar: %s\n", slot->prompt.checkpoints.size(), ckpt_path.c_str()); @@ -3638,6 +3777,12 @@ struct server_context_impl { // preparation is already one measured blocking segment; do not append a // graph segment before the queued NEXT_RESPONSE lets serving run again. if (learning_at_start) { advance_learning(); } + + if constexpr (server_test_faults::enabled) { + if (test_faults.step_delay_ms > 0 && batch.size() > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(test_faults.step_delay_ms)); + } + } } void pre_decode() { @@ -5189,6 +5334,7 @@ std::unique_ptr server_routes::handle_completions_impl( task.params = server_schema::eval_llama_cmpl_schema( ctx_server.vocab, params, + ctx_server.live_n_seq(), meta->logit_bias_eog, data); @@ -5740,8 +5886,44 @@ void server_routes::init_routes() { std::unique_lock lock(mutex_cache); res->ok(cached_props); } else { - res->ok(get_res_props(*meta, params, false)); + json props = get_res_props(*meta, params, false); + // Continuum fork: the geometry is live (POST /reconfigure); meta is the load's snapshot + const json geometry = ctx_server.geometry_report_snapshot(); + if (!geometry.empty()) { + props["reconfigure"] = geometry; + props["context_geometry"] = geometry.at("context_geometry"); + props["total_slots"] = geometry.at("n_seq"); + props["default_generation_settings"]["n_ctx"] = geometry.at("n_ctx_seq"); + } + res->ok(props); + } + return res; + }; + + // POST /reconfigure (Continuum fork): grow or shrink the serving contexts live while the weights stay + // loaded (server-reconfigure.h). {"n_ctx": int, "n_seq": int}; the outcome names every slot paged out. + this->post_reconfigure = [this](const server_http_req & req) { + auto res = create_response(); + const auto parsed = server_reconfigure_from_json(json::parse(req.body), + ggml_type_name(params.cache_type_k), ggml_type_name(params.cache_type_v)); + if (const auto * refused = std::get_if(&parsed)) { + res->error(server_reconfigure_to_json(*refused)); + return res; } + server_task task(SERVER_TASK_TYPE_RECONFIGURE); + task.id = res->rd.get_new_id(); + task.reconfigure = std::get(parsed); + res->rd.post_task(std::move(task), true); // high-priority 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; + } + res->ok(result->to_json()); return res; }; @@ -5821,7 +6003,7 @@ void server_routes::init_routes() { data.at("input_extra"), params.n_batch, params.n_predict, - meta->slot_n_ctx, + ctx_server.live_n_ctx_seq(), // the live window (POST /reconfigure), not the load's params.spm_infill, tokenized_prompts[0].get_tokens() // TODO: this could maybe be multimodal. ); @@ -6553,7 +6735,7 @@ void server_routes::init_routes() { rows = render_learn_example(body).reply.size(); } else if (body.contains("learn")) { // Completion-attached corrections share the completion schema, including EOT. - rows = server_schema::eval_llama_cmpl_schema(ctx_server.vocab, params, meta->logit_bias_eog, + rows = server_schema::eval_llama_cmpl_schema(ctx_server.vocab, params, ctx_server.live_n_seq(), meta->logit_bias_eog, json{{"learn", body.at("learn")}}).learn->reply.size(); } else { throw std::invalid_argument("preflight requires messages or a completion learn example"); @@ -6631,7 +6813,7 @@ void server_routes::init_routes() { // prefill path, one learn path; the slot is the caller's choice (whose cache it replaces). server_task serve(SERVER_TASK_TYPE_COMPLETION); serve.tokens = server_tokens(task.learn.prompt, false); - serve.params = server_schema::eval_llama_cmpl_schema(ctx_server.vocab, params, meta->logit_bias_eog, + serve.params = server_schema::eval_llama_cmpl_schema(ctx_server.vocab, params, ctx_server.live_n_seq(), meta->logit_bias_eog, json { { "n_predict", 1 }, { "temperature", 0.0 }, { "cache_prompt", true } }); serve.params.learn = task.learn; serve.params.learn_only = true; @@ -6847,7 +7029,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 >= ctx_server.live_n_seq()) { res->error(format_error_response("Invalid restore operation identity or slot", ERROR_TYPE_INVALID_REQUEST)); return res; } diff --git a/tools/server/server-context.h b/tools/server/server-context.h index a810c3e3a1b9..8e606832bd2f 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -150,6 +150,7 @@ struct server_routes { server_http_context::handler_t post_slots; server_http_context::handler_t get_props; server_http_context::handler_t post_props; + server_http_context::handler_t post_reconfigure; // Continuum fork (server-reconfigure.h) server_http_context::handler_t post_infill; server_http_context::handler_t post_completions; server_http_context::handler_t post_completions_oai; diff --git a/tools/server/server-geometry.h b/tools/server/server-geometry.h new file mode 100644 index 000000000000..cc1ddf6246ea --- /dev/null +++ b/tools/server/server-geometry.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +// Continuum fork: the serving contexts' geometry, what a live reconfigure changes +// (server-reconfigure.h). Its own header so a task can carry one (server-task.h). +struct server_geometry { + int32_t n_ctx = 0; // tokens the target context allocates, every sequence together (-c) + int32_t n_seq = 0; // sequences, one per slot (-np) + + bool operator==(const server_geometry & other) const { return n_ctx == other.n_ctx && n_seq == other.n_seq; } + bool operator!=(const server_geometry & other) const { return !(*this == other); } +}; diff --git a/tools/server/server-reconfigure.cpp b/tools/server/server-reconfigure.cpp new file mode 100644 index 000000000000..32e01dae6752 --- /dev/null +++ b/tools/server/server-reconfigure.cpp @@ -0,0 +1,309 @@ +#include "server-reconfigure.h" + +#include "server-common.h" + +#include "ggml.h" + +#include +#include + +const char * server_reconfigure_refusal_name(server_reconfigure_refusal_reason reason) { + switch (reason) { + case SERVER_RECONFIGURE_INVALID_GEOMETRY: return "invalid_geometry"; + case SERVER_RECONFIGURE_CACHE_TYPE_CHANGE: return "cache_type_change"; + case SERVER_RECONFIGURE_LEARN_IN_PROGRESS: return "learn_in_progress"; + case SERVER_RECONFIGURE_LEARNING_BOUND: return "learning_bound"; + case SERVER_RECONFIGURE_FAN_OUT: return "fan_out_in_progress"; + case SERVER_RECONFIGURE_PAGE_OUT_FAILED: return "page_out_failed"; + case SERVER_RECONFIGURE_CONTEXT_BUILD_FAILED: return "context_build_failed"; + } + GGML_ABORT("unknown reconfigure refusal"); +} + +namespace { + +// a slot's state while the contexts are rebuilt +struct paged_slot { + server_slot_page page; + int32_t tokens_held = 0; + bool busy = false; +}; + +server_reconfigure_refused refuse(server_reconfigure_refusal_reason reason, std::string detail, int32_t slot = -1) { + return server_reconfigure_refused { reason, std::move(detail), slot }; +} + +} // namespace + +server_reconfigure_outcome server_reconfigure(server_reconfigure_engine & e, server_reconfigure_state & st, const server_geometry & target) { + auto refused = [&](server_reconfigure_refused r) -> server_reconfigure_outcome { + SRV_WRN("reconfigure to n_ctx = %d, n_seq = %d refused (%s): %s\n", + target.n_ctx, target.n_seq, server_reconfigure_refusal_name(r.reason), r.detail.c_str()); + st.n_refused++; + st.last = r; + return r; + }; + + // every sequence needs at least one cell of its own + if (target.n_seq < 1 || target.n_ctx < target.n_seq) { + return refused(refuse(SERVER_RECONFIGURE_INVALID_GEOMETRY, "n_seq must be at least 1 and n_ctx at least n_seq")); + } + if (auto blocked = e.blocked()) { + return refused(*blocked); + } + + const int64_t t_start = ggml_time_us(); + const server_geometry before = e.live(); + + server_reconfigure_applied out; + out.rebuilt = target != before; + + std::map pages; // slots in service whose state crosses the rebuild + std::vector parked_now; + + // the slot leaves service whole: its host state to the engine's park, its page here + auto park = [&](int32_t id, paged_slot && p) { + e.park(id); + st.parked[id] = server_parked_slot { std::move(p.page), p.tokens_held, p.busy }; + parked_now.push_back(id); + }; + + if (out.rebuilt) { + // 1. page out every slot holding state, unfinished generations included. Reading never changes + // a context, so a failure here leaves everything as it was. + for (int32_t id = 0; id < before.n_seq; ++id) { + const int32_t held = e.tokens_held(id); + if (held == 0) { + continue; + } + paged_slot p; + p.tokens_held = held; + p.busy = e.busy(id); + if (!e.page_out(id, p.page)) { + return refused(refuse(SERVER_RECONFIGURE_PAGE_OUT_FAILED, "the slot's state could not be read", id)); + } + out.tokens_held += held; + out.bytes_paged += p.page.data.size(); + pages.emplace(id, std::move(p)); + } + + // 2. the ids a smaller n_seq removes leave service with their state, busy or not + for (int32_t id = target.n_seq; id < before.n_seq; ++id) { + auto it = pages.find(id); + if (it == pages.end() && !e.busy(id)) { + continue; // nothing there: the id just goes + } + paged_slot p; + if (it != pages.end()) { + p = std::move(it->second); + pages.erase(it); + } else { + p.busy = true; // a task that has not filled its prompt yet + } + out.paged_out.push_back({ id, p.tokens_held, 0 }); + park(id, std::move(p)); + } + + // 3. the contexts, rebuilt on the same weights + e.destroy_contexts(); + if (!e.build_contexts(target)) { + // put back exactly what was: the previous geometry, the slots this call parked, every state + if (!e.build_contexts(before)) { + GGML_ABORT("reconfigure: the previous geometry (n_ctx = %d, n_seq = %d) could not be rebuilt either; " + "the server has no context to serve with", before.n_ctx, before.n_seq); + } + for (int32_t id : parked_now) { + auto it = st.parked.find(id); + e.unpark(id); + pages[id] = paged_slot { std::move(it->second.page), it->second.tokens_held, it->second.busy }; + st.parked.erase(it); + } + for (auto & [id, p] : pages) { + if (p.tokens_held > 0 && !e.page_in(id, p.page)) { + // the geometry that held it a moment ago does not take it back: keep it paged, by name + SRV_ERR("reconfigure: slot %d's state did not return to the previous geometry; it stays parked\n", id); + e.park(id); + st.parked[id] = server_parked_slot { std::move(p.page), p.tokens_held, p.busy }; + } + } + e.resume(); + return refused(refuse(SERVER_RECONFIGURE_CONTEXT_BUILD_FAILED, + "the contexts for the target geometry could not be created; the previous geometry was rebuilt")); + } + } + + out.live = e.live(); + out.n_ctx_seq = e.n_ctx_seq(); + + const bool unified = e.kv_unified(); + const int64_t n_cells = e.n_cells(); + // with one pool every sequence draws on, the states placed so far use part of it + int64_t used = 0; + if (!out.rebuilt) { + for (int32_t id = 0; id < out.live.n_seq; ++id) { + used += e.tokens_held(id); + } + } + // a state fits when its sequence keeps room for the next token + auto room = [&]() -> int32_t { + const int64_t free = unified ? n_cells - used : std::numeric_limits::max(); + return (int32_t) std::min(out.n_ctx_seq, free); + }; + + // 4. page in: busy slots first (work in progress), then idle ones, each in id order. A state the + // target cannot hold is not cut: it is parked whole and named. + std::vector order; + for (const auto & [id, p] : pages) { + if (p.busy) { + order.push_back(id); + } + } + for (const auto & [id, p] : pages) { + if (!p.busy) { + order.push_back(id); + } + } + for (int32_t id : order) { + paged_slot & p = pages.at(id); + const int32_t fit = room(); + if (p.tokens_held < fit && e.page_in(id, p.page)) { + used += p.tokens_held; + continue; + } + out.paged_out.push_back({ id, p.tokens_held, std::max(fit, 0) }); + park(id, std::move(p)); + } + + // 5. parked slots come back where the live geometry holds them and their id is free + for (auto it = st.parked.begin(); it != st.parked.end();) { + const int32_t id = it->first; + const server_parked_slot & p = it->second; + const bool this_call = std::find(parked_now.begin(), parked_now.end(), id) != parked_now.end(); + if (this_call || id >= out.live.n_seq || p.tokens_held >= room() || e.busy(id) || e.tokens_held(id) > 0) { + ++it; + continue; + } + e.unpark(id); + if (p.tokens_held > 0 && !e.page_in(id, p.page)) { + e.park(id); // the page stays where it is, parked + ++it; + continue; + } + used += p.tokens_held; + out.paged_in.push_back(id); + it = st.parked.erase(it); + } + + e.resume(); + + for (const auto & [id, p] : st.parked) { + out.parked.push_back(id); + } + out.pause_ms = (ggml_time_us() - t_start) / 1000.0; + + SRV_INF("reconfigured n_ctx = %d -> %d, n_seq = %d -> %d (%d tokens per sequence): paused %.2f ms for %lld tokens " + "(%.3f MiB) of state, %zu slots paged out, %zu placed back, %zu parked\n", + before.n_ctx, out.live.n_ctx, before.n_seq, out.live.n_seq, out.n_ctx_seq, out.pause_ms, + (long long) out.tokens_held, out.bytes_paged / (1024.0 * 1024.0), + out.paged_out.size(), out.paged_in.size(), out.parked.size()); + + st.n_applied++; + if (out.rebuilt) { + st.n_rebuilt++; + st.pause_ms_total += out.pause_ms; + st.pause_ms_max = std::max(st.pause_ms_max, out.pause_ms); + st.tokens_held_total += out.tokens_held; + st.bytes_paged_total += out.bytes_paged; + } + st.last = out; + return out; +} + +json server_reconfigure_cost_json(const server_reconfigure_state & st) { + json last = nullptr; + if (st.last) { + if (const auto * a = std::get_if(&*st.last)) { + last = { { "pause_ms", a->pause_ms }, { "tokens_held", a->tokens_held }, { "bytes_paged", a->bytes_paged }, { "rebuilt", a->rebuilt } }; + } + } + return json { + { "n_rebuilt", st.n_rebuilt }, + { "pause_ms_total", st.pause_ms_total }, + { "pause_ms_max", st.pause_ms_max }, + { "tokens_held_total", st.tokens_held_total }, + { "bytes_paged_total", st.bytes_paged_total }, + // pause per thousand tokens held, over every rebuild that held any (null until one did) + { "pause_ms_per_1k_tokens", st.tokens_held_total > 0 ? json(st.pause_ms_total * 1000.0 / st.tokens_held_total) : json(nullptr) }, + { "last_applied", last }, + }; +} + +std::variant server_reconfigure_from_json( + const json & body, const char * cache_type_k, const char * cache_type_v) { + if (!body.is_object() + || !body.contains("n_ctx") || !body.at("n_ctx").is_number_integer() + || !body.contains("n_seq") || !body.at("n_seq").is_number_integer()) { + return refuse(SERVER_RECONFIGURE_INVALID_GEOMETRY, "\"n_ctx\" and \"n_seq\" (integers) are required"); + } + const int64_t n_ctx = body.at("n_ctx").get(); + const int64_t n_seq = body.at("n_seq").get(); + if (n_ctx < 1 || n_seq < 1 || n_ctx > std::numeric_limits::max() || n_seq > std::numeric_limits::max()) { + return refuse(SERVER_RECONFIGURE_INVALID_GEOMETRY, "\"n_ctx\" and \"n_seq\" must be positive 32-bit integers"); + } + // the KV type is the cache's, never the geometry's: a change of it is the intentional cold start + for (const auto & [key, live] : { std::make_pair("cache_type_k", cache_type_k), std::make_pair("cache_type_v", cache_type_v) }) { + if (!body.contains(key)) { + continue; + } + if (!body.at(key).is_string() || body.at(key).get() != live) { + return refuse(SERVER_RECONFIGURE_CACHE_TYPE_CHANGE, + std::string("\"") + key + "\" is " + live + "; a KV cache type change is a cold start, not a reconfigure"); + } + } + return server_geometry { (int32_t) n_ctx, (int32_t) n_seq }; +} + +json server_reconfigure_to_json(const server_reconfigure_outcome & outcome) { + if (const auto * r = std::get_if(&outcome)) { + // an error body the HTTP edge sends as is ({"error": ...}), with the refusal typed beside the message + error_type type = ERROR_TYPE_SERVER; // the server failed at something it should do + switch (r->reason) { + case SERVER_RECONFIGURE_INVALID_GEOMETRY: + case SERVER_RECONFIGURE_CACHE_TYPE_CHANGE: + type = ERROR_TYPE_INVALID_REQUEST; // the request asks for what no reconfigure is + break; + case SERVER_RECONFIGURE_LEARN_IN_PROGRESS: + case SERVER_RECONFIGURE_LEARNING_BOUND: + case SERVER_RECONFIGURE_FAN_OUT: + type = ERROR_TYPE_UNAVAILABLE; // the server's state refuses it now; asked again later it may not + break; + case SERVER_RECONFIGURE_PAGE_OUT_FAILED: + case SERVER_RECONFIGURE_CONTEXT_BUILD_FAILED: + break; + } + json j = format_error_response(std::string("reconfigure refused (") + server_reconfigure_refusal_name(r->reason) + "): " + r->detail, type); + j["applied"] = false; + j["refused"] = server_reconfigure_refusal_name(r->reason); + j["detail"] = r->detail; + if (r->slot >= 0) { + j["slot"] = r->slot; + } + return j; + } + const auto & a = std::get(outcome); + json paged_out = json::array(); + for (const auto & s : a.paged_out) { + paged_out.push_back({ { "slot", s.slot }, { "tokens_held", s.tokens_held }, { "tokens_fit", s.tokens_fit } }); + } + return json { + { "applied", true }, + { "rebuilt", a.rebuilt }, + { "live", { { "n_ctx", a.live.n_ctx }, { "n_seq", a.live.n_seq }, { "n_ctx_seq", a.n_ctx_seq } } }, + { "paged_out", paged_out }, + { "paged_in", a.paged_in }, + { "parked", a.parked }, + { "pause_ms", a.pause_ms }, + { "tokens_held", a.tokens_held }, + { "bytes_paged", a.bytes_paged }, + }; +} diff --git a/tools/server/server-reconfigure.h b/tools/server/server-reconfigure.h new file mode 100644 index 000000000000..4111717916bd --- /dev/null +++ b/tools/server/server-reconfigure.h @@ -0,0 +1,142 @@ +#pragma once + +// Continuum fork: RECONFIGURE, GROW OR SHRINK LIVE (SERVING-ENGINE.md, "S2 in detail"). The serving +// contexts take a new geometry (context size, sequence count) at a safe point between batched decode +// steps while the weights and their mmap stay loaded: every slot holding state is paged out through +// the one state mover (server-slot-page.h), unfinished generations included, the target and draft +// contexts are rebuilt, and every state that fits is paged back in, so each generation continues from +// its next token with its host-side sampler. Nothing is cancelled and nothing waits for a generation to +// end; new work waits in the queue. Grow and shrink are the same path. +// +// Shrink never truncates: a slot whose state the target cannot hold (an id past a smaller n_seq, or +// more tokens than a sequence of the target holds) is paged out whole and PARKED, named in the outcome; +// a later reconfigure places it back when the geometry holds it and its id is free. A KV cache TYPE +// change is not a reconfigure: it is the intentional cold start (#5188), and is refused by name. +// +// The logic lives here, against server_reconfigure_engine; the server context implements the engine +// (its context lifetimes and slots) and calls server_reconfigure from SERVER_TASK_TYPE_RECONFIGURE. + +#include "server-geometry.h" +#include "server-slot-page.h" + +#include +#include +#include +#include +#include +#include + +// a slot whose state the applied geometry cannot hold: paged out whole and parked, never cut +struct server_slot_refusal { + int32_t slot = -1; + int32_t tokens_held = 0; // tokens its state holds + int32_t tokens_fit = 0; // tokens the applied geometry had room for in its sequence +}; + +struct server_reconfigure_applied { + server_geometry live; // the geometry the contexts hold now + int32_t n_ctx_seq = 0; // tokens one sequence of it holds + std::vector paged_out; // slots this reconfigure parked + std::vector paged_in; // parked slots it placed back + std::vector parked; // every slot parked after it + // the cost: how long serving paused, against the state that crossed the rebuild + double pause_ms = 0.0; + int64_t tokens_held = 0; + uint64_t bytes_paged = 0; + bool rebuilt = false; // false: the geometry was already live, only parked slots moved +}; + +enum server_reconfigure_refusal_reason { + SERVER_RECONFIGURE_INVALID_GEOMETRY, // n_seq below 1, or fewer cells than sequences + SERVER_RECONFIGURE_CACHE_TYPE_CHANGE, // a KV type change is the intentional cold start, never a reconfigure + SERVER_RECONFIGURE_LEARN_IN_PROGRESS, // a learn's graph and accumulators live in the context until its restoration ends: retry after + SERVER_RECONFIGURE_LEARNING_BOUND, // the context holds shadow optimizers or the training workspace: they do not cross a rebuild + SERVER_RECONFIGURE_FAN_OUT, // a parent prompt is being copied into its child slots + SERVER_RECONFIGURE_PAGE_OUT_FAILED, // a slot's state could not be read; nothing changed + SERVER_RECONFIGURE_CONTEXT_BUILD_FAILED, // the target contexts could not be created; the previous geometry was rebuilt +}; + +const char * server_reconfigure_refusal_name(server_reconfigure_refusal_reason reason); + +struct server_reconfigure_refused { + server_reconfigure_refusal_reason reason; + std::string detail; + int32_t slot = -1; // the slot it concerns, when one does +}; + +using server_reconfigure_outcome = std::variant; + +// What reconfigure needs of the server. Slots are addressed by id, 0 .. live().n_seq - 1. +struct server_reconfigure_engine { + virtual ~server_reconfigure_engine() = default; + + virtual server_geometry live() const = 0; + virtual int32_t n_ctx_seq() const = 0; // tokens one sequence of the live contexts holds + virtual int32_t n_cells() const = 0; // cells the live target context allocated (padded, so >= n_ctx asked) + virtual bool kv_unified() const = 0; // every sequence draws on that one pool of cells + + // why the contexts cannot be rebuilt right now; nullopt when they can + virtual std::optional blocked() const = 0; + + virtual int32_t tokens_held(int32_t slot) const = 0; // tokens whose state its KV holds (0: none) + virtual bool busy (int32_t slot) const = 0; // it has a task, finished or not + + virtual bool page_out(int32_t slot, server_slot_page & page) = 0; + virtual bool page_in (int32_t slot, const server_slot_page & page) = 0; + + // park: the slot's host state (prompt, checkpoints, task, sampler) leaves service and a fresh, + // empty slot takes its id; unpark: the parked host state takes its id back + virtual void park (int32_t slot) = 0; + virtual void unpark(int32_t slot) = 0; + + // the contexts' lifetime apart from the weights': destroy frees the target and draft contexts and + // the speculative state; build creates them for `g` and sizes the slots to g.n_seq (the highest + // ids go, and a new id comes in empty) + virtual void destroy_contexts() = 0; + virtual bool build_contexts(const server_geometry & g) = 0; + + // after the states are in: samplers re-attached, drafting restarted, queued work offered the free slots + virtual void resume() = 0; +}; + +// a parked slot's page, with what its state held when it left +struct server_parked_slot { + server_slot_page page; + int32_t tokens_held = 0; + bool busy = false; +}; + +// what reconfigure keeps between calls: the parked pages and the counters /props reports +struct server_reconfigure_state { + std::map parked; + int64_t n_applied = 0; + int64_t n_refused = 0; + std::optional last; + + // THE COST, measured over every reconfigure that rebuilt the contexts, so the core can price the next + // one: how long serving paused, against the state that crossed (server_reconfigure_cost_json) + int64_t n_rebuilt = 0; + double pause_ms_total = 0.0; + double pause_ms_max = 0.0; + int64_t tokens_held_total = 0; + uint64_t bytes_paged_total = 0; +}; + +// the measured cost as /props reports it (the server adds the buffers the contexts hold now) +json server_reconfigure_cost_json(const server_reconfigure_state & state); + +server_reconfigure_outcome server_reconfigure(server_reconfigure_engine & engine, server_reconfigure_state & state, const server_geometry & target); + +// THE HTTP EDGE (POST /reconfigure): {"n_ctx": int, "n_seq": int} and, optionally, "cache_type_k" and +// "cache_type_v", which must name the live types. Either the geometry or the refusal. +std::variant server_reconfigure_from_json( + const json & body, const char * cache_type_k, const char * cache_type_v); + +json server_reconfigure_to_json(const server_reconfigure_outcome & outcome); + +struct server_task_result_reconfigure : server_task_result { + server_reconfigure_outcome outcome; + + virtual json to_json() override { return server_reconfigure_to_json(outcome); } + virtual bool is_error() override { return std::holds_alternative(outcome); } +}; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index c858c7e0bc26..3722d1527d0b 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -8,7 +8,7 @@ namespace server_schema { // llama.cpp-specific completion schema // -std::vector> make_llama_cmpl_schema(const common_params & params_base, task_params & params) { +std::vector> make_llama_cmpl_schema(const common_params & params_base, int32_t n_seq_live, task_params & params) { std::vector> fields; auto add = [&](field * f) { fields.emplace_back(f); @@ -60,7 +60,7 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_desc("Number of tokens after n_keep that may be discarded when shifting context (0 = half context)")); add((new field_num("n_cmpl", params.n_cmpl)) - ->set_hard_limits(1, params_base.n_parallel) + ->set_hard_limits(1, n_seq_live) ->add_alias("n") // alias "n" as fallback (OpenAI completions API) ->set_desc("Number of completions to generate. If the input has multiple prompts, total outputs will be N prompts times n_cmpl")); @@ -515,6 +515,7 @@ std::vector> make_llama_cmpl_schema(const common_params & task_params eval_llama_cmpl_schema( const llama_vocab * vocab, const common_params & params_base, + int32_t n_seq_live, const std::vector & logit_bias_eog, const json & data) { task_params params; @@ -539,7 +540,7 @@ task_params eval_llama_cmpl_schema( ctx.vocab = vocab; ctx.logit_bias_eog = &logit_bias_eog; - auto schema = make_llama_cmpl_schema(params_base, params); + auto schema = make_llama_cmpl_schema(params_base, n_seq_live, params); // eval all fields in the schema for (const auto & f : schema) { diff --git a/tools/server/server-schema.h b/tools/server/server-schema.h index d0a81431bc6b..e411b11373b2 100644 --- a/tools/server/server-schema.h +++ b/tools/server/server-schema.h @@ -91,13 +91,17 @@ struct field_nested : public field { virtual void eval(field_eval_context & ctx, const json & data) override; }; +// n_seq_live: the sequences the server serves now (POST /reconfigure changes it; params_base.n_parallel +// is the launch's count), which bounds n_cmpl std::vector> make_llama_cmpl_schema( const common_params & params_base, + int32_t n_seq_live, task_params & params); task_params eval_llama_cmpl_schema( const llama_vocab * vocab, const common_params & params_base, + int32_t n_seq_live, const std::vector & logit_bias_eog, const json & data); diff --git a/tools/server/server-slot-page.cpp b/tools/server/server-slot-page.cpp new file mode 100644 index 000000000000..ef8fc0316dda --- /dev/null +++ b/tools/server/server-slot-page.cpp @@ -0,0 +1,62 @@ +#include "server-slot-page.h" + +#include "server-common.h" + +server_slot_page_size server_slot_page_measure(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq) { + server_slot_page_size size; + size.main = llama_state_seq_get_size_ext(ctx_tgt, seq, LLAMA_STATE_SEQ_FLAGS_NONE); + size.drft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + return size; +} + +bool server_slot_page_read(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, server_prompt_data & data) { + const size_t n_main = llama_state_seq_get_data_ext(ctx_tgt, data.main.data(), data.main.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE); + if (n_main != data.main.size()) { + SRV_ERR("seq %d: target state gave %zu of %zu bytes\n", seq, n_main, data.main.size()); + return false; + } + if (data.drft.empty()) { + return true; + } + if (ctx_dft == nullptr) { + SRV_ERR("seq %d: a draft state was sized with no draft context\n", seq); + return false; + } + const size_t n_drft = llama_state_seq_get_data_ext(ctx_dft, data.drft.data(), data.drft.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE); + if (n_drft != data.drft.size()) { + SRV_ERR("seq %d: draft state gave %zu of %zu bytes\n", seq, n_drft, data.drft.size()); + return false; + } + return true; +} + +bool server_slot_page_write(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, const server_prompt_data & data) { + auto empty_both = [&]() { + llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq, -1, -1); + if (ctx_dft) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq, -1, -1); + } + }; + + const size_t n_main = llama_state_seq_set_data_ext(ctx_tgt, data.main.data(), data.main.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE); + if (n_main != data.main.size()) { + SRV_ERR("seq %d: failed to restore the target state with size %zu\n", seq, data.main.size()); + empty_both(); + return false; + } + if (data.drft.empty()) { + return true; + } + if (ctx_dft == nullptr) { + SRV_ERR("seq %d: the page holds a draft state and there is no draft context\n", seq); + empty_both(); + return false; + } + const size_t n_drft = llama_state_seq_set_data_ext(ctx_dft, data.drft.data(), data.drft.size(), seq, LLAMA_STATE_SEQ_FLAGS_NONE); + if (n_drft != data.drft.size()) { + SRV_WRN("seq %d: failed to restore the draft state with size %zu\n", seq, data.drft.size()); + empty_both(); + return false; + } + return true; +} diff --git a/tools/server/server-slot-page.h b/tools/server/server-slot-page.h new file mode 100644 index 000000000000..f759c28d50d4 --- /dev/null +++ b/tools/server/server-slot-page.h @@ -0,0 +1,101 @@ +#pragma once + +// Continuum fork: THE ONE STATE MOVER between a sequence in a llama_context and host memory +// (SERVING-ENGINE.md, S2). A page is a server_prompt_cache_state: the prompt's tokens with their +// host-memory checkpoints, the {main, drft} state bytes of the target and draft contexts, and the +// adapter set they were computed under. The prompt cache (server_slot::prompt_save, +// server_prompt_cache::load), the live reconfigure (server-reconfigure.h) and the saved page's +// checkpoint sidecar (/slots?action=save|restore) all move state through this file; nothing else +// calls llama_state_seq_get_data/set_data for a slot. + +#include "server-task.h" + +#include "llama.h" + +#include +#include +#include +#include + +using server_slot_page = server_prompt_cache_state; + +// the bytes a sequence's state takes, in the target context and in the draft (0 without one) +struct server_slot_page_size { + size_t main = 0; + size_t drft = 0; + + size_t total() const { return main + drft; } +}; + +server_slot_page_size server_slot_page_measure(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq); + +// page out: the sequence's target and draft state into `data`, whose vectors are already sized by +// server_slot_page_measure. Reading never changes the context. False if a context gave back fewer +// bytes than measured (the page is then not a copy of the sequence and must not be used). +bool server_slot_page_read(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, server_prompt_data & data); + +// page in: `data` into the sequence. False if a context could not take it (no room in its cache, a +// draft page with no draft context); the sequence then holds nothing in either context, so a +// failed page-in never leaves half a state behind. +bool server_slot_page_write(llama_context * ctx_tgt, llama_context * ctx_dft, llama_seq_id seq, const server_prompt_data & data); + +// THE CHECKPOINT SIDECAR a saved page carries beside its state file (.ckpt): 'CKPT', version 1, +// the count, then each checkpoint's positions, sizes and bytes. The size quote and the write share +// the writer (a counting `put` and a file `put`); the restore reads it back with the reader. +constexpr uint32_t SERVER_SLOT_PAGE_CKPT_MAGIC = 0x434B5054u; // 'CKPT' +constexpr uint32_t SERVER_SLOT_PAGE_CKPT_VERSION = 1; + +// put(const void * ptr, size_t n) +template +void server_slot_page_put_checkpoints(const std::vector & kept, Put && put) { + const uint32_t magic = SERVER_SLOT_PAGE_CKPT_MAGIC; + const uint32_t version = SERVER_SLOT_PAGE_CKPT_VERSION; + 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()); + } +} + +// get(void * ptr, size_t n) -> bool (false on a short read); on_checkpoint() after each one read. +// False at the first short read, a foreign header or an implausible size; `out` is then partial and +// the caller drops it. +template +bool server_slot_page_get_checkpoints(Get && get, std::list & out, OnCheckpoint && on_checkpoint) { + uint32_t magic = 0, version = 0, n_ckpt = 0; + bool ok = get(&magic, sizeof(magic)) && get(&version, sizeof(version)) && get(&n_ckpt, sizeof(n_ckpt)); + // a page holds at most a few hundred checkpoints; 4096 refuses a corrupt count before allocating + ok = ok && magic == SERVER_SLOT_PAGE_CKPT_MAGIC && version == SERVER_SLOT_PAGE_CKPT_VERSION && n_ckpt <= 4096; + for (uint32_t i = 0; ok && i < n_ckpt; ++i) { + int64_t n_tokens = 0; int32_t pos_min = 0, pos_max = 0; + uint64_t sz_tgt = 0, sz_dft = 0, sz_spec = 0; + ok = get(&n_tokens, sizeof(n_tokens)) && get(&pos_min, sizeof(pos_min)) && get(&pos_max, sizeof(pos_max)) + && get(&sz_tgt, sizeof(sz_tgt)) && get(&sz_dft, sizeof(sz_dft)) && get(&sz_spec, sizeof(sz_spec)); + // a checkpoint is at most one sequence state; refuse absurd sizes + ok = ok && sz_tgt < (1ull << 36) && sz_dft < (1ull << 36) && sz_spec < (1ull << 30); + if (!ok) { + break; + } + auto & c = out.emplace_back(); + c.update_pos(n_tokens, pos_min, pos_max); + c.id_task = -1; + c.data_tgt.resize(sz_tgt); + c.data_dft.resize(sz_dft); + c.data_spec.resize(sz_spec); + ok = get(c.data_tgt.data(), sz_tgt) && get(c.data_dft.data(), sz_dft) && get(c.data_spec.data(), sz_spec); + if (ok) { + on_checkpoint(); + } + } + return ok; +} diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index b531ba041379..a41e33e61190 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1,4 +1,5 @@ #include "server-task.h" +#include "server-slot-page.h" #include "build-info.h" #include "server-chat.h" @@ -1885,39 +1886,15 @@ bool server_prompt_cache::load(server_prompt & prompt, const lora_active_set & l if (it_best != states.end()) { SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best); - { - auto & data = it_best->data.main; - - const size_t size = data.size(); - const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0); - if (n != size) { - SRV_ERR("failed to restore state with size %zu\n", size); - - return false; - } - - data.clear(); - data.shrink_to_fit(); + // the one state mover (server-slot-page.h) + if (!server_slot_page_write(ctx_tgt, ctx_dft, id_slot, it_best->data)) { + return false; } - { - auto & data = it_best->data.drft; - - if (!data.empty()) { - GGML_ASSERT(ctx_dft); - - const size_t size = data.size(); - const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0); - if (n != size) { - SRV_WRN("failed to restore state with size %zu\n", size); - - return false; - } - - data.clear(); - data.shrink_to_fit(); - } - } + it_best->data.main.clear(); + it_best->data.main.shrink_to_fit(); + it_best->data.drft.clear(); + it_best->data.drft.shrink_to_fit(); prompt = std::move(it_best->prompt); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 652697ba4d2b..bc40b7560f6d 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -12,6 +12,7 @@ // TODO: prevent including the whole server-common.h as we only use server_tokens #include "server-common.h" +#include "server-geometry.h" // One immutable operation identity with cheap counters. The route holds at most @@ -63,6 +64,7 @@ enum server_task_type { SERVER_TASK_TYPE_UNLOAD_LORA, SERVER_TASK_TYPE_LEARN, SERVER_TASK_TYPE_LEARN_GRANTS, + SERVER_TASK_TYPE_RECONFIGURE, // Continuum fork: grow or shrink the contexts live (server-reconfigure.h) }; // TODO: change this to more generic "response_format" to replace the "format_response_*" in server-common @@ -293,6 +295,9 @@ struct server_task { learn_request learn; json restore_checkpoint; + // used by SERVER_TASK_TYPE_RECONFIGURE: the geometry to take + server_geometry reconfigure; + server_task() = default; server_task(server_task_type type) : type(type) {} diff --git a/tools/server/server-test-faults.h b/tools/server/server-test-faults.h index a46bfc8bd86e..336050032297 100644 --- a/tools/server/server-test-faults.h +++ b/tools/server/server-test-faults.h @@ -20,6 +20,10 @@ struct server_test_faults { int32_t prompt_checkpoints = 0; // LLAMA_SERVER_TEST_FAIL_CHECKPOINTS: the next N prompt checkpoint copies fail to allocate int32_t spec_checkpoints = 0; // LLAMA_SERVER_TEST_FAIL_SPEC_CHECKPOINTS: the next N speculative checkpoint copies fail to allocate int32_t prompt_fills = 0; // LLAMA_SERVER_TEST_FAIL_PROMPT_FILLS: the next N prompt fills throw after adding their tokens + // LLAMA_SERVER_TEST_STEP_DELAY_MS: every decode step that ran a batch then waits N ms, so a test + // can act while generations are in progress on a model that would otherwise finish in milliseconds + // (the reconfigure test); a pace, not a failure + int32_t step_delay_ms = 0; static server_test_faults from_env() { server_test_faults faults; @@ -27,12 +31,13 @@ struct server_test_faults { faults.prompt_checkpoints = read("LLAMA_SERVER_TEST_FAIL_CHECKPOINTS"); faults.spec_checkpoints = read("LLAMA_SERVER_TEST_FAIL_SPEC_CHECKPOINTS"); faults.prompt_fills = read("LLAMA_SERVER_TEST_FAIL_PROMPT_FILLS"); + faults.step_delay_ms = read("LLAMA_SERVER_TEST_STEP_DELAY_MS"); #endif return faults; } bool any() const { - return prompt_checkpoints || spec_checkpoints || prompt_fills; + return prompt_checkpoints || spec_checkpoints || prompt_fills || step_delay_ms; } // true when this occurrence is one to fail, counting it diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 098cc26c2004..7448ebdc7e76 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -1055,6 +1055,11 @@ void server_trainer::run(json req, examples_data ex) { // the single-context gradient): one step per window on the window's whole loss /*walk_exact =*/ req.value("exact", false), /*walk_horizon =*/ (uint32_t) req.value("walk_horizon", (int64_t) 0), + // the trainer's own context is outside serving's leases: no admission callbacks + /*static_admission =*/ nullptr, + /*static_admission_data =*/ nullptr, + /*workspace_admission =*/ nullptr, + /*workspace_admission_data =*/ nullptr, }; // THE BUDGET COVERS THE WHOLE RUN (continuum, the 5090, 2026-10-11): the training context's own // memory (its Q8_0 K/V for the whole window, and any recurrent state) is allocated when it is diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 4a4ec63ab39f..201aa65c6eec 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -198,6 +198,7 @@ int llama_server(common_params & params, int argc, char ** argv) { // note: routes.get_health stays the same routes.get_metrics = models_routes->proxy_get; routes.post_props = models_routes->proxy_post; + routes.post_reconfigure = models_routes->proxy_post; routes.post_completions = models_routes->proxy_post; routes.post_completions_oai = models_routes->proxy_post; routes.post_chat_completions = models_routes->proxy_post; @@ -250,6 +251,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.get ("/metrics", ex_wrapper(routes.get_metrics)); ctx_http.get ("/props", ex_wrapper(routes.get_props)); ctx_http.post("/props", ex_wrapper(routes.post_props)); + ctx_http.post("/reconfigure", ex_wrapper(routes.post_reconfigure)); ctx_http.get ("/models", ex_wrapper(routes.get_models)); ctx_http.get ("/v1/models", ex_wrapper(routes.get_models)); ctx_http.post("/completion", ex_wrapper(routes.post_completions)); // legacy diff --git a/tools/server/tests/unit/test_reconfigure.py b/tools/server/tests/unit/test_reconfigure.py new file mode 100644 index 000000000000..6ad35a9ec8ea --- /dev/null +++ b/tools/server/tests/unit/test_reconfigure.py @@ -0,0 +1,380 @@ +import os +import tempfile +import threading +import time +import pytest +from utils import * + +# POST /reconfigure (Continuum fork, server-reconfigure.h): the serving contexts take a new geometry +# (n_ctx, n_seq) live, between batched decode steps, while the weights stay loaded. Every slot holding +# state is paged out, unfinished generations included; the contexts are rebuilt; every state that fits +# is paged back in and each generation continues from its next token. A state the target cannot hold +# is parked whole, named in the outcome, and placed back by a later reconfigure that holds it. +# +# what these catch: a reconfigure that reloads the weights, cuts or loses a sequence's state, changes +# a generation in flight, drops a queued request, or shrinks without naming the slot it paged out. + +server: ServerProcess + +PROMPTS = [ + "Once upon a time, there was a little girl named", + "The big brown dog ran to the park and", + "One day, a small bird found a shiny", + "Tom and his mom went to the store to buy", +] + +N_CTX = 1024 +N_SLOTS = 4 + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_ctx = N_CTX + server.n_slots = N_SLOTS + server.temperature = 0.0 + server.server_slots = True + server.cache_ram = 0 # a state that survives must survive in its slot, never restored from the RAM cache + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + # pytest shows a test's output only when it fails: then the server's own log says why + print(log_text()) + + +def log_text() -> str: + with open(server.log_path) as log: + return log.read() + + +def start_with_step_delay(monkeypatch, delay_ms: int): + monkeypatch.setenv("LLAMA_SERVER_TEST_STEP_DELAY_MS", str(delay_ms)) + server.start() + deadline = time.time() + 5 + while "test faults:" not in log_text(): + if time.time() > deadline: + pytest.skip("llama-server was built without -DLLAMA_SERVER_TEST_FAULTS=ON") + time.sleep(0.05) + + +def reconfigure(n_ctx: int, n_seq: int, **extra): + return server.make_request("POST", "/reconfigure", data={"n_ctx": n_ctx, "n_seq": n_seq, **extra}) + + +def applied(n_ctx: int, n_seq: int) -> dict: + res = reconfigure(n_ctx, n_seq) + assert res.status_code == 200, res.body + assert res.body["applied"] is True + assert res.body["live"]["n_seq"] == n_seq + return res.body + + +def tokenize(text: str) -> list: + res = server.make_request("POST", "/tokenize", data={"content": text, "add_special": True}) + assert res.status_code == 200 + return res.body["tokens"] + + +def complete(tokens: list, id_slot: int, n_predict: int) -> dict: + res = server.make_request("POST", "/completion", data={ + "prompt": tokens, + "id_slot": id_slot, + "n_predict": n_predict, + "temperature": 0.0, + "cache_prompt": True, + "return_tokens": True, + "ignore_eos": True, # every reply its full length, so every state holds what the test counts + }) + assert res.status_code == 200, res.body + return res.body + + +def hold_states(n_predict: int) -> list: + """every slot holds its prompt and n_predict generated tokens; returns what each slot holds""" + held = [] + for i, prompt in enumerate(PROMPTS): + tokens = tokenize(prompt) + res = complete(tokens, i, n_predict) + assert len(res["tokens"]) == n_predict + held.append(tokens + res["tokens"]) + return held + + +def continuations(held: list, n_predict: int) -> list: + """each slot's next tokens from the state it holds, with how much of it was reused""" + out = [] + for i, tokens in enumerate(held): + res = complete(tokens, i, n_predict) + out.append((res["tokens"], res["timings"]["cache_n"])) + return out + + +def weight_loads() -> int: + props = server.make_request("GET", "/props") + assert props.status_code == 200 + return props.body["reconfigure"]["n_weight_loads"] + + +def assert_never_reloaded(): + assert weight_loads() == 1 + # the load's own line, where the log level prints it (the counter above is the gate) + assert log_text().count("loading model '") <= 1 + + +def rss_kib() -> int | None: + # the footprint, where the platform shows it (Linux); a measurement printed, not a gate + try: + with open(f"/proc/{server.process.pid}/status") as status: + for line in status: + if line.startswith("VmRSS:"): + return int(line.split()[1]) + except OSError: + return None + return None + + +def test_generations_in_progress_finish_identical_across_a_reconfigure(monkeypatch): + # one batch takes all four prompts, so the control and the reconfigured run group every step alike + server.n_batch = 512 + # a pace of 40 ms per decode step keeps 48-token generations in progress for about 2 s + start_with_step_delay(monkeypatch, 40) + request = { + "prompt": PROMPTS, + "n_predict": 48, + "temperature": 0.0, + "cache_prompt": False, + "return_tokens": True, + "ignore_eos": True, + } + + control = server.make_request("POST", "/completion", data=request) + assert control.status_code == 200 + expected = [r["tokens"] for r in control.body] + assert all(len(t) == 48 for t in expected) + + results = {} + def run(): + results["res"] = server.make_request("POST", "/completion", data=request) + thread = threading.Thread(target=run) + thread.start() + + # all four generating, each part way through its reply + deadline = time.time() + 10 + while True: + slots = server.make_request("GET", "/slots").body + decoding = [s for s in slots if s["is_processing"] and s["next_token"][0]["n_decoded"] >= 4] + if len(decoding) == N_SLOTS: + break + assert time.time() < deadline, slots + time.sleep(0.02) + + grown = applied(2 * N_CTX, N_SLOTS) + assert grown["rebuilt"] is True + assert grown["paged_out"] == [] + assert grown["tokens_held"] > 0 + shrunk = applied(N_CTX, N_SLOTS) + assert shrunk["paged_out"] == [] + + thread.join(timeout=30) + assert not thread.is_alive() + res = results["res"] + assert res.status_code == 200 + assert [r["tokens"] for r in res.body] == expected + assert_never_reloaded() + + +def test_many_grow_and_shrink_cycles_keep_every_state_and_never_reload(): + server.start() + n_held = 24 + + # control: each slot's next tokens from its state, with no reconfigure + control = continuations(hold_states(n_held), 8) + assert all(cache_n > 0 for _, cache_n in control) + + # the same states again, then 100 cycles alternating grow and shrink of both n_ctx and n_seq + held = hold_states(n_held) + rss_before = rss_kib() + geometries = [(3 * N_CTX, N_SLOTS + 2), (N_CTX, N_SLOTS)] + for cycle in range(100): + n_ctx, n_seq = geometries[cycle % 2] + body = applied(n_ctx, n_seq) + assert body["paged_out"] == [], body + assert body["parked"] == [] + rss_after = rss_kib() + print(f"footprint over 100 reconfigures: VmRSS {rss_before} -> {rss_after} KiB") + + props = server.make_request("GET", "/props").body + assert props["reconfigure"]["n_applied"] == 100 + assert props["total_slots"] == N_SLOTS + assert_never_reloaded() + + # each state came back byte for byte: the same reuse and the same next tokens as the control + assert continuations(held, 8) == control + + +def test_requests_queued_across_reconfigures_all_complete(monkeypatch): + start_with_step_delay(monkeypatch, 20) + n_requests = 12 + results = [None] * n_requests + + def run(i): + results[i] = server.make_request("POST", "/completion", data={ + "prompt": PROMPTS[i % len(PROMPTS)], + "n_predict": 16, + "temperature": 0.0, + "return_tokens": True, + "ignore_eos": True, + }) + threads = [threading.Thread(target=run, args=(i,)) for i in range(n_requests)] + for t in threads: + t.start() + + # alternate a grow and a shrink while they run (a shrink parks a busy slot past its n_seq, the next + # grow places it back); end on the grow, so nothing stays parked + geometries = [(3 * N_CTX, N_SLOTS + 2), (N_CTX, N_SLOTS)] + n_reconfigures = 0 + while any(t.is_alive() for t in threads) and n_reconfigures < 200: + applied(*geometries[n_reconfigures % 2]) + n_reconfigures += 1 + time.sleep(0.01) + final = applied(*geometries[0]) + assert final["parked"] == [] + + for t in threads: + t.join(timeout=30) + assert not t.is_alive() + assert n_reconfigures > 0 + for res in results: + assert res.status_code == 200, res.body + assert len(res.body["tokens"]) == 16 + assert_never_reloaded() + + +def test_a_smaller_n_seq_pages_out_the_removed_slots_and_names_them(): + server.start() + n_held = 24 + control = continuations(hold_states(n_held), 8) + held = hold_states(n_held) + + shrunk = applied(N_CTX, 2) + named = {s["slot"]: s for s in shrunk["paged_out"]} + assert sorted(named) == [2, 3] + for slot in (2, 3): + # what its KV holds: the prompt and the reply, but for the last token, which was sampled and never decoded + assert len(held[slot]) - 1 <= named[slot]["tokens_held"] <= len(held[slot]) + assert named[slot]["tokens_fit"] == 0 + assert shrunk["parked"] == [2, 3] + assert server.make_request("GET", "/props").body["total_slots"] == 2 + + grown = applied(N_CTX, N_SLOTS) + assert sorted(grown["paged_in"]) == [2, 3] + assert grown["parked"] == [] + assert continuations(held, 8) == control + assert_never_reloaded() + + +def test_a_pool_too_small_for_every_state_pages_out_the_rest_and_names_them(): + server.kv_unified = True # one pool of cells every sequence draws on + server.start() + n_held = 100 + control = continuations(hold_states(n_held), 8) + held = hold_states(n_held) + + # 256 cells (the context's padding unit) hold two of these ~110-token states, not four + shrunk = applied(256, N_SLOTS) + assert len(shrunk["paged_out"]) >= 1 + for s in shrunk["paged_out"]: + assert s["tokens_fit"] <= s["tokens_held"] + paged = sorted(s["slot"] for s in shrunk["paged_out"]) + assert shrunk["parked"] == paged + + grown = applied(N_CTX, N_SLOTS) + assert sorted(grown["paged_in"]) == paged + assert grown["parked"] == [] + assert continuations(held, 8) == control + assert_never_reloaded() + + +def test_a_cache_type_change_and_a_bad_geometry_are_refused_by_name(): + server.start() + res = reconfigure(N_CTX, N_SLOTS, cache_type_k="q4_0") + assert res.status_code == 400 + assert res.body["error"]["refused"] == "cache_type_change" + + res = server.make_request("POST", "/reconfigure", data={"n_ctx": N_CTX}) + assert res.status_code == 400 + assert res.body["error"]["refused"] == "invalid_geometry" + + res = reconfigure(N_CTX, 0) + assert res.status_code == 400 + assert res.body["error"]["refused"] == "invalid_geometry" + + # the live types, named, are not a change + props = server.make_request("GET", "/props").body + assert props["reconfigure"]["n_applied"] == 0 + same = reconfigure(N_CTX, N_SLOTS, cache_type_k="f16", cache_type_v="f16") + assert same.status_code == 200, same.body + assert same.body["rebuilt"] is False + + +def test_a_request_for_a_parked_slot_waits_for_it_and_never_runs_on_another(): + # the upstream wrap-around (id % n_slots) ran a request for slot 3 on slot 1 after a shrink to 2: + # another mind's state, silently. Now it waits for its own slot; an id no slot has is refused + server.start() + n_held = 24 + control = continuations(hold_states(n_held), 8) + held = hold_states(n_held) + + applied(N_CTX, 2) + slot1_before = [s for s in server.make_request("GET", "/slots").body if s["id"] == 1][0]["kv_pos_max"] + + results = {} + def run(): + results["res"] = complete(held[3], 3, 8) + thread = threading.Thread(target=run) + thread.start() + time.sleep(1.0) + assert thread.is_alive() # waiting for slot 3, not served by slot 1 + slot1_after = [s for s in server.make_request("GET", "/slots").body if s["id"] == 1][0]["kv_pos_max"] + assert slot1_after == slot1_before + + unknown = server.make_request("POST", "/completion", data={"prompt": "hello", "id_slot": 9, "n_predict": 1}) + assert unknown.status_code == 400 + + grown = applied(N_CTX, N_SLOTS) + assert 3 in grown["paged_in"] + thread.join(timeout=30) + assert not thread.is_alive() + res = results["res"] + # it ran on slot 3, from slot 3's own state: the control's tokens and reuse + assert (res["tokens"], res["timings"]["cache_n"]) == control[3] + + +def test_n_cmpl_is_bounded_by_the_live_slot_count(): + server.start() + # n_cmpl children each take a slot: the bound is the slots served now, not the launch's + request = {"prompt": PROMPTS[0], "n_predict": 4, "temperature": 0.0} + applied(N_CTX * 2, N_SLOTS + 2) + wide = server.make_request("POST", "/completion", data={**request, "n_cmpl": N_SLOTS + 2}) + assert wide.status_code == 200, wide.body + applied(N_CTX, 2) + narrow = server.make_request("POST", "/completion", data={**request, "n_cmpl": 3}) + assert narrow.status_code == 400, narrow.body + + +def test_props_report_the_measured_cost_of_a_reconfigure(): + server.start() + hold_states(24) + applied(2 * N_CTX, N_SLOTS) + applied(N_CTX, N_SLOTS) + applied(N_CTX, N_SLOTS) # already live: no rebuild, no cost + cost = server.make_request("GET", "/props").body["reconfigure"]["cost"] + assert cost["n_rebuilt"] == 2 + assert cost["pause_ms_total"] > 0 + assert cost["pause_ms_max"] <= cost["pause_ms_total"] + assert cost["tokens_held_total"] > 0 + assert cost["bytes_paged_total"] > 0 + assert cost["pause_ms_per_1k_tokens"] > 0 + assert cost["last_applied"]["rebuilt"] is False + assert len(cost["context_buffers"]) > 0