diff --git a/common/arg.cpp b/common/arg.cpp index eec92f5af..f733c2470 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3708,6 +3708,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.models_max = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MAX")); + add_opt(common_arg( + {"--models-memory-margin"}, "N", + string_format("for router server, MiB of device memory to keep free per device, unload LRU models until a new one fits (default: %d, 0 = disabled)", params.models_memory_margin), + [](common_params & params, int value) { + params.models_memory_margin = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MODELS_MEMORY_MARGIN")); add_opt(common_arg( {"--models-autoload"}, {"--no-models-autoload"}, diff --git a/common/common.cpp b/common/common.cpp index f69f4db93..ca85f66d2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1253,44 +1253,50 @@ common_gguf_info common_get_gguf_info(const std::string & fname) { return info; } +void common_fit_apply(common_params & params, llama_model_params & mparams, llama_context_params & cparams) { + if (!params.fit_params) { + return; + } + + 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"); + + // the draft context is created from the same base params and follows the main context, fit both together + 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(); + + common_params params_dft = common_base_params_to_speculative(params); + + auto mparams_dft = common_model_params_to_llama(params_dft); + auto cparams_dft = common_context_params_to_llama(params_dft); + if (spec_mtp) { + cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + cparams_dft.n_rs_seq = 0; + + const common_fit_extra_model extra = { + /*.path_model =*/ params_dft.model.path.c_str(), + /*.mparams =*/ &mparams_dft, + /*.cparams =*/ &cparams_dft, + /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model + }; + + common_fit_params(params.model.path.c_str(), &mparams, &cparams, + params.tensor_split, + params.tensor_buft_overrides.data(), + params.fit_params_target.data(), + params.fit_params_min_ctx, + has_draft || spec_mtp ? &extra : nullptr, + params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); +} + 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) { - 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"); - - // the draft context is created from the same base params and follows the main context, fit both together - 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(); - - common_params params_dft = common_base_params_to_speculative(params); - - auto mparams_dft = common_model_params_to_llama(params_dft); - auto cparams_dft = common_context_params_to_llama(params_dft); - if (spec_mtp) { - cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - } - cparams_dft.n_rs_seq = 0; - - const common_fit_extra_model extra = { - /*.path_model =*/ params_dft.model.path.c_str(), - /*.mparams =*/ &mparams_dft, - /*.cparams =*/ &cparams_dft, - /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model - }; - - common_fit_params(params.model.path.c_str(), &mparams, &cparams, - params.tensor_split, - params.tensor_buft_overrides.data(), - params.fit_params_target.data(), - params.fit_params_min_ctx, - has_draft || spec_mtp ? &extra : nullptr, - params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); - } + common_fit_apply(params, mparams, cparams); llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); if (model == NULL) { diff --git a/common/common.h b/common/common.h index e0f226fae..a0d960c5e 100644 --- a/common/common.h +++ b/common/common.h @@ -701,6 +701,7 @@ struct common_params { std::string models_dir = ""; // directory containing models for the router server std::string models_preset = ""; // directory containing model presets for the router server int models_max = 4; // maximum number of models to load simultaneously + int models_memory_margin = 0; // MiB of device memory to keep free when loading models (0 = no memory accounting) bool models_autoload = true; // automatically load models when requested via the router server std::string models_preset_hf = ""; // show a warning about remote presets on router loaded (if not empty) @@ -1018,6 +1019,9 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode 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); +// apply --fit to mparams and cparams, no-op if params.fit_params is false +void common_fit_apply(common_params & params, struct llama_model_params & mparams, struct llama_context_params & cparams); + // clear LoRA adapters from context, then apply new list of adapters void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora); diff --git a/common/fit.cpp b/common/fit.cpp index 7a0300829..55b457e10 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -1,6 +1,8 @@ #include "fit.h" +#include "common.h" #include "json.h" +#include "llama-cpp.h" #include "log.h" #include "../src/llama-ext.h" @@ -27,6 +29,32 @@ class common_params_fit_exception : public std::runtime_error { using std::runtime_error::runtime_error; }; +// RAII: send llama logs below min_level to the debug log +struct common_fit_log_scope { + common_fit_log_scope(ggml_log_level min_level) : min_level(min_level) { + llama_log_get(&original_logger.callback, &original_logger.user_data); + llama_log_set([](ggml_log_level level, const char * text, void * user_data) { + const common_fit_log_scope * self = (const common_fit_log_scope *) user_data; + const ggml_log_level level_eff = level >= self->min_level ? level : GGML_LOG_LEVEL_DEBUG; + self->original_logger.callback(level_eff, text, self->original_logger.user_data); + }, this); + } + + ~common_fit_log_scope() { + llama_log_set(original_logger.callback, original_logger.user_data); + } + + common_fit_log_scope(const common_fit_log_scope &) = delete; + common_fit_log_scope & operator=(const common_fit_log_scope &) = delete; + +private: + struct { + ggml_log_callback callback; + void * user_data; + } original_logger; + ggml_log_level min_level; +}; + static std::vector common_get_device_memory_data_impl( const char * path_model, const llama_model_params * mparams, @@ -36,22 +64,7 @@ static std::vector common_get_device_memory_data_impl( uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, ggml_log_level log_level) { - struct user_data_t { - struct { - ggml_log_callback callback; - void * user_data; - } original_logger; - ggml_log_level min_level; // prints below this log level go to debug log - }; - user_data_t ud; - llama_log_get(&ud.original_logger.callback, &ud.original_logger.user_data); - ud.min_level = log_level; - - llama_log_set([](ggml_log_level level, const char * text, void * user_data) { - const user_data_t * ud = (const user_data_t *) user_data; - const ggml_log_level level_eff = level >= ud->min_level ? level : GGML_LOG_LEVEL_DEBUG; - ud->original_logger.callback(level_eff, text, ud->original_logger.user_data); - }, &ud); + common_fit_log_scope log_scope(log_level); llama_model_params mparams_copy = *mparams; mparams_copy.no_alloc = true; @@ -59,14 +72,12 @@ static std::vector common_get_device_memory_data_impl( llama_model * model = llama_model_load_from_file(path_model, mparams_copy); if (model == nullptr) { - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); throw std::runtime_error("failed to load model"); } llama_context * ctx = llama_init_from_model(model, *cparams); if (ctx == nullptr) { llama_model_free(model); - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); throw std::runtime_error("failed to create llama_context from model"); } @@ -148,11 +159,81 @@ static std::vector common_get_device_memory_data_impl( llama_free(ctx); llama_model_free(model); - llama_log_set(ud.original_logger.callback, ud.original_logger.user_data); return ret; } +std::vector common_fit_measure( + const char * path_model, + const llama_model_params * mparams, + const llama_context_params * cparams, + uint32_t * n_ctx_train, + uint32_t * n_ctx, + ggml_log_level log_level) { + common_fit_log_scope log_scope(log_level); + + llama_model_params mparams_copy = *mparams; + mparams_copy.no_alloc = true; + mparams_copy.load_mode = LLAMA_LOAD_MODE_NONE; + + llama_model_ptr model(llama_model_load_from_file(path_model, mparams_copy)); + if (!model) { + throw std::runtime_error("failed to load model"); + } + + llama_context_ptr ctx(llama_init_from_model(model.get(), *cparams)); + if (!ctx) { + throw std::runtime_error("failed to create llama_context from model"); + } + + if (n_ctx_train) { + *n_ctx_train = llama_model_n_ctx_train(model.get()); + } + if (n_ctx) { + *n_ctx = llama_n_ctx(ctx.get()); + } + + return common_memory_breakdown_rows(ctx.get()); +} + +common_memory_reports common_measure_memory(common_params & params) { + const ggml_log_level log_level = params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR; + const char * path_model = params.model.path.c_str(); + + common_memory_reports res; + + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + // max: what --fit picks now, no-op with -fit off + { + auto mparams_fit = mparams; + auto cparams_fit = cparams; + common_fit_apply(params, mparams_fit, cparams_fit); + res.info_max = common_fit_measure(path_model, &mparams_fit, &cparams_fit, &res.n_ctx_train, &res.n_ctx_max, log_level); + } + + // min: the smallest context --fit may pick (--fit only changes n_ctx when it is 0) + uint32_t n_ctx_min = cparams.n_ctx; + if (n_ctx_min == 0 && params.fit_params) { + const uint32_t n_streams = cparams.kv_unified ? 1 : std::max(1, cparams.n_seq_max); + // -c 0 sets fit_params_min_ctx to UINT32_MAX, cap at n_ctx_train like --fit + n_ctx_min = (uint32_t) std::min( + uint64_t(params.fit_params_min_ctx) * n_streams, + uint64_t(res.n_ctx_train) * n_streams); + } + if (n_ctx_min == 0 || n_ctx_min >= res.n_ctx_max) { + // nothing to reduce, or -fit off with the trained context + res.n_ctx_min = res.n_ctx_max; + res.info_min = res.info_max; + } else { + cparams.n_ctx = n_ctx_min; + res.info_min = common_fit_measure(path_model, &mparams, &cparams, nullptr, &res.n_ctx_min, log_level); + } + + return res; +} + common_device_memory_data_vec common_get_device_memory_data( const char * path_model, const llama_model_params * mparams, @@ -903,8 +984,7 @@ enum common_params_fit_status common_fit_params( return status; } -void common_memory_breakdown_print(const struct llama_context * ctx) { - //const auto & devices = ctx->get_model().devices; +std::vector common_memory_breakdown_rows(const struct llama_context * ctx) { const auto * model = llama_get_model(ctx); std::vector devices; @@ -914,18 +994,6 @@ void common_memory_breakdown_print(const struct llama_context * ctx) { llama_memory_breakdown memory_breakdown = llama_get_memory_breakdown(ctx); - std::vector> table_data; - table_data.reserve(devices.size()); - - // same data as the table below, for --log-jsonl consumers - common_json rows = common_json::array(); - const std::string template_header = "%s: | %s | %s %s %s %s %s %s %s |\n"; - const std::string template_gpu = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n"; - const std::string template_other = "%s: | %s | %s %s %s = %s + %s + %s %s |\n"; - - table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"}); - - constexpr size_t MiB = 1024 * 1024; const std::vector desc_prefixes_strip = {"NVIDIA ", "GeForce ", "Tesla ", "AMD ", "Radeon ", "Instinct "}; // track seen buffer types to avoid double counting: @@ -964,104 +1032,130 @@ void common_memory_breakdown_print(const struct llama_context * ctx) { } } - // print memory breakdown for each device: + std::vector rows; + rows.reserve(devices.size() + 1 + memory_breakdown.size()); + for (size_t i = 0; i < devices.size(); i++) { ggml_backend_dev_t dev = devices[i]; - llama_memory_breakdown_data mb = mb_dev[i]; - const std::string name = ggml_backend_dev_name(dev); - std::string desc = ggml_backend_dev_description(dev); + common_memory_breakdown_row row; + row.kind = common_memory_breakdown_row::DEVICE; + row.name = ggml_backend_dev_name(dev); + row.description = ggml_backend_dev_description(dev); for (const std::string & prefix : desc_prefixes_strip) { - if (desc.length() >= prefix.length() && desc.substr(0, prefix.length()) == prefix) { - desc = desc.substr(prefix.length()); + if (row.description.length() >= prefix.length() && row.description.substr(0, prefix.length()) == prefix) { + row.description = row.description.substr(prefix.length()); } } size_t free, total; ggml_backend_dev_memory(dev, &free, &total); - - const size_t self = mb.model + mb.context + mb.compute; - const int64_t unaccounted = static_cast(total) - static_cast(free) - static_cast(self); - - table_data.push_back({ - template_gpu, - " - " + name + " (" + desc + ")", - std::to_string(total / MiB), - std::to_string(free / MiB), - std::to_string(self / MiB), - std::to_string(mb.model / MiB), - std::to_string(mb.context / MiB), - std::to_string(mb.compute / MiB), - std::to_string(unaccounted / static_cast(MiB))}); - - rows.push_back({ - {"kind", "device"}, - {"name", name}, - {"description", desc}, - {"total", total / MiB}, - {"free", free / MiB}, - {"self", self / MiB}, - {"model", mb.model / MiB}, - {"context", mb.context / MiB}, - {"compute", mb.compute / MiB}, - {"unaccounted", unaccounted / static_cast(MiB)}, - }); + row.total = total; + row.free = free; + row.model = mb_dev[i].model; + row.context = mb_dev[i].context; + row.compute = mb_dev[i].compute; + rows.push_back(std::move(row)); } - // print memory breakdown for host: { - const size_t self = mb_host.model + mb_host.context + mb_host.compute; - table_data.push_back({ - template_other, - " - Host", - "", // total - "", // free - std::to_string(self / MiB), - std::to_string(mb_host.model / MiB), - std::to_string(mb_host.context / MiB), - std::to_string(mb_host.compute / MiB), - ""}); // unaccounted - - rows.push_back({ - {"kind", "host"}, - {"name", "Host"}, - {"self", self / MiB}, - {"model", mb_host.model / MiB}, - {"context", mb_host.context / MiB}, - {"compute", mb_host.compute / MiB}, - }); + common_memory_breakdown_row row; + row.kind = common_memory_breakdown_row::HOST; + row.name = "Host"; + row.model = mb_host.model; + row.context = mb_host.context; + row.compute = mb_host.compute; + rows.push_back(std::move(row)); } - // print memory breakdown for all remaining buffer types: + // all remaining buffer types: for (const auto & buft_mb : memory_breakdown) { ggml_backend_buffer_type_t buft = buft_mb.first; const llama_memory_breakdown_data & mb = buft_mb.second; if (seen_buffer_types.count(buft) == 1) { continue; } - const std::string name = ggml_backend_buft_name(buft); - const size_t self = mb.model + mb.context + mb.compute; + common_memory_breakdown_row row; + row.kind = common_memory_breakdown_row::BUFFER_TYPE; + row.name = ggml_backend_buft_name(buft); + row.model = mb.model; + row.context = mb.context; + row.compute = mb.compute; + rows.push_back(std::move(row)); + + seen_buffer_types.insert(buft); + } + + return rows; +} + +void common_memory_breakdown_print(const struct llama_context * ctx) { + const std::vector rows = common_memory_breakdown_rows(ctx); + + std::vector> table_data; + table_data.reserve(rows.size() + 1); + + // same data as the table below, for --log-jsonl consumers + common_json json_rows = common_json::array(); + const std::string template_header = "%s: | %s | %s %s %s %s %s %s %s |\n"; + const std::string template_gpu = "%s: | %s | %s = %s + (%s = %s + %s + %s) + %s |\n"; + const std::string template_other = "%s: | %s | %s %s %s = %s + %s + %s %s |\n"; + + table_data.push_back({template_header, "memory breakdown [MiB]", "total", "free", "self", "model", "context", "compute", "unaccounted"}); + + constexpr size_t MiB = 1024 * 1024; + + for (const auto & row : rows) { + const size_t self = row.model + row.context + row.compute; + + if (row.kind == common_memory_breakdown_row::DEVICE) { + const int64_t unaccounted = row.total - row.free - static_cast(self); + + table_data.push_back({ + template_gpu, + " - " + row.name + " (" + row.description + ")", + std::to_string(row.total / MiB), + std::to_string(row.free / MiB), + std::to_string(self / MiB), + std::to_string(row.model / MiB), + std::to_string(row.context / MiB), + std::to_string(row.compute / MiB), + std::to_string(unaccounted / static_cast(MiB))}); + + json_rows.push_back({ + {"kind", "device"}, + {"name", row.name}, + {"description", row.description}, + {"total", row.total / MiB}, + {"free", row.free / MiB}, + {"self", self / MiB}, + {"model", row.model / MiB}, + {"context", row.context / MiB}, + {"compute", row.compute / MiB}, + {"unaccounted", unaccounted / static_cast(MiB)}, + }); + continue; + } + table_data.push_back({ template_other, - " - " + name, + " - " + row.name, "", // total "", // free std::to_string(self / MiB), - std::to_string(mb.model / MiB), - std::to_string(mb.context / MiB), - std::to_string(mb.compute / MiB), + std::to_string(row.model / MiB), + std::to_string(row.context / MiB), + std::to_string(row.compute / MiB), ""}); // unaccounted - rows.push_back({ - {"kind", "buffer_type"}, - {"name", name}, + json_rows.push_back({ + {"kind", row.kind == common_memory_breakdown_row::HOST ? "host" : "buffer_type"}, + {"name", row.name}, {"self", self / MiB}, - {"model", mb.model / MiB}, - {"context", mb.context / MiB}, - {"compute", mb.compute / MiB}, + {"model", row.model / MiB}, + {"context", row.context / MiB}, + {"compute", row.compute / MiB}, }); - - seen_buffer_types.insert(buft); } for (size_t j = 1; j < table_data[0].size(); j++) { @@ -1081,7 +1175,7 @@ void common_memory_breakdown_print(const struct llama_context * ctx) { LOG_JSON("fit_memory_breakdown", common_json({ {"unit", "MiB"}, - {"rows", rows}, + {"rows", json_rows}, })); } diff --git a/common/fit.h b/common/fit.h index 824d386b0..47478bba1 100644 --- a/common/fit.h +++ b/common/fit.h @@ -3,6 +3,7 @@ #include "ggml.h" #include "llama.h" +#include #include enum common_params_fit_status { @@ -45,6 +46,47 @@ void common_fit_print( void common_memory_breakdown_print(const llama_context * ctx); +// one row of the memory breakdown: a device, the host or another buffer type +struct common_memory_breakdown_row { + enum kind_t { DEVICE, HOST, BUFFER_TYPE }; + + kind_t kind = DEVICE; + std::string name; // device or buffer type name, "Host" for the host row + std::string description; // device description, empty otherwise + int64_t total = 0; // device memory in bytes, 0 unless kind == DEVICE + int64_t free = 0; // free device memory in bytes, 0 unless kind == DEVICE + size_t model = 0; // bytes held by this context, by kind of buffer + size_t context = 0; + size_t compute = 0; +}; + +// memory a model needs at the min and max context --fit may pick +struct common_memory_reports { + uint32_t n_ctx_train = 0; + uint32_t n_ctx_max = 0; // the largest context that fits the free device memory right now + uint32_t n_ctx_min = 0; // --fit-ctx (the smallest context --fit may pick), with the placement set in params + std::vector info_max; // breakdown at n_ctx_max + std::vector info_min; // breakdown at n_ctx_min +}; + +// rows of common_memory_breakdown_print: model devices first, then host +std::vector common_memory_breakdown_rows(const llama_context * ctx); + +// load model + context with no_alloc and return the memory breakdown, params are used as-is +// n_ctx_train and n_ctx are optional outputs; throws on failure +std::vector common_fit_measure( + const char * path_model, + const llama_model_params * mparams, + const llama_context_params * cparams, + uint32_t * n_ctx_train, + uint32_t * n_ctx, + ggml_log_level log_level); + +struct common_params; + +// measure min and max without allocating, params may be modified as with --fit; throws on failure +common_memory_reports common_measure_memory(common_params & params); + struct common_device_memory_data { int64_t total; int64_t free; diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt index e0b9d0ec1..43d9027df 100644 --- a/tools/server/CMakeLists.txt +++ b/tools/server/CMakeLists.txt @@ -47,6 +47,8 @@ add_library(${TARGET} server-http.h server-models.cpp server-models.h + server-models-sched.cpp + server-models-sched.h ) set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON) diff --git a/tools/server/README.md b/tools/server/README.md index e7487c56f..fe4012c4a 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -230,6 +230,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--models-dir PATH` | directory containing models for the router server (default: disabled)
(env: LLAMA_ARG_MODELS_DIR) | | `--models-preset PATH` | path to INI file containing model presets for the router server (default: disabled)
(env: LLAMA_ARG_MODELS_PRESET) | | `--models-max N` | for router server, maximum number of models to load simultaneously (default: 4, 0 = unlimited)
(env: LLAMA_ARG_MODELS_MAX) | +| `--models-memory-margin N` | for router server, MiB of device memory to keep free per device, unload LRU models until a new one fits (default: 0, 0 = disabled)
(env: LLAMA_ARG_MODELS_MEMORY_MARGIN) | | `--models-autoload, --no-models-autoload` | for router server, whether to automatically load models (default: enabled)
(env: LLAMA_ARG_MODELS_AUTOLOAD) | | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:
- none: leaves thoughts unparsed in `message.content`
- deepseek: puts thoughts in `message.reasoning_content`
- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`
(default: auto)
(env: LLAMA_ARG_THINK) | @@ -2117,6 +2118,40 @@ For the **WebSocket** endpoint (`/v1/realtime`) The router also uses the `model` By default, the model will be loaded automatically if it's not loaded. To disable this, add `--no-models-autoload` when starting the server. Additionally, you can include `?autoload=true|false` in the query param to control this behavior per-request. +### Memory-based unloading + +By default, the router only limits the *number* of loaded models (`--models-max`, least recently used first). With `--models-memory-margin N`, it also keeps track of device memory and unloads models (LRU) until the model being loaded fits, while leaving `N` MiB free on each device. + +How it works: +- On startup, and whenever a model is added or its preset changes, the router spawns one short-lived `llama-server` child per model (one at a time, in the background) that reads the model metadata without allocating memory and reports two measurements: + - `min`: the memory the model needs at the smallest context that `--fit` may pick, i.e. `--fit-ctx` (default 4096) times the number of sequences, with the device placement from its preset. If the preset pins the context with `-c`, or sets `--fit off`, this is the memory at that context. + - `max`: the memory it needs at the largest context that `--fit` picks given the device memory free at that moment. This is a snapshot: it changes as other models load and unload. +- When an instance is loaded, it reports the memory it actually allocated (`effective`). With `--fit` on (the default) an instance takes the largest context that fits at load time, so it typically ends up close to `max`. Pin the context with `-c` in the preset of a model if you want it to leave room for other models. +- Before spawning an instance, the router unloads idle models (LRU) until the `min` of the new model fits in the free memory of every device minus the margin. A model that is serving requests is never unloaded. If nothing is left to unload and the model still does not fit, it is loaded anyway and `--fit` in the instance decides. +- A load request for a model whose measurement has not finished yet waits for it (the model is moved to the front of the measurement queue). + +Notes: +- The router itself never touches a GPU, all device memory figures come from the child processes. +- Host (CPU) memory is assumed to be unlimited, as with `--fit`. On a build without GPU devices the memory limit has no effect. +- A model whose measurement fails (for example a broken GGUF, or an `-hf` model that is not downloaded yet) is loaded without memory checks; the measurement is retried once the model has been loaded successfully. +- Multimodal projectors (`mmproj`) are not included in the measurements yet, such a model needs more than its `min` reports. +- The measurements are exposed in the `memory` field of each model in `GET /models` (sizes in bytes), and a `model_memory` event is sent on `/models/sse` when a measurement completes: + +```js +"memory": { + "status": "measured", // unmeasured | measuring | measured | failed + "min": { + "n_ctx": 4096, + "devices": [ + { "kind": "device", "name": "CUDA0", "total": 25757220864, "free": 24387584000, "model": 4986635264, "context": 536870912, "compute": 301989888 }, + { "kind": "host", "name": "Host", "model": 263454720, "context": 0, "compute": 16793600 } + ] + }, + "max": { "n_ctx": 131072, "devices": [ /* ... */ ] }, + "effective": { "n_ctx": 131072, "devices": [ /* ... */ ] } // null unless the model is loaded +} +``` + ### GET `/models`: List available models Listing all models in cache. The model metadata will also include a field to indicate the status of the model: @@ -2367,6 +2402,18 @@ Example events: "event": "model_remove" } +// only with --models-memory-margin, when the measurement of a model completes (see "Memory-based unloading") +{ + "model": "...", + "event": "model_memory", + "data": { + "status": "measured", + "min": { /* ... */ }, + "max": { /* ... */ }, + "effective": null + } +} + // special event: reload of the list of all models { "model": "*", diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 0925c508e..01bf12e61 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -737,6 +737,7 @@ enum server_role_type { SERVER_ROLE_ROUTER, // no model: lists the models, spawns instances on demand and proxies requests to them SERVER_ROLE_INSTANCE, // serves one model, standalone or as a child of a router SERVER_ROLE_DOWNLOAD, // child of a router: download a model, report the progress, exit + SERVER_ROLE_MEASURE, // child of a router: report the memory the model needs without loading it, exit }; static inline const char * server_role_type_to_str(server_role_type role) { @@ -744,6 +745,7 @@ static inline const char * server_role_type_to_str(server_role_type role) { case SERVER_ROLE_ROUTER: return "router"; case SERVER_ROLE_INSTANCE: return "instance"; case SERVER_ROLE_DOWNLOAD: return "download"; + case SERVER_ROLE_MEASURE: return "measure"; default: GGML_ASSERT(false && "invalid server_role_type"); } } @@ -753,6 +755,9 @@ static inline server_role_type server_role_type_from_str(const std::string & str if (str == "download") { return SERVER_ROLE_DOWNLOAD; } + if (str == "measure") { + return SERVER_ROLE_MEASURE; + } return SERVER_ROLE_INSTANCE; } @@ -772,6 +777,7 @@ static inline server_role_type server_role_type_from_str(const std::string & str // the states a child reports to the router enum server_state { SERVER_STATE_DOWNLOADING, + SERVER_STATE_MEASURING, // sent by a measure child, see server_run_measure SERVER_STATE_LOADING, SERVER_STATE_READY, SERVER_STATE_SLEEPING, @@ -780,6 +786,7 @@ enum server_state { static inline std::string server_state_to_str(server_state state) { switch (state) { case SERVER_STATE_DOWNLOADING: return "downloading"; + case SERVER_STATE_MEASURING: return "measuring"; case SERVER_STATE_LOADING: return "loading"; case SERVER_STATE_READY: return "ready"; case SERVER_STATE_SLEEPING: return "sleeping"; @@ -789,6 +796,7 @@ static inline std::string server_state_to_str(server_state state) { static inline server_state server_state_from_str(const std::string & str) { if (str == "downloading") return SERVER_STATE_DOWNLOADING; + if (str == "measuring") return SERVER_STATE_MEASURING; if (str == "loading") return SERVER_STATE_LOADING; if (str == "ready") return SERVER_STATE_READY; if (str == "sleeping") return SERVER_STATE_SLEEPING; diff --git a/tools/server/server-models-sched.cpp b/tools/server/server-models-sched.cpp new file mode 100644 index 000000000..08c1c4c57 --- /dev/null +++ b/tools/server/server-models-sched.cpp @@ -0,0 +1,593 @@ +#include "server-models-sched.h" +#include "server-models.h" + +#include "fit.h" + +#include +#include + +static constexpr int MEASURE_TIMEOUT = 300; // seconds, then the measure child is killed + +// +// server_model_memory +// + +server_model_memory::server_model_memory(uint32_t n_ctx, const std::vector & rows) : n_ctx(n_ctx) { + for (const auto & row : rows) { + device_t d; + switch (row.kind) { + case common_memory_breakdown_row::DEVICE: d.kind = "device"; break; + case common_memory_breakdown_row::HOST: d.kind = "host"; break; + case common_memory_breakdown_row::BUFFER_TYPE: d.kind = "buffer_type"; break; + } + d.name = row.name; + d.total = row.total; + d.free = row.free; + d.model = row.model; + d.context = row.context; + d.compute = row.compute; + devices.push_back(std::move(d)); + } +} + +json server_model_memory::to_json() const { + json devs = json::array(); + for (const auto & d : devices) { + json jd = { + {"kind", d.kind}, + {"name", d.name}, + {"model", d.model}, + {"context", d.context}, + {"compute", d.compute}, + }; + if (d.is_device()) { + jd["total"] = d.total; + jd["free"] = d.free; + } + devs.push_back(std::move(jd)); + } + return { + {"n_ctx", n_ctx}, + {"devices", devs}, + }; +} + +server_model_memory server_model_memory::from_json(const json & j) { + server_model_memory mem; + mem.n_ctx = j.at("n_ctx").get(); + for (const auto & jd : j.at("devices")) { + device_t d; + d.kind = jd.at("kind").get(); + d.name = jd.at("name").get(); + d.total = json_value(jd, "total", (size_t) 0); + d.free = json_value(jd, "free", (size_t) 0); + d.model = jd.at("model").get(); + d.context = jd.at("context").get(); + d.compute = jd.at("compute").get(); + mem.devices.push_back(std::move(d)); + } + return mem; +} + +json server_model_memory_info::to_json() const { + return { + {"status", server_model_memory_status_to_string(status)}, + {"min", min ? min->to_json() : json(nullptr)}, + {"max", max ? max->to_json() : json(nullptr)}, + {"effective", effective ? effective->to_json() : json(nullptr)}, + }; +} + +// +// request queue +// + +bool server_models_sched::has_capacity(std::unique_lock & lk) { + check_lock(lk); + return models.base_params.models_max <= 0 + || count_running() < (size_t) models.base_params.models_max; +} + +std::string server_models_sched::pick_victim(std::unique_lock & lk) { + check_lock(lk); + std::string victim; + int64_t victim_last_used = 0; + for (const auto & m : models.mapping) { + // a busy model is mid-request, one still coming up has no request to finish + if (m.second.req_count != 0 || !m.second.meta.is_ready_or_sleep()) { + continue; + } + // already on its way out, or a queued request wants it + if (models.stopping_models.count(m.first) || find(m.first)) { + continue; + } + if (victim.empty() || m.second.meta.last_used < victim_last_used) { + victim = m.first; + victim_last_used = m.second.meta.last_used; + } + } + return victim; +} + +void server_models_sched::join(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + if (entry_t * e = find(model_id)) { + e->n_waiters++; + SRV_INF("request for name=%s joined the queue, %d waiting\n", model_id.c_str(), e->n_waiters); + return; + } + queue.push_back({ model_id, 1, false }); + SRV_INF("request for name=%s queued at position %zu\n", + model_id.c_str(), queue.size()); +} + +void server_models_sched::leave(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + for (auto it = queue.begin(); it != queue.end(); ++it) { + if (it->model_id == model_id) { + if (--it->n_waiters <= 0) { + queue.erase(it); // last one waiting for this model went away + } + return; + } + } +} + +bool server_models_sched::queue_empty(std::unique_lock & lk) { + check_lock(lk); + return queue.empty(); +} + +bool server_models_sched::try_claim(std::unique_lock & lk, const std::string & model_id) { + check_lock(lk); + if (queue.empty() || queue.front().model_id != model_id || queue.front().loading) { + return false; + } + if (!has_capacity(lk)) { + return false; + } + // over the memory margin with every running model busy: wait for one to go idle + if (memory_enabled() && limits_exceeded(lk, model_id, /* verbose */ false) && pick_victim(lk).empty()) { + return false; + } + queue.front().loading = true; + return true; +} + +void server_models_sched::claim_done(std::unique_lock & lk, const std::string & model_id, bool ok) { + check_lock(lk); + if (ok) { + return; + } + for (auto it = queue.begin(); it != queue.end(); ++it) { + if (it->model_id == model_id) { + it->loading = false; + return; + } + } +} + +void server_models_sched::tick(std::unique_lock & lk) { + check_lock(lk); + if (models.base_params.models_max <= 0 || queue.empty()) { + return; + } + int n_running = 0; + int n_stopping = 0; + for (const auto & m : models.mapping) { + if (m.second.meta.is_running()) { + n_running++; + if (models.stopping_models.count(m.first)) { + n_stopping++; + } + } + } + int n_needed = 0; + int n_claimed = 0; // claimed the slot, but load() has not spawned yet + for (const auto & e : queue) { + if (!e.loading) { + n_needed++; + continue; + } + auto it = models.mapping.find(e.model_id); + if (it != models.mapping.end() && !it->second.meta.is_running()) { + n_claimed++; + } + } + int n_free = models.base_params.models_max - n_running + n_stopping - n_claimed; + while (n_free < n_needed) { + std::string victim = pick_victim(lk); + if (victim.empty()) { + return; // all remaining models are busy, wait for a request to end + } + SRV_INF("evicting idle LRU name=%s for a queued request\n", victim.c_str()); + models.request_stop(victim); + n_free++; + } +} + +server_models_sched::entry_t * server_models_sched::find(const std::string & model_id) { + for (auto & e : queue) { + if (e.model_id == model_id) { + return &e; + } + } + return nullptr; +} + +void server_models_sched::check_lock(std::unique_lock & lk) { + GGML_ASSERT(lk.owns_lock() && lk.mutex() == &models.mutex); +} + +size_t server_models_sched::count_running() { + size_t count = 0; + for (const auto & m : models.mapping) { + if (m.second.meta.is_running()) { + count++; + } + } + return count; +} + +// +// loading a model +// + +bool server_models_sched::limits_exceeded(std::unique_lock & lk, const std::string & name, bool verbose) { + if (!has_capacity(lk)) { + return true; + } + if (!memory_enabled()) { + return false; + } + const auto missing = deficit(name); + if (missing.empty()) { + return false; + } + bool any_running = false; + for (const auto & m : models.mapping) { + if (m.second.meta.is_running()) { + any_running = true; + break; + } + } + for (const auto & [dev, bytes] : missing) { + if (!verbose) { + break; + } + if (any_running) { + SRV_DBG("model name=%s is short by %zu MiB on %s\n", name.c_str(), bytes / (1024 * 1024), dev.c_str()); + } else { + // nothing to unload: min is pessimistic, let --fit in the instance try + SRV_WRN("model name=%s is projected to be short by %zu MiB on %s with nothing else loaded, loading anyway\n", + name.c_str(), bytes / (1024 * 1024), dev.c_str()); + } + } + return any_running; +} + +void server_models_sched::make_room(const std::string & name) { + if (memory_enabled()) { + // the measure result decides what to unload + std::unique_lock lk(models.mutex); + enqueue_measure(name, /* front */ true); + auto measured = [this, &name]() { + auto it = models.mapping.find(name); + return it == models.mapping.end() || it->second.meta.memory.is_settled(); + }; + if (!measured()) { + SRV_INF("waiting for the memory measurement of model name=%s\n", name.c_str()); + models.cv.wait(lk, measured); + } + } + + while (true) { + std::string lru_model_name; + { + std::unique_lock lk(models.mutex); + if (!limits_exceeded(lk, name)) { + return; + } + lru_model_name = pick_victim(lk); + } + if (lru_model_name.empty()) { + return; // every running model is busy, load() reports the limit + } + SRV_INF("limits reached, removing LRU name=%s\n", lru_model_name.c_str()); + models.unload(lru_model_name); + // wait for unload to complete + { + std::unique_lock lk(models.mutex); + models.cv.wait(lk, [this, &lru_model_name]() { + return models.mapping[lru_model_name].meta.status == SERVER_MODEL_STATUS_UNLOADED; + }); + } + } +} + +// +// memory accounting +// + +bool server_models_sched::memory_enabled() const { + return models.base_params.models_memory_margin > 0; +} + +void server_models_sched::measure_unmeasured(std::unique_lock & lk) { + check_lock(lk); + for (const auto & [name, inst] : models.mapping) { + enqueue_measure(name); + } +} + +void server_models_sched::on_spawned(std::unique_lock & lk, const std::string & name) { + check_lock(lk); + reserve(name); +} + +void server_models_sched::on_ready(std::unique_lock & lk, const std::string & name, const std::optional & memory) { + check_lock(lk); + if (!memory_enabled()) { + return; + } + auto it = models.mapping.find(name); + if (it == models.mapping.end()) { + return; + } + if (!memory.has_value()) { + reserve(name); // woke up from sleep + return; + } + // use the real usage, not the estimate + auto & info = it->second.meta.memory; + info.effective = memory; + release(name); + reserve(name); + calibrate(*memory, name); + if (info.status == SERVER_MODEL_MEMORY_FAILED && info.n_attempts < 2) { + // the model was maybe not downloaded at measure time, retry + info.status = SERVER_MODEL_MEMORY_UNMEASURED; + enqueue_measure(name); + } +} + +void server_models_sched::on_sleeping(std::unique_lock & lk, const std::string & name) { + check_lock(lk); + release(name); +} + +void server_models_sched::on_unloaded(std::unique_lock & lk, const std::string & name) { + check_lock(lk); + release(name); + auto it = models.mapping.find(name); + if (it != models.mapping.end()) { + it->second.meta.memory.effective.reset(); + } +} + +void server_models_sched::cancel_measurements(std::unique_lock & lk) { + check_lock(lk); + measure_queue.clear(); + if (measure_proc) { + SRV_INF("stopping measure process for model name=%s\n", measuring.c_str()); + measure_proc->terminate(); + // the monitor reaps it and clears measuring + } +} + +bool server_models_sched::is_measuring(std::unique_lock & lk) { + check_lock(lk); + return !measuring.empty(); +} + +void server_models_sched::enqueue_measure(const std::string & name, bool front) { + if (!memory_enabled()) { + return; + } + auto it = models.mapping.find(name); + if (it == models.mapping.end() || it->second.meta.memory.status != SERVER_MODEL_MEMORY_UNMEASURED) { + return; + } + auto pos = std::find(measure_queue.begin(), measure_queue.end(), name); + if (pos != measure_queue.end()) { + if (!front) { + return; // already queued + } + measure_queue.erase(pos); + } + if (front) { + measure_queue.push_front(name); + } else { + measure_queue.push_back(name); + } + measure_next(); +} + +void server_models_sched::measure_next() { + if (!measuring.empty()) { + return; // one measure child at a time + } + while (!measure_queue.empty()) { + std::string name = std::move(measure_queue.front()); + measure_queue.pop_front(); + auto it = models.mapping.find(name); + if (it == models.mapping.end() || it->second.meta.memory.status != SERVER_MODEL_MEMORY_UNMEASURED) { + continue; // removed or measured meanwhile + } + auto & meta = it->second.meta; + + // same args as the instance, --offline: never download + std::vector child_args = meta.args; + child_args.push_back("--offline"); + std::vector child_env = models.base_env; + child_env.push_back(std::string(SERVER_ENV_ROUTER_PORT) + "=" + std::to_string(models.base_params.port)); + child_env.push_back(std::string(SERVER_ENV_CHILD_ROLE) + "=" + server_role_type_to_str(SERVER_ROLE_MEASURE)); + + auto proc = std::make_shared(); + if (!proc->sproc.create(child_args, subprocess_option_no_window, child_env)) { + SRV_ERR("failed to spawn measure process for model name=%s\n", name.c_str()); + meta.memory.status = SERVER_MODEL_MEMORY_FAILED; + continue; + } + SRV_INF("measuring memory for model name=%s\n", name.c_str()); + meta.memory.status = SERVER_MODEL_MEMORY_MEASURING; + meta.memory.n_attempts++; + measuring = name; + measure_proc = proc; + models.watch_child(name, proc, SERVER_ROLE_MEASURE, 0, MEASURE_TIMEOUT); + return; + } +} + +void server_models_sched::on_measure_exit(const std::string & name, const std::shared_ptr & proc, int exit_code) { + std::lock_guard lk(models.mutex); + if (measure_proc != proc) { + return; // not the child we are waiting for + } + measure_proc.reset(); + measuring.clear(); + auto it = models.mapping.find(name); + if (it != models.mapping.end() && it->second.meta.memory.status == SERVER_MODEL_MEMORY_MEASURING) { + // exited without a result: crashed, or killed after the timeout + SRV_WRN("measure process for model name=%s exited with status %d before reporting, memory limits will not apply to it\n", + name.c_str(), exit_code); + it->second.meta.memory.status = SERVER_MODEL_MEMORY_FAILED; + models.notify_sse("model_memory", name, it->second.meta.memory.to_json()); + } + models.cv.notify_all(); + measure_next(); +} + +void server_models_sched::handle_measured(const std::string & name, const json & payload) { + std::lock_guard lk(models.mutex); + auto it = models.mapping.find(name); + if (it == models.mapping.end()) { + return; + } + auto & info = it->second.meta.memory; + const std::string result = json_value(payload, "result", std::string()); + + bool ok = result == "measured"; + if (ok) { + try { + info.min = server_model_memory::from_json(payload.at("min")); + info.max = server_model_memory::from_json(payload.at("max")); + } catch (const std::exception & e) { + SRV_ERR("malformed memory measurement for model name=%s: %s\n", name.c_str(), e.what()); + info.min.reset(); + info.max.reset(); + ok = false; + } + } else { + SRV_WRN("measuring memory for model name=%s failed: %s\n", + name.c_str(), json_value(payload, "error", std::string("unknown error")).c_str()); + } + + if (ok) { + info.status = SERVER_MODEL_MEMORY_MEASURED; + calibrate(*info.max, ""); // latest view of the devices + const auto & mn = *info.min; + const auto & mx = *info.max; + for (size_t i = 0; i < mn.devices.size() && i < mx.devices.size(); i++) { + if (!mn.devices[i].is_device() || mn.devices[i].name != mx.devices[i].name) { + continue; + } + SRV_INF("model name=%s needs on %s: %zu MiB at n_ctx=%u (min), %zu MiB at n_ctx=%u (max), %zu MiB free\n", + name.c_str(), mn.devices[i].name.c_str(), + mn.devices[i].self() / (1024 * 1024), mn.n_ctx, + mx.devices[i].self() / (1024 * 1024), mx.n_ctx, + mx.devices[i].free / (1024 * 1024)); + } + } else { + info.status = SERVER_MODEL_MEMORY_FAILED; + } + models.notify_sse("model_memory", name, info.to_json()); + models.cv.notify_all(); +} + +void server_models_sched::calibrate(const server_model_memory & mem, const std::string & reporter) { + for (const auto & d : mem.devices) { + if (!d.is_device() || d.total == 0) { + continue; + } + // other instances are tracked, so they are not external (CUDA counts them in free, Metal does not) + int64_t tracked = 0; + for (const auto & [name, devs] : held) { + if (name == reporter) { + continue; + } + auto h = devs.find(d.name); + if (h != devs.end()) { + tracked += (int64_t) h->second; + } + } + // a measure child allocates nothing, its free does not include self + const int64_t self = reporter.empty() ? 0 : (int64_t) d.self(); + auto & dm = device_memory[d.name]; + dm.total = (int64_t) d.total; + dm.external = std::max(0, dm.total - (int64_t) d.free - self - tracked); + SRV_DBG("device %s: total=%" PRId64 " MiB, external=%" PRId64 " MiB, available=%" PRId64 " MiB\n", + d.name.c_str(), dm.total / (1024 * 1024), dm.external / (1024 * 1024), available(d.name) / (1024 * 1024)); + } +} + +void server_models_sched::reserve(const std::string & name) { + auto it = models.mapping.find(name); + if (it == models.mapping.end() || held.count(name)) { + return; + } + const auto & info = it->second.meta.memory; + // prefer real usage over the estimate + const auto & mem = info.effective ? info.effective : info.min; + if (!mem) { + return; + } + auto & devs = held[name]; + for (const auto & d : mem->devices) { + if (d.is_device() && d.total > 0) { + devs[d.name] = d.self(); + } + } +} + +void server_models_sched::release(const std::string & name) { + held.erase(name); +} + +int64_t server_models_sched::available(const std::string & dev) const { + auto dm = device_memory.find(dev); + if (dm == device_memory.end()) { + return INT64_MAX; // never reported, nothing to check against + } + int64_t n_held = 0; + for (const auto & [name, devs] : held) { + auto h = devs.find(dev); + if (h != devs.end()) { + n_held += (int64_t) h->second; + } + } + return dm->second.total - dm->second.external - n_held; +} + +std::map server_models_sched::deficit(const std::string & name) const { + std::map res; + auto it = models.mapping.find(name); + if (it == models.mapping.end() || !it->second.meta.memory.min) { + return res; + } + const int64_t margin = (int64_t) models.base_params.models_memory_margin * 1024 * 1024; + for (const auto & d : it->second.meta.memory.min->devices) { + // host memory is assumed unlimited, same as --fit + if (!d.is_device() || d.total == 0) { + continue; + } + if (device_memory.find(d.name) == device_memory.end()) { + continue; + } + const int64_t avail = available(d.name) - margin; + if ((int64_t) d.self() > avail) { + res[d.name] = (size_t) ((int64_t) d.self() - avail); + } + } + return res; +} diff --git a/tools/server/server-models-sched.h b/tools/server/server-models-sched.h new file mode 100644 index 000000000..dda80fca0 --- /dev/null +++ b/tools/server/server-models-sched.h @@ -0,0 +1,181 @@ +#pragma once + +#include "server-common.h" + +#include +#include +#include +#include +#include +#include +#include + +// when to load and unload model instances: request queue (--models-max) and memory accounting (--models-memory-margin) +// see "Memory-based unloading" in the README + +struct server_models; + +// memory of a model instance per device, in bytes; from a measure child or from the running instance +struct server_model_memory { + struct device_t { + std::string kind; // "device", "host" or "buffer_type" + std::string name; // ggml device name, "Host" for host memory + size_t total = 0; // device memory, 0 unless kind == "device" + size_t free = 0; // free device memory when this was measured, 0 unless kind == "device" + size_t model = 0; // held by the instance, by kind of buffer + size_t context = 0; + size_t compute = 0; + + bool is_device() const { return kind == "device"; } + size_t self() const { return model + context + compute; } + }; + + uint32_t n_ctx = 0; // context size this was measured at + std::vector devices; + + server_model_memory() = default; + server_model_memory(uint32_t n_ctx, const std::vector & rows); + + json to_json() const; + static server_model_memory from_json(const json & j); // throws on a malformed input +}; + +// memory measurement state of a model +enum server_model_memory_status { + SERVER_MODEL_MEMORY_UNMEASURED, // queued for a measure child + SERVER_MODEL_MEMORY_MEASURING, // its measure child is running + SERVER_MODEL_MEMORY_MEASURED, + SERVER_MODEL_MEMORY_FAILED, // measure failed, no memory limit for this model +}; + +static std::string server_model_memory_status_to_string(server_model_memory_status status) { + switch (status) { + case SERVER_MODEL_MEMORY_UNMEASURED: return "unmeasured"; + case SERVER_MODEL_MEMORY_MEASURING: return "measuring"; + case SERVER_MODEL_MEMORY_MEASURED: return "measured"; + case SERVER_MODEL_MEMORY_FAILED: return "failed"; + default: return "unknown"; + } +} + +struct server_model_memory_info { + server_model_memory_status status = SERVER_MODEL_MEMORY_UNMEASURED; + int n_attempts = 0; // to limit the retries + std::optional min = std::nullopt; // at the smallest context --fit may pick + std::optional max = std::nullopt; // at the largest context --fit picks, a snapshot + std::optional effective = std::nullopt; // what the running instance allocated, reset when it exits + + bool is_settled() const { + return status == SERVER_MODEL_MEMORY_MEASURED || status == SERVER_MODEL_MEMORY_FAILED; + } + + json to_json() const; +}; + +// functions taking lk need models.mutex held by lk +struct server_models_sched { + server_models_sched(server_models & models) : models(models) {} + + // + // request queue (--models-max) + // + + bool has_capacity(std::unique_lock & lk); + + // returns "" if no model can be given up + std::string pick_victim(std::unique_lock & lk); + + // requests wanting the same model share one entry, so they all need only one slot + // and all get unblocked by the single load that entry performs + void join(std::unique_lock & lk, const std::string & model_id); + void leave(std::unique_lock & lk, const std::string & model_id); + bool queue_empty(std::unique_lock & lk); + + // true if it is this model's turn to load, and nobody is loading it yet + bool try_claim(std::unique_lock & lk, const std::string & model_id); + + // on failure the entry is back in line; on success it stays until its waiters leave, + // so the model coming up is never picked as a victim before they use it + void claim_done(std::unique_lock & lk, const std::string & model_id, bool ok); + + // evict idle models while queued requests outnumber the slots that are free or being freed + // never blocks, so it is safe from any thread + void tick(std::unique_lock & lk); + + // + // loading a model + // + + // true if loading name now exceeds models_max or the memory margin + bool limits_exceeded(std::unique_lock & lk, const std::string & name, bool verbose = true); + + // wait for the measure of name, then unload idle LRU models until it fits; call without models.mutex + void make_room(const std::string & name); + + // + // memory accounting (--models-memory-margin), the router never touches a GPU + // track what each instance holds, child reports only calibrate the rest: works for per-process (Metal) and system-wide (CUDA) free memory + // + + bool memory_enabled() const; + + // queue a measure child for every model not measured yet + void measure_unmeasured(std::unique_lock & lk); + + // count a new instance at its min until it reports + void on_spawned(std::unique_lock & lk, const std::string & name); + + // memory: what the instance allocated, nullopt on wake up from sleep + void on_ready(std::unique_lock & lk, const std::string & name, const std::optional & memory); + // a sleeping instance frees its model and context + void on_sleeping(std::unique_lock & lk, const std::string & name); + void on_unloaded(std::unique_lock & lk, const std::string & name); + + // stop the measure child and drop the queue, for shutdown + void cancel_measurements(std::unique_lock & lk); + bool is_measuring(std::unique_lock & lk); + + // called from the monitor thread, take models.mutex themselves + void handle_measured(const std::string & name, const json & payload); + void on_measure_exit(const std::string & name, const std::shared_ptr & proc, int exit_code); + + private: + struct entry_t { + std::string model_id; + int n_waiters; // requests waiting for this model + bool loading; // one of the waiters is doing the load right now + }; + + struct device_memory_t { + int64_t total = 0; // bytes + int64_t external = 0; // bytes not tracked by the router: other processes, backend contexts + }; + + server_models & models; + + // requests waiting for a slot + std::deque queue; + + // one measure child runs at a time + std::deque measure_queue; + std::string measuring; // model whose measure child is running, empty when idle + std::shared_ptr measure_proc; + + std::map device_memory; // by device name, empty until the first child reports + std::map> held; // model -> device -> bytes counted for its instance + + entry_t * find(const std::string & model_id); + void check_lock(std::unique_lock & lk); + size_t count_running(); + + // all below need models.mutex + void enqueue_measure(const std::string & name, bool front = false); + void measure_next(); + // update device_memory from a child report; reporter: the instance that sent it, empty for a measure child + void calibrate(const server_model_memory & mem, const std::string & reporter); + void reserve(const std::string & name); // count the instance's memory (effective, else min) + void release(const std::string & name); + int64_t available(const std::string & dev) const; // total - external - what the instances hold + // bytes missing per device to load name at its min, empty if it fits + std::map deficit(const std::string & name) const; +}; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 575140132..9f583b738 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -7,6 +7,7 @@ #include "build-info.h" #include "preset.h" #include "download.h" +#include "fit.h" #include "hf-cache.h" #include "http.h" #include "subproc.h" @@ -24,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -58,12 +60,16 @@ struct server_monitor { } // thread-safe - void watch(const std::string & name, std::shared_ptr proc, server_role_type role, int port) { + // kill_after > 0: force-kill the child after that many seconds + void watch(const std::string & name, std::shared_ptr proc, server_role_type role, int port, int kill_after = 0) { child_t c; c.name = name; c.proc = std::move(proc); c.role = role; c.port = port; + if (kill_after > 0) { + c.deadline = ggml_time_ms() + (int64_t) kill_after * 1000; + } if (!c.proc->has_output()) { SRV_ERR("failed to get stdout/stderr of child process for name=%s\n", name.c_str()); c.eof = true; @@ -118,7 +124,8 @@ private: case cmd_t::STOP: // the newest child with this name is the one the registry knows for (auto it = children.rbegin(); it != children.rend(); ++it) { - if (it->name != cmd.name) { + // a measure child has the same name, skip it + if (it->name != cmd.name || it->role == SERVER_ROLE_MEASURE) { continue; } if (cmd.send_exit && !it->eof) { @@ -227,7 +234,7 @@ private: now = ggml_time_ms(); for (auto it = children.begin(); it != children.end();) { if (it->deadline && now >= it->deadline && !it->proc->stopped.load(std::memory_order_acquire)) { - SRV_WRN("force-killing model instance name=%s after timeout\n", it->name.c_str()); + SRV_WRN("force-killing child process for name=%s after timeout\n", it->name.c_str()); it->proc->terminate(); it->deadline = 0; } @@ -254,171 +261,6 @@ private: std::thread th; }; -struct server_lru_sched { - server_lru_sched(server_models & models) : models(models) {} - - bool has_capacity(std::unique_lock & lk) { - check_lock(lk); - return models.base_params.models_max <= 0 - || count_running() < (size_t) models.base_params.models_max; - } - - // returns "" if no model can be given up - std::string pick_victim(std::unique_lock & lk) { - check_lock(lk); - std::string victim; - int64_t victim_last_used = 0; - for (const auto & m : models.mapping) { - // a busy model is mid-request, one still coming up has no request to finish - if (m.second.req_count != 0 || !m.second.meta.is_ready_or_sleep()) { - continue; - } - // already on its way out, or a queued request wants it - if (models.stopping_models.count(m.first) || find(m.first)) { - continue; - } - if (victim.empty() || m.second.meta.last_used < victim_last_used) { - victim = m.first; - victim_last_used = m.second.meta.last_used; - } - } - return victim; - } - - // requests wanting the same model share one entry, so they all need only one slot - // and all get unblocked by the single load that entry performs - void join(std::unique_lock & lk, const std::string & model_id) { - check_lock(lk); - if (entry_t * e = find(model_id)) { - e->n_waiters++; - SRV_INF("request for name=%s joined the queue, %d waiting\n", model_id.c_str(), e->n_waiters); - return; - } - queue.push_back({ model_id, 1, false }); - SRV_INF("request for name=%s queued at position %zu\n", - model_id.c_str(), queue.size()); - } - - void leave(std::unique_lock & lk, const std::string & model_id) { - check_lock(lk); - for (auto it = queue.begin(); it != queue.end(); ++it) { - if (it->model_id == model_id) { - if (--it->n_waiters <= 0) { - queue.erase(it); // last one waiting for this model went away - } - return; - } - } - } - - bool queue_empty(std::unique_lock & lk) { - check_lock(lk); - return queue.empty(); - } - - // true if it is this model's turn to load, and nobody is loading it yet - bool try_claim(std::unique_lock & lk, const std::string & model_id) { - check_lock(lk); - if (queue.empty() || queue.front().model_id != model_id || queue.front().loading) { - return false; - } - if (!has_capacity(lk)) { - return false; - } - queue.front().loading = true; - return true; - } - - // on failure the entry is back in line; on success it stays until its waiters leave, - // so the model coming up is never picked as a victim before they use it - void claim_done(std::unique_lock & lk, const std::string & model_id, bool ok) { - check_lock(lk); - if (ok) { - return; - } - for (auto it = queue.begin(); it != queue.end(); ++it) { - if (it->model_id == model_id) { - it->loading = false; - return; - } - } - } - - // evict idle models while queued requests outnumber the slots that are free or being freed - // caller must hold models.mutex; never blocks, so it is safe from any thread - void tick(std::unique_lock & lk) { - check_lock(lk); - if (models.base_params.models_max <= 0 || queue.empty()) { - return; - } - int n_running = 0; - int n_stopping = 0; - for (const auto & m : models.mapping) { - if (m.second.meta.is_running()) { - n_running++; - if (models.stopping_models.count(m.first)) { - n_stopping++; - } - } - } - int n_needed = 0; - int n_claimed = 0; // claimed the slot, but load() has not spawned yet - for (const auto & e : queue) { - if (!e.loading) { - n_needed++; - continue; - } - auto it = models.mapping.find(e.model_id); - if (it != models.mapping.end() && !it->second.meta.is_running()) { - n_claimed++; - } - } - int n_free = models.base_params.models_max - n_running + n_stopping - n_claimed; - while (n_free < n_needed) { - std::string victim = pick_victim(lk); - if (victim.empty()) { - return; // all remaining models are busy, wait for a request to end - } - SRV_INF("evicting idle LRU name=%s for a queued request\n", victim.c_str()); - models.request_stop(victim); - n_free++; - } - } - - private: - struct entry_t { - std::string model_id; - int n_waiters; // requests waiting for this model - bool loading; // one of the waiters is doing the load right now - }; - - entry_t * find(const std::string & model_id) { - for (auto & e : queue) { - if (e.model_id == model_id) { - return &e; - } - } - return nullptr; - } - - void check_lock(std::unique_lock & lk) { - GGML_ASSERT(lk.owns_lock() && lk.mutex() == &models.mutex); - } - - size_t count_running() { - size_t count = 0; - for (const auto & m : models.mapping) { - if (m.second.meta.is_running()) { - count++; - } - } - return count; - } - - server_models & models; - std::deque queue; -}; - // short loopback budget for the resumable stream router to child JSON calls (probe, lookup, // delete). distinct from params.timeout_read/write which only applies to the generation proxy static constexpr int STREAM_LOOKUP_TIMEOUT_MS = 250; @@ -430,6 +272,7 @@ static void unset_reserved_args(common_preset & preset, bool unset_model_args) { preset.unset_option("LLAMA_ARG_API_KEY_FILE"); preset.unset_option("LLAMA_ARG_MODELS_DIR"); preset.unset_option("LLAMA_ARG_MODELS_MAX"); + preset.unset_option("LLAMA_ARG_MODELS_MEMORY_MARGIN"); preset.unset_option("LLAMA_ARG_MODELS_PRESET"); preset.unset_option("LLAMA_ARG_MODELS_AUTOLOAD"); preset.unset_option("LLAMA_ARG_CONNECT"); @@ -555,7 +398,7 @@ server_models::server_models( base_params(params), base_env(get_environment()), base_preset(ctx_preset.load_from_args(argc, argv)), - sched(std::make_unique(*this)), + sched(std::make_unique(*this)), monitor(std::make_unique(*this)) { // propagate base params to child unset_reserved_args(base_preset, true); @@ -568,9 +411,10 @@ server_models::server_models( bin_path = fs_path_to_utf8(get_server_exec_path()); } catch (const std::exception & e) { bin_path = argv[0]; - LOG_WRN("failed to get server executable path: %s\n", e.what()); - LOG_WRN("using original argv[0] as fallback: %s\n", argv[0]); + SRV_WRN("failed to get server executable path: %s\n", e.what()); + SRV_WRN("using original argv[0] as fallback: %s\n", argv[0]); } + unfinished_downloads.load(); load_models(); debug_fake_timing = !common_get_env("LLAMA_SERVER_DEBUG_FAKE_TIMING").empty(); @@ -1056,6 +900,7 @@ void server_models::load_models() { apply_stop_timeout(); apply_hidden(); log_available_models(); + sched->measure_unmeasured(lk); // skipped on reload, see startup_models if (startup_models.has_value()) { @@ -1131,6 +976,10 @@ void server_models::load_models() { auto it = final_presets.find(name); if (it == final_presets.end()) continue; // erased above + if (preset_options_for_compare(inst.meta.preset) != preset_options_for_compare(it->second)) { + // different args, different memory: measure again + inst.meta.memory = {}; + } inst.meta.preset = it->second; // re-parse aliases, then validate against other models @@ -1198,6 +1047,7 @@ void server_models::load_models() { apply_stop_timeout(); apply_hidden(); + sched->measure_unmeasured(lk); // clear reload flag under the lock, this releases the load() calls waiting on !is_reloading is_reloading = false; @@ -1285,32 +1135,6 @@ std::vector server_models::get_all_meta() { return result; } -void server_models::unload_lru() { - if (base_params.models_max <= 0) { - return; // no limit - } - // remove one of the servers if we passed the models_max (least recently used - LRU) - std::string lru_model_name; - { - std::unique_lock lk(mutex); - if (sched->has_capacity(lk)) { - return; - } - lru_model_name = sched->pick_victim(lk); - } - if (!lru_model_name.empty()) { - SRV_INF("models_max limit reached, removing LRU name=%s\n", lru_model_name.c_str()); - unload(lru_model_name); - // wait for unload to complete - { - std::unique_lock lk(mutex); - cv.wait(lk, [this, &lru_model_name]() { - return mapping[lru_model_name].meta.status == SERVER_MODEL_STATUS_UNLOADED; - }); - } - } -} - void server_models::load(const std::string & name) { load(name, load_options{}); } @@ -1326,7 +1150,9 @@ void server_models::load(const std::string & name, const load_options & opts) { if (!has_model(name)) { throw std::runtime_error("model name=" + name + " is not found"); } - unload_lru(); + if (opts.role == SERVER_ROLE_INSTANCE) { + sched->make_room(name); + } } std::unique_lock lk(mutex); @@ -1378,20 +1204,12 @@ void server_models::load(const std::string & name, const load_options & opts) { } // Re-check capacity under the lock to prevent concurrent loads from - // exceeding models_max. Without this, the window between unload_lru() + // exceeding models_max. Without this, the window between make_room() // releasing its lock and this lock_guard acquiring allows multiple // threads to each observe capacity and all proceed to load. // Download workers do not use models_max slots. - if (opts.role == SERVER_ROLE_INSTANCE && base_params.models_max > 0) { - size_t count_active = 0; - for (const auto & m : mapping) { - if (m.second.meta.is_running()) { - count_active++; - } - } - if (count_active >= (size_t)base_params.models_max) { - throw std::runtime_error("model limit reached, try again later"); - } + if (opts.role == SERVER_ROLE_INSTANCE && sched->limits_exceeded(lk, name, /* verbose */ false)) { + throw std::runtime_error("model limit reached, try again later"); } // prepare new instance info @@ -1473,6 +1291,9 @@ void server_models::load(const std::string & name, const load_options & opts) { auto proc = inst.subproc; int port = inst.meta.port; mapping[name] = std::move(inst); + if (opts.role == SERVER_ROLE_INSTANCE) { + sched->on_spawned(lk, name); + } monitor->watch(name, proc, opts.role, port); cv.notify_all(); } @@ -1486,7 +1307,15 @@ void server_models::request_stop(const std::string & name, bool send_exit) { monitor->stop(name, it->second.meta.stop_timeout, send_exit); } +void server_models::watch_child(const std::string & name, std::shared_ptr proc, server_role_type role, int port, int kill_after) { + monitor->watch(name, std::move(proc), role, port, kill_after); +} + void server_models::on_child_exit(const std::string & name, const std::shared_ptr & proc, server_role_type role, int exit_code) { + if (role == SERVER_ROLE_MEASURE) { + sched->on_measure_exit(name, proc, exit_code); + return; + } { std::lock_guard lk(mutex); auto it = mapping.find(name); @@ -1537,6 +1366,7 @@ void server_models::unload(const std::string & name) { void server_models::unload_all() { std::unique_lock lk(mutex); + sched->cancel_measurements(lk); for (auto & [name, inst] : mapping) { if (inst.meta.status == SERVER_MODEL_STATUS_DOWNLOADING) { SRV_INF("cancelling download for model name=%s\n", name.c_str()); @@ -1551,7 +1381,10 @@ void server_models::unload_all() { } } // wait for every child to exit, the monitor force-kills the ones that ignore the exit command - cv.wait(lk, [this]() { + cv.wait(lk, [this, &lk]() { + if (sched->is_measuring(lk)) { + return false; + } for (const auto & [name, inst] : mapping) { if (inst.meta.is_running() || inst.meta.status == SERVER_MODEL_STATUS_DOWNLOADING) { return false; @@ -1598,6 +1431,12 @@ void server_models::update_status(const std::string & name, const update_status_ if (!args.progress.is_null()) { meta.progress = args.progress; } + switch (args.status) { + case SERVER_MODEL_STATUS_LOADED: sched->on_ready(lk, name, args.memory); break; + case SERVER_MODEL_STATUS_SLEEPING: sched->on_sleeping(lk, name); break; + case SERVER_MODEL_STATUS_UNLOADED: sched->on_unloaded(lk, name); break; + default: break; + } // a model that comes up idle or goes down changes the slot count for queued requests sched->tick(lk); } @@ -1615,6 +1454,9 @@ void server_models::update_status(const std::string & name, const update_status_ if (!args.progress.is_null()) { data["progress"] = args.progress; } + if (args.memory.has_value()) { + data["memory"] = args.memory->to_json(); + } // note: notify_sse doesn't acquire the lock, so no deadlock here notify_sse("status_change", name, data); } @@ -2019,14 +1861,29 @@ void server_models::handle_child_state(const std::string & name, const std::stri payload, }); } break; + case SERVER_STATE_MEASURING: + { + sched->handle_measured(name, payload); + } break; case SERVER_STATE_READY: { + std::optional memory; + json info = payload; + if (info.contains("memory")) { + try { + memory = server_model_memory::from_json(info.at("memory")); + } catch (const std::exception & e) { + SRV_WRN("malformed memory report from model name=%s: %s\n", name.c_str(), e.what()); + } + info.erase("memory"); + } update_status(name, { SERVER_MODEL_STATUS_LOADED, 0, // note: payload can be empty if this is a wakeup from sleep - payload.size() > 0 ? payload : nullptr, + info.size() > 0 ? info : nullptr, {}, // reset progress info + memory, }); } break; case SERVER_STATE_SLEEPING: @@ -2326,6 +2183,10 @@ void server_models_routes::init_routes() { model_info["context_length"] = meta.n_ctx_train; } + if (models.sched->memory_enabled()) { + model_info["memory"] = meta.memory.to_json(); + } + // the sidecars a model has on disk, resolved as -hf would load them, for a // model that resolves from a HF repo. A download on its way reports the // kinds it pulls, as an unfinished one does diff --git a/tools/server/server-models.h b/tools/server/server-models.h index bfaefaf3f..0db948773 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -5,8 +5,11 @@ #include "preset.h" #include "server-common.h" #include "server-http.h" +#include "server-models-sched.h" #include "server-queue.h" +#include +#include #include #include #include @@ -86,6 +89,9 @@ struct server_model_meta { json architecture = server_model_architecture_json(false, false, false, {"text"}); uint32_t n_ctx_train = 0; // trained context, read from the GGUF metadata; 0 when unknown + // only populated with --models-memory-margin, see server_models_sched + server_model_memory_info memory = {}; + bool is_ready() const { return status == SERVER_MODEL_STATUS_LOADED; } @@ -152,12 +158,11 @@ private: }; struct server_models_routes; -struct server_lru_sched; // defined in server-models.cpp -struct server_monitor; // defined in server-models.cpp +struct server_monitor; // defined in server-models.cpp struct server_models { friend struct server_models_routes; - friend struct server_lru_sched; + friend struct server_models_sched; friend struct server_monitor; private: @@ -252,17 +257,14 @@ private: std::vector base_env; common_preset base_preset; // base preset from llama-server CLI args - // queue of requests waiting for a models_max slot - std::unique_ptr sched; + // request queue (models_max) and memory accounting + std::unique_ptr sched; // if true, add some delay to simulate works (useful for testing) bool debug_fake_timing = false; void update_meta(const std::string & name, const server_model_meta & meta); - // unload least recently used models if the limit is reached - void unload_lru(); - // not thread-safe, caller must hold mutex void add_model(server_model_meta && meta); @@ -270,6 +272,9 @@ private: // not thread-safe, caller must hold mutex void request_stop(const std::string & name, bool send_exit = true); + // let the monitor watch a child; kill_after > 0 force-kills it after that many seconds + void watch_child(const std::string & name, std::shared_ptr proc, server_role_type role, int port, int kill_after = 0); + // called by the monitor once a child exited and was reaped void on_child_exit(const std::string & name, const std::shared_ptr & proc, server_role_type role, int exit_code); @@ -336,6 +341,7 @@ public: int exit_code = 0; // only valid if status == UNLOADED json loaded_info = nullptr; json progress = nullptr; + std::optional memory = std::nullopt; // what the instance allocated, sent with READY }; // update the status of a model instance (thread-safe) // also send SSE notification to /models/sse endpoint @@ -369,8 +375,9 @@ public: // called from the monitor thread // payload per state: // state = loading -> payload = {} (TODO: add progress info) - // state = ready -> payload = model_info (json), or {} if wakeup from sleeping + // state = ready -> payload = model_info (json) with a "memory" field, or {} if wakeup from sleeping // state = sleeping -> payload = {} + // state = measuring -> payload = {result: "measured", min, max} or {result: "failed", error} void handle_child_state(const std::string & name, const std::string & raw_input); private: diff --git a/tools/server/server.cpp b/tools/server/server.cpp index b6ebf7523..8b1d42650 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -424,6 +424,45 @@ static int server_run_download(common_params & params, server_role & role) { return 0; } +// +// measure worker: a child of the router, reports the memory a model needs without loading it +// + +static int server_run_measure(common_params & params, server_role & role, common_models_handler & models_handler) { + // the router sets --offline: only resolve the local path, never download + try { + common_models_handler_apply(models_handler, params); + } catch (const std::exception & e) { + SRV_ERR("failed to resolve model: %s\n", e.what()); + return 1; + } + + const std::string model_name = params.model.get_name(); + + // TODO: mmproj is not in llama_get_memory_breakdown, a multimodal model is underestimated + json payload; + try { + common_memory_reports m = common_measure_memory(params); + payload = { + {"result", "measured"}, + {"n_ctx_train", m.n_ctx_train}, + {"min", server_model_memory(m.n_ctx_min, m.info_min).to_json()}, + {"max", server_model_memory(m.n_ctx_max, m.info_max).to_json()}, + }; + SRV_INF("memory measured for model name=%s at n_ctx=%u (min) and n_ctx=%u (max)\n", + model_name.c_str(), m.n_ctx_min, m.n_ctx_max); + } catch (const std::exception & e) { + SRV_ERR("failed to measure memory for model name=%s: %s\n", model_name.c_str(), e.what()); + payload = { + {"result", "failed"}, + {"error", e.what()}, + }; + } + + role.notify_to_router(server_state_to_str(SERVER_STATE_MEASURING), payload); + return 0; +} + // // router // @@ -664,7 +703,12 @@ static int server_run_instance(common_params & params, server_role & role, commo std::thread monitor_thread; if (role.is_managed_by_router()) { monitor_thread = role.setup(shutdown_handler); - role.notify_to_router(server_state_to_str(SERVER_STATE_READY), routes.get_model_info()); + json info = routes.get_model_info(); + // the router tracks device memory from these reports + if (const auto * ll_ctx = ctx_server.get_llama_context()) { + info["memory"] = server_model_memory(llama_n_ctx(ll_ctx), common_memory_breakdown_rows(ll_ctx)).to_json(); + } + role.notify_to_router(server_state_to_str(SERVER_STATE_READY), info); } // this call blocks the main thread until queue_tasks.terminate() is called @@ -717,6 +761,7 @@ static int server_run(common_params & params, int argc, char ** argv, server_rol switch (role_type) { case SERVER_ROLE_DOWNLOAD: return server_run_download(params, role); + case SERVER_ROLE_MEASURE: return server_run_measure(params, role, models_handler); case SERVER_ROLE_ROUTER: return server_run_router(params, argc, argv, opts); case SERVER_ROLE_INSTANCE: return server_run_instance(params, role, models_handler, opts); } diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 6df41f8d6..98f6aab9e 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -318,6 +318,72 @@ def test_router_queue_two_waiters_share_one_eviction(): assert _get_model_status(MODEL_A) == "unloaded" +# --models-memory-margin tests + +def _get_model_memory(model_id: str) -> dict: + res = server.make_request("GET", "/models") + assert res.status_code == 200 + for item in res.body.get("data", []): + if item.get("id") == model_id: + return item["memory"] + raise AssertionError(f"Model {model_id} not found in /models response") + + +def _wait_for_memory_status(model_id: str, desired: set[str], timeout: int = 120) -> dict: + deadline = time.time() + timeout + mem = None + while time.time() < deadline: + mem = _get_model_memory(model_id) + if mem["status"] in desired: + return mem + time.sleep(0.5) + raise AssertionError(f"memory status of {model_id} is {mem and mem['status']}, expected one of {desired}") + + +def test_router_memory_margin_measures_models(): + """every model is measured in the background on startup, a loaded instance reports what it allocated""" + global server + server.models_memory_margin = 1 + server.start() + + mem = _wait_for_memory_status(MODEL_B, {"measured", "failed"}) + assert mem["status"] == "measured" + assert mem["effective"] is None + assert mem["min"]["n_ctx"] > 0 + assert mem["max"]["n_ctx"] >= mem["min"]["n_ctx"] + kinds = {d["kind"] for d in mem["min"]["devices"]} + assert "host" in kinds + for d in mem["min"]["devices"]: + assert d["model"] + d["context"] + d["compute"] > 0 or d["kind"] != "host" + + _load_model_and_wait(MODEL_B, timeout=120) + eff = _get_model_memory(MODEL_B)["effective"] + assert eff is not None + assert eff["n_ctx"] == mem["min"]["n_ctx"] # the test preset pins -c, so min == max == effective + + server.make_request("POST", "/models/unload", data={"model": MODEL_B}) + _wait_for_model_status(MODEL_B, {"unloaded"}) + assert _get_model_memory(MODEL_B)["effective"] is None + + +def test_router_memory_margin_evicts_lru(): + """with a margin larger than any device, two models never fit together: the idle one is unloaded""" + global server + server.models_max = 4 + server.models_memory_margin = 1 << 30 # MiB + server.start() + + # nothing else is loaded, so the model is loaded although it does not fit the margin + _load_model_and_wait(MODEL_B, timeout=120) + eff = _get_model_memory(MODEL_B)["effective"] + if not any(d["kind"] == "device" for d in eff["devices"]): + pytest.skip("no GPU device, the memory limit does not apply") # ty: ignore[too-many-positional-arguments] + + _load_model_and_wait(MODEL_C, timeout=120) + assert _get_model_status(MODEL_C) == "loaded" + assert _get_model_status(MODEL_B) == "unloaded" + + def test_router_no_models_autoload(): global server server.no_models_autoload = True diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a712ca9ee..0ca49192a 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -94,6 +94,7 @@ class ServerProcess: api_key: str | None = None models_dir: str | None = None models_max: int | None = None + models_memory_margin: int | None = None models_preset: str | None = None no_models_autoload: bool | None = None lora_files: List[str] | None = None @@ -188,6 +189,8 @@ class ServerProcess: server_args.extend(["--models-dir", self.models_dir]) if self.models_max is not None: server_args.extend(["--models-max", self.models_max]) + if self.models_memory_margin is not None: + server_args.extend(["--models-memory-margin", self.models_memory_margin]) if self.models_preset: server_args.extend(["--models-preset", self.models_preset]) if self.cors_origins: