mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-28 17:18:49 -04:00
feat(llama-cpp): add slot prompt-cache save and restore API (#12730)
* feat(llama-cpp): add slot prompt-cache save and restore API Add POST /v1/slots/save and POST /v1/slots/restore. They write one llama.cpp slot's prompt cache (KV state) to a file and load it back, so a client can warm a long system prompt in a few hundred milliseconds instead of re-running prefill. Refs #12685. The new SlotCache gRPC call posts the same slot save and restore tasks the llama-server /slots endpoint uses. It is off unless the model sets the new slot_save_path option, so existing models are not affected. The /slots monitoring endpoint is not needed here: the backend has no HTTP server. The file name comes from the caller, so it is checked twice, in Go and again in the backend. It must be one path component: no empty name, "/", "\", "..", NUL, control characters or absolute path. The resolved path, with symbolic links followed, must be a direct child of slot_save_path. The routes are admin only because they write files on the backend host. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(llama-cpp): answer 501 when a backend lacks slot cache A backend that does not implement the SlotCache rpc returns gRPC UNIMPLEMENTED. The API passed that text through, so the client saw "unknown method SlotCache". Map the status to a typed error in core/backend and return 501 with a fixed message instead. Every other error path is unchanged. Make the SlotCache comments in backend.proto backend-neutral: slot is the backend's state slot, filename must stay inside the save directory, and an empty model_identity means the backend does not check it. Only comments change, no field or name. Add tests that the identity sent is the configured model, and for the error mapping. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(llama-cpp): issue slot cache IDs on the server The slot save and restore API let the client pick the file name. Restore could load any file in slot_save_path, and the symlink check ran before llama.cpp opened the file. The backend now issues the name. A save returns a random 128-bit ID (32 lowercase hex characters, read from /dev/urandom) and never overwrites an existing file. Files live in a 0700 directory per model, named from the SHA-256 of the model identity, with a manifest.json that records each save. The manifest is written atomically and a corrupt one stops both save and restore. A restore takes the ID, finds it in the manifest, refuses an entry of another model, builds the path from the entry and checks with lstat that it is a plain file of the recorded size. The directory mode is the main defence against another local user, since llama.cpp still opens the path itself. Each model keeps at most 32 saves and the oldest is deleted first. A file over 8 GiB is deleted after the save, because llama.cpp does not report the size before it writes. Responses carry no host path. The HTTP bodies change: save takes {model, slot}, restore takes {model, slot, id}, and an unknown ID returns 404. In the proto, SlotCacheRequest.filename and SlotCacheResponse.filename become id. SlotCache no longer calls checkModelIdentity, which reads a field this request does not have. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(llama-cpp): copy slot_cache files in prepare.sh grpc-server.cpp includes slot_cache.h, but prepare.sh copies the helper headers into llama.cpp/tools/grpc-server/ one by one and skipped it. The llama-cpp, turboquant and bonsai builds share this step, so they all failed with "slot_cache.h: No such file or directory". Copy slot_cache.h and slot_cache_test.cpp, which CMakeLists.txt builds when LLAMA_GRPC_BUILD_TESTS is on. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * chore(deps): bump LocalAGI to pick up the scheduler start/stop fix LocalAGI's scheduler could hit a nil pointer in run when Start and Stop raced. That made the core/services/agentpool tests panic at random and turned tests-linux red. Move to mudler/LocalAGI@7aa0c8fea0, the merge of mudler/LocalAGI#519 ("make Start and Stop safe to call concurrently"). The new module version also adds a few indirect requirements to go.mod and go.sum. Assisted-by: Claude:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
This commit is contained in:
29 files changed
+1989
-1
No files matched your search
@@ -62,6 +62,7 @@
|
||||
#include "model_load_error.h"
|
||||
#include "thread_params.h"
|
||||
#include "parallel_params.h"
|
||||
#include "slot_cache.h"
|
||||
#include "message_content.h"
|
||||
#include "moe_cache_option.h"
|
||||
#include "passthrough_options.h"
|
||||
@@ -78,10 +79,12 @@
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <iterator>
|
||||
#include <list>
|
||||
#include <map>
|
||||
#include <ctime>
|
||||
#include <mutex>
|
||||
#include <signal.h>
|
||||
#include <thread>
|
||||
@@ -678,6 +681,35 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
|
||||
// If conversion fails, keep the default value
|
||||
}
|
||||
}
|
||||
} else if (!strcmp(optname, "slot_save_path")) {
|
||||
// Directory for SlotCache save/restore files. Take everything after
|
||||
// the first ':' so a path that contains one is kept whole. Without
|
||||
// this option the RPC stays disabled, as in llama-server.
|
||||
const size_t sep = opt.find(':');
|
||||
const std::string dir = (sep == std::string::npos) ? "" : opt.substr(sep + 1);
|
||||
if (!dir.empty()) {
|
||||
// A missing directory is created with mode 0700. An existing
|
||||
// one is left as it is, with a warning when others can read it.
|
||||
std::string warn, derr;
|
||||
#ifdef _WIN32
|
||||
std::error_code ec;
|
||||
std::filesystem::create_directories(std::filesystem::u8path(dir), ec);
|
||||
const bool dir_ok = !ec && std::filesystem::is_directory(std::filesystem::u8path(dir));
|
||||
#else
|
||||
const bool dir_ok = llama_grpc::slot_ensure_private_dir(dir, warn, derr);
|
||||
#endif
|
||||
if (!warn.empty()) {
|
||||
fprintf(stderr, "slot_save_path: %s\n", warn.c_str());
|
||||
}
|
||||
if (!dir_ok) {
|
||||
fprintf(stderr, "slot_save_path %s is not a usable directory (%s), slot save/restore stays disabled\n", dir.c_str(), derr.c_str());
|
||||
} else {
|
||||
params.slot_save_path = dir;
|
||||
if (params.slot_save_path.back() != '/' && params.slot_save_path.back() != '\\') {
|
||||
params.slot_save_path += DIRECTORY_SEPARATOR;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (!strcmp(optname, "parallel") || !strcmp(optname, "n_parallel")) {
|
||||
if (optval != NULL) {
|
||||
try {
|
||||
@@ -3742,6 +3774,157 @@ public:
|
||||
return grpc::Status::OK;
|
||||
}
|
||||
|
||||
// slotStatusToGrpc maps a slot cache helper status to a gRPC status.
|
||||
static grpc::Status slotStatusToGrpc(llama_grpc::SlotStatus s, const std::string& msg) {
|
||||
switch (s) {
|
||||
case llama_grpc::SlotStatus::NotFound: return grpc::Status(grpc::StatusCode::NOT_FOUND, msg);
|
||||
case llama_grpc::SlotStatus::Invalid: return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, msg);
|
||||
case llama_grpc::SlotStatus::Corrupt: return grpc::Status(grpc::StatusCode::DATA_LOSS, msg);
|
||||
default: return grpc::Status(grpc::StatusCode::INTERNAL, msg);
|
||||
}
|
||||
}
|
||||
|
||||
// SlotCache saves or restores the prompt cache of one slot. It reuses the
|
||||
// task the llama-server /slots endpoint posts, so the file format is the
|
||||
// same. The caller never names a file. A save writes
|
||||
// <slot_save_path>/<hash of the model identity>/<id> with an ID the backend
|
||||
// generates, and records it in that directory's manifest.json. A restore
|
||||
// takes an ID, looks it up in the manifest and builds the path from the
|
||||
// manifest entry (see slot_cache.h for the checks and their limits).
|
||||
grpc::Status SlotCache(ServerContext* context, const backend::SlotCacheRequest* request, backend::SlotCacheResponse* response) override {
|
||||
#ifdef _WIN32
|
||||
(void)context; (void)request; (void)response;
|
||||
return grpc::Status(grpc::StatusCode::UNIMPLEMENTED, "slot save/restore is not supported on Windows");
|
||||
#else
|
||||
auto auth = checkAuth(context);
|
||||
if (!auth.ok()) return auth;
|
||||
// Same rule as checkModelIdentity; written out because this request
|
||||
// names the field model_identity, so it has no modelidentity().
|
||||
if (!request->model_identity().empty() && !loaded_model_identity.empty() &&
|
||||
loaded_model_identity != request->model_identity()) {
|
||||
return grpc::Status(grpc::StatusCode::NOT_FOUND,
|
||||
"llama-cpp: model identity mismatch: loaded \"" + loaded_model_identity +
|
||||
"\", requested \"" + request->model_identity() + "\"");
|
||||
}
|
||||
if (params_base.model.path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
|
||||
}
|
||||
if (params_base.slot_save_path.empty()) {
|
||||
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
|
||||
"slot save/restore is disabled: set the slot_save_path option for this model");
|
||||
}
|
||||
|
||||
const bool is_save = request->action() == backend::SLOT_CACHE_ACTION_SAVE;
|
||||
if (!is_save && request->action() != backend::SLOT_CACHE_ACTION_RESTORE) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "action must be save or restore");
|
||||
}
|
||||
if (request->slot() < 0) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "slot must not be negative");
|
||||
}
|
||||
|
||||
// The identity that names the directory and is written to the manifest.
|
||||
std::string model_id = !loaded_model_identity.empty() ? loaded_model_identity : request->model_identity();
|
||||
if (model_id.empty()) {
|
||||
model_id = params_base.model.path;
|
||||
}
|
||||
|
||||
std::string err, dir, id, filepath;
|
||||
llama_grpc::SlotEntry entry;
|
||||
if (is_save) {
|
||||
if (!request->id().empty()) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "id must be empty on save: the backend issues it");
|
||||
}
|
||||
auto s = llama_grpc::slot_identity_dir(params_base.slot_save_path, model_id, true, dir, err);
|
||||
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
|
||||
// Fail before writing a large file if the manifest is unusable.
|
||||
std::vector<llama_grpc::SlotEntry> existing;
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(llama_grpc::slot_cache_mutex());
|
||||
s = llama_grpc::slot_manifest_load(dir, existing, err);
|
||||
}
|
||||
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
|
||||
s = llama_grpc::slot_reserve_id(dir, id, filepath, err);
|
||||
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
|
||||
} else {
|
||||
auto s = llama_grpc::slot_lookup_restore(params_base.slot_save_path, model_id, request->id(), entry, filepath, err);
|
||||
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
|
||||
id = entry.id;
|
||||
}
|
||||
|
||||
// On a failed save the reserved (empty or partial) file is removed.
|
||||
auto drop_reserved = [&]() { if (is_save) { std::error_code rec; std::filesystem::remove(std::filesystem::u8path(filepath), rec); } };
|
||||
|
||||
server_task task(is_save ? SERVER_TASK_TYPE_SLOT_SAVE : SERVER_TASK_TYPE_SLOT_RESTORE);
|
||||
auto rd = ctx_server.get_response_reader();
|
||||
task.id = rd.get_new_id();
|
||||
task.slot_action.id_slot = request->slot();
|
||||
task.slot_action.filename = id;
|
||||
task.slot_action.filepath = filepath;
|
||||
rd.post_task(std::move(task));
|
||||
|
||||
auto res = rd.next([context]() { return context->IsCancelled(); });
|
||||
if (!res) {
|
||||
drop_reserved();
|
||||
return grpc::Status(grpc::StatusCode::CANCELLED, "slot cache request cancelled");
|
||||
}
|
||||
if (res->is_error()) {
|
||||
drop_reserved();
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, res->to_json().dump());
|
||||
}
|
||||
auto * slot_res = dynamic_cast<server_task_result_slot_save_load *>(res.get());
|
||||
if (slot_res == nullptr) {
|
||||
drop_reserved();
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "unexpected result type for a slot cache task");
|
||||
}
|
||||
|
||||
uint64_t n_bytes = slot_res->n_bytes;
|
||||
std::string message;
|
||||
if (is_save) {
|
||||
// The size is only known after llama.cpp wrote the file. Use the
|
||||
// size on disk: restore compares against it.
|
||||
std::error_code sec;
|
||||
const auto on_disk = std::filesystem::file_size(std::filesystem::u8path(filepath), sec);
|
||||
if (sec) {
|
||||
drop_reserved();
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "cannot read the size of the saved file");
|
||||
}
|
||||
n_bytes = static_cast<uint64_t>(on_disk);
|
||||
if (n_bytes > llama_grpc::SLOT_CACHE_MAX_FILE_BYTES) {
|
||||
drop_reserved();
|
||||
return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
|
||||
"the saved state is larger than the " + std::to_string(llama_grpc::SLOT_CACHE_MAX_FILE_BYTES) +
|
||||
" byte limit and was deleted");
|
||||
}
|
||||
entry.id = id;
|
||||
entry.model_identity = model_id;
|
||||
entry.slot = request->slot();
|
||||
entry.n_tokens = slot_res->n_tokens;
|
||||
entry.n_bytes = n_bytes;
|
||||
entry.created_unix = static_cast<int64_t>(std::time(nullptr));
|
||||
std::vector<std::string> pruned;
|
||||
auto s = llama_grpc::slot_commit_save(dir, entry, pruned, err);
|
||||
if (s != llama_grpc::SlotStatus::Ok) {
|
||||
drop_reserved();
|
||||
return slotStatusToGrpc(s, err);
|
||||
}
|
||||
if (!pruned.empty()) {
|
||||
message = "pruned the oldest saved state to keep at most " +
|
||||
std::to_string(llama_grpc::SLOT_CACHE_MAX_ENTRIES) + ":";
|
||||
for (const auto & p : pruned) message += " " + p;
|
||||
}
|
||||
}
|
||||
|
||||
response->set_success(true);
|
||||
response->set_message(message);
|
||||
response->set_slot(slot_res->id_slot);
|
||||
response->set_id(id);
|
||||
response->set_n_tokens(slot_res->n_tokens);
|
||||
response->set_n_bytes(n_bytes);
|
||||
response->set_duration_ms(slot_res->t_ms);
|
||||
return grpc::Status::OK;
|
||||
#endif
|
||||
}
|
||||
|
||||
grpc::Status Detokenize(ServerContext* context, const backend::DetokenizeRequest* request, backend::DetokenizeResponse* response) override {
|
||||
auto auth = checkAuth(context);
|
||||
if (!auth.ok()) return auth;
|
||||
|
||||
Reference in new issue
Block a user