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: