mirror/LocalAI
1
0 Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2026-10-28 09:08:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat(llama-cpp): add slot prompt-cache save and restore API (#12730)

Browse Source 
* feat(llama-cpp): add slot prompt-cache save and restore API

Add POST /v1/slots/save and POST /v1/slots/restore. They write one
llama.cpp slot's prompt cache (KV state) to a file and load it back,
so a client can warm a long system prompt in a few hundred
milliseconds instead of re-running prefill. Refs #12685.

The new SlotCache gRPC call posts the same slot save and restore
tasks the llama-server /slots endpoint uses. It is off unless the
model sets the new slot_save_path option, so existing models are not
affected. The /slots monitoring endpoint is not needed here: the
backend has no HTTP server.

The file name comes from the caller, so it is checked twice, in Go
and again in the backend. It must be one path component: no empty
name, "/", "\", "..", NUL, control characters or absolute path. The
resolved path, with symbolic links followed, must be a direct child
of slot_save_path. The routes are admin only because they write
files on the backend host.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(llama-cpp): answer 501 when a backend lacks slot cache

A backend that does not implement the SlotCache rpc returns gRPC
UNIMPLEMENTED. The API passed that text through, so the client saw
"unknown method SlotCache". Map the status to a typed error in
core/backend and return 501 with a fixed message instead. Every other
error path is unchanged.

Make the SlotCache comments in backend.proto backend-neutral: slot is
the backend's state slot, filename must stay inside the save directory,
and an empty model_identity means the backend does not check it. Only
comments change, no field or name.

Add tests that the identity sent is the configured model, and for the
error mapping.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(llama-cpp): issue slot cache IDs on the server

The slot save and restore API let the client pick the file name.
Restore could load any file in slot_save_path, and the symlink check
ran before llama.cpp opened the file.

The backend now issues the name. A save returns a random 128-bit ID
(32 lowercase hex characters, read from /dev/urandom) and never
overwrites an existing file. Files live in a 0700 directory per model,
named from the SHA-256 of the model identity, with a manifest.json
that records each save. The manifest is written atomically and a
corrupt one stops both save and restore.

A restore takes the ID, finds it in the manifest, refuses an entry of
another model, builds the path from the entry and checks with lstat
that it is a plain file of the recorded size. The directory mode is
the main defence against another local user, since llama.cpp still
opens the path itself.

Each model keeps at most 32 saves and the oldest is deleted first. A
file over 8 GiB is deleted after the save, because llama.cpp does not
report the size before it writes. Responses carry no host path.

The HTTP bodies change: save takes {model, slot}, restore takes
{model, slot, id}, and an unknown ID returns 404. In the proto,
SlotCacheRequest.filename and SlotCacheResponse.filename become id.
SlotCache no longer calls checkModelIdentity, which reads a field this
request does not have.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(llama-cpp): copy slot_cache files in prepare.sh

grpc-server.cpp includes slot_cache.h, but prepare.sh copies the helper
headers into llama.cpp/tools/grpc-server/ one by one and skipped it. The
llama-cpp, turboquant and bonsai builds share this step, so they all
failed with "slot_cache.h: No such file or directory".

Copy slot_cache.h and slot_cache_test.cpp, which CMakeLists.txt builds
when LLAMA_GRPC_BUILD_TESTS is on.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* chore(deps): bump LocalAGI to pick up the scheduler start/stop fix

LocalAGI's scheduler could hit a nil pointer in run when Start and Stop
raced. That made the core/services/agentpool tests panic at random and
turned tests-linux red.

Move to mudler/LocalAGI@7aa0c8fea0, the merge of
mudler/LocalAGI#519 ("make Start and Stop safe to call concurrently").
The new module version also adds a few indirect requirements to go.mod
and go.sum.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
This commit is contained in:
mudler-agentEttore Di GiacintoEttore Di Giacinto
 authored and GitHub committed 2026-10-17 18:23:02 +02:00
1 parent 9bceeb1e5d
commit e715af1def
29 files changed
+1989 -1
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/cpp/llama-cpp/CMakeLists.txt
+6
View File
@@ -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)
backend/cpp/llama-cpp/grpc-server.cpp
+183
View File
@@ -62,6 +62,7 @@
#include "model_load_error.h"
#include "thread_params.h"
#include "parallel_params.h"
#include "slot_cache.h"
#include "message_content.h"
#include "moe_cache_option.h"
#include "passthrough_options.h"
@@ -78,10 +79,12 @@
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iterator>
#include <list>
#include <map>
#include <ctime>
#include <mutex>
#include <signal.h>
#include <thread>
@@ -678,6 +681,35 @@ static void params_parse(server_context& /*ctx_server*/, const backend::ModelOpt
// If conversion fails, keep the default value
}
}
} else if (!strcmp(optname, "slot_save_path")) {
// Directory for SlotCache save/restore files. Take everything after
// the first ':' so a path that contains one is kept whole. Without
// this option the RPC stays disabled, as in llama-server.
const size_t sep = opt.find(':');
const std::string dir = (sep == std::string::npos) ? "" : opt.substr(sep + 1);
if (!dir.empty()) {
// A missing directory is created with mode 0700. An existing
// one is left as it is, with a warning when others can read it.
std::string warn, derr;
#ifdef _WIN32
std::error_code ec;
std::filesystem::create_directories(std::filesystem::u8path(dir), ec);
const bool dir_ok = !ec && std::filesystem::is_directory(std::filesystem::u8path(dir));
#else
const bool dir_ok = llama_grpc::slot_ensure_private_dir(dir, warn, derr);
#endif
if (!warn.empty()) {
fprintf(stderr, "slot_save_path: %s\n", warn.c_str());
}
if (!dir_ok) {
fprintf(stderr, "slot_save_path %s is not a usable directory (%s), slot save/restore stays disabled\n", dir.c_str(), derr.c_str());
} else {
params.slot_save_path = dir;
if (params.slot_save_path.back() != '/' && params.slot_save_path.back() != '\\') {
params.slot_save_path += DIRECTORY_SEPARATOR;
}
}
}
} else if (!strcmp(optname, "parallel") || !strcmp(optname, "n_parallel")) {
if (optval != NULL) {
try {
@@ -3742,6 +3774,157 @@ public:
return grpc::Status::OK;
}
// slotStatusToGrpc maps a slot cache helper status to a gRPC status.
static grpc::Status slotStatusToGrpc(llama_grpc::SlotStatus s, const std::string& msg) {
switch (s) {
case llama_grpc::SlotStatus::NotFound: return grpc::Status(grpc::StatusCode::NOT_FOUND, msg);
case llama_grpc::SlotStatus::Invalid: return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, msg);
case llama_grpc::SlotStatus::Corrupt: return grpc::Status(grpc::StatusCode::DATA_LOSS, msg);
default: return grpc::Status(grpc::StatusCode::INTERNAL, msg);
}
}
// SlotCache saves or restores the prompt cache of one slot. It reuses the
// task the llama-server /slots endpoint posts, so the file format is the
// same. The caller never names a file. A save writes
// <slot_save_path>/<hash of the model identity>/<id> with an ID the backend
// generates, and records it in that directory's manifest.json. A restore
// takes an ID, looks it up in the manifest and builds the path from the
// manifest entry (see slot_cache.h for the checks and their limits).
grpc::Status SlotCache(ServerContext* context, const backend::SlotCacheRequest* request, backend::SlotCacheResponse* response) override {
#ifdef _WIN32
(void)context; (void)request; (void)response;
return grpc::Status(grpc::StatusCode::UNIMPLEMENTED, "slot save/restore is not supported on Windows");
#else
auto auth = checkAuth(context);
if (!auth.ok()) return auth;
// Same rule as checkModelIdentity; written out because this request
// names the field model_identity, so it has no modelidentity().
if (!request->model_identity().empty() && !loaded_model_identity.empty() &&
loaded_model_identity != request->model_identity()) {
return grpc::Status(grpc::StatusCode::NOT_FOUND,
"llama-cpp: model identity mismatch: loaded \"" + loaded_model_identity +
"\", requested \"" + request->model_identity() + "\"");
}
if (params_base.model.path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
}
if (params_base.slot_save_path.empty()) {
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
"slot save/restore is disabled: set the slot_save_path option for this model");
}
const bool is_save = request->action() == backend::SLOT_CACHE_ACTION_SAVE;
if (!is_save && request->action() != backend::SLOT_CACHE_ACTION_RESTORE) {
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "action must be save or restore");
}
if (request->slot() < 0) {
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "slot must not be negative");
}
// The identity that names the directory and is written to the manifest.
std::string model_id = !loaded_model_identity.empty() ? loaded_model_identity : request->model_identity();
if (model_id.empty()) {
model_id = params_base.model.path;
}
std::string err, dir, id, filepath;
llama_grpc::SlotEntry entry;
if (is_save) {
if (!request->id().empty()) {
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "id must be empty on save: the backend issues it");
}
auto s = llama_grpc::slot_identity_dir(params_base.slot_save_path, model_id, true, dir, err);
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
// Fail before writing a large file if the manifest is unusable.
std::vector<llama_grpc::SlotEntry> existing;
{
std::lock_guard<std::mutex> lock(llama_grpc::slot_cache_mutex());
s = llama_grpc::slot_manifest_load(dir, existing, err);
}
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
s = llama_grpc::slot_reserve_id(dir, id, filepath, err);
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
} else {
auto s = llama_grpc::slot_lookup_restore(params_base.slot_save_path, model_id, request->id(), entry, filepath, err);
if (s != llama_grpc::SlotStatus::Ok) return slotStatusToGrpc(s, err);
id = entry.id;
}
// On a failed save the reserved (empty or partial) file is removed.
auto drop_reserved = [&]() { if (is_save) { std::error_code rec; std::filesystem::remove(std::filesystem::u8path(filepath), rec); } };
server_task task(is_save ? SERVER_TASK_TYPE_SLOT_SAVE : SERVER_TASK_TYPE_SLOT_RESTORE);
auto rd = ctx_server.get_response_reader();
task.id = rd.get_new_id();
task.slot_action.id_slot = request->slot();
task.slot_action.filename = id;
task.slot_action.filepath = filepath;
rd.post_task(std::move(task));
auto res = rd.next([context]() { return context->IsCancelled(); });
if (!res) {
drop_reserved();
return grpc::Status(grpc::StatusCode::CANCELLED, "slot cache request cancelled");
}
if (res->is_error()) {
drop_reserved();
return grpc::Status(grpc::StatusCode::INTERNAL, res->to_json().dump());
}
auto * slot_res = dynamic_cast<server_task_result_slot_save_load *>(res.get());
if (slot_res == nullptr) {
drop_reserved();
return grpc::Status(grpc::StatusCode::INTERNAL, "unexpected result type for a slot cache task");
}
uint64_t n_bytes = slot_res->n_bytes;
std::string message;
if (is_save) {
// The size is only known after llama.cpp wrote the file. Use the
// size on disk: restore compares against it.
std::error_code sec;
const auto on_disk = std::filesystem::file_size(std::filesystem::u8path(filepath), sec);
if (sec) {
drop_reserved();
return grpc::Status(grpc::StatusCode::INTERNAL, "cannot read the size of the saved file");
}
n_bytes = static_cast<uint64_t>(on_disk);
if (n_bytes > llama_grpc::SLOT_CACHE_MAX_FILE_BYTES) {
drop_reserved();
return grpc::Status(grpc::StatusCode::RESOURCE_EXHAUSTED,
"the saved state is larger than the " + std::to_string(llama_grpc::SLOT_CACHE_MAX_FILE_BYTES) +
" byte limit and was deleted");
}
entry.id = id;
entry.model_identity = model_id;
entry.slot = request->slot();
entry.n_tokens = slot_res->n_tokens;
entry.n_bytes = n_bytes;
entry.created_unix = static_cast<int64_t>(std::time(nullptr));
std::vector<std::string> pruned;
auto s = llama_grpc::slot_commit_save(dir, entry, pruned, err);
if (s != llama_grpc::SlotStatus::Ok) {
drop_reserved();
return slotStatusToGrpc(s, err);
}
if (!pruned.empty()) {
message = "pruned the oldest saved state to keep at most " +
std::to_string(llama_grpc::SLOT_CACHE_MAX_ENTRIES) + ":";
for (const auto & p : pruned) message += " " + p;
}
}
response->set_success(true);
response->set_message(message);
response->set_slot(slot_res->id_slot);
response->set_id(id);
response->set_n_tokens(slot_res->n_tokens);
response->set_n_bytes(n_bytes);
response->set_duration_ms(slot_res->t_ms);
return grpc::Status::OK;
#endif
}
grpc::Status Detokenize(ServerContext* context, const backend::DetokenizeRequest* request, backend::DetokenizeResponse* response) override {
auto auth = checkAuth(context);
if (!auth.ok()) return auth;
backend/cpp/llama-cpp/prepare.sh
+4
View File
@@ -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/
backend/cpp/llama-cpp/slot_cache.h
+556
View File
@@ -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 <slot_save_path>/<dir>/<id>, where <dir> comes from
// the model identity. A manifest.json in <dir> lists every ID the backend
// issued. A restore accepts an ID only if the manifest lists it.
#include <nlohmann/json.hpp>
#include <cerrno>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <filesystem>
#include <functional>
#include <mutex>
#include <set>
#include <stdexcept>
#include <string>
#include <vector>
#ifndef _WIN32
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#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<size_t>(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<uint64_t>(msg.size()) * 8;
data.push_back(static_cast<char>(0x80));
while (data.size() % 64 != 56) {
data.push_back('\0');
}
for (int i = 7; i >= 0; i--) {
data.push_back(static_cast<char>((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<const unsigned char*>(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<unsigned char>(h[i] >> 24);
digest[i * 4 + 1] = static_cast<unsigned char>(h[i] >> 16);
digest[i * 4 + 2] = static_cast<unsigned char>(h[i] >> 8);
digest[i * 4 + 3] = static_cast<unsigned char>(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<unsigned>(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 <root>/<hash of identity>. 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 <dir>/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<SlotEntry>& 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<size_t>(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<std::string>();
e.model_identity = j.at("model_identity").get<std::string>();
e.slot = j.at("slot").get<int32_t>();
e.n_tokens = j.at("n_tokens").get<uint64_t>();
e.n_bytes = j.at("n_bytes").get<uint64_t>();
e.created_unix = j.at("created_unix").get<int64_t>();
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<SlotEntry>& 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<size_t>(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<bool(std::string&)>& 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<std::string>& pruned,
std::string& err, size_t max_entries = SLOT_CACHE_MAX_ENTRIES) {
namespace fs = std::filesystem;
pruned.clear();
std::lock_guard<std::mutex> lock(slot_cache_mutex());
std::vector<SlotEntry> 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<std::ptrdiff_t>(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<SlotEntry> entries;
{
std::lock_guard<std::mutex> 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<uint64_t>(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
backend/cpp/llama-cpp/slot_cache_test.cpp
+229
View File
@@ -0,0 +1,229 @@
#include "slot_cache.h"
#include <cstdio>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <set>
#include <string>
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<std::string> 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<SlotEntry> 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<std::string> 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<std::string> 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<std::string> 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<SlotEntry> 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<SlotEntry> 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<char>(in)), std::istreambuf_iterator<char>());
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;
}