From e715af1defc92eb66110450083b780df774c75ff Mon Sep 17 00:00:00 2001 From: mudler-agent Date: Sat, 17 Oct 2026 18:23:02 +0200 Subject: [PATCH] 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 * 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 * 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 * 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 * 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@7aa0c8fea076, 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 --------- Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- backend/backend.proto | 35 ++ backend/cpp/llama-cpp/CMakeLists.txt | 6 + backend/cpp/llama-cpp/grpc-server.cpp | 183 ++++++ backend/cpp/llama-cpp/prepare.sh | 4 + backend/cpp/llama-cpp/slot_cache.h | 556 ++++++++++++++++++ backend/cpp/llama-cpp/slot_cache_test.cpp | 229 ++++++++ core/backend/slot_cache.go | 119 ++++ core/backend/slot_cache_test.go | 75 +++ .../localai/export_slot_cache_test.go | 4 + core/http/endpoints/localai/slot_cache.go | 97 +++ .../http/endpoints/localai/slot_cache_test.go | 124 ++++ core/http/routes/localai.go | 15 + core/schema/tokenize.go | 20 + core/services/nodes/health_mock_test.go | 3 + core/services/nodes/inflight_test.go | 3 + docs/content/features/authentication.md | 1 + docs/content/features/text-generation.md | 45 ++ docs/content/reference/api-errors.md | 12 + go.mod | 10 +- go.sum | 15 + pkg/grpc/backend.go | 3 + pkg/grpc/client.go | 17 + pkg/grpc/embed.go | 4 + pkg/slotcache/slotcache.go | 49 ++ pkg/slotcache/slotcache_suite_test.go | 13 + pkg/slotcache/slotcache_test.go | 56 ++ swagger/docs.go | 107 ++++ swagger/swagger.json | 107 ++++ swagger/swagger.yaml | 78 +++ 29 files changed, 1989 insertions(+), 1 deletion(-) create mode 100644 backend/cpp/llama-cpp/slot_cache.h create mode 100644 backend/cpp/llama-cpp/slot_cache_test.cpp create mode 100644 core/backend/slot_cache.go create mode 100644 core/backend/slot_cache_test.go create mode 100644 core/http/endpoints/localai/export_slot_cache_test.go create mode 100644 core/http/endpoints/localai/slot_cache.go create mode 100644 core/http/endpoints/localai/slot_cache_test.go create mode 100644 pkg/slotcache/slotcache.go create mode 100644 pkg/slotcache/slotcache_suite_test.go create mode 100644 pkg/slotcache/slotcache_test.go diff --git a/backend/backend.proto b/backend/backend.proto index f6a111121..1c8b273f8 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -76,6 +76,11 @@ service Backend { rpc GetMetrics(MetricsRequest) returns (MetricsResponse); + // SlotCache saves or restores the backend's per-slot prompt state (KV or + // session state). On save the backend issues an ID; on restore the client + // sends an ID the backend issued. The client never names a file. A backend that does not support it returns gRPC UNIMPLEMENTED. + rpc SlotCache(SlotCacheRequest) returns (SlotCacheResponse) {} + rpc VAD(VADRequest) returns (VADResponse) {} rpc Diarize(DiarizeRequest) returns (DiarizeResponse) {} @@ -929,6 +934,36 @@ message TokenizationResponse { repeated int32 tokens = 2; } +enum SlotCacheAction { + SLOT_CACHE_ACTION_UNSPECIFIED = 0; + SLOT_CACHE_ACTION_SAVE = 1; + SLOT_CACHE_ACTION_RESTORE = 2; +} + +message SlotCacheRequest { + // Ties a save file to the model it came from. Empty means the backend does + // not check it. + string model_identity = 1; + SlotCacheAction action = 2; + // The backend's state slot id, 0 is the first slot. For llama.cpp it is the + // server slot. Other backends may map it to a session id. + int32 slot = 3; + // Save: leave empty, the backend issues the ID. Restore: an ID that the + // backend issued for this model. The client never chooses a file name or a + // path; the backend must refuse any ID it did not issue for this model. + string id = 4; +} + +message SlotCacheResponse { + bool success = 1; + string message = 2; // error detail when success is false + int32 slot = 3; + string id = 4; // save: the ID the backend issued. restore: the ID restored + uint64 n_tokens = 5; // tokens saved or restored + uint64 n_bytes = 6; // bytes written or read + double duration_ms = 7; +} + message DetokenizeRequest { repeated int32 tokens = 1; } diff --git a/backend/cpp/llama-cpp/CMakeLists.txt b/backend/cpp/llama-cpp/CMakeLists.txt index 8f4c324a1..773a13b5e 100644 --- a/backend/cpp/llama-cpp/CMakeLists.txt +++ b/backend/cpp/llama-cpp/CMakeLists.txt @@ -143,6 +143,12 @@ if(LLAMA_GRPC_BUILD_TESTS) target_compile_features(parallel_params_test PRIVATE cxx_std_17) add_test(NAME parallel_params_test COMMAND parallel_params_test) + add_executable(slot_cache_test slot_cache_test.cpp slot_cache.h) + target_include_directories(slot_cache_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + target_link_libraries(slot_cache_test PRIVATE ${_LLAMA_COMMON_TARGET}) # vendored nlohmann/json + target_compile_features(slot_cache_test PRIVATE cxx_std_17) + add_test(NAME slot_cache_test COMMAND slot_cache_test) + add_executable(model_load_error_test model_load_error_test.cpp model_load_error.h) target_include_directories(model_load_error_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) target_compile_features(model_load_error_test PRIVATE cxx_std_17) diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index 9f7862b51..f8ba18f40 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -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 #include #include +#include #include #include #include #include +#include #include #include #include @@ -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 + // // 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 existing; + { + std::lock_guard 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(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(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(std::time(nullptr)); + std::vector 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; diff --git a/backend/cpp/llama-cpp/prepare.sh b/backend/cpp/llama-cpp/prepare.sh index 49afc16c8..cad79e2c3 100644 --- a/backend/cpp/llama-cpp/prepare.sh +++ b/backend/cpp/llama-cpp/prepare.sh @@ -80,6 +80,10 @@ cp -r parent_watch_test.cpp llama.cpp/tools/grpc-server/ # test (run via backend/cpp/run-unit-tests.sh; also buildable under ctest). cp -r stream_peer.h llama.cpp/tools/grpc-server/ cp -r stream_peer_test.cpp llama.cpp/tools/grpc-server/ +# Slot cache save/restore helpers (included by grpc-server.cpp) and their +# unit test (compiled only when -DLLAMA_GRPC_BUILD_TESTS=ON). +cp -r slot_cache.h llama.cpp/tools/grpc-server/ +cp -r slot_cache_test.cpp llama.cpp/tools/grpc-server/ cp -rfv llama.cpp/vendor/nlohmann/json.hpp llama.cpp/tools/grpc-server/ cp -rfv llama.cpp/vendor/cpp-httplib/httplib.h llama.cpp/tools/grpc-server/ diff --git a/backend/cpp/llama-cpp/slot_cache.h b/backend/cpp/llama-cpp/slot_cache.h new file mode 100644 index 000000000..2dbd10adb --- /dev/null +++ b/backend/cpp/llama-cpp/slot_cache.h @@ -0,0 +1,556 @@ +#pragma once + +// Helpers for the SlotCache RPC. They are header-only and free of llama.cpp +// and gRPC types so they can be tested on their own. +// +// The client never chooses a file name or a path. A save gets a random ID from +// the backend and writes //, where comes from +// the model identity. A manifest.json in lists every ID the backend +// issued. A restore accepts an ID only if the manifest lists it. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +#include +#include +#include +#include +#endif + +namespace llama_grpc { + +// Saved files kept per model directory. The oldest entry is deleted when a +// save would go over this number. +constexpr size_t SLOT_CACHE_MAX_ENTRIES = 32; + +// Largest file one save may leave on disk. llama.cpp does not report the size +// before it writes, so the check runs after the save and the file is deleted +// when it is too large. 8 GiB keeps one model directory at or below 256 GiB +// (32 files) and is far above a typical saved prompt (hundreds of MiB). The +// disk can briefly hold one file over the limit while the check runs. +constexpr uint64_t SLOT_CACHE_MAX_FILE_BYTES = 8ULL * 1024 * 1024 * 1024; + +constexpr const char* SLOT_CACHE_MANIFEST = "manifest.json"; + +enum class SlotStatus { + Ok, + NotFound, // unknown ID, or an entry that belongs to another model + Invalid, // bad ID, or the file on disk is not what the manifest says + Corrupt, // the manifest cannot be read as written by this code + Internal, // disk, permission or randomness failure +}; + +struct SlotEntry { + std::string id; + std::string model_identity; + int32_t slot = 0; + uint64_t n_tokens = 0; + uint64_t n_bytes = 0; + int64_t created_unix = 0; +}; + +// slot_id_valid accepts exactly 32 lowercase hex characters. Nothing else can +// contain a separator, a dot or a NUL, so a valid ID is always one plain path +// component. +inline bool slot_id_valid(const std::string& id) { + if (id.size() != 32) { + return false; + } + for (const char c : id) { + if (!((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'))) { + return false; + } + } + return true; +} + +inline std::string slot_hex(const unsigned char* data, size_t n) { + static const char* digits = "0123456789abcdef"; + std::string out; + out.reserve(n * 2); + for (size_t i = 0; i < n; i++) { + out.push_back(digits[data[i] >> 4]); + out.push_back(digits[data[i] & 0x0F]); + } + return out; +} + +// slot_random_id draws 128 bits from the operating system and renders them as +// 32 lowercase hex characters. /dev/urandom is the kernel CSPRNG; +// std::random_device is not guaranteed to be one, so it is not used. Returns +// false when no secure source is available (always on Windows, where this +// feature is not supported). +inline bool slot_random_id(std::string& out) { +#ifdef _WIN32 + (void)out; + return false; +#else + unsigned char buf[16]; + const int fd = ::open("/dev/urandom", O_RDONLY | O_CLOEXEC); + if (fd < 0) { + return false; + } + size_t got = 0; + while (got < sizeof(buf)) { + const ssize_t r = ::read(fd, buf + got, sizeof(buf) - got); + if (r <= 0) { + ::close(fd); + return false; + } + got += static_cast(r); + } + ::close(fd); + out = slot_hex(buf, sizeof(buf)); + return true; +#endif +} + +// SHA-256 (FIPS 180-4), used only to turn a model identity into a directory +// name. Kept here so the header stays free of other dependencies. +inline std::string slot_sha256_hex(const std::string& msg) { + static const uint32_t K[64] = { + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5, + 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174, + 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, + 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, + 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, + 0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, + 0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, + 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2}; + uint32_t h[8] = {0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19}; + + std::string data = msg; + const uint64_t bits = static_cast(msg.size()) * 8; + data.push_back(static_cast(0x80)); + while (data.size() % 64 != 56) { + data.push_back('\0'); + } + for (int i = 7; i >= 0; i--) { + data.push_back(static_cast((bits >> (i * 8)) & 0xFF)); + } + + auto rotr = [](uint32_t x, int n) { return (x >> n) | (x << (32 - n)); }; + for (size_t off = 0; off < data.size(); off += 64) { + uint32_t w[64]; + for (int i = 0; i < 16; i++) { + const unsigned char* p = reinterpret_cast(data.data()) + off + i * 4; + w[i] = (uint32_t(p[0]) << 24) | (uint32_t(p[1]) << 16) | (uint32_t(p[2]) << 8) | uint32_t(p[3]); + } + for (int i = 16; i < 64; i++) { + const uint32_t s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >> 3); + const uint32_t s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >> 10); + w[i] = w[i - 16] + s0 + w[i - 7] + s1; + } + uint32_t a = h[0], b = h[1], c = h[2], d = h[3], e = h[4], f = h[5], g = h[6], hh = h[7]; + for (int i = 0; i < 64; i++) { + const uint32_t S1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25); + const uint32_t ch = (e & f) ^ (~e & g); + const uint32_t t1 = hh + S1 + ch + K[i] + w[i]; + const uint32_t S0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22); + const uint32_t mj = (a & b) ^ (a & c) ^ (b & c); + const uint32_t t2 = S0 + mj; + hh = g; g = f; f = e; e = d + t1; d = c; c = b; b = a; a = t1 + t2; + } + h[0] += a; h[1] += b; h[2] += c; h[3] += d; h[4] += e; h[5] += f; h[6] += g; h[7] += hh; + } + + unsigned char digest[32]; + for (int i = 0; i < 8; i++) { + digest[i * 4] = static_cast(h[i] >> 24); + digest[i * 4 + 1] = static_cast(h[i] >> 16); + digest[i * 4 + 2] = static_cast(h[i] >> 8); + digest[i * 4 + 3] = static_cast(h[i]); + } + return slot_hex(digest, sizeof(digest)); +} + +// slot_identity_dir_name is the first 32 hex characters of the SHA-256 of the +// model identity. +inline std::string slot_identity_dir_name(const std::string& identity) { + return slot_sha256_hex(identity).substr(0, 32); +} + +// One lock for every manifest read-modify-write in this process. Two backend +// processes that share one slot_save_path and one model are not coordinated. +inline std::mutex& slot_cache_mutex() { + static std::mutex m; + return m; +} + +#ifndef _WIN32 + +// slot_ensure_private_dir makes sure path is a directory. A missing directory +// is created with mode 0700 (parents get the default mode). An existing one is +// left as it is; if group or others can access it, warn is set so the caller +// can log it. +inline bool slot_ensure_private_dir(const std::string& path, std::string& warn, std::string& err) { + namespace fs = std::filesystem; + warn.clear(); + std::error_code ec; + const fs::path p = fs::u8path(path); + struct stat st; + if (::lstat(p.c_str(), &st) != 0) { + if (errno != ENOENT) { + err = "cannot inspect " + path + ": " + std::strerror(errno); + return false; + } + if (p.has_parent_path()) { + fs::create_directories(p.parent_path(), ec); + if (ec) { + err = "cannot create the parent of " + path + ": " + ec.message(); + return false; + } + } + if (::mkdir(p.c_str(), 0700) != 0 && errno != EEXIST) { + err = "cannot create " + path + ": " + std::strerror(errno); + return false; + } + if (::lstat(p.c_str(), &st) != 0) { + err = "cannot inspect " + path + ": " + std::strerror(errno); + return false; + } + } + // stat, not lstat, for the final check: a configured slot_save_path may be + // a link to a directory. The identity directories below it are checked + // with lstat. + struct stat target; + if (::stat(p.c_str(), &target) != 0 || !S_ISDIR(target.st_mode)) { + err = path + " is not a directory"; + return false; + } + if ((target.st_mode & 077) != 0) { + warn = path + " is accessible by group or others (mode " + [&] { + char buf[8]; + std::snprintf(buf, sizeof(buf), "%04o", static_cast(target.st_mode & 07777)); + return std::string(buf); + }() + "); mode 0700 is advised because saved prompts can contain private text"; + } + return true; +} + +// slot_identity_dir returns /. With create it makes the +// directory (mode 0700) when it is missing. The directory must be a real +// directory, not a link, and owned by the running user. +inline SlotStatus slot_identity_dir(const std::string& root, const std::string& identity, bool create, + std::string& dir, std::string& err) { + namespace fs = std::filesystem; + if (root.empty()) { + err = "slot_save_path is not set for this model"; + return SlotStatus::Internal; + } + dir = (fs::u8path(root) / slot_identity_dir_name(identity)).string(); + struct stat st; + if (::lstat(dir.c_str(), &st) != 0) { + if (errno != ENOENT) { + err = std::string("cannot inspect the slot directory: ") + std::strerror(errno); + return SlotStatus::Internal; + } + if (!create) { + return SlotStatus::NotFound; + } + if (::mkdir(dir.c_str(), 0700) != 0 && errno != EEXIST) { + err = std::string("cannot create the slot directory: ") + std::strerror(errno); + return SlotStatus::Internal; + } + if (::lstat(dir.c_str(), &st) != 0) { + err = std::string("cannot inspect the slot directory: ") + std::strerror(errno); + return SlotStatus::Internal; + } + } + if (!S_ISDIR(st.st_mode)) { + err = "the slot directory is not a plain directory"; + return SlotStatus::Internal; + } + if (st.st_uid != ::geteuid() || (st.st_mode & 077) != 0) { + err = "the slot directory must be owned by the backend user and have mode 0700"; + return SlotStatus::Internal; + } + return SlotStatus::Ok; +} + +inline nlohmann::json slot_entry_to_json(const SlotEntry& e) { + return nlohmann::json{{"id", e.id}, + {"model_identity", e.model_identity}, + {"slot", e.slot}, + {"n_tokens", e.n_tokens}, + {"n_bytes", e.n_bytes}, + {"created_unix", e.created_unix}}; +} + +// slot_manifest_load reads /manifest.json. A missing file is an empty +// list. A file that does not parse, or has an entry that is not well formed, +// returns Corrupt and is never reset: the caller must stop. +inline SlotStatus slot_manifest_load(const std::string& dir, std::vector& entries, std::string& err) { + namespace fs = std::filesystem; + entries.clear(); + const std::string path = (fs::u8path(dir) / SLOT_CACHE_MANIFEST).string(); + struct stat st; + if (::lstat(path.c_str(), &st) != 0) { + if (errno == ENOENT) { + return SlotStatus::Ok; + } + err = std::string("cannot read the slot manifest: ") + std::strerror(errno); + return SlotStatus::Internal; + } + if (!S_ISREG(st.st_mode)) { + err = "the slot manifest is not a regular file"; + return SlotStatus::Corrupt; + } + const int fd = ::open(path.c_str(), O_RDONLY | O_CLOEXEC | O_NOFOLLOW); + if (fd < 0) { + err = std::string("cannot read the slot manifest: ") + std::strerror(errno); + return SlotStatus::Internal; + } + std::string text; + char buf[4096]; + for (;;) { + const ssize_t r = ::read(fd, buf, sizeof(buf)); + if (r < 0) { + ::close(fd); + err = std::string("cannot read the slot manifest: ") + std::strerror(errno); + return SlotStatus::Internal; + } + if (r == 0) { + break; + } + text.append(buf, static_cast(r)); + if (text.size() > (16u << 20)) { + ::close(fd); + err = "the slot manifest is corrupt: too large"; + return SlotStatus::Corrupt; + } + } + ::close(fd); + + try { + const nlohmann::json doc = nlohmann::json::parse(text); + if (!doc.is_object() || !doc.contains("entries") || !doc["entries"].is_array()) { + throw std::runtime_error("no entries array"); + } + for (const auto& j : doc["entries"]) { + SlotEntry e; + e.id = j.at("id").get(); + e.model_identity = j.at("model_identity").get(); + e.slot = j.at("slot").get(); + e.n_tokens = j.at("n_tokens").get(); + e.n_bytes = j.at("n_bytes").get(); + e.created_unix = j.at("created_unix").get(); + if (!slot_id_valid(e.id)) { + throw std::runtime_error("bad id"); + } + entries.push_back(std::move(e)); + } + } catch (const std::exception& ex) { + entries.clear(); + err = std::string("the slot manifest is corrupt: ") + ex.what(); + return SlotStatus::Corrupt; + } + return SlotStatus::Ok; +} + +// slot_manifest_write replaces the manifest atomically: write a temp file in +// the same directory, fsync it, rename it over the manifest, fsync the +// directory. +inline SlotStatus slot_manifest_write(const std::string& dir, const std::vector& entries, std::string& err) { + namespace fs = std::filesystem; + nlohmann::json doc; + doc["version"] = 1; + doc["entries"] = nlohmann::json::array(); + for (const auto& e : entries) { + doc["entries"].push_back(slot_entry_to_json(e)); + } + const std::string text = doc.dump(2); + + std::string suffix; + if (!slot_random_id(suffix)) { + err = "no secure random source available"; + return SlotStatus::Internal; + } + const std::string final_path = (fs::u8path(dir) / SLOT_CACHE_MANIFEST).string(); + const std::string tmp_path = final_path + ".tmp." + suffix; + const int fd = ::open(tmp_path.c_str(), O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC | O_NOFOLLOW, 0600); + if (fd < 0) { + err = std::string("cannot write the slot manifest: ") + std::strerror(errno); + return SlotStatus::Internal; + } + size_t done = 0; + bool ok = true; + while (done < text.size()) { + const ssize_t w = ::write(fd, text.data() + done, text.size() - done); + if (w <= 0) { + ok = false; + break; + } + done += static_cast(w); + } + if (ok && ::fsync(fd) != 0) { + ok = false; + } + ::close(fd); + if (!ok || ::rename(tmp_path.c_str(), final_path.c_str()) != 0) { + err = std::string("cannot write the slot manifest: ") + std::strerror(errno); + ::unlink(tmp_path.c_str()); + return SlotStatus::Internal; + } + const int dfd = ::open(dir.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC); + if (dfd >= 0) { + ::fsync(dfd); + ::close(dfd); + } + return SlotStatus::Ok; +} + +// slot_reserve_id picks an ID that is not in use and creates an empty file +// with that name (O_EXCL, mode 0600), so a save never overwrites anything. +// The generator can be replaced in tests. +inline SlotStatus slot_reserve_id(const std::string& dir, std::string& id, std::string& path, std::string& err, + const std::function& gen = slot_random_id) { + namespace fs = std::filesystem; + for (int attempt = 0; attempt < 16; attempt++) { + std::string cand; + if (!gen(cand) || !slot_id_valid(cand)) { + err = "no secure random source available"; + return SlotStatus::Internal; + } + const std::string p = (fs::u8path(dir) / cand).string(); + const int fd = ::open(p.c_str(), O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC | O_NOFOLLOW, 0600); + if (fd >= 0) { + ::close(fd); + id = cand; + path = p; + return SlotStatus::Ok; + } + if (errno != EEXIST) { + err = std::string("cannot create the slot file: ") + std::strerror(errno); + return SlotStatus::Internal; + } + } + err = "could not find an unused slot ID"; + return SlotStatus::Internal; +} + +// slot_commit_save appends the entry to the manifest and, when the list is +// over SLOT_CACHE_MAX_ENTRIES, removes the oldest entries and their files. The +// IDs removed are returned in pruned. A corrupt manifest stops the save. +inline SlotStatus slot_commit_save(const std::string& dir, const SlotEntry& entry, std::vector& pruned, + std::string& err, size_t max_entries = SLOT_CACHE_MAX_ENTRIES) { + namespace fs = std::filesystem; + pruned.clear(); + std::lock_guard lock(slot_cache_mutex()); + std::vector entries; + SlotStatus s = slot_manifest_load(dir, entries, err); + if (s != SlotStatus::Ok) { + return s; + } + entries.push_back(entry); + while (entries.size() > max_entries) { + size_t oldest = 0; + for (size_t i = 1; i < entries.size(); i++) { + if (entries[i].created_unix < entries[oldest].created_unix) { + oldest = i; + } + } + // Never prune the entry that was just added. + if (entries[oldest].id == entry.id) { + break; + } + pruned.push_back(entries[oldest].id); + entries.erase(entries.begin() + static_cast(oldest)); + } + s = slot_manifest_write(dir, entries, err); + if (s != SlotStatus::Ok) { + pruned.clear(); + return s; + } + for (const auto& id : pruned) { + std::error_code ec; + fs::remove(fs::u8path(dir) / id, ec); + } + return SlotStatus::Ok; +} + +// slot_lookup_restore decides whether a restore may go ahead and returns the +// path to hand to llama.cpp. The path is built from the manifest entry's ID, +// never from request text. +// +// The checks are: the ID has the right shape; the manifest lists it; the entry +// belongs to this model; the path is a regular file (lstat, so a link or a +// directory is refused); and its size is the size recorded at save time, so a +// replaced or truncated file is refused. +// +// This is a check of a path, and llama.cpp then opens that path itself. A +// small window remains between the lstat and that open. The 0700 mode of the +// directory is the main defence: it stops another local user from swapping the +// file in that window. This code does not close the window against the same +// user or against root. +inline SlotStatus slot_lookup_restore(const std::string& root, const std::string& identity, const std::string& id, + SlotEntry& entry, std::string& path, std::string& err) { + namespace fs = std::filesystem; + if (!slot_id_valid(id)) { + err = "invalid id: expected 32 lowercase hex characters"; + return SlotStatus::Invalid; + } + std::string dir; + SlotStatus s = slot_identity_dir(root, identity, false, dir, err); + if (s == SlotStatus::NotFound) { + err = "no saved state with this id for this model"; + return s; + } + if (s != SlotStatus::Ok) { + return s; + } + + std::vector entries; + { + std::lock_guard lock(slot_cache_mutex()); + s = slot_manifest_load(dir, entries, err); + } + if (s != SlotStatus::Ok) { + return s; + } + const SlotEntry* found = nullptr; + for (const auto& e : entries) { + if (e.id == id) { + found = &e; + break; + } + } + if (found == nullptr || found->model_identity != identity) { + err = "no saved state with this id for this model"; + return SlotStatus::NotFound; + } + + path = (fs::u8path(dir) / found->id).string(); + struct stat st; + if (::lstat(path.c_str(), &st) != 0) { + err = "no saved state with this id for this model"; + return SlotStatus::NotFound; + } + if (!S_ISREG(st.st_mode)) { + err = "the saved file is not a regular file"; + return SlotStatus::Invalid; + } + if (static_cast(st.st_size) != found->n_bytes) { + err = "the saved file does not match its record: it was changed or truncated"; + return SlotStatus::Invalid; + } + entry = *found; + return SlotStatus::Ok; +} + +#endif // !_WIN32 + +} // namespace llama_grpc diff --git a/backend/cpp/llama-cpp/slot_cache_test.cpp b/backend/cpp/llama-cpp/slot_cache_test.cpp new file mode 100644 index 000000000..98a9f7551 --- /dev/null +++ b/backend/cpp/llama-cpp/slot_cache_test.cpp @@ -0,0 +1,229 @@ +#include "slot_cache.h" + +#include +#include +#include +#include +#include +#include + +static int failures = 0; + +static void expect(bool cond, const char* what) { + if (!cond) { + std::fprintf(stderr, "FAIL: %s\n", what); + failures++; + } +} + +using namespace llama_grpc; +namespace fs = std::filesystem; + +static SlotEntry make_entry(const std::string& id, const std::string& identity, int64_t created, uint64_t bytes) { + SlotEntry e; + e.id = id; + e.model_identity = identity; + e.slot = 1; + e.n_tokens = 10; + e.n_bytes = bytes; + e.created_unix = created; + return e; +} + +static std::string write_file(const std::string& path, size_t n) { + std::ofstream(path, std::ios::binary) << std::string(n, 'x'); + return path; +} + +int main() { + // ID validation. + expect(slot_id_valid("0123456789abcdef0123456789abcdef"), "valid id"); + expect(!slot_id_valid("0123456789ABCDEF0123456789abcdef"), "uppercase rejected"); + expect(!slot_id_valid("0123456789abcdef"), "too short"); + expect(!slot_id_valid("0123456789abcdef0123456789abcdef0"), "too long"); + expect(!slot_id_valid(""), "empty"); + expect(!slot_id_valid("0123456789abcdef0123456789abcdeg"), "non-hex"); + expect(!slot_id_valid("..............................ab"), "dots"); + expect(!slot_id_valid("../../../../etc/passwd/../../xx"), "traversal"); + expect(!slot_id_valid("0123456789abcdef/123456789abcdef"), "slash"); + expect(!slot_id_valid(std::string("0123456789abcdef0123456789abcd\0f", 32)), "NUL byte"); + + // SHA-256 known answers. + expect(slot_sha256_hex("abc") == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", "sha256 abc"); + expect(slot_sha256_hex("") == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "sha256 empty"); + expect(slot_sha256_hex("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq") == + "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1", "sha256 two blocks"); + expect(slot_identity_dir_name("abc") == "ba7816bf8f01cfea414140de5dae2223", "dir name is 32 hex of sha256"); + + // ID generation: format and uniqueness. + { + std::set seen; + bool fmt = true; + for (int i = 0; i < 20000; i++) { + std::string id; + if (!slot_random_id(id) || !slot_id_valid(id)) { + fmt = false; + break; + } + seen.insert(id); + } + expect(fmt, "generated ids have the valid format"); + expect(seen.size() == 20000, "20000 generated ids are unique"); + } + + std::error_code ec; + const fs::path base = fs::temp_directory_path() / "localai_slot_cache_test"; + fs::remove_all(base, ec); + const std::string root = (base / "save").string(); + std::string warn, err; + + // Root creation: missing root gets mode 0700. + expect(slot_ensure_private_dir(root, warn, err), "create root"); + expect((fs::status(root).permissions() & (fs::perms::group_all | fs::perms::others_all)) == fs::perms::none, + "new root is 0700"); + expect(warn.empty(), "no warning for 0700 root"); + // A looser existing root is kept and reported. + const std::string loose = (base / "loose").string(); + fs::create_directories(loose); + ::chmod(loose.c_str(), 0755); + expect(slot_ensure_private_dir(loose, warn, err), "existing root accepted"); + expect(!warn.empty(), "warning for loose root"); + expect((fs::status(loose).permissions() & fs::perms::group_read) != fs::perms::none, "loose root mode unchanged"); + + const std::string identity = "model-a"; + const std::string other_identity = "model-b"; + std::string dir; + expect(slot_identity_dir(root, identity, false, dir, err) == SlotStatus::NotFound, "no dir before first save"); + expect(slot_identity_dir(root, identity, true, dir, err) == SlotStatus::Ok, "identity dir created"); + expect(fs::path(dir).filename().string() == slot_identity_dir_name(identity), "dir named by identity hash"); + expect((fs::status(dir).permissions() & (fs::perms::group_all | fs::perms::others_all)) == fs::perms::none, + "identity dir is 0700"); + + // Reserve never overwrites: the generator repeats a used ID first. + { + std::string id1, path1, id2, path2; + expect(slot_reserve_id(dir, id1, path1, err) == SlotStatus::Ok, "reserve first id"); + write_file(path1, 5); + int calls = 0; + auto gen = [&](std::string& out) { + out = (calls++ < 2) ? id1 : "ffffffffffffffffffffffffffffffff"; + return true; + }; + expect(slot_reserve_id(dir, id2, path2, err, gen) == SlotStatus::Ok, "reserve after collision"); + expect(id2 == "ffffffffffffffffffffffffffffffff" && calls == 3, "collision retried with a new id"); + expect(fs::file_size(path1) == 5, "existing file untouched by collision"); + fs::remove(path1); + fs::remove(path2); + } + + // Missing manifest is empty; round trip. + std::vector entries; + expect(slot_manifest_load(dir, entries, err) == SlotStatus::Ok && entries.empty(), "missing manifest is empty"); + const std::string id_a = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + const std::string id_b = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + std::vector pruned; + write_file(dir + "/" + id_a, 100); + expect(slot_commit_save(dir, make_entry(id_a, identity, 1000, 100), pruned, err) == SlotStatus::Ok, "commit a"); + write_file(dir + "/" + id_b, 200); + expect(slot_commit_save(dir, make_entry(id_b, identity, 2000, 200), pruned, err) == SlotStatus::Ok, "commit b"); + expect(slot_manifest_load(dir, entries, err) == SlotStatus::Ok, "load manifest"); + expect(entries.size() == 2 && entries[0].id == id_a && entries[1].id == id_b, "manifest round trip order"); + expect(entries[1].model_identity == identity && entries[1].slot == 1 && entries[1].n_tokens == 10 && + entries[1].n_bytes == 200 && entries[1].created_unix == 2000, "manifest round trip fields"); + { + bool tmp_left = false; + for (const auto& de : fs::directory_iterator(dir)) { + if (de.path().filename().string().find(".tmp.") != std::string::npos) tmp_left = true; + } + expect(!tmp_left, "no temp manifest left behind"); + } + + // Pruning: oldest by created time goes first, with its file. + { + const std::string d2 = (base / "prune").string(); + fs::create_directories(d2); + std::vector ids; + for (int i = 0; i < 4; i++) { + std::string id = std::string(31, 'c') + char('0' + i); + ids.push_back(id); + write_file(d2 + "/" + id, 1); + } + // Commit out of time order so "oldest" is not "first in the list". + const int64_t times[4] = {300, 100, 400, 200}; + std::vector p; + for (int i = 0; i < 3; i++) { + expect(slot_commit_save(d2, make_entry(ids[i], identity, times[i], 1), p, err, 3) == SlotStatus::Ok, "fill"); + expect(p.empty(), "no prune under the limit"); + } + expect(slot_commit_save(d2, make_entry(ids[3], identity, times[3], 1), p, err, 3) == SlotStatus::Ok, "over limit"); + expect(p.size() == 1 && p[0] == ids[1], "oldest entry pruned"); + expect(!fs::exists(d2 + "/" + ids[1]), "pruned file deleted"); + expect(fs::exists(d2 + "/" + ids[0]), "newer file kept"); + slot_manifest_load(d2, entries, err); + expect(entries.size() == 3, "manifest holds the limit"); + expect(SLOT_CACHE_MAX_ENTRIES == 32, "limit is 32"); + } + + // Restore: success. + SlotEntry got; + std::string path; + expect(slot_lookup_restore(root, identity, id_b, got, path, err) == SlotStatus::Ok, "restore known id"); + expect(path == dir + "/" + id_b && got.n_bytes == 200, "path built from the entry id"); + // Restore: refusals. + expect(slot_lookup_restore(root, identity, "dddddddddddddddddddddddddddddddd", got, path, err) == SlotStatus::NotFound, + "unknown id refused"); + expect(slot_lookup_restore(root, identity, "../" + id_b, got, path, err) == SlotStatus::Invalid, "bad id refused"); + expect(slot_lookup_restore(root, other_identity, id_b, got, path, err) == SlotStatus::NotFound, + "other model has no such id"); + + // Identity mismatch inside a manifest (edited by hand) is refused. + { + std::vector forged = {make_entry(id_a, "someone-else", 1000, 100)}; + slot_manifest_write(dir, forged, err); + expect(slot_lookup_restore(root, identity, id_a, got, path, err) == SlotStatus::NotFound, + "identity mismatch refused"); + std::vector restored = {make_entry(id_a, identity, 1000, 100), make_entry(id_b, identity, 2000, 200)}; + slot_manifest_write(dir, restored, err); + } + + // Size mismatch is refused. + write_file(dir + "/" + id_a, 99); + expect(slot_lookup_restore(root, identity, id_a, got, path, err) == SlotStatus::Invalid, "size mismatch refused"); + write_file(dir + "/" + id_a, 100); + expect(slot_lookup_restore(root, identity, id_a, got, path, err) == SlotStatus::Ok, "size match accepted"); + + // A symlink with the right target size is refused. + { + write_file((base / "target.bin").string(), 200); + fs::remove(dir + "/" + id_b); + fs::create_symlink(base / "target.bin", dir + "/" + id_b, ec); + if (!ec) { + expect(slot_lookup_restore(root, identity, id_b, got, path, err) == SlotStatus::Invalid, "symlink refused"); + } else { + expect(false, "could not create the test symlink"); + } + fs::remove(dir + "/" + id_b); + fs::create_directory(dir + "/" + id_b); + expect(slot_lookup_restore(root, identity, id_b, got, path, err) == SlotStatus::Invalid, "directory refused"); + fs::remove(dir + "/" + id_b); + } + + // A corrupt manifest is an error for save and restore, and is not reset. + { + std::ofstream(dir + "/manifest.json") << "{not json"; + expect(slot_lookup_restore(root, identity, id_a, got, path, err) == SlotStatus::Corrupt, "corrupt manifest: restore"); + expect(slot_commit_save(dir, make_entry(id_b, identity, 3000, 1), pruned, err) == SlotStatus::Corrupt, + "corrupt manifest: save"); + std::ifstream in(dir + "/manifest.json"); + std::string text((std::istreambuf_iterator(in)), std::istreambuf_iterator()); + expect(text == "{not json", "corrupt manifest left as it was"); + std::ofstream(dir + "/manifest.json") << R"({"entries":[{"id":"../x","model_identity":"m","slot":0,"n_tokens":0,"n_bytes":0,"created_unix":0}]})"; + expect(slot_manifest_load(dir, entries, err) == SlotStatus::Corrupt, "entry with a bad id is corrupt"); + } + + fs::remove_all(base, ec); + if (failures == 0) { + std::puts("slot_cache_test: ok"); + } + return failures == 0 ? 0 : 1; +} diff --git a/core/backend/slot_cache.go b/core/backend/slot_cache.go new file mode 100644 index 000000000..5d7628981 --- /dev/null +++ b/core/backend/slot_cache.go @@ -0,0 +1,119 @@ +package backend + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/trace" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// ErrSlotCacheUnsupported is returned when the model's backend does not +// implement the SlotCache rpc (gRPC UNIMPLEMENTED). +var ErrSlotCacheUnsupported = errors.New("slot save and restore is not supported by this model's backend") + +// slotCacheCaller is the part of a loaded backend that slot cache needs. +type slotCacheCaller interface { + SlotCache(ctx context.Context, in *pb.SlotCacheRequest, opts ...grpc.CallOption) (*pb.SlotCacheResponse, error) +} + +// callSlotCache sends the request and maps an UNIMPLEMENTED status to +// ErrSlotCacheUnsupported. Every other error is returned unchanged. +func callSlotCache(ctx context.Context, c slotCacheCaller, req *pb.SlotCacheRequest) (*pb.SlotCacheResponse, error) { + resp, err := c.SlotCache(ctx, req) + if err != nil { + if st, ok := status.FromError(err); ok && st.Code() == codes.Unimplemented { + return nil, ErrSlotCacheUnsupported + } + return nil, err + } + return resp, nil +} + +// newSlotCacheRequest always sends the configured model as ModelIdentity, so a +// backend that checks it can refuse a request meant for another model. +func newSlotCacheRequest(action pb.SlotCacheAction, slot int32, id string, modelConfig config.ModelConfig) *pb.SlotCacheRequest { + return &pb.SlotCacheRequest{ + ModelIdentity: modelConfig.Model, + Action: action, + Slot: slot, + Id: id, + } +} + +// ModelSlotCache saves or restores the prompt cache of one slot of a loaded +// model. On save id is empty and the backend issues one. On restore the caller +// must have validated id (see pkg/slotcache); the backend checks it again and +// accepts only an ID it issued for this model. +func ModelSlotCache(action pb.SlotCacheAction, slot int32, id string, loader *model.ModelLoader, modelConfig config.ModelConfig, appConfig *config.ApplicationConfig) (schema.SlotCacheResponse, error) { + name := "restore" + if action == pb.SlotCacheAction_SLOT_CACHE_ACTION_SAVE { + name = "save" + } + + inferenceModel, err := loader.Load(ModelOptions(modelConfig, appConfig)...) + if err != nil { + recordModelLoadFailure(appConfig, modelConfig.Name, modelConfig.Backend, err, nil) + return schema.SlotCacheResponse{}, err + } + + release, err := AcquireGlobalBackendSlot() + if err != nil { + return schema.SlotCacheResponse{}, err + } + defer release() + + var startTime time.Time + var traceID string + if appConfig.EnableTracing { + trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) + startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTokenize, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "slot cache " + name}) + } + defer trace.CancelBackendTrace(traceID) + + resp, err := callSlotCache(appConfig.Context, inferenceModel, newSlotCacheRequest(action, slot, id, modelConfig)) + + if appConfig.EnableTracing { + errStr := "" + if err != nil { + errStr = err.Error() + } + trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, + Timestamp: startTime, + Duration: time.Since(startTime), + Type: trace.BackendTraceTokenize, + ModelName: modelConfig.Name, + Backend: modelConfig.Backend, + Summary: fmt.Sprintf("slot cache %s", name), + Error: errStr, + Data: map[string]any{"action": name, "slot": slot, "id": id}, + }) + } + if err != nil { + return schema.SlotCacheResponse{}, err + } + if !resp.GetSuccess() { + return schema.SlotCacheResponse{}, fmt.Errorf("slot cache %s failed: %s", name, resp.GetMessage()) + } + + return schema.SlotCacheResponse{ + Model: modelConfig.Name, + Action: name, + Slot: resp.GetSlot(), + ID: resp.GetId(), + Message: resp.GetMessage(), + NTokens: resp.GetNTokens(), + NBytes: resp.GetNBytes(), + DurationMS: resp.GetDurationMs(), + }, nil +} diff --git a/core/backend/slot_cache_test.go b/core/backend/slot_cache_test.go new file mode 100644 index 000000000..a036e2281 --- /dev/null +++ b/core/backend/slot_cache_test.go @@ -0,0 +1,75 @@ +package backend + +import ( + "context" + "errors" + + "github.com/mudler/LocalAI/core/config" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type fakeSlotCache struct { + got *pb.SlotCacheRequest + resp *pb.SlotCacheResponse + err error +} + +func (f *fakeSlotCache) SlotCache(_ context.Context, in *pb.SlotCacheRequest, _ ...grpc.CallOption) (*pb.SlotCacheResponse, error) { + f.got = in + return f.resp, f.err +} + +var _ = Describe("slot cache backend call", func() { + It("maps UNIMPLEMENTED to ErrSlotCacheUnsupported", func() { + f := &fakeSlotCache{err: status.Error(codes.Unimplemented, "method SlotCache not implemented")} + _, err := callSlotCache(context.Background(), f, &pb.SlotCacheRequest{}) + Expect(errors.Is(err, ErrSlotCacheUnsupported)).To(BeTrue()) + }) + + It("leaves other gRPC errors unchanged", func() { + orig := status.Error(codes.FailedPrecondition, "slot save/restore is disabled") + _, err := callSlotCache(context.Background(), &fakeSlotCache{err: orig}, &pb.SlotCacheRequest{}) + Expect(err).To(Equal(orig)) + Expect(errors.Is(err, ErrSlotCacheUnsupported)).To(BeFalse()) + }) + + It("leaves non-gRPC errors unchanged", func() { + orig := errors.New("dial failed") + _, err := callSlotCache(context.Background(), &fakeSlotCache{err: orig}, &pb.SlotCacheRequest{}) + Expect(err).To(Equal(orig)) + }) + + It("returns the response on success", func() { + want := &pb.SlotCacheResponse{Success: true, NTokens: 7} + got, err := callSlotCache(context.Background(), &fakeSlotCache{resp: want}, &pb.SlotCacheRequest{}) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal(want)) + }) + + It("returns the issued ID and the message", func() { + f := &fakeSlotCache{resp: &pb.SlotCacheResponse{Success: true, Id: "0123456789abcdef0123456789abcdef", Message: "pruned x", Slot: 1, NTokens: 3, NBytes: 9}} + cfg := config.ModelConfig{Name: "m"} + resp, err := callSlotCache(context.Background(), f, newSlotCacheRequest(pb.SlotCacheAction_SLOT_CACHE_ACTION_SAVE, 1, "", cfg)) + Expect(err).ToNot(HaveOccurred()) + Expect(f.got.GetId()).To(BeEmpty()) + Expect(resp.GetId()).To(Equal("0123456789abcdef0123456789abcdef")) + }) + + It("sends the configured model as the identity", func() { + cfg := config.ModelConfig{Name: "alias"} + cfg.Model = "weights.gguf" + req := newSlotCacheRequest(pb.SlotCacheAction_SLOT_CACHE_ACTION_RESTORE, 2, "0123456789abcdef0123456789abcdef", cfg) + f := &fakeSlotCache{resp: &pb.SlotCacheResponse{Success: true}} + _, err := callSlotCache(context.Background(), f, req) + Expect(err).ToNot(HaveOccurred()) + Expect(f.got.GetModelIdentity()).To(Equal("weights.gguf")) + Expect(f.got.GetAction()).To(Equal(pb.SlotCacheAction_SLOT_CACHE_ACTION_RESTORE)) + Expect(f.got.GetSlot()).To(Equal(int32(2))) + Expect(f.got.GetId()).To(Equal("0123456789abcdef0123456789abcdef")) + }) +}) diff --git a/core/http/endpoints/localai/export_slot_cache_test.go b/core/http/endpoints/localai/export_slot_cache_test.go new file mode 100644 index 000000000..26d925497 --- /dev/null +++ b/core/http/endpoints/localai/export_slot_cache_test.go @@ -0,0 +1,4 @@ +package localai + +// SlotCacheErrorForTest exposes slotCacheError to the external test package. +var SlotCacheErrorForTest = slotCacheError diff --git a/core/http/endpoints/localai/slot_cache.go b/core/http/endpoints/localai/slot_cache.go new file mode 100644 index 000000000..485efb2e4 --- /dev/null +++ b/core/http/endpoints/localai/slot_cache.go @@ -0,0 +1,97 @@ +package localai + +import ( + "errors" + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/slotcache" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// SlotCacheSaveEndpoint saves the prompt cache of a model slot. +// @Summary Save a slot's prompt cache. +// @Description Writes the prompt cache (KV state) of one slot of a llama-cpp model under the directory set by the model's slot_save_path option, and returns an ID for it. The server issues the ID; the client cannot choose a name or a path. The server keeps at most 32 saves per model and deletes the oldest one beyond that. Admin only. +// @Tags inference +// @Param request body schema.SlotCacheRequest true "Request" +// @Success 200 {object} schema.SlotCacheResponse "Response" +// @Router /v1/slots/save [post] +func SlotCacheSaveEndpoint(ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return slotCacheEndpoint(pb.SlotCacheAction_SLOT_CACHE_ACTION_SAVE, ml, appConfig) +} + +// SlotCacheRestoreEndpoint loads a saved prompt cache back into a model slot. +// @Summary Restore a slot's prompt cache by ID. +// @Description Loads the prompt cache that a save call returned an ID for into one slot of a llama-cpp model. The ID must have been issued for this model. The state is only as fresh as the save. Admin only. +// @Tags inference +// @Param request body schema.SlotCacheRequest true "Request" +// @Success 200 {object} schema.SlotCacheResponse "Response" +// @Router /v1/slots/restore [post] +func SlotCacheRestoreEndpoint(ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return slotCacheEndpoint(pb.SlotCacheAction_SLOT_CACHE_ACTION_RESTORE, ml, appConfig) +} + +func slotCacheEndpoint(action pb.SlotCacheAction, ml *model.ModelLoader, appConfig *config.ApplicationConfig) echo.HandlerFunc { + return func(c echo.Context) error { + input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.SlotCacheRequest) + if !ok || input.Model == "" { + return echo.NewHTTPError(http.StatusBadRequest, "model is required") + } + cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + if !ok || cfg == nil { + return echo.NewHTTPError(http.StatusBadRequest, "model not found") + } + + if input.Slot < 0 { + return echo.NewHTTPError(http.StatusBadRequest, "slot must not be negative") + } + // Refuse bad input before any backend is loaded or called. + if action == pb.SlotCacheAction_SLOT_CACHE_ACTION_SAVE { + if input.ID != "" { + return echo.NewHTTPError(http.StatusBadRequest, "id must not be set on save: the server issues it") + } + } else if err := slotcache.ValidateID(input.ID); err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + if slotcache.DirFromOptions(cfg.Options) == "" { + return echo.NewHTTPError(http.StatusBadRequest, "slot save and restore is disabled for this model: set the slot_save_path option") + } + + resp, err := backend.ModelSlotCache(action, input.Slot, input.ID, ml, *cfg, appConfig) + if err != nil { + return slotCacheError(err) + } + return c.JSON(http.StatusOK, resp) + } +} + +// errSlotNotFound is the fixed text for an ID the backend does not know. +const errSlotNotFound = "no saved state with this id for this model" + +// slotCacheError turns a backend error into an HTTP error. A backend without +// slot cache support gets 501. NOT_FOUND gets 404 with a fixed text, so the +// response says nothing about the disk; a model identity mismatch is not an +// unknown ID and keeps the existing mapping. Everything else follows +// mapBackendError. +func slotCacheError(err error) error { + if errors.Is(err, backend.ErrSlotCacheUnsupported) { + return echo.NewHTTPError(http.StatusNotImplemented, backend.ErrSlotCacheUnsupported.Error()) + } + if st, ok := status.FromError(err); ok { + switch { + case st.Code() == codes.NotFound && !grpcerrors.IsModelMismatch(err): + return echo.NewHTTPError(http.StatusNotFound, errSlotNotFound) + case st.Code() == codes.ResourceExhausted: + return echo.NewHTTPError(http.StatusInsufficientStorage, st.Message()) + } + } + return mapBackendError(err) +} diff --git a/core/http/endpoints/localai/slot_cache_test.go b/core/http/endpoints/localai/slot_cache_test.go new file mode 100644 index 000000000..bc65f59d0 --- /dev/null +++ b/core/http/endpoints/localai/slot_cache_test.go @@ -0,0 +1,124 @@ +package localai_test + +import ( + "net/http" + "net/http/httptest" + "strings" + + "errors" + "fmt" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/http/endpoints/localai" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// These cases all fail before a backend is loaded, so no model loader is +// needed. They cover the checks that keep a request away from the disk. +var _ = Describe("slot cache endpoints", func() { + call := func(h func() echo.HandlerFunc, options []string, body schema.SlotCacheRequest) int { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/v1/slots/save", strings.NewReader("{}")) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + body.Model = "m" + c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, &body) + c.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{Name: "m", Options: options}) + err := h()(c) + if err == nil { + return rec.Code + } + if he, ok := err.(*echo.HTTPError); ok { + return he.Code + } + return http.StatusInternalServerError + } + + It("answers 501 when the backend does not support slot cache", func() { + err := SlotCacheErrorForTest(fmt.Errorf("wrapped: %w", backend.ErrSlotCacheUnsupported)) + var he *echo.HTTPError + Expect(errors.As(err, &he)).To(BeTrue()) + Expect(he.Code).To(Equal(http.StatusNotImplemented)) + Expect(he.Message).To(Equal("slot save and restore is not supported by this model's backend")) + }) + + It("keeps other backend errors on the existing mapping", func() { + plain := errors.New("boom") + Expect(SlotCacheErrorForTest(plain)).To(Equal(plain)) + }) + + It("answers 404 with a fixed text for an unknown ID", func() { + err := SlotCacheErrorForTest(status.Error(codes.NotFound, "/srv/slots/abc: secret path detail")) + var he *echo.HTTPError + Expect(errors.As(err, &he)).To(BeTrue()) + Expect(he.Code).To(Equal(http.StatusNotFound)) + Expect(he.Message).To(Equal("no saved state with this id for this model")) + }) + + It("keeps the model identity mismatch off the 404 text", func() { + err := SlotCacheErrorForTest(status.Error(codes.NotFound, "llama-cpp: model identity mismatch: loaded \"a\", requested \"b\"")) + var he *echo.HTTPError + Expect(errors.As(err, &he)).To(BeTrue()) + Expect(he.Message).ToNot(Equal("no saved state with this id for this model")) + }) + + It("answers 412 when the saved file no longer matches its record", func() { + err := SlotCacheErrorForTest(status.Error(codes.FailedPrecondition, "the saved file does not match its record")) + var he *echo.HTTPError + Expect(errors.As(err, &he)).To(BeTrue()) + Expect(he.Code).To(Equal(http.StatusPreconditionFailed)) + }) + + handlers := map[string]func() echo.HandlerFunc{ + "save": func() echo.HandlerFunc { return SlotCacheSaveEndpoint(nil, nil) }, + "restore": func() echo.HandlerFunc { return SlotCacheRestoreEndpoint(nil, nil) }, + } + const goodID = "0123456789abcdef0123456789abcdef" + + for name, h := range handlers { + Describe(name, func() { + id := goodID + if name == "save" { + id = "" + } + + It("is disabled when slot_save_path is not set", func() { + Expect(call(h, []string{"cache_ram:1024"}, schema.SlotCacheRequest{ID: id})).To(Equal(http.StatusBadRequest)) + }) + + It("rejects a negative slot", func() { + Expect(call(h, []string{"slot_save_path:/srv/slots"}, schema.SlotCacheRequest{ID: id, Slot: -1})).To(Equal(http.StatusBadRequest)) + }) + + It("requires a model", func() { + e := echo.New() + c := e.NewContext(httptest.NewRequest(http.MethodPost, "/", nil), httptest.NewRecorder()) + err := h()(c) + Expect(err).To(HaveOccurred()) + }) + }) + } + + Describe("restore", func() { + h := handlers["restore"] + It("rejects an ID that is not 32 lowercase hex characters", func() { + opts := []string{"slot_save_path:/srv/slots"} + for _, id := range []string{"", ".", "..", "../a", "a/b", "prompt.bin", "0123456789ABCDEF0123456789abcdef", "0123456789abcdef", "a\x00b"} { + Expect(call(h, opts, schema.SlotCacheRequest{ID: id})).To(Equal(http.StatusBadRequest), id) + } + }) + }) + + Describe("save", func() { + It("rejects a client-chosen ID", func() { + Expect(call(handlers["save"], []string{"slot_save_path:/srv/slots"}, schema.SlotCacheRequest{ID: goodID})).To(Equal(http.StatusBadRequest)) + }) + }) +}) diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index e0889004e..ddf0cb3bd 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -417,6 +417,8 @@ func RegisterLocalAIRoutes(router *echo.Echo, "detection": "/v1/detection", "tokenize": "/v1/tokenize", "detokenize": "/v1/detokenize", + "slots_save": "/v1/slots/save", + "slots_restore": "/v1/slots/restore", }, "monitoring": monitoringRoutes, "mcp": map[string]string{ @@ -494,6 +496,19 @@ func RegisterLocalAIRoutes(router *echo.Echo, requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_TOKENIZE)), requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.DetokenizeRequest) })) + // Slot prompt-cache save and restore (llama-cpp). Admin only: it writes + // files on the host that runs the backend. + router.POST("/v1/slots/save", + localai.SlotCacheSaveEndpoint(ml, appConfig), + adminMiddleware, + requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), + requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.SlotCacheRequest) })) + router.POST("/v1/slots/restore", + localai.SlotCacheRestoreEndpoint(ml, appConfig), + adminMiddleware, + requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)), + requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.SlotCacheRequest) })) + // MCP endpoint - supports both streaming and non-streaming modes // Note: streaming mode is NOT compatible with the OpenAI apis. We have a set which streams more states. if evaluator != nil && !appConfig.DisableMCP { diff --git a/core/schema/tokenize.go b/core/schema/tokenize.go index f9a2fb13b..b68a0f0c4 100644 --- a/core/schema/tokenize.go +++ b/core/schema/tokenize.go @@ -17,3 +17,23 @@ type DetokenizeRequest struct { type DetokenizeResponse struct { Content string `json:"content"` // detokenized text } + +// SlotCacheRequest names a model slot. Restore also sends the ID that a +// save returned. The client never chooses a file name. Slot defaults to 0. +type SlotCacheRequest struct { + BasicModelRequest + ID string `json:"id,omitempty"` // restore only: an ID issued by a save for this model + Slot int32 `json:"slot"` // slot id, 0 is the first slot +} + +// SlotCacheResponse reports what was saved or restored. +type SlotCacheResponse struct { + Model string `json:"model"` + Action string `json:"action"` // "save" or "restore" + Slot int32 `json:"slot"` + ID string `json:"id"` // issued by the save, sent back to restore + Message string `json:"message,omitempty"` // for example the IDs a save pruned + NTokens uint64 `json:"n_tokens"` // tokens saved or restored + NBytes uint64 `json:"n_bytes"` // bytes written or read + DurationMS float64 `json:"duration_ms"` // time the backend spent +} diff --git a/core/services/nodes/health_mock_test.go b/core/services/nodes/health_mock_test.go index 95ea866d7..330871570 100644 --- a/core/services/nodes/health_mock_test.go +++ b/core/services/nodes/health_mock_test.go @@ -233,6 +233,9 @@ func (c *fakeBackendClient) Rerank(_ context.Context, _ *pb.RerankRequest, _ ... func (c *fakeBackendClient) GetTokenMetrics(_ context.Context, _ *pb.MetricsRequest, _ ...ggrpc.CallOption) (*pb.MetricsResponse, error) { return nil, nil } +func (c *fakeBackendClient) SlotCache(_ context.Context, _ *pb.SlotCacheRequest, _ ...ggrpc.CallOption) (*pb.SlotCacheResponse, error) { + return nil, nil +} func (c *fakeBackendClient) VAD(_ context.Context, _ *pb.VADRequest, _ ...ggrpc.CallOption) (*pb.VADResponse, error) { return nil, nil } diff --git a/core/services/nodes/inflight_test.go b/core/services/nodes/inflight_test.go index 1c2f6df20..0952c6831 100644 --- a/core/services/nodes/inflight_test.go +++ b/core/services/nodes/inflight_test.go @@ -182,6 +182,9 @@ func (f *fakeGRPCBackend) Rerank(_ context.Context, _ *pb.RerankRequest, _ ...gg return &pb.RerankResult{}, nil } +func (f *fakeGRPCBackend) SlotCache(_ context.Context, _ *pb.SlotCacheRequest, _ ...ggrpc.CallOption) (*pb.SlotCacheResponse, error) { + return &pb.SlotCacheResponse{}, nil +} func (f *fakeGRPCBackend) GetTokenMetrics(_ context.Context, _ *pb.MetricsRequest, _ ...ggrpc.CallOption) (*pb.MetricsResponse, error) { return &pb.MetricsResponse{}, nil } diff --git a/docs/content/features/authentication.md b/docs/content/features/authentication.md index 13660dd27..40e528a09 100644 --- a/docs/content/features/authentication.md +++ b/docs/content/features/authentication.md @@ -261,6 +261,7 @@ When authentication is enabled, the following endpoints require admin role: - `GET /api/backend-logs/*`, `POST /api/backend-logs/*/clear` - `GET /api/resources`, `GET /api/settings`, `POST /api/settings` - `GET /system`, `GET /backend/monitor`, `POST /backend/shutdown`, `POST /backend/load` +- `POST /v1/slots/save`, `POST /v1/slots/restore` (write and read prompt-cache files on the host) **P2P:** - `GET /api/p2p/*` diff --git a/docs/content/features/text-generation.md b/docs/content/features/text-generation.md index 26f25170b..60972b222 100644 --- a/docs/content/features/text-generation.md +++ b/docs/content/features/text-generation.md @@ -592,6 +592,7 @@ The `llama.cpp` backend supports additional configuration options that can be sp | `use_jinja` or `jinja` | boolean | Enable Jinja2 template processing for chat templates. When enabled, the backend uses Jinja2-based chat templates from the model for formatting messages. | `use_jinja:true` | | `context_shift` | boolean | Enable context shifting, which allows the model to dynamically adjust context window usage. | `context_shift:true` | | `cache_ram` | integer | Size budget in MiB for the **server-side prompt cache** (a host-RAM store of idle slot KV states that's reloaded on a prompt-prefix hit, see [upstream PR #16391](https://github.com/ggml-org/llama.cpp/pull/16391)). Default: `8192` MiB (llama.cpp default). `-1` removes the limit. `0` disables the prompt cache entirely. Together with `kv_unified` and `cache_idle_slots` this is what makes a repeated system prompt skip prefill on subsequent calls. | `cache_ram:4096` | +| `slot_save_path` | string | Directory where the slot save and restore API writes and reads prompt-cache files. Save and restore are off unless this is set. LocalAI creates the directory with mode `0700` if it is missing. See [Saving and restoring a prompt cache](#saving-and-restoring-a-prompt-cache). | `slot_save_path:/var/lib/localai/slots` | | `parallel` or `n_parallel` | integer | Enable parallel request processing. When set to a value greater than 1, enables continuous batching for handling multiple requests concurrently. | `parallel:4` | | `grpc_servers` or `rpc_servers` | string | Comma-separated list of gRPC server addresses for distributed inference. Allows distributing workload across multiple llama.cpp workers. | `grpc_servers:localhost:50051,localhost:50052` | | `fit_params` or `fit` | boolean | Enable auto-adjustment of model/context parameters to fit available device memory. Default: `true`. | `fit_params:true` | @@ -667,6 +668,50 @@ Set `cache_ram:0` to opt out of the prompt cache entirely (saves host RAM at the `cache_ram:-1` removes the limit. With idle-slot saving on, every distinct prompt then leaves its slot state in host RAM, so a workload with many different prompts (classification, ingestion) grows the backend by roughly the KV size of each prompt until the host runs out of memory. +##### Saving and restoring a prompt cache + +The server-side prompt cache lives in RAM and is lost when the model unloads. To keep the prompt state of one slot on disk, set `slot_save_path` and use the slot API. A typical use is a voice assistant with a long system prompt: save the state once, then restore it when the wake word fires, so the first request does not wait for prefill. + +```yaml +options: + - slot_save_path:/var/lib/localai/slots +``` + +Both calls need an admin API key when authentication is on. Without `slot_save_path` they return `400`. + +The server chooses the name of every saved file. A save returns an `id`, and a restore takes that `id`. A client cannot pick a file name or a path. + +```bash +# Save slot 0 of the model. The reply contains the id. +curl http://localhost:8080/v1/slots/save \ + -H "Content-Type: application/json" \ + -d '{"model": "my-model", "slot": 0}' + +# Load that state back into slot 0 +curl http://localhost:8080/v1/slots/restore \ + -H "Content-Type: application/json" \ + -d '{"model": "my-model", "slot": 0, "id": "0123456789abcdef0123456789abcdef"}' +``` + +| Field | Description | +|---|---| +| `model` | Name of a `llama-cpp` model. | +| `slot` | Slot id. Default `0`. With `parallel:N` the ids are `0` to `N-1`. | +| `id` | Restore only. The `id` that a save returned for this model: 32 lowercase hex characters. Do not send it on save. | + +The save reply gives the `id`, the slot, the number of tokens and bytes, and the time taken in `duration_ms`. The `message` field lists the IDs of any older saves that were deleted to make room. The restore reply has the same fields. Neither reply contains a path on the host. + +Keep these points in mind: + +- A client cannot choose names. The backend creates a directory for each model under `slot_save_path` (mode `0700`) and writes a `manifest.json` in it. A restore works only for an `id` in the manifest of that model. Any other `id` returns `404` with `no saved state with this id for this model`. +- The backend keeps at most 32 saves for each model. When a save would go over that, it deletes the oldest one and lists its `id` in `message`. A save larger than 8 GiB is deleted and fails. +- A restore also checks that the file is a plain file and has the size that was recorded at save time. A file that was replaced or cut short is refused. +- Keep `slot_save_path` private to the user that runs LocalAI. LocalAI creates it with mode `0700` when it is missing, and writes a warning to the log when an existing directory is open to other users. The checks above cannot stop a local user who can write to the directory. +- A saved state is only as fresh as its save. To change the cached prompt, save again and use the new `id`. +- A saved state holds the prompt state for the model and quantisation it came from. Do not restore it into a different model, quantisation or context size. The restore call fails if the state does not fit, but a state from another model is not always detected. +- A slot that is busy with a request is served after that request ends. +- The API is for `llama-cpp` only. Other backends return `501`. + #### Reference - [llama](https://github.com/ggerganov/llama.cpp) diff --git a/docs/content/reference/api-errors.md b/docs/content/reference/api-errors.md index cb3b65ca2..abbe18bdf 100644 --- a/docs/content/reference/api-errors.md +++ b/docs/content/reference/api-errors.md @@ -372,6 +372,18 @@ See also: [Object Detection]({{%relref "features/object-detection" %}}) | 400 | Invalid or malformed request body | `Bad Request` | | 400 | Model not found in configuration | `Bad Request` | +### Slot cache - `POST /v1/slots/save`, `POST /v1/slots/restore` + +| Status | Cause | Example Message | +|--------|--------------------------------|---------------------------------------| +| 400 | Invalid `id` on restore, an `id` sent on save, negative slot, or model not found | `invalid id: expected 32 lowercase hex characters` | +| 404 | The `id` is unknown, or was issued for another model | `no saved state with this id for this model` | +| 400 | `slot_save_path` is not set for the model | `slot save and restore is disabled for this model: set the slot_save_path option` | +| 403 | Caller is not an admin (authentication on) | `Forbidden` | +| 412 | Backend has save and restore turned off, or the saved file no longer matches its record | `slot save/restore is disabled` | +| 507 | The saved state is larger than the 8 GiB limit and was deleted | `the saved state is larger than the 8589934592 byte limit and was deleted` | +| 501 | Backend does not support the call | `slot save and restore is not supported by this model's backend` | + ### Models - `GET /v1/models`, `GET /models` | Status | Cause | Example Message | diff --git a/go.mod b/go.mod index 127e9c973..30e94c430 100644 --- a/go.mod +++ b/go.mod @@ -109,10 +109,12 @@ require ( github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/charmbracelet/bubbles v0.21.0 // indirect github.com/charmbracelet/bubbletea v1.3.10 // indirect + github.com/chasefleming/elem-go v0.30.0 // indirect github.com/chromedp/cdproto v0.0.0-20260321001828-e3e3800016bc // indirect github.com/chromedp/chromedp v0.15.1 // indirect github.com/chromedp/sysutil v1.1.0 // indirect github.com/cyberphone/json-canonicalization v0.0.0-20241213102144-19d51d7fe467 // indirect + github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2 // indirect github.com/digitorus/pkcs7 v0.0.0-20230818184609-3a137a874352 // indirect github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7 // indirect github.com/dunglas/httpsfv v1.1.0 // indirect @@ -141,10 +143,14 @@ require ( github.com/gobwas/httphead v0.1.0 // indirect github.com/gobwas/pool v0.2.1 // indirect github.com/gobwas/ws v1.4.0 // indirect + github.com/gofiber/template v1.8.3 // indirect + github.com/gofiber/template/html/v2 v2.1.3 // indirect + github.com/gofiber/utils v1.1.0 // indirect github.com/google/certificate-transparency-go v1.3.2 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/in-toto/attestation v1.1.2 // indirect github.com/in-toto/in-toto-golang v0.9.0 // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect @@ -184,6 +190,8 @@ require ( github.com/sigstore/rekor-tiles/v2 v2.0.1 // indirect github.com/sigstore/sigstore v1.10.0 // indirect github.com/sigstore/timestamp-authority/v2 v2.0.3 // indirect + github.com/spf13/cobra v1.10.2 // indirect + github.com/spf13/pflag v1.0.10 // indirect github.com/standard-webhooks/standard-webhooks/libraries v0.0.0-20260508151727-1282bb917829 // indirect github.com/stretchr/testify v1.11.1 // indirect github.com/sv-tools/openapi v0.2.1 // indirect @@ -258,7 +266,7 @@ require ( github.com/kevinburke/ssh_config v1.2.0 // indirect github.com/labstack/gommon v0.4.2 // indirect github.com/mschoch/smat v0.2.0 // indirect - github.com/mudler/LocalAGI v0.0.0-20260927202351-7e0947d7ebca + github.com/mudler/LocalAGI v0.0.0-20261016084905-7aa0c8fea076 github.com/mudler/localrecall v0.6.7-0.20261012155221-88432ad54488 github.com/mudler/skillserver v0.0.7-0.20260520220837-a7317cbf9145 github.com/olekukonko/tablewriter v0.0.5 // indirect diff --git a/go.sum b/go.sum index 8599b18f5..a9c0d779d 100644 --- a/go.sum +++ b/go.sum @@ -277,6 +277,8 @@ github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Y github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ= github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg= +github.com/chasefleming/elem-go v0.30.0 h1:BlhV1ekv1RbFiM8XZUQeln1Ikb4D+bu2eDO4agREvok= +github.com/chasefleming/elem-go v0.30.0/go.mod h1:hz73qILBIKnTgOujnSMtEj20/epI+f6vg71RUilJAA4= github.com/chengxilo/virtualterm v1.0.4 h1:Z6IpERbRVlfB8WkOmtbHiDbBANU7cimRIof7mk9/PwM= github.com/chengxilo/virtualterm v1.0.4/go.mod h1:DyxxBZz/x1iqJjFxTFcr6/x+jSpqN0iwWCOK1q10rlY= github.com/chromedp/cdproto v0.0.0-20260321001828-e3e3800016bc h1:wkN/LMi5vc60pBRWx6qpbk/aEvq3/ZVNpnMvsw8PVVU= @@ -319,6 +321,7 @@ github.com/cpuguy83/dockercfg v0.3.2 h1:DlJTyZGBDlXqUZ2Dk2Q3xHs/FtnooJJVaad2S9GK github.com/cpuguy83/dockercfg v0.3.2/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHfjj5/jFyUJc= github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= github.com/cpuguy83/go-md2man/v2 v2.0.0/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/creachadair/mds v0.21.3 h1:RRgEAPIb52cU0q7UxGyN+13QlCVTZIL4slRr0cYYQfA= github.com/creachadair/mds v0.21.3/go.mod h1:1ltMWZd9yXhaHEoZwBialMaviWVUpRPvMwVP7saFAzM= github.com/creachadair/otp v0.5.0 h1:q3Th7CXm2zlmCdBjw5tEPFOj4oWJMnVL5HXlq0sNKS0= @@ -331,6 +334,8 @@ github.com/cyphar/filepath-securejoin v0.6.1 h1:5CeZ1jPXEiYt3+Z6zqprSAgSWiggmpVy github.com/cyphar/filepath-securejoin v0.6.1/go.mod h1:A8hd4EnAeyujCJRrICiOWqjS1AX0a9kM5XL+NwKoYSc= github.com/danieljoos/wincred v1.2.2 h1:774zMFJrqaeYCK2W57BgAem/MLi6mtSE47MB6BOJ0i0= github.com/danieljoos/wincred v1.2.2/go.mod h1:w7w4Utbrz8lqeMbDAK0lkNJUv5sAOkFi7nd/ogr0Uh8= +github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2 h1:flLYmnQFZNo04x2NPehMbf30m7Pli57xwZ0NFqR/hb0= +github.com/dave-gray101/v2keyauth v0.0.0-20240624150259-c45d584d25e2/go.mod h1:NtWqRzAp/1tw+twkW8uuBenEVVYndEAZACWU3F3xdoQ= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= @@ -551,6 +556,12 @@ github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk= github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/gofiber/fiber/v2 v2.52.13 h1:TOKP64iqC9b5P49VrBW5tHhUOvDyrtJ0xePEfzJbCbk= github.com/gofiber/fiber/v2 v2.52.13/go.mod h1:YEcBbO/FB+5M1IZNBP9FO3J9281zgPAreiI1oqg8nDw= +github.com/gofiber/template v1.8.3 h1:hzHdvMwMo/T2kouz2pPCA0zGiLCeMnoGsQZBTSYgZxc= +github.com/gofiber/template v1.8.3/go.mod h1:bs/2n0pSNPOkRa5VJ8zTIvedcI/lEYxzV3+YPXdBvq8= +github.com/gofiber/template/html/v2 v2.1.3 h1:n1LYBtmr9C0V/k/3qBblXyMxV5B0o/gpb6dFLp8ea+o= +github.com/gofiber/template/html/v2 v2.1.3/go.mod h1:U5Fxgc5KpyujU9OqKzy6Kn6Qup6Tm7zdsISR+VpnHRE= +github.com/gofiber/utils v1.1.0 h1:vdEBpn7AzIUJRhe+CiTOJdUcTg4Q9RK+pEa0KPbLdrM= +github.com/gofiber/utils v1.1.0/go.mod h1:poZpsnhBykfnY1Mc0KeEa6mSHrS3dV0+oBWyeQmb2e0= github.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw= github.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= @@ -1036,6 +1047,8 @@ github.com/msuozzo/bonsai/bonsai-yaml v0.4.0 h1:PJfyfjMQrcWU5e3ib0xGsVuOafNuimf5 github.com/msuozzo/bonsai/bonsai-yaml v0.4.0/go.mod h1:z8jc0tjXDSOQRKRG0G3uymvVxHdBbrVsWd149nPOyHE= github.com/mudler/LocalAGI v0.0.0-20260927202351-7e0947d7ebca h1:bHlzSuOc5cKvHGF21ZjFFUV7IlkS3wt9YkBsWtkcaC4= github.com/mudler/LocalAGI v0.0.0-20260927202351-7e0947d7ebca/go.mod h1:nk6zt1s5ANgchJYTWGY1jfFPuITSn1gB5oHZ/uFeFDg= +github.com/mudler/LocalAGI v0.0.0-20261016084905-7aa0c8fea076 h1:IHWoTN1T3LEfI08u1M2Mr4Zb7PWETC8xw/hnKUPc6lQ= +github.com/mudler/LocalAGI v0.0.0-20261016084905-7aa0c8fea076/go.mod h1:Zt29aPgdHzUgXnD3mmvkAkWSmVAy3dWEx30Pyp6yT9A= github.com/mudler/cogito v0.11.1-0.20260928072733-b40513ef5d1a h1:b3bZ15XGd3kH+gje6ggFAciT4uC83+24xO/qc3o4viQ= github.com/mudler/cogito v0.11.1-0.20260928072733-b40513ef5d1a/go.mod h1:UxGNMBRakV0A2uVHv2JQIHjpnVPkEWFaS7SDhyynZqQ= github.com/mudler/edgevpn v0.34.0 h1:qDrD/rCPFY/FdURbXudIZWihVKY4VOX3nMn3CcbeQEU= @@ -1243,6 +1256,7 @@ github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6 github.com/russross/blackfriday v1.6.0 h1:KqfZb0pUVN2lYqZUYRddxF4OR8ZMURnJIG5Y3VRLtww= github.com/russross/blackfriday v1.6.0/go.mod h1:ti0ldHuxg49ri4ksnFxlkCfN+hvslNlmVHqNRXXJNAY= github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/ruudk/golang-pdf417 v0.0.0-20181029194003-1af4ab5afa58/go.mod h1:6lfFZQK844Gfx8o5WFuvpxWRwnSoipWe/p622j1v06w= github.com/ryanuber/columnize v0.0.0-20160712163229-9b3edd62028f/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts= github.com/ryanuber/go-glob v1.0.0 h1:iQh3xXAumdQ+4Ufa5b25cRpC5TYKlno6hsv6Cb3pkBk= @@ -1335,6 +1349,7 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= github.com/spf13/jwalterweatherman v1.1.0/go.mod h1:aNWZUN0dPAAO/Ljvb5BEdw96iTZ0EXowPYD95IqWIGo= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.8.1/go.mod h1:o0Pch8wJ9BVSWGQMbra6iw0oQ5oktSIBaujf1rJH9Ns= diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 86f4b8628..7ffa8384b 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -118,6 +118,9 @@ type ControlBackend interface { GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opts ...grpc.CallOption) (*pb.MetricsResponse, error) + // SlotCache saves or restores one slot's prompt cache (llama.cpp only). + SlotCache(ctx context.Context, in *pb.SlotCacheRequest, opts ...grpc.CallOption) (*pb.SlotCacheResponse, error) + // Streaming constructors: these return a stream client immediately; the // actual inference spans the stream's lifetime, not this call, so they are // NOT tracked as a single in-flight unit. diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index f46b9dcf0..f73da7382 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -633,6 +633,23 @@ func (c *Client) GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opt return client.GetMetrics(ctx, in, opts...) } +func (c *Client) SlotCache(ctx context.Context, in *pb.SlotCacheRequest, opts ...grpc.CallOption) (*pb.SlotCacheResponse, error) { + if !c.parallel { + c.opMutex.Lock() + defer c.opMutex.Unlock() + } + c.setBusy(true) + defer c.setBusy(false) + defer c.wdMark()() + conn, err := c.dial() + if err != nil { + return nil, err + } + defer func() { _ = conn.Close() }() + client := pb.NewBackendClient(conn) + return client.SlotCache(ctx, in, opts...) +} + func (c *Client) VAD(ctx context.Context, in *pb.VADRequest, opts ...grpc.CallOption) (*pb.VADResponse, error) { if !c.parallel { c.opMutex.Lock() diff --git a/pkg/grpc/embed.go b/pkg/grpc/embed.go index 59b8b3f4a..fab97e4f7 100644 --- a/pkg/grpc/embed.go +++ b/pkg/grpc/embed.go @@ -215,6 +215,10 @@ func (e *embedBackend) GetTokenMetrics(ctx context.Context, in *pb.MetricsReques return e.s.GetMetrics(ctx, in) } +func (e *embedBackend) SlotCache(ctx context.Context, in *pb.SlotCacheRequest, opts ...grpc.CallOption) (*pb.SlotCacheResponse, error) { + return e.s.SlotCache(ctx, in) +} + func (e *embedBackend) StartFineTune(ctx context.Context, in *pb.FineTuneRequest, opts ...grpc.CallOption) (*pb.FineTuneJobResult, error) { return e.s.StartFineTune(ctx, in) } diff --git a/pkg/slotcache/slotcache.go b/pkg/slotcache/slotcache.go new file mode 100644 index 000000000..64b9be071 --- /dev/null +++ b/pkg/slotcache/slotcache.go @@ -0,0 +1,49 @@ +// Package slotcache holds the checks for the slot prompt-cache save and +// restore API. The backend issues an opaque ID on save and the client sends it +// back on restore. This package decides whether an ID has the right shape +// before a request reaches a backend. +package slotcache + +import ( + "errors" + "strings" +) + +// OptionName is the llama-cpp model option that sets the save directory. +const OptionName = "slot_save_path" + +// idLen is the length of an ID: 128 random bits as lowercase hex. +const idLen = 32 + +// ErrInvalidID is returned for an ID that is not 32 lowercase hex characters. +var ErrInvalidID = errors.New("invalid id: expected 32 lowercase hex characters") + +// ValidateID accepts exactly 32 lowercase hex characters, the format the +// backend issues. Anything else cannot be an ID the backend issued, and cannot +// contain a separator, a dot or a NUL byte. The backend checks again and also +// requires the ID to be in the model's manifest. +func ValidateID(id string) error { + if len(id) != idLen { + return ErrInvalidID + } + for i := 0; i < len(id); i++ { + c := id[i] + if (c < '0' || c > '9') && (c < 'a' || c > 'f') { + return ErrInvalidID + } + } + return nil +} + +// DirFromOptions returns the slot_save_path value from a model's options, or +// "" when it is not set. Options have the form name:value. Everything after +// the first colon is the value, so a path that contains a colon is kept whole. +func DirFromOptions(options []string) string { + dir := "" + for _, o := range options { + if k, v, ok := strings.Cut(o, ":"); ok && k == OptionName { + dir = strings.TrimSpace(v) + } + } + return dir +} diff --git a/pkg/slotcache/slotcache_suite_test.go b/pkg/slotcache/slotcache_suite_test.go new file mode 100644 index 000000000..428e3a9b3 --- /dev/null +++ b/pkg/slotcache/slotcache_suite_test.go @@ -0,0 +1,13 @@ +package slotcache_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestSlotCache(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Slot cache test suite") +} diff --git a/pkg/slotcache/slotcache_test.go b/pkg/slotcache/slotcache_test.go new file mode 100644 index 000000000..714297a82 --- /dev/null +++ b/pkg/slotcache/slotcache_test.go @@ -0,0 +1,56 @@ +package slotcache_test + +import ( + "strings" + + . "github.com/mudler/LocalAI/pkg/slotcache" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("ValidateID", func() { + DescribeTable("accepts an ID the backend can issue", + func(id string) { Expect(ValidateID(id)).To(Succeed()) }, + Entry("mixed digits and letters", "0123456789abcdef0123456789abcdef"), + Entry("all digits", strings.Repeat("7", 32)), + Entry("all letters", strings.Repeat("f", 32)), + ) + + DescribeTable("rejects anything else", + func(id string) { Expect(ValidateID(id)).To(MatchError(ErrInvalidID)) }, + Entry("empty", ""), + Entry("uppercase", "0123456789ABCDEF0123456789abcdef"), + Entry("too short", "0123456789abcdef"), + Entry("too long", "0123456789abcdef0123456789abcdef0"), + Entry("non-hex letter", "0123456789abcdef0123456789abcdeg"), + Entry("dot dot", strings.Repeat(".", 32)), + Entry("parent traversal", "../../../../../../../../etc/pass"), + Entry("slash", "0123456789abcdef/123456789abcdef"), + Entry("backslash", `0123456789abcdef\123456789abcdef`), + Entry("NUL byte", "0123456789abcdef0123456789abcd\x00f"), + Entry("newline", "0123456789abcdef0123456789abcde\n"), + Entry("space", "0123456789abcdef0123456789abcde "), + Entry("non-ASCII digit", "0123456789abcdef0123456789abcd\u0663"), + Entry("file name", "prompt.bin"), + ) +}) + +var _ = Describe("DirFromOptions", func() { + It("returns empty when the option is absent", func() { + Expect(DirFromOptions([]string{"cache_ram:1024", "use_jinja"})).To(BeEmpty()) + Expect(DirFromOptions(nil)).To(BeEmpty()) + }) + + It("returns the value, keeping any colon in it", func() { + Expect(DirFromOptions([]string{"slot_save_path:/var/cache/slots"})).To(Equal("/var/cache/slots")) + Expect(DirFromOptions([]string{`slot_save_path:C:\slots`})).To(Equal(`C:\slots`)) + }) + + It("lets a later option win, as the backend does", func() { + Expect(DirFromOptions([]string{"slot_save_path:/a", "slot_save_path:/b"})).To(Equal("/b")) + }) + + It("treats the bare flag form as unset", func() { + Expect(DirFromOptions([]string{"slot_save_path"})).To(BeEmpty()) + }) +}) diff --git a/swagger/docs.go b/swagger/docs.go index a7fe807b3..6ad4eaa34 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -4229,6 +4229,62 @@ const docTemplate = `{ } } }, + "/v1/slots/restore": { + "post": { + "description": "Loads the prompt cache that a save call returned an ID for into one slot of a llama-cpp model. The ID must have been issued for this model. The state is only as fresh as the save. Admin only.", + "tags": [ + "inference" + ], + "summary": "Restore a slot's prompt cache by ID.", + "parameters": [ + { + "description": "Request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.SlotCacheRequest" + } + } + ], + "responses": { + "200": { + "description": "Response", + "schema": { + "$ref": "#/definitions/schema.SlotCacheResponse" + } + } + } + } + }, + "/v1/slots/save": { + "post": { + "description": "Writes the prompt cache (KV state) of one slot of a llama-cpp model under the directory set by the model's slot_save_path option, and returns an ID for it. The server issues the ID; the client cannot choose a name or a path. The server keeps at most 32 saves per model and deletes the oldest one beyond that. Admin only.", + "tags": [ + "inference" + ], + "summary": "Save a slot's prompt cache.", + "parameters": [ + { + "description": "Request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.SlotCacheRequest" + } + } + ], + "responses": { + "200": { + "description": "Response", + "schema": { + "$ref": "#/definitions/schema.SlotCacheResponse" + } + } + } + } + }, "/v1/sound-generation": { "post": { "tags": [ @@ -8890,6 +8946,57 @@ const docTemplate = `{ } } }, + "schema.SlotCacheRequest": { + "type": "object", + "properties": { + "id": { + "description": "restore only: an ID issued by a save for this model", + "type": "string" + }, + "model": { + "type": "string" + }, + "slot": { + "description": "slot id, 0 is the first slot", + "type": "integer" + } + } + }, + "schema.SlotCacheResponse": { + "type": "object", + "properties": { + "action": { + "description": "\"save\" or \"restore\"", + "type": "string" + }, + "duration_ms": { + "description": "time the backend spent", + "type": "number" + }, + "id": { + "description": "issued by the save, sent back to restore", + "type": "string" + }, + "message": { + "description": "for example the IDs a save pruned", + "type": "string" + }, + "model": { + "type": "string" + }, + "n_bytes": { + "description": "bytes written or read", + "type": "integer" + }, + "n_tokens": { + "description": "tokens saved or restored", + "type": "integer" + }, + "slot": { + "type": "integer" + } + } + }, "schema.SoundClassification": { "type": "object", "properties": { diff --git a/swagger/swagger.json b/swagger/swagger.json index 218478fc0..461ef12bb 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -4226,6 +4226,62 @@ } } }, + "/v1/slots/restore": { + "post": { + "description": "Loads the prompt cache that a save call returned an ID for into one slot of a llama-cpp model. The ID must have been issued for this model. The state is only as fresh as the save. Admin only.", + "tags": [ + "inference" + ], + "summary": "Restore a slot's prompt cache by ID.", + "parameters": [ + { + "description": "Request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.SlotCacheRequest" + } + } + ], + "responses": { + "200": { + "description": "Response", + "schema": { + "$ref": "#/definitions/schema.SlotCacheResponse" + } + } + } + } + }, + "/v1/slots/save": { + "post": { + "description": "Writes the prompt cache (KV state) of one slot of a llama-cpp model under the directory set by the model's slot_save_path option, and returns an ID for it. The server issues the ID; the client cannot choose a name or a path. The server keeps at most 32 saves per model and deletes the oldest one beyond that. Admin only.", + "tags": [ + "inference" + ], + "summary": "Save a slot's prompt cache.", + "parameters": [ + { + "description": "Request", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/schema.SlotCacheRequest" + } + } + ], + "responses": { + "200": { + "description": "Response", + "schema": { + "$ref": "#/definitions/schema.SlotCacheResponse" + } + } + } + } + }, "/v1/sound-generation": { "post": { "tags": [ @@ -8887,6 +8943,57 @@ } } }, + "schema.SlotCacheRequest": { + "type": "object", + "properties": { + "id": { + "description": "restore only: an ID issued by a save for this model", + "type": "string" + }, + "model": { + "type": "string" + }, + "slot": { + "description": "slot id, 0 is the first slot", + "type": "integer" + } + } + }, + "schema.SlotCacheResponse": { + "type": "object", + "properties": { + "action": { + "description": "\"save\" or \"restore\"", + "type": "string" + }, + "duration_ms": { + "description": "time the backend spent", + "type": "number" + }, + "id": { + "description": "issued by the save, sent back to restore", + "type": "string" + }, + "message": { + "description": "for example the IDs a save pruned", + "type": "string" + }, + "model": { + "type": "string" + }, + "n_bytes": { + "description": "bytes written or read", + "type": "integer" + }, + "n_tokens": { + "description": "tokens saved or restored", + "type": "integer" + }, + "slot": { + "type": "integer" + } + } + }, "schema.SoundClassification": { "type": "object", "properties": { diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index a7ec22800..f2fe92f7e 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -3019,6 +3019,42 @@ definitions: classifier-side confidence signal). type: number type: object + schema.SlotCacheRequest: + properties: + id: + description: 'restore only: an ID issued by a save for this model' + type: string + model: + type: string + slot: + description: slot id, 0 is the first slot + type: integer + type: object + schema.SlotCacheResponse: + properties: + action: + description: '"save" or "restore"' + type: string + duration_ms: + description: time the backend spent + type: number + id: + description: issued by the save, sent back to restore + type: string + message: + description: for example the IDs a save pruned + type: string + model: + type: string + n_bytes: + description: bytes written or read + type: integer + n_tokens: + description: tokens saved or restored + type: integer + slot: + type: integer + type: object schema.SoundClassification: properties: index: @@ -6602,6 +6638,48 @@ paths: summary: Cancel a response tags: - inference + /v1/slots/restore: + post: + description: Loads the prompt cache that a save call returned an ID for into + one slot of a llama-cpp model. The ID must have been issued for this model. + The state is only as fresh as the save. Admin only. + parameters: + - description: Request + in: body + name: request + required: true + schema: + $ref: '#/definitions/schema.SlotCacheRequest' + responses: + "200": + description: Response + schema: + $ref: '#/definitions/schema.SlotCacheResponse' + summary: Restore a slot's prompt cache by ID. + tags: + - inference + /v1/slots/save: + post: + description: Writes the prompt cache (KV state) of one slot of a llama-cpp model + under the directory set by the model's slot_save_path option, and returns + an ID for it. The server issues the ID; the client cannot choose a name or + a path. The server keeps at most 32 saves per model and deletes the oldest + one beyond that. Admin only. + parameters: + - description: Request + in: body + name: request + required: true + schema: + $ref: '#/definitions/schema.SlotCacheRequest' + responses: + "200": + description: Response + schema: + $ref: '#/definitions/schema.SlotCacheResponse' + summary: Save a slot's prompt cache. + tags: + - inference /v1/sound-generation: post: parameters: