diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b03adacd3d3..1005d386f9d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3323,6 +3323,26 @@ struct server_context_impl { case SERVER_TASK_TYPE_LEARN_GRANTS: process_learn_grants(task); break; + case SERVER_TASK_TYPE_PROMPT_CACHE_LIMIT: + { + // The host prompt cache gives its bytes back: a cache is never need (card 494c99f1). + // The limit is lowered or raised, and update() evicts the oldest states until the + // cache fits it. The reply is what the cache holds now, so the core frees exactly + // what it measured, never what it hoped. + auto result = std::make_unique(); + result->id = task.id; + if (!prompt_cache) { + result->receipt = {{"refusal", "prompt_cache_disabled"}}; + } else { + const size_t before = prompt_cache->size(); + prompt_cache->limit_size = 1024ull * 1024ull * (size_t) task.prompt_cache_limit_mib; + prompt_cache->update(); + result->receipt = {{"refusal", "none"}, {"limit_mib", task.prompt_cache_limit_mib}, + {"held_bytes", prompt_cache->size()}, {"freed_bytes", before - std::min(before, prompt_cache->size())}, + {"states", prompt_cache->states.size()}}; + } + queue_results.send(std::move(result)); + } break; case SERVER_TASK_TYPE_LEARN: { if (!task.restore_checkpoint.is_null()) { @@ -6578,6 +6598,28 @@ void server_routes::init_routes() { this->release_learn_residency = [grant_route](const server_http_req & req) { return grant_route(req, "residency_release"); }; this->revoke_learn_residency = [grant_route](const server_http_req & req) { return grant_route(req, "residency_revoke"); }; this->get_learn_grants = [grant_route](const server_http_req & req) { return grant_route(req, "list"); }; + // A cache is never need: the core lowers the host prompt cache's limit to give its bytes to an + // active consumer (a learner's step), and raises it again after (card 494c99f1). + this->post_prompt_cache_limit = [this](const server_http_req & req) { + auto res = create_response(); + const json body = req.body.empty() ? json::object() : json::parse(req.body, nullptr, false); + // 0 would mean "no limit" to the cache itself, so a release to nothing is 1 MiB, never 0; and the + // limit in bytes must fit a size_t, so a value past it is refused, never wrapped (Cormac on #110). + constexpr int64_t max_mib = (int64_t) (SIZE_MAX >> 20) < INT64_MAX ? (int64_t) (SIZE_MAX >> 20) : INT64_MAX; + if (!body.is_object() || !body.contains("limit_mib") || !body.at("limit_mib").is_number_integer() + || body.at("limit_mib").get() < 1 || body.at("limit_mib").get() > max_mib) { + res->error(format_error_response("limit_mib must be an integer from 1 to the bytes a size_t holds, in MiB", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + server_task task(SERVER_TASK_TYPE_PROMPT_CACHE_LIMIT); + task.prompt_cache_limit_mib = body.at("limit_mib").get(); + auto & rd = res->rd; + task.id = rd.get_new_id(); + rd.post_task(std::move(task)); + auto result = rd.next(req.should_stop); + if (result) { res->ok(result->to_json()); } + return res; + }; this->post_learn_grants = [grant_route](const server_http_req & req) { return grant_route(req, "register"); }; this->revoke_learn_grant = [grant_route](const server_http_req & req) { return grant_route(req, "revoke"); }; diff --git a/tools/server/server-context.h b/tools/server/server-context.h index a810c3e3a1b..d6acce4840e 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -177,6 +177,7 @@ struct server_routes { server_http_context::handler_t post_learn; server_http_context::handler_t post_learn_residency, release_learn_residency, revoke_learn_residency; server_http_context::handler_t get_learn_grants, post_learn_grants, revoke_learn_grant, post_learn_preflight; + server_http_context::handler_t post_prompt_cache_limit; server_http_context::handler_t get_train; server_http_context::handler_t post_train; server_http_context::handler_t post_train_cancel; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 652697ba4d2..681ee333c18 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -63,6 +63,7 @@ enum server_task_type { SERVER_TASK_TYPE_UNLOAD_LORA, SERVER_TASK_TYPE_LEARN, SERVER_TASK_TYPE_LEARN_GRANTS, + SERVER_TASK_TYPE_PROMPT_CACHE_LIMIT, }; // TODO: change this to more generic "response_format" to replace the "format_response_*" in server-common @@ -288,6 +289,9 @@ struct server_task { std::string unload_engine; json learn_grants; // authority operation, processed on the serving queue + // Continuum fork: the host prompt cache's new limit in MiB (POST /prompt-cache/limit). A cache is + // never need: the core gives its bytes back when an active consumer needs them (card 494c99f1). + int64_t prompt_cache_limit_mib = -1; // used by SERVER_TASK_TYPE_LEARN learn_request learn; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 4a4ec63ab39..dfecafd7389 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -226,6 +226,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.get_learn_grants = models_routes->proxy_get; routes.post_learn_grants = models_routes->proxy_post; routes.post_learn_preflight = models_routes->proxy_post; + routes.post_prompt_cache_limit = models_routes->proxy_post; routes.post_learn_residency = models_routes->proxy_post; routes.release_learn_residency = models_routes->proxy_post; routes.revoke_learn_residency = models_routes->proxy_post; @@ -295,6 +296,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/learn/residency/revoke", ex_wrapper(routes.revoke_learn_residency)); ctx_http.post("/learn/preflight", ex_wrapper(routes.post_learn_preflight)); ctx_http.post("/learn/grants", ex_wrapper(routes.post_learn_grants)); + ctx_http.post("/prompt-cache/limit", ex_wrapper(routes.post_prompt_cache_limit)); ctx_http.post("/learn/grants/revoke", ex_wrapper(routes.revoke_learn_grant)); ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); diff --git a/tools/server/tests/prompt-cache-limit.py b/tools/server/tests/prompt-cache-limit.py new file mode 100755 index 00000000000..243ee1d5bfe --- /dev/null +++ b/tools/server/tests/prompt-cache-limit.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +"""Real server test; requires an isolated server launched with a host prompt cache (--cache-ram N, +one slot). Usage: prompt-cache-limit.py URL + +What this catches (card 494c99f1): the host prompt cache counted as serving's need, so a 4 MB learner +waited forever on the M5 ('0 of 4096000 governed Vram bytes are free') while ~13 GB of cache held +states nothing predicted reusing. The cache gives its bytes back: POST /prompt-cache/limit lowers its +limit and evicts the oldest states until it fits, and replies what it now holds and what it freed, so +the core frees exactly what was measured. Raising the limit again lets it refill. +""" +import os +import sys + +sys.path.insert(0, os.path.dirname(__file__)) +from learn_fixture import Engine + +engine = Engine(sys.argv[1], "cache-limit", timeout=120) + +# Two different prompts on one slot: the first is saved to the host cache when the second replaces it. +for prompt in ("Describe a lighthouse at dawn in one sentence:", "List three uses for a paperclip:"): + engine.ok("/completion", {"prompt": prompt, "n_predict": 8, "temperature": 0}) + +status, refused = engine.post("/prompt-cache/limit", {"limit_mib": 0}) +assert status == 400, ("0 means 'no limit' to the cache, so it is refused, never taken as unlimited", status, refused) + +status, overflow = engine.post("/prompt-cache/limit", {"limit_mib": 2**62}) +assert status == 400, ("a limit whose bytes overflow a size_t is refused, never wrapped", status, overflow) + +lowered = engine.ok("/prompt-cache/limit", {"limit_mib": 1}) +assert lowered["refusal"] == "none", lowered +assert lowered["freed_bytes"] > 0, ("the filled cache gave bytes back", lowered) +assert lowered["held_bytes"] <= 1024 * 1024, ("the cache fits its new limit", lowered) + +raised = engine.ok("/prompt-cache/limit", {"limit_mib": 512}) +assert raised["refusal"] == "none" and raised["limit_mib"] == 512, raised + +engine.close() +print("ok: the host prompt cache gives its bytes back on request and reports what it freed")