mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-10 13:08:55 -04:00
Compare commits
56
Commits
No files matched your search
+1
-65
@@ -29,10 +29,6 @@ updates:
|
||||
schedule:
|
||||
# Check for updates to GitHub Actions every weekday
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/bark"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/common/template"
|
||||
schedule:
|
||||
@@ -55,30 +51,10 @@ updates:
|
||||
ignore:
|
||||
- dependency-name: "torch"
|
||||
- dependency-name: "transformers"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/exllama"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/exllama2"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/mamba"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/openvoice"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/rerankers"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/sentencetransformers"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/transformers"
|
||||
schedule:
|
||||
@@ -86,44 +62,4 @@ updates:
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/backend/python/vllm"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/examples/chainlit"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/examples/functions"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/examples/langchain/langchainpy-localai-example"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/examples/langchain-chroma"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/examples/streamlit-bot"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "docker"
|
||||
directory: "/examples/k8sgpt"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "docker"
|
||||
directory: "/examples/kubernetes"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "docker"
|
||||
directory: "/examples/langchain"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "gomod"
|
||||
directory: "/examples/semantic-todo"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
- package-ecosystem: "docker"
|
||||
directory: "/examples/telegram-bot"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
interval: "weekly"
|
||||
@@ -31,13 +31,14 @@ jobs:
|
||||
messages: [
|
||||
{
|
||||
role: "system",
|
||||
content: "Write a discord message with a bullet point summary of the release notes."
|
||||
content: "Write a Discord message with a bullet point summary of the release notes. Keep the complete message under 1800 characters."
|
||||
},
|
||||
{
|
||||
role: "user",
|
||||
content: $input
|
||||
}
|
||||
]
|
||||
],
|
||||
max_tokens: 450
|
||||
}')
|
||||
|
||||
# Send the request to LocalAI API
|
||||
@@ -46,7 +47,7 @@ jobs:
|
||||
-d "$json_payload")
|
||||
|
||||
# Extract the summary from the response
|
||||
summary=$(echo $response | jq -r '.choices[0].message.content')
|
||||
summary=$(printf '%s' "$response" | jq -er '.choices[0].message.content | strings | .[0:1800]')
|
||||
|
||||
# Print the summary
|
||||
# -H "Authorization: Bearer $API_KEY" \
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
# recipe is a make target (not a prepare.sh) so 'make purge && make' is a clean
|
||||
# rebuild and so the bump bot can see the pin.
|
||||
|
||||
AUDIO_CPP_VERSION?=43001a7e0f452d80f4588e613f13332940dd4d3a
|
||||
AUDIO_CPP_VERSION?=c79e58899bf13db4d78fd06372da23cc13f55b28
|
||||
AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -41,6 +41,7 @@ define bonsai-build
|
||||
# and are applied by apply-patches.sh below.
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/patches
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build purge
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-tts-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-$(1)-build/grpc-server.cpp
|
||||
$(info $(GREEN)I bonsai build info:$(1)$(RESET))
|
||||
@@ -79,6 +80,7 @@ bonsai-cpu-all:
|
||||
# and are applied by apply-patches.sh below.
|
||||
rm -rf $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/patches
|
||||
$(MAKE) -C $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build purge
|
||||
bash $(CURRENT_MAKEFILE_DIR)/patch-grpc-server.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-score-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/grpc-server.cpp
|
||||
bash $(LLAMA_CPP_DIR)/disable-tts-task.sh $(CURRENT_MAKEFILE_DIR)/../bonsai-cpu-all-build/grpc-server.cpp
|
||||
$(info $(GREEN)I bonsai build info:cpu-all-variants$(RESET))
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
# Adapt the shared llama.cpp gRPC source to the older JSON API in Bonsai.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 1 ]]; then
|
||||
echo "usage: $0 <grpc-server.cpp>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
SRC=$1
|
||||
if [[ ! -f "$SRC" ]]; then
|
||||
echo "grpc-server.cpp not found at $SRC" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
if grep -q 'common_json_error' "$SRC"; then
|
||||
echo "==> patching $SRC to use the Bonsai JSON exception type"
|
||||
awk '{ gsub(/common_json_error/, "json::parse_error"); print }' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
echo "==> Bonsai JSON exception patch OK"
|
||||
else
|
||||
echo "==> $SRC already uses a Bonsai-compatible JSON exception type, skipping"
|
||||
fi
|
||||
@@ -1,10 +1,10 @@
|
||||
# ds4 backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DS4_VERSION?=84cc882352757baf628a1776badf7cc54d584e28
|
||||
# Upstream pin lives below as DS4_VERSION?=c1d4597a80e300b803dc642519718f2c999589da
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=84cc882352757baf628a1776badf7cc54d584e28
|
||||
DS4_VERSION?=c1d4597a80e300b803dc642519718f2c999589da
|
||||
DS4_REPO?=https://github.com/antirez/ds4
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
IK_LLAMA_VERSION?=8337e4cd3861406fc04e0854b1409cd1b027fbc9
|
||||
IK_LLAMA_VERSION?=08b500b958a3f1102e6500e5c425e65517d6fb7e
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=d59d455fd8ea09e5a2e87ce2a9d668267ffb5ccd
|
||||
LLAMA_VERSION?=eab8ee41f889ef7823af517e8098fb8a9b3cf601
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -294,7 +294,7 @@ json parse_options(bool streaming, const backend::PredictOptions* predict, const
|
||||
} else {
|
||||
SRV_WRN("[TOOLS DEBUG] parse_options: Parsed tools JSON is not an array: %s\n", tools_json.dump().c_str());
|
||||
}
|
||||
} catch (const json::parse_error& e) {
|
||||
} catch (const common_json_error& e) {
|
||||
SRV_WRN("Failed to parse tools JSON from proto: %s\n", e.what());
|
||||
SRV_WRN("[TOOLS DEBUG] parse_options: Tools string that failed to parse: %s\n", predict->tools().c_str());
|
||||
}
|
||||
@@ -324,7 +324,7 @@ json parse_options(bool streaming, const backend::PredictOptions* predict, const
|
||||
SRV_DBG("[TOOLS DEBUG] Received tool_choice object from Go layer: %s\n", tool_choice_json.dump().c_str());
|
||||
}
|
||||
SRV_INF("Extracted tool_choice from proto: %s\n", predict->toolchoice().c_str());
|
||||
} catch (const json::parse_error& e) {
|
||||
} catch (const common_json_error& e) {
|
||||
// If parsing fails, treat as string
|
||||
data["tool_choice"] = predict->toolchoice();
|
||||
SRV_INF("Extracted tool_choice as string: %s\n", predict->toolchoice().c_str());
|
||||
@@ -353,7 +353,7 @@ json parse_options(bool streaming, const backend::PredictOptions* predict, const
|
||||
// Add to data - llama.cpp server expects it as an object (map)
|
||||
data["logit_bias"] = logit_bias_json;
|
||||
SRV_INF("Using logit_bias: %s\n", predict->logitbias().c_str());
|
||||
} catch (const json::parse_error& e) {
|
||||
} catch (const common_json_error& e) {
|
||||
SRV_ERR("Failed to parse logit_bias JSON from proto: %s\n", e.what());
|
||||
}
|
||||
}
|
||||
@@ -398,7 +398,10 @@ json parse_options(bool streaming, const backend::PredictOptions* predict, const
|
||||
});
|
||||
}
|
||||
|
||||
data["stop"] = predict->stopprompts();
|
||||
data["stop"] = json::array();
|
||||
for (const auto & stop : predict->stopprompts()) {
|
||||
data["stop"].push_back(stop);
|
||||
}
|
||||
// data["n_probs"] = predict->nprobs();
|
||||
//TODO: images,
|
||||
|
||||
@@ -1795,7 +1798,7 @@ public:
|
||||
for (int j = 0; j < request->audios_size(); j++) rin.audios.push_back(request->audios(j));
|
||||
for (int j = 0; j < request->videos_size(); j++) rin.videos.push_back(request->videos(j));
|
||||
}
|
||||
messages_json.push_back(llama_grpc::build_reconstructed_message(rin));
|
||||
messages_json.push_back(json::parse(llama_grpc::build_reconstructed_message(rin).dump()));
|
||||
}
|
||||
|
||||
// Final safety check: Ensure no message has null content (Jinja templates require strings)
|
||||
@@ -1988,7 +1991,7 @@ public:
|
||||
if (!body_json.contains("chat_template_kwargs")) {
|
||||
body_json["chat_template_kwargs"] = json::object();
|
||||
}
|
||||
for (auto& el : ctk.items()) {
|
||||
for (auto el : ctk.items()) {
|
||||
body_json["chat_template_kwargs"][el.key()] = el.value();
|
||||
}
|
||||
}
|
||||
@@ -2074,30 +2077,27 @@ public:
|
||||
// If not using chat templates, extract files from image_data/audio_data fields
|
||||
// (If using chat templates, files were already extracted by oaicompat_chat_params_parse)
|
||||
if (!request->usetokenizertemplate() || request->messages_size() == 0 || ctx_server.impl->chat_params.tmpls == nullptr) {
|
||||
const auto &images_data = data.find("image_data");
|
||||
if (images_data != data.end() && images_data->is_array())
|
||||
if (data.contains("image_data") && data.at("image_data").is_array())
|
||||
{
|
||||
for (const auto &img : *images_data)
|
||||
for (const auto &img : data.at("image_data"))
|
||||
{
|
||||
auto decoded_data = base64_decode(img["data"].get<std::string>());
|
||||
files.push_back(decoded_data);
|
||||
}
|
||||
}
|
||||
|
||||
const auto &audio_data = data.find("audio_data");
|
||||
if (audio_data != data.end() && audio_data->is_array())
|
||||
if (data.contains("audio_data") && data.at("audio_data").is_array())
|
||||
{
|
||||
for (const auto &audio : *audio_data)
|
||||
for (const auto &audio : data.at("audio_data"))
|
||||
{
|
||||
auto decoded_data = base64_decode(audio["data"].get<std::string>());
|
||||
files.push_back(decoded_data);
|
||||
}
|
||||
}
|
||||
|
||||
const auto &video_data = data.find("video_data");
|
||||
if (video_data != data.end() && video_data->is_array())
|
||||
if (data.contains("video_data") && data.at("video_data").is_array())
|
||||
{
|
||||
for (const auto &video : *video_data)
|
||||
for (const auto &video : data.at("video_data"))
|
||||
{
|
||||
auto decoded_data = base64_decode(video["data"].get<std::string>());
|
||||
files.push_back(decoded_data);
|
||||
@@ -2370,7 +2370,7 @@ public:
|
||||
for (int j = 0; j < request->audios_size(); j++) rin.audios.push_back(request->audios(j));
|
||||
for (int j = 0; j < request->videos_size(); j++) rin.videos.push_back(request->videos(j));
|
||||
}
|
||||
messages_json.push_back(llama_grpc::build_reconstructed_message(rin));
|
||||
messages_json.push_back(json::parse(llama_grpc::build_reconstructed_message(rin).dump()));
|
||||
}
|
||||
|
||||
// Final safety check: Ensure no message has null content (Jinja templates require strings)
|
||||
@@ -2563,7 +2563,7 @@ public:
|
||||
if (!body_json.contains("chat_template_kwargs")) {
|
||||
body_json["chat_template_kwargs"] = json::object();
|
||||
}
|
||||
for (auto& el : ctk.items()) {
|
||||
for (auto el : ctk.items()) {
|
||||
body_json["chat_template_kwargs"][el.key()] = el.value();
|
||||
}
|
||||
}
|
||||
@@ -2649,11 +2649,10 @@ public:
|
||||
// If not using chat templates, extract files from image_data/audio_data fields
|
||||
// (If using chat templates, files were already extracted by oaicompat_chat_params_parse)
|
||||
if (!request->usetokenizertemplate() || request->messages_size() == 0 || ctx_server.impl->chat_params.tmpls == nullptr) {
|
||||
const auto &images_data = data.find("image_data");
|
||||
if (images_data != data.end() && images_data->is_array())
|
||||
if (data.contains("image_data") && data.at("image_data").is_array())
|
||||
{
|
||||
std::cout << "[PREDICT] Processing " << images_data->size() << " images" << std::endl;
|
||||
for (const auto &img : *images_data)
|
||||
std::cout << "[PREDICT] Processing " << data.at("image_data").size() << " images" << std::endl;
|
||||
for (const auto &img : data.at("image_data"))
|
||||
{
|
||||
std::cout << "[PREDICT] Processing image" << std::endl;
|
||||
auto decoded_data = base64_decode(img["data"].get<std::string>());
|
||||
@@ -2661,20 +2660,18 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
const auto &audio_data = data.find("audio_data");
|
||||
if (audio_data != data.end() && audio_data->is_array())
|
||||
if (data.contains("audio_data") && data.at("audio_data").is_array())
|
||||
{
|
||||
for (const auto &audio : *audio_data)
|
||||
for (const auto &audio : data.at("audio_data"))
|
||||
{
|
||||
auto decoded_data = base64_decode(audio["data"].get<std::string>());
|
||||
files.push_back(decoded_data);
|
||||
}
|
||||
}
|
||||
|
||||
const auto &video_data = data.find("video_data");
|
||||
if (video_data != data.end() && video_data->is_array())
|
||||
if (data.contains("video_data") && data.at("video_data").is_array())
|
||||
{
|
||||
for (const auto &video : *video_data)
|
||||
for (const auto &video : data.at("video_data"))
|
||||
{
|
||||
auto decoded_data = base64_decode(video["data"].get<std::string>());
|
||||
files.push_back(decoded_data);
|
||||
@@ -3005,7 +3002,7 @@ public:
|
||||
}
|
||||
|
||||
// Collect responses
|
||||
json responses = json::array();
|
||||
std::vector<json> responses;
|
||||
for (auto & res : all_results.results) {
|
||||
GGML_ASSERT(dynamic_cast<server_task_result_rerank*>(res.get()) != nullptr);
|
||||
responses.push_back(res->to_json());
|
||||
@@ -3018,7 +3015,7 @@ public:
|
||||
// Crop results by request.top_n if specified
|
||||
int top_n = request->top_n();
|
||||
if (top_n > 0 && top_n < static_cast<int>(responses.size())) {
|
||||
responses = json(responses.begin(), responses.begin() + top_n);
|
||||
responses.resize(top_n);
|
||||
}
|
||||
// Set usage information
|
||||
backend::Usage* usage = rerankResult->mutable_usage();
|
||||
|
||||
@@ -52,14 +52,15 @@ inline nlohmann::ordered_json normalize_message_content(const std::string& role,
|
||||
// (#7528). A multimodal user message legitimately carries a typed-part array
|
||||
// ({type:text}, {type:image_url}, ...), which must be left intact. Shared by the
|
||||
// streaming and non-streaming paths so this invariant cannot drift between them.
|
||||
inline void normalize_template_message(nlohmann::ordered_json& msg) {
|
||||
template <typename Json>
|
||||
inline void normalize_template_message(Json& msg) {
|
||||
if (!msg.contains("content")) {
|
||||
msg["content"] = ""; // templates expect the field to exist
|
||||
return;
|
||||
}
|
||||
nlohmann::ordered_json& content = msg["content"];
|
||||
auto& content = msg["content"];
|
||||
const std::string role = (msg.contains("role") && msg["role"].is_string())
|
||||
? msg["role"].get<std::string>()
|
||||
? msg["role"].template get<std::string>()
|
||||
: std::string();
|
||||
if (content.is_null()) {
|
||||
content = ""; // #7324: null would crash content[:N] slicing
|
||||
|
||||
@@ -6,10 +6,9 @@ Subject: [PATCH 1/2] score-patch
|
||||
---
|
||||
common/common.cpp | 6 +-
|
||||
common/common.h | 3 +
|
||||
tools/CMakeLists.txt | 1 +
|
||||
tools/server/server-context.cpp | 358 +++++++++++++++++++++++++++++++-
|
||||
tools/server/server-task.h | 47 +++++
|
||||
5 files changed, 406 insertions(+), 9 deletions(-)
|
||||
4 files changed, 405 insertions(+), 9 deletions(-)
|
||||
|
||||
diff --git a/common/common.cpp b/common/common.cpp
|
||||
index 2e3f14c..0cec0dc 100644
|
||||
@@ -42,15 +41,6 @@ index 878534d..4001df2 100644
|
||||
int32_t n_sequences = 1; // number of sequences to decode
|
||||
int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch)
|
||||
int32_t n_outputs_max_per_seq = 1; // max outputs per sequence
|
||||
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
|
||||
index 780df32..1d2fe8f 100644
|
||||
--- a/tools/CMakeLists.txt
|
||||
+++ b/tools/CMakeLists.txt
|
||||
@@ -41,3 +41,4 @@ else()
|
||||
add_subdirectory(fit-params)
|
||||
add_subdirectory(results)
|
||||
endif()
|
||||
+add_subdirectory(grpc-server)
|
||||
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
|
||||
index 3b5f6a1..d0e18e6 100644
|
||||
--- a/tools/server/server-context.cpp
|
||||
|
||||
@@ -8,6 +8,8 @@
|
||||
# so the grpc-server option parser skips the two references to
|
||||
# common_params::checkpoint_min_step (the default and the option handler).
|
||||
# That field does not exist in the fork yet; drop this once it does.
|
||||
# 3. Use nlohmann's parse_error type in JSON catch clauses because the fork
|
||||
# predates upstream's common_json_error wrapper.
|
||||
#
|
||||
# The fork used to lag upstream on the whole common_params_speculative refactor
|
||||
# (ggml-org/llama.cpp#22397/#22838/#22964), the model_tgt rename (#22838) and
|
||||
@@ -100,4 +102,16 @@ else
|
||||
echo "==> LOCALAI_TURBOQUANT_NO_CHECKPOINT_MIN_STEP define OK"
|
||||
fi
|
||||
|
||||
# 3. The shared source follows current upstream and catches common_json_error.
|
||||
# TurboQuant still exposes nlohmann::json directly, so its equivalent parse
|
||||
# failures use json::parse_error instead.
|
||||
if grep -q 'common_json_error' "$SRC"; then
|
||||
echo "==> patching $SRC to use the TurboQuant JSON exception type"
|
||||
awk '{ gsub(/common_json_error/, "json::parse_error"); print }' "$SRC" > "$SRC.tmp"
|
||||
mv "$SRC.tmp" "$SRC"
|
||||
echo "==> TurboQuant JSON exception patch OK"
|
||||
else
|
||||
echo "==> $SRC already uses a TurboQuant-compatible JSON exception type, skipping"
|
||||
fi
|
||||
|
||||
echo "==> all patches applied"
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# CrispASR version (release tag)
|
||||
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
|
||||
CRISPASR_VERSION?=a153b09b37c90cd55cd9336fccbdf3ba7a289596
|
||||
CRISPASR_VERSION?=ae4474dd8306384a0e697183d863dfc52e69a2fb
|
||||
SO_TARGET?=libgocrispasr.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=97d2990807fe6d558e395f8764198d7c7e7b411c
|
||||
STABLEDIFFUSION_GGML_VERSION?=50d640568388f876b0d63ee6ddb6bc86d997ec64
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ JOBS?=$(shell nproc --ignore=1 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || e
|
||||
|
||||
# vllm.cpp version
|
||||
VLLM_CPP_REPO?=https://github.com/mudler/vllm.cpp
|
||||
VLLM_CPP_VERSION?=438305e1577768ec0f75729456a4c8b9f425e2ee
|
||||
VLLM_CPP_VERSION?=6738e0b4639199f3ff0998815e4d32bfa7fe5be2
|
||||
|
||||
# MLX GEMM provider (darwin/metal only; see the metal branch below for why).
|
||||
# Consumed as the prebuilt pip wheel: building MLX from source needs `xcrun
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
package main
|
||||
|
||||
// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v21).
|
||||
// purego bindings for the vllm.cpp stable C ABI (include/vllm.h, ABI v23).
|
||||
//
|
||||
// The structs below are hand-mirrored PODs of the C declarations, with
|
||||
// explicit padding so the Go layout matches the C layout on linux/darwin
|
||||
@@ -21,7 +21,7 @@ import (
|
||||
// the header of the VLLM_CPP_VERSION pinned in the Makefile: the build checks
|
||||
// the two against each other, because a mismatch is only caught at runtime by
|
||||
// registerLib, where it takes the backend down on every load (issue #11379).
|
||||
const abiVersion = 21
|
||||
const abiVersion = 23
|
||||
|
||||
// The ABI's tri-state toggles (enable_prefix_caching ABI v7,
|
||||
// enable_jump_forward ABI v10) share one encoding: 0 is NOT "off", it is
|
||||
@@ -83,6 +83,7 @@ type cModelParams struct {
|
||||
LanguageModelOnly int32 // 0 = multimodal inputs enabled (ABI v19)
|
||||
_ [4]byte
|
||||
LimitMMPerPrompt uintptr // const char* JSON; NULL = default limits (ABI v19)
|
||||
MMProjPath uintptr // const char*; NULL = no GGUF projector (ABI v22)
|
||||
}
|
||||
|
||||
// cSamplingParams mirrors vllm_sampling_params (structured fields included).
|
||||
|
||||
@@ -16,7 +16,7 @@ func TestVllmCpp(t *testing.T) {
|
||||
RunSpecs(t, "vllm-cpp suite")
|
||||
}
|
||||
|
||||
// The Go POD mirrors must match the C struct layout of vllm.h (ABI v21)
|
||||
// The Go POD mirrors must match the C struct layout of vllm.h (ABI v23)
|
||||
// byte-for-byte: these offsets are the C offsets on LP64 (linux/darwin
|
||||
// amd64+arm64). A failure here means govllmcpp.go drifted from vllm.h.
|
||||
var _ = Describe("C ABI struct mirrors", func() {
|
||||
@@ -24,7 +24,7 @@ var _ = Describe("C ABI struct mirrors", func() {
|
||||
// VLLM_ABI_VERSION in the vllm.h of VLLM_CPP_VERSION (Makefile).
|
||||
// Moving the pin past this without growing the mirrors below ships a
|
||||
// backend that refuses every load at startup (issue #11379).
|
||||
Expect(abiVersion).To(Equal(21))
|
||||
Expect(abiVersion).To(Equal(23))
|
||||
})
|
||||
|
||||
It("cModelParams matches vllm_model_params", func() {
|
||||
@@ -51,7 +51,8 @@ var _ = Describe("C ABI struct mirrors", func() {
|
||||
Expect(unsafe.Offsetof(p.KVCacheMemoryBytes)).To(Equal(uintptr(104)))
|
||||
Expect(unsafe.Offsetof(p.LanguageModelOnly)).To(Equal(uintptr(112)))
|
||||
Expect(unsafe.Offsetof(p.LimitMMPerPrompt)).To(Equal(uintptr(120)))
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(128)))
|
||||
Expect(unsafe.Offsetof(p.MMProjPath)).To(Equal(uintptr(128)))
|
||||
Expect(unsafe.Sizeof(p)).To(Equal(uintptr(136)))
|
||||
})
|
||||
|
||||
It("cSamplingParams matches vllm_sampling_params (ABI v8)", func() {
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=4834a2327d008ace3ec5a9ed00f51454bcabbc1c
|
||||
WHISPER_CPP_VERSION?=978113305b2ead22249b881deafa131dc8884911
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -363,8 +363,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
template_kwargs["tools"] = json.loads(request.Tools)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
||||
template_kwargs["enable_thinking"] = True
|
||||
_thinking = request.Metadata.get("enable_thinking", "").lower()
|
||||
if _thinking in ("true", "false"):
|
||||
template_kwargs["enable_thinking"] = (_thinking == "true")
|
||||
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
|
||||
@@ -96,6 +96,38 @@ class TestSglangHelpers(unittest.TestCase):
|
||||
servicer._apply_engine_args({}, "[1,2,3]")
|
||||
self.assertIn("must be a JSON object", str(ctx.exception))
|
||||
|
||||
def test_build_prompt_forwards_enable_thinking(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
class Tok:
|
||||
def __init__(self):
|
||||
self.kwargs = None
|
||||
|
||||
def apply_chat_template(self, messages, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
return "PROMPT"
|
||||
|
||||
def kwargs_for(metadata):
|
||||
servicer = self._servicer()
|
||||
tok = Tok()
|
||||
servicer.tokenizer = tok
|
||||
msg = SimpleNamespace(
|
||||
role="user", content="hi", name="",
|
||||
tool_call_id="", reasoning_content="", tool_calls="",
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
Prompt="", UseTokenizerTemplate=True,
|
||||
Messages=[msg], Tools="", Metadata=metadata,
|
||||
)
|
||||
self.assertEqual(servicer._build_prompt(req), "PROMPT")
|
||||
return tok.kwargs
|
||||
|
||||
self.assertIs(kwargs_for({"enable_thinking": "true"})["enable_thinking"], True)
|
||||
# "false" used to be dropped, so Qwen3 kept thinking on
|
||||
self.assertIs(kwargs_for({"enable_thinking": "false"})["enable_thinking"], False)
|
||||
self.assertNotIn("enable_thinking", kwargs_for({}))
|
||||
self.assertIs(kwargs_for({"enable_thinking": "FALSE"})["enable_thinking"], False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -587,9 +587,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Enable thinking mode if requested
|
||||
if request.Metadata.get("enable_thinking", "").lower() == "true":
|
||||
template_kwargs["enable_thinking"] = True
|
||||
_thinking = request.Metadata.get("enable_thinking", "").lower()
|
||||
if _thinking in ("true", "false"):
|
||||
template_kwargs["enable_thinking"] = (_thinking == "true")
|
||||
|
||||
try:
|
||||
prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
|
||||
@@ -267,6 +267,10 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
}
|
||||
|
||||
// Initialize distributed mode services (NATS, object storage, node registry)
|
||||
// revisionStore is built inside the distributed block below but used after
|
||||
// the model configs are loaded, so it is declared out here.
|
||||
var revisionStore modeladmin.RevisionStore
|
||||
|
||||
distSvc, err := initDistributed(options, application.authDB, application.ModelConfigLoader())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("distributed mode initialization failed: %w", err)
|
||||
@@ -373,6 +377,9 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
cfgLoaderOpts := options.ToConfigLoaderOptions()
|
||||
modelRevisionLifecycle := modeladmin.NewDistributedModelRevisionLifecycle(distSvc.Registry, distSvc.ModelCleanup)
|
||||
gs.SetModelRevisionLifecycle(modelRevisionLifecycle)
|
||||
// Captured here, used after the model configs are loaded below: the
|
||||
// resync reads the loader, which is still empty at this point.
|
||||
revisionStore = modeladmin.NewRevisionStore(distSvc.Registry, modelRevisionLifecycle)
|
||||
gs.OnModelsChanged = func(evt messaging.CacheInvalidateEvent) {
|
||||
// ApplyRemoteChange honors the op: a "delete" prunes the element
|
||||
// (a reload-from-path is additive and cannot drop it), anything
|
||||
@@ -419,6 +426,18 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
xlog.Error("error loading config files", "error", err)
|
||||
}
|
||||
|
||||
// Bring the controller's stored revisions back in line with the
|
||||
// configuration just loaded. An inference request may only establish a
|
||||
// revision, never replace one, so a model whose stored value has drifted
|
||||
// stays unroutable until something republishes it. This has to run after
|
||||
// the load above: the loader is empty until then, and a resync against an
|
||||
// empty loader silently reconciles nothing.
|
||||
if revisionStore != nil {
|
||||
if err := modeladmin.ResyncModelConfigRevisions(options.Context, application.ModelConfigLoader(), options, revisionStore); err != nil {
|
||||
xlog.Warn("Failed to resync model config revisions", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := gallery.RegisterBackends(options.SystemState, application.ModelLoader()); err != nil {
|
||||
xlog.Error("error registering external backends", "error", err)
|
||||
}
|
||||
@@ -446,13 +465,20 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
// Wire gallery generation counter into VRAM caches so they invalidate
|
||||
// when gallery data refreshes instead of using a fixed TTL.
|
||||
vram.SetGalleryGenerationFunc(gallery.GalleryGeneration)
|
||||
if options.AutoloadGalleries {
|
||||
if options.VRAMPersistentCache {
|
||||
// Remote GGUF probes can transfer substantial metadata. Keep successful
|
||||
// results across restarts so the startup warmer does not repeat that work.
|
||||
vram.ConfigurePersistentCache(filepath.Join(options.SystemState.Model.ModelsPath, "..", "cache", "vram"), 24*time.Hour)
|
||||
}
|
||||
|
||||
// Fill those caches ahead of the first visitor. An estimate for an entry
|
||||
// nobody has asked about yet costs a remote probe of its weight files, and
|
||||
// the model gallery asks for one per row, so without this the first page
|
||||
// spends seconds filling in its own sizes while somebody watches it.
|
||||
// Non-blocking, and bounded: see DefaultEstimateWarmConfig.
|
||||
gallery.WarmEstimateCache(options.Context, options.Galleries, options.SystemState, gallery.EstimateWarmConfigFromEnv())
|
||||
// Fill those caches ahead of the first visitor. An estimate for an entry
|
||||
// nobody has asked about yet costs a remote probe of its weight files, and
|
||||
// the model gallery asks for one per row, so without this the first page
|
||||
// spends seconds filling in its own sizes while somebody watches it.
|
||||
// Non-blocking, and bounded: see DefaultEstimateWarmConfig.
|
||||
gallery.WarmEstimateCache(options.Context, options.Galleries, options.SystemState, gallery.EstimateWarmConfigFromEnv())
|
||||
}
|
||||
|
||||
if options.ConfigFile != "" {
|
||||
if err := application.ModelConfigLoader().LoadMultipleModelConfigsSingleFile(options.ConfigFile, configLoaderOpts...); err != nil {
|
||||
|
||||
+10
-2
@@ -202,10 +202,18 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo
|
||||
model.WithContext(so.Context),
|
||||
model.WithModelID(c.ModelID()),
|
||||
}
|
||||
if revision, err := config.ModelConfigRevision(&c); err == nil {
|
||||
// Use the revision stamped when the configuration was parsed, and only
|
||||
// that. By this point c has been merged with the request's prediction
|
||||
// parameters and had SetDefaults applied, so hashing it here would produce
|
||||
// a revision that depends on the request body and on whether the model file
|
||||
// parsed, which the controller reads as a config change and rejects. Every
|
||||
// config the loader hands out is stamped; an unstamped one was synthesized
|
||||
// elsewhere and is routed without a revision rather than with a wrong one.
|
||||
if revision := c.PersistedConfigRevision(); revision != "" {
|
||||
defOpts = append(defOpts, model.WithConfigRevision(revision))
|
||||
} else {
|
||||
xlog.Warn("Failed to compute model configuration revision", "model", c.ModelID(), "error", err)
|
||||
xlog.Warn("Model configuration carries no revision stamp; routing without one",
|
||||
"model", c.ModelID())
|
||||
}
|
||||
managedPrimary := len(c.Artifacts) > 0 && c.Artifacts[0].Resolved != nil
|
||||
if managedPrimary {
|
||||
|
||||
@@ -53,6 +53,7 @@ type RunCMD struct {
|
||||
BackendGalleries string `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"`
|
||||
Galleries string `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"`
|
||||
AutoloadGalleries bool `env:"LOCALAI_AUTOLOAD_GALLERIES,AUTOLOAD_GALLERIES" group:"models" default:"true"`
|
||||
VRAMPersistentCache bool `env:"LOCALAI_VRAM_PERSISTENT_CACHE,VRAM_PERSISTENT_CACHE" group:"models" default:"true" help:"Persist successful remote VRAM metadata probes across restarts"`
|
||||
AutoloadBackendGalleries bool `env:"LOCALAI_AUTOLOAD_BACKEND_GALLERIES,AUTOLOAD_BACKEND_GALLERIES" group:"backends" default:"true"`
|
||||
BackendImagesReleaseTag string `env:"LOCALAI_BACKEND_IMAGES_RELEASE_TAG,BACKEND_IMAGES_RELEASE_TAG" help:"Fallback release tag for backend images" group:"backends" default:"latest"`
|
||||
BackendImagesBranchTag string `env:"LOCALAI_BACKEND_IMAGES_BRANCH_TAG,BACKEND_IMAGES_BRANCH_TAG" help:"Fallback branch tag for backend images" group:"backends" default:"master"`
|
||||
@@ -302,6 +303,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
config.WithF16(r.F16),
|
||||
config.WithStringGalleries(r.Galleries),
|
||||
config.WithBackendGalleries(r.BackendGalleries),
|
||||
config.WithVRAMPersistentCache(r.VRAMPersistentCache),
|
||||
config.WithCors(r.CORS),
|
||||
config.WithCorsAllowOrigins(r.CORSAllowOrigins),
|
||||
config.WithDisableCSRF(r.DisableCSRF),
|
||||
|
||||
@@ -125,6 +125,7 @@ type ApplicationConfig struct {
|
||||
ExternalGRPCBackends map[string]string
|
||||
|
||||
AutoloadGalleries, AutoloadBackendGalleries bool
|
||||
VRAMPersistentCache bool
|
||||
AutoUpgradeBackends bool
|
||||
PreferDevelopmentBackends bool
|
||||
|
||||
@@ -284,6 +285,7 @@ func NewApplicationConfig(o ...AppOption) *ApplicationConfig {
|
||||
// toggle can still turn it off (a persisted false wins - see
|
||||
// loadRuntimeSettingsFromFile).
|
||||
EnableBackendLogging: true,
|
||||
VRAMPersistentCache: true,
|
||||
ArtifactDownloadConcurrency: modelartifacts.DefaultDownloadConcurrency,
|
||||
AgentJobRetentionDays: 30, // Default: 30 days
|
||||
LRUEvictionMaxRetries: 30, // Default: 30 retries
|
||||
@@ -596,6 +598,10 @@ func WithAutoUpgradeBackends(v bool) AppOption {
|
||||
return func(o *ApplicationConfig) { o.AutoUpgradeBackends = v }
|
||||
}
|
||||
|
||||
func WithVRAMPersistentCache(v bool) AppOption {
|
||||
return func(o *ApplicationConfig) { o.VRAMPersistentCache = v }
|
||||
}
|
||||
|
||||
func WithRequireBackendIntegrity(v bool) AppOption {
|
||||
return func(o *ApplicationConfig) { o.RequireBackendIntegrity = v }
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
@@ -9,6 +10,15 @@ import (
|
||||
|
||||
var _ = Describe("ApplicationConfig RuntimeSettings Conversion", func() {
|
||||
Describe("ToRuntimeSettings", func() {
|
||||
It("includes the persistent VRAM cache toggle", func() {
|
||||
encoded, err := json.Marshal(NewApplicationConfig().ToRuntimeSettings())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
var settings map[string]any
|
||||
Expect(json.Unmarshal(encoded, &settings)).To(Succeed())
|
||||
Expect(settings).To(HaveKeyWithValue("vram_persistent_cache", true))
|
||||
})
|
||||
|
||||
It("should convert all fields correctly", func() {
|
||||
appConfig := &ApplicationConfig{
|
||||
WatchDog: true,
|
||||
|
||||
@@ -43,8 +43,17 @@ type TTSConfig struct {
|
||||
|
||||
// @Description ModelConfig represents a model configuration
|
||||
type ModelConfig struct {
|
||||
modelConfigFile string `yaml:"-" json:"-"`
|
||||
modelTemplate string `yaml:"-" json:"-"`
|
||||
modelConfigFile string `yaml:"-" json:"-"`
|
||||
modelTemplate string `yaml:"-" json:"-"`
|
||||
// persistedConfigRevision is the revision of this model's persisted
|
||||
// configuration, stamped when the loader materializes it and therefore
|
||||
// before any per-request override is merged in. The request pipeline
|
||||
// mutates its copy of a ModelConfig with the caller's sampling parameters
|
||||
// (temperature, top_p, stop, ...), so hashing the config at load time is
|
||||
// the only way the controller sees one revision per configuration rather
|
||||
// than one per request body. Unexported, so it never enters the hash it
|
||||
// describes and never reaches YAML or JSON.
|
||||
persistedConfigRevision string `yaml:"-" json:"-"`
|
||||
schema.PredictionOptions `yaml:"parameters,omitempty" json:"parameters,omitempty"`
|
||||
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||||
Artifacts []modelartifacts.Spec `yaml:"artifacts,omitempty" json:"artifacts,omitempty"`
|
||||
@@ -1360,6 +1369,12 @@ func (c *ModelConfig) syncKnownUsecasesFromString() {
|
||||
c.KnownUsecaseStrings = append(c.KnownUsecaseStrings, k)
|
||||
}
|
||||
}
|
||||
// GetAllModelConfigUsecases returns a map, and ranging one yields a random
|
||||
// order per call. KnownUsecaseStrings is part of the serialized config, so
|
||||
// an unsorted list gives the same file a different config revision on every
|
||||
// load. In distributed mode that reads as a config change and the router
|
||||
// rejects the request with ErrStaleModelConfigRevision.
|
||||
slices.Sort(c.KnownUsecaseStrings)
|
||||
}
|
||||
|
||||
func (c *ModelConfig) UnmarshalYAML(value *yaml.Node) error {
|
||||
@@ -1836,6 +1851,27 @@ func (c *ModelConfig) GetModelConfigFile() string {
|
||||
return c.modelConfigFile
|
||||
}
|
||||
|
||||
// PersistedConfigRevision returns the revision stamped when this configuration
|
||||
// was loaded, or "" when it was never stamped (a config synthesized outside the
|
||||
// loader). Callers that need a revision for a request must prefer this over
|
||||
// recomputing one from the config they hold: by then the request pipeline has
|
||||
// merged the caller's prediction parameters into it.
|
||||
func (c *ModelConfig) PersistedConfigRevision() string {
|
||||
return c.persistedConfigRevision
|
||||
}
|
||||
|
||||
// StampPersistedConfigRevision records the revision of this configuration as
|
||||
// persisted. It is computed from the receiver as-is, so callers must invoke it
|
||||
// only on a configuration that has not been merged with request overrides.
|
||||
func (c *ModelConfig) StampPersistedConfigRevision() error {
|
||||
revision, err := modelConfigRevision(c)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.persistedConfigRevision = revision
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetModelTemplate returns the model's chat template if available
|
||||
func (c *ModelConfig) GetModelTemplate() string {
|
||||
return c.modelTemplate
|
||||
|
||||
@@ -168,6 +168,14 @@ func readModelConfigsFromFile(file string, opts ...ConfigLoaderOption) ([]*Model
|
||||
if err := yaml.Unmarshal(f, &configs); err == nil && len(configs) > 0 {
|
||||
for _, cc := range configs {
|
||||
cc.modelConfigFile = file
|
||||
// Stamp before SetDefaults: the revision describes what is on disk.
|
||||
// SetDefaults folds in the GGUF guess, hardware defaults and
|
||||
// app-level options, none of which are persisted configuration, and
|
||||
// the GGUF guess in particular depends on whether the model file
|
||||
// parses at that moment.
|
||||
if err := cc.StampPersistedConfigRevision(); err != nil {
|
||||
return nil, fmt.Errorf("stamping config revision for %q: %w", cc.Name, err)
|
||||
}
|
||||
cc.SetDefaults(opts...)
|
||||
cc.syncKnownUsecasesFromString()
|
||||
}
|
||||
@@ -182,6 +190,9 @@ func readModelConfigsFromFile(file string, opts ...ConfigLoaderOption) ([]*Model
|
||||
|
||||
c.modelConfigFile = file
|
||||
c.syncKnownUsecasesFromString()
|
||||
if err := c.StampPersistedConfigRevision(); err != nil {
|
||||
return nil, fmt.Errorf("stamping config revision for %q: %w", c.Name, err)
|
||||
}
|
||||
c.SetDefaults(opts...)
|
||||
|
||||
return []*ModelConfig{c}, nil
|
||||
@@ -218,6 +229,16 @@ func (bcl *ModelConfigLoader) LoadModelConfigFileByName(modelName, modelPath str
|
||||
}
|
||||
}
|
||||
|
||||
// Stamp before SetDefaults, and only when this config did not come from
|
||||
// disk already carrying one (a name with no config file on disk is
|
||||
// synthesized above). Re-stamping a loaded config here would hash it after
|
||||
// SetDefaults and reintroduce the dependency on the GGUF guess.
|
||||
if cfg.PersistedConfigRevision() == "" {
|
||||
if err := cfg.StampPersistedConfigRevision(); err != nil {
|
||||
return nil, fmt.Errorf("stamping config revision for %q: %w", modelName, err)
|
||||
}
|
||||
}
|
||||
|
||||
cfg.SetDefaults(append(opts, ModelPath(modelPath))...)
|
||||
|
||||
return cfg, nil
|
||||
@@ -944,3 +965,38 @@ func hasAnyMappingKey(mapping *yaml.Node, keys ...string) bool {
|
||||
func nonemptyScalar(node *yaml.Node) bool {
|
||||
return node != nil && node.Kind == yaml.ScalarNode && node.Tag == "!!str" && strings.TrimSpace(node.Value) != ""
|
||||
}
|
||||
|
||||
// RevisionFor returns the config revision for modelName: the one an inference
|
||||
// request for that model will carry.
|
||||
//
|
||||
// This is the only way to obtain a revision outside this package. Every
|
||||
// publisher must use it, so that what is published and what is checked are
|
||||
// the same value by construction rather than by two implementations happening
|
||||
// to agree. Hashing a ModelConfig directly is not available to callers, because
|
||||
// a config that has been through SetDefaults or the request middleware hashes
|
||||
// to something no request will ever present.
|
||||
func (bcl *ModelConfigLoader) RevisionFor(modelName string, appConfig *ApplicationConfig) (string, error) {
|
||||
cfg, err := bcl.LoadModelConfigFileByNameDefaultOptions(modelName, appConfig)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolving config revision for %q: %w", modelName, err)
|
||||
}
|
||||
return stampedRevision(cfg, modelName)
|
||||
}
|
||||
|
||||
// RevisionForPath is RevisionFor for callers that hold loader options and a
|
||||
// models path rather than an ApplicationConfig.
|
||||
func (bcl *ModelConfigLoader) RevisionForPath(modelName, modelPath string, opts ...ConfigLoaderOption) (string, error) {
|
||||
cfg, err := bcl.LoadModelConfigFileByName(modelName, modelPath, opts...)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolving config revision for %q: %w", modelName, err)
|
||||
}
|
||||
return stampedRevision(cfg, modelName)
|
||||
}
|
||||
|
||||
func stampedRevision(cfg *ModelConfig, modelName string) (string, error) {
|
||||
revision := cfg.PersistedConfigRevision()
|
||||
if revision == "" {
|
||||
return "", fmt.Errorf("no config revision stamped for %q", modelName)
|
||||
}
|
||||
return revision, nil
|
||||
}
|
||||
@@ -10,10 +10,17 @@ import (
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
// ModelConfigRevision returns a stable revision of the persisted semantic
|
||||
// modelConfigRevision returns a stable revision of the persisted semantic
|
||||
// configuration. ModelConfig's JSON tags exclude runtime-derived state and
|
||||
// source bookkeeping, while encoding/json orders map keys deterministically.
|
||||
func ModelConfigRevision(cfg *ModelConfig) (string, error) {
|
||||
//
|
||||
// Deliberately unexported. It must only ever be called on a configuration as
|
||||
// parsed from disk, before SetDefaults folds in the GGUF guess, the hardware
|
||||
// defaults and app-level options. Callers outside this package cannot tell
|
||||
// which they hold, and every time one hashed a defaulted or request-merged
|
||||
// config it published a revision no inference request would carry, which makes
|
||||
// the model unroutable. Use ModelConfigLoader.RevisionFor instead.
|
||||
func modelConfigRevision(cfg *ModelConfig) (string, error) {
|
||||
if cfg == nil {
|
||||
return "", errors.New("model config is nil")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
package config_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// The distributed controller pins a model's replicas to its config revision and
|
||||
// rejects any request carrying a different one. A revision that is not stable
|
||||
// for one unchanged file on disk therefore wedges the model.
|
||||
var _ = Describe("Model config revision stability", func() {
|
||||
// A chat model with an mmproj derives two usecase flags, FLAG_CHAT and
|
||||
// FLAG_VISION. syncKnownUsecasesFromString builds that list by ranging a
|
||||
// map, so an unstable order shows up with two or more flags and stays
|
||||
// hidden with one.
|
||||
const multiUsecaseModel = `backend: llama-cpp
|
||||
context_size: 50000
|
||||
known_usecases:
|
||||
- chat
|
||||
mmproj: llama-cpp/mmproj/example/mmproj.gguf
|
||||
name: example
|
||||
options:
|
||||
- use_jinja:true
|
||||
- parallel:2
|
||||
parameters:
|
||||
model: llama-cpp/models/example/example.gguf
|
||||
template:
|
||||
use_tokenizer_template: true
|
||||
`
|
||||
|
||||
var (
|
||||
dir string
|
||||
appConfig *config.ApplicationConfig
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(filepath.Join(dir, "example.yaml"), []byte(multiUsecaseModel), 0o600)).To(Succeed())
|
||||
appConfig = config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
})
|
||||
|
||||
loadRevision := func() string {
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
cfg, ok := loader.GetModelConfig("example")
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(cfg.PersistedConfigRevision()).ToNot(BeEmpty())
|
||||
return cfg.PersistedConfigRevision()
|
||||
}
|
||||
|
||||
It("does not change when the same file is loaded repeatedly", func() {
|
||||
baseline := loadRevision()
|
||||
for i := 0; i < 20; i++ {
|
||||
Expect(loadRevision()).To(Equal(baseline), "revision changed between two loads of one unchanged file")
|
||||
}
|
||||
})
|
||||
|
||||
It("orders the derived usecases deterministically", func() {
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
cfg, ok := loader.GetModelConfig("example")
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(len(cfg.KnownUsecaseStrings)).To(BeNumerically(">=", 2), "fixture must derive several usecases to expose ordering")
|
||||
Expect(cfg.KnownUsecaseStrings).To(Equal([]string{"FLAG_CHAT", "FLAG_VISION"}))
|
||||
})
|
||||
|
||||
// The request pipeline reloads the config through LoadModelConfigFileByName,
|
||||
// which applies SetDefaults a second time. The stamp is taken before those
|
||||
// defaults, so both the stored config and the one a request resolves carry
|
||||
// the same revision.
|
||||
It("survives the extra SetDefaults the request path applies", func() {
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
stored, ok := loader.GetModelConfig("example")
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(stored.PersistedConfigRevision()).ToNot(BeEmpty())
|
||||
|
||||
requestCfg, err := loader.LoadModelConfigFileByNameDefaultOptions("example", appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(requestCfg.PersistedConfigRevision()).To(Equal(stored.PersistedConfigRevision()))
|
||||
})
|
||||
})
|
||||
|
||||
// The revision must describe the configuration as persisted, and nothing else.
|
||||
// SetDefaults folds in values that are not persisted config: the GGUF guess
|
||||
// (which reads the model file and can fail on slow or remote storage), the
|
||||
// hardware defaults, and app-level options like threads. Hashing after that
|
||||
// made the revision a function of whether a multi-gigabyte file happened to
|
||||
// parse, so one unchanged YAML produced two different revisions depending on
|
||||
// the moment, and the controller rejected every request carrying the other one.
|
||||
var _ = Describe("Model config revision independence from runtime defaults", func() {
|
||||
It("does not change when SetDefaults is applied", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
body := "backend: llama-cpp\ncontext_size: 50000\nknown_usecases:\n - chat\n" +
|
||||
"mmproj: llama-cpp/mmproj/example/mmproj.gguf\nname: example\n" +
|
||||
"parameters:\n model: llama-cpp/models/example/example.gguf\n"
|
||||
Expect(os.WriteFile(filepath.Join(dir, "example.yaml"), []byte(body), 0o600)).To(Succeed())
|
||||
|
||||
appConfig := config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
|
||||
stored, ok := loader.GetModelConfig("example")
|
||||
Expect(ok).To(BeTrue())
|
||||
before := stored.PersistedConfigRevision()
|
||||
Expect(before).ToNot(BeEmpty())
|
||||
|
||||
// Applying defaults again is what the request path does.
|
||||
stored.SetDefaults(appConfig.ToConfigLoaderOptions()...)
|
||||
Expect(stored.PersistedConfigRevision()).To(Equal(before))
|
||||
|
||||
resolved, err := loader.LoadModelConfigFileByNameDefaultOptions("example", appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(resolved.PersistedConfigRevision()).To(Equal(before),
|
||||
"the request path must carry the same revision as the stored config")
|
||||
})
|
||||
|
||||
It("does not change when app-level defaults differ", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
Expect(os.WriteFile(filepath.Join(dir, "example.yaml"),
|
||||
[]byte("name: example\nbackend: llama-cpp\nparameters:\n model: m.gguf\n"), 0o600)).To(Succeed())
|
||||
|
||||
revWith := func(threads int, f16 bool) string {
|
||||
appConfig := config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
appConfig.Threads = threads
|
||||
appConfig.F16 = f16
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
cfg, err := loader.LoadModelConfigFileByNameDefaultOptions("example", appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return cfg.PersistedConfigRevision()
|
||||
}
|
||||
|
||||
Expect(revWith(8, false)).To(Equal(revWith(1, true)),
|
||||
"an operator changing threads must not make every model unroutable")
|
||||
})
|
||||
})
|
||||
@@ -19,10 +19,11 @@ var _ = Describe("Model configuration revisions", func() {
|
||||
return cfg
|
||||
}
|
||||
|
||||
// The raw hash is unexported on purpose, so these specs exercise it the way
|
||||
// every caller now must: by stamping the parsed config.
|
||||
revision := func(cfg *config.ModelConfig) string {
|
||||
value, err := config.ModelConfigRevision(cfg)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return value
|
||||
Expect(cfg.StampPersistedConfigRevision()).To(Succeed())
|
||||
return cfg.PersistedConfigRevision()
|
||||
}
|
||||
|
||||
It("is stable across equivalent YAML formatting and map order", func() {
|
||||
|
||||
@@ -59,6 +59,7 @@ type RuntimeSettings struct {
|
||||
BackendGalleries *[]Gallery `json:"backend_galleries,omitempty"`
|
||||
AutoloadGalleries *bool `json:"autoload_galleries,omitempty"`
|
||||
AutoloadBackendGalleries *bool `json:"autoload_backend_galleries,omitempty"`
|
||||
VRAMPersistentCache *bool `json:"vram_persistent_cache,omitempty"`
|
||||
|
||||
// API keys - No omitempty as we need to save empty arrays to clear keys
|
||||
ApiKeys *[]string `json:"api_keys"`
|
||||
|
||||
@@ -328,6 +328,10 @@ var runtimeSettingsFields = []fieldSpec{
|
||||
func(s *RuntimeSettings) **bool { return &s.AutoloadBackendGalleries },
|
||||
func(o *ApplicationConfig) bool { return o.AutoloadBackendGalleries },
|
||||
func(o *ApplicationConfig, v bool) { o.AutoloadBackendGalleries = v }),
|
||||
field("vram_persistent_cache",
|
||||
func(s *RuntimeSettings) **bool { return &s.VRAMPersistentCache },
|
||||
func(o *ApplicationConfig) bool { return o.VRAMPersistentCache },
|
||||
func(o *ApplicationConfig, v bool) { o.VRAMPersistentCache = v }),
|
||||
|
||||
// API keys: echoed for the UI, but the apply loops never touch them.
|
||||
// The settings endpoint and the file watcher own the env+runtime merge
|
||||
|
||||
@@ -45,6 +45,7 @@ func DefaultRuntimeBaseline() *ApplicationConfig {
|
||||
o.BackendGalleries = mustGalleries(DefaultBackendGalleriesJSON)
|
||||
o.AutoloadGalleries = true
|
||||
o.AutoloadBackendGalleries = true
|
||||
o.VRAMPersistentCache = true
|
||||
// core/cli/run.go injects WithMemoryReclaimer(enabled, threshold)
|
||||
// unconditionally, so the kong threshold default (0.95) reaches the
|
||||
// config even when the reclaimer flag is off - this overlay must match
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"github.com/mudler/LocalAI/pkg/downloader"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
"github.com/mudler/LocalAI/pkg/vram"
|
||||
"github.com/mudler/LocalAI/pkg/xsync"
|
||||
"github.com/mudler/xlog"
|
||||
|
||||
@@ -457,6 +458,9 @@ func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.Syste
|
||||
galleryGeneration.Add(1)
|
||||
}
|
||||
availableModelsMu.Unlock()
|
||||
if changed {
|
||||
vram.InvalidatePersistentCache()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
|
||||
@@ -294,9 +294,9 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(client.published[0]).To(Equal(messaging.CacheInvalidateEvent{
|
||||
Element: "old", Op: "delete", ConfigRevision: modeladmin.DeletedModelConfigRevision("old"),
|
||||
}))
|
||||
newConfig, ok := loader.GetModelConfig("new")
|
||||
_, ok := loader.GetModelConfig("new")
|
||||
Expect(ok).To(BeTrue())
|
||||
newRevision, err := config.ModelConfigRevision(&newConfig)
|
||||
newRevision, err := loader.RevisionForPath("new", tempDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(client.published[1]).To(Equal(messaging.CacheInvalidateEvent{
|
||||
Element: "new", Op: "install", ConfigRevision: newRevision,
|
||||
@@ -313,9 +313,9 @@ var _ = Describe("Edit Model test", func() {
|
||||
}
|
||||
_, oldOnPeer := peerLoader.GetModelConfig("old")
|
||||
Expect(oldOnPeer).To(BeFalse())
|
||||
peerConfig, newOnPeer := peerLoader.GetModelConfig("new")
|
||||
_, newOnPeer := peerLoader.GetModelConfig("new")
|
||||
Expect(newOnPeer).To(BeTrue())
|
||||
peerRevision, err := config.ModelConfigRevision(&peerConfig)
|
||||
peerRevision, err := peerLoader.RevisionForPath("new", tempDir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(peerRevision).To(Equal(newRevision))
|
||||
Expect(peerLifecycle.batches).To(Equal([][]modeladmin.ModelRevisionTransition{
|
||||
|
||||
@@ -42,6 +42,8 @@ func (s *stubNodeCommandSender) StopBackend(_, _ string) error { return nil }
|
||||
|
||||
func (s *stubNodeCommandSender) UnloadModelOnNode(_, _ string) error { return nil }
|
||||
|
||||
func (s *stubNodeCommandSender) PingNode(_ string) error { return nil }
|
||||
|
||||
var _ = Describe("ListBackendsOnNodeEndpoint", func() {
|
||||
var registry *nodes.NodeRegistry
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
@@ -12,6 +13,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openresponses"
|
||||
"github.com/mudler/LocalAI/core/p2p"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/pkg/vram"
|
||||
"github.com/mudler/LocalAI/pkg/vrambudget"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
@@ -185,6 +187,13 @@ func UpdateSettingsEndpoint(app *application.Application) echo.HandlerFunc {
|
||||
|
||||
// Apply settings using centralized method
|
||||
watchdogChanged := appConfig.ApplyRuntimeSettings(&settings)
|
||||
if settings.VRAMPersistentCache != nil || settings.AutoloadGalleries != nil {
|
||||
if appConfig.VRAMPersistentCache && appConfig.AutoloadGalleries {
|
||||
vram.ConfigurePersistentCache(filepath.Join(appConfig.SystemState.Model.ModelsPath, "..", "cache", "vram"), 24*time.Hour)
|
||||
} else {
|
||||
vram.DisablePersistentCache()
|
||||
}
|
||||
}
|
||||
|
||||
// Handle API keys specially (merge with startup keys)
|
||||
if settings.ApiKeys != nil {
|
||||
|
||||
@@ -146,6 +146,20 @@ parameters:
|
||||
Expect(resp.Details.Format).To(Equal("gguf"))
|
||||
Expect(resp.Details.Families).ToNot(BeEmpty())
|
||||
})
|
||||
|
||||
It("looks up the model when the Ollama :latest tag is included", func() {
|
||||
writeConfig("chat", `
|
||||
name: chat
|
||||
backend: llama-cpp
|
||||
template:
|
||||
chat: "{{ .Input }}"
|
||||
parameters:
|
||||
model: Llama-3-8B-Q4_K_M.gguf
|
||||
`)
|
||||
resp := callShow("chat:latest")
|
||||
Expect(resp.Details.Format).To(Equal("gguf"))
|
||||
Expect(resp.Capabilities).To(ContainElement("completion"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("ListModelsEndpoint", func() {
|
||||
|
||||
@@ -141,6 +141,12 @@ func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIR
|
||||
}
|
||||
|
||||
modelName := input.ModelName(nil)
|
||||
// Ollama-compat /api/tags appends ":latest" to untagged names.
|
||||
// Strip it for lookup so the listed name works on /api/chat,
|
||||
// /v1/chat/completions, and the other model-bearing endpoints.
|
||||
if strings.HasSuffix(modelName, ":latest") {
|
||||
modelName = strings.TrimSuffix(modelName, ":latest")
|
||||
}
|
||||
cfg, err := re.modelConfigLoader.LoadModelConfigFileByNameDefaultOptions(modelName, re.applicationConfig)
|
||||
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
package middleware_test
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// The distributed controller pins a model's replicas to the revision of its
|
||||
// persisted configuration. Inference requests only ever *establish* that
|
||||
// revision, so a revision that varies per request permanently wedges the model:
|
||||
// the first request's value is stored, and every later request carrying a
|
||||
// different one is rejected with "stale model config revision".
|
||||
var _ = Describe("Model config revision seen by inference requests", func() {
|
||||
var (
|
||||
app *echo.Echo
|
||||
modelDir string
|
||||
)
|
||||
|
||||
// revisionFor drives the real request pipeline (SetModelAndConfig ->
|
||||
// SetOpenAIRequest) and returns the config revision the handler is left
|
||||
// holding: the value core/backend.ModelOptions forwards to the model
|
||||
// router, and that the controller stores as the model's revision.
|
||||
revisionFor := func(body string) string {
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
app.ServeHTTP(rec, req)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK), "request pipeline rejected the request: %s", rec.Body.String())
|
||||
// An unstamped config would make every comparison below trivially true.
|
||||
Expect(rec.Body.String()).ToNot(BeEmpty(), "no config revision reached the handler")
|
||||
return rec.Body.String()
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
var err error
|
||||
modelDir, err = os.MkdirTemp("", "localai-revision-models-*")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(os.WriteFile(
|
||||
filepath.Join(modelDir, "test-model.yaml"),
|
||||
// The mmproj makes this derive several usecase flags. A single-flag
|
||||
// model hides any instability in how that derived list is ordered.
|
||||
[]byte("name: test-model\nbackend: llama-cpp\ncontext_size: 4096\n"+
|
||||
"mmproj: llama-cpp/mmproj/test-model/mmproj.gguf\n"+
|
||||
"known_usecases:\n - chat\n"),
|
||||
0o600,
|
||||
)).To(Succeed())
|
||||
|
||||
ss := &system.SystemState{Model: system.Model{ModelsPath: modelDir}}
|
||||
appConfig := config.NewApplicationConfig()
|
||||
appConfig.SystemState = ss
|
||||
|
||||
mcl := config.NewModelConfigLoader(modelDir)
|
||||
ml := model.NewModelLoader(ss)
|
||||
re := NewRequestExtractor(mcl, ml, appConfig)
|
||||
|
||||
app = echo.New()
|
||||
app.POST("/v1/chat/completions",
|
||||
func(c echo.Context) error {
|
||||
if err := re.SetOpenAIRequest(c); err != nil {
|
||||
return err
|
||||
}
|
||||
cfg, ok := c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig)
|
||||
Expect(ok).To(BeTrue())
|
||||
return c.String(http.StatusOK, cfg.PersistedConfigRevision())
|
||||
},
|
||||
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
|
||||
)
|
||||
})
|
||||
|
||||
AfterEach(func() { Expect(os.RemoveAll(modelDir)).To(Succeed()) })
|
||||
|
||||
It("is identical for requests that differ only in sampling parameters", func() {
|
||||
baseline := revisionFor(`{"model":"test-model","messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
Expect(revisionFor(`{"model":"test-model","temperature":0.9,"messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(baseline), "temperature must not change the persisted config revision")
|
||||
Expect(revisionFor(`{"model":"test-model","top_p":0.5,"messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(baseline), "top_p must not change the persisted config revision")
|
||||
Expect(revisionFor(`{"model":"test-model","top_k":20,"messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(baseline), "top_k must not change the persisted config revision")
|
||||
Expect(revisionFor(`{"model":"test-model","max_tokens":128,"messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(baseline), "max_tokens must not change the persisted config revision")
|
||||
Expect(revisionFor(`{"model":"test-model","stop":"STOP","messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(baseline), "stop words must not change the persisted config revision")
|
||||
})
|
||||
|
||||
It("is identical for repeated requests carrying the same sampling parameters", func() {
|
||||
body := `{"model":"test-model","temperature":0.2,"stop":"END","messages":[{"role":"user","content":"hi"}]}`
|
||||
Expect(revisionFor(body)).To(Equal(revisionFor(body)))
|
||||
})
|
||||
|
||||
// The controller compares the revision an inference request establishes
|
||||
// against the one model administration publishes when a YAML changes. If
|
||||
// the two paths hash different things, an edited model can never be routed
|
||||
// again, so they must agree on the same persisted configuration.
|
||||
It("matches the revision model administration computes for the same config", func() {
|
||||
ss := &system.SystemState{Model: system.Model{ModelsPath: modelDir}}
|
||||
appConfig := config.NewApplicationConfig()
|
||||
appConfig.SystemState = ss
|
||||
|
||||
admin := config.NewModelConfigLoader(modelDir)
|
||||
Expect(admin.LoadModelConfigsFromPath(modelDir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
loaded, ok := admin.GetModelConfig("test-model")
|
||||
Expect(ok).To(BeTrue())
|
||||
adminRevision := loaded.PersistedConfigRevision()
|
||||
Expect(adminRevision).ToNot(BeEmpty())
|
||||
|
||||
Expect(revisionFor(`{"model":"test-model","temperature":0.7,"messages":[{"role":"user","content":"hi"}]}`)).
|
||||
To(Equal(adminRevision))
|
||||
})
|
||||
})
|
||||
@@ -82,6 +82,13 @@ var _ = Describe("SetModelAndConfig middleware", func() {
|
||||
Expect(resp.Error.Message).To(ContainSubstring("not found"))
|
||||
Expect(resp.Error.Type).To(Equal("invalid_request_error"))
|
||||
})
|
||||
|
||||
It("still 404s when :latest is appended to an unknown model", func() {
|
||||
rec := postJSON(app, "/v1/chat/completions",
|
||||
`{"model":"nonexistent-model:latest","messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
Expect(rec.Code).To(Equal(http.StatusNotFound))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when the model exists as a config file", func() {
|
||||
@@ -97,6 +104,13 @@ var _ = Describe("SetModelAndConfig middleware", func() {
|
||||
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
})
|
||||
|
||||
It("accepts the Ollama :latest tag that /api/tags appends", func() {
|
||||
rec := postJSON(app, "/v1/chat/completions",
|
||||
`{"model":"test-model:latest","messages":[{"role":"user","content":"hi"}]}`)
|
||||
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when the model exists as a pre-loaded config", func() {
|
||||
|
||||
@@ -38,4 +38,13 @@ test.describe('Agents page', () => {
|
||||
// when the render won, swinging total UI coverage ~1pp run-to-run.
|
||||
await expect(page.getByRole('heading', { name: 'Create Agent' })).toBeVisible()
|
||||
})
|
||||
|
||||
test('opens agent status without leaving an active chat', async ({ page }) => {
|
||||
await page.goto('/app/agents/demo/chat?user_id=test-user')
|
||||
|
||||
const status = page.getByRole('link', { name: 'Status' })
|
||||
await expect(status).toHaveAttribute('href', '/app/agents/demo/status?user_id=test-user')
|
||||
await expect(status).toHaveAttribute('target', '_blank')
|
||||
await expect(page).toHaveURL(/\/app\/agents\/demo\/chat\?user_id=test-user$/)
|
||||
})
|
||||
})
|
||||
@@ -18,6 +18,34 @@ test.describe('Settings - Backend Logging', () => {
|
||||
await expect(input).toHaveValue('4')
|
||||
})
|
||||
|
||||
test('persistent VRAM cache can be toggled', async ({ page }) => {
|
||||
const row = page.locator('.form-row', { hasText: 'Persist remote VRAM estimates' })
|
||||
await expect(row).toBeVisible()
|
||||
|
||||
const checkbox = row.locator('input[type="checkbox"]')
|
||||
const wasChecked = await checkbox.isChecked()
|
||||
await checkbox.locator('..').click()
|
||||
if (wasChecked) {
|
||||
await expect(checkbox).not.toBeChecked()
|
||||
} else {
|
||||
await expect(checkbox).toBeChecked()
|
||||
}
|
||||
})
|
||||
|
||||
test('gallery startup loading and pre-warming can be toggled together', async ({ page }) => {
|
||||
const row = page.locator('.form-row', { hasText: 'Load and pre-warm galleries on boot' })
|
||||
await expect(row).toBeVisible()
|
||||
|
||||
const checkbox = row.locator('input[type="checkbox"]')
|
||||
const wasChecked = await checkbox.isChecked()
|
||||
await checkbox.locator('..').click()
|
||||
if (wasChecked) {
|
||||
await expect(checkbox).not.toBeChecked()
|
||||
} else {
|
||||
await expect(checkbox).toBeChecked()
|
||||
}
|
||||
})
|
||||
|
||||
test('backend logging toggle can be toggled', async ({ page }) => {
|
||||
// Find the checkbox associated with backend logging
|
||||
const section = page.locator('div', { has: page.locator('text=Enable Backend Logging') })
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { useState, useEffect, useRef, useCallback, useMemo } from 'react'
|
||||
import { useParams, useNavigate, useOutletContext, useSearchParams } from 'react-router-dom'
|
||||
import { useParams, useOutletContext, useSearchParams } from 'react-router-dom'
|
||||
import { agentsApi } from '../utils/api'
|
||||
import { apiUrl } from '../utils/basePath'
|
||||
import { renderMarkdown, highlightAll, enhanceCodeBlocks } from '../utils/markdown'
|
||||
@@ -72,7 +72,6 @@ function AgentActivityGroup({ items }) {
|
||||
|
||||
export default function AgentChat() {
|
||||
const { name } = useParams()
|
||||
const navigate = useNavigate()
|
||||
const { addToast } = useOutletContext()
|
||||
const [searchParams] = useSearchParams()
|
||||
const userId = searchParams.get('user_id') || undefined
|
||||
@@ -600,9 +599,15 @@ export default function AgentChat() {
|
||||
<i className="fas fa-layer-group" /> {artifacts.length}
|
||||
</button>
|
||||
)}
|
||||
<button className="btn btn-secondary btn-sm" onClick={() => navigate(`/app/agents/${encodeURIComponent(name)}/status${userId ? `?user_id=${encodeURIComponent(userId)}` : ''}`)} title="View status & observables">
|
||||
<a
|
||||
className="btn btn-secondary btn-sm"
|
||||
href={`/app/agents/${encodeURIComponent(name)}/status${userId ? `?user_id=${encodeURIComponent(userId)}` : ''}`}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
title="View status & observables in a new tab"
|
||||
>
|
||||
<i className="fas fa-chart-bar" /> Status
|
||||
</button>
|
||||
</a>
|
||||
<button className="btn btn-secondary btn-sm" onClick={() => clearMessages()} disabled={messages.length === 0} title="Clear chat history">
|
||||
<i className="fas fa-eraser" /> Clear
|
||||
</button>
|
||||
|
||||
@@ -482,12 +482,15 @@ export default function Settings() {
|
||||
<i className="fas fa-images text-accent" /> Galleries
|
||||
</h3>
|
||||
<div className="card">
|
||||
<SettingRow label="Autoload Galleries" description="Automatically load model galleries on startup">
|
||||
<SettingRow label="Load and pre-warm galleries on boot" description="Load model galleries and pre-warm their remote size and VRAM estimates when LocalAI starts">
|
||||
<Toggle checked={settings.autoload_galleries} onChange={(v) => update('autoload_galleries', v)} />
|
||||
</SettingRow>
|
||||
<SettingRow label="Autoload Backend Galleries" description="Automatically load backend galleries on startup">
|
||||
<Toggle checked={settings.autoload_backend_galleries} onChange={(v) => update('autoload_backend_galleries', v)} />
|
||||
</SettingRow>
|
||||
<SettingRow label="Persist remote VRAM estimates" description="Reuse successful remote model metadata probes across restarts; disabled when gallery autoload is off">
|
||||
<Toggle checked={settings.vram_persistent_cache} onChange={(v) => update('vram_persistent_cache', v)} />
|
||||
</SettingRow>
|
||||
<div className="mt-sm">
|
||||
<label className="form-label">Model Galleries (JSON)</label>
|
||||
<textarea
|
||||
|
||||
@@ -177,13 +177,17 @@ func (s *ConfigService) patchConfig(ctx context.Context, name string, patch map[
|
||||
if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil {
|
||||
return fmt.Errorf("reload configs: %w", err)
|
||||
}
|
||||
loaded, ok := s.Loader.GetModelConfig(updated.Name)
|
||||
if !ok {
|
||||
if _, ok := s.Loader.GetModelConfig(updated.Name); !ok {
|
||||
return fmt.Errorf("reload configs: model %q missing", updated.Name)
|
||||
}
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
// Resolve the revision the way an inference request does. Hashing the
|
||||
// stored config instead publishes a value no request will ever carry,
|
||||
// because SetDefaults runs again on the request path and is not
|
||||
// idempotent for every model, and the edit would leave the model
|
||||
// unroutable.
|
||||
revision, err := s.Loader.RevisionFor(updated.Name, s.AppConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("compute config revision: %w", err)
|
||||
return err
|
||||
}
|
||||
_ = s.Loader.Preload(s.modelsPath())
|
||||
pending, err := s.applyRevision(ctx, name, updated.Name, revision, updated.IsDisabled())
|
||||
@@ -343,13 +347,12 @@ func (s *ConfigService) editYAML(ctx context.Context, name string, body []byte)
|
||||
if err := s.Loader.LoadModelConfigsFromPath(modelsPath, s.AppConfig.ToConfigLoaderOptions()...); err != nil {
|
||||
return fmt.Errorf("reload configs: %w", err)
|
||||
}
|
||||
loaded, ok := s.Loader.GetModelConfig(req.Name)
|
||||
if !ok {
|
||||
if _, ok := s.Loader.GetModelConfig(req.Name); !ok {
|
||||
return fmt.Errorf("reload configs: model %q missing", req.Name)
|
||||
}
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := s.Loader.RevisionFor(req.Name, s.AppConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("compute config revision: %w", err)
|
||||
return err
|
||||
}
|
||||
if err := s.Loader.Preload(modelsPath); err != nil {
|
||||
return fmt.Errorf("preload after edit: %w", err)
|
||||
|
||||
@@ -2,6 +2,7 @@ package modeladmin
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
@@ -9,5 +10,12 @@ import (
|
||||
|
||||
func TestModelAdmin(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
// Several specs in this suite coordinate goroutines through
|
||||
// Eventually/Consistently on unbuffered-ish channels (e.g. the
|
||||
// blockingRevisionLifecycle helper). Gomega's 1s default timeout can be
|
||||
// too tight on slower or loaded CI runners (notably macOS runners),
|
||||
// causing spurious "Timed out after 1.005s" failures even though the
|
||||
// goroutines eventually make progress. Give them more headroom.
|
||||
SetDefaultEventuallyTimeout(5 * time.Second)
|
||||
RunSpecs(t, "modeladmin test suite")
|
||||
}
|
||||
@@ -49,9 +49,9 @@ func applyRemoteChange(ctx context.Context, cl *config.ModelConfigLoader, models
|
||||
disabled := true
|
||||
if exists {
|
||||
var err error
|
||||
revision, err = config.ModelConfigRevision(&cfg)
|
||||
revision, err = authoritative.RevisionForPath(name, modelsPath, opts...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("compute authoritative model config revision for %q: %w", name, err)
|
||||
return fmt.Errorf("resolve authoritative model config revision for %q: %w", name, err)
|
||||
}
|
||||
disabled = cfg.IsDisabled()
|
||||
}
|
||||
@@ -83,15 +83,9 @@ func changedConfigNames(current, snapshot map[string]config.ModelConfig, named s
|
||||
changed[name] = struct{}{}
|
||||
continue
|
||||
}
|
||||
previousRevision, err := config.ModelConfigRevision(&previous)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compute current model config revision for %q: %w", name, err)
|
||||
}
|
||||
revision, err := config.ModelConfigRevision(&cfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("compute authoritative model config revision for %q: %w", name, err)
|
||||
}
|
||||
if previousRevision != revision {
|
||||
// Both sides come from a loader, so both carry the revision stamped
|
||||
// when their file was parsed. Comparing the stamps compares the files.
|
||||
if previous.PersistedConfigRevision() != cfg.PersistedConfigRevision() {
|
||||
changed[name] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -58,9 +58,9 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
Expect(ApplyRemoteChange(context.Background(), loader, dir, evt, lifecycle)).To(Succeed())
|
||||
Expect(ApplyRemoteChange(context.Background(), loader, dir, evt, lifecycle)).To(Succeed())
|
||||
Expect(lifecycle.calls).To(HaveLen(2))
|
||||
loaded, ok := loader.GetModelConfig("peer-alias")
|
||||
_, ok := loader.GetModelConfig("peer-alias")
|
||||
Expect(ok).To(BeTrue())
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := loader.RevisionForPath("peer-alias", dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.calls[0].revision).To(Equal(revision))
|
||||
Expect(lifecycle.calls[1].revision).To(Equal(revision))
|
||||
@@ -84,7 +84,7 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
loaded, ok := loader.GetModelConfig("peer-alias")
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(loaded.ContextSize).To(HaveValue(Equal(10000)))
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := loader.RevisionForPath(loaded.Name, dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.calls).To(HaveLen(3))
|
||||
Expect(lifecycle.calls[1].revision).To(Equal(revision))
|
||||
@@ -100,7 +100,7 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
|
||||
loaded, ok := loader.GetModelConfig("reinstalled")
|
||||
Expect(ok).To(BeTrue())
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := loader.RevisionForPath(loaded.Name, dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.calls).To(HaveLen(1))
|
||||
Expect(lifecycle.calls[0].revision).To(Equal(revision))
|
||||
@@ -172,7 +172,7 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
Expect(loaded.ContextSize).To(HaveValue(Equal(10000)))
|
||||
_, ok = loader.GetModelConfig("deleted")
|
||||
Expect(ok).To(BeFalse())
|
||||
changedRevision, err := config.ModelConfigRevision(&loaded)
|
||||
changedRevision, err := loader.RevisionForPath(loaded.Name, dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.calls).To(ConsistOf(
|
||||
revisionLifecycleCall{oldName: "changed", newName: "changed", revision: changedRevision},
|
||||
@@ -228,7 +228,7 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
loaded, ok := loader.GetModelConfig("ordered")
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(loaded.ContextSize).To(HaveValue(Equal(10000)))
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := loader.RevisionForPath(loaded.Name, dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.revisions()).To(HaveLen(2))
|
||||
Expect(lifecycle.revisions()[1]).To(Equal(revision))
|
||||
@@ -263,7 +263,7 @@ var _ = Describe("ApplyRemoteChange", func() {
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(loaded.ContextSize).To(HaveValue(Equal(10000)))
|
||||
Expect(readMap(filepath.Join(dir, "ordered.yaml"))).To(HaveKeyWithValue("context_size", 10000))
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
revision, err := loader.RevisionForPath(loaded.Name, dir)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(lifecycle.revisions()).To(HaveLen(2))
|
||||
Expect(lifecycle.revisions()[1]).To(Equal(revision))
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
package modeladmin
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// A model's revision is published by administration and checked against on
|
||||
// every inference request. Those were computed by different code, and each time
|
||||
// they drifted the model became unroutable until someone deleted the row by
|
||||
// hand: the request path resolves through the loader, while publishers hashed
|
||||
// whatever ModelConfig they were holding, which by then had SetDefaults applied.
|
||||
//
|
||||
// There is now one resolver, ModelConfigLoader.RevisionFor, and the raw hash is
|
||||
// unexported so a new publisher cannot reintroduce the split. This pins the
|
||||
// property that mattered: whatever a publisher writes is what a request brings.
|
||||
var _ = Describe("Published and requested revisions agree", func() {
|
||||
var (
|
||||
dir string
|
||||
appConfig *config.ApplicationConfig
|
||||
loader *config.ModelConfigLoader
|
||||
)
|
||||
|
||||
// Several shapes, because the divergence only ever showed up on configs
|
||||
// rich enough for SetDefaults to change something: a model file to guess
|
||||
// from, several derived usecases, explicit options.
|
||||
models := map[string]string{
|
||||
"plain": "name: plain\nbackend: llama-cpp\nparameters:\n model: plain.gguf\n",
|
||||
"multimodal": "name: multimodal\nbackend: llama-cpp\ncontext_size: 50000\nknown_usecases:\n - chat\nmmproj: mm/mmproj.gguf\noptions:\n - use_jinja:true\n - parallel:2\nparameters:\n model: mm/model.gguf\n",
|
||||
"auto-ctx": "name: auto-ctx\nbackend: llama-cpp\ncontext_size: -1\nparameters:\n model: auto.gguf\n",
|
||||
"no-backend": "name: no-backend\nparameters:\n model: bare.gguf\n",
|
||||
"with-thread": "name: with-thread\nbackend: llama-cpp\nthreads: 3\nparameters:\n model: t.gguf\n",
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
for name, body := range models {
|
||||
Expect(os.WriteFile(filepath.Join(dir, name+".yaml"), []byte(body), 0o600)).To(Succeed())
|
||||
}
|
||||
appConfig = config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
appConfig.Threads = 8
|
||||
loader = config.NewModelConfigLoader(dir)
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
})
|
||||
|
||||
// requestRevision mirrors what core/backend.ModelOptions forwards to the
|
||||
// router: the stamp on the config the request pipeline resolved.
|
||||
requestRevision := func(name string) string {
|
||||
cfg, err := loader.LoadModelConfigFileByNameDefaultOptions(name, appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return cfg.PersistedConfigRevision()
|
||||
}
|
||||
|
||||
It("resolves the same revision a request will carry, for every model shape", func() {
|
||||
for name := range models {
|
||||
published, err := loader.RevisionFor(name, appConfig)
|
||||
Expect(err).ToNot(HaveOccurred(), "model %s", name)
|
||||
Expect(published).To(Equal(requestRevision(name)), "model %s: publisher and request disagree", name)
|
||||
}
|
||||
})
|
||||
|
||||
It("resolves the same revision through the path-based form", func() {
|
||||
for name := range models {
|
||||
byAppConfig, err := loader.RevisionFor(name, appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
byPath, err := loader.RevisionForPath(name, dir, appConfig.ToConfigLoaderOptions()...)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(byPath).To(Equal(byAppConfig), "model %s", name)
|
||||
}
|
||||
})
|
||||
|
||||
It("does not move when the app-level defaults change", func() {
|
||||
before := map[string]string{}
|
||||
for name := range models {
|
||||
r, err := loader.RevisionFor(name, appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
before[name] = r
|
||||
}
|
||||
|
||||
other := config.NewApplicationConfig()
|
||||
other.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
other.Threads = 1
|
||||
other.F16 = true
|
||||
other.ContextSize = 4096
|
||||
fresh := config.NewModelConfigLoader(dir)
|
||||
Expect(fresh.LoadModelConfigsFromPath(dir, other.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
|
||||
for name := range models {
|
||||
r, err := fresh.RevisionFor(name, other)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(r).To(Equal(before[name]),
|
||||
"model %s: changing an app-level setting must not make every model unroutable", name)
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,133 @@
|
||||
package modeladmin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
)
|
||||
|
||||
// ErrNoStoredRevision reports that the controller holds no revision for a
|
||||
// model, which is the normal state for one that has never been served.
|
||||
var ErrNoStoredRevision = gorm.ErrRecordNotFound
|
||||
|
||||
// RevisionStore is the controller state this resync reads and corrects.
|
||||
type RevisionStore interface {
|
||||
GetModelConfigRevision(ctx context.Context, modelName string) (string, error)
|
||||
ApplyConfigRevisions(ctx context.Context, transitions []ModelRevisionTransition) (int, error)
|
||||
}
|
||||
|
||||
// RevisionReader is the read half, satisfied by the node registry.
|
||||
type RevisionReader interface {
|
||||
GetModelConfigRevision(ctx context.Context, modelName string) (string, error)
|
||||
}
|
||||
|
||||
type revisionStore struct {
|
||||
RevisionReader
|
||||
lifecycle ModelRevisionLifecycle
|
||||
}
|
||||
|
||||
func (s revisionStore) ApplyConfigRevisions(ctx context.Context, t []ModelRevisionTransition) (int, error) {
|
||||
return s.lifecycle.ApplyConfigRevisions(ctx, t)
|
||||
}
|
||||
|
||||
// NewRevisionStore pairs the registry that holds the stored revisions with the
|
||||
// lifecycle that publishes new ones. Returns nil when either half is missing,
|
||||
// which ResyncModelConfigRevisions treats as "nothing to reconcile".
|
||||
func NewRevisionStore(reader RevisionReader, lifecycle ModelRevisionLifecycle) RevisionStore {
|
||||
if reader == nil || lifecycle == nil {
|
||||
return nil
|
||||
}
|
||||
return revisionStore{RevisionReader: reader, lifecycle: lifecycle}
|
||||
}
|
||||
|
||||
// ResyncModelConfigRevisions makes the controller's stored revision for each
|
||||
// model agree with what this build computes from the configuration on disk.
|
||||
//
|
||||
// The stored revision is what every inference request is checked against, but
|
||||
// nothing ever re-derived it from the persisted configuration: it moved only on
|
||||
// an edit, a gallery install, or a peer's change broadcast. Any other way for
|
||||
// the two to diverge left the model permanently unroutable, because an
|
||||
// inference request may only establish a revision, never replace one. A
|
||||
// configuration edited while this frontend was down, or a change in what the
|
||||
// revision is computed over, both landed there, and the only recovery was
|
||||
// deleting the row by hand.
|
||||
//
|
||||
// Running this at startup makes that self-correcting. Only a model whose stored
|
||||
// revision disagrees is republished, so replicas of models that did not drift
|
||||
// keep serving: republishing is not free, it quarantines every replica loaded
|
||||
// under the old revision.
|
||||
//
|
||||
// A model with no stored revision is left alone. It has never been served, and
|
||||
// inventing controller state for it here would quarantine nothing and describe
|
||||
// a model that may never be requested.
|
||||
func ResyncModelConfigRevisions(ctx context.Context, loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, store RevisionStore) error {
|
||||
if loader == nil || store == nil || appConfig == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
configs := loader.GetAllModelsConfigs()
|
||||
if len(configs) == 0 {
|
||||
// Reconciling nothing is indistinguishable from reconciling correctly,
|
||||
// which is how a caller that ran this before the configs were loaded
|
||||
// went unnoticed. Say so rather than report success.
|
||||
xlog.Warn("Skipping model config revision resync: no model configurations are loaded")
|
||||
return nil
|
||||
}
|
||||
|
||||
var transitions []ModelRevisionTransition
|
||||
for _, cfg := range configs {
|
||||
// Resolve the revision the way an inference request does, through the
|
||||
// loader, rather than hashing the stored config directly. SetDefaults
|
||||
// is applied again on that path and is not idempotent for every model
|
||||
// (it re-runs the GGUF guess and hardware defaults), so hashing the
|
||||
// stored config yields a value no request will ever carry, and
|
||||
// publishing it would wedge the model this resync exists to unwedge.
|
||||
want, err := loader.RevisionFor(cfg.Name, appConfig)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stored, err := store.GetModelConfigRevision(ctx, cfg.Name)
|
||||
if errors.Is(err, ErrNoStoredRevision) {
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("read stored config revision for %q: %w", cfg.Name, err)
|
||||
}
|
||||
if stored == want {
|
||||
continue
|
||||
}
|
||||
|
||||
xlog.Warn("Stored model config revision disagrees with the configuration on disk, republishing",
|
||||
"model", cfg.Name, "stored", shortRevision(stored), "computed", shortRevision(want))
|
||||
transitions = append(transitions, ModelRevisionTransition{
|
||||
ModelName: cfg.Name, ConfigRevision: want, Disabled: cfg.IsDisabled(),
|
||||
})
|
||||
}
|
||||
|
||||
if len(transitions) == 0 {
|
||||
return nil
|
||||
}
|
||||
if _, err := store.ApplyConfigRevisions(ctx, transitions); err != nil {
|
||||
return fmt.Errorf("republish model config revisions: %w", err)
|
||||
}
|
||||
xlog.Info("Republished model config revisions to match the configuration on disk", "models", len(transitions))
|
||||
return nil
|
||||
}
|
||||
|
||||
// shortRevision trims a revision for log output; the leading bytes identify it
|
||||
// well enough to tell two apart.
|
||||
func shortRevision(revision string) string {
|
||||
if revision == "" {
|
||||
return "(none)"
|
||||
}
|
||||
if len(revision) > 12 {
|
||||
return revision[:12]
|
||||
}
|
||||
return revision
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package modeladmin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// stubRevisionStore stands in for the controller's stored revisions.
|
||||
type stubRevisionStore struct {
|
||||
stored map[string]string
|
||||
getErr error
|
||||
applied []ModelRevisionTransition
|
||||
applyEr error
|
||||
}
|
||||
|
||||
func (s *stubRevisionStore) GetModelConfigRevision(_ context.Context, name string) (string, error) {
|
||||
if s.getErr != nil {
|
||||
return "", s.getErr
|
||||
}
|
||||
rev, ok := s.stored[name]
|
||||
if !ok {
|
||||
return "", ErrNoStoredRevision
|
||||
}
|
||||
return rev, nil
|
||||
}
|
||||
|
||||
func (s *stubRevisionStore) ApplyConfigRevisions(_ context.Context, t []ModelRevisionTransition) (int, error) {
|
||||
s.applied = append(s.applied, t...)
|
||||
return 0, s.applyEr
|
||||
}
|
||||
|
||||
// The controller pins a model's replicas to a stored revision and rejects any
|
||||
// request carrying a different one. Nothing ever re-derived that stored value
|
||||
// from the configuration on disk: it only moved on an edit, a gallery install
|
||||
// or a peer's change event. So whenever the stored value stopped matching what
|
||||
// this build computes for an unchanged file, every request for that model was
|
||||
// rejected until an operator deleted the row by hand.
|
||||
var _ = Describe("ResyncModelConfigRevisions", func() {
|
||||
var (
|
||||
dir string
|
||||
loader *config.ModelConfigLoader
|
||||
store *stubRevisionStore
|
||||
appConfig *config.ApplicationConfig
|
||||
)
|
||||
|
||||
write := func(name, body string) {
|
||||
Expect(os.WriteFile(filepath.Join(dir, name+".yaml"), []byte(body), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
// revisionOf resolves the revision the way an inference request does, which
|
||||
// is the value the resync must publish.
|
||||
revisionOf := func(name string) string {
|
||||
cfg, err := loader.LoadModelConfigFileByNameDefaultOptions(name, appConfig)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(cfg.PersistedConfigRevision()).ToNot(BeEmpty())
|
||||
return cfg.PersistedConfigRevision()
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir = GinkgoT().TempDir()
|
||||
appConfig = config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
loader = config.NewModelConfigLoader(dir)
|
||||
store = &stubRevisionStore{stored: map[string]string{}}
|
||||
})
|
||||
|
||||
load := func() {
|
||||
Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed())
|
||||
}
|
||||
|
||||
It("republishes the revision when the stored one no longer matches the config on disk", func() {
|
||||
write("drifted", "name: drifted\nbackend: llama-cpp\ncontext_size: 4096\n")
|
||||
load()
|
||||
store.stored["drifted"] = "a-revision-from-an-earlier-build"
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).To(Succeed())
|
||||
|
||||
Expect(store.applied).To(HaveLen(1))
|
||||
Expect(store.applied[0].ModelName).To(Equal("drifted"))
|
||||
Expect(store.applied[0].ConfigRevision).To(Equal(revisionOf("drifted")))
|
||||
})
|
||||
|
||||
It("leaves a model alone when the stored revision already matches", func() {
|
||||
write("agreed", "name: agreed\nbackend: llama-cpp\n")
|
||||
load()
|
||||
store.stored["agreed"] = revisionOf("agreed")
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).To(Succeed())
|
||||
|
||||
Expect(store.applied).To(BeEmpty(), "republishing an unchanged revision would quarantine live replicas for nothing")
|
||||
})
|
||||
|
||||
// A model nobody has served has no stored revision. Creating one here would
|
||||
// invent controller state for a model that may never be requested; the first
|
||||
// request establishes it.
|
||||
It("does not create state for a model that has never been served", func() {
|
||||
write("never-served", "name: never-served\nbackend: llama-cpp\n")
|
||||
load()
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).To(Succeed())
|
||||
|
||||
Expect(store.applied).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("republishes only the models that actually drifted", func() {
|
||||
write("drifted", "name: drifted\nbackend: llama-cpp\n")
|
||||
write("agreed", "name: agreed\nbackend: llama-cpp\ncontext_size: 2048\n")
|
||||
load()
|
||||
store.stored["drifted"] = "stale"
|
||||
store.stored["agreed"] = revisionOf("agreed")
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).To(Succeed())
|
||||
|
||||
Expect(store.applied).To(HaveLen(1))
|
||||
Expect(store.applied[0].ModelName).To(Equal("drifted"))
|
||||
})
|
||||
|
||||
It("reports a store failure instead of continuing silently", func() {
|
||||
write("drifted", "name: drifted\nbackend: llama-cpp\n")
|
||||
load()
|
||||
store.stored["drifted"] = "stale"
|
||||
store.applyEr = errors.New("database is down")
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).ToNot(Succeed())
|
||||
})
|
||||
|
||||
It("skips a model whose stored revision cannot be read rather than guessing", func() {
|
||||
write("unreadable", "name: unreadable\nbackend: llama-cpp\n")
|
||||
load()
|
||||
store.getErr = errors.New("connection reset")
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).ToNot(Succeed())
|
||||
Expect(store.applied).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
// Running the resync before the model configs are loaded reconciled nothing
|
||||
// while reporting success, which is how a mis-ordered startup call went
|
||||
// unnoticed. An empty loader is now called out instead of looking like a
|
||||
// clean run.
|
||||
var _ = Describe("ResyncModelConfigRevisions with nothing loaded", func() {
|
||||
It("does not touch stored revisions when no configs are loaded", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
loader := config.NewModelConfigLoader(dir)
|
||||
appConfig := config.NewApplicationConfig()
|
||||
appConfig.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
store := &stubRevisionStore{stored: map[string]string{"served-before": "stale"}}
|
||||
|
||||
Expect(ResyncModelConfigRevisions(context.Background(), loader, appConfig, store)).To(Succeed())
|
||||
Expect(store.applied).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -7,7 +7,6 @@ import (
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
)
|
||||
|
||||
@@ -61,13 +60,17 @@ func (s *ConfigService) toggleState(ctx context.Context, name string, action Act
|
||||
if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil {
|
||||
return fmt.Errorf("reload configs: %w", err)
|
||||
}
|
||||
loaded, ok := s.Loader.GetModelConfig(name)
|
||||
if !ok {
|
||||
if _, ok := s.Loader.GetModelConfig(name); !ok {
|
||||
return fmt.Errorf("reload configs: model %q missing", name)
|
||||
}
|
||||
revision, err := config.ModelConfigRevision(&loaded)
|
||||
// Resolve the revision the way an inference request does. Hashing the
|
||||
// stored config instead publishes a value no request will ever carry,
|
||||
// because SetDefaults runs again on the request path and is not
|
||||
// idempotent for every model, and the edit would leave the model
|
||||
// unroutable.
|
||||
revision, err := s.Loader.RevisionFor(name, s.AppConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("compute config revision: %w", err)
|
||||
return err
|
||||
}
|
||||
pending, err := s.applyRevision(ctx, name, name, revision, action == ActionDisable)
|
||||
if err != nil {
|
||||
|
||||
@@ -60,6 +60,7 @@ type ModelRouter interface {
|
||||
GetNodeLabels(ctx context.Context, nodeID string) ([]NodeLabel, error)
|
||||
FindNodesWithModel(ctx context.Context, modelName string) ([]BackendNode, error)
|
||||
LoadedReplicaStats(ctx context.Context, modelName string, candidateNodeIDs []string) ([]ReplicaCandidate, error)
|
||||
MarkUnhealthy(ctx context.Context, nodeID string) error
|
||||
LoadJobStore
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// Agent workers do not subscribe to the backend.* subjects, so asking one to
|
||||
// list its backends can only answer "no responders". ListBackends read that as
|
||||
// a node that had gone away and marked it unhealthy; the node's next heartbeat
|
||||
// marked it healthy again. Every poll of the backends view therefore flapped
|
||||
// every agent node in the cluster, and while it was unhealthy the router would
|
||||
// not schedule onto it.
|
||||
var _ = Describe("Backend listing across mixed node types", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
mc *scriptedMessagingClient
|
||||
mgr *DistributedBackendManager
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
mc = newScriptedMessagingClient()
|
||||
mgr = &DistributedBackendManager{
|
||||
local: stubLocalBackendManager{},
|
||||
adapter: NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute),
|
||||
registry: registry,
|
||||
}
|
||||
ctx = context.Background()
|
||||
})
|
||||
|
||||
register := func(name, nodeType string) *BackendNode {
|
||||
node := &BackendNode{Name: name, NodeType: nodeType, Address: name + ":50051"}
|
||||
Expect(registry.Register(ctx, node, true)).To(Succeed())
|
||||
fetched, err := registry.GetByName(ctx, name)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(fetched.Status).To(Equal(StatusHealthy))
|
||||
return fetched
|
||||
}
|
||||
|
||||
statusOf := func(id string) string {
|
||||
n, err := registry.Get(ctx, id)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return n.Status
|
||||
}
|
||||
|
||||
It("leaves an agent node healthy instead of flapping it", func() {
|
||||
agent := register("agent-worker-1", NodeTypeAgent)
|
||||
mc.scriptNoResponders(messaging.SubjectNodeBackendList(agent.ID))
|
||||
|
||||
_, err := mgr.ListBackends()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(statusOf(agent.ID)).To(Equal(StatusHealthy),
|
||||
"an agent node cannot answer backend.list and must not be judged on it")
|
||||
})
|
||||
|
||||
It("still marks a backend node unhealthy when it does not answer", func() {
|
||||
backendNode := register("worker-a", NodeTypeBackend)
|
||||
mc.scriptNoResponders(messaging.SubjectNodeBackendList(backendNode.ID))
|
||||
|
||||
_, err := mgr.ListBackends()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(statusOf(backendNode.ID)).To(Equal(StatusUnhealthy),
|
||||
"a backend worker that does not answer is genuinely gone")
|
||||
})
|
||||
})
|
||||
@@ -331,8 +331,9 @@ func (d *DistributedBackendManager) DeleteBackendDetailed(ctx context.Context, n
|
||||
// populated from the first node seen so single-node-minded callers still work.
|
||||
//
|
||||
// Pending/offline/draining nodes are skipped because they aren't expected to
|
||||
// answer NATS requests; unhealthy nodes are still queried — ErrNoResponders
|
||||
// then marks them unhealthy and the loop continues.
|
||||
// answer NATS requests, and so are non-backend workers, which do not subscribe
|
||||
// to backend.list at all; unhealthy backend nodes are still queried —
|
||||
// ErrNoResponders then marks them unhealthy and the loop continues.
|
||||
func (d *DistributedBackendManager) ListBackends() (gallery.SystemBackends, error) {
|
||||
result := make(gallery.SystemBackends)
|
||||
allNodes, err := d.registry.List(context.Background())
|
||||
@@ -344,6 +345,14 @@ func (d *DistributedBackendManager) ListBackends() (gallery.SystemBackends, erro
|
||||
if node.Status == StatusPending || node.Status == StatusOffline || node.Status == StatusDraining {
|
||||
continue
|
||||
}
|
||||
// Only backend workers subscribe to backend.list. Asking an agent
|
||||
// worker can only answer "no responders", which the error handling
|
||||
// below reads as a node that has gone away, so every poll of this view
|
||||
// marked every agent node unhealthy and its next heartbeat marked it
|
||||
// healthy again. The backend-op fan-out skips them for the same reason.
|
||||
if node.NodeType != "" && node.NodeType != NodeTypeBackend {
|
||||
continue
|
||||
}
|
||||
reply, err := d.adapter.ListBackends(node.ID)
|
||||
if err != nil {
|
||||
if errors.Is(err, nats.ErrNoResponders) {
|
||||
|
||||
@@ -250,3 +250,7 @@ var _ = Describe("ModelRouterAdapter", func() {
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
func (f *fakeModelRouterForSmartRouter) MarkUnhealthy(_ context.Context, _ string) error {
|
||||
return nil
|
||||
}
|
||||
@@ -278,8 +278,9 @@ func (rc *ReplicaReconciler) reconcileOnce(ctx context.Context) {
|
||||
|
||||
// reconcileState runs the state-reconciliation passes: drain pending backend
|
||||
// ops for freshly-healthy nodes, reconcile registry rows against what workers
|
||||
// report they are running, then port-probe whatever is left. All passes are
|
||||
// best-effort: a failure on one node doesn't stop the rest.
|
||||
// report they are running, port-probe whatever is left, then reclaim replica
|
||||
// slots held by loads nobody is driving. All passes are best-effort: a failure
|
||||
// on one node doesn't stop the rest.
|
||||
//
|
||||
// Order matters. The worker pass runs first and refreshes updated_at for every
|
||||
// model a worker vouches for, which takes those rows out of the port prober's
|
||||
@@ -292,6 +293,9 @@ func (rc *ReplicaReconciler) reconcileState(ctx context.Context) {
|
||||
rc.reconcileNodeProcesses(ctx)
|
||||
rc.probeLoadedModels(ctx)
|
||||
rc.sweepLeakedInFlight(ctx)
|
||||
// Runs last: the passes above can move a row into a serving state, and a
|
||||
// row that just became loaded is no longer this sweeper's business.
|
||||
rc.reclaimAbandonedLoads(ctx)
|
||||
}
|
||||
|
||||
// drainPendingBackendOps retries queued backend ops whose next_retry_at has
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
// abandonedLoadGrace is how long a replica row may sit in a pre-serving
|
||||
// state before the sweeper will consider it at all.
|
||||
//
|
||||
// It exists to cover the window between creating the replica row and
|
||||
// writing the load job that vouches for it. Without it a load could be
|
||||
// reclaimed in the moment before its own job row exists. It is not the
|
||||
// thing that protects a long transfer: the job heartbeat does that.
|
||||
abandonedLoadGrace = 5 * time.Minute
|
||||
)
|
||||
|
||||
// preServingStates are the replica states that hold a slot without being able
|
||||
// to serve a request. NextFreeReplicaIndex counts every state except
|
||||
// "unloading", so a row parked in one of these occupies capacity while
|
||||
// answering nothing.
|
||||
var preServingStates = []string{"loading", "staging"}
|
||||
|
||||
// reclaimAbandonedLoads removes replica rows whose load will never finish.
|
||||
//
|
||||
// The other reconciler passes and the router's eviction query all filter
|
||||
// state = "loaded", and the per-model probe skips rows without an address, so
|
||||
// nothing reclaimed a row that never got that far. On a node with one replica
|
||||
// slot per model, a single interrupted transfer made the model unschedulable
|
||||
// there until an operator intervened: scheduling saw no free slot, and eviction
|
||||
// found nothing it was allowed to evict.
|
||||
//
|
||||
// A row is only reclaimed when something proves the load is not progressing:
|
||||
// either a load job that has failed or stopped heartbeating, or, for a row with
|
||||
// no job at all, a node that is no longer healthy.
|
||||
//
|
||||
// The no-job case has to be conservative. Only the request path creates load
|
||||
// jobs; the reconciler's own scale-up loads a replica without one. Treating a
|
||||
// missing job as proof of abandonment would let this sweeper delete a healthy
|
||||
// reconciler-driven transfer the moment it ran past the grace period, which for
|
||||
// a multi-gigabyte checkpoint is every time. A healthy node with no job is
|
||||
// therefore left alone; when the node is gone, nothing can be progressing and
|
||||
// the row is safe to reclaim.
|
||||
func (rc *ReplicaReconciler) reclaimAbandonedLoads(ctx context.Context) {
|
||||
if rc.db == nil {
|
||||
return
|
||||
}
|
||||
|
||||
cutoff := time.Now().Add(-abandonedLoadGrace)
|
||||
var stuck []NodeModel
|
||||
if err := rc.db.WithContext(ctx).
|
||||
Where("state IN ? AND updated_at < ?", preServingStates, cutoff).
|
||||
Find(&stuck).Error; err != nil {
|
||||
xlog.Warn("Reconciler: failed to list replicas stuck before serving", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
for _, row := range stuck {
|
||||
if !rc.loadAbandoned(ctx, row, now) {
|
||||
continue
|
||||
}
|
||||
if err := rc.registry.RemoveNodeModel(ctx, row.NodeID, row.ModelName, row.ReplicaIndex); err != nil {
|
||||
xlog.Warn("Reconciler: failed to reclaim abandoned load",
|
||||
"node", row.NodeID, "model", row.ModelName, "replica", row.ReplicaIndex,
|
||||
"state", row.State, "error", err)
|
||||
continue
|
||||
}
|
||||
xlog.Warn("Reconciler: reclaimed a replica slot held by a load nobody is driving",
|
||||
"node", row.NodeID, "model", row.ModelName, "replica", row.ReplicaIndex, "state", row.State)
|
||||
}
|
||||
}
|
||||
|
||||
// loadAbandoned reports whether this row's load has demonstrably stopped.
|
||||
//
|
||||
// Every uncertain case answers false. Leaving a slot held for another pass
|
||||
// costs one scheduling opportunity; reclaiming a row out from under a live
|
||||
// transfer restarts a multi-gigabyte load and, on a single-slot node, makes the
|
||||
// model unschedulable there for as long as the retry loop runs.
|
||||
func (rc *ReplicaReconciler) loadAbandoned(ctx context.Context, row NodeModel, now time.Time) bool {
|
||||
job, err := rc.registry.GetLoadJob(ctx, row.ModelName)
|
||||
switch {
|
||||
case errors.Is(err, gorm.ErrRecordNotFound), err == nil && job == nil:
|
||||
// No job: only the request path creates them, so this may be a healthy
|
||||
// reconciler-driven load. Reclaim only once its node is gone.
|
||||
return !rc.nodeHealthy(ctx, row.NodeID)
|
||||
case err != nil:
|
||||
xlog.Warn("Reconciler: cannot read load job, leaving the replica slot held",
|
||||
"model", row.ModelName, "error", err)
|
||||
return false
|
||||
case job.State == LoadJobStateFailed:
|
||||
return true
|
||||
default:
|
||||
return job.IsOrphaned(now)
|
||||
}
|
||||
}
|
||||
|
||||
// nodeHealthy reports whether the row's node is still healthy. An unreadable
|
||||
// node counts as healthy so a database blip cannot trigger a reclaim.
|
||||
func (rc *ReplicaReconciler) nodeHealthy(ctx context.Context, nodeID string) bool {
|
||||
node, err := rc.registry.Get(ctx, nodeID)
|
||||
if err != nil || node == nil {
|
||||
xlog.Warn("Reconciler: cannot read node for a stuck replica, leaving the slot held",
|
||||
"node", nodeID, "error", err)
|
||||
return true
|
||||
}
|
||||
return node.Status == StatusHealthy
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// A replica row in loading or staging holds its slot: NextFreeReplicaIndex
|
||||
// counts every state except unloading. Nothing reclaimed such a row. Every
|
||||
// reconciler sweep and the router's eviction query filter state = "loaded", and
|
||||
// the per-model health probe skips rows with no address, which is exactly what a
|
||||
// row that never finished loading has. So a worker that dropped out mid-transfer
|
||||
// left a row that pinned the only replica slot on that node for that model, and
|
||||
// the next request failed with "no replica slot ... all models busy".
|
||||
//
|
||||
// Elapsed time alone cannot decide this: staging a large checkpoint legitimately
|
||||
// runs for tens of minutes. The load job's LastProgress heartbeat is the
|
||||
// discriminator, the same signal job takeover already trusts.
|
||||
var _ = Describe("ReplicaReconciler — abandoned load sweeper", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
node *BackendNode
|
||||
rc *ReplicaReconciler
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
node = &BackendNode{Name: "n1", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
|
||||
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
||||
rc = NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db})
|
||||
})
|
||||
|
||||
// seedReplica creates a replica row in the given state, aged so it is past
|
||||
// the sweeper's grace period unless stated otherwise.
|
||||
seedReplica := func(model, state string, age time.Duration) {
|
||||
Expect(db.Create(&NodeModel{
|
||||
ID: model + "-row",
|
||||
NodeID: node.ID,
|
||||
ModelName: model,
|
||||
State: state,
|
||||
UpdatedAt: time.Now().Add(-age),
|
||||
}).Error).To(Succeed())
|
||||
}
|
||||
|
||||
seedJob := func(model, state string, sinceProgress time.Duration) {
|
||||
Expect(db.Create(&ModelLoadJob{
|
||||
TrackingKey: model,
|
||||
State: state,
|
||||
OwnerReplica: "someone",
|
||||
LastProgress: time.Now().Add(-sinceProgress),
|
||||
CreatedAt: time.Now().Add(-sinceProgress),
|
||||
UpdatedAt: time.Now().Add(-sinceProgress),
|
||||
}).Error).To(Succeed())
|
||||
}
|
||||
|
||||
rowExists := func(model string) bool {
|
||||
var count int64
|
||||
Expect(db.Model(&NodeModel{}).Where("model_name = ?", model).Count(&count).Error).To(Succeed())
|
||||
return count > 0
|
||||
}
|
||||
|
||||
It("reclaims a staging row whose load job has stopped heartbeating", func() {
|
||||
seedReplica("abandoned", "staging", time.Hour)
|
||||
seedJob("abandoned", LoadJobStateStaging, 30*time.Minute)
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("abandoned")).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reclaims a jobless row once its node is gone", func() {
|
||||
seedReplica("orphan", "loading", time.Hour)
|
||||
Expect(registry.MarkUnhealthy(context.Background(), node.ID)).To(Succeed())
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("orphan")).To(BeFalse())
|
||||
})
|
||||
|
||||
// Only the request path creates load jobs. The reconciler's own scale-up
|
||||
// loads a replica without one, so treating a missing job as abandonment
|
||||
// deleted healthy transfers the moment they outran the grace period, which
|
||||
// for a multi-gigabyte checkpoint is every time. That is what made a replica
|
||||
// appear to hop between nodes instead of finishing anywhere.
|
||||
It("keeps a jobless row while its node is still healthy", func() {
|
||||
seedReplica("scaling-up", "staging", time.Hour)
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("scaling-up")).To(BeTrue(),
|
||||
"a reconciler-driven load has no job row and must not be reclaimed for it")
|
||||
})
|
||||
|
||||
It("keeps a long transfer whose job is still heartbeating", func() {
|
||||
// The row itself is old, because staging does not touch it. Only the
|
||||
// job proves the transfer is alive.
|
||||
seedReplica("big-model", "staging", time.Hour)
|
||||
seedJob("big-model", LoadJobStateStaging, time.Second)
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("big-model")).To(BeTrue(), "a live transfer must never be reclaimed")
|
||||
})
|
||||
|
||||
It("leaves a freshly created row alone while its job row is still being written", func() {
|
||||
seedReplica("just-started", "loading", time.Second)
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("just-started")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("does not touch loaded replicas, which the other sweeps own", func() {
|
||||
seedReplica("serving", "loaded", time.Hour)
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
Expect(rowExists("serving")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("frees the slot so the model can be scheduled on that node again", func() {
|
||||
seedReplica("wedged", "staging", time.Hour)
|
||||
seedJob("wedged", LoadJobStateFailed, time.Minute)
|
||||
|
||||
_, err := registry.NextFreeReplicaIndex(context.Background(), node.ID, "wedged", 1)
|
||||
Expect(err).To(MatchError(ErrNoFreeSlot), "precondition: the stuck row holds the only slot")
|
||||
|
||||
rc.reclaimAbandonedLoads(context.Background())
|
||||
|
||||
idx, err := registry.NextFreeReplicaIndex(context.Background(), node.ID, "wedged", 1)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(idx).To(Equal(0))
|
||||
})
|
||||
})
|
||||
@@ -1159,11 +1159,28 @@ func requireCurrentRevision(tx *gorm.DB, modelName, revision string) error {
|
||||
return err
|
||||
}
|
||||
if state.ConfigRevision != revision {
|
||||
return ErrStaleModelConfigRevision
|
||||
// Name both sides. "stale model config revision" on its own says only
|
||||
// that two hashes differ, which leaves an operator no way to tell an
|
||||
// edited configuration from a revision that is not reproducible for one
|
||||
// unchanged file.
|
||||
return fmt.Errorf("%w (request carries %s, controller holds %s)",
|
||||
ErrStaleModelConfigRevision, shortRevision(revision), shortRevision(state.ConfigRevision))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// shortRevision trims a revision for log and error output. The full value is a
|
||||
// sha256 hex digest; the leading bytes identify it well enough to compare two.
|
||||
func shortRevision(revision string) string {
|
||||
if revision == "" {
|
||||
return "(none)"
|
||||
}
|
||||
if len(revision) > 12 {
|
||||
return revision[:12]
|
||||
}
|
||||
return revision
|
||||
}
|
||||
|
||||
func validateRevisionWrite(modelName, revision string, revisionRequired bool) error {
|
||||
if modelName == "" {
|
||||
return fmt.Errorf("model name is required")
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Stale revision error detail", func() {
|
||||
It("keeps errors.Is matching so callers can still classify it", func() {
|
||||
err := fmt.Errorf("%w (request carries %s, controller holds %s)",
|
||||
ErrStaleModelConfigRevision, shortRevision("aaaabbbbccccdddd"), shortRevision("1111222233334444"))
|
||||
Expect(errors.Is(err, ErrStaleModelConfigRevision)).To(BeTrue())
|
||||
Expect(err.Error()).To(ContainSubstring("stale model config revision"))
|
||||
})
|
||||
|
||||
It("names both revisions so an operator can tell which side moved", func() {
|
||||
Expect(shortRevision("aaaabbbbccccdddd")).To(Equal("aaaabbbbcccc"))
|
||||
Expect(shortRevision("short")).To(Equal("short"))
|
||||
Expect(shortRevision("")).To(Equal("(none)"))
|
||||
})
|
||||
})
|
||||
@@ -1089,37 +1089,47 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID
|
||||
// If freeSlotNodes is empty (everyone full), candidateNodeIDs is whatever
|
||||
// it was — we'll fall through to eviction below.
|
||||
|
||||
var node *BackendNode
|
||||
|
||||
if estimatedVRAM > 0 {
|
||||
if candidateNodeIDs != nil {
|
||||
node, err = r.registry.FindNodeWithVRAMFromSet(ctx, estimatedVRAM, candidateNodeIDs)
|
||||
} else {
|
||||
node, err = r.registry.FindNodeWithVRAM(ctx, estimatedVRAM)
|
||||
}
|
||||
if err != nil {
|
||||
xlog.Warn("No nodes with enough VRAM, falling back to standard scheduling",
|
||||
"required_vram", vram.FormatBytes(estimatedVRAM), "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
if node == nil {
|
||||
if candidateNodeIDs != nil {
|
||||
node, err = r.registry.FindIdleNodeFromSet(ctx, candidateNodeIDs)
|
||||
if err != nil {
|
||||
node, err = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs)
|
||||
// Node choice is wrapped in a liveness check: a node's stored status comes
|
||||
// from its HTTP heartbeat, which is a different channel from the bus that
|
||||
// carries the install. A worker that has died stops answering on the bus at
|
||||
// once but stays healthy in the database until its heartbeat ages out, so
|
||||
// without this the scheduler could commit to a node it cannot reach.
|
||||
selectNode := func() *BackendNode {
|
||||
var candidate *BackendNode
|
||||
var selErr error
|
||||
if estimatedVRAM > 0 {
|
||||
if candidateNodeIDs != nil {
|
||||
candidate, selErr = r.registry.FindNodeWithVRAMFromSet(ctx, estimatedVRAM, candidateNodeIDs)
|
||||
} else {
|
||||
candidate, selErr = r.registry.FindNodeWithVRAM(ctx, estimatedVRAM)
|
||||
}
|
||||
} else {
|
||||
node, err = r.registry.FindIdleNode(ctx)
|
||||
if err != nil {
|
||||
node, err = r.registry.FindLeastLoadedNode(ctx)
|
||||
if selErr != nil {
|
||||
xlog.Warn("No nodes with enough VRAM, falling back to standard scheduling",
|
||||
"required_vram", vram.FormatBytes(estimatedVRAM), "error", selErr)
|
||||
}
|
||||
}
|
||||
|
||||
if candidate == nil {
|
||||
if candidateNodeIDs != nil {
|
||||
candidate, selErr = r.registry.FindIdleNodeFromSet(ctx, candidateNodeIDs)
|
||||
if selErr != nil {
|
||||
candidate, _ = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs)
|
||||
}
|
||||
} else {
|
||||
candidate, selErr = r.registry.FindIdleNode(ctx)
|
||||
if selErr != nil {
|
||||
candidate, _ = r.registry.FindLeastLoadedNode(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
return candidate
|
||||
}
|
||||
|
||||
node := r.pickReachableNode(ctx, selectNode)
|
||||
|
||||
// 4. Preemptive eviction: if no suitable node found, evict the LRU model with zero in-flight
|
||||
if node == nil {
|
||||
evictedNode, evictErr := r.evictLRUAndFreeNode(ctx)
|
||||
evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs)
|
||||
if evictErr != nil {
|
||||
if errors.Is(evictErr, ErrEvictionBusy) {
|
||||
return nil, "", 0, fmt.Errorf("no healthy nodes available: %w", evictErr)
|
||||
@@ -1143,7 +1153,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID
|
||||
// it can race with another concurrent scheduler.
|
||||
xlog.Warn("Chosen node has no free replica slot, evicting LRU",
|
||||
"node", node.Name, "model", modelID, "max_slots", maxSlots)
|
||||
evictedNode, evictErr := r.evictLRUAndFreeNode(ctx)
|
||||
evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs)
|
||||
if evictErr != nil {
|
||||
return nil, "", 0, fmt.Errorf("no replica slot on %s and eviction failed: %w", node.Name, evictErr)
|
||||
}
|
||||
@@ -1969,7 +1979,23 @@ var ErrEvictionBusy = errors.New("all models busy, cannot evict")
|
||||
// Uses SELECT FOR UPDATE inside a transaction to prevent two frontends from
|
||||
// simultaneously picking the same eviction target. The NodeModel row is deleted
|
||||
// inside the transaction; the NATS unload command is sent after commit.
|
||||
// evictLRUAndFreeNode evicts across every healthy node. Callers that hold a
|
||||
// candidate set must use evictLRUAndFreeNodeFrom instead.
|
||||
func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, error) {
|
||||
return r.evictLRUAndFreeNodeFrom(ctx, nil)
|
||||
}
|
||||
|
||||
// evictLRUAndFreeNodeFrom evicts the least-recently-used idle model from one of
|
||||
// candidateNodeIDs, or from any healthy node when the set is nil.
|
||||
//
|
||||
// Restricting eviction to the candidate set matters whenever the model being
|
||||
// scheduled has a node selector. Evicting globally freed a slot on a node the
|
||||
// selector forbids, so the model was then placed there anyway, on hardware it
|
||||
// was explicitly pinned away from, and an unrelated model was dropped to make
|
||||
// the room. On a cluster where the selector-matching node was momentarily
|
||||
// unavailable this repeated, and the evicted model appeared to bounce between
|
||||
// nodes.
|
||||
func (r *SmartRouter) evictLRUAndFreeNodeFrom(ctx context.Context, candidateNodeIDs []string) (*BackendNode, error) {
|
||||
const maxEvictionRetries = 5
|
||||
const evictionRetryInterval = 500 * time.Millisecond
|
||||
|
||||
@@ -1981,7 +2007,7 @@ func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, er
|
||||
var lru NodeModel
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Lock the row so no other frontend can evict the same model
|
||||
if err := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})).
|
||||
q := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})).
|
||||
Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id").
|
||||
Where(`node_models.in_flight = 0 AND node_models.state = ? AND backend_nodes.status = ?
|
||||
AND (
|
||||
@@ -1990,7 +2016,11 @@ func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, er
|
||||
AND (NOT EXISTS (SELECT 1 FROM model_config_states mcs2 WHERE mcs2.model_name = nm2.model_name)
|
||||
OR nm2.config_revision = (SELECT mcs3.config_revision FROM model_config_states mcs3 WHERE mcs3.model_name = nm2.model_name)))
|
||||
> COALESCE((SELECT sc2.min_replicas FROM model_scheduling_configs sc2 WHERE sc2.model_name = node_models.model_name), 1)
|
||||
)`, "loaded", StatusHealthy).
|
||||
)`, "loaded", StatusHealthy)
|
||||
if len(candidateNodeIDs) > 0 {
|
||||
q = q.Where("node_models.node_id IN ?", candidateNodeIDs)
|
||||
}
|
||||
if err := q.
|
||||
Order("node_models.last_used ASC").
|
||||
First(&lru).Error; err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// When no node the selector allows has a free slot, scheduling falls back to
|
||||
// evicting the globally least-recently-used model. That eviction knew nothing
|
||||
// about the selector, so a model pinned to one class of hardware would evict an
|
||||
// unrelated model from a node it is not allowed to run on, and then be placed
|
||||
// there. Two models lose: the pinned one runs on the wrong hardware, and the
|
||||
// evicted one is dropped for nothing and has to reload elsewhere.
|
||||
var _ = Describe("Eviction under a node selector", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
router *SmartRouter
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
router = NewSmartRouter(registry, SmartRouterOptions{DB: db})
|
||||
ctx = context.Background()
|
||||
})
|
||||
|
||||
register := func(name string) *BackendNode {
|
||||
node := &BackendNode{Name: name, NodeType: NodeTypeBackend, Address: name + ":50051"}
|
||||
Expect(registry.Register(ctx, node, true)).To(Succeed())
|
||||
fetched, err := registry.GetByName(ctx, name)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return fetched
|
||||
}
|
||||
|
||||
rowID := 0
|
||||
seed := func(node *BackendNode, model string, idleFor time.Duration, inFlight int) {
|
||||
rowID++
|
||||
Expect(db.Create(&NodeModel{
|
||||
ID: fmt.Sprintf("row-%d", rowID), NodeID: node.ID, ModelName: model,
|
||||
Address: node.Address, State: "loaded", InFlight: inFlight,
|
||||
LastUsed: time.Now().Add(-idleFor), UpdatedAt: time.Now(),
|
||||
}).Error).To(Succeed())
|
||||
}
|
||||
seedLoaded := func(node *BackendNode, model string, idleFor time.Duration) {
|
||||
seed(node, model, idleFor, 0)
|
||||
}
|
||||
|
||||
rowExists := func(model string) bool {
|
||||
var n int64
|
||||
Expect(db.Model(&NodeModel{}).Where("model_name = ?", model).Count(&n).Error).To(Succeed())
|
||||
return n > 0
|
||||
}
|
||||
|
||||
It("does not evict from a node the selector excludes", func() {
|
||||
allowed := register("allowed-node")
|
||||
excluded := register("excluded-node")
|
||||
// The only eviction candidate sits on the excluded node and is the
|
||||
// global LRU, so an unconstrained eviction would take it.
|
||||
seedLoaded(excluded, "innocent-bystander", time.Hour)
|
||||
// In-flight, so it is not an eviction candidate: the allowed node has
|
||||
// nothing that can be freed.
|
||||
seed(allowed, "busy-here", time.Minute, 1)
|
||||
|
||||
_, err := router.evictLRUAndFreeNodeFrom(ctx, []string{allowed.ID})
|
||||
|
||||
Expect(err).To(HaveOccurred(), "no eviction candidate exists on an allowed node")
|
||||
Expect(rowExists("innocent-bystander")).To(BeTrue(),
|
||||
"a model on a node the selector excludes must not be evicted to make room")
|
||||
})
|
||||
|
||||
It("evicts the LRU among the allowed nodes only", func() {
|
||||
allowed := register("allowed-node")
|
||||
excluded := register("excluded-node")
|
||||
seedLoaded(excluded, "older-elsewhere", 2*time.Hour)
|
||||
seedLoaded(allowed, "newer-but-allowed", time.Hour)
|
||||
|
||||
node, err := router.evictLRUAndFreeNodeFrom(ctx, []string{allowed.ID})
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(node.ID).To(Equal(allowed.ID))
|
||||
Expect(rowExists("newer-but-allowed")).To(BeFalse())
|
||||
Expect(rowExists("older-elsewhere")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("keeps evicting globally when the model has no selector", func() {
|
||||
a := register("node-a")
|
||||
seedLoaded(a, "anything", time.Hour)
|
||||
|
||||
node, err := router.evictLRUAndFreeNodeFrom(ctx, nil)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(node.ID).To(Equal(a.ID))
|
||||
Expect(rowExists("anything")).To(BeFalse())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,60 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
"github.com/nats-io/nats.go"
|
||||
)
|
||||
|
||||
// maxNodeLivenessRetries bounds how many unreachable nodes a single scheduling
|
||||
// attempt discards before giving up. Each discarded node is marked unhealthy,
|
||||
// so the bound only has to cover one burst of dead workers rather than the
|
||||
// whole fleet.
|
||||
const maxNodeLivenessRetries = 3
|
||||
|
||||
// nodeAnswersOnBus reports whether a node still has a live subscription.
|
||||
//
|
||||
// Only nats.ErrNoResponders means "absent". Any other outcome, a timeout or a
|
||||
// transport hiccup, leaves the node eligible: wrongly excluding a node that is
|
||||
// merely slow costs real capacity, while the install that follows already
|
||||
// reports its own failure. When no command sender is configured there is no bus
|
||||
// to consult and every node is treated as reachable, which preserves the
|
||||
// behaviour of deployments that do not run one.
|
||||
func (r *SmartRouter) nodeAnswersOnBus(node *BackendNode) bool {
|
||||
if r.unloader == nil || node == nil {
|
||||
return true
|
||||
}
|
||||
err := r.unloader.PingNode(node.ID)
|
||||
return !errors.Is(err, nats.ErrNoResponders)
|
||||
}
|
||||
|
||||
// pickReachableNode calls selectNode until it yields a node that still answers
|
||||
// on the bus, and returns nil when it cannot find one.
|
||||
//
|
||||
// A node that does not answer is marked unhealthy before the next attempt. That
|
||||
// both removes it from the next selection, which queries only healthy nodes,
|
||||
// and tells every other scheduler in the cluster what this one just learned, so
|
||||
// the discovery is not repeated one failed request at a time.
|
||||
func (r *SmartRouter) pickReachableNode(ctx context.Context, selectNode func() *BackendNode) *BackendNode {
|
||||
for range maxNodeLivenessRetries {
|
||||
node := selectNode()
|
||||
if node == nil {
|
||||
return nil
|
||||
}
|
||||
if r.nodeAnswersOnBus(node) {
|
||||
return node
|
||||
}
|
||||
xlog.Warn("Scheduled node is not answering on the bus, marking unhealthy and re-scheduling",
|
||||
"node", node.Name, "nodeID", node.ID)
|
||||
if err := r.registry.MarkUnhealthy(ctx, node.ID); err != nil {
|
||||
// Without the demotion the next selection would hand back the same
|
||||
// node, so stop rather than spin.
|
||||
xlog.Warn("Failed to mark unreachable node unhealthy",
|
||||
"node", node.Name, "nodeID", node.ID, "error", err)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// A node's stored status comes from its HTTP heartbeat, but work is dispatched
|
||||
// over NATS. A worker that dies stops answering on the bus immediately and
|
||||
// keeps its healthy status until the heartbeat ages out, so the scheduler could
|
||||
// commit to a node it could not reach. The request then failed outright with
|
||||
// "no responders available" rather than moving to a node that was actually up.
|
||||
var _ = Describe("Scheduling past a node that left the bus", func() {
|
||||
var (
|
||||
reg *fakeModelRouter
|
||||
fake *fakeUnloader
|
||||
router *SmartRouter
|
||||
)
|
||||
|
||||
newNode := func(id string) *BackendNode {
|
||||
return &BackendNode{ID: id, Name: id, Address: id + ":50051"}
|
||||
}
|
||||
|
||||
// selectorReturning hands back each node in turn, mimicking a scheduler
|
||||
// that re-picks after the previous choice was demoted.
|
||||
selectorReturning := func(nodes ...*BackendNode) func() *BackendNode {
|
||||
i := 0
|
||||
return func() *BackendNode {
|
||||
if i >= len(nodes) {
|
||||
return nil
|
||||
}
|
||||
n := nodes[i]
|
||||
i++
|
||||
return n
|
||||
}
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
reg = &fakeModelRouter{}
|
||||
fake = &fakeUnloader{deadNodes: map[string]bool{}}
|
||||
router = NewSmartRouter(reg, SmartRouterOptions{Unloader: fake})
|
||||
})
|
||||
|
||||
It("passes over a node that no longer answers and takes one that does", func() {
|
||||
dead, alive := newNode("dead-node"), newNode("alive-node")
|
||||
fake.deadNodes["dead-node"] = true
|
||||
|
||||
picked := router.pickReachableNode(context.Background(), selectorReturning(dead, alive))
|
||||
|
||||
Expect(picked).ToNot(BeNil())
|
||||
Expect(picked.ID).To(Equal("alive-node"))
|
||||
Expect(fake.pingCalls).To(Equal([]string{"dead-node", "alive-node"}))
|
||||
})
|
||||
|
||||
It("demotes the absent node so other schedulers stop choosing it", func() {
|
||||
dead, alive := newNode("dead-node"), newNode("alive-node")
|
||||
fake.deadNodes["dead-node"] = true
|
||||
|
||||
router.pickReachableNode(context.Background(), selectorReturning(dead, alive))
|
||||
|
||||
Expect(reg.markedUnhealthy).To(Equal([]string{"dead-node"}))
|
||||
})
|
||||
|
||||
It("takes the first node when it answers, without probing further", func() {
|
||||
first, second := newNode("first"), newNode("second")
|
||||
|
||||
picked := router.pickReachableNode(context.Background(), selectorReturning(first, second))
|
||||
|
||||
Expect(picked.ID).To(Equal("first"))
|
||||
Expect(fake.pingCalls).To(Equal([]string{"first"}))
|
||||
})
|
||||
|
||||
It("gives up rather than spinning when every node is gone", func() {
|
||||
a, b, c, d := newNode("a"), newNode("b"), newNode("c"), newNode("d")
|
||||
for _, id := range []string{"a", "b", "c", "d"} {
|
||||
fake.deadNodes[id] = true
|
||||
}
|
||||
|
||||
picked := router.pickReachableNode(context.Background(), selectorReturning(a, b, c, d))
|
||||
|
||||
Expect(picked).To(BeNil())
|
||||
Expect(len(fake.pingCalls)).To(BeNumerically("<=", maxNodeLivenessRetries))
|
||||
})
|
||||
|
||||
It("stops when the demotion itself fails, so it cannot loop on one node", func() {
|
||||
dead := newNode("dead-node")
|
||||
fake.deadNodes["dead-node"] = true
|
||||
reg.markUnhealthyErr = errors.New("database is down")
|
||||
|
||||
picked := router.pickReachableNode(context.Background(), selectorReturning(dead, dead, dead))
|
||||
|
||||
Expect(picked).To(BeNil())
|
||||
Expect(fake.pingCalls).To(Equal([]string{"dead-node"}))
|
||||
})
|
||||
|
||||
// Only a no-responders answer proves absence. Excluding a node that is
|
||||
// merely slow would cost real capacity.
|
||||
It("keeps a node that answers slowly or errors for another reason", func() {
|
||||
slow := newNode("slow-node")
|
||||
fake.pingErr = errors.New("timeout waiting for reply")
|
||||
|
||||
picked := router.pickReachableNode(context.Background(), selectorReturning(slow))
|
||||
|
||||
Expect(picked).ToNot(BeNil())
|
||||
Expect(picked.ID).To(Equal("slow-node"))
|
||||
Expect(reg.markedUnhealthy).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("treats every node as reachable when no command sender is configured", func() {
|
||||
plain := NewSmartRouter(reg, SmartRouterOptions{})
|
||||
node := newNode("only-node")
|
||||
|
||||
Expect(plain.pickReachableNode(context.Background(), selectorReturning(node))).To(Equal(node))
|
||||
})
|
||||
})
|
||||
@@ -249,8 +249,8 @@ var _ = Describe("revision-bound load publication", func() {
|
||||
LLMConfig: config.LLMConfig{ContextSize: &contextSize},
|
||||
}
|
||||
cfg.Model = "models/full-flow.gguf"
|
||||
expectedRevision, err := config.ModelConfigRevision(&cfg)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(cfg.StampPersistedConfigRevision()).To(Succeed())
|
||||
expectedRevision := cfg.PersistedConfigRevision()
|
||||
|
||||
router := NewSmartRouter(registry, SmartRouterOptions{
|
||||
Unloader: unloader,
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"github.com/mudler/LocalAI/pkg/distributedhdr"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/nats-io/nats.go"
|
||||
ggrpc "google.golang.org/grpc"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"gorm.io/gorm"
|
||||
@@ -61,6 +62,10 @@ func (f *fakeFileStager) ListRemoteDir(_ context.Context, _, _ string) ([]string
|
||||
|
||||
// fakeModelRouter implements ModelRouter with configurable return values.
|
||||
type fakeModelRouter struct {
|
||||
// markedUnhealthy records nodes demoted by the scheduler's liveness check.
|
||||
markedUnhealthy []string
|
||||
markUnhealthyErr error
|
||||
|
||||
fakeLoadJobStore
|
||||
|
||||
// FindAndLockNodeWithModel returns
|
||||
@@ -474,7 +479,15 @@ type fakeUnloader struct {
|
||||
stopCalls []string // "nodeID:model"
|
||||
stopErr error
|
||||
unloadCalls []string
|
||||
unloadErr error
|
||||
|
||||
// deadNodes names the nodes PingNode reports as absent from the bus, and
|
||||
// pingCalls records every node it was asked about, in order.
|
||||
deadNodes map[string]bool
|
||||
pingCalls []string
|
||||
// pingErr is returned for nodes not in deadNodes, so a spec can model a
|
||||
// node that is reachable but answering badly.
|
||||
pingErr error
|
||||
unloadErr error
|
||||
}
|
||||
|
||||
// installCall captures the args we care about when asserting that the
|
||||
@@ -532,6 +545,22 @@ func (f *fakeUnloader) UnloadModelOnNode(nodeID, modelName string) error {
|
||||
return f.unloadErr
|
||||
}
|
||||
|
||||
func (f *fakeModelRouter) MarkUnhealthy(_ context.Context, nodeID string) error {
|
||||
f.markedUnhealthy = append(f.markedUnhealthy, nodeID)
|
||||
return f.markUnhealthyErr
|
||||
}
|
||||
|
||||
func (f *fakeUnloader) PingNode(nodeID string) error {
|
||||
f.mu.Lock()
|
||||
f.pingCalls = append(f.pingCalls, nodeID)
|
||||
dead := f.deadNodes[nodeID]
|
||||
f.mu.Unlock()
|
||||
if dead {
|
||||
return nats.ErrNoResponders
|
||||
}
|
||||
return f.pingErr
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -36,6 +36,10 @@ type NodeCommandSender interface {
|
||||
ListBackends(nodeID string) (*messaging.BackendListReply, error)
|
||||
StopBackend(nodeID, backend string) error
|
||||
UnloadModelOnNode(nodeID, modelName string) error
|
||||
// PingNode reports whether the node is still subscribed on the bus. It
|
||||
// returns nats.ErrNoResponders when nothing answers for the node, which is
|
||||
// the only condition callers may read as "this node cannot be given work".
|
||||
PingNode(nodeID string) error
|
||||
}
|
||||
|
||||
// RemoteUnloaderAdapter implements NodeCommandSender and model.RemoteModelUnloader
|
||||
@@ -360,6 +364,46 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendL
|
||||
return messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply](a.nats, subject, messaging.BackendListRequest{}, 30*time.Second)
|
||||
}
|
||||
|
||||
// PingNode checks that a worker still has a live subscription on the bus.
|
||||
//
|
||||
// A node's status in the database comes from its HTTP heartbeat, which is a
|
||||
// separate channel from NATS. A worker that has died stops answering on NATS
|
||||
// at once but keeps its healthy status until the heartbeat ages out, so the
|
||||
// scheduler could pick a node that could not be given work and the request
|
||||
// failed with "no responders available".
|
||||
//
|
||||
// The subject asked has to be one every worker in the fleet subscribes to, or
|
||||
// this check condemns the workers that do not. models.running was the obvious
|
||||
// choice and the wrong one: it arrived in 4.6, so a 4.5 worker that is alive
|
||||
// and serving never answers it, and a model pinned to that node could never be
|
||||
// scheduled. backend.list has been part of the worker protocol far longer, so
|
||||
// it is the safer question to ask.
|
||||
//
|
||||
// A worker that answers anything is alive. Only when every subject reports no
|
||||
// responders is the node treated as absent, so adding a newer subject here can
|
||||
// never condemn an older worker.
|
||||
func (a *RemoteUnloaderAdapter) PingNode(nodeID string) error {
|
||||
subjects := []string{
|
||||
messaging.SubjectNodeBackendList(nodeID),
|
||||
messaging.SubjectNodeModelsRunning(nodeID),
|
||||
}
|
||||
var lastErr error
|
||||
for _, subject := range subjects {
|
||||
_, err := messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply](
|
||||
a.nats, subject, messaging.BackendListRequest{}, 5*time.Second)
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if !errors.Is(err, nats.ErrNoResponders) {
|
||||
// Reached someone, or failed for a reason that is not absence.
|
||||
// Either way the node is not proven gone.
|
||||
return nil
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
// ListRunningModels asks a worker node which model backend processes it
|
||||
// currently has running, via NATS request-reply.
|
||||
//
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"github.com/nats-io/nats.go"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
)
|
||||
|
||||
// The scheduler's liveness probe asks a worker a question over NATS and treats
|
||||
// "no responders" as proof the worker is gone. That is only sound if every
|
||||
// worker in the fleet subscribes to the subject asked.
|
||||
//
|
||||
// It originally asked models.running, which arrived in 4.6. A 4.5 worker is
|
||||
// perfectly alive and serving, answers backend.list, and never subscribes to
|
||||
// models.running, so the probe condemned it on every scheduling attempt. A
|
||||
// model pinned to such a node could then never be placed at all.
|
||||
var _ = Describe("Node liveness probe subject", func() {
|
||||
var (
|
||||
mc *scriptedMessagingClient
|
||||
adapter *RemoteUnloaderAdapter
|
||||
)
|
||||
|
||||
const nodeID = "11111111-2222-3333-4444-555555555555"
|
||||
|
||||
BeforeEach(func() {
|
||||
mc = newScriptedMessagingClient()
|
||||
adapter = NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute)
|
||||
})
|
||||
|
||||
It("treats a worker that answers backend.list as alive", func() {
|
||||
// A worker old enough to predate models.running: it answers the
|
||||
// long-standing backend.list subject and nothing else.
|
||||
mc.scriptReply(messaging.SubjectNodeBackendList(nodeID), messaging.BackendListReply{})
|
||||
mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID))
|
||||
|
||||
Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeFalse(),
|
||||
"a worker answering backend.list is alive regardless of newer subjects")
|
||||
})
|
||||
|
||||
It("still reports a worker that answers nothing as absent", func() {
|
||||
mc.scriptNoResponders(messaging.SubjectNodeBackendList(nodeID))
|
||||
mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID))
|
||||
|
||||
Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeTrue())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,112 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
const (
|
||||
// defaultEphemeralStagingTTL bounds how long a staged request input can
|
||||
// outlive the request that needed it. Inference reads these files while the
|
||||
// request runs, so the window has to cover a slow multimodal request; it
|
||||
// does not have to cover anything longer.
|
||||
defaultEphemeralStagingTTL = 6 * time.Hour
|
||||
// defaultEphemeralStagingSweep is how often the worker sweeps.
|
||||
defaultEphemeralStagingSweep = 30 * time.Minute
|
||||
)
|
||||
|
||||
// StartEphemeralStagingCleanup sweeps the worker's own staging directory for
|
||||
// per-request input files left behind by finished requests.
|
||||
//
|
||||
// The frontend already expires ephemeral keys from object storage
|
||||
// (services/storage.StartEphemeralCleanup), but a worker receives these files
|
||||
// over the file-transfer server and writes them to its local disk, where
|
||||
// nothing expired them. They accumulated for as long as the worker lived and
|
||||
// eventually filled the volume, at which point every backend start failed
|
||||
// because the process manager could no longer create a state directory.
|
||||
func StartEphemeralStagingCleanup(ctx context.Context, stagingDir string, ttl, interval time.Duration) {
|
||||
if stagingDir == "" {
|
||||
return
|
||||
}
|
||||
if ttl <= 0 {
|
||||
ttl = defaultEphemeralStagingTTL
|
||||
}
|
||||
if interval <= 0 {
|
||||
interval = defaultEphemeralStagingSweep
|
||||
}
|
||||
|
||||
go func() {
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
// Sweep once at startup: a worker that crashed with staged files leaves
|
||||
// them behind, and waiting a full interval to reclaim that space is the
|
||||
// case that hurts on a volume that is already close to full.
|
||||
CleanEphemeralStaging(stagingDir, ttl)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
CleanEphemeralStaging(stagingDir, ttl)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
xlog.Info("Ephemeral staging cleanup started", "dir", stagingDir, "ttl", ttl, "interval", interval)
|
||||
}
|
||||
|
||||
// CleanEphemeralStaging removes staged per-request directories older than ttl.
|
||||
// It only ever descends into <stagingDir>/ephemeral, so staged model weights,
|
||||
// which live alongside it and are not scratch, are never considered.
|
||||
func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
|
||||
root := filepath.Join(stagingDir, "ephemeral")
|
||||
categories, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
// A worker that has never served a file-bearing request has no
|
||||
// ephemeral directory at all. That is the normal case, not a fault.
|
||||
if !os.IsNotExist(err) {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot read staging root", "dir", root, "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
cutoff := time.Now().Add(-ttl)
|
||||
removed := 0
|
||||
for _, category := range categories {
|
||||
if !category.IsDir() {
|
||||
continue
|
||||
}
|
||||
categoryDir := filepath.Join(root, category.Name())
|
||||
entries, err := os.ReadDir(categoryDir)
|
||||
if err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot read category", "dir", categoryDir, "error", err)
|
||||
continue
|
||||
}
|
||||
for _, entry := range entries {
|
||||
path := filepath.Join(categoryDir, entry.Name())
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot stat entry", "path", path, "error", err)
|
||||
continue
|
||||
}
|
||||
// A request rewrites nothing after staging, so the entry's own
|
||||
// modification time is when its request was served.
|
||||
if !info.ModTime().Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if err := os.RemoveAll(path); err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", path, "error", err)
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
}
|
||||
}
|
||||
|
||||
if removed > 0 {
|
||||
xlog.Info("Ephemeral staging cleanup removed stale request files", "count", removed, "dir", root)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Worker ephemeral staging cleanup", func() {
|
||||
var stagingDir string
|
||||
|
||||
// mkEphemeral creates one staged request directory holding a file, and
|
||||
// backdates both so the sweeper sees it as `age` old.
|
||||
mkEphemeral := func(id string, age time.Duration) string {
|
||||
dir := filepath.Join(stagingDir, "ephemeral", "inputs", id)
|
||||
Expect(os.MkdirAll(dir, 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(filepath.Join(dir, "payload.bin"), []byte("x"), 0o600)).To(Succeed())
|
||||
stamp := time.Now().Add(-age)
|
||||
Expect(os.Chtimes(filepath.Join(dir, "payload.bin"), stamp, stamp)).To(Succeed())
|
||||
Expect(os.Chtimes(dir, stamp, stamp)).To(Succeed())
|
||||
return dir
|
||||
}
|
||||
|
||||
BeforeEach(func() { stagingDir = GinkgoT().TempDir() })
|
||||
|
||||
It("removes staged request directories older than the TTL", func() {
|
||||
old := mkEphemeral("aaaa1111", 48*time.Hour)
|
||||
CleanEphemeralStaging(stagingDir, time.Hour)
|
||||
Expect(old).ToNot(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("keeps directories a running request may still be reading", func() {
|
||||
fresh := mkEphemeral("bbbb2222", 5*time.Minute)
|
||||
CleanEphemeralStaging(stagingDir, time.Hour)
|
||||
Expect(fresh).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("leaves staged models and everything outside ephemeral alone", func() {
|
||||
modelDir := filepath.Join(stagingDir, "models", "some-model")
|
||||
Expect(os.MkdirAll(modelDir, 0o750)).To(Succeed())
|
||||
weights := filepath.Join(modelDir, "weights.gguf")
|
||||
Expect(os.WriteFile(weights, []byte("w"), 0o600)).To(Succeed())
|
||||
stamp := time.Now().Add(-90 * 24 * time.Hour)
|
||||
Expect(os.Chtimes(weights, stamp, stamp)).To(Succeed())
|
||||
Expect(os.Chtimes(modelDir, stamp, stamp)).To(Succeed())
|
||||
|
||||
CleanEphemeralStaging(stagingDir, time.Hour)
|
||||
|
||||
Expect(weights).To(BeAnExistingFile(), "a staged model is not ephemeral scratch")
|
||||
})
|
||||
|
||||
It("does nothing when no ephemeral directory exists", func() {
|
||||
Expect(func() { CleanEphemeralStaging(stagingDir, time.Hour) }).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
@@ -159,6 +159,10 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
return fmt.Errorf("starting HTTP file transfer server: %w", err)
|
||||
}
|
||||
|
||||
// Per-request input files land in stagingDir over that server and nothing
|
||||
// used to remove them, so a long-lived worker filled its own disk.
|
||||
StartEphemeralStagingCleanup(shutdownCtx, stagingDir, 0, 0)
|
||||
|
||||
// Connect to NATS
|
||||
xlog.Info("Connecting to NATS", "url", sanitize.URL(cfg.NatsURL))
|
||||
natsClient, err := connectNats()
|
||||
|
||||
@@ -460,8 +460,16 @@ context lengths, so you can see whether something will run before installing it.
|
||||
Working that out means reading the metadata of a model's weight files, which for
|
||||
a model you have not installed is a request to the host that serves them. It
|
||||
takes a second or two the first time, and the gallery needs one per row. LocalAI
|
||||
caches the result, and warms that cache in the background at startup so the
|
||||
gallery reads instantly rather than filling in its own numbers while you watch.
|
||||
caches successful remote probes for 24 hours under the LocalAI data directory,
|
||||
and warms that cache in the background at startup so the gallery reads instantly
|
||||
rather than filling in its own numbers while you watch. The on-disk cache is
|
||||
reused after a restart, so frequent restarts do not download the same metadata
|
||||
again. Local model files are always inspected directly. The cache keeps at most
|
||||
4,096 entries and removes the oldest entries when it reaches that limit.
|
||||
Disable **Persist remote VRAM estimates** under **Settings > Galleries**, or set
|
||||
`LOCALAI_VRAM_PERSISTENT_CACHE=false`, to keep estimates in memory only. Setting
|
||||
`LOCALAI_AUTOLOAD_GALLERIES=false` also disables the startup warmer and the
|
||||
persistent cache.
|
||||
|
||||
The same warm-up also describes each entry's **variants** - the alternative
|
||||
builds of the same weights that the picker offers - because that costs the same
|
||||
|
||||
@@ -1020,6 +1020,31 @@ Notes:
|
||||
- Upgrade the worker when it does not support the exact model-stop request.
|
||||
- Stop and restart the stale backend only as an operational recovery action. LocalAI keeps it non-routable while durable cleanup is pending.
|
||||
|
||||
**A model cannot be scheduled on a node that looks free (`no replica slot ... all models busy, cannot evict`):**
|
||||
- A replica row in `staging` or `loading` holds its slot: slot allocation counts every state except `unloading`. If a worker drops out mid-transfer, that row never reaches `loaded`, and eviction only ever considers `loaded` replicas, so on a node with one replica slot per model the model became unschedulable there.
|
||||
- The reconciler now reclaims a replica row stuck before serving when no load job is still driving it, and the freed slot is immediately reusable.
|
||||
- Liveness is decided by the load job's progress heartbeat, not by elapsed time. Staging a large checkpoint legitimately runs for a long time without touching the replica row, so a transfer that is still progressing is never reclaimed however long it takes.
|
||||
- `Reconciler: reclaimed a replica slot held by a load nobody is driving` names each row reclaimed this way.
|
||||
|
||||
**A request fails with `nats: no responders available for request`:**
|
||||
- The chosen worker was not subscribed on the bus when the frontend tried to install the backend on it. A node's status comes from its HTTP heartbeat, which is a separate channel: a worker that stops stays `healthy` until that heartbeat ages out.
|
||||
- The scheduler now checks that a node still answers on the bus before it commits to it, marks one that does not as unhealthy, and picks another. A request should therefore see this only when no reachable node is left.
|
||||
- Only a no-responders answer counts as absent. A worker that answers slowly stays eligible, because excluding it would cost capacity that is really there.
|
||||
- Check the worker process is running and its NATS connection is up. `Scheduled node is not answering on the bus` in the frontend log names each node demoted this way.
|
||||
|
||||
**A worker fills its own disk over time:**
|
||||
- A request that carries a file (an image, an audio clip, a video) stages that file to the worker under `<models>/../staging/ephemeral/`. The worker deletes these 6 hours after the request that needed them, and sweeps every 30 minutes plus once at startup, so a worker that crashed mid-request still reclaims the space.
|
||||
- Releases before this sweep existed kept every staged input for the lifetime of the worker. Delete `<models>/../staging/ephemeral/` on an affected worker once, as the user the worker runs as; the sweep keeps it bounded from then on.
|
||||
- Staged **model** files are not touched by this. They live beside the ephemeral directory and are not per-request scratch.
|
||||
- A worker whose volume is genuinely full reports `creating backend process state directory under ...: no space left on device` when a backend starts.
|
||||
|
||||
**Requests fail with `stale model config revision` although nobody edited the model:**
|
||||
- A model's stored revision must describe its persisted configuration. Releases before this fix also hashed the per-request prediction parameters, so the first request after a restart pinned the revision to its own `temperature`, `top_p`, `stop` and similar values. Every later request that sent different values was then rejected.
|
||||
- Upgrade the frontend replicas first. After the upgrade the revision is stamped when the configuration is loaded, so it no longer depends on the request body.
|
||||
- Each frontend now reconciles the stored revisions against the configuration on disk at startup, and republishes any that disagree, so a drifted revision heals on the next restart. Only models that actually drifted are republished, because republishing quarantines the replicas loaded under the old revision.
|
||||
- A model that has never been served has no stored revision and is left alone; its first request establishes one.
|
||||
- On a release without that reconciliation, clear the row once per affected model so the next request establishes the correct revision: `DELETE FROM model_config_states WHERE model_name = '<model>';` Saving any edit through the API or the WebUI has the same effect.
|
||||
|
||||
**Port conflicts on workers:**
|
||||
- Each model gets its own gRPC process on an incrementing port (50051, 50052, ...)
|
||||
- The HTTP file transfer server runs on the base port - 1 (default: 50050)
|
||||
|
||||
@@ -82,7 +82,7 @@ Manage model and backend galleries:
|
||||
|
||||
- **Model Galleries**: JSON array of gallery objects with `url` and `name` fields, plus an optional `mirrors` list of fallback URLs (see [Gallery mirrors]({{%relref "features/model-gallery#gallery-mirrors" %}}))
|
||||
- **Backend Galleries**: JSON array of backend gallery objects, which accept the same `mirrors` key
|
||||
- **Autoload Galleries**: Automatically load model galleries on startup
|
||||
- **Load and pre-warm galleries on boot**: Load model galleries and pre-warm their remote size and VRAM estimates when LocalAI starts. Disable this setting to skip both startup operations.
|
||||
- **Autoload Backend Galleries**: Automatically load backend galleries on startup
|
||||
|
||||
### Agent Pool Settings
|
||||
@@ -164,6 +164,7 @@ The `runtime_settings.json` file follows this structure:
|
||||
],
|
||||
"autoload_galleries": true,
|
||||
"autoload_backend_galleries": true,
|
||||
"vram_persistent_cache": true,
|
||||
"api_keys": []
|
||||
}
|
||||
```
|
||||
|
||||
@@ -61,6 +61,7 @@ For more information on VRAM management, see [VRAM and Memory Management]({{%rel
|
||||
|-----------|---------|-------------|----------------------|
|
||||
| `--galleries` | | JSON list of galleries | `$LOCALAI_GALLERIES`, `$GALLERIES` |
|
||||
| `--autoload-galleries` | `true` | Automatically load galleries on startup | `$LOCALAI_AUTOLOAD_GALLERIES`, `$AUTOLOAD_GALLERIES` |
|
||||
| `--vram-persistent-cache` | `true` | Persist successful remote VRAM metadata probes across restarts | `$LOCALAI_VRAM_PERSISTENT_CACHE`, `$VRAM_PERSISTENT_CACHE` |
|
||||
| `--preload-models` | | A list of models to apply in JSON at start | `$LOCALAI_PRELOAD_MODELS`, `$PRELOAD_MODELS` |
|
||||
| `--models` | | A list of model configuration URLs to load | `$LOCALAI_MODELS`, `$MODELS` |
|
||||
| `--preload-models-config` | | A list of models to apply at startup. Path to a YAML config file | `$LOCALAI_PRELOAD_MODELS_CONFIG`, `$PRELOAD_MODELS_CONFIG` |
|
||||
|
||||
+1076
-8
File diff suppressed because it is too large.
Load diff
@@ -38,7 +38,7 @@ require (
|
||||
github.com/modelcontextprotocol/go-sdk v1.5.0
|
||||
github.com/mudler/cogito v0.11.1-0.20260721122412-6eece18a6bb6
|
||||
github.com/mudler/edgevpn v0.34.0
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260720195933-3d64f5c974fc
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260823202314-dfa0ed852db6
|
||||
github.com/mudler/memory v0.0.0-20260406210934-424c1ecf2cf8
|
||||
github.com/mudler/nib v0.6.0
|
||||
github.com/mudler/xlog v0.0.6
|
||||
|
||||
@@ -1027,6 +1027,8 @@ github.com/mudler/go-piper v0.0.0-20241023091659-2494246fd9fc h1:RxwneJl1VgvikiX
|
||||
github.com/mudler/go-piper v0.0.0-20241023091659-2494246fd9fc/go.mod h1:O7SwdSWMilAWhBZMK9N9Y/oBDyMMzshE3ju8Xkexwig=
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260720195933-3d64f5c974fc h1:NEFmd7+JoImN5dZI81/vcBRjtMg+GfEa7nEjQ029hd0=
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260720195933-3d64f5c974fc/go.mod h1:h6kmHUZeafr+k5hRYpGLMzJFH4hItHffgpRo2QIkP+o=
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260823202314-dfa0ed852db6 h1:/nFm1Ttf8g1BnWtEth986JR34pCh9rzae5A2vKBZosc=
|
||||
github.com/mudler/go-processmanager v0.1.2-0.20260823202314-dfa0ed852db6/go.mod h1:h6kmHUZeafr+k5hRYpGLMzJFH4hItHffgpRo2QIkP+o=
|
||||
github.com/mudler/localrecall v0.6.3 h1:uXOrP9JmetzxgVKzSrawviyBHZfAcvPBBIrvVUdZjDA=
|
||||
github.com/mudler/localrecall v0.6.3/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU=
|
||||
github.com/mudler/memory v0.0.0-20260406210934-424c1ecf2cf8 h1:Ry8RiWy8fZ6Ff4E7dPmjRsBrnHOnPeOOj2LhCgyjQu0=
|
||||
|
||||
+22
-1
@@ -231,6 +231,16 @@ func (ml *ModelLoader) StartProcess(grpcProcess, id string, serverAddress string
|
||||
return ml.startProcess(grpcProcess, id, serverAddress, args...)
|
||||
}
|
||||
|
||||
// newProcessStateDir creates the directory a backend process uses for its pid,
|
||||
// state and log files, and reports why when it cannot.
|
||||
func newProcessStateDir() (string, error) {
|
||||
dir, err := os.MkdirTemp(os.TempDir(), "go-processmanager")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("creating backend process state directory under %s: %w", os.TempDir(), err)
|
||||
}
|
||||
return dir, nil
|
||||
}
|
||||
|
||||
func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string, args ...string) (*process.Process, error) {
|
||||
// Make sure the process is executable
|
||||
// Check first if it has executable permissions
|
||||
@@ -261,8 +271,19 @@ func (ml *ModelLoader) startProcess(grpcProcess, id string, serverAddress string
|
||||
// and the GPU would silently fall back to CPU). No-op for other backends.
|
||||
env = append(env, vulkanICDEnv(workDir)...)
|
||||
|
||||
// Resolve the state directory here rather than through
|
||||
// process.WithTemporaryStateDir(). process.New applies its options but
|
||||
// discards the error they return, so a temp directory that cannot be
|
||||
// created leaves StateDir empty and every later option unapplied. Run()
|
||||
// then reported "mkdir : no such file or directory" with no path, hiding
|
||||
// the real cause (a full volume, or a TMPDIR that no longer resolves).
|
||||
stateDir, err := newProcessStateDir()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
grpcControlProcess := process.New(
|
||||
process.WithTemporaryStateDir(),
|
||||
process.WithStateDir(stateDir),
|
||||
process.WithName(filepath.Base(grpcProcess)),
|
||||
process.WithArgs(append(args, []string{"--addr", serverAddress}...)...),
|
||||
process.WithEnvironment(env...),
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Backend process state directory", func() {
|
||||
It("reports why the state directory could not be created", func() {
|
||||
// A worker whose volume is full, or whose TMPDIR no longer resolves,
|
||||
// cannot get a state directory. go-processmanager's New() drops the
|
||||
// option error, leaving StateDir empty, and Run() then failed with
|
||||
// "mkdir : no such file or directory" naming no path at all. Resolving
|
||||
// the directory here keeps the real cause attached.
|
||||
GinkgoT().Setenv("TMPDIR", filepath.Join(GinkgoT().TempDir(), "does-not-exist"))
|
||||
|
||||
dir, err := newProcessStateDir()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(dir).To(BeEmpty())
|
||||
Expect(err.Error()).To(ContainSubstring("backend process state directory"))
|
||||
Expect(err.Error()).To(ContainSubstring("does-not-exist"),
|
||||
"the error must name the directory it could not create")
|
||||
})
|
||||
|
||||
It("returns a usable directory when the temp location works", func() {
|
||||
GinkgoT().Setenv("TMPDIR", GinkgoT().TempDir())
|
||||
|
||||
dir, err := newProcessStateDir()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(dir).ToNot(BeEmpty())
|
||||
info, statErr := os.Stat(dir)
|
||||
Expect(statErr).ToNot(HaveOccurred())
|
||||
Expect(info.IsDir()).To(BeTrue())
|
||||
})
|
||||
})
|
||||
+183
-24
@@ -80,27 +80,133 @@ var layerRetryBackoff = func(attempt int) time.Duration {
|
||||
return d
|
||||
}
|
||||
|
||||
// blobRangeOpener re-opens a layer blob at a byte offset. It returns the
|
||||
// stream and the offset it actually starts at: the requested offset when the
|
||||
// server honoured the Range request, or 0 when it ignored it and is sending
|
||||
// the blob from the first byte again.
|
||||
type blobRangeOpener func(ctx context.Context, offset int64) (io.ReadCloser, int64, error)
|
||||
|
||||
// newBlobRangeOpener returns a blobRangeOpener that re-fetches the layer's
|
||||
// blob from its registry with an HTTP Range request. Registries like quay.io
|
||||
// redirect blob downloads to pre-signed S3/CDN URLs that expire after ~10
|
||||
// minutes; on a slow connection a multi-GiB layer cannot finish inside that
|
||||
// window, so restarting from byte zero can never succeed while resuming from
|
||||
// the current offset can (docker pull survives the same expiry this way).
|
||||
// Each call goes back to the registry, so it obtains a fresh redirect URL and
|
||||
// a fresh auth token. Returns nil when imageRef does not name a registry blob
|
||||
// (e.g. local tarballs), which disables resuming. See issue #10577.
|
||||
func newBlobRangeOpener(imageRef string, layer v1.Layer, auth *registrytypes.AuthConfig, base http.RoundTripper) blobRangeOpener {
|
||||
ref, err := name.ParseReference(imageRef)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
digest, err := layer.Digest()
|
||||
if err != nil || digest.Hex == "" {
|
||||
return nil
|
||||
}
|
||||
repo := ref.Context()
|
||||
if base == nil {
|
||||
base = http.DefaultTransport
|
||||
}
|
||||
var authenticator authn.Authenticator
|
||||
if auth != nil {
|
||||
authenticator = staticAuth{auth}
|
||||
} else if authenticator, err = authn.DefaultKeychain.Resolve(repo.Registry); err != nil {
|
||||
authenticator = authn.Anonymous
|
||||
}
|
||||
blobURL := fmt.Sprintf("%s://%s/v2/%s/blobs/%s", repo.Registry.Scheme(), repo.RegistryStr(), repo.RepositoryStr(), digest.String())
|
||||
|
||||
return func(ctx context.Context, offset int64) (io.ReadCloser, int64, error) {
|
||||
tr, err := transport.NewWithContext(ctx, repo.Registry, authenticator, base, []string{repo.Scope(transport.PullScope)})
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, blobURL, nil)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
if offset > 0 {
|
||||
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", offset))
|
||||
}
|
||||
req.Header.Set("User-Agent", UserAgent())
|
||||
resp, err := (&http.Client{Transport: tr}).Do(req)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
switch resp.StatusCode {
|
||||
case http.StatusPartialContent:
|
||||
return resp.Body, offset, nil
|
||||
case http.StatusOK:
|
||||
return resp.Body, 0, nil
|
||||
default:
|
||||
_ = resp.Body.Close()
|
||||
return nil, 0, fmt.Errorf("unexpected status %d resuming blob %s", resp.StatusCode, digest.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// verifyLayerFile proves the assembled layer file matches the digest the
|
||||
// registry advertised. A resumed download splices bytes from independent HTTP
|
||||
// responses and bypasses the verified reader layer.Compressed() provides, so
|
||||
// the whole file must be re-checked before it is trusted.
|
||||
func verifyLayerFile(layer v1.Layer, f *os.File) error {
|
||||
digest, err := layer.Digest()
|
||||
if err != nil || digest.Hex == "" || digest.Algorithm != "sha256" {
|
||||
return nil
|
||||
}
|
||||
if _, err := f.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
got, _, err := v1.SHA256(f)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if got.Hex != digest.Hex {
|
||||
return fmt.Errorf("resumed layer digest mismatch: got %s, want %s", got, digest)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadLayerToFile streams a single compressed layer into dst, retrying on
|
||||
// transient network errors (unexpected EOF, connection reset, ...). Large
|
||||
// backend images (e.g. vLLM) are several GiB and a single dropped connection
|
||||
// mid-stream previously failed the whole install with "unexpected EOF" and no
|
||||
// recovery. The registry transport already retries manifest fetches via
|
||||
// defaultRetryPredicate (see GetImage/GetImageDigest); this extends the same
|
||||
// behaviour to the layer data stream. See issue #10577.
|
||||
func downloadLayerToFile(ctx context.Context, layer v1.Layer, dst *os.File, progress *progressWriter) error {
|
||||
// recovery. When resume is non-nil, a retry keeps the bytes already on disk
|
||||
// and continues from that offset instead of starting over: registries that
|
||||
// serve blobs through expiring pre-signed URLs (quay.io + S3/Akamai) cut off
|
||||
// every full-length transfer on slow connections, so restarting can never
|
||||
// finish while resuming makes progress each round. The retry budget only
|
||||
// counts attempts that made no forward progress, so a download that keeps
|
||||
// advancing keeps going. See issue #10577.
|
||||
func downloadLayerToFile(ctx context.Context, layer v1.Layer, dst *os.File, progress *progressWriter, resume blobRangeOpener) error {
|
||||
var lastErr error
|
||||
// written tracks the valid bytes currently in dst across attempts, and
|
||||
// bestWritten the furthest offset any attempt has reached: only beating
|
||||
// it counts as forward progress for the retry budget, so a server that
|
||||
// ignores Range requests and keeps dropping mid-stream still runs out
|
||||
// of attempts instead of looping forever.
|
||||
var written, bestWritten int64
|
||||
// resumed records whether any byte in dst came from a resumed raw blob
|
||||
// fetch, which requires re-verifying the assembled file at the end.
|
||||
resumed := false
|
||||
|
||||
truncate := func() error {
|
||||
if _, err := dst.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := dst.Truncate(0); err != nil {
|
||||
return err
|
||||
}
|
||||
written = 0
|
||||
resumed = false
|
||||
if progress != nil {
|
||||
progress.written = 0
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
for attempt := 0; attempt <= layerDownloadRetries; attempt++ {
|
||||
if attempt > 0 {
|
||||
// Discard any partial data from the previous failed attempt.
|
||||
if _, err := dst.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := dst.Truncate(0); err != nil {
|
||||
return err
|
||||
}
|
||||
if progress != nil {
|
||||
progress.written = 0
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
@@ -108,19 +214,69 @@ func downloadLayerToFile(ctx context.Context, layer v1.Layer, dst *os.File, prog
|
||||
}
|
||||
}
|
||||
|
||||
var w io.Writer = dst
|
||||
if progress != nil {
|
||||
w = io.MultiWriter(dst, progress)
|
||||
var reader io.ReadCloser
|
||||
if attempt > 0 && resume != nil && written > 0 {
|
||||
r, offset, rerr := resume(ctx, written)
|
||||
switch {
|
||||
case rerr != nil:
|
||||
// Keep the partial bytes: opening the resume stream can
|
||||
// fail transiently (token refresh, connection refused)
|
||||
// and the next attempt can still continue from here.
|
||||
lastErr = rerr
|
||||
case offset != written:
|
||||
// The server ignored the Range request and is sending
|
||||
// the blob from the first byte: drop the partial data.
|
||||
if err := truncate(); err != nil {
|
||||
_ = r.Close()
|
||||
return err
|
||||
}
|
||||
reader = r
|
||||
resumed = true
|
||||
default:
|
||||
reader = r
|
||||
resumed = true
|
||||
}
|
||||
} else {
|
||||
// First attempt, or no way to resume: restart from scratch
|
||||
// through the digest-verifying layer reader.
|
||||
if err := truncate(); err != nil {
|
||||
return err
|
||||
}
|
||||
reader, lastErr = layer.Compressed()
|
||||
}
|
||||
|
||||
var reader io.ReadCloser
|
||||
reader, lastErr = layer.Compressed()
|
||||
if lastErr == nil {
|
||||
_, lastErr = xio.Copy(ctx, w, reader)
|
||||
if reader != nil {
|
||||
var w io.Writer = dst
|
||||
if progress != nil {
|
||||
w = io.MultiWriter(dst, progress)
|
||||
}
|
||||
var n int64
|
||||
n, lastErr = xio.Copy(ctx, w, reader)
|
||||
written += n
|
||||
_ = reader.Close()
|
||||
if written > bestWritten {
|
||||
// Forward progress: don't charge this round against the
|
||||
// retry budget, or slow links would still exhaust it.
|
||||
bestWritten = written
|
||||
attempt = 0
|
||||
}
|
||||
}
|
||||
|
||||
if lastErr == nil {
|
||||
return nil
|
||||
if !resumed {
|
||||
return nil
|
||||
}
|
||||
verr := verifyLayerFile(layer, dst)
|
||||
if verr == nil {
|
||||
return nil
|
||||
}
|
||||
// The spliced file is corrupt: discard it and retry cleanly.
|
||||
logs.Warn.Printf("discarding resumed layer download: %v", verr)
|
||||
lastErr = verr
|
||||
if err := truncate(); err != nil {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Stop early on context cancellation or non-retryable errors.
|
||||
@@ -382,8 +538,11 @@ func DownloadOCIImageTar(ctx context.Context, img v1.Image, imageRef string, tar
|
||||
}
|
||||
}
|
||||
|
||||
// Download the compressed layer, retrying on transient network errors.
|
||||
err = downloadLayerToFile(ctx, layer, file, progress)
|
||||
// Download the compressed layer, retrying on transient network
|
||||
// errors and resuming from the last byte received where possible.
|
||||
// Anonymous/default-keychain credentials match what GetImage uses
|
||||
// for every in-tree caller (they all pass a nil auth).
|
||||
err = downloadLayerToFile(ctx, layer, file, progress, newBlobRangeOpener(imageRef, layer, nil, nil))
|
||||
file.Close()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to download layer %d: %v", i, err)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
package oci
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-containerregistry/pkg/name"
|
||||
"github.com/google/go-containerregistry/pkg/registry"
|
||||
"github.com/google/go-containerregistry/pkg/v1/random"
|
||||
"github.com/google/go-containerregistry/pkg/v1/remote"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// droppingBlobRegistry emulates how quay.io serves layer blobs from S3/Akamai
|
||||
// with a short-lived pre-signed URL: a full-blob GET on a slow connection is
|
||||
// always cut off mid-transfer, so a client that restarts from byte zero can
|
||||
// never complete the download. Only a client that resumes with a Range request
|
||||
// (like docker pull does) receives the remaining bytes and can finish.
|
||||
type droppingBlobRegistry struct {
|
||||
inner http.Handler
|
||||
|
||||
mu sync.Mutex
|
||||
rangeRequests []int64
|
||||
fullRequests int
|
||||
}
|
||||
|
||||
// dropThreshold separates real layer blobs from small metadata blobs (image
|
||||
// config), which are served untouched.
|
||||
const dropThreshold = 1024
|
||||
|
||||
func (h *droppingBlobRegistry) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet || !strings.Contains(r.URL.Path, "/blobs/sha256:") {
|
||||
h.inner.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
// Fetch the full blob from the inner registry (which does not speak
|
||||
// Range) and apply the Range semantics here.
|
||||
inner := r.Clone(r.Context())
|
||||
inner.Header.Del("Range")
|
||||
rec := httptest.NewRecorder()
|
||||
h.inner.ServeHTTP(rec, inner)
|
||||
body := rec.Body.Bytes()
|
||||
if rec.Code != http.StatusOK || len(body) <= dropThreshold {
|
||||
for k, vv := range rec.Header() {
|
||||
for _, v := range vv {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(rec.Code)
|
||||
_, _ = w.Write(body)
|
||||
return
|
||||
}
|
||||
|
||||
if rh := r.Header.Get("Range"); rh != "" {
|
||||
offset, err := strconv.ParseInt(strings.TrimSuffix(strings.TrimPrefix(rh, "bytes="), "-"), 10, 64)
|
||||
if err != nil || offset < 0 || offset >= int64(len(body)) {
|
||||
w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.rangeRequests = append(h.rangeRequests, offset)
|
||||
h.mu.Unlock()
|
||||
w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", offset, len(body)-1, len(body)))
|
||||
w.Header().Set("Content-Length", strconv.Itoa(len(body)-int(offset)))
|
||||
w.WriteHeader(http.StatusPartialContent)
|
||||
_, _ = w.Write(body[offset:])
|
||||
return
|
||||
}
|
||||
|
||||
h.mu.Lock()
|
||||
h.fullRequests++
|
||||
h.mu.Unlock()
|
||||
|
||||
// Announce the full size but deliver only half, then sever the
|
||||
// connection, like a pre-signed URL expiring mid-download.
|
||||
w.Header().Set("Content-Length", strconv.Itoa(len(body)))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(body[:len(body)/2])
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
panic(http.ErrAbortHandler)
|
||||
}
|
||||
|
||||
var _ = Describe("DownloadOCIImageTar resume", func() {
|
||||
var (
|
||||
server *httptest.Server
|
||||
reg *droppingBlobRegistry
|
||||
tmpDir string
|
||||
restoreWait func()
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
reg = &droppingBlobRegistry{inner: registry.New()}
|
||||
server = httptest.NewServer(reg)
|
||||
|
||||
var err error
|
||||
tmpDir, err = os.MkdirTemp("", "oci-resume-e2e-*")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
prev := layerRetryBackoff
|
||||
layerRetryBackoff = func(int) time.Duration { return 0 }
|
||||
restoreWait = func() { layerRetryBackoff = prev }
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
restoreWait()
|
||||
server.Close()
|
||||
_ = os.RemoveAll(tmpDir)
|
||||
})
|
||||
|
||||
It("completes the download by resuming interrupted layer transfers with Range requests", func() {
|
||||
img, err := random.Image(4096, 1)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
imageRef := strings.TrimPrefix(server.URL, "http://") + "/testrepo/backend:latest"
|
||||
ref, err := name.ParseReference(imageRef)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(remote.Write(ref, img)).To(Succeed())
|
||||
|
||||
pulled, err := GetImage(imageRef, "", nil, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
tarPath := filepath.Join(tmpDir, "image.tar")
|
||||
err = DownloadOCIImageTar(context.Background(), pulled, imageRef, tarPath, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
// The full-blob attempt was cut off, so success is only possible
|
||||
// through at least one Range request picking up where it stopped.
|
||||
reg.mu.Lock()
|
||||
defer reg.mu.Unlock()
|
||||
Expect(reg.rangeRequests).NotTo(BeEmpty())
|
||||
for _, off := range reg.rangeRequests {
|
||||
Expect(off).To(BeNumerically(">", 0))
|
||||
}
|
||||
|
||||
fi, err := os.Stat(tarPath)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(fi.Size()).To(BeNumerically(">", 0))
|
||||
})
|
||||
})
|
||||
@@ -33,14 +33,18 @@ func (r *failingReader) Read(p []byte) (int, error) {
|
||||
|
||||
// fakeLayer is a minimal v1.Layer whose Compressed() fails failUntil times with
|
||||
// err (after emitting a partial prefix) before finally returning data in full.
|
||||
// The failing attempts emit prefix when set, or placeholder garbage otherwise.
|
||||
// digest, when set, is what Digest() reports.
|
||||
type fakeLayer struct {
|
||||
data []byte
|
||||
prefix []byte
|
||||
digest v1.Hash
|
||||
failUntil int
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
func (f *fakeLayer) Digest() (v1.Hash, error) { return v1.Hash{}, nil }
|
||||
func (f *fakeLayer) Digest() (v1.Hash, error) { return f.digest, nil }
|
||||
func (f *fakeLayer) DiffID() (v1.Hash, error) { return v1.Hash{}, nil }
|
||||
func (f *fakeLayer) Size() (int64, error) { return int64(len(f.data)), nil }
|
||||
func (f *fakeLayer) MediaType() (types.MediaType, error) { return types.DockerLayer, nil }
|
||||
@@ -51,7 +55,11 @@ func (f *fakeLayer) Uncompressed() (io.ReadCloser, error) {
|
||||
func (f *fakeLayer) Compressed() (io.ReadCloser, error) {
|
||||
f.calls++
|
||||
if f.calls <= f.failUntil {
|
||||
return io.NopCloser(&failingReader{prefix: []byte("partial-garbage"), err: f.err}), nil
|
||||
prefix := f.prefix
|
||||
if prefix == nil {
|
||||
prefix = []byte("partial-garbage")
|
||||
}
|
||||
return io.NopCloser(&failingReader{prefix: prefix, err: f.err}), nil
|
||||
}
|
||||
return io.NopCloser(bytes.NewReader(f.data)), nil
|
||||
}
|
||||
@@ -86,7 +94,7 @@ var _ = Describe("downloadLayerToFile", func() {
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil)
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(layer.calls).To(Equal(3))
|
||||
|
||||
@@ -104,7 +112,7 @@ var _ = Describe("downloadLayerToFile", func() {
|
||||
err: errors.New("permission denied"),
|
||||
}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil)
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(layer.calls).To(Equal(1))
|
||||
})
|
||||
@@ -116,7 +124,7 @@ var _ = Describe("downloadLayerToFile", func() {
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil)
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, nil)
|
||||
Expect(err).To(MatchError(io.ErrUnexpectedEOF))
|
||||
Expect(layer.calls).To(Equal(layerDownloadRetries + 1))
|
||||
})
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
package oci
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
v1 "github.com/google/go-containerregistry/pkg/v1"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// recordingOpener is a test blobRangeOpener that records the offsets it was
|
||||
// asked to resume from and delegates the stream to open.
|
||||
type recordingOpener struct {
|
||||
offsets []int64
|
||||
open func(offset int64) (io.ReadCloser, int64, error)
|
||||
}
|
||||
|
||||
func (o *recordingOpener) opener() blobRangeOpener {
|
||||
return func(_ context.Context, offset int64) (io.ReadCloser, int64, error) {
|
||||
o.offsets = append(o.offsets, offset)
|
||||
return o.open(offset)
|
||||
}
|
||||
}
|
||||
|
||||
func sha256Of(data []byte) v1.Hash {
|
||||
h, _, err := v1.SHA256(bytes.NewReader(data))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return h
|
||||
}
|
||||
|
||||
var _ = Describe("downloadLayerToFile resume", func() {
|
||||
var (
|
||||
dst *os.File
|
||||
data []byte
|
||||
restoreWait func()
|
||||
)
|
||||
|
||||
readDst := func() string {
|
||||
got, err := os.ReadFile(dst.Name())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return string(got)
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
var err error
|
||||
dst, err = os.CreateTemp("", "layer-resume-*.tar.gz")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
data = []byte("0123456789abcdefghijklmnopqrstuvwxyzABCD")
|
||||
|
||||
prev := layerRetryBackoff
|
||||
layerRetryBackoff = func(int) time.Duration { return 0 }
|
||||
restoreWait = func() { layerRetryBackoff = prev }
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
restoreWait()
|
||||
_ = dst.Close()
|
||||
_ = os.Remove(dst.Name())
|
||||
})
|
||||
|
||||
It("continues from the interruption offset instead of restarting", func() {
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:15],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(offset int64) (io.ReadCloser, int64, error) {
|
||||
return io.NopCloser(bytes.NewReader(data[offset:])), offset, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, rec.opener())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(readDst()).To(Equal(string(data)))
|
||||
// The interrupted first attempt left 15 bytes; the resume must ask
|
||||
// for exactly the rest, without a second full-stream attempt.
|
||||
Expect(rec.offsets).To(Equal([]int64{15}))
|
||||
Expect(layer.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("restarts cleanly when the server ignores the Range request", func() {
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:15],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(int64) (io.ReadCloser, int64, error) {
|
||||
// A 200 response: the whole blob from the first byte.
|
||||
return io.NopCloser(bytes.NewReader(data)), 0, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, rec.opener())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
// The partial bytes must have been discarded, not prepended.
|
||||
Expect(readDst()).To(Equal(string(data)))
|
||||
Expect(rec.offsets).To(HaveLen(1))
|
||||
Expect(layer.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("discards a resumed download whose digest does not match", func() {
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:15],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(offset int64) (io.ReadCloser, int64, error) {
|
||||
corrupt := bytes.Repeat([]byte("x"), len(data)-int(offset))
|
||||
return io.NopCloser(bytes.NewReader(corrupt)), offset, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, rec.opener())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
// The spliced file failed verification, so the download must have
|
||||
// started over through the verified layer reader and succeeded.
|
||||
Expect(readDst()).To(Equal(string(data)))
|
||||
Expect(rec.offsets).To(Equal([]int64{15}))
|
||||
Expect(layer.calls).To(Equal(2))
|
||||
})
|
||||
|
||||
It("keeps retrying beyond the budget while each resume makes progress", func() {
|
||||
const step = 5
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:step],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(offset int64) (io.ReadCloser, int64, error) {
|
||||
if offset+step >= int64(len(data)) {
|
||||
return io.NopCloser(bytes.NewReader(data[offset:])), offset, nil
|
||||
}
|
||||
return io.NopCloser(&failingReader{prefix: data[offset : offset+step], err: io.ErrUnexpectedEOF}), offset, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(context.Background(), layer, dst, nil, rec.opener())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(readDst()).To(Equal(string(data)))
|
||||
// 40 bytes delivered 5 at a time: 7 resumes, far more rounds than
|
||||
// the retry budget allows for stalled attempts.
|
||||
Expect(len(rec.offsets)).To(BeNumerically(">", layerDownloadRetries))
|
||||
})
|
||||
|
||||
It("gives up when resumes stop making progress", func(ctx SpecContext) {
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:15],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1000,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(offset int64) (io.ReadCloser, int64, error) {
|
||||
// Resume accepted but the connection dies before any byte.
|
||||
return io.NopCloser(&failingReader{err: io.ErrUnexpectedEOF}), offset, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(ctx, layer, dst, nil, rec.opener())
|
||||
Expect(err).To(MatchError(io.ErrUnexpectedEOF))
|
||||
Expect(len(rec.offsets)).To(Equal(layerDownloadRetries))
|
||||
}, NodeTimeout(10*time.Second))
|
||||
|
||||
It("terminates when the server ignores Range and keeps dropping mid-stream", func(ctx SpecContext) {
|
||||
// Each round delivers some bytes from the start and dies: the file
|
||||
// never gets further than before, so this must exhaust the budget
|
||||
// rather than count the repeated partial bytes as progress.
|
||||
layer := &fakeLayer{
|
||||
data: data,
|
||||
prefix: data[:15],
|
||||
digest: sha256Of(data),
|
||||
failUntil: 1000,
|
||||
err: io.ErrUnexpectedEOF,
|
||||
}
|
||||
rec := &recordingOpener{open: func(int64) (io.ReadCloser, int64, error) {
|
||||
return io.NopCloser(&failingReader{prefix: data[:15], err: io.ErrUnexpectedEOF}), 0, nil
|
||||
}}
|
||||
|
||||
err := downloadLayerToFile(ctx, layer, dst, nil, rec.opener())
|
||||
Expect(err).To(MatchError(io.ErrUnexpectedEOF))
|
||||
Expect(len(rec.offsets)).To(Equal(layerDownloadRetries))
|
||||
}, NodeTimeout(10*time.Second))
|
||||
})
|
||||
+333
-2
@@ -2,9 +2,24 @@ package vram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const persistentCacheEntryLimit = 4096
|
||||
const persistentCacheVersion = 1
|
||||
|
||||
var defaultPersistentGuard = &persistentGenerationGuard{}
|
||||
var defaultCacheMu sync.RWMutex
|
||||
|
||||
// galleryGenFunc returns the current gallery generation counter.
|
||||
// When set, cache entries are invalidated when the generation changes.
|
||||
// When nil (e.g., in tests or non-gallery contexts), entries never expire.
|
||||
@@ -23,6 +38,149 @@ func currentGeneration() uint64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// ConfigurePersistentCache replaces the process-wide estimator caches with
|
||||
// instances that reuse successful remote probes across server restarts.
|
||||
func ConfigurePersistentCache(dir string, ttl time.Duration) {
|
||||
defaultCacheMu.Lock()
|
||||
defer defaultCacheMu.Unlock()
|
||||
removeAbandonedPersistentTemps(dir)
|
||||
prunePersistentEntries(dir, ttl, persistentCacheEntryLimit)
|
||||
guard := &persistentGenerationGuard{dir: dir}
|
||||
defaultPersistentGuard = guard
|
||||
defaultCachedSizeResolver = newCachedSizeResolverWithGuard(defaultSizeResolver{}, dir, ttl, guard)
|
||||
defaultCachedGGUFReader = newCachedGGUFReaderWithGuard(defaultGGUFReader{}, dir, ttl, guard)
|
||||
}
|
||||
|
||||
// DisablePersistentCache keeps process-local caching but stops disk reads and writes.
|
||||
func DisablePersistentCache() {
|
||||
defaultCacheMu.Lock()
|
||||
defer defaultCacheMu.Unlock()
|
||||
defaultPersistentGuard = &persistentGenerationGuard{}
|
||||
defaultCachedSizeResolver = newCachedSizeResolver(defaultSizeResolver{}, "", 0)
|
||||
defaultCachedGGUFReader = newCachedGGUFReader(defaultGGUFReader{}, "", 0)
|
||||
}
|
||||
|
||||
// InvalidatePersistentCache removes remote probe results after the gallery
|
||||
// changes, including when no estimate is requested before the next restart.
|
||||
func InvalidatePersistentCache() {
|
||||
defaultCacheMu.RLock()
|
||||
guard := defaultPersistentGuard
|
||||
defaultCacheMu.RUnlock()
|
||||
guard.invalidate(currentGeneration())
|
||||
}
|
||||
|
||||
func removeAbandonedPersistentTemps(dir string) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.Type().IsRegular() && strings.HasPrefix(entry.Name(), ".vram-") {
|
||||
_ = os.Remove(filepath.Join(dir, entry.Name()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func removePersistentEntries(dir string) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.Type().IsRegular() && (strings.HasPrefix(name, "size-") || strings.HasPrefix(name, "gguf-")) {
|
||||
_ = os.Remove(filepath.Join(dir, name))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type persistentGenerationGuard struct {
|
||||
mu sync.Mutex
|
||||
dir string
|
||||
generation uint64
|
||||
set bool
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) invalidate(generation uint64) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
removePersistentEntries(g.dir)
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) canRead(generation uint64) bool {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if !g.set {
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
return true
|
||||
}
|
||||
if g.generation == generation {
|
||||
return true
|
||||
}
|
||||
removePersistentEntries(g.dir)
|
||||
g.generation = generation
|
||||
return false
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) persist(generation uint64, write func()) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if !g.set {
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
}
|
||||
if g.generation == generation {
|
||||
write()
|
||||
}
|
||||
}
|
||||
|
||||
func prunePersistentEntries(dir string, ttl time.Duration, limit int) {
|
||||
if dir == "" || ttl <= 0 {
|
||||
return
|
||||
}
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
type cacheFile struct {
|
||||
path string
|
||||
modTime time.Time
|
||||
}
|
||||
files := make([]cacheFile, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.Type().IsRegular() && (strings.HasPrefix(name, "size-") || strings.HasPrefix(name, "gguf-")) {
|
||||
if info, err := entry.Info(); err == nil {
|
||||
path := filepath.Join(dir, name)
|
||||
if time.Since(info.ModTime()) > ttl {
|
||||
_ = os.Remove(path)
|
||||
continue
|
||||
}
|
||||
files = append(files, cacheFile{path: path, modTime: info.ModTime()})
|
||||
}
|
||||
}
|
||||
}
|
||||
if limit <= 0 || len(files) <= limit {
|
||||
return
|
||||
}
|
||||
sort.Slice(files, func(i, j int) bool { return files[i].modTime.Before(files[j].modTime) })
|
||||
for _, file := range files[:len(files)-limit] {
|
||||
_ = os.Remove(file.path)
|
||||
}
|
||||
}
|
||||
|
||||
func persistentRemoteURI(uri string) bool {
|
||||
parsed, err := url.Parse(uri)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
scheme := strings.ToLower(parsed.Scheme)
|
||||
return (scheme == "http" || scheme == "https") && parsed.Host != ""
|
||||
}
|
||||
|
||||
type sizeCacheEntry struct {
|
||||
size int64
|
||||
err error
|
||||
@@ -33,6 +191,28 @@ type cachedSizeResolver struct {
|
||||
underlying SizeResolver
|
||||
mu sync.Mutex
|
||||
cache map[string]sizeCacheEntry
|
||||
diskDir string
|
||||
diskTTL time.Duration
|
||||
diskGuard *persistentGenerationGuard
|
||||
}
|
||||
|
||||
type persistentSizeEntry struct {
|
||||
Version int `json:"version"`
|
||||
Size int64 `json:"size"`
|
||||
}
|
||||
|
||||
func newCachedSizeResolver(underlying SizeResolver, diskDir string, diskTTL time.Duration) *cachedSizeResolver {
|
||||
return newCachedSizeResolverWithGuard(underlying, diskDir, diskTTL, &persistentGenerationGuard{dir: diskDir})
|
||||
}
|
||||
|
||||
func newCachedSizeResolverWithGuard(underlying SizeResolver, diskDir string, diskTTL time.Duration, guard *persistentGenerationGuard) *cachedSizeResolver {
|
||||
return &cachedSizeResolver{
|
||||
underlying: underlying,
|
||||
cache: make(map[string]sizeCacheEntry),
|
||||
diskDir: diskDir,
|
||||
diskTTL: diskTTL,
|
||||
diskGuard: guard,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) ContentLength(ctx context.Context, uri string) (int64, error) {
|
||||
@@ -43,13 +223,62 @@ func (c *cachedSizeResolver) ContentLength(ctx context.Context, uri string) (int
|
||||
if ok && e.generation == gen {
|
||||
return e.size, e.err
|
||||
}
|
||||
if persistentRemoteURI(uri) && c.canReadPersistent(gen) {
|
||||
if size, ok := c.readPersistent(uri); ok {
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = sizeCacheEntry{size: size, generation: gen}
|
||||
c.mu.Unlock()
|
||||
return size, nil
|
||||
}
|
||||
}
|
||||
size, err := c.underlying.ContentLength(ctx, uri)
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = sizeCacheEntry{size: size, err: err, generation: gen}
|
||||
c.mu.Unlock()
|
||||
if err == nil && persistentRemoteURI(uri) {
|
||||
c.writePersistent(uri, size, gen)
|
||||
}
|
||||
return size, err
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) canReadPersistent(generation uint64) bool {
|
||||
return c.diskGuard.canRead(generation)
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) persistentPath(uri string) string {
|
||||
digest := sha256.Sum256([]byte(uri))
|
||||
return filepath.Join(c.diskDir, "size-"+hex.EncodeToString(digest[:])+".json")
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) readPersistent(uri string) (int64, bool) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
path := c.persistentPath(uri)
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || time.Since(info.ModTime()) > c.diskTTL {
|
||||
return 0, false
|
||||
}
|
||||
data, err := os.ReadFile(path) // #nosec G304 -- path is a hash under the configured cache directory.
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
var entry persistentSizeEntry
|
||||
if json.Unmarshal(data, &entry) != nil || entry.Version != persistentCacheVersion || entry.Size <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
return entry.Size, true
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) writePersistent(uri string, size int64, generation uint64) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 || os.MkdirAll(c.diskDir, 0o750) != nil {
|
||||
return
|
||||
}
|
||||
c.diskGuard.persist(generation, func() {
|
||||
writePersistentJSON(c.persistentPath(uri), persistentSizeEntry{Version: persistentCacheVersion, Size: size}, c.diskTTL)
|
||||
})
|
||||
}
|
||||
|
||||
type ggufCacheEntry struct {
|
||||
meta *GGUFMeta
|
||||
err error
|
||||
@@ -60,6 +289,28 @@ type cachedGGUFReader struct {
|
||||
underlying GGUFMetadataReader
|
||||
mu sync.Mutex
|
||||
cache map[string]ggufCacheEntry
|
||||
diskDir string
|
||||
diskTTL time.Duration
|
||||
diskGuard *persistentGenerationGuard
|
||||
}
|
||||
|
||||
type persistentGGUFEntry struct {
|
||||
Version int `json:"version"`
|
||||
Meta *GGUFMeta `json:"meta"`
|
||||
}
|
||||
|
||||
func newCachedGGUFReader(underlying GGUFMetadataReader, diskDir string, diskTTL time.Duration) *cachedGGUFReader {
|
||||
return newCachedGGUFReaderWithGuard(underlying, diskDir, diskTTL, &persistentGenerationGuard{dir: diskDir})
|
||||
}
|
||||
|
||||
func newCachedGGUFReaderWithGuard(underlying GGUFMetadataReader, diskDir string, diskTTL time.Duration, guard *persistentGenerationGuard) *cachedGGUFReader {
|
||||
return &cachedGGUFReader{
|
||||
underlying: underlying,
|
||||
cache: make(map[string]ggufCacheEntry),
|
||||
diskDir: diskDir,
|
||||
diskTTL: diskTTL,
|
||||
diskGuard: guard,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFMeta, error) {
|
||||
@@ -70,26 +321,106 @@ func (c *cachedGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFM
|
||||
if ok && e.generation == gen {
|
||||
return e.meta, e.err
|
||||
}
|
||||
if persistentRemoteURI(uri) && c.canReadPersistent(gen) {
|
||||
if meta, ok := c.readPersistent(uri); ok {
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = ggufCacheEntry{meta: meta, generation: gen}
|
||||
c.mu.Unlock()
|
||||
return meta, nil
|
||||
}
|
||||
}
|
||||
meta, err := c.underlying.ReadMetadata(ctx, uri)
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = ggufCacheEntry{meta: meta, err: err, generation: gen}
|
||||
c.mu.Unlock()
|
||||
if err == nil && meta != nil && persistentRemoteURI(uri) {
|
||||
c.writePersistent(uri, meta, gen)
|
||||
}
|
||||
return meta, err
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) canReadPersistent(generation uint64) bool {
|
||||
return c.diskGuard.canRead(generation)
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) persistentPath(uri string) string {
|
||||
digest := sha256.Sum256([]byte(uri))
|
||||
return filepath.Join(c.diskDir, "gguf-"+hex.EncodeToString(digest[:])+".json")
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) readPersistent(uri string) (*GGUFMeta, bool) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 {
|
||||
return nil, false
|
||||
}
|
||||
path := c.persistentPath(uri)
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || time.Since(info.ModTime()) > c.diskTTL {
|
||||
return nil, false
|
||||
}
|
||||
data, err := os.ReadFile(path) // #nosec G304 -- path is a hash under the configured cache directory.
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
var entry persistentGGUFEntry
|
||||
if json.Unmarshal(data, &entry) != nil || entry.Version != persistentCacheVersion || !validPersistentGGUFMeta(entry.Meta) {
|
||||
return nil, false
|
||||
}
|
||||
return entry.Meta, true
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) writePersistent(uri string, meta *GGUFMeta, generation uint64) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 || os.MkdirAll(c.diskDir, 0o750) != nil {
|
||||
return
|
||||
}
|
||||
c.diskGuard.persist(generation, func() {
|
||||
writePersistentJSON(c.persistentPath(uri), persistentGGUFEntry{Version: persistentCacheVersion, Meta: meta}, c.diskTTL)
|
||||
})
|
||||
}
|
||||
|
||||
func validPersistentGGUFMeta(meta *GGUFMeta) bool {
|
||||
return meta != nil && meta.BlockCount > 0 && meta.EmbeddingLength > 0 && meta.HeadCount > 0 && meta.HeadCountKV > 0
|
||||
}
|
||||
|
||||
func writePersistentJSON(path string, value any, ttl time.Duration) {
|
||||
data, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
tmp, err := os.CreateTemp(filepath.Dir(path), ".vram-*.tmp")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
tmpPath := tmp.Name()
|
||||
defer func() { _ = os.Remove(tmpPath) }()
|
||||
if _, err = tmp.Write(data); err != nil {
|
||||
_ = tmp.Close()
|
||||
return
|
||||
}
|
||||
if err = tmp.Close(); err != nil {
|
||||
return
|
||||
}
|
||||
if os.Rename(tmpPath, path) == nil {
|
||||
prunePersistentEntries(filepath.Dir(path), ttl, persistentCacheEntryLimit)
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultCachedSizeResolver returns a cached SizeResolver using the default implementation.
|
||||
// Entries are invalidated when the gallery generation changes.
|
||||
func DefaultCachedSizeResolver() SizeResolver {
|
||||
defaultCacheMu.RLock()
|
||||
defer defaultCacheMu.RUnlock()
|
||||
return defaultCachedSizeResolver
|
||||
}
|
||||
|
||||
// DefaultCachedGGUFReader returns a cached GGUFMetadataReader using the default implementation.
|
||||
// Entries are invalidated when the gallery generation changes.
|
||||
func DefaultCachedGGUFReader() GGUFMetadataReader {
|
||||
defaultCacheMu.RLock()
|
||||
defer defaultCacheMu.RUnlock()
|
||||
return defaultCachedGGUFReader
|
||||
}
|
||||
|
||||
var (
|
||||
defaultCachedSizeResolver = &cachedSizeResolver{underlying: defaultSizeResolver{}, cache: make(map[string]sizeCacheEntry)}
|
||||
defaultCachedGGUFReader = &cachedGGUFReader{underlying: defaultGGUFReader{}, cache: make(map[string]ggufCacheEntry)}
|
||||
defaultCachedSizeResolver = newCachedSizeResolver(defaultSizeResolver{}, "", 0)
|
||||
defaultCachedGGUFReader = newCachedGGUFReader(defaultGGUFReader{}, "", 0)
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
package vram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type countingSizeResolver struct {
|
||||
size int64
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
type countingGGUFReader struct {
|
||||
meta *GGUFMeta
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
type blockingSizeResolver struct {
|
||||
started chan struct{}
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (r *blockingSizeResolver) ContentLength(context.Context, string) (int64, error) {
|
||||
close(r.started)
|
||||
<-r.release
|
||||
return 42, nil
|
||||
}
|
||||
|
||||
func (r *countingGGUFReader) ReadMetadata(context.Context, string) (*GGUFMeta, error) {
|
||||
r.calls++
|
||||
return r.meta, r.err
|
||||
}
|
||||
|
||||
func (r *countingSizeResolver) ContentLength(context.Context, string) (int64, error) {
|
||||
r.calls++
|
||||
return r.size, r.err
|
||||
}
|
||||
|
||||
var _ = Describe("persistent VRAM metadata cache", func() {
|
||||
AfterEach(func() {
|
||||
ConfigurePersistentCache("", 0)
|
||||
SetGalleryGenerationFunc(nil)
|
||||
})
|
||||
|
||||
It("reuses a successful size probe after the in-memory cache is replaced", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
firstSource := &countingSizeResolver{size: 42}
|
||||
first := newCachedSizeResolver(firstSource, cacheDir, time.Hour)
|
||||
|
||||
size, err := first.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
Expect(firstSource.calls).To(Equal(1))
|
||||
|
||||
secondSource := &countingSizeResolver{err: errors.New("unexpected remote probe")}
|
||||
second := newCachedSizeResolver(secondSource, cacheDir, time.Hour)
|
||||
size, err = second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
Expect(secondSource.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("reuses successful GGUF metadata after the in-memory cache is replaced", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
want := &GGUFMeta{BlockCount: 32, EmbeddingLength: 4096, HeadCount: 32, HeadCountKV: 8, MaximumContextLength: 131072}
|
||||
firstSource := &countingGGUFReader{meta: want}
|
||||
first := newCachedGGUFReader(firstSource, cacheDir, time.Hour)
|
||||
|
||||
meta, err := first.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
Expect(firstSource.calls).To(Equal(1))
|
||||
|
||||
secondSource := &countingGGUFReader{err: errors.New("unexpected remote probe")}
|
||||
second := newCachedGGUFReader(secondSource, cacheDir, time.Hour)
|
||||
meta, err = second.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
Expect(secondSource.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("configures the default caches used by model estimates", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
|
||||
Expect(defaultCachedSizeResolver.diskDir).To(Equal(cacheDir))
|
||||
Expect(defaultCachedGGUFReader.diskDir).To(Equal(cacheDir))
|
||||
Expect(defaultCachedSizeResolver.diskTTL).To(Equal(time.Hour))
|
||||
Expect(defaultCachedGGUFReader.diskTTL).To(Equal(time.Hour))
|
||||
Expect(defaultCachedSizeResolver.diskGuard).To(BeIdenticalTo(defaultCachedGGUFReader.diskGuard))
|
||||
})
|
||||
|
||||
It("removes expired VRAM entries when the persistent cache is configured", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
stale := filepath.Join(cacheDir, "size-stale.json")
|
||||
abandoned := filepath.Join(cacheDir, ".vram-abandoned.tmp")
|
||||
unrelated := filepath.Join(cacheDir, "keep.txt")
|
||||
Expect(os.WriteFile(stale, []byte("{}"), 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(abandoned, []byte("partial"), 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(unrelated, []byte("keep"), 0o600)).To(Succeed())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
Expect(os.Chtimes(stale, old, old)).To(Succeed())
|
||||
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
|
||||
Expect(stale).NotTo(BeAnExistingFile())
|
||||
Expect(abandoned).NotTo(BeAnExistingFile())
|
||||
Expect(unrelated).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("does not reuse a persistent entry after the gallery generation changes", func() {
|
||||
var generation uint64 = 1
|
||||
SetGalleryGenerationFunc(func() uint64 { return generation })
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
first := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
_, err := first.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
freshSource := &countingSizeResolver{size: 84}
|
||||
second := newCachedSizeResolver(freshSource, cacheDir, time.Hour)
|
||||
size, err := second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
|
||||
generation = 2
|
||||
size, err = second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(84)))
|
||||
Expect(freshSource.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("does not persist probes for local model files", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
first := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
_, err := first.ContentLength(context.Background(), "file:///models/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
freshSource := &countingSizeResolver{size: 84}
|
||||
second := newCachedSizeResolver(freshSource, cacheDir, time.Hour)
|
||||
size, err := second.ContentLength(context.Background(), "file:///models/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(84)))
|
||||
Expect(freshSource.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("falls back to the remote probe for an invalid persisted size", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
resolver := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
Expect(os.WriteFile(resolver.persistentPath("https://example.com/model.gguf"), []byte(`{"size":-1}`), 0o600)).To(Succeed())
|
||||
|
||||
size, err := resolver.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
})
|
||||
|
||||
It("falls back to the remote probe for empty persisted GGUF metadata", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
want := &GGUFMeta{BlockCount: 32, EmbeddingLength: 4096, HeadCount: 32, HeadCountKV: 8}
|
||||
reader := newCachedGGUFReader(&countingGGUFReader{meta: want}, cacheDir, time.Hour)
|
||||
Expect(os.WriteFile(reader.persistentPath("https://example.com/model.gguf"), []byte(`{"meta":{}}`), 0o600)).To(Succeed())
|
||||
|
||||
meta, err := reader.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
})
|
||||
|
||||
It("keeps the persistent cache within its entry limit", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
for _, name := range []string{"size-a.json", "size-b.json", "gguf-c.json"} {
|
||||
Expect(os.WriteFile(filepath.Join(cacheDir, name), []byte("{}"), 0o600)).To(Succeed())
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
prunePersistentEntries(cacheDir, time.Hour, 2)
|
||||
|
||||
entries, err := os.ReadDir(cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(entries).To(HaveLen(2))
|
||||
Expect(filepath.Join(cacheDir, "size-a.json")).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("removes persistent entries when gallery data is invalidated", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
stale := filepath.Join(cacheDir, "size-stale.json")
|
||||
Expect(os.WriteFile(stale, []byte(`{"version":1,"size":42}`), 0o600)).To(Succeed())
|
||||
|
||||
InvalidatePersistentCache()
|
||||
|
||||
Expect(stale).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("does not persist a probe that finishes after invalidation", func() {
|
||||
var generation uint64 = 1
|
||||
SetGalleryGenerationFunc(func() uint64 { return generation })
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
guard := &persistentGenerationGuard{dir: cacheDir}
|
||||
source := &blockingSizeResolver{started: make(chan struct{}), release: make(chan struct{})}
|
||||
resolver := newCachedSizeResolverWithGuard(source, cacheDir, time.Hour, guard)
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := resolver.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
done <- err
|
||||
}()
|
||||
Eventually(source.started).Should(BeClosed())
|
||||
|
||||
generation = 2
|
||||
guard.invalidate(generation)
|
||||
close(source.release)
|
||||
|
||||
Eventually(done).Should(Receive(BeNil()))
|
||||
Expect(resolver.persistentPath("https://example.com/model.gguf")).NotTo(BeAnExistingFile())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
|
||||
PATCHER="$ROOT/backend/cpp/bonsai/patch-grpc-server.sh"
|
||||
WORK=$(mktemp -d)
|
||||
trap 'rm -rf "$WORK"' EXIT
|
||||
|
||||
cat > "$WORK/grpc-server.cpp" <<'EOF'
|
||||
try {
|
||||
json::parse("{");
|
||||
} catch (const common_json_error& e) {
|
||||
}
|
||||
EOF
|
||||
|
||||
bash "$PATCHER" "$WORK/grpc-server.cpp"
|
||||
grep -q 'catch (const json::parse_error& e)' "$WORK/grpc-server.cpp"
|
||||
! grep -q 'common_json_error' "$WORK/grpc-server.cpp"
|
||||
|
||||
# A repeated preparation pass must not change the generated source.
|
||||
cp "$WORK/grpc-server.cpp" "$WORK/once.cpp"
|
||||
bash "$PATCHER" "$WORK/grpc-server.cpp"
|
||||
cmp "$WORK/once.cpp" "$WORK/grpc-server.cpp"
|
||||
|
||||
echo "PASS: Bonsai uses its fork-compatible JSON exception"
|
||||
@@ -4580,6 +4580,9 @@ const docTemplate = `{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"config": {},
|
||||
"config_revision": {
|
||||
"type": "string"
|
||||
},
|
||||
"details": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
@@ -4595,6 +4598,9 @@ const docTemplate = `{
|
||||
"message": {
|
||||
"type": "string"
|
||||
},
|
||||
"pending_cleanup": {
|
||||
"type": "integer"
|
||||
},
|
||||
"success": {
|
||||
"type": "boolean"
|
||||
}
|
||||
@@ -4735,9 +4741,24 @@ const docTemplate = `{
|
||||
"description": "e.g. \"llama-cpp\"; used by reconciler to replicate loads",
|
||||
"type": "string"
|
||||
},
|
||||
"cleanup_attempts": {
|
||||
"type": "integer"
|
||||
},
|
||||
"cleanup_error": {
|
||||
"type": "string"
|
||||
},
|
||||
"cleanup_next_retry_at": {
|
||||
"type": "string"
|
||||
},
|
||||
"config_revision": {
|
||||
"type": "string"
|
||||
},
|
||||
"created_at": {
|
||||
"type": "string"
|
||||
},
|
||||
"effective_options_hash": {
|
||||
"type": "string"
|
||||
},
|
||||
"id": {
|
||||
"type": "string"
|
||||
},
|
||||
|
||||
@@ -4577,6 +4577,9 @@
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"config": {},
|
||||
"config_revision": {
|
||||
"type": "string"
|
||||
},
|
||||
"details": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
@@ -4592,6 +4595,9 @@
|
||||
"message": {
|
||||
"type": "string"
|
||||
},
|
||||
"pending_cleanup": {
|
||||
"type": "integer"
|
||||
},
|
||||
"success": {
|
||||
"type": "boolean"
|
||||
}
|
||||
@@ -4732,9 +4738,24 @@
|
||||
"description": "e.g. \"llama-cpp\"; used by reconciler to replicate loads",
|
||||
"type": "string"
|
||||
},
|
||||
"cleanup_attempts": {
|
||||
"type": "integer"
|
||||
},
|
||||
"cleanup_error": {
|
||||
"type": "string"
|
||||
},
|
||||
"cleanup_next_retry_at": {
|
||||
"type": "string"
|
||||
},
|
||||
"config_revision": {
|
||||
"type": "string"
|
||||
},
|
||||
"created_at": {
|
||||
"type": "string"
|
||||
},
|
||||
"effective_options_hash": {
|
||||
"type": "string"
|
||||
},
|
||||
"id": {
|
||||
"type": "string"
|
||||
},
|
||||
|
||||
@@ -414,6 +414,8 @@ definitions:
|
||||
localai.ModelResponse:
|
||||
properties:
|
||||
config: {}
|
||||
config_revision:
|
||||
type: string
|
||||
details:
|
||||
items:
|
||||
type: string
|
||||
@@ -424,6 +426,8 @@ definitions:
|
||||
type: string
|
||||
message:
|
||||
type: string
|
||||
pending_cleanup:
|
||||
type: integer
|
||||
success:
|
||||
type: boolean
|
||||
type: object
|
||||
@@ -518,8 +522,18 @@ definitions:
|
||||
backend_type:
|
||||
description: e.g. "llama-cpp"; used by reconciler to replicate loads
|
||||
type: string
|
||||
cleanup_attempts:
|
||||
type: integer
|
||||
cleanup_error:
|
||||
type: string
|
||||
cleanup_next_retry_at:
|
||||
type: string
|
||||
config_revision:
|
||||
type: string
|
||||
created_at:
|
||||
type: string
|
||||
effective_options_hash:
|
||||
type: string
|
||||
id:
|
||||
type: string
|
||||
in_flight:
|
||||
|
||||
@@ -260,6 +260,12 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func()
|
||||
data, _ := json.Marshal(reply)
|
||||
msg.Respond(data)
|
||||
})
|
||||
_, err := infra.NC.Conn().Subscribe("nodes.*.models.running", func(msg *nats.Msg) {
|
||||
data, _ := json.Marshal(messaging.ModelsRunningReply{})
|
||||
_ = msg.Respond(data)
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
FlushNATS(infra.NC)
|
||||
|
||||
return router
|
||||
}
|
||||
|
||||
@@ -66,6 +66,12 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() {
|
||||
data, _ := json.Marshal(reply)
|
||||
msg.Respond(data)
|
||||
})
|
||||
_, err = infra.NC.Conn().Subscribe("nodes.*.models.running", func(msg *nats.Msg) {
|
||||
data, _ := json.Marshal(messaging.ModelsRunningReply{})
|
||||
_ = msg.Respond(data)
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
FlushNATS(infra.NC)
|
||||
|
||||
// Start a mock gRPC backend using the same helper as full flow tests
|
||||
llm := &trackingTestLLM{}
|
||||
|
||||
@@ -3,10 +3,10 @@
|
||||
# The four GitHub fields are rewritten by .github/ci/refresh-site-counters.sh,
|
||||
# which runs weekly from .github/workflows/refresh-site-counters.yml. Editing
|
||||
# them by hand works but will be overwritten on the next run.
|
||||
stars: 48067
|
||||
forks: 4320
|
||||
contributors: 225
|
||||
releases: 133
|
||||
stars: 48646
|
||||
forks: 4377
|
||||
contributors: 230
|
||||
releases: 136
|
||||
|
||||
# The GitHub API cannot answer for this one, so it is maintained by hand and
|
||||
# the refresh script carries it through untouched.
|
||||
|
||||
Reference in new issue
Block a user