mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
chore: merge master into backend stop fix
Preserve runtime cleanup and stopped-process acknowledgements. Assisted-by: Codex:gpt-6
This commit is contained in:
commit
32edf0810e
126 files changed
+9359
-546
No files matched your search
@@ -5,7 +5,7 @@ This PR fixes #
|
||||
**Notes for Reviewers**
|
||||
|
||||
|
||||
**[Signed commits](../CONTRIBUTING.md#signing-off-on-commits-developer-certificate-of-origin)**
|
||||
**[Signed commits](../CONTRIBUTING.md#commit-messages)**
|
||||
- [ ] Yes, I signed my commits.
|
||||
- [ ] Documentation updated (docs/content/) for user-facing changes, or not applicable
|
||||
|
||||
|
||||
@@ -240,7 +240,7 @@ test-ci-scripts:
|
||||
## pure stdlib on purpose so they run without any backend venv; the list is
|
||||
## explicit because their siblings (model_identity_test) import grpc and the
|
||||
## generated protobufs, which only exist inside a built backend.
|
||||
PYTHON_HELPER_TESTS?=python_utils_test vllm_utils_test model_utils_test mlx_utils_test parent_watch_test
|
||||
PYTHON_HELPER_TESTS?=python_utils_test vllm_utils_test model_utils_test mlx_utils_test parent_watch_test temp_utils_test
|
||||
test-python-helpers:
|
||||
cd backend/python/common && python3 -m unittest $(PYTHON_HELPER_TESTS)
|
||||
|
||||
|
||||
@@ -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?=c18b7f737aac0a2855e9f963a427498739ad40fe
|
||||
AUDIO_CPP_VERSION?=9c6a282337cc83f227cc10428867a478947706ad
|
||||
AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -84,9 +84,10 @@ elseif(DS4_GPU STREQUAL "cpu")
|
||||
set(DS4_OBJS "${DS4_DIR}/ds4_cpu.o")
|
||||
endif()
|
||||
|
||||
# Upstream splits distributed inference, tensor-parallel transport, the SSD
|
||||
# expert cache, and layer placement into GPU-agnostic translation units. Link
|
||||
# them regardless of DS4_GPU.
|
||||
# Upstream splits image preprocessing, distributed inference, tensor-parallel
|
||||
# transport, the SSD expert cache, and layer placement into GPU-agnostic
|
||||
# translation units. Link them regardless of DS4_GPU.
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_image.o")
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_distributed.o")
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_tp.o")
|
||||
list(APPEND DS4_OBJS "${DS4_DIR}/ds4_ssd.o")
|
||||
|
||||
+11
-10
@@ -1,10 +1,10 @@
|
||||
# ds4 backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DS4_VERSION?=8db89fe083ae4d17c9a2428ccd29803d3ae8f577
|
||||
# Upstream pin lives below as DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=8db89fe083ae4d17c9a2428ccd29803d3ae8f577
|
||||
DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
|
||||
DS4_REPO?=https://github.com/antirez/ds4
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
@@ -79,21 +79,22 @@ else
|
||||
endif
|
||||
endif
|
||||
|
||||
# Upstream splits distributed inference, tensor-parallel transport, the SSD
|
||||
# expert cache, and layer placement into GPU-agnostic translation units. They
|
||||
# are shared by every GPU mode, so append them unconditionally below.
|
||||
# Upstream splits image preprocessing, distributed inference, tensor-parallel
|
||||
# transport, the SSD expert cache, and layer placement into GPU-agnostic
|
||||
# translation units. They are shared by every GPU mode, so append them
|
||||
# unconditionally below.
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CMAKE_ARGS += -DDS4_GPU=cuda
|
||||
DS4_OBJ_TARGET := ds4.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o \
|
||||
DS4_OBJ_TARGET := ds4.o ds4_image.o ds4_cuda.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o \
|
||||
cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o \
|
||||
cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o
|
||||
else ifeq ($(UNAME_S),Darwin)
|
||||
CMAKE_ARGS += -DDS4_GPU=metal
|
||||
DS4_OBJ_TARGET := ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
DS4_OBJ_TARGET := ds4.o ds4_image.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
else
|
||||
# CPU reference path (Linux only - macOS CPU path is broken by VM bug per ds4 README).
|
||||
CMAKE_ARGS += -DDS4_GPU=cpu
|
||||
DS4_OBJ_TARGET := ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
DS4_OBJ_TARGET := ds4_cpu.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
endif
|
||||
|
||||
ifneq ($(NATIVE),true)
|
||||
@@ -120,9 +121,9 @@ ds4/ds4.o: ds4
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
+$(MAKE) -C ds4 $(DS4_ARCH_MAKEVARS) $(DS4_OBJ_TARGET)
|
||||
else ifeq ($(UNAME_S),Darwin)
|
||||
+$(MAKE) -C ds4 ds4.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
+$(MAKE) -C ds4 ds4.o ds4_image.o ds4_metal.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
else
|
||||
+$(MAKE) -C ds4 ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
+$(MAKE) -C ds4 ds4_cpu.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
|
||||
endif
|
||||
|
||||
grpc-server: ds4/ds4.o
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
IK_LLAMA_VERSION?=caf7eae5282d840d77e9f91a56df7d2ef28fa612
|
||||
IK_LLAMA_VERSION?=fe215a8ccdce6b844d2a3a3bbde08ae76a6284bf
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=9cffdcc801582616250520966699cb5b25d28243
|
||||
LLAMA_VERSION?=67672dc5b76f8bc17785a19d3dc6d1463fc2902c
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# CrispASR version (release tag)
|
||||
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
|
||||
CRISPASR_VERSION?=ff3945c94cab9191199a5d531a32c4e9535c094b
|
||||
CRISPASR_VERSION?=301acd87b036764973b8bfba71e0a21818036d33
|
||||
SO_TARGET?=libgocrispasr.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -615,10 +615,10 @@ func (w *CrispASR) TTSStream(req *pb.TTSRequest, results chan []byte) error {
|
||||
return fmt.Errorf("crispasr: tempfile: %w", err)
|
||||
}
|
||||
dst := tmp.Name()
|
||||
defer func() { _ = os.Remove(dst) }()
|
||||
if err := tmp.Close(); err != nil {
|
||||
return fmt.Errorf("crispasr: close tempfile: %w", err)
|
||||
}
|
||||
defer func() { _ = os.Remove(dst) }()
|
||||
|
||||
if err := writeWAV(dst, pcm, w.sampleRate); err != nil {
|
||||
return err
|
||||
|
||||
@@ -14,7 +14,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
# It is kept alive by the upstream tag da2-support (survives a squash-merge);
|
||||
# repoint to the master merge commit once mudler/depth-anything.cpp PR #1 lands.
|
||||
DEPTHANYTHING_REPO?=https://github.com/mudler/depth-anything.cpp.git
|
||||
DEPTHANYTHING_VERSION?=02ba082274e001a63e50de5a1eb0ccc50c6af4b1
|
||||
DEPTHANYTHING_VERSION?=14f7461d1f704761a038ac9f50dbde8fdb7275e2
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"unsafe"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
@@ -109,30 +110,25 @@ func (r *LocateAnythingCpp) Detect(opts *pb.DetectOptions) (pb.DetectResponse, e
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: a text prompt is required (open-vocabulary detection)")
|
||||
}
|
||||
|
||||
// Decode base64 image and write to temp file.
|
||||
imgData, err := base64.StdEncoding.DecodeString(opts.Src)
|
||||
if err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: failed to decode base64 image: %w", err)
|
||||
}
|
||||
|
||||
tmpFile, err := os.CreateTemp("", "locate-anything-*.img")
|
||||
if err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: failed to create temp file: %w", err)
|
||||
}
|
||||
defer func() { _ = os.Remove(tmpFile.Name()) }()
|
||||
|
||||
if _, err := tmpFile.Write(imgData); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: failed to write temp file: %w", err)
|
||||
}
|
||||
if err := tmpFile.Close(); err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: failed to close temp file: %w", err)
|
||||
if len(imgData) == 0 {
|
||||
return pb.DetectResponse{}, fmt.Errorf("locate-anything-cpp: decoded image is empty")
|
||||
}
|
||||
|
||||
// mode 0 = hybrid (Parallel Box Decoding). The JSON return value is unused:
|
||||
// structured detections are read via the accessor functions. Still must
|
||||
// free the returned string.
|
||||
jsonPtr := CapiLocatePath(r.handle, tmpFile.Name(), prompt, 0)
|
||||
jsonPtr := CapiLocateBuffer(
|
||||
r.handle,
|
||||
uintptr(unsafe.Pointer(unsafe.SliceData(imgData))),
|
||||
uintptr(len(imgData)),
|
||||
prompt,
|
||||
0,
|
||||
)
|
||||
runtime.KeepAlive(imgData)
|
||||
if jsonPtr != 0 {
|
||||
CapiFreeString(jsonPtr)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"path/filepath"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("LocateAnythingCpp detection input", func() {
|
||||
It("detects from memory when the temporary directory is unavailable", func() {
|
||||
originalLocateBuffer := CapiLocateBuffer
|
||||
originalLocatePath := CapiLocatePath
|
||||
originalGetNDetections := CapiGetNDetections
|
||||
defer func() {
|
||||
CapiLocateBuffer = originalLocateBuffer
|
||||
CapiLocatePath = originalLocatePath
|
||||
CapiGetNDetections = originalGetNDetections
|
||||
}()
|
||||
|
||||
image := []byte("encoded-image")
|
||||
var receivedData uintptr
|
||||
var receivedLength uintptr
|
||||
CapiLocateBuffer = func(_ uintptr, data uintptr, length uintptr, _ string, _ int32) uintptr {
|
||||
receivedData = data
|
||||
receivedLength = length
|
||||
return 0
|
||||
}
|
||||
CapiLocatePath = func(_ uintptr, _ string, _ string, _ int32) uintptr {
|
||||
Fail("path-based detection must not be called")
|
||||
return 0
|
||||
}
|
||||
CapiGetNDetections = func(uintptr) int32 { return 0 }
|
||||
GinkgoT().Setenv("TMPDIR", filepath.Join(GinkgoT().TempDir(), "missing"))
|
||||
|
||||
result, err := (&LocateAnythingCpp{handle: 1}).Detect(&pb.DetectOptions{
|
||||
Src: base64.StdEncoding.EncodeToString(image),
|
||||
Prompt: "the object",
|
||||
})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result.Detections).To(BeEmpty())
|
||||
Expect(receivedData).NotTo(BeZero())
|
||||
Expect(receivedLength).To(Equal(uintptr(len(image))))
|
||||
})
|
||||
|
||||
It("rejects an empty decoded image", func() {
|
||||
_, err := (&LocateAnythingCpp{handle: 1}).Detect(&pb.DetectOptions{Prompt: "the object"})
|
||||
|
||||
Expect(err).To(MatchError("locate-anything-cpp: decoded image is empty"))
|
||||
})
|
||||
})
|
||||
@@ -12,7 +12,7 @@
|
||||
# runs 'make -C backend/go/$(BACKEND) build' and then copies package/), so it
|
||||
# has to produce the binary and the package, not just the shared libraries.
|
||||
|
||||
NEMO_SPEECH_VERSION?=56b60d432f1731d6d5b28a4c5a31cbaf871daba1
|
||||
NEMO_SPEECH_VERSION?=ffa38cb2408f1e832a36d46fef5e3e1e80d07e6c
|
||||
NEMO_SPEECH_REPO?=https://github.com/NVIDIA/NeMo-Speech.cpp
|
||||
|
||||
GOCMD?=go
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# omnivoice.cpp version
|
||||
OMNIVOICE_REPO?=https://github.com/ServeurpersoCom/omnivoice.cpp
|
||||
OMNIVOICE_VERSION?=4f33af825d66e6ef1cb185e87b4589cacf747291
|
||||
OMNIVOICE_VERSION?=040c8b344d8c670ce1475194751d119b5ef82c78
|
||||
SO_TARGET?=libgomnivoicecpp.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"unsafe"
|
||||
|
||||
@@ -102,24 +103,12 @@ func (r *RFDetrCpp) Detect(opts *pb.DetectOptions) (pb.DetectResponse, error) {
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: model not loaded")
|
||||
}
|
||||
|
||||
// Decode base64 image and write to temp file.
|
||||
imgData, err := base64.StdEncoding.DecodeString(opts.Src)
|
||||
if err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: failed to decode base64 image: %w", err)
|
||||
}
|
||||
|
||||
tmpFile, err := os.CreateTemp("", "rfdetr-*.img")
|
||||
if err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: failed to create temp file: %w", err)
|
||||
}
|
||||
defer func() { _ = os.Remove(tmpFile.Name()) }()
|
||||
|
||||
if _, err := tmpFile.Write(imgData); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: failed to write temp file: %w", err)
|
||||
}
|
||||
if err := tmpFile.Close(); err != nil {
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: failed to close temp file: %w", err)
|
||||
if len(imgData) == 0 {
|
||||
return pb.DetectResponse{}, fmt.Errorf("rfdetr-cpp: decoded image is empty")
|
||||
}
|
||||
|
||||
threshold := opts.Threshold
|
||||
@@ -127,10 +116,18 @@ func (r *RFDetrCpp) Detect(opts *pb.DetectOptions) (pb.DetectResponse, error) {
|
||||
threshold = 0.5
|
||||
}
|
||||
|
||||
// JSON output from detect_path is unused: we read structured detections via
|
||||
// JSON output from the detection ABI is unused: we read structured detections via
|
||||
// the accessor functions. Still must free the returned string.
|
||||
var jsonPtr uintptr
|
||||
rc := CapiDetectPath(r.handle, tmpFile.Name(), threshold, uint32(defaultTopK), &jsonPtr)
|
||||
rc := CapiDetectBuffer(
|
||||
r.handle,
|
||||
uintptr(unsafe.Pointer(unsafe.SliceData(imgData))),
|
||||
uintptr(len(imgData)),
|
||||
threshold,
|
||||
uint32(defaultTopK),
|
||||
&jsonPtr,
|
||||
)
|
||||
runtime.KeepAlive(imgData)
|
||||
if jsonPtr != 0 {
|
||||
CapiFreeString(jsonPtr)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"path/filepath"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("RFDetrCpp detection input", func() {
|
||||
It("detects from memory when the temporary directory is unavailable", func() {
|
||||
originalDetectBuffer := CapiDetectBuffer
|
||||
originalDetectPath := CapiDetectPath
|
||||
originalFreeString := CapiFreeString
|
||||
originalGetNDetections := CapiGetNDetections
|
||||
defer func() {
|
||||
CapiDetectBuffer = originalDetectBuffer
|
||||
CapiDetectPath = originalDetectPath
|
||||
CapiFreeString = originalFreeString
|
||||
CapiGetNDetections = originalGetNDetections
|
||||
}()
|
||||
|
||||
image := []byte("encoded-image")
|
||||
var receivedData uintptr
|
||||
var receivedLength uintptr
|
||||
CapiDetectBuffer = func(_ uintptr, data uintptr, length uintptr, _ float32, _ uint32, _ *uintptr) int32 {
|
||||
receivedData = data
|
||||
receivedLength = length
|
||||
return 0
|
||||
}
|
||||
CapiDetectPath = func(_ uintptr, _ string, _ float32, _ uint32, _ *uintptr) int32 {
|
||||
Fail("path-based detection must not be called")
|
||||
return -1
|
||||
}
|
||||
CapiFreeString = func(uintptr) {}
|
||||
CapiGetNDetections = func(uintptr) int32 { return 0 }
|
||||
GinkgoT().Setenv("TMPDIR", filepath.Join(GinkgoT().TempDir(), "missing"))
|
||||
|
||||
result, err := (&RFDetrCpp{handle: 1}).Detect(&pb.DetectOptions{
|
||||
Src: base64.StdEncoding.EncodeToString(image),
|
||||
})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result.Detections).To(BeEmpty())
|
||||
Expect(receivedData).NotTo(BeZero())
|
||||
Expect(receivedLength).To(Equal(uintptr(len(image))))
|
||||
})
|
||||
|
||||
It("rejects an empty decoded image", func() {
|
||||
_, err := (&RFDetrCpp{handle: 1}).Detect(&pb.DetectOptions{})
|
||||
|
||||
Expect(err).To(MatchError("rfdetr-cpp: decoded image is empty"))
|
||||
})
|
||||
})
|
||||
@@ -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?=6b3edaaf32cc19e5bb2d819c788bd557eddc8eba
|
||||
STABLEDIFFUSION_GGML_VERSION?=d04e8950c1ec8d30248cbe996682b3182fb1adf6
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
@@ -401,7 +401,6 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
const char *params_backend_arg = "";
|
||||
const char *rpc_servers_arg = "";
|
||||
const char *max_vram_arg = "";
|
||||
bool stream_layers = false;
|
||||
|
||||
int n_threads = threads;
|
||||
enum sd_type_t wtype = SD_TYPE_COUNT;
|
||||
@@ -510,7 +509,10 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
if (!strcmp(optname, "params_backend")) params_backend_arg = strdup(optval);
|
||||
if (!strcmp(optname, "rpc_servers")) rpc_servers_arg = strdup(optval);
|
||||
if (!strcmp(optname, "max_vram")) max_vram_arg = strdup(optval);
|
||||
if (!strcmp(optname, "stream_layers")) stream_layers = (strcmp(optval, "true") == 0 || strcmp(optval, "1") == 0);
|
||||
if (!strcmp(optname, "stream_layers")) {
|
||||
// Retained as a no-op for existing configurations. Upstream now
|
||||
// selects segmented weight streaming automatically.
|
||||
}
|
||||
|
||||
// vae_decode_only is still accepted for backwards compatibility with
|
||||
// existing gallery configs, but upstream dropped the option (the model
|
||||
@@ -650,11 +652,9 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
ctx_params.rpc_servers = env_rpc_servers;
|
||||
}
|
||||
}
|
||||
// max_vram: GiB budget or per-backend spec for graph-cut segmented param
|
||||
// offload ("0" = disabled, "-1" = auto). stream_layers only has effect when
|
||||
// max_vram is set.
|
||||
// max_vram is an optional GiB budget or per-backend spec for automatic
|
||||
// graph-cut execution. A zero value uses the live free-VRAM budget.
|
||||
if (strlen(max_vram_arg) > 0) ctx_params.max_vram = max_vram_arg;
|
||||
ctx_params.stream_layers = stream_layers;
|
||||
ctx_params.diffusion_flash_attn = diffusion_flash_attn;
|
||||
ctx_params.tae_preview_only = tae_preview_only;
|
||||
ctx_params.diffusion_conv_direct = diffusion_conv_direct;
|
||||
@@ -1144,17 +1144,25 @@ static uint8_t* load_and_resize_image(const char* path, int target_width, int ta
|
||||
// Write sd.cpp's audio buffer to a temp WAV file (IEEE float, interleaved).
|
||||
// sd_audio_t.data is planar (all channel 0 samples, then channel 1, etc.) — we
|
||||
// interleave on the fly so ffmpeg's standard wav demuxer can read it directly.
|
||||
// Returns 0 on success and fills wav_path (must be at least 64 bytes).
|
||||
// Returns 0 on success and fills wav_path.
|
||||
static int write_planar_float_wav(const sd_audio_t* a, char* wav_path, size_t wav_path_sz) {
|
||||
if (!a || !a->data || a->sample_count == 0 || a->channels == 0 || a->sample_rate == 0) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
snprintf(wav_path, wav_path_sz, "/tmp/gosd-audio-XXXXXX.wav");
|
||||
const char* temp_dir = getenv("TMPDIR");
|
||||
if (!temp_dir || temp_dir[0] == '\0') {
|
||||
temp_dir = "/tmp";
|
||||
}
|
||||
int path_len = snprintf(wav_path, wav_path_sz, "%s/gosd-audio-XXXXXX.wav", temp_dir);
|
||||
if (path_len < 0 || (size_t)path_len >= wav_path_sz) {
|
||||
fprintf(stderr, "temporary directory path is too long\n");
|
||||
return -1;
|
||||
}
|
||||
int fd = mkstemps(wav_path, 4);
|
||||
if (fd < 0) { perror("mkstemps wav"); return -1; }
|
||||
FILE* f = fdopen(fd, "wb");
|
||||
if (!f) { perror("fdopen wav"); close(fd); return -1; }
|
||||
if (!f) { perror("fdopen wav"); close(fd); unlink(wav_path); return -1; }
|
||||
|
||||
uint64_t frames = a->sample_count;
|
||||
uint32_t channels = a->channels;
|
||||
@@ -1221,7 +1229,7 @@ static int ffmpeg_mux_raw_to_mp4(sd_image_t* frames, int num_frames, int fps,
|
||||
snprintf(fps_str, sizeof(fps_str), "%d", fps);
|
||||
|
||||
// Optional audio: write a temp WAV file if the model produced audio.
|
||||
char wav_path[64] = {0};
|
||||
char wav_path[4096] = {0};
|
||||
bool have_audio = false;
|
||||
if (audio && audio->data && audio->sample_count > 0 && audio->channels > 0 && audio->sample_rate > 0) {
|
||||
if (write_planar_float_wav(audio, wav_path, sizeof(wav_path)) == 0) {
|
||||
@@ -1438,4 +1446,3 @@ int unload() {
|
||||
free_sd_ctx(sd_c);
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -128,9 +128,40 @@ func parseOptions(opts *pb.ModelOptions) loadOptions {
|
||||
lo := loadOptions{}
|
||||
applyOptionsList(&lo, opts.GetOptions())
|
||||
applyEngineArgs(&lo, opts.GetEngineArgs())
|
||||
applyDraftModelOption(&lo, opts.GetOptions())
|
||||
return lo
|
||||
}
|
||||
|
||||
// applyDraftModelOption binds a managed companion snapshot after engine_args
|
||||
// has supplied the speculative document. Companion paths do not exist until
|
||||
// LocalAI materializes the artifact, so they must replace the gallery's static
|
||||
// repository reference without disturbing the method or token budget.
|
||||
func applyDraftModelOption(lo *loadOptions, options []string) {
|
||||
if strings.TrimSpace(lo.speculativeConfig) == "" {
|
||||
return
|
||||
}
|
||||
var draftModel string
|
||||
for _, option := range options {
|
||||
key, value, found := strings.Cut(option, ":")
|
||||
if found && strings.TrimSpace(key) == "draft_model" {
|
||||
draftModel = strings.TrimSpace(value)
|
||||
}
|
||||
}
|
||||
if draftModel == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var spec map[string]any
|
||||
if err := json.Unmarshal([]byte(lo.speculativeConfig), &spec); err != nil {
|
||||
return
|
||||
}
|
||||
spec["model"] = draftModel
|
||||
encoded, err := json.Marshal(spec)
|
||||
if err == nil {
|
||||
lo.speculativeConfig = string(encoded)
|
||||
}
|
||||
}
|
||||
|
||||
// applyOptionsList reads the legacy free-form "key:value" list. strings.Cut
|
||||
// splits on the FIRST colon only, so a JSON object value survives intact.
|
||||
func applyOptionsList(lo *loadOptions, options []string) {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
var _ = Describe("managed DFlash companion options", func() {
|
||||
It("replaces only the draft model in an existing speculative configuration", func() {
|
||||
managedPath := ".artifacts/huggingface/0123456789abcdef/snapshot"
|
||||
lo := parseOptions(&pb.ModelOptions{
|
||||
Options: []string{"draft_model:" + managedPath},
|
||||
EngineArgs: `{
|
||||
"speculative_config": {
|
||||
"method": "dflash",
|
||||
"model": "Mia-AiLab/Qwen3.8-27B-DFlash2-EXL3-5.0bpw",
|
||||
"num_speculative_tokens": 7
|
||||
}
|
||||
}`,
|
||||
})
|
||||
|
||||
Expect(lo.speculativeConfig).To(MatchJSON(`{
|
||||
"method": "dflash",
|
||||
"model": ".artifacts/huggingface/0123456789abcdef/snapshot",
|
||||
"num_speculative_tokens": 7
|
||||
}`))
|
||||
})
|
||||
|
||||
It("ignores a draft companion when speculative decoding is not configured", func() {
|
||||
lo := parseOptions(&pb.ModelOptions{
|
||||
Options: []string{"draft_model:.artifacts/huggingface/0123456789abcdef/snapshot"},
|
||||
})
|
||||
|
||||
Expect(lo.speculativeConfig).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=eacbd8234c6654cdbf2c377f72b2106875479bdc
|
||||
WHISPER_CPP_VERSION?=52a939a2a762224e255d366c1182b2af4dd1a032
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -19,6 +19,7 @@ import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from temp_utils import cleanup_paths
|
||||
|
||||
import tempfile
|
||||
|
||||
@@ -115,11 +116,6 @@ def merge_audio_files(audio_files, output_path, sample_rate):
|
||||
# Save the merged audio
|
||||
ta.save(output_path, merged_waveform, sample_rate)
|
||||
|
||||
# Clean up temporary files
|
||||
for audio_file in audio_files:
|
||||
if os.path.exists(audio_file):
|
||||
os.remove(audio_file)
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
@@ -226,19 +222,20 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
text_chunks = split_text_at_word_boundary(request.text, max_length=250)
|
||||
print(f"Splitting text into chunks of 250 characters: {len(text_chunks)}", file=sys.stderr)
|
||||
# Generate audio for each chunk
|
||||
temp_audio_files = []
|
||||
for i, chunk in enumerate(text_chunks):
|
||||
# Generate audio for this chunk
|
||||
wav = self.model.generate(chunk, **kwargs)
|
||||
|
||||
# Create temporary file for this chunk
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
|
||||
temp_file.close()
|
||||
ta.save(temp_file.name, wav, self.model.sr)
|
||||
temp_audio_files.append(temp_file.name)
|
||||
|
||||
# Merge all audio files
|
||||
merge_audio_files(temp_audio_files, request.dst, self.model.sr)
|
||||
with cleanup_paths() as temp_audio_files:
|
||||
for i, chunk in enumerate(text_chunks):
|
||||
# Generate audio for this chunk
|
||||
wav = self.model.generate(chunk, **kwargs)
|
||||
|
||||
# Register ownership before saving so a partial write is
|
||||
# removed too when generation or encoding fails.
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
|
||||
temp_file.close()
|
||||
temp_audio_files.append(temp_file.name)
|
||||
ta.save(temp_file.name, wav, self.model.sr)
|
||||
|
||||
# Merge all audio files
|
||||
merge_audio_files(temp_audio_files, request.dst, self.model.sr)
|
||||
else:
|
||||
# Generate audio using ChatterboxTTS for short text
|
||||
wav = self.model.generate(request.text, **kwargs)
|
||||
|
||||
@@ -37,6 +37,46 @@ def parse_options(options_list):
|
||||
return opts
|
||||
|
||||
|
||||
def attach_media_parts(messages_dicts, n_images=0, n_videos=0):
|
||||
"""Rebuild the last user message as content *parts* carrying media markers.
|
||||
|
||||
Backends that let the tokenizer do the templating hand plain string content
|
||||
to ``apply_chat_template``, but a chat template only emits the model's own
|
||||
media tokens (``<|vision_start|><|image_pad|><|vision_end|>`` for the
|
||||
Qwen-VL family, and the equivalents elsewhere) when the content is a list
|
||||
of parts. Without those markers the engine's multimodal processor finds
|
||||
nothing to substitute and silently discards the pixels, even though they
|
||||
were forwarded correctly out of band.
|
||||
|
||||
Returns a new list whose last user message has
|
||||
``[{"type": "image"} * n_images, {"type": "video"} * n_videos, text]`` as
|
||||
its content, or ``None`` when there is nothing to attach - no media, no
|
||||
user turn, or content that is already a list of parts - so the caller can
|
||||
keep using the original string-content list.
|
||||
"""
|
||||
if not n_images and not n_videos:
|
||||
return None
|
||||
idx = next(
|
||||
(
|
||||
i
|
||||
for i in reversed(range(len(messages_dicts)))
|
||||
if messages_dicts[i].get("role") == "user"
|
||||
),
|
||||
None,
|
||||
)
|
||||
if idx is None:
|
||||
return None
|
||||
text = messages_dicts[idx].get("content") or ""
|
||||
if not isinstance(text, str):
|
||||
return None
|
||||
parts = [{"type": "image"}] * n_images + [{"type": "video"}] * n_videos
|
||||
if text:
|
||||
parts.append({"type": "text", "text": text})
|
||||
patched = list(messages_dicts)
|
||||
patched[idx] = dict(patched[idx], content=parts)
|
||||
return patched
|
||||
|
||||
|
||||
def messages_to_dicts(proto_messages):
|
||||
"""Convert proto ``Message`` objects to dicts suitable for ``apply_chat_template``.
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ import json
|
||||
import types
|
||||
import unittest
|
||||
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
from python_utils import attach_media_parts, messages_to_dicts, parse_options
|
||||
|
||||
|
||||
def _msg(**fields):
|
||||
@@ -118,5 +118,63 @@ class TestMessagesToDicts(unittest.TestCase):
|
||||
self.assertNotIn("tool_calls", out[0])
|
||||
|
||||
|
||||
class TestAttachMediaParts(unittest.TestCase):
|
||||
def test_image_marker_added_to_last_user_turn(self):
|
||||
messages = [
|
||||
{"role": "system", "content": "be brief"},
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "assistant", "content": "ok"},
|
||||
{"role": "user", "content": "how high is the water?"},
|
||||
]
|
||||
out = attach_media_parts(messages, n_images=1)
|
||||
self.assertEqual(
|
||||
out[3]["content"],
|
||||
[{"type": "image"}, {"type": "text", "text": "how high is the water?"}],
|
||||
)
|
||||
# Earlier turns and the input list itself are untouched.
|
||||
self.assertEqual(out[:3], messages[:3])
|
||||
self.assertEqual(messages[3]["content"], "how high is the water?")
|
||||
|
||||
def test_counts_and_order_images_then_videos(self):
|
||||
out = attach_media_parts(
|
||||
[{"role": "user", "content": "describe"}], n_images=2, n_videos=1
|
||||
)
|
||||
self.assertEqual(
|
||||
out[0]["content"],
|
||||
[
|
||||
{"type": "image"},
|
||||
{"type": "image"},
|
||||
{"type": "video"},
|
||||
{"type": "text", "text": "describe"},
|
||||
],
|
||||
)
|
||||
|
||||
def test_empty_text_yields_media_only_parts(self):
|
||||
out = attach_media_parts([{"role": "user", "content": ""}], n_images=1)
|
||||
self.assertEqual(out[0]["content"], [{"type": "image"}])
|
||||
|
||||
def test_other_message_keys_are_preserved(self):
|
||||
out = attach_media_parts(
|
||||
[{"role": "user", "content": "hi", "name": "bob"}], n_images=1
|
||||
)
|
||||
self.assertEqual(out[0]["name"], "bob")
|
||||
|
||||
def test_no_media_is_a_no_op(self):
|
||||
self.assertIsNone(attach_media_parts([{"role": "user", "content": "hi"}]))
|
||||
|
||||
def test_no_user_turn_is_a_no_op(self):
|
||||
self.assertIsNone(
|
||||
attach_media_parts([{"role": "system", "content": "hi"}], n_images=1)
|
||||
)
|
||||
|
||||
def test_content_already_parts_is_a_no_op(self):
|
||||
self.assertIsNone(
|
||||
attach_media_parts(
|
||||
[{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||||
n_images=1,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,36 @@
|
||||
import base64
|
||||
import contextlib
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def materialize_base64(data, suffix=""):
|
||||
"""Materialize base64 data for a path-only library and always remove it."""
|
||||
descriptor, path = tempfile.mkstemp(prefix="localai-media-", suffix=suffix)
|
||||
try:
|
||||
with os.fdopen(descriptor, "wb") as output:
|
||||
descriptor = None
|
||||
output.write(base64.b64decode(data))
|
||||
yield path
|
||||
finally:
|
||||
if descriptor is not None:
|
||||
os.close(descriptor)
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def cleanup_paths():
|
||||
"""Collect temporary paths and remove them on success or failure."""
|
||||
paths = []
|
||||
try:
|
||||
yield paths
|
||||
finally:
|
||||
for path in paths:
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from temp_utils import cleanup_paths, materialize_base64
|
||||
|
||||
|
||||
class MaterializeBase64Test(unittest.TestCase):
|
||||
def test_removes_materialized_file_after_success(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
with mock.patch.object(tempfile, "tempdir", directory):
|
||||
with materialize_base64("aGVsbG8=", suffix=".data") as path:
|
||||
with open(path, "rb") as materialized:
|
||||
self.assertEqual(materialized.read(), b"hello")
|
||||
self.assertFalse(os.path.exists(path))
|
||||
|
||||
def test_removes_materialized_file_when_consumer_fails(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
with mock.patch.object(tempfile, "tempdir", directory):
|
||||
with self.assertRaisesRegex(RuntimeError, "decode failed"):
|
||||
with materialize_base64("aGVsbG8="):
|
||||
raise RuntimeError("decode failed")
|
||||
self.assertEqual(os.listdir(directory), [])
|
||||
|
||||
|
||||
class CleanupPathsTest(unittest.TestCase):
|
||||
def test_removes_every_registered_path_after_failure(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
paths = [os.path.join(directory, name) for name in ("one.wav", "two.wav")]
|
||||
with self.assertRaisesRegex(RuntimeError, "merge failed"):
|
||||
with cleanup_paths() as registered:
|
||||
for path in paths:
|
||||
open(path, "wb").close()
|
||||
registered.append(path)
|
||||
raise RuntimeError("merge failed")
|
||||
self.assertEqual(os.listdir(directory), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,4 +1,4 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/xpu
|
||||
torch==2.13.0+xpu
|
||||
torch==2.14.0+xpu
|
||||
oneccl_bind_pt==2.8.0+xpu
|
||||
optimum[openvino]
|
||||
@@ -1,3 +1,3 @@
|
||||
grpcio==1.82.1
|
||||
grpcio==1.83.1
|
||||
protobuf
|
||||
grpcio-tools
|
||||
@@ -1,4 +1,4 @@
|
||||
grpcio==1.83.0
|
||||
grpcio==1.83.1
|
||||
protobuf
|
||||
certifi
|
||||
packaging==26.3
|
||||
@@ -122,6 +122,21 @@ from diffusers.schedulers import (
|
||||
UniPCMultistepScheduler,
|
||||
)
|
||||
|
||||
def select_device(request_cuda, device_option, cuda_available, xpu, mps_available):
|
||||
"""Pick the pipeline device. An explicit `device:` model option wins;
|
||||
otherwise CUDA is used whenever torch reports it available (ROCm
|
||||
builds included) or the model config forces it with `cuda: true`,
|
||||
keeping the pre-existing XPU/MPS overrides. CPU is the fallback, not
|
||||
the default."""
|
||||
if device_option:
|
||||
return device_option
|
||||
device = "cuda" if (request_cuda or cuda_available) else "cpu"
|
||||
if xpu:
|
||||
device = "xpu"
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
return device
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
@@ -627,12 +642,13 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
# modify LoraAdapter to be relative to modelFileBase
|
||||
request.LoraAdapter = os.path.join(request.ModelPath, request.LoraAdapter)
|
||||
|
||||
device = "cpu" if not request.CUDA else "cuda"
|
||||
if XPU:
|
||||
device = "xpu"
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
device = select_device(
|
||||
request.CUDA,
|
||||
self.options.pop("device", None),
|
||||
torch.cuda.is_available(),
|
||||
XPU,
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available(),
|
||||
)
|
||||
self.device = device
|
||||
if request.LoraAdapter:
|
||||
# Check if its a local file and not a directory ( we load lora differently for a safetensor file )
|
||||
@@ -800,12 +816,12 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
image = image.resize((1024, 576))
|
||||
|
||||
generator = torch.manual_seed(request.seed)
|
||||
frames = self.pipe(image, guidance_scale=self.cfg_scale, decode_chunk_size=CHUNK_SIZE, generator=generator).frames[0]
|
||||
frames = self.pipe(image=image, guidance_scale=self.cfg_scale, decode_chunk_size=CHUNK_SIZE, generator=generator).frames[0]
|
||||
export_to_video(frames, request.dst, fps=FPS)
|
||||
return backend_pb2.Result(message="Media generated successfully", success=True)
|
||||
|
||||
if self.txt2vid:
|
||||
video_frames = self.pipe(prompt, guidance_scale=self.cfg_scale, num_inference_steps=steps, num_frames=int(FRAMES)).frames
|
||||
video_frames = self.pipe(prompt=prompt, guidance_scale=self.cfg_scale, num_inference_steps=steps, num_frames=int(FRAMES)).frames
|
||||
export_to_video(video_frames, request.dst)
|
||||
return backend_pb2.Result(message="Media generated successfully", success=True)
|
||||
|
||||
@@ -868,7 +884,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
else:
|
||||
# pass the kwargs dictionary to the self.pipe method
|
||||
image = self.pipe(
|
||||
prompt,
|
||||
prompt=prompt,
|
||||
guidance_scale=self.cfg_scale,
|
||||
**kwargs
|
||||
).images[0]
|
||||
|
||||
@@ -7,6 +7,7 @@ import time
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Import dynamic loader for testing (these don't need gRPC)
|
||||
import backend
|
||||
import diffusers_dynamic_loader as loader
|
||||
from diffusers import DiffusionPipeline, StableDiffusionPipeline
|
||||
|
||||
@@ -373,3 +374,74 @@ class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
|
||||
finally:
|
||||
os.unlink(src_file.name)
|
||||
os.unlink(dst_file.name)
|
||||
|
||||
def test_text_to_image_prompt_is_passed_by_keyword(self):
|
||||
"""Test compatibility with pipelines that take image before prompt."""
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from backend import BackendServicer
|
||||
|
||||
class Flux2CompatiblePipeline:
|
||||
"""Model the FLUX.2 call signature: image is before prompt."""
|
||||
|
||||
def __call__(self, image=None, prompt=None, **kwargs):
|
||||
if prompt is None:
|
||||
raise ValueError("prompt was not passed by keyword")
|
||||
self.prompt = prompt
|
||||
self.kwargs = kwargs
|
||||
return MagicMock(images=[Image.new("RGB", (4, 4))])
|
||||
|
||||
pipeline = Flux2CompatiblePipeline()
|
||||
svc = BackendServicer.__new__(BackendServicer)
|
||||
svc.pipe = pipeline
|
||||
svc.cfg_scale = 7.5
|
||||
svc.controlnet = None
|
||||
svc.img2vid = False
|
||||
svc.txt2vid = False
|
||||
svc.clip_skip = 0
|
||||
svc.PipelineType = "Flux2KleinPipeline"
|
||||
svc.options = {}
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as dst_file:
|
||||
dst_path = dst_file.name
|
||||
|
||||
try:
|
||||
request = MagicMock()
|
||||
request.positive_prompt = "a red apple on a wooden table"
|
||||
request.negative_prompt = ""
|
||||
request.step = 4
|
||||
request.seed = 0
|
||||
request.width = 0
|
||||
request.height = 0
|
||||
request.src = ""
|
||||
request.ref_images = []
|
||||
request.dst = dst_path
|
||||
|
||||
svc.GenerateImage(request, context=None)
|
||||
|
||||
self.assertEqual(pipeline.prompt, request.positive_prompt)
|
||||
self.assertEqual(pipeline.kwargs["num_inference_steps"], 4)
|
||||
finally:
|
||||
os.unlink(dst_path)
|
||||
|
||||
|
||||
class TestDeviceSelection(unittest.TestCase):
|
||||
"""Unit tests for backend.select_device (no GPU required)."""
|
||||
|
||||
def test_autodetect_cuda(self):
|
||||
self.assertEqual(backend.select_device(False, None, True, False, False), "cuda")
|
||||
|
||||
def test_cpu_fallback(self):
|
||||
self.assertEqual(backend.select_device(False, None, False, False, False), "cpu")
|
||||
|
||||
def test_forced_cuda(self):
|
||||
self.assertEqual(backend.select_device(True, None, False, False, False), "cuda")
|
||||
|
||||
def test_device_option_wins(self):
|
||||
self.assertEqual(backend.select_device(True, "cpu", True, True, True), "cpu")
|
||||
|
||||
def test_mps_overrides(self):
|
||||
self.assertEqual(backend.select_device(False, None, True, False, True), "mps")
|
||||
@@ -11,7 +11,7 @@ RPC. It supports:
|
||||
systems such as NVIDIA DGX Spark.
|
||||
|
||||
Install the `longcat-video` or `longcat-video-avatar-1.5` recipe from the
|
||||
LocalAI Model Gallery. See the [LongCat user guide](../../../docs/content/features/longcat-video.md)
|
||||
LocalAI Model Gallery. LongCat video backend
|
||||
for Studio and API examples, hardware requirements, and manual configuration.
|
||||
|
||||
The upstream source is pinned in `Makefile` and patched at build time. The
|
||||
|
||||
@@ -6,6 +6,7 @@ import datetime
|
||||
import gc
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -888,6 +889,13 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
def _release_model(self):
|
||||
self.pipeline = None
|
||||
self.model_kind = None
|
||||
try:
|
||||
if hasattr(self, "dist") and self.dist.is_initialized():
|
||||
self.dist.destroy_process_group()
|
||||
finally:
|
||||
if self._dist_store_dir is not None:
|
||||
shutil.rmtree(self._dist_store_dir, ignore_errors=True)
|
||||
self._dist_store_dir = None
|
||||
gc.collect()
|
||||
if hasattr(self, "torch") and self.torch.cuda.is_available():
|
||||
self.torch.cuda.empty_cache()
|
||||
|
||||
@@ -18,6 +18,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from model_utils import resolve_model_reference
|
||||
from device_utils import device_map_for, select_device
|
||||
|
||||
|
||||
|
||||
@@ -95,13 +96,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
||||
|
||||
def LoadModel(self, request, context):
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
else:
|
||||
device = "cpu"
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
device = select_device(torch)
|
||||
if not torch.cuda.is_available() and request.CUDA:
|
||||
return backend_pb2.Result(success=False, message="CUDA is not available")
|
||||
|
||||
@@ -123,7 +118,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
model_path, local_only = resolve_model_reference(
|
||||
request, "Qwen/Qwen3-ASR-1.7B"
|
||||
)
|
||||
default_dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
|
||||
default_dtype = torch.bfloat16 if self.device in ("cuda", "xpu") else torch.float32
|
||||
load_dtype = default_dtype
|
||||
if "torch_dtype" in self.options:
|
||||
d = str(self.options["torch_dtype"]).lower()
|
||||
@@ -145,12 +140,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
if attn_implementation is not None and isinstance(attn_implementation, str):
|
||||
attn_implementation = attn_implementation.strip() or None
|
||||
|
||||
if self.device == "mps":
|
||||
device_map = None
|
||||
elif self.device == "cuda":
|
||||
device_map = "cuda:0"
|
||||
else:
|
||||
device_map = "cpu"
|
||||
device_map = device_map_for(self.device)
|
||||
|
||||
load_kwargs = dict(
|
||||
dtype=load_dtype,
|
||||
@@ -423,4 +413,4 @@ if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run the gRPC server.")
|
||||
parser.add_argument("--addr", default="localhost:50051", help="The address to bind the server to.")
|
||||
args = parser.parse_args()
|
||||
serve(args.addr)
|
||||
serve(args.addr)
|
||||
@@ -0,0 +1,18 @@
|
||||
def select_device(torch_module):
|
||||
mps = getattr(getattr(torch_module, "backends", None), "mps", None)
|
||||
if mps is not None and mps.is_available():
|
||||
return "mps"
|
||||
if torch_module.cuda.is_available():
|
||||
return "cuda"
|
||||
xpu = getattr(torch_module, "xpu", None)
|
||||
if xpu is not None and xpu.is_available():
|
||||
return "xpu"
|
||||
return "cpu"
|
||||
|
||||
|
||||
def device_map_for(device):
|
||||
if device == "mps":
|
||||
return None
|
||||
if device in ("cuda", "xpu"):
|
||||
return f"{device}:0"
|
||||
return "cpu"
|
||||
@@ -0,0 +1,58 @@
|
||||
import unittest
|
||||
|
||||
from device_utils import device_map_for, select_device
|
||||
|
||||
|
||||
class Availability:
|
||||
def __init__(self, available):
|
||||
self._available = available
|
||||
|
||||
def is_available(self):
|
||||
return self._available
|
||||
|
||||
|
||||
class TorchStub:
|
||||
def __init__(self, *, cuda=False, mps=False, xpu=False):
|
||||
self.cuda = Availability(cuda)
|
||||
self.backends = type("Backends", (), {"mps": Availability(mps)})()
|
||||
self.xpu = Availability(xpu)
|
||||
|
||||
|
||||
class SelectDeviceTest(unittest.TestCase):
|
||||
def test_preserves_cuda_selection(self):
|
||||
torch_module = TorchStub(cuda=True)
|
||||
|
||||
self.assertEqual(select_device(torch_module), "cuda")
|
||||
|
||||
def test_preserves_mps_selection(self):
|
||||
torch_module = TorchStub(mps=True)
|
||||
|
||||
self.assertEqual(select_device(torch_module), "mps")
|
||||
|
||||
def test_selects_xpu_when_intel_gpu_is_available(self):
|
||||
torch_module = TorchStub(xpu=True)
|
||||
|
||||
self.assertEqual(select_device(torch_module), "xpu")
|
||||
|
||||
def test_falls_back_to_cpu(self):
|
||||
torch_module = TorchStub()
|
||||
|
||||
self.assertEqual(select_device(torch_module), "cpu")
|
||||
|
||||
|
||||
class DeviceMapTest(unittest.TestCase):
|
||||
def test_preserves_cuda_model_placement(self):
|
||||
self.assertEqual(device_map_for("cuda"), "cuda:0")
|
||||
|
||||
def test_preserves_mps_model_placement(self):
|
||||
self.assertIsNone(device_map_for("mps"))
|
||||
|
||||
def test_places_the_model_on_the_first_xpu(self):
|
||||
self.assertEqual(device_map_for("xpu"), "xpu:0")
|
||||
|
||||
def test_preserves_cpu_model_placement(self):
|
||||
self.assertEqual(device_map_for("cpu"), "cpu")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,3 +1,3 @@
|
||||
grpcio==1.82.1
|
||||
grpcio==1.83.1
|
||||
protobuf
|
||||
certifi
|
||||
@@ -40,6 +40,7 @@ import grpc
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from python_utils import attach_media_parts
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from model_utils import resolve_model_reference
|
||||
|
||||
@@ -90,6 +91,14 @@ except Exception:
|
||||
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# proto3 has no field presence, so an explicit 0 is indistinguishable from
|
||||
# "unset" and the zero-filter below would drop it. These two fields have a
|
||||
# meaningful zero a caller can actually intend: temperature 0 is greedy
|
||||
# decoding, and 0 is a valid seed. Silently substituting a default for either
|
||||
# turns a reproducible request into a random one.
|
||||
_EXPLICIT_ZERO_FIELDS = ("Temperature", "Seed")
|
||||
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
|
||||
@@ -323,7 +332,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
if not hasattr(request, proto_field):
|
||||
continue
|
||||
value = getattr(request, proto_field)
|
||||
if proto_field != "Temperature" and value in (None, 0, 0.0, [], False, ""):
|
||||
if proto_field not in _EXPLICIT_ZERO_FIELDS and value in (None, 0, 0.0, [], False, ""):
|
||||
continue
|
||||
# repeated fields come back as RepeatedScalarContainer — convert
|
||||
if hasattr(value, "__iter__") and not isinstance(value, (str, bytes)):
|
||||
@@ -367,6 +376,24 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
if _thinking in ("true", "false"):
|
||||
template_kwargs["enable_thinking"] = (_thinking == "true")
|
||||
|
||||
# sglang locates the attached images/videos by scanning the rendered
|
||||
# prompt for the model's own media token, so the template has to be
|
||||
# given content *parts* - string content renders a prompt with no
|
||||
# placeholder and the media are dropped without a word (#11621).
|
||||
media_dicts = attach_media_parts(
|
||||
messages_dicts, len(request.Images), len(request.Videos)
|
||||
)
|
||||
if media_dicts is not None:
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(media_dicts, **template_kwargs)
|
||||
except Exception as e:
|
||||
# A text-only template cannot iterate content parts; fall
|
||||
# through to the text-only prompt instead of failing.
|
||||
print(
|
||||
f"chat template rejected multimodal content parts: {e!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
try:
|
||||
return self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
except TypeError:
|
||||
@@ -374,10 +401,67 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
messages_dicts, tokenize=False, add_generation_prompt=True,
|
||||
)
|
||||
|
||||
def _make_parsers(self, request):
|
||||
def _new_reasoning_parser(self, stream_reasoning: bool, prompt: str = "",
|
||||
grammar_constrained: bool = False):
|
||||
"""Build a ReasoningParser for one request, or None.
|
||||
|
||||
Reasoning templates come in two flavours. Some let the model emit the
|
||||
opening tag, others put it into the *prompt* — Qwen3's template appends
|
||||
``<think>`` when thinking is on, so the completion starts straight in
|
||||
the reasoning block and only the closing ``</think>`` ever shows up.
|
||||
sglang's detector keys off the opening tag, so in that second case it
|
||||
classifies the whole completion as normal content and
|
||||
``reasoning_content`` stays empty.
|
||||
|
||||
sglang's own OpenAI server covers this with
|
||||
``template_manager.force_reasoning``; this backend has no template
|
||||
manager, so it derives the same signal from the rendered prompt.
|
||||
``force_reasoning`` is only passed when we mean True, leaving detector
|
||||
defaults (e.g. DeepSeek-R1's built-in True) untouched.
|
||||
|
||||
``grammar_constrained`` suppresses the prefill heuristic. A structured
|
||||
decoding constraint applies from the first token, so the model cannot
|
||||
emit the closing tag even though the template opened the block: the
|
||||
whole completion is schema output and belongs in ``content``. Forcing
|
||||
there files the answer as reasoning and leaves content empty. sglang's
|
||||
own server keeps the two apart for the same reason — its grammar
|
||||
backend owns the reasoning prefix when a reasoning parser is set.
|
||||
"""
|
||||
if grammar_constrained:
|
||||
prompt = ""
|
||||
|
||||
if not (HAS_REASONING_PARSERS and self.reasoning_parser_name):
|
||||
return None
|
||||
|
||||
kwargs = {
|
||||
"model_type": self.reasoning_parser_name,
|
||||
"stream_reasoning": stream_reasoning,
|
||||
}
|
||||
try:
|
||||
parser = ReasoningParser(**kwargs)
|
||||
except Exception as e:
|
||||
print(f"ReasoningParser init failed: {e!r}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
start = getattr(getattr(parser, "detector", None), "think_start_token", None)
|
||||
if start and prompt and prompt.rstrip().endswith(start):
|
||||
try:
|
||||
parser = ReasoningParser(force_reasoning=True, **kwargs)
|
||||
except TypeError:
|
||||
# sglang without the force_reasoning kwarg: keep the default
|
||||
# parser rather than failing the request.
|
||||
pass
|
||||
except Exception as e:
|
||||
print(
|
||||
f"ReasoningParser(force_reasoning=True) failed: {e!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
def _make_parsers(self, request, prompt: str = ""):
|
||||
"""Construct fresh per-request parser instances (stateful)."""
|
||||
tool_parser = None
|
||||
reasoning_parser = None
|
||||
|
||||
if HAS_TOOL_PARSERS and self.tool_parser_name and request.Tools:
|
||||
try:
|
||||
@@ -389,14 +473,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
except Exception as e:
|
||||
print(f"FunctionCallParser init failed: {e!r}", file=sys.stderr)
|
||||
|
||||
if HAS_REASONING_PARSERS and self.reasoning_parser_name:
|
||||
try:
|
||||
reasoning_parser = ReasoningParser(
|
||||
model_type=self.reasoning_parser_name,
|
||||
stream_reasoning=True,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"ReasoningParser init failed: {e!r}", file=sys.stderr)
|
||||
reasoning_parser = self._new_reasoning_parser(
|
||||
True, prompt, bool(getattr(request, "Grammar", "")),
|
||||
)
|
||||
|
||||
return tool_parser, reasoning_parser
|
||||
|
||||
@@ -404,7 +483,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
sampling_params = self._build_sampling_params(request)
|
||||
prompt = self._build_prompt(request)
|
||||
|
||||
tool_parser, reasoning_parser = self._make_parsers(request)
|
||||
tool_parser, reasoning_parser = self._make_parsers(request, prompt)
|
||||
|
||||
image_data = list(request.Images) if request.Images else None
|
||||
video_data = list(request.Videos) if request.Videos else None
|
||||
@@ -500,15 +579,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
final_tool_calls: List[backend_pb2.ToolCallDelta] = []
|
||||
|
||||
if not streaming:
|
||||
final_reasoning_parser = None
|
||||
if HAS_REASONING_PARSERS and self.reasoning_parser_name:
|
||||
try:
|
||||
final_reasoning_parser = ReasoningParser(
|
||||
model_type=self.reasoning_parser_name,
|
||||
stream_reasoning=False,
|
||||
)
|
||||
except Exception:
|
||||
final_reasoning_parser = None
|
||||
final_reasoning_parser = self._new_reasoning_parser(
|
||||
False, prompt, bool(getattr(request, "Grammar", "")),
|
||||
)
|
||||
|
||||
if final_reasoning_parser is not None:
|
||||
try:
|
||||
|
||||
@@ -128,11 +128,66 @@ class TestSglangHelpers(unittest.TestCase):
|
||||
self.assertNotIn("enable_thinking", kwargs_for({}))
|
||||
self.assertIs(kwargs_for({"enable_thinking": "FALSE"})["enable_thinking"], False)
|
||||
|
||||
def test_explicit_zero_temperature_is_preserved(self):
|
||||
"""Temperature=0 is valid greedy decoding, not an unset value."""
|
||||
def test_reasoning_parser_forced_when_template_prefills_think_tag(self):
|
||||
"""Qwen3's template puts ``<think>`` in the prompt, so the completion
|
||||
never contains it. Without force_reasoning the detector treats the whole
|
||||
completion as normal text and reasoning_content stays empty."""
|
||||
servicer = self._servicer()
|
||||
servicer.reasoning_parser_name = "qwen3"
|
||||
|
||||
# What the model actually emits when the prompt ends in "<think>".
|
||||
completion = "adding two and two</think>4"
|
||||
|
||||
forced = servicer._new_reasoning_parser(False, prompt="user: hi\n<think>\n")
|
||||
reasoning, content = forced.parse_non_stream(completion)
|
||||
self.assertEqual(reasoning, "adding two and two")
|
||||
self.assertEqual(content, "4")
|
||||
|
||||
# No prefilled tag in the prompt: detector default, unchanged behaviour.
|
||||
unforced = servicer._new_reasoning_parser(False, prompt="user: hi\n")
|
||||
reasoning, content = unforced.parse_non_stream(completion)
|
||||
self.assertFalse(reasoning)
|
||||
self.assertEqual(content, completion)
|
||||
|
||||
def test_reasoning_parser_not_forced_when_thinking_is_off(self):
|
||||
"""Thinking off means no ``<think>`` in the prompt either, so the answer
|
||||
must not be swallowed into reasoning_content."""
|
||||
servicer = self._servicer()
|
||||
servicer.reasoning_parser_name = "qwen3"
|
||||
|
||||
parser = servicer._new_reasoning_parser(False, prompt="user: primes?\n")
|
||||
reasoning, content = parser.parse_non_stream("2,3,5,7,11")
|
||||
self.assertFalse(reasoning)
|
||||
self.assertEqual(content, "2,3,5,7,11")
|
||||
|
||||
def test_grammar_constrained_output_is_not_forced_into_reasoning(self):
|
||||
"""Structured decoding applies from the first token, so the model cannot
|
||||
emit the closing tag even though the template opened the block. The whole
|
||||
completion is schema output and must stay in content."""
|
||||
servicer = self._servicer()
|
||||
servicer.reasoning_parser_name = "qwen3"
|
||||
|
||||
schema_out = '{"findings": [{"line": 42, "issue": "off-by-one"}]}'
|
||||
parser = servicer._new_reasoning_parser(
|
||||
False, prompt="audit this\n<think>\n", grammar_constrained=True,
|
||||
)
|
||||
reasoning, content = parser.parse_non_stream(schema_out)
|
||||
self.assertFalse(reasoning)
|
||||
self.assertEqual(content, schema_out)
|
||||
|
||||
def test_reasoning_parser_absent_without_configured_parser(self):
|
||||
servicer = self._servicer()
|
||||
servicer.reasoning_parser_name = None
|
||||
self.assertIsNone(servicer._new_reasoning_parser(False, prompt="<think>"))
|
||||
|
||||
def test_explicit_zero_temperature_and_seed_are_preserved(self):
|
||||
"""Temperature=0 is greedy decoding and 0 is a valid seed — neither is
|
||||
an unset value. A dropped seed turns a reproducible request random."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
servicer = self._servicer()
|
||||
import sys as _sys
|
||||
_SEED_KEY_FOR_TEST = _sys.modules["backend"]._SEED_KEY
|
||||
request = SimpleNamespace(
|
||||
Temperature=0,
|
||||
N=0,
|
||||
@@ -154,8 +209,12 @@ class TestSglangHelpers(unittest.TestCase):
|
||||
|
||||
params = servicer._build_sampling_params(request)
|
||||
self.assertEqual(params["temperature"], 0)
|
||||
# Other protobuf-default scalar fields must remain filtered.
|
||||
self.assertEqual(params[_SEED_KEY_FOR_TEST], 0)
|
||||
# Other protobuf-default scalar fields must remain filtered. top_k=0 in
|
||||
# particular is not a value sglang accepts (-1 disables it), so it must
|
||||
# keep falling through to the engine default.
|
||||
self.assertNotIn("top_p", params)
|
||||
self.assertNotIn("top_k", params)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -19,7 +19,6 @@ import base64
|
||||
import io
|
||||
import json
|
||||
import gc
|
||||
import tempfile
|
||||
|
||||
from PIL import Image
|
||||
import torch
|
||||
@@ -34,6 +33,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from model_utils import resolve_model_reference
|
||||
from temp_utils import materialize_base64
|
||||
from vllm_utils import parse_options, messages_to_dicts, setup_parsers
|
||||
|
||||
|
||||
@@ -118,13 +118,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
return video_to_ndarrays(video_path, num_frames=16)
|
||||
# Try base64 decode
|
||||
try:
|
||||
timestamp = str(int(time.time() * 1000))
|
||||
p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data")
|
||||
with open(p, "wb") as f:
|
||||
f.write(base64.b64decode(video_path))
|
||||
video = VideoAsset(name=p).np_ndarrays
|
||||
os.remove(p)
|
||||
return video
|
||||
with materialize_base64(video_path, suffix=".data") as path:
|
||||
return VideoAsset(name=path).np_ndarrays
|
||||
except:
|
||||
return None
|
||||
|
||||
@@ -136,15 +131,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
return (audio_signal.astype(np.float32), sr)
|
||||
# Try base64 decode
|
||||
try:
|
||||
audio_data = base64.b64decode(audio_path)
|
||||
# Save to temp file and load
|
||||
timestamp = str(int(time.time() * 1000))
|
||||
p = os.path.join(tempfile.gettempdir(), f"audio-{timestamp}.wav")
|
||||
with open(p, "wb") as f:
|
||||
f.write(audio_data)
|
||||
audio_signal, sr = librosa.load(p, sr=16000)
|
||||
os.remove(p)
|
||||
return (audio_signal.astype(np.float32), sr)
|
||||
with materialize_base64(audio_path, suffix=".wav") as path:
|
||||
audio_signal, sr = librosa.load(path, sr=16000)
|
||||
return (audio_signal.astype(np.float32), sr)
|
||||
except:
|
||||
return None
|
||||
|
||||
|
||||
+115
-20
@@ -10,7 +10,6 @@ import os
|
||||
import json
|
||||
import time
|
||||
import gc
|
||||
import tempfile
|
||||
from typing import List
|
||||
from PIL import Image
|
||||
|
||||
@@ -20,8 +19,10 @@ import backend_pb2_grpc
|
||||
import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from python_utils import attach_media_parts
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from model_utils import resolve_model_reference
|
||||
from temp_utils import materialize_base64
|
||||
from vllm_utils import apply_options_to_engine_args, normalize_option_key
|
||||
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
@@ -60,6 +61,12 @@ except ImportError:
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# proto3 has no field presence, so an explicit 0 is indistinguishable from
|
||||
# "unset". These two fields have a meaningful zero a caller can intend:
|
||||
# temperature 0 is greedy decoding, and 0 is a valid seed.
|
||||
_EXPLICIT_ZERO_FIELDS = ("Temperature", "Seed")
|
||||
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
@@ -553,11 +560,80 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
for request_field, param_field in request_to_sampling_params.items():
|
||||
if hasattr(request, request_field):
|
||||
value = getattr(request, request_field)
|
||||
if request_field == "Temperature" or value not in (None, 0, [], False):
|
||||
# See _EXPLICIT_ZERO_FIELDS: temperature 0 is greedy decoding
|
||||
# and 0 is a valid seed, so neither may be filtered out.
|
||||
if request_field in _EXPLICIT_ZERO_FIELDS or value not in (None, 0, [], False):
|
||||
setattr(sampling_params, param_field, value)
|
||||
|
||||
return sampling_params
|
||||
|
||||
def _new_reasoning_parser(self, chat_template_kwargs):
|
||||
"""Build the reasoning parser, telling it whether thinking is on.
|
||||
|
||||
vLLM's newer parser engines decide their *initial state* from
|
||||
``chat_template_kwargs``: ``Qwen3Parser`` reads
|
||||
``chat_template_kwargs["enable_thinking"]`` and defaults to ``True``,
|
||||
starting in the REASONING state. Constructed without it, a completion
|
||||
produced with thinking disabled is classified as reasoning end to end,
|
||||
and the answer is reported in both ``reasoning_content`` and
|
||||
``content``.
|
||||
|
||||
vLLM's own OpenAI server forwards the request's chat template kwargs
|
||||
here; this backend renders the template itself, so it forwards the
|
||||
same dict. Older parsers do not accept the argument — fall back to the
|
||||
plain constructor for those.
|
||||
"""
|
||||
try:
|
||||
return self.reasoning_parser_cls(
|
||||
self.tokenizer, chat_template_kwargs=chat_template_kwargs or {},
|
||||
)
|
||||
except TypeError:
|
||||
return self.reasoning_parser_cls(self.tokenizer)
|
||||
|
||||
@staticmethod
|
||||
def _split_reasoning(rp, generated_text, prompt, reasoning, content):
|
||||
"""Decide what the reasoning parser's output actually means.
|
||||
|
||||
Covers the *older* parser shape, which has no initial state to set:
|
||||
``BaseThinkingReasoningParser.extract_reasoning`` documents its own
|
||||
fallback — "For models that may not generate start token, assume the
|
||||
reasoning content is always at the start." When no end token is
|
||||
present it returns *everything* as reasoning and ``None`` as content,
|
||||
which is right for a truncated reasoning run and wrong for a
|
||||
completion that never contained reasoning at all.
|
||||
|
||||
Taking ``None`` content to mean "keep the raw text" then duplicates
|
||||
the answer into both fields.
|
||||
|
||||
The prompt says which case it is. A template with thinking on leaves
|
||||
the reasoning block open (the prompt ends with the start token); with
|
||||
thinking off it closes the block in the prompt, so the completion is
|
||||
plain content. Parsers that expose no token pair (the engine-based
|
||||
adapters, which take the ``chat_template_kwargs`` route above) keep
|
||||
the parser's verdict unchanged.
|
||||
"""
|
||||
start = getattr(rp, "start_token", None)
|
||||
end = getattr(rp, "end_token", None)
|
||||
|
||||
if end and end in generated_text:
|
||||
# The parser split on the end token. Empty content here means the
|
||||
# model stopped right after it, not that parsing failed.
|
||||
return reasoning or "", content or ""
|
||||
|
||||
if not start:
|
||||
# Unknown token layout — keep the previous behaviour rather than
|
||||
# guess.
|
||||
return reasoning or "", content if content is not None else generated_text
|
||||
|
||||
if not (start in generated_text or (prompt or "").rstrip().endswith(start)):
|
||||
# No end token and the block was never open: the "reasoning starts
|
||||
# at the beginning" fallback does not apply to this completion.
|
||||
return "", generated_text
|
||||
|
||||
# Block was open and the end token never arrived — reasoning ran out of
|
||||
# budget. It is all reasoning, and there is no answer to report.
|
||||
return reasoning or "", content or ""
|
||||
|
||||
async def _predict(self, request, context, streaming=False):
|
||||
# Build the sampling parameters
|
||||
sampling_params = self._build_sampling_params(request)
|
||||
@@ -572,6 +648,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
# Extract image paths and process images
|
||||
prompt = request.Prompt
|
||||
# Kept in scope: the reasoning parser needs to know which chat
|
||||
# template kwargs produced this prompt.
|
||||
template_kwargs = {}
|
||||
|
||||
image_paths = request.Images
|
||||
image_data = [self.load_image(img_path) for img_path in image_paths]
|
||||
@@ -582,7 +661,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
# If tokenizer template is enabled and messages are provided instead of prompt, apply the tokenizer template
|
||||
if not request.Prompt and request.UseTokenizerTemplate and request.Messages:
|
||||
messages_dicts = self._messages_to_dicts(request.Messages)
|
||||
template_kwargs = {"tokenize": False, "add_generation_prompt": True}
|
||||
template_kwargs.update({"tokenize": False, "add_generation_prompt": True})
|
||||
|
||||
# Pass tools for tool calling
|
||||
if request.Tools:
|
||||
@@ -595,13 +674,33 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
if _thinking in ("true", "false"):
|
||||
template_kwargs["enable_thinking"] = (_thinking == "true")
|
||||
|
||||
try:
|
||||
prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
except TypeError:
|
||||
# Some tokenizers don't support tools/enable_thinking kwargs — retry without them
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
messages_dicts, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
# vLLM substitutes multi_modal_data into the model's own media
|
||||
# token, so the template has to be given content *parts* - string
|
||||
# content renders a prompt with no placeholder and the media are
|
||||
# dropped without a word (#11621).
|
||||
prompt = None
|
||||
media_dicts = attach_media_parts(
|
||||
messages_dicts, len(image_data), len(video_data)
|
||||
)
|
||||
if media_dicts is not None:
|
||||
try:
|
||||
prompt = self.tokenizer.apply_chat_template(media_dicts, **template_kwargs)
|
||||
except Exception as e:
|
||||
# A text-only template cannot iterate content parts; fall
|
||||
# through to the text-only prompt instead of failing.
|
||||
print(
|
||||
f"chat template rejected multimodal content parts: {e!r}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if prompt is None:
|
||||
try:
|
||||
prompt = self.tokenizer.apply_chat_template(messages_dicts, **template_kwargs)
|
||||
except TypeError:
|
||||
# Some tokenizers don't support tools/enable_thinking kwargs — retry without them
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
messages_dicts, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Generate text using the LLM engine
|
||||
request_id = random_uuid()
|
||||
@@ -757,10 +856,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
if self.reasoning_parser_cls:
|
||||
try:
|
||||
rp = self.reasoning_parser_cls(self.tokenizer)
|
||||
rp = self._new_reasoning_parser(template_kwargs)
|
||||
r, c = rp.extract_reasoning(generated_text, request=None)
|
||||
reasoning_content = r or ""
|
||||
content = c if c is not None else generated_text
|
||||
reasoning_content, content = self._split_reasoning(
|
||||
rp, generated_text, prompt, r, c,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Reasoning parser error: {e}", file=sys.stderr)
|
||||
|
||||
@@ -905,13 +1005,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
Video: The loaded video.
|
||||
"""
|
||||
try:
|
||||
timestamp = str(int(time.time() * 1000)) # Generate timestamp
|
||||
p = os.path.join(tempfile.gettempdir(), f"vl-{timestamp}.data")
|
||||
with open(p, "wb") as f:
|
||||
f.write(base64.b64decode(video_path))
|
||||
video = VideoAsset(name=p).np_ndarrays
|
||||
os.remove(p)
|
||||
return video
|
||||
with materialize_base64(video_path, suffix=".data") as path:
|
||||
return VideoAsset(name=path).np_ndarrays
|
||||
except Exception as e:
|
||||
print(f"Error loading video {video_path}: {e}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
grpcio==1.83.0
|
||||
grpcio==1.83.1
|
||||
protobuf
|
||||
certifi
|
||||
setuptools
|
||||
|
||||
+118
-3
@@ -121,16 +121,18 @@ class TestBackendServicer(unittest.TestCase):
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
def test_explicit_zero_temperature_is_preserved(self):
|
||||
"""Temperature=0 is valid greedy decoding, not an unset value."""
|
||||
def test_explicit_zero_temperature_and_seed_are_preserved(self):
|
||||
"""Temperature=0 is greedy decoding and 0 is a valid seed — neither is
|
||||
an unset value. A dropped seed turns a reproducible request random."""
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from backend import BackendServicer
|
||||
|
||||
servicer = BackendServicer()
|
||||
request = backend_pb2.PredictOptions(Prompt="hello", Temperature=0)
|
||||
request = backend_pb2.PredictOptions(Prompt="hello", Temperature=0, Seed=0)
|
||||
sampling_params = servicer._build_sampling_params(request)
|
||||
self.assertEqual(sampling_params.temperature, 0)
|
||||
self.assertEqual(sampling_params.seed, 0)
|
||||
# Other protobuf-default scalar fields must remain filtered.
|
||||
self.assertEqual(sampling_params.top_p, 0.9)
|
||||
|
||||
@@ -549,3 +551,116 @@ class TestStreamingToolParser(unittest.TestCase):
|
||||
intermediate, ["Hello ", "world", "!"],
|
||||
f"plain streaming changed; got {intermediate!r}",
|
||||
)
|
||||
|
||||
|
||||
class TestReasoningSplit(unittest.TestCase):
|
||||
"""Server-less tests for BackendServicer._split_reasoning.
|
||||
|
||||
vLLM's BaseThinkingReasoningParser returns the whole completion as
|
||||
reasoning and None as content whenever the end token is missing. Taken
|
||||
literally that duplicates a thinking-disabled answer into both fields.
|
||||
"""
|
||||
|
||||
class _Parser:
|
||||
start_token = "<think>"
|
||||
end_token = "</think>"
|
||||
|
||||
def _split(self, generated, prompt, reasoning, content):
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from backend import BackendServicer
|
||||
return BackendServicer._split_reasoning(
|
||||
self._Parser(), generated, prompt, reasoning, content,
|
||||
)
|
||||
|
||||
def test_thinking_off_is_not_duplicated_into_reasoning(self):
|
||||
"""No tags anywhere: the answer is content, and only content."""
|
||||
r, c = self._split(
|
||||
"391", "user: 17*23?\n<think>\n\n</think>\n\n",
|
||||
reasoning="391", content=None,
|
||||
)
|
||||
self.assertEqual(r, "")
|
||||
self.assertEqual(c, "391")
|
||||
|
||||
def test_prefilled_start_tag_keeps_truncated_reasoning(self):
|
||||
"""Prompt left the block open and the end token never arrived
|
||||
(budget exhausted): that really is all reasoning."""
|
||||
r, c = self._split(
|
||||
"thinking and thinking", "user: hi\n<think>\n",
|
||||
reasoning="thinking and thinking", content=None,
|
||||
)
|
||||
self.assertEqual(r, "thinking and thinking")
|
||||
self.assertEqual(c, "")
|
||||
|
||||
def test_end_token_present_keeps_parser_split(self):
|
||||
r, c = self._split(
|
||||
"adding two and two</think>4", "user: hi\n<think>\n",
|
||||
reasoning="adding two and two", content="4",
|
||||
)
|
||||
self.assertEqual(r, "adding two and two")
|
||||
self.assertEqual(c, "4")
|
||||
|
||||
def test_stop_right_after_end_token_yields_empty_content(self):
|
||||
"""Content must not fall back to the raw text — that would put the
|
||||
reasoning into the answer."""
|
||||
r, c = self._split(
|
||||
"reasoned</think>", "user: hi\n<think>\n",
|
||||
reasoning="reasoned", content=None,
|
||||
)
|
||||
self.assertEqual(r, "reasoned")
|
||||
self.assertEqual(c, "")
|
||||
|
||||
def test_unknown_token_layout_keeps_previous_behaviour(self):
|
||||
class _Bare:
|
||||
pass
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from backend import BackendServicer
|
||||
r, c = BackendServicer._split_reasoning(
|
||||
_Bare(), "raw", "prompt", "raw", None,
|
||||
)
|
||||
self.assertEqual(r, "raw")
|
||||
self.assertEqual(c, "raw")
|
||||
|
||||
|
||||
class TestReasoningParserConstruction(unittest.TestCase):
|
||||
"""The parser must learn whether thinking was on for this request.
|
||||
|
||||
vLLM's engine-based parsers (Qwen3Parser and friends) read
|
||||
chat_template_kwargs["enable_thinking"] and default to True, so a parser
|
||||
built without it treats a thinking-disabled completion as pure reasoning.
|
||||
"""
|
||||
|
||||
def _servicer(self):
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from backend import BackendServicer
|
||||
s = BackendServicer()
|
||||
s.tokenizer = object()
|
||||
return s
|
||||
|
||||
def test_chat_template_kwargs_are_forwarded(self):
|
||||
seen = {}
|
||||
|
||||
class _Parser:
|
||||
def __init__(self, tokenizer, **kwargs):
|
||||
seen.update(kwargs)
|
||||
|
||||
s = self._servicer()
|
||||
s.reasoning_parser_cls = _Parser
|
||||
s._new_reasoning_parser({"enable_thinking": False})
|
||||
self.assertEqual(
|
||||
seen.get("chat_template_kwargs"), {"enable_thinking": False},
|
||||
)
|
||||
|
||||
def test_parser_without_the_kwarg_still_builds(self):
|
||||
"""Older parsers take only the tokenizer — must not break them."""
|
||||
class _Old:
|
||||
def __init__(self, tokenizer):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
s = self._servicer()
|
||||
s.reasoning_parser_cls = _Old
|
||||
self.assertIsInstance(
|
||||
s._new_reasoning_parser({"enable_thinking": False}), _Old,
|
||||
)
|
||||
@@ -16,6 +16,7 @@ import grpc
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
|
||||
from grpc_auth import get_auth_interceptors
|
||||
from transcript_utils import require_diarization_token, seconds_to_nanoseconds
|
||||
|
||||
|
||||
|
||||
@@ -81,6 +82,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
import whisperx
|
||||
from whisperx.diarize import DiarizationPipeline
|
||||
|
||||
try:
|
||||
require_diarization_token(request.diarize, self.hf_token)
|
||||
except ValueError as err:
|
||||
context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(err))
|
||||
|
||||
resultSegments = []
|
||||
text = ""
|
||||
try:
|
||||
@@ -117,8 +123,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
# Build result segments
|
||||
for idx, seg in enumerate(transcript["segments"]):
|
||||
seg_text = seg.get("text", "")
|
||||
start = int(seg.get("start", 0))
|
||||
end = int(seg.get("end", 0))
|
||||
start = seconds_to_nanoseconds(seg.get("start", 0))
|
||||
end = seconds_to_nanoseconds(seg.get("end", 0))
|
||||
speaker = seg.get("speaker", "")
|
||||
|
||||
resultSegments.append(backend_pb2.TranscriptSegment(
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
import unittest
|
||||
|
||||
import transcript_utils
|
||||
|
||||
|
||||
class TestTranscriptUtils(unittest.TestCase):
|
||||
def test_diarization_requires_hugging_face_token(self):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"HF_TOKEN is required for WhisperX diarization",
|
||||
):
|
||||
transcript_utils.require_diarization_token(True, None)
|
||||
|
||||
def test_diarization_does_not_require_token_when_disabled(self):
|
||||
transcript_utils.require_diarization_token(False, None)
|
||||
|
||||
def test_seconds_are_serialized_as_nanoseconds(self):
|
||||
self.assertEqual(
|
||||
transcript_utils.seconds_to_nanoseconds(3.25),
|
||||
3_250_000_000,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Helpers for WhisperX transcript responses."""
|
||||
|
||||
|
||||
def require_diarization_token(diarize, token):
|
||||
"""Reject diarization when WhisperX cannot load its gated pipeline."""
|
||||
if diarize and not token:
|
||||
raise ValueError("HF_TOKEN is required for WhisperX diarization")
|
||||
|
||||
|
||||
def seconds_to_nanoseconds(seconds):
|
||||
"""Convert WhisperX timestamps to the duration unit used by LocalAI."""
|
||||
return int(seconds * 1_000_000_000)
|
||||
@@ -0,0 +1,101 @@
|
||||
package gallery_test
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
func remarshal(value any, target any) error {
|
||||
data, err := yaml.Marshal(value)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return yaml.Unmarshal(data, target)
|
||||
}
|
||||
|
||||
var _ = Describe("EXL3 gallery entries", func() {
|
||||
It("pins the four full repositories and configures the Qwen DFlash companion", func() {
|
||||
entries, err := gallery.ReadConfigFile[[]gallery.GalleryModel](filepath.Join("..", "..", "gallery", "index.yaml"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
byName := make(map[string]gallery.GalleryModel, len(*entries))
|
||||
for _, entry := range *entries {
|
||||
byName[entry.Name] = entry
|
||||
}
|
||||
|
||||
expected := map[string]struct {
|
||||
repo string
|
||||
revision string
|
||||
}{
|
||||
"qwen3.8-27b-exl3-vllm-cpp": {
|
||||
repo: "Mia-AiLab/Qwen3.8-27B-EXL3-3.5bpw", revision: "19441ac874c4018295da848e250f23511361cda4",
|
||||
},
|
||||
"qwen3.8-27b-dflash2-exl3-vllm-cpp": {
|
||||
repo: "Mia-AiLab/Qwen3.8-27B-EXL3-3.5bpw", revision: "19441ac874c4018295da848e250f23511361cda4",
|
||||
},
|
||||
"deepseek-v4-flash-spark-exl3-vllm-cpp": {
|
||||
repo: "0xSero/deepseek-v4-flash-0731-spark", revision: "ce5ff0f1efb2e184aafc759d281bfae47d3a359c",
|
||||
},
|
||||
"deepseek-v4-flash-exl3-3bpw-vllm-cpp": {
|
||||
repo: "0xSero/DeepSeek-V4-Flash-0731-EXL3-3.0bpw", revision: "e0bf84ac76a5100e8790c22ad10b70b1e2d06d71",
|
||||
},
|
||||
}
|
||||
|
||||
for name, want := range expected {
|
||||
entry, found := byName[name]
|
||||
Expect(found).To(BeTrue(), "missing gallery entry %q", name)
|
||||
Expect(entry.Tags).To(ContainElements("vllm-cpp", "exl3", "gpu", "cuda"), name)
|
||||
Expect(entry.Overrides).To(HaveKeyWithValue("backend", "vllm-cpp"), name)
|
||||
cfg := config.ModelConfig{}
|
||||
Expect(remarshal(entry.Overrides, &cfg)).To(Succeed(), name)
|
||||
Expect(cfg.Artifacts).ToNot(BeEmpty(), name)
|
||||
Expect(cfg.Artifacts[0].Source.Repo).To(Equal(want.repo), name)
|
||||
Expect(cfg.Artifacts[0].Source.Revision).To(Equal(want.revision), name)
|
||||
}
|
||||
|
||||
plain := byName["qwen3.8-27b-exl3-vllm-cpp"]
|
||||
Expect(plain.Tags).ToNot(ContainElement("dflash"))
|
||||
dflashTags := 0
|
||||
for name := range expected {
|
||||
if contains(byName[name].Tags, "dflash") {
|
||||
dflashTags++
|
||||
}
|
||||
}
|
||||
Expect(dflashTags).To(Equal(1))
|
||||
|
||||
dflash := byName["qwen3.8-27b-dflash2-exl3-vllm-cpp"]
|
||||
Expect(dflash.Tags).To(ContainElement("dflash"))
|
||||
Expect(dflash.Variants).To(ConsistOf(gallery.Variant{Model: "qwen3.8-27b-exl3-vllm-cpp"}))
|
||||
cfg := config.ModelConfig{}
|
||||
Expect(remarshal(dflash.Overrides, &cfg)).To(Succeed())
|
||||
Expect(cfg.ContextSize).To(HaveValue(Equal(8192)))
|
||||
Expect(cfg.EngineArgs).To(HaveKeyWithValue("num_blocks", 2048))
|
||||
Expect(cfg.EngineArgs).To(HaveKeyWithValue("max_num_seqs", 8))
|
||||
Expect(cfg.EngineArgs).To(HaveKeyWithValue("max_num_batched_tokens", 16384))
|
||||
Expect(cfg.EngineArgs).To(HaveKeyWithValue("enable_prefix_caching", false))
|
||||
spec, ok := cfg.EngineArgs["speculative_config"].(map[string]any)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(spec).To(HaveKeyWithValue("method", "dflash"))
|
||||
Expect(spec).To(HaveKeyWithValue("num_speculative_tokens", 7))
|
||||
Expect(cfg.Artifacts).To(HaveLen(2))
|
||||
Expect(cfg.Artifacts[1].Name).To(Equal("draft_model"))
|
||||
Expect(cfg.Artifacts[1].Target).To(Equal("companion"))
|
||||
Expect(cfg.Artifacts[1].Source.Repo).To(Equal("Mia-AiLab/Qwen3.8-27B-DFlash2-EXL3-5.0bpw"))
|
||||
Expect(cfg.Artifacts[1].Source.Revision).To(Equal("4f0436269bca761b071f05319e8e04a87cc633f9"))
|
||||
})
|
||||
})
|
||||
|
||||
func contains(values []string, target string) bool {
|
||||
for _, value := range values {
|
||||
if value == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package localai
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
@@ -33,22 +34,31 @@ func FaceRegisterEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, a
|
||||
return echo.NewHTTPError(http.StatusBadRequest, "name is required")
|
||||
}
|
||||
|
||||
img, err := decodeImageInput(input.Img)
|
||||
if err != nil {
|
||||
return err
|
||||
if (input.Img == "") == (len(input.Embedding) == 0) {
|
||||
return echo.NewHTTPError(http.StatusBadRequest, "provide exactly one of img or embedding")
|
||||
}
|
||||
|
||||
xlog.Debug("FaceRegister", "model", cfg.Name, "name", input.Name)
|
||||
embedding, err := backend.FaceEmbed(c.Request().Context(), img, ml, appConfig, *cfg)
|
||||
if err != nil {
|
||||
return mapBackendError(err)
|
||||
embedding := input.Embedding
|
||||
if len(embedding) == 0 {
|
||||
img, err := decodeImageInput(input.Img)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
xlog.Debug("FaceRegister", "model", cfg.Name, "name", input.Name)
|
||||
embedding, err = backend.FaceEmbed(c.Request().Context(), img, ml, appConfig, *cfg)
|
||||
if err != nil {
|
||||
return mapBackendError(err)
|
||||
}
|
||||
}
|
||||
|
||||
stored, err := registry.Register(c.Request().Context(), embedding, facerecognition.Metadata{
|
||||
Name: input.Name,
|
||||
Labels: input.Labels,
|
||||
Name: input.Name,
|
||||
RegisteredAt: input.RegisteredAt,
|
||||
Labels: input.Labels,
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, facerecognition.ErrInvalidEmbedding) || errors.Is(err, facerecognition.ErrDimensionMismatch) {
|
||||
return echo.NewHTTPError(http.StatusBadRequest, err.Error())
|
||||
}
|
||||
return err
|
||||
}
|
||||
return c.JSON(http.StatusOK, schema.FaceRegisterResponse{
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package localai_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"time"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/facerecognition"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type registrationRecorder struct {
|
||||
facerecognition.Registry
|
||||
vector []float32
|
||||
meta facerecognition.Metadata
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *registrationRecorder) Register(_ context.Context, v []float32, m facerecognition.Metadata) (facerecognition.Metadata, error) {
|
||||
r.vector = v
|
||||
r.meta = m
|
||||
m.ID = "saved-id"
|
||||
return m, r.err
|
||||
}
|
||||
|
||||
var _ = Describe("Face registration replay", func() {
|
||||
var reg *registrationRecorder
|
||||
call := func(in schema.FaceRegisterRequest) (*httptest.ResponseRecorder, error) {
|
||||
e := echo.New()
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(httptest.NewRequest(http.MethodPost, "/v1/face/register", nil), rec)
|
||||
c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, &in)
|
||||
c.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &config.ModelConfig{})
|
||||
// No model loader: replay must not call the embedding backend.
|
||||
err := FaceRegisterEndpoint(nil, nil, nil, reg)(c)
|
||||
return rec, err
|
||||
}
|
||||
BeforeEach(func() { reg = ®istrationRecorder{} })
|
||||
It("accepts the saved vector and timestamp without running inference", func() {
|
||||
at := time.Now().UTC()
|
||||
in := schema.FaceRegisterRequest{Name: "Alice", Embedding: []float32{1, 0}, RegisteredAt: at, Labels: map[string]string{"client_id": "alice"}}
|
||||
in.Model = "faces"
|
||||
rec, err := call(in)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
Expect(reg.vector).To(Equal(in.Embedding))
|
||||
Expect(reg.meta.RegisteredAt).To(Equal(at))
|
||||
Expect(reg.meta.Labels).To(Equal(in.Labels))
|
||||
Expect(rec.Body.String()).To(ContainSubstring("saved-id"))
|
||||
})
|
||||
It("rejects ambiguous and missing inputs before inference", func() {
|
||||
for _, in := range []schema.FaceRegisterRequest{
|
||||
{Name: "Alice"},
|
||||
{Name: "Alice", Img: "image", Embedding: []float32{1, 0}},
|
||||
} {
|
||||
in.Model = "faces"
|
||||
_, err := call(in)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.(*echo.HTTPError).Code).To(Equal(http.StatusBadRequest))
|
||||
Expect(reg.vector).To(BeNil())
|
||||
}
|
||||
})
|
||||
It("reports invalid vectors as a client error", func() {
|
||||
reg.err = facerecognition.ErrInvalidEmbedding
|
||||
in := schema.FaceRegisterRequest{Name: "Alice", Embedding: []float32{0, 0}}
|
||||
in.Model = "faces"
|
||||
_, err := call(in)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.(*echo.HTTPError).Code).To(Equal(http.StatusBadRequest))
|
||||
})
|
||||
})
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"math"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -633,6 +634,17 @@ func runRealtimeSession(application *application.Application, t Transport, model
|
||||
sendError(t, "model_load_error", "Failed to load model", "", "")
|
||||
return
|
||||
}
|
||||
if wrapped, ok := m.(*wrappedModel); ok {
|
||||
resolvedVoice, params, release, resolveErr := resolveRealtimeVoice(context.Background(), session.Voice, wrapped.TTSConfig, application.VoiceProfileStore())
|
||||
if resolveErr != nil {
|
||||
xlog.Error("failed to resolve realtime voice", "error", resolveErr)
|
||||
sendError(t, "voice_profile_error", resolveErr.Error(), "", "")
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
session.Voice = resolvedVoice
|
||||
wrapped.ttsParams = params
|
||||
}
|
||||
session.ModelInterface = m
|
||||
// A pipeline-seeded option list gets its scoring prompt prewarmed
|
||||
// alongside the model warm-up below, so the session's first turn
|
||||
@@ -1923,7 +1935,7 @@ func commitUtteranceWithTranscript(ctx context.Context, utt []byte, live *liveUt
|
||||
// Generate an LLM response only when there is a transcript to feed it. A
|
||||
// sound-detection-only session (no transcription) has no LLM stage, so it
|
||||
// stops here after emitting the sound-detection event.
|
||||
if session.InputAudioTranscription != nil && !session.TranscriptionOnly {
|
||||
if session.InputAudioTranscription != nil && !session.TranscriptionOnly && strings.TrimSpace(transcript) != "" {
|
||||
generateResponse(ctx, session, utt, transcript, speaker, conv, t)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,9 @@ import (
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"maps"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -18,6 +20,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/routing/router"
|
||||
"github.com/mudler/LocalAI/core/services/voiceprofile"
|
||||
"github.com/mudler/LocalAI/core/templates"
|
||||
"github.com/mudler/LocalAI/pkg/functions"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
@@ -35,6 +38,7 @@ var (
|
||||
// which are for Any-To-Any models, but instead we will call a pipeline (for e.g STT->LLM->TTS)
|
||||
type wrappedModel struct {
|
||||
TTSConfig *config.ModelConfig
|
||||
ttsParams map[string]string
|
||||
TranscriptionConfig *config.ModelConfig
|
||||
LLMConfig *config.ModelConfig
|
||||
VADConfig *config.ModelConfig
|
||||
@@ -391,11 +395,35 @@ func newRealtimeDecisionID() string {
|
||||
}
|
||||
|
||||
func (m *wrappedModel) TTS(ctx context.Context, text, voice, language string) (string, *proto.Result, error) {
|
||||
return backend.ModelTTS(ctx, text, voice, language, "", nil, m.modelLoader, m.appConfig, *m.TTSConfig)
|
||||
return backend.ModelTTS(ctx, text, voice, language, "", maps.Clone(m.ttsParams), m.modelLoader, m.appConfig, *m.TTSConfig)
|
||||
}
|
||||
|
||||
func (m *wrappedModel) TTSStream(ctx context.Context, text, voice, language string, onAudio func(pcm []byte, sampleRate int) error) error {
|
||||
return ttsStream(ctx, m.modelLoader, m.appConfig, *m.TTSConfig, text, voice, language, onAudio)
|
||||
return ttsStream(ctx, m.modelLoader, m.appConfig, *m.TTSConfig, text, voice, language, maps.Clone(m.ttsParams), onAudio)
|
||||
}
|
||||
|
||||
func resolveRealtimeVoice(ctx context.Context, configuredVoice string, ttsConfig *config.ModelConfig, profiles *voiceprofile.Store) (string, map[string]string, func(), error) {
|
||||
if !voiceprofile.IsReference(configuredVoice) {
|
||||
return configuredVoice, nil, func() {}, nil
|
||||
}
|
||||
profileID, valid := voiceprofile.ParseReference(configuredVoice)
|
||||
if !valid {
|
||||
return "", nil, nil, fmt.Errorf("invalid voice profile reference %q", configuredVoice)
|
||||
}
|
||||
if config.VoiceCloningForModel(ttsConfig) == nil {
|
||||
return "", nil, nil, fmt.Errorf("selected TTS model does not support reference-audio voice cloning")
|
||||
}
|
||||
if profiles == nil {
|
||||
return "", nil, nil, fmt.Errorf("voice profile store is unavailable")
|
||||
}
|
||||
profile, referencePath, release, err := profiles.LeaseAudio(ctx, profileID)
|
||||
if err != nil {
|
||||
if errors.Is(err, voiceprofile.ErrNotFound) {
|
||||
return "", nil, nil, fmt.Errorf("voice profile not found: %w", err)
|
||||
}
|
||||
return "", nil, nil, fmt.Errorf("resolve voice profile: %w", err)
|
||||
}
|
||||
return referencePath, map[string]string{"ref_text": profile.Transcript}, release, nil
|
||||
}
|
||||
|
||||
func (m *wrappedModel) TranscribeStream(ctx context.Context, audio, language string, translate, diarize bool, prompt string, onDelta func(text string)) (*schema.TranscriptionResult, error) {
|
||||
@@ -674,11 +702,11 @@ const wavStreamHeaderBytes = 44
|
||||
// callback, which wants raw PCM plus the sample rate. The header is buffered
|
||||
// until complete, the sample rate is read from it, and subsequent bytes are
|
||||
// forwarded as PCM.
|
||||
func ttsStream(ctx context.Context, ml *model.ModelLoader, appConfig *config.ApplicationConfig, ttsConfig config.ModelConfig, text, voice, language string, onAudio func(pcm []byte, sampleRate int) error) error {
|
||||
func ttsStream(ctx context.Context, ml *model.ModelLoader, appConfig *config.ApplicationConfig, ttsConfig config.ModelConfig, text, voice, language string, params map[string]string, onAudio func(pcm []byte, sampleRate int) error) error {
|
||||
var header []byte
|
||||
headerDone := false
|
||||
sampleRate := 0
|
||||
return backend.ModelTTSStream(ctx, text, voice, language, "", nil, ml, appConfig, ttsConfig, func(b []byte) error {
|
||||
return backend.ModelTTSStream(ctx, text, voice, language, "", params, ml, appConfig, ttsConfig, func(b []byte) error {
|
||||
if headerDone {
|
||||
if len(b) == 0 {
|
||||
return nil
|
||||
|
||||
@@ -355,6 +355,19 @@ var _ = Describe("commitUtteranceWithTranscript", func() {
|
||||
|
||||
Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1))
|
||||
})
|
||||
|
||||
It("does not generate a response for a blank transcript", func() {
|
||||
session, model := itSession(nil)
|
||||
model.transcribeFinal = &schema.TranscriptionResult{Text: " \t\n"}
|
||||
tr := &fakeTransport{}
|
||||
conv := &Conversation{}
|
||||
|
||||
commitUtterance(context.Background(), []byte{1, 2}, session, conv, tr)
|
||||
|
||||
Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1))
|
||||
Expect(conv.Items).To(BeEmpty())
|
||||
Expect(tr.countEvents(types.ServerEventTypeResponseCreated)).To(Equal(0))
|
||||
})
|
||||
})
|
||||
|
||||
// transcribeUtterance is the retranscribe gate's offline decode of the
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
grpcPkg "github.com/mudler/LocalAI/pkg/grpc"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
"google.golang.org/grpc"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/voiceprofile"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func realtimeProfileWAV(duration time.Duration) []byte {
|
||||
const (
|
||||
sampleRate = 16000
|
||||
channels = 1
|
||||
bitsPerSample = 16
|
||||
)
|
||||
dataSize := int(duration.Seconds() * sampleRate * channels * bitsPerSample / 8)
|
||||
buf := bytes.NewBuffer(nil)
|
||||
buf.WriteString("RIFF")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(36+dataSize))
|
||||
buf.WriteString("WAVEfmt ")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(16))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(1))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(channels))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(sampleRate))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(sampleRate*channels*bitsPerSample/8))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(channels*bitsPerSample/8))
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint16(bitsPerSample))
|
||||
buf.WriteString("data")
|
||||
_ = binary.Write(buf, binary.LittleEndian, uint32(dataSize))
|
||||
buf.Write(make([]byte, dataSize))
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
var _ = Describe("realtime pipeline voice profiles", func() {
|
||||
It("resolves a saved profile to an immutable lease and transcript", func(ctx SpecContext) {
|
||||
store := voiceprofile.NewStore(GinkgoT().TempDir())
|
||||
DeferCleanup(func() { Expect(store.Close()).To(Succeed()) })
|
||||
profile, err := store.Create(ctx, voiceprofile.CreateInput{
|
||||
Name: "Narrator",
|
||||
Language: "en-US",
|
||||
Transcript: "The reference transcript.",
|
||||
ConsentConfirmed: true,
|
||||
}, bytes.NewReader(realtimeProfileWAV(time.Second)))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
voice, params, release, err := resolveRealtimeVoice(ctx, profile.Voice, &config.ModelConfig{
|
||||
Name: "clone-base",
|
||||
Backend: "qwen3-tts-cpp",
|
||||
TTSConfig: config.TTSConfig{VoiceCloning: ptrTo(true)},
|
||||
}, store)
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(voice).To(BeAnExistingFile())
|
||||
Expect(params).To(Equal(map[string]string{"ref_text": "The reference transcript."}))
|
||||
release()
|
||||
release()
|
||||
Expect(voice).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("leaves an ordinary backend voice unchanged with no parameters", func() {
|
||||
voice, params, release, err := resolveRealtimeVoice(context.Background(), "speaker-7", &config.ModelConfig{}, nil)
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(voice).To(Equal("speaker-7"))
|
||||
Expect(params).To(BeNil())
|
||||
Expect(release).NotTo(BeNil())
|
||||
Expect(func() { release(); release() }).NotTo(Panic())
|
||||
})
|
||||
|
||||
DescribeTable("returns actionable reference errors",
|
||||
func(configuredVoice string, cfg *config.ModelConfig, store *voiceprofile.Store, expected string) {
|
||||
_, _, release, err := resolveRealtimeVoice(context.Background(), configuredVoice, cfg, store)
|
||||
Expect(err).To(MatchError(ContainSubstring(expected)))
|
||||
Expect(release).To(BeNil())
|
||||
},
|
||||
Entry("malformed reference", "localai://voice-profiles/not-a-uuid", &config.ModelConfig{}, nil, "invalid voice profile reference"),
|
||||
Entry("unsupported model", "localai://voice-profiles/00000000-0000-0000-0000-000000000001", &config.ModelConfig{Backend: "piper"}, nil, "does not support reference-audio voice cloning"),
|
||||
Entry("unavailable store", "localai://voice-profiles/00000000-0000-0000-0000-000000000001", &config.ModelConfig{Name: "clone-base", Backend: "qwen3-tts-cpp", TTSConfig: config.TTSConfig{VoiceCloning: ptrTo(true)}}, nil, "voice profile store is unavailable"),
|
||||
)
|
||||
|
||||
It("reports a missing profile", func() {
|
||||
store := voiceprofile.NewStore(GinkgoT().TempDir())
|
||||
DeferCleanup(func() { Expect(store.Close()).To(Succeed()) })
|
||||
_, _, release, err := resolveRealtimeVoice(context.Background(), "localai://voice-profiles/00000000-0000-0000-0000-000000000001", &config.ModelConfig{
|
||||
Name: "clone-base", Backend: "qwen3-tts-cpp", TTSConfig: config.TTSConfig{VoiceCloning: ptrTo(true)},
|
||||
}, store)
|
||||
Expect(errors.Is(err, voiceprofile.ErrNotFound)).To(BeTrue())
|
||||
Expect(err.Error()).To(ContainSubstring("voice profile not found"))
|
||||
Expect(release).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
type recordingTTSBackend struct {
|
||||
grpcPkg.Backend
|
||||
requests []*proto.TTSRequest
|
||||
}
|
||||
|
||||
func (b *recordingTTSBackend) HealthCheck(context.Context) (bool, error) { return true, nil }
|
||||
func (b *recordingTTSBackend) IsBusy() bool { return false }
|
||||
|
||||
func (b *recordingTTSBackend) record(req *proto.TTSRequest) {
|
||||
b.requests = append(b.requests, req)
|
||||
req.Params["ref_text"] = "backend mutation"
|
||||
}
|
||||
|
||||
func (b *recordingTTSBackend) TTS(_ context.Context, req *proto.TTSRequest, _ ...grpc.CallOption) (*proto.Result, error) {
|
||||
b.record(req)
|
||||
return &proto.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *recordingTTSBackend) TTSStream(_ context.Context, req *proto.TTSRequest, callback func(*proto.Reply), _ ...grpc.CallOption) error {
|
||||
b.record(req)
|
||||
header := make([]byte, wavStreamHeaderBytes)
|
||||
binary.LittleEndian.PutUint32(header[24:28], 24000)
|
||||
callback(&proto.Reply{Audio: header})
|
||||
return nil
|
||||
}
|
||||
|
||||
var _ = Describe("wrappedModel voice profile parameters", func() {
|
||||
var (
|
||||
wrapped *wrappedModel
|
||||
backendRecorder *recordingTTSBackend
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
state, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir()))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
appConfig := config.NewApplicationConfig(config.WithSystemState(state))
|
||||
appConfig.GeneratedContentDir = GinkgoT().TempDir()
|
||||
loader := model.NewModelLoader(state)
|
||||
backendRecorder = &recordingTTSBackend{}
|
||||
cfg := &config.ModelConfig{Name: "tts-test", Backend: "test"}
|
||||
cfg.Model = "weights"
|
||||
loaded := model.NewModelWithClient(cfg.ModelID(), "in-process", backendRecorder)
|
||||
loaded.MarkHealthy()
|
||||
_, err = loader.LoadModel(cfg.ModelID(), cfg.Model, func(_, _, _ string) (*model.Model, error) { return loaded, nil })
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
wrapped = &wrappedModel{
|
||||
TTSConfig: cfg,
|
||||
ttsParams: map[string]string{"ref_text": "Original transcript"},
|
||||
modelLoader: loader,
|
||||
appConfig: appConfig,
|
||||
}
|
||||
})
|
||||
|
||||
It("forwards a fresh transcript parameter map to every unary request", func() {
|
||||
_, _, err := wrapped.TTS(context.Background(), "one", "voice.wav", "en")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
_, _, err = wrapped.TTS(context.Background(), "two", "voice.wav", "en")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(backendRecorder.requests).To(HaveLen(2))
|
||||
Expect(backendRecorder.requests[0].Params).To(HaveKeyWithValue("ref_text", "backend mutation"))
|
||||
Expect(backendRecorder.requests[1].Params).To(HaveKeyWithValue("ref_text", "backend mutation"))
|
||||
Expect(wrapped.ttsParams).To(HaveKeyWithValue("ref_text", "Original transcript"))
|
||||
})
|
||||
|
||||
It("forwards a copied transcript parameter map to streaming requests", func() {
|
||||
err := wrapped.TTSStream(context.Background(), "one", "voice.wav", "en", func([]byte, int) error { return nil })
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(backendRecorder.requests).To(HaveLen(1))
|
||||
Expect(backendRecorder.requests[0].Params).To(HaveKeyWithValue("ref_text", "backend mutation"))
|
||||
Expect(wrapped.ttsParams).To(HaveKeyWithValue("ref_text", "Original transcript"))
|
||||
})
|
||||
})
|
||||
|
||||
func ptrTo[T any](value T) *T { return &value }
|
||||
@@ -53,6 +53,7 @@
|
||||
"overrides": {
|
||||
"hono": "4.12.34",
|
||||
"ip-address": "10.3.1",
|
||||
"path-to-regexp": "^8.4.0",
|
||||
},
|
||||
"packages": {
|
||||
"@babel/code-frame": ["@babel/code-frame@7.29.0", "", { "dependencies": { "@babel/helper-validator-identifier": "^7.28.5", "js-tokens": "^4.0.0", "picocolors": "^1.1.1" } }, "sha512-9NhCeYjq9+3uxgdtp20LSiJXJvN0FeCtNGpJxuMFZ1Kv3cWUNb6DOhJwUvcVCzKGR66cw4njwM6hrJLqgOwbcw=="],
|
||||
@@ -807,7 +808,7 @@
|
||||
|
||||
"path-scurry": ["path-scurry@2.0.2", "", { "dependencies": { "lru-cache": "^11.0.0", "minipass": "^7.1.2" } }, "sha512-3O/iVVsJAPsOnpwWIeD+d6z/7PmqApyQePUtCndjatj/9I5LylHvt5qluFaBT3I5h3r1ejfR056c+FCv+NnNXg=="],
|
||||
|
||||
"path-to-regexp": ["path-to-regexp@8.3.0", "", {}, "sha512-7jdwVIRtsP8MYpdXSwOS0YdD0Du+qOoF/AEPIt88PcCFrZCzx41oxku1jD88hZBwbNUIEfpqvuhjFaMAqMTWnA=="],
|
||||
"path-to-regexp": ["path-to-regexp@8.4.0", "", {}, "sha512-PuseHIvAnz3bjrM2rGJtSgo1zjgxapTLZ7x2pjhzWwlp4SJQgK3f3iZIQwkpEnBaKz6seKBADpM4B4ySkuYypg=="],
|
||||
|
||||
"picocolors": ["picocolors@1.1.1", "", {}, "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA=="],
|
||||
|
||||
|
||||
@@ -19,4 +19,33 @@ test.describe('Collections page', () => {
|
||||
await input.fill('my-kb')
|
||||
await expect(input).toHaveValue('my-kb')
|
||||
})
|
||||
|
||||
test('posts the source update interval as a JSON number', async ({ page }) => {
|
||||
const collectionName = 'interval-regression'
|
||||
const collectionPath = encodeURIComponent(collectionName)
|
||||
let postedBody
|
||||
|
||||
await page.route(`**/api/agents/collections/${collectionPath}/entries`, route =>
|
||||
route.fulfill({ contentType: 'application/json', body: JSON.stringify({ entries: [] }) }))
|
||||
await page.route(`**/api/agents/collections/${collectionPath}/sources`, async route => {
|
||||
if (route.request().method() === 'POST') {
|
||||
postedBody = route.request().postDataJSON()
|
||||
await route.fulfill({ contentType: 'application/json', body: JSON.stringify({ status: 'ok' }) })
|
||||
} else {
|
||||
await route.fulfill({ contentType: 'application/json', body: JSON.stringify({ sources: [] }) })
|
||||
}
|
||||
})
|
||||
|
||||
await page.goto(`/app/collections/${collectionPath}`)
|
||||
await page.getByRole('button', { name: 'Sources' }).click()
|
||||
await page.locator('#source-url').fill('https://example.com/feed')
|
||||
await page.locator('#source-interval').fill('3600')
|
||||
await page.getByRole('button', { name: 'Add Source' }).click()
|
||||
|
||||
await expect.poll(() => postedBody).toEqual({
|
||||
url: 'https://example.com/feed',
|
||||
update_interval: 3600,
|
||||
})
|
||||
expect(typeof postedBody.update_interval).toBe('number')
|
||||
})
|
||||
})
|
||||
@@ -79,4 +79,30 @@ test.describe('Traces - bounded list and on-demand detail', () => {
|
||||
await expect(page.locator('text=hello from the response body')).toBeVisible()
|
||||
await expect(page.locator('text=203.0.113.9').first()).toBeVisible()
|
||||
})
|
||||
|
||||
test('keeps the expanded trace open when a refresh prepends a new row', async ({ page }) => {
|
||||
await page.locator('tr', { hasText: '/v1/chat/completions' }).first().click()
|
||||
await expect(page.locator('text=hello from the request body')).toBeVisible()
|
||||
|
||||
await page.route('**/api/traces?*', (route) => {
|
||||
route.fulfill({
|
||||
contentType: 'application/json',
|
||||
headers: { 'X-Total-Count': '843' },
|
||||
body: JSON.stringify([
|
||||
{
|
||||
id: '8',
|
||||
request: { method: 'GET', path: '/v1/models', body: null },
|
||||
response: { status: 200, body: null },
|
||||
},
|
||||
...LIST_BODY,
|
||||
]),
|
||||
})
|
||||
})
|
||||
|
||||
await page.getByRole('button', { name: 'Refresh' }).click()
|
||||
|
||||
await expect(page.locator('text=hello from the request body')).toBeVisible()
|
||||
const originalRow = page.locator('tr', { hasText: '/v1/chat/completions' }).first()
|
||||
await expect(originalRow.locator('i.fa-chevron-down')).toBeVisible()
|
||||
})
|
||||
})
|
||||
Generated
+74
-29
@@ -32,7 +32,7 @@
|
||||
"react": "^19.1.0",
|
||||
"react-dom": "^19.1.0",
|
||||
"react-i18next": "^17.0.6",
|
||||
"react-router": "7.18.2",
|
||||
"react-router": "8.3.1",
|
||||
"react-router-dom": "7.18.2",
|
||||
"yaml": "^2.8.3"
|
||||
},
|
||||
@@ -649,27 +649,43 @@
|
||||
}
|
||||
},
|
||||
"node_modules/@humanfs/core": {
|
||||
"version": "0.19.1",
|
||||
"resolved": "https://registry.npmjs.org/@humanfs/core/-/core-0.19.1.tgz",
|
||||
"integrity": "sha512-5DyQ4+1JEUzejeK1JGICcideyfUbGixgS9jNgex5nqkW+cY7WZhxBigmieN5Qnw9ZosSNVC9KQKyb+GUaGyKUA==",
|
||||
"version": "0.19.2",
|
||||
"resolved": "https://registry.npmjs.org/@humanfs/core/-/core-0.19.2.tgz",
|
||||
"integrity": "sha512-UhXNm+CFMWcbChXywFwkmhqjs3PRCmcSa/hfBgLIb7oQ5HNb1wS0icWsGtSAUNgefHeI+eBrA8I1fxmbHsGdvA==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@humanfs/types": "^0.15.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=18.18.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@humanfs/node": {
|
||||
"version": "0.16.7",
|
||||
"resolved": "https://registry.npmjs.org/@humanfs/node/-/node-0.16.7.tgz",
|
||||
"integrity": "sha512-/zUx+yOsIrG4Y43Eh2peDeKCxlRt/gET6aHfaKpuq267qXdYDFViVHfMaLyygZOnl0kGWxFIgsBy8QFuTLUXEQ==",
|
||||
"version": "0.16.8",
|
||||
"resolved": "https://registry.npmjs.org/@humanfs/node/-/node-0.16.8.tgz",
|
||||
"integrity": "sha512-gE1eQNZ3R++kTzFUpdGlpmy8kDZD/MLyHqDwqjkVQI0JMdI1D51sy1H958PNXYkM2rAac7e5/CnIKZrHtPh3BQ==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@humanfs/core": "^0.19.1",
|
||||
"@humanfs/core": "^0.19.2",
|
||||
"@humanfs/types": "^0.15.0",
|
||||
"@humanwhocodes/retry": "^0.4.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=18.18.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@humanfs/types": {
|
||||
"version": "0.15.0",
|
||||
"resolved": "https://registry.npmjs.org/@humanfs/types/-/types-0.15.0.tgz",
|
||||
"integrity": "sha512-ZZ1w0aoQkwuUuC7Yf+7sdeaNfqQiiLcSRbfI08oAxqLtpXQr9AIVX7Ay7HLDuiLYAaFPu8oBYNq/QIi9URHJ3Q==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"engines": {
|
||||
"node": ">=18.18.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@humanwhocodes/module-importer": {
|
||||
"version": "1.0.1",
|
||||
"resolved": "https://registry.npmjs.org/@humanwhocodes/module-importer/-/module-importer-1.0.1.tgz",
|
||||
@@ -2198,6 +2214,12 @@
|
||||
"url": "https://opencollective.com/express"
|
||||
}
|
||||
},
|
||||
"node_modules/cookie-es": {
|
||||
"version": "3.1.1",
|
||||
"resolved": "https://registry.npmjs.org/cookie-es/-/cookie-es-3.1.1.tgz",
|
||||
"integrity": "sha512-UaXxwISYJPTr9hwQxMFYZ7kNhSXboMXP+Z3TRX6f1/NyaGPfuNUZOWP1pUEb75B2HjfklIYLVRfWiFZJyC6Npg==",
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/cookie-signature": {
|
||||
"version": "1.2.2",
|
||||
"resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.2.2.tgz",
|
||||
@@ -2889,9 +2911,9 @@
|
||||
"dev": true
|
||||
},
|
||||
"node_modules/fast-uri": {
|
||||
"version": "3.1.5",
|
||||
"resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz",
|
||||
"integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==",
|
||||
"version": "3.1.7",
|
||||
"resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.7.tgz",
|
||||
"integrity": "sha512-dOvZVzjdZdz7phd9v6jCbwxrBW3fK6n8Rc0CtdmM4bumzMnxywBYhuph6J819RRw/ku+rLbelwfMunktuzVVHg==",
|
||||
"funding": [
|
||||
{
|
||||
"type": "github",
|
||||
@@ -5951,9 +5973,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/qs": {
|
||||
"version": "6.15.3",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.15.3.tgz",
|
||||
"integrity": "sha512-O9gl3zCl5h5blw1KGUzQKhA5oUXSl8rwUIM5o0S3nCXMliSvy5Dzx7/DJcI+SwgICv+IneSZwhBh1oSyEHA71A==",
|
||||
"version": "6.16.0",
|
||||
"resolved": "https://registry.npmjs.org/qs/-/qs-6.16.0.tgz",
|
||||
"integrity": "sha512-h6fhOIaRrID2CbEY2fqs+7t+UXZo+MLAnU5gRIq85uFtdiUPCdsApMlHhXogKVM4HM2DVbIjGNTTYH2OcmP1vA==",
|
||||
"license": "BSD-3-Clause",
|
||||
"dependencies": {
|
||||
"es-define-property": "^1.0.1",
|
||||
@@ -6061,22 +6083,24 @@
|
||||
}
|
||||
},
|
||||
"node_modules/react": {
|
||||
"version": "19.2.4",
|
||||
"resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz",
|
||||
"integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==",
|
||||
"version": "19.2.8",
|
||||
"resolved": "https://registry.npmjs.org/react/-/react-19.2.8.tgz",
|
||||
"integrity": "sha512-PWaYA1L/q9u2u7xYQi+Y3L3Yfnie7XyLeaJICV1MGD6LprsBxcAqGjYyr0eY3p+QdsA+x/Irkt4Qif8D63+Sbw==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">=0.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/react-dom": {
|
||||
"version": "19.2.4",
|
||||
"resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.4.tgz",
|
||||
"integrity": "sha512-AXJdLo8kgMbimY95O2aKQqsz2iWi9jMgKJhRBAxECE4IFxfcazB2LmzloIoibJI3C12IlY20+KFaLv+71bUJeQ==",
|
||||
"version": "19.2.8",
|
||||
"resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.8.tgz",
|
||||
"integrity": "sha512-rVprimfGBG3DR+Tq0IQG2DT5PxKth1WIGDmj5yPmlzr4YBe7uyE+Du4oVqTDXZSHGGGXRtTJEGSSePyQCMBglQ==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"scheduler": "^0.27.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": "^19.2.4"
|
||||
"react": "^19.2.8"
|
||||
}
|
||||
},
|
||||
"node_modules/react-i18next": {
|
||||
@@ -6106,20 +6130,19 @@
|
||||
}
|
||||
},
|
||||
"node_modules/react-router": {
|
||||
"version": "7.18.2",
|
||||
"resolved": "https://registry.npmjs.org/react-router/-/react-router-7.18.2.tgz",
|
||||
"integrity": "sha512-aUVMjFm3GAPTTZL7oYr5E7ETiqfQCHRLH+B+5afnICvf0r7kkK4eR6SMuwbSTJw/7t+12khT/Kahij49fqOCIg==",
|
||||
"version": "8.3.1",
|
||||
"resolved": "https://registry.npmjs.org/react-router/-/react-router-8.3.1.tgz",
|
||||
"integrity": "sha512-TEOpiO2g0TJHEOJeRVv4amUFun9v1npCKszvcquNvzETUtJ8udV86ah5eFoHT7g26bsBvT6EiIhqulR8eDF++A==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"cookie": "^1.0.1",
|
||||
"set-cookie-parser": "^2.6.0"
|
||||
"cookie-es": "^3.1.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
"node": ">=22.22.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": ">=18",
|
||||
"react-dom": ">=18"
|
||||
"react": ">=19.2.7",
|
||||
"react-dom": ">=19.2.7"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"react-dom": {
|
||||
@@ -6143,6 +6166,28 @@
|
||||
"react-dom": ">=18"
|
||||
}
|
||||
},
|
||||
"node_modules/react-router-dom/node_modules/react-router": {
|
||||
"version": "7.18.2",
|
||||
"resolved": "https://registry.npmjs.org/react-router/-/react-router-7.18.2.tgz",
|
||||
"integrity": "sha512-aUVMjFm3GAPTTZL7oYr5E7ETiqfQCHRLH+B+5afnICvf0r7kkK4eR6SMuwbSTJw/7t+12khT/Kahij49fqOCIg==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"cookie": "^1.0.1",
|
||||
"set-cookie-parser": "^2.6.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": ">=18",
|
||||
"react-dom": ">=18"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"react-dom": {
|
||||
"optional": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/readable-stream": {
|
||||
"version": "2.3.8",
|
||||
"resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-2.3.8.tgz",
|
||||
|
||||
@@ -20,7 +20,8 @@
|
||||
},
|
||||
"overrides": {
|
||||
"hono": "4.12.34",
|
||||
"ip-address": "10.3.1"
|
||||
"ip-address": "10.3.1",
|
||||
"path-to-regexp": "^8.4.0"
|
||||
},
|
||||
"dependencies": {
|
||||
"@codemirror/autocomplete": "^6.18.6",
|
||||
@@ -47,7 +48,7 @@
|
||||
"react": "^19.1.0",
|
||||
"react-dom": "^19.1.0",
|
||||
"react-i18next": "^17.0.6",
|
||||
"react-router": "7.18.2",
|
||||
"react-router": "8.3.1",
|
||||
"react-router-dom": "7.18.2",
|
||||
"yaml": "^2.8.3"
|
||||
},
|
||||
|
||||
@@ -432,10 +432,12 @@ export default function CollectionDetails() {
|
||||
<input
|
||||
id="source-interval"
|
||||
className="input"
|
||||
type="text"
|
||||
type="number"
|
||||
min="1"
|
||||
step="1"
|
||||
value={newSourceInterval}
|
||||
onChange={(e) => setNewSourceInterval(e.target.value)}
|
||||
placeholder="e.g. 1h, 30m"
|
||||
placeholder="e.g. 60 (minutes)"
|
||||
/>
|
||||
</div>
|
||||
<button className="btn btn-primary" type="submit" disabled={!newSourceUrl.trim() || addingSource}>
|
||||
|
||||
@@ -363,7 +363,7 @@ export default function Traces() {
|
||||
const [apiCount, setApiCount] = useState(0)
|
||||
const [backendCount, setBackendCount] = useState(0)
|
||||
const [loading, setLoading] = useState(true)
|
||||
const [expandedRow, setExpandedRow] = useState(null)
|
||||
const [expandedTraceId, setExpandedTraceId] = useState(null)
|
||||
// detail holds the full record for the currently expanded row, fetched on
|
||||
// demand from /api/traces/:id (the list response omits the bodies).
|
||||
const [detail, setDetail] = useState(null)
|
||||
@@ -381,7 +381,8 @@ export default function Traces() {
|
||||
duration: (a, b) => (a.duration || 0) - (b.duration || 0),
|
||||
}
|
||||
const toggleSort = (key) => {
|
||||
setExpandedRow(null)
|
||||
setExpandedTraceId(null)
|
||||
setDetail(null)
|
||||
setSort(s => s.key === key ? { key, dir: s.dir === 'asc' ? 'desc' : 'asc' } : { key, dir: 'asc' })
|
||||
}
|
||||
const sortableTh = (key, label, props = {}) => (
|
||||
@@ -454,20 +455,21 @@ export default function Traces() {
|
||||
|
||||
useEffect(() => {
|
||||
setLoading(true)
|
||||
setExpandedRow(null)
|
||||
setExpandedTraceId(null)
|
||||
setDetail(null)
|
||||
fetchTraces()
|
||||
}, [fetchTraces])
|
||||
|
||||
// Expanding a row pulls the full record (bodies, data fields, audio
|
||||
// snippets) that the list response deliberately omits.
|
||||
const toggleRow = useCallback(async (index, row) => {
|
||||
if (expandedRow === index) {
|
||||
setExpandedRow(null)
|
||||
const toggleRow = useCallback(async (row, index) => {
|
||||
const traceKey = row?.id ?? index
|
||||
if (expandedTraceId === traceKey) {
|
||||
setExpandedTraceId(null)
|
||||
setDetail(null)
|
||||
return
|
||||
}
|
||||
setExpandedRow(index)
|
||||
setExpandedTraceId(traceKey)
|
||||
setDetail(null)
|
||||
if (!row?.id) return
|
||||
try {
|
||||
@@ -478,7 +480,7 @@ export default function Traces() {
|
||||
} catch {
|
||||
// Fall back to the summary view; the row still renders what it has.
|
||||
}
|
||||
}, [expandedRow, activeTab])
|
||||
}, [expandedTraceId, activeTab])
|
||||
|
||||
// Auto-refresh every 5 seconds
|
||||
useEffect(() => {
|
||||
@@ -491,7 +493,7 @@ export default function Traces() {
|
||||
if (activeTab === 'api') await tracesApi.clear()
|
||||
else await tracesApi.clearBackend()
|
||||
setTraces([])
|
||||
setExpandedRow(null)
|
||||
setExpandedTraceId(null)
|
||||
setDetail(null)
|
||||
addToast('Traces cleared', 'success')
|
||||
} catch (err) {
|
||||
@@ -521,7 +523,7 @@ export default function Traces() {
|
||||
}
|
||||
|
||||
// Reset sort + expansion when switching trace tabs (columns differ).
|
||||
useEffect(() => { setSort({ key: null, dir: 'asc' }); setExpandedRow(null); setDetail(null) }, [activeTab])
|
||||
useEffect(() => { setSort({ key: null, dir: 'asc' }); setExpandedTraceId(null); setDetail(null) }, [activeTab])
|
||||
|
||||
const sortedTraces = sort.key && TRACE_SORT[sort.key]
|
||||
? [...traces].sort((a, b) => sort.dir === 'asc' ? TRACE_SORT[sort.key](a, b) : TRACE_SORT[sort.key](b, a))
|
||||
@@ -659,9 +661,9 @@ export default function Traces() {
|
||||
</thead>
|
||||
<tbody>
|
||||
{sortedTraces.map((trace, i) => (
|
||||
<React.Fragment key={i}>
|
||||
<tr onClick={() => toggleRow(i, trace)} className="clickable">
|
||||
<td><i className={`fas fa-chevron-${expandedRow === i ? 'down' : 'right'} text-xs`} /></td>
|
||||
<React.Fragment key={trace.id ?? i}>
|
||||
<tr onClick={() => toggleRow(trace, i)} className="clickable">
|
||||
<td><i className={`fas fa-chevron-${expandedTraceId === (trace.id ?? i) ? 'down' : 'right'} text-xs`} /></td>
|
||||
<td><span className="badge badge-info">{trace.request?.method || '-'}</span></td>
|
||||
<td className="text-mono text-sm">{trace.request?.path || '-'}</td>
|
||||
<td className="text-sub cell-clip" title={trace.user_name || trace.user_id || ''}>{trace.user_name || trace.user_id || '-'}</td>
|
||||
@@ -681,7 +683,7 @@ export default function Traces() {
|
||||
: <i className="fas fa-check-circle text-success" />}
|
||||
</td>
|
||||
</tr>
|
||||
{expandedRow === i && (
|
||||
{expandedTraceId === (trace.id ?? i) && (
|
||||
<tr>
|
||||
<td colSpan="7" className="p-0">
|
||||
<ApiTraceDetail trace={detail && detail.id === trace.id ? detail : trace} />
|
||||
@@ -707,9 +709,9 @@ export default function Traces() {
|
||||
</thead>
|
||||
<tbody>
|
||||
{sortedTraces.map((trace, i) => (
|
||||
<React.Fragment key={i}>
|
||||
<tr onClick={() => toggleRow(i, trace)} className="clickable">
|
||||
<td><i className={`fas fa-chevron-${expandedRow === i ? 'down' : 'right'} text-xs`} /></td>
|
||||
<React.Fragment key={trace.id ?? i}>
|
||||
<tr onClick={() => toggleRow(trace, i)} className="clickable">
|
||||
<td><i className={`fas fa-chevron-${expandedTraceId === (trace.id ?? i) ? 'down' : 'right'} text-xs`} /></td>
|
||||
<td><span style={typeBadgeStyle(trace.type)}>{trace.type || '-'}</span></td>
|
||||
<td className="text-sub nowrap">{formatDateTime(trace.timestamp)}</td>
|
||||
<td className="text-mono text-sm">{trace.model_name || '-'}</td>
|
||||
@@ -725,7 +727,7 @@ export default function Traces() {
|
||||
: <i className="fas fa-check-circle text-success" />}
|
||||
</td>
|
||||
</tr>
|
||||
{expandedRow === i && (
|
||||
{expandedTraceId === (trace.id ?? i) && (
|
||||
<tr>
|
||||
<td colSpan="7" className="p-0">
|
||||
<BackendTraceDetail trace={detail && detail.id === trace.id ? detail : trace} />
|
||||
|
||||
Vendored
+4
-1
@@ -457,7 +457,10 @@ export const agentCollectionsApi = {
|
||||
reset: (name, userId) => postJSON(`/api/agents/collections/${enc(name)}/reset${userQ(userId)}`),
|
||||
deleteEntry: (name, entry, userId) => fetchJSON(`/api/agents/collections/${enc(name)}/entry/delete${userQ(userId)}`, { method: 'DELETE', body: JSON.stringify({ entry }), headers: { 'Content-Type': 'application/json' } }),
|
||||
sources: (name, userId) => fetchJSON(`/api/agents/collections/${enc(name)}/sources${userQ(userId)}`),
|
||||
addSource: (name, url, interval, userId) => postJSON(`/api/agents/collections/${enc(name)}/sources${userQ(userId)}`, { url, update_interval: interval }),
|
||||
addSource: (name, url, interval, userId) => postJSON(`/api/agents/collections/${enc(name)}/sources${userQ(userId)}`, {
|
||||
url,
|
||||
update_interval: interval === undefined ? undefined : Number(interval),
|
||||
}),
|
||||
removeSource: (name, url, userId) => fetchJSON(`/api/agents/collections/${enc(name)}/sources${userQ(userId)}`, { method: 'DELETE', body: JSON.stringify({ url }), headers: { 'Content-Type': 'application/json' } }),
|
||||
}
|
||||
|
||||
|
||||
@@ -353,10 +353,12 @@ type FaceEmbedResponse struct {
|
||||
// FaceRegisterRequest enrolls a face into the 1:N recognition store.
|
||||
type FaceRegisterRequest struct {
|
||||
BasicModelRequest
|
||||
Img string `json:"img"`
|
||||
Name string `json:"name"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
Store string `json:"store,omitempty"` // vector store model; empty = local-store default
|
||||
RegisteredAt time.Time `json:"registered_at,omitempty"` // original enrollment time when replaying a saved embedding
|
||||
Embedding []float32 `json:"embedding,omitempty"`
|
||||
Img string `json:"img"`
|
||||
Name string `json:"name"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
Store string `json:"store,omitempty"` // vector store model; empty = local-store default
|
||||
}
|
||||
|
||||
type FaceRegisterResponse struct {
|
||||
|
||||
@@ -56,5 +56,6 @@ type Match struct {
|
||||
var (
|
||||
ErrNotFound = errors.New("facerecognition: id not found")
|
||||
ErrEmptyEmbedding = errors.New("facerecognition: embedding is empty")
|
||||
ErrInvalidEmbedding = errors.New("facerecognition: embedding must be finite and nonzero")
|
||||
ErrDimensionMismatch = errors.New("facerecognition: embedding dimension mismatch")
|
||||
)
|
||||
@@ -0,0 +1,67 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package facerecognition
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
ggrpc "google.golang.org/grpc"
|
||||
)
|
||||
|
||||
func TestEnrollmentReplay(t *testing.T) { RegisterFailHandler(Fail); RunSpecs(t, "Enrollment replay") }
|
||||
|
||||
type replayStore struct {
|
||||
grpc.Backend
|
||||
mu sync.Mutex
|
||||
entries map[string][]byte
|
||||
}
|
||||
|
||||
func (s *replayStore) StoresSet(_ context.Context, in *pb.StoresSetOptions, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i, k := range in.Keys {
|
||||
b, _ := json.Marshal(k.Floats)
|
||||
s.entries[string(b)] = append([]byte(nil), in.Values[i].Bytes...)
|
||||
}
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
var _ = Describe("Enrollment replay", func() {
|
||||
It("keeps the identity across registry instances and a cleared store", func(ctx SpecContext) {
|
||||
storage := &replayStore{entries: map[string][]byte{}}
|
||||
newRegistry := func() Registry {
|
||||
return NewStoreRegistry(func(context.Context, string) (grpc.Backend, error) { return storage, nil }, "faces", 0)
|
||||
}
|
||||
vector := []float32{1, 0, 0, 0}
|
||||
meta := Metadata{Name: "Alice", RegisteredAt: time.Now().UTC()}
|
||||
first, err := newRegistry().Register(ctx, vector, meta)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
again, err := newRegistry().Register(ctx, vector, meta)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(again).To(Equal(first))
|
||||
Expect(storage.entries).To(HaveLen(1))
|
||||
storage.entries = map[string][]byte{}
|
||||
restored, err := newRegistry().Register(ctx, vector, meta)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(restored).To(Equal(first))
|
||||
Expect(storage.entries).To(HaveLen(1))
|
||||
})
|
||||
It("rejects zero and non-finite embeddings before writing", func(ctx SpecContext) {
|
||||
for _, v := range [][]float32{{0, 0}, {float32(math.NaN()), 1}, {float32(math.Inf(1)), 1}} {
|
||||
storage := &replayStore{entries: map[string][]byte{}}
|
||||
reg := NewStoreRegistry(func(context.Context, string) (grpc.Backend, error) { return storage, nil }, "faces", 0)
|
||||
_, err := reg.Register(ctx, v, Metadata{Name: "Alice"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(storage.entries).To(BeEmpty())
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -2,8 +2,10 @@ package facerecognition
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -57,13 +59,32 @@ func (r *storeRegistry) Register(ctx context.Context, embedding []float32, meta
|
||||
if r.dim != 0 && len(embedding) != r.dim {
|
||||
return Metadata{}, fmt.Errorf("%w: expected %d, got %d", ErrDimensionMismatch, r.dim, len(embedding))
|
||||
}
|
||||
var norm float64
|
||||
key := make([]byte, 4*len(embedding))
|
||||
for i, value := range embedding {
|
||||
if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) {
|
||||
return Metadata{}, ErrInvalidEmbedding
|
||||
}
|
||||
norm += float64(value) * float64(value)
|
||||
// The store treats negative and positive zero as the same key.
|
||||
if value == 0 {
|
||||
value = 0
|
||||
}
|
||||
binary.LittleEndian.PutUint32(key[i*4:], math.Float32bits(value))
|
||||
}
|
||||
if norm == 0 {
|
||||
return Metadata{}, ErrInvalidEmbedding
|
||||
}
|
||||
|
||||
backend, err := r.resolve(ctx, r.storeName)
|
||||
if err != nil {
|
||||
return Metadata{}, fmt.Errorf("facerecognition: resolve store: %w", err)
|
||||
}
|
||||
|
||||
meta.ID = uuid.NewString()
|
||||
// The vector store upserts by the exact embedding. Derive the ID from the
|
||||
// same key so replaying a saved vector preserves identity across replicas
|
||||
// and after the in-memory store restarts.
|
||||
meta.ID = uuid.NewSHA1(uuid.NewSHA1(uuid.NameSpaceOID, []byte(r.storeName)), key).String()
|
||||
if meta.RegisteredAt.IsZero() {
|
||||
meta.RegisteredAt = time.Now().UTC()
|
||||
}
|
||||
|
||||
@@ -467,6 +467,12 @@ func SubjectNodeFilesStage(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.stage"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesRelease tells a serve-backend node to evict one request's ephemeral cache keys.
|
||||
// Reply: {error}
|
||||
func SubjectNodeFilesRelease(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.release"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesTemp tells a serve-backend node to allocate a temp file.
|
||||
// Reply: {local_path, error}
|
||||
func SubjectNodeFilesTemp(nodeID string) string {
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
package nodes
|
||||
|
||||
import "context"
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FileStager abstracts file transfer between frontend and backend nodes
|
||||
// in distributed mode. Two implementations exist:
|
||||
@@ -29,7 +34,59 @@ type FileStager interface {
|
||||
// StageRemoteToStore uploads a remote file to shared storage.
|
||||
StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error
|
||||
|
||||
// ReleaseRemote removes one ephemeral key from the remote node.
|
||||
ReleaseRemote(ctx context.Context, nodeID, key string) error
|
||||
|
||||
// ListRemoteDir returns relative file paths within a directory on the remote node.
|
||||
// keyPrefix is a storage-style key prefix (e.g. "models/mymodel").
|
||||
ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error)
|
||||
}
|
||||
|
||||
// RequestFileReleaser removes all ephemeral keys staged for one inference in
|
||||
// one transport operation. FileStagingClient falls back to ReleaseRemote for
|
||||
// stagers that do not implement this optional rolling-upgrade extension.
|
||||
type RequestFileReleaser interface {
|
||||
ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error
|
||||
}
|
||||
|
||||
func validateEphemeralRequestRelease(requestID string, keys []string) error {
|
||||
if err := validateEphemeralRequestID(requestID); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(keys) == 0 {
|
||||
return fmt.Errorf("release batch must contain at least one key")
|
||||
}
|
||||
for _, key := range keys {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
parts := strings.Split(key, "/")
|
||||
if parts[2] != requestID {
|
||||
return fmt.Errorf("release batch mixes request IDs %q and %q", requestID, parts[2])
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateEphemeralRequestID(requestID string) error {
|
||||
if requestID == "" || strings.ContainsAny(requestID, "/\\") || path.Clean(requestID) != requestID || requestID == "." || requestID == ".." {
|
||||
return fmt.Errorf("invalid ephemeral request ID %q", requestID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateEphemeralReleaseKey(key string) error {
|
||||
if strings.Contains(key, "\\") || path.Clean(key) != key {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
parts := strings.Split(key, "/")
|
||||
if len(parts) != 4 || parts[0] != "ephemeral" {
|
||||
return fmt.Errorf("release key %q must identify one file below ephemeral/", key)
|
||||
}
|
||||
for _, part := range parts[1:] {
|
||||
if part == "" || part == "." || part == ".." {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
@@ -10,6 +11,7 @@ import (
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
@@ -81,6 +83,86 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st
|
||||
}
|
||||
}
|
||||
|
||||
// ReleaseRemote removes one exact ephemeral key from a backend node.
|
||||
func (h *HTTPFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
addr, err := h.httpAddrFor(nodeID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err)
|
||||
}
|
||||
releaseURL := (&url.URL{Scheme: "http", Host: addr, Path: "/v1/files/" + key}).String()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, releaseURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating release request for %q: %w", key, err)
|
||||
}
|
||||
if h.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+h.token)
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("releasing %q from node %s: %w", key, nodeID, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return fmt.Errorf("releasing %q from node %s: status %d: %s", key, nodeID, resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseRemoteRequest removes one inference's staged inputs with one HTTP
|
||||
// request. Older workers return 404 for the batch endpoint, so the client
|
||||
// retries through the exact-key API during rolling upgrades.
|
||||
func (h *HTTPFileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error {
|
||||
if err := validateEphemeralRequestRelease(requestID, keys); err != nil {
|
||||
return err
|
||||
}
|
||||
addr, err := h.httpAddrFor(nodeID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err)
|
||||
}
|
||||
payload, err := json.Marshal(struct {
|
||||
RequestID string `json:"request_id"`
|
||||
}{RequestID: requestID})
|
||||
if err != nil {
|
||||
return fmt.Errorf("encoding request release: %w", err)
|
||||
}
|
||||
releaseURL := (&url.URL{Scheme: "http", Host: addr, Path: "/v1/files-release"}).String()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, releaseURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating request release: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if h.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+h.token)
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("releasing request inputs from node %s: %w", nodeID, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
|
||||
return h.releaseRemoteKeys(ctx, nodeID, keys)
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return fmt.Errorf("releasing request inputs from node %s: status %d: %s", nodeID, resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *HTTPFileStager) releaseRemoteKeys(ctx context.Context, nodeID string, keys []string) error {
|
||||
var releaseErrors []error
|
||||
for _, key := range keys {
|
||||
if err := h.ReleaseRemote(ctx, nodeID, key); err != nil {
|
||||
releaseErrors = append(releaseErrors, err)
|
||||
}
|
||||
}
|
||||
return errors.Join(releaseErrors...)
|
||||
}
|
||||
|
||||
func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
|
||||
xlog.Debug("Staging file to remote node via HTTP", "node", nodeID, "localPath", localPath, "key", key)
|
||||
|
||||
@@ -90,7 +172,9 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke
|
||||
}
|
||||
|
||||
// Probe: check if the remote already has the file with matching content hash.
|
||||
if remotePath, ok := h.probeExisting(ctx, addr, localPath, key); ok {
|
||||
if remotePath, ok, probeErr := h.probeExisting(ctx, addr, localPath, key); probeErr != nil {
|
||||
return "", fmt.Errorf("claiming existing file on node %s: %w", nodeID, probeErr)
|
||||
} else if ok {
|
||||
xlog.Info("Upload skipped (file already exists with matching hash)", "node", nodeID, "key", key, "remotePath", remotePath)
|
||||
return remotePath, nil
|
||||
}
|
||||
@@ -439,14 +523,15 @@ func isTransientError(err error) bool {
|
||||
|
||||
// probeExisting sends a HEAD request to check if the remote already has the
|
||||
// file with a matching SHA-256 hash. Returns the remote path and true if the
|
||||
// upload can be skipped. Any errors (including 405 from older servers) silently
|
||||
// fall through so the caller proceeds with a normal PUT.
|
||||
func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key string) (string, bool) {
|
||||
// upload can be skipped. HEAD and hash errors fall through to a normal PUT.
|
||||
// Matching ephemeral files are claimed first; a 404 or 405 claim response
|
||||
// identifies an older worker and also falls through to PUT.
|
||||
func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key string) (string, bool, error) {
|
||||
url := fmt.Sprintf("http://%s/v1/files/%s", addr, key)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
|
||||
if err != nil {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
if h.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+h.token)
|
||||
@@ -454,18 +539,18 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key
|
||||
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
remotePath := resp.Header.Get(HeaderLocalPath)
|
||||
remoteHash := resp.Header.Get(HeaderContentSHA256)
|
||||
if remotePath == "" || remoteHash == "" {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
// A 200 with a content hash is proof the worker is alive and serving right
|
||||
@@ -475,14 +560,53 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key
|
||||
|
||||
localHash, err := hashLocalCached(ctx, localPath)
|
||||
if err != nil {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
if localHash != remoteHash {
|
||||
return "", false
|
||||
return "", false, nil
|
||||
}
|
||||
|
||||
return remotePath, true
|
||||
if strings.HasPrefix(key, "ephemeral/") {
|
||||
claimed, err := h.claimExisting(ctx, addr, key)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
if !claimed {
|
||||
return "", false, nil
|
||||
}
|
||||
}
|
||||
|
||||
return remotePath, true, nil
|
||||
}
|
||||
|
||||
func (h *HTTPFileStager) claimExisting(ctx context.Context, addr, key string) (bool, error) {
|
||||
claimURL := (&url.URL{
|
||||
Scheme: "http",
|
||||
Host: addr,
|
||||
Path: "/v1/files/" + key,
|
||||
RawQuery: "claim=1",
|
||||
}).String()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, claimURL, nil)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("creating claim request for %q: %w", key, err)
|
||||
}
|
||||
if h.token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+h.token)
|
||||
}
|
||||
resp, err := h.client.Do(req)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("claiming %q: %w", key, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
|
||||
return false, nil
|
||||
}
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return false, fmt.Errorf("claiming %q: status %d: %s", key, resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// hashChunkSize is how much of a file is hashed between activity ticks and
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type releaseTestSubscription struct{}
|
||||
|
||||
func (releaseTestSubscription) Unsubscribe() error { return nil }
|
||||
|
||||
type releaseTestMessaging struct {
|
||||
subject string
|
||||
payload []byte
|
||||
onRequest func()
|
||||
requestCalled bool
|
||||
requestCount int
|
||||
timeout time.Duration
|
||||
replies [][]byte
|
||||
}
|
||||
|
||||
func (m *releaseTestMessaging) Publish(string, any) error { return nil }
|
||||
func (m *releaseTestMessaging) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseTestSubscription{}, nil
|
||||
}
|
||||
func (m *releaseTestMessaging) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) {
|
||||
m.subject = subject
|
||||
m.payload = append([]byte(nil), data...)
|
||||
m.requestCalled = true
|
||||
m.requestCount++
|
||||
m.timeout = timeout
|
||||
if m.onRequest != nil {
|
||||
m.onRequest()
|
||||
}
|
||||
if m.requestCount <= len(m.replies) {
|
||||
return append([]byte(nil), m.replies[m.requestCount-1]...), nil
|
||||
}
|
||||
return []byte(`{}`), nil
|
||||
}
|
||||
func (m *releaseTestMessaging) IsConnected() bool { return true }
|
||||
func (m *releaseTestMessaging) Close() {}
|
||||
|
||||
var _ = Describe("File stager exact-key release", func() {
|
||||
startReleaseServer := func(stagingDir, token string) (*HTTPFileStager, func()) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
server, err := StartFileTransferServerWithListener(
|
||||
listener,
|
||||
stagingDir,
|
||||
GinkgoT().TempDir(),
|
||||
GinkgoT().TempDir(),
|
||||
token,
|
||||
0,
|
||||
)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return NewHTTPFileStager(func(string) (string, error) {
|
||||
return listener.Addr().String(), nil
|
||||
}, token), func() {
|
||||
Expect(server.Shutdown(context.Background())).To(Succeed())
|
||||
}
|
||||
}
|
||||
|
||||
It("transmits URL metacharacters as the exact key", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(stagingDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
|
||||
key := "ephemeral/request-id/audio/name ?#%2F.wav"
|
||||
exactPath := filepath.Join(categoryDir, "name ?#%2F.wav")
|
||||
wrongPath := filepath.Join(categoryDir, "name ")
|
||||
for _, path := range []string{exactPath, exactPath + hashSidecarSuffix, exactPath + targetSidecarSuffix, wrongPath} {
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", key)).To(Succeed())
|
||||
|
||||
Expect(exactPath).NotTo(BeAnExistingFile())
|
||||
Expect(exactPath + hashSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(exactPath + targetSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(wrongPath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("is idempotent and prunes empty category and request directories", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
path := filepath.Join(stagingDir, "ephemeral", "request-id", "audio", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
for range 2 {
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", "ephemeral/request-id/audio/input.wav")).To(Succeed())
|
||||
}
|
||||
|
||||
Expect(filepath.Join(stagingDir, "ephemeral", "request-id", "audio")).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(stagingDir, "ephemeral", "request-id")).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(stagingDir, "ephemeral")).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("releases one request's HTTP inputs in one batch", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-id/input.wav",
|
||||
"ephemeral/images/request-id/frame.jpg",
|
||||
}
|
||||
for _, key := range keys {
|
||||
path := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
Expect(stager.ReleaseRemoteRequest(context.Background(), "node-1", "request-id", keys)).To(Succeed())
|
||||
|
||||
for _, key := range keys {
|
||||
Expect(filepath.Join(stagingDir, filepath.FromSlash(key))).NotTo(BeAnExistingFile())
|
||||
}
|
||||
})
|
||||
|
||||
It("falls back to exact HTTP releases for an older worker", func() {
|
||||
batchCalls := 0
|
||||
exactCalls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/v1/files-release" {
|
||||
batchCalls++
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/v1/files/") && r.Method == http.MethodDelete {
|
||||
exactCalls++
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
http.Error(w, "unexpected request", http.StatusInternalServerError)
|
||||
}))
|
||||
DeferCleanup(server.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(server.URL, "http://"), nil
|
||||
}, "")
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-id/input.wav",
|
||||
"ephemeral/images/request-id/frame.jpg",
|
||||
}
|
||||
|
||||
Expect(stager.ReleaseRemoteRequest(context.Background(), "node-1", "request-id", keys)).To(Succeed())
|
||||
|
||||
Expect(batchCalls).To(Equal(1))
|
||||
Expect(exactCalls).To(Equal(2))
|
||||
})
|
||||
|
||||
It("rejects non-ephemeral and traversing keys before making a request", func() {
|
||||
resolved := false
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
resolved = true
|
||||
return "127.0.0.1:1", nil
|
||||
}, "token")
|
||||
|
||||
for _, key := range []string{
|
||||
"models/model.gguf",
|
||||
"ephemeral/../models/model.gguf",
|
||||
"ephemeral/request-id/../../model.gguf",
|
||||
"/ephemeral/request-id/audio/input.wav",
|
||||
} {
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", key)).NotTo(Succeed(), key)
|
||||
}
|
||||
Expect(resolved).To(BeFalse())
|
||||
})
|
||||
|
||||
It("rejects symlink escapes", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
outsideDir := GinkgoT().TempDir()
|
||||
outsidePath := filepath.Join(outsideDir, "input.wav")
|
||||
Expect(os.WriteFile(outsidePath, []byte("keep"), 0640)).To(Succeed())
|
||||
requestDir := filepath.Join(stagingDir, "ephemeral", "request-id")
|
||||
Expect(os.MkdirAll(requestDir, 0750)).To(Succeed())
|
||||
Expect(os.Symlink(outsideDir, filepath.Join(requestDir, "audio"))).To(Succeed())
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
DeferCleanup(stop)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", "ephemeral/request-id/audio/input.wav")).NotTo(Succeed())
|
||||
Expect(outsidePath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("rejects symlinked files and sidecars without deleting their targets", func() {
|
||||
for _, linkedName := range []string{"input.wav", "input.wav" + hashSidecarSuffix, "input.wav" + targetSidecarSuffix} {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(stagingDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
target := filepath.Join(categoryDir, "input.wav")
|
||||
if linkedName != "input.wav" {
|
||||
Expect(os.WriteFile(target, []byte("input"), 0640)).To(Succeed())
|
||||
}
|
||||
preserved := filepath.Join(stagingDir, "ephemeral", "preserved-"+linkedName)
|
||||
Expect(os.WriteFile(preserved, []byte("keep"), 0640)).To(Succeed())
|
||||
Expect(os.Symlink(preserved, filepath.Join(categoryDir, linkedName))).To(Succeed())
|
||||
|
||||
stager, stop := startReleaseServer(stagingDir, "release-token")
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node-1", "ephemeral/request-id/audio/input.wav")).NotTo(Succeed(), linkedName)
|
||||
stop()
|
||||
Expect(preserved).To(BeAnExistingFile(), linkedName)
|
||||
}
|
||||
})
|
||||
|
||||
It("evicts the worker cache before deleting the shared object", func() {
|
||||
storeRoot := GinkgoT().TempDir()
|
||||
cacheRoot := GinkgoT().TempDir()
|
||||
store, err := storage.NewFilesystemStore(storeRoot)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, cacheRoot)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
key := "ephemeral/request-id/audio/input.wav"
|
||||
Expect(store.Put(context.Background(), key, strings.NewReader("shared"))).To(Succeed())
|
||||
|
||||
client := &releaseTestMessaging{}
|
||||
client.onRequest = func() {
|
||||
exists, existsErr := store.Exists(context.Background(), key)
|
||||
Expect(existsErr).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeTrue())
|
||||
}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
Expect(stager.ReleaseRemote(context.Background(), "node.one", key)).To(Succeed())
|
||||
|
||||
Expect(client.requestCalled).To(BeTrue())
|
||||
Expect(client.subject).To(Equal(messaging.SubjectNodeFilesRelease("node.one")))
|
||||
var payload fileReleaseRequest
|
||||
Expect(json.Unmarshal(client.payload, &payload)).To(Succeed())
|
||||
Expect(payload.Key).To(Equal(key))
|
||||
exists, err := store.Exists(context.Background(), key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeFalse())
|
||||
})
|
||||
|
||||
It("evicts a request's S3 inputs with one NATS round trip", func() {
|
||||
storeRoot := GinkgoT().TempDir()
|
||||
store, err := storage.NewFilesystemStore(storeRoot)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-id/input.wav",
|
||||
"ephemeral/images/request-id/frame.jpg",
|
||||
}
|
||||
for _, key := range keys {
|
||||
Expect(store.Put(context.Background(), key, strings.NewReader("shared"))).To(Succeed())
|
||||
}
|
||||
client := &releaseTestMessaging{}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
|
||||
Expect(stager.ReleaseRemoteRequest(context.Background(), "node.one", "request-id", keys)).To(Succeed())
|
||||
|
||||
Expect(client.requestCount).To(Equal(1))
|
||||
var payload fileReleaseRequest
|
||||
Expect(json.Unmarshal(client.payload, &payload)).To(Succeed())
|
||||
Expect(payload.Key).To(BeEmpty())
|
||||
Expect(payload.RequestID).To(Equal("request-id"))
|
||||
for _, key := range keys {
|
||||
exists, existsErr := store.Exists(context.Background(), key)
|
||||
Expect(existsErr).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeFalse())
|
||||
}
|
||||
})
|
||||
|
||||
It("keeps worker coordination fixed-size for large requests", func() {
|
||||
store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
keys := make([]string, 2048)
|
||||
for i := range keys {
|
||||
keys[i] = fmt.Sprintf("ephemeral/inputs/request-id/input-%d.bin", i)
|
||||
}
|
||||
client := &releaseTestMessaging{}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
|
||||
Expect(stager.ReleaseRemoteRequest(context.Background(), "node.one", "request-id", keys)).To(Succeed())
|
||||
|
||||
Expect(client.requestCount).To(Equal(1))
|
||||
Expect(len(client.payload)).To(BeNumerically("<", 128))
|
||||
var payload fileReleaseRequest
|
||||
Expect(json.Unmarshal(client.payload, &payload)).To(Succeed())
|
||||
Expect(payload.RequestID).To(Equal("request-id"))
|
||||
})
|
||||
|
||||
It("falls back to exact NATS releases for an older worker", func() {
|
||||
store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-id/input.wav",
|
||||
"ephemeral/images/request-id/frame.jpg",
|
||||
}
|
||||
for _, key := range keys {
|
||||
Expect(store.Put(context.Background(), key, strings.NewReader("shared"))).To(Succeed())
|
||||
}
|
||||
client := &releaseTestMessaging{replies: [][]byte{
|
||||
[]byte(`{"error":"batch payload unsupported"}`),
|
||||
[]byte(`{}`),
|
||||
[]byte(`{}`),
|
||||
}}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
|
||||
Expect(stager.ReleaseRemoteRequest(context.Background(), "node.one", "request-id", keys)).To(Succeed())
|
||||
|
||||
Expect(client.requestCount).To(Equal(3))
|
||||
for _, key := range keys {
|
||||
exists, existsErr := store.Exists(context.Background(), key)
|
||||
Expect(existsErr).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeFalse())
|
||||
}
|
||||
})
|
||||
|
||||
It("rejects release batches that mix request IDs", func() {
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-one/input.wav",
|
||||
"ephemeral/images/request-two/frame.jpg",
|
||||
}
|
||||
Expect(validateEphemeralRequestRelease("request-one", keys)).To(MatchError(ContainSubstring("mixes request IDs")))
|
||||
})
|
||||
|
||||
It("does not send a release request after cleanup is canceled", func() {
|
||||
store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseTestMessaging{}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
Expect(stager.ReleaseRemote(ctx, "node.one", "ephemeral/request-id/audio/input.wav")).To(MatchError(context.Canceled))
|
||||
Expect(client.requestCalled).To(BeFalse())
|
||||
})
|
||||
|
||||
It("bounds the NATS release wait by the remaining cleanup deadline", func() {
|
||||
store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseTestMessaging{}
|
||||
stager := NewS3NATSFileStager(fm, client)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
Expect(stager.ReleaseRemote(ctx, "node.one", "ephemeral/request-id/audio/input.wav")).To(Succeed())
|
||||
Expect(client.timeout).To(BeNumerically(">", time.Second))
|
||||
Expect(client.timeout).To(BeNumerically("<=", 2*time.Second))
|
||||
})
|
||||
})
|
||||
@@ -2,6 +2,7 @@ package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
@@ -48,6 +49,15 @@ type fileStageReply struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileReleaseRequest struct {
|
||||
Key string `json:"key,omitempty"`
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
}
|
||||
|
||||
type fileReleaseReply struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileTempRequest struct{}
|
||||
|
||||
type fileTempReply struct {
|
||||
@@ -181,3 +191,79 @@ func (s *S3NATSFileStager) StageRemoteToStore(ctx context.Context, nodeID, remot
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseRemote evicts one exact ephemeral key from the worker before deleting
|
||||
// the shared object.
|
||||
func (s *S3NATSFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.releaseWorkerKeys(ctx, nodeID, fileReleaseRequest{Key: key}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.fm.Delete(ctx, key); err != nil {
|
||||
return fmt.Errorf("deleting shared object %q: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseRemoteRequest evicts one inference's inputs with one NATS round trip.
|
||||
// A worker that only understands the exact-key payload returns an error, so the
|
||||
// frontend retries each key during a rolling upgrade.
|
||||
func (s *S3NATSFileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error {
|
||||
if err := validateEphemeralRequestRelease(requestID, keys); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.releaseWorkerKeys(ctx, nodeID, fileReleaseRequest{RequestID: requestID}); err != nil {
|
||||
var fallbackErrors []error
|
||||
for _, key := range keys {
|
||||
if fallbackErr := s.ReleaseRemote(ctx, nodeID, key); fallbackErr != nil {
|
||||
fallbackErrors = append(fallbackErrors, fallbackErr)
|
||||
}
|
||||
}
|
||||
if fallbackErr := errors.Join(fallbackErrors...); fallbackErr != nil {
|
||||
return errors.Join(err, fmt.Errorf("exact-key release fallback: %w", fallbackErr))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var deleteErrors []error
|
||||
for _, key := range keys {
|
||||
if err := s.fm.Delete(ctx, key); err != nil {
|
||||
deleteErrors = append(deleteErrors, fmt.Errorf("deleting shared object %q: %w", key, err))
|
||||
}
|
||||
}
|
||||
return errors.Join(deleteErrors...)
|
||||
}
|
||||
|
||||
func (s *S3NATSFileStager) releaseWorkerKeys(ctx context.Context, nodeID string, request fileReleaseRequest) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
timeout := 30 * time.Second
|
||||
if deadline, ok := ctx.Deadline(); ok {
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
timeout = min(timeout, remaining)
|
||||
}
|
||||
reply, err := messaging.RequestJSON[fileReleaseRequest, fileReleaseReply](
|
||||
s.nats,
|
||||
messaging.SubjectNodeFilesRelease(nodeID),
|
||||
request,
|
||||
timeout,
|
||||
)
|
||||
if err != nil {
|
||||
if contextErr := ctx.Err(); contextErr != nil {
|
||||
return contextErr
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
return fmt.Errorf("backend release failed: %s", reply.Error)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -16,8 +16,11 @@ import (
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/xlog"
|
||||
ggrpc "google.golang.org/grpc"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const stagedInputReleaseTimeout = 30 * time.Second
|
||||
|
||||
// FileStagingClient wraps a grpc.Backend to transparently handle file transfer
|
||||
// for distributed mode. Input files are staged on the backend node before the
|
||||
// gRPC call. Output files are retrieved from the backend after the call.
|
||||
@@ -48,21 +51,70 @@ func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string)
|
||||
|
||||
// requestID generates a unique ID for ephemeral file keys.
|
||||
func requestID() string {
|
||||
return uuid.New().String()[:8]
|
||||
return uuid.NewString()
|
||||
}
|
||||
|
||||
type stagedInputLifecycle struct {
|
||||
client *FileStagingClient
|
||||
requestID string
|
||||
keys []string
|
||||
seen map[string]struct{}
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) newStagedInputLifecycle() *stagedInputLifecycle {
|
||||
return &stagedInputLifecycle{
|
||||
client: f,
|
||||
requestID: requestID(),
|
||||
keys: []string{},
|
||||
seen: map[string]struct{}{},
|
||||
}
|
||||
}
|
||||
|
||||
func (l *stagedInputLifecycle) track(key string) {
|
||||
if _, ok := l.seen[key]; ok {
|
||||
return
|
||||
}
|
||||
l.seen[key] = struct{}{}
|
||||
l.keys = append(l.keys, key)
|
||||
}
|
||||
|
||||
func (l *stagedInputLifecycle) release() {
|
||||
if len(l.keys) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), stagedInputReleaseTimeout)
|
||||
defer cancel()
|
||||
if releaser, ok := l.client.stager.(RequestFileReleaser); ok {
|
||||
if err := releaser.ReleaseRemoteRequest(ctx, l.client.nodeID, l.requestID, l.keys); err != nil {
|
||||
xlog.Warn("Failed to release staged request inputs", "node", l.client.nodeID, "requestID", l.requestID, "keyCount", len(l.keys), "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
for _, key := range l.keys {
|
||||
if err := l.client.stager.ReleaseRemote(ctx, l.client.nodeID, key); err != nil {
|
||||
xlog.Warn("Failed to release staged input", "node", l.client.nodeID, "key", key, "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stageInputFile uploads a local file to the remote node via the FileStager.
|
||||
// Returns the remote-local path and the ephemeral key.
|
||||
func (f *FileStagingClient) stageInputFile(ctx context.Context, reqID, localPath, category string) (string, string, error) {
|
||||
func (f *FileStagingClient) stageInputFile(
|
||||
ctx context.Context,
|
||||
lifecycle *stagedInputLifecycle,
|
||||
localPath,
|
||||
category string,
|
||||
) (string, error) {
|
||||
basename := filepath.Base(localPath)
|
||||
key := storage.EphemeralKey(reqID, category, basename)
|
||||
key := storage.EphemeralKey(lifecycle.requestID, category, basename)
|
||||
lifecycle.track(key)
|
||||
|
||||
remotePath, err := f.stager.EnsureRemote(ctx, f.nodeID, localPath, key)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("staging input file: %w", err)
|
||||
return "", fmt.Errorf("staging input file: %w", err)
|
||||
}
|
||||
|
||||
return remotePath, key, nil
|
||||
return remotePath, nil
|
||||
}
|
||||
|
||||
// retrieveOutputFile retrieves an output file from the backend to a local path.
|
||||
@@ -100,23 +152,37 @@ func (f *FileStagingClient) translateModelPath(frontendPath string) string {
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) Predict(ctx context.Context, in *pb.PredictOptions, opts ...ggrpc.CallOption) (*pb.Reply, error) {
|
||||
reqID := requestID()
|
||||
in, _ = f.stageMultimodalInputs(ctx, reqID, in)
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.PredictOptions)
|
||||
var err error
|
||||
in, err = f.stageMultimodalInputs(ctx, lifecycle, in)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return f.Backend.Predict(ctx, in, opts...)
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) PredictStream(ctx context.Context, in *pb.PredictOptions, fn func(reply *pb.Reply), opts ...ggrpc.CallOption) error {
|
||||
reqID := requestID()
|
||||
in, _ = f.stageMultimodalInputs(ctx, reqID, in)
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.PredictOptions)
|
||||
var err error
|
||||
in, err = f.stageMultimodalInputs(ctx, lifecycle, in)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return f.Backend.PredictStream(ctx, in, fn, opts...)
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.GenerateImageRequest)
|
||||
|
||||
// Stage input source image if present
|
||||
if in.Src != "" && isFilePath(in.Src) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging image src: %w", err)
|
||||
}
|
||||
@@ -126,7 +192,7 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm
|
||||
// Stage reference images
|
||||
for i, img := range in.RefImages {
|
||||
if isFilePath(img) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, img, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, img, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging ref image: %w", err)
|
||||
}
|
||||
@@ -160,25 +226,27 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.GenerateVideoRequest)
|
||||
|
||||
// Stage start/end images and optional audio conditioning.
|
||||
if in.StartImage != "" && isFilePath(in.StartImage) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.StartImage, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.StartImage, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging start image: %w", err)
|
||||
}
|
||||
in.StartImage = backendPath
|
||||
}
|
||||
if in.EndImage != "" && isFilePath(in.EndImage) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.EndImage, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.EndImage, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging end image: %w", err)
|
||||
}
|
||||
in.EndImage = backendPath
|
||||
}
|
||||
if in.Audio != "" && isFilePath(in.Audio) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Audio, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Audio, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging video audio: %w", err)
|
||||
}
|
||||
@@ -210,11 +278,13 @@ func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVi
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.Generate3DRequest)
|
||||
|
||||
// Stage the conditioning image or existing GLB used by 3D post-processing.
|
||||
if in.Src != "" && isFilePath(in.Src) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging 3D input asset: %w", err)
|
||||
}
|
||||
@@ -246,7 +316,9 @@ func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DReq
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.TTSRequest)
|
||||
|
||||
// Translate model path from frontend to remote worker path.
|
||||
// The model and its companion files (e.g. .onnx.json) were already staged
|
||||
@@ -257,7 +329,7 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...
|
||||
// Voice may be a named backend speaker or a request-scoped reference WAV.
|
||||
// Only path-shaped values are staged; speaker IDs pass through unchanged.
|
||||
if in.Voice != "" && isFilePath(in.Voice) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging TTS voice reference: %w", err)
|
||||
}
|
||||
@@ -289,14 +361,16 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn func(*pb.Reply), opts ...ggrpc.CallOption) error {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.TTSRequest)
|
||||
|
||||
// Translate model path from frontend to remote worker path (same as TTS above)
|
||||
if in.Model != "" && isFilePath(in.Model) {
|
||||
in.Model = f.translateModelPath(in.Model)
|
||||
}
|
||||
if in.Voice != "" && isFilePath(in.Voice) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs")
|
||||
if err != nil {
|
||||
return fmt.Errorf("staging streaming TTS voice reference: %w", err)
|
||||
}
|
||||
@@ -307,11 +381,13 @@ func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest, opts ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.SoundGenerationRequest)
|
||||
|
||||
// Stage input source
|
||||
if in.Src != nil && *in.Src != "" && isFilePath(*in.Src) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, *in.Src, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, *in.Src, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging sound src: %w", err)
|
||||
}
|
||||
@@ -342,12 +418,28 @@ func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGen
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) SoundDetection(ctx context.Context, in *pb.SoundDetectionRequest, opts ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.SoundDetectionRequest)
|
||||
if in.Src != "" && isFilePath(in.Src) {
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging audio for sound detection: %w", err)
|
||||
}
|
||||
in.Src = backendPath
|
||||
}
|
||||
return f.Backend.SoundDetection(ctx, in, opts...)
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.TranscriptRequest, opts ...ggrpc.CallOption) (*pb.TranscriptResult, error) {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.TranscriptRequest)
|
||||
|
||||
// Stage input audio file
|
||||
if in.Dst != "" && isFilePath(in.Dst) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging audio for transcription: %w", err)
|
||||
}
|
||||
@@ -358,11 +450,13 @@ func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.Trans
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) AudioTranscriptionStream(ctx context.Context, in *pb.TranscriptRequest, fn func(chunk *pb.TranscriptStreamResponse), opts ...ggrpc.CallOption) error {
|
||||
reqID := requestID()
|
||||
lifecycle := f.newStagedInputLifecycle()
|
||||
defer lifecycle.release()
|
||||
in = proto.Clone(in).(*pb.TranscriptRequest)
|
||||
|
||||
// Stage input audio file
|
||||
if in.Dst != "" && isFilePath(in.Dst) {
|
||||
backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs")
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs")
|
||||
if err != nil {
|
||||
return fmt.Errorf("staging audio for transcription stream: %w", err)
|
||||
}
|
||||
@@ -444,31 +538,46 @@ func (f *FileStagingClient) QuantizationProgress(ctx context.Context, in *pb.Qua
|
||||
|
||||
// stageMultimodalInputs stages Images, Videos, Audios fields in PredictOptions
|
||||
// if they are file paths (not base64 or URLs).
|
||||
func (f *FileStagingClient) stageMultimodalInputs(ctx context.Context, reqID string, in *pb.PredictOptions) (*pb.PredictOptions, []string) {
|
||||
var keys []string
|
||||
in.Images = f.stagePathSlice(ctx, reqID, in.Images, "inputs", &keys)
|
||||
in.Videos = f.stagePathSlice(ctx, reqID, in.Videos, "inputs", &keys)
|
||||
in.Audios = f.stagePathSlice(ctx, reqID, in.Audios, "inputs", &keys)
|
||||
return in, keys
|
||||
func (f *FileStagingClient) stageMultimodalInputs(
|
||||
ctx context.Context,
|
||||
lifecycle *stagedInputLifecycle,
|
||||
in *pb.PredictOptions,
|
||||
) (*pb.PredictOptions, error) {
|
||||
var err error
|
||||
in.Images, err = f.stagePathSlice(ctx, lifecycle, in.Images, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging predict images: %w", err)
|
||||
}
|
||||
in.Videos, err = f.stagePathSlice(ctx, lifecycle, in.Videos, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging predict videos: %w", err)
|
||||
}
|
||||
in.Audios, err = f.stagePathSlice(ctx, lifecycle, in.Audios, "inputs")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("staging predict audios: %w", err)
|
||||
}
|
||||
return in, nil
|
||||
}
|
||||
|
||||
func (f *FileStagingClient) stagePathSlice(ctx context.Context, reqID string, paths []string, category string, keys *[]string) []string {
|
||||
func (f *FileStagingClient) stagePathSlice(
|
||||
ctx context.Context,
|
||||
lifecycle *stagedInputLifecycle,
|
||||
paths []string,
|
||||
category string,
|
||||
) ([]string, error) {
|
||||
result := make([]string, len(paths))
|
||||
for i, p := range paths {
|
||||
if isFilePath(p) {
|
||||
backendPath, key, err := f.stageInputFile(ctx, reqID, p, category)
|
||||
backendPath, err := f.stageInputFile(ctx, lifecycle, p, category)
|
||||
if err != nil {
|
||||
xlog.Warn("Failed to stage multimodal file, passing through", "path", p, "error", err)
|
||||
result[i] = p
|
||||
continue
|
||||
return nil, fmt.Errorf("staging %q: %w", p, err)
|
||||
}
|
||||
result[i] = backendPath
|
||||
*keys = append(*keys, key)
|
||||
} else {
|
||||
result[i] = p
|
||||
}
|
||||
}
|
||||
return result
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// isFilePath checks if a string looks like a local file path (not base64 or URL).
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
ggrpc "google.golang.org/grpc"
|
||||
"google.golang.org/protobuf/proto"
|
||||
)
|
||||
|
||||
const fullUUIDPattern = `[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`
|
||||
|
||||
type lifecycleStager struct {
|
||||
fakeFileStager
|
||||
ensureErr error
|
||||
ensureErrAt int
|
||||
releaseErr error
|
||||
releasedKeys []string
|
||||
releaseBatches [][]string
|
||||
releaseCtxErr []error
|
||||
releaseHasDeadline []bool
|
||||
releaseDeadlines []time.Time
|
||||
}
|
||||
|
||||
func (s *lifecycleStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
|
||||
s.fakeFileStager.EnsureRemote(ctx, nodeID, localPath, key)
|
||||
if s.ensureErr != nil && (s.ensureErrAt == 0 || len(s.ensureCalls) == s.ensureErrAt) {
|
||||
return "", s.ensureErr
|
||||
}
|
||||
return "/remote/" + key, nil
|
||||
}
|
||||
|
||||
func (s *lifecycleStager) ReleaseRemote(ctx context.Context, _ string, key string) error {
|
||||
s.releasedKeys = append(s.releasedKeys, key)
|
||||
s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err())
|
||||
deadline, ok := ctx.Deadline()
|
||||
s.releaseHasDeadline = append(s.releaseHasDeadline, ok)
|
||||
s.releaseDeadlines = append(s.releaseDeadlines, deadline)
|
||||
return s.releaseErr
|
||||
}
|
||||
|
||||
func (s *lifecycleStager) ReleaseRemoteRequest(ctx context.Context, _, _ string, keys []string) error {
|
||||
s.releaseBatches = append(s.releaseBatches, append([]string(nil), keys...))
|
||||
s.releasedKeys = append(s.releasedKeys, keys...)
|
||||
s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err())
|
||||
deadline, ok := ctx.Deadline()
|
||||
s.releaseHasDeadline = append(s.releaseHasDeadline, ok)
|
||||
s.releaseDeadlines = append(s.releaseDeadlines, deadline)
|
||||
return s.releaseErr
|
||||
}
|
||||
|
||||
type lifecycleBackend struct {
|
||||
grpc.Backend
|
||||
predictResult *pb.Reply
|
||||
predictErr error
|
||||
predictCalls int
|
||||
predictInput *pb.PredictOptions
|
||||
streamCalls int
|
||||
streamBlock <-chan struct{}
|
||||
streamStarted chan<- struct{}
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) Predict(_ context.Context, in *pb.PredictOptions, _ ...ggrpc.CallOption) (*pb.Reply, error) {
|
||||
b.predictCalls++
|
||||
b.predictInput = proto.Clone(in).(*pb.PredictOptions)
|
||||
if b.predictResult == nil {
|
||||
b.predictResult = &pb.Reply{}
|
||||
}
|
||||
return b.predictResult, b.predictErr
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) PredictStream(_ context.Context, _ *pb.PredictOptions, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
|
||||
b.streamCalls++
|
||||
if b.streamStarted != nil {
|
||||
b.streamStarted <- struct{}{}
|
||||
}
|
||||
if b.streamBlock != nil {
|
||||
<-b.streamBlock
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) GenerateImage(_ context.Context, _ *pb.GenerateImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) Generate3D(_ context.Context, _ *pb.Generate3DRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) TTS(_ context.Context, _ *pb.TTSRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) TTSStream(_ context.Context, _ *pb.TTSRequest, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) SoundGeneration(_ context.Context, _ *pb.SoundGenerationRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
|
||||
return &pb.Result{Success: true}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) SoundDetection(_ context.Context, _ *pb.SoundDetectionRequest, _ ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
|
||||
return &pb.SoundDetectionResponse{}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) AudioTranscription(_ context.Context, _ *pb.TranscriptRequest, _ ...ggrpc.CallOption) (*pb.TranscriptResult, error) {
|
||||
return &pb.TranscriptResult{}, nil
|
||||
}
|
||||
|
||||
func (b *lifecycleBackend) AudioTranscriptionStream(_ context.Context, _ *pb.TranscriptRequest, _ func(*pb.TranscriptStreamResponse), _ ...ggrpc.CallOption) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
var _ = Describe("FileStagingClient request lifecycle", func() {
|
||||
It("uses a full UUID for ephemeral request keys", func() {
|
||||
Expect(requestID()).To(MatchRegexp(`^` + fullUUIDPattern + `$`))
|
||||
})
|
||||
|
||||
It("releases every staged key and preserves caller requests", func(ctx SpecContext) {
|
||||
tests := []struct {
|
||||
name string
|
||||
keyCount int
|
||||
invoke func(*FileStagingClient) proto.Message
|
||||
}{
|
||||
{name: "predict", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}, Videos: []string{"/tmp/video.mp4"}, Audios: []string{"/tmp/audio.wav"}}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.Predict(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "predict stream", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}}
|
||||
original := proto.Clone(request)
|
||||
Expect(client.PredictStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "image generation", keyCount: 2, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.GenerateImageRequest{Src: "/tmp/source.png", RefImages: []string{"/tmp/reference.png"}}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.GenerateImage(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "video generation", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.GenerateVideoRequest{StartImage: "/tmp/start.png", EndImage: "/tmp/end.png", Audio: "/tmp/audio.wav"}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.GenerateVideo(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "3D generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.Generate3DRequest{Src: "/tmp/source.glb"}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.Generate3D(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.TTS(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "streaming TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
|
||||
original := proto.Clone(request)
|
||||
Expect(client.TTSStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "sound generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
source := "/tmp/source.wav"
|
||||
request := &pb.SoundGenerationRequest{Src: &source}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.SoundGeneration(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "sound detection", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.SoundDetectionRequest{Src: "/tmp/source.wav"}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.SoundDetection(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
|
||||
original := proto.Clone(request)
|
||||
_, err := client.AudioTranscription(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
{name: "streaming transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
|
||||
request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
|
||||
original := proto.Clone(request)
|
||||
Expect(client.AudioTranscriptionStream(ctx, request, func(*pb.TranscriptStreamResponse) {})).To(Succeed())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
return request
|
||||
}},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
By(test.name)
|
||||
stager := &lifecycleStager{}
|
||||
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
|
||||
test.invoke(client)
|
||||
Expect(stager.ensureCalls).To(HaveLen(test.keyCount))
|
||||
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
|
||||
Expect(stager.releaseBatches).To(Equal([][]string{keysFromEnsureCalls(stager.ensureCalls)}))
|
||||
Expect(stager.releaseCtxErr).To(Equal([]error{nil}))
|
||||
Expect(stager.releaseHasDeadline).To(HaveLen(1))
|
||||
for _, hasDeadline := range stager.releaseHasDeadline {
|
||||
Expect(hasDeadline).To(BeTrue())
|
||||
}
|
||||
for _, deadline := range stager.releaseDeadlines {
|
||||
Expect(time.Until(deadline)).To(BeNumerically(">", 0))
|
||||
Expect(time.Until(deadline)).To(BeNumerically("<=", time.Minute))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
It("tracks a key before staging so a partial upload failure is released", func(ctx SpecContext) {
|
||||
uploadErr := errors.New("upload failed")
|
||||
stager := &lifecycleStager{ensureErr: uploadErr, releaseErr: errors.New("release failed")}
|
||||
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
|
||||
|
||||
_, err := client.GenerateImage(ctx, &pb.GenerateImageRequest{Src: "/tmp/source.png"})
|
||||
|
||||
Expect(err).To(MatchError(ContainSubstring("upload failed")))
|
||||
Expect(stager.ensureCalls).To(HaveLen(1))
|
||||
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
|
||||
})
|
||||
|
||||
It("does not invoke predict when multimodal staging fails", func(ctx SpecContext) {
|
||||
uploadErr := errors.New("ephemeral capacity exceeded")
|
||||
stager := &lifecycleStager{ensureErr: uploadErr, ensureErrAt: 2}
|
||||
backend := &lifecycleBackend{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
request := &pb.PredictOptions{Images: []string{"/tmp/first.png", "/tmp/second.png"}}
|
||||
original := proto.Clone(request)
|
||||
|
||||
result, err := client.Predict(ctx, request)
|
||||
|
||||
Expect(result).To(BeNil())
|
||||
Expect(err).To(MatchError(ContainSubstring("ephemeral capacity exceeded")))
|
||||
Expect(backend.predictCalls).To(BeZero())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
Expect(stager.ensureCalls).To(HaveLen(2))
|
||||
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
|
||||
})
|
||||
|
||||
It("does not invoke streaming predict when multimodal staging fails", func(ctx SpecContext) {
|
||||
uploadErr := errors.New("ephemeral capacity exceeded")
|
||||
stager := &lifecycleStager{ensureErr: uploadErr}
|
||||
backend := &lifecycleBackend{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
request := &pb.PredictOptions{Audios: []string{"/tmp/audio.wav"}}
|
||||
original := proto.Clone(request)
|
||||
|
||||
err := client.PredictStream(ctx, request, func(*pb.Reply) {})
|
||||
|
||||
Expect(err).To(MatchError(ContainSubstring("ephemeral capacity exceeded")))
|
||||
Expect(backend.streamCalls).To(BeZero())
|
||||
Expect(proto.Equal(request, original)).To(BeTrue())
|
||||
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
|
||||
})
|
||||
|
||||
It("uses an active bounded cleanup context after caller cancellation", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
stager := &lifecycleStager{}
|
||||
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
|
||||
|
||||
_, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(stager.releaseCtxErr).To(Equal([]error{nil}))
|
||||
})
|
||||
|
||||
It("does not release streaming inputs before the backend completes", func(ctx SpecContext) {
|
||||
block := make(chan struct{})
|
||||
started := make(chan struct{}, 1)
|
||||
backend := &lifecycleBackend{streamBlock: block, streamStarted: started}
|
||||
stager := &lifecycleStager{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
done := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
done <- client.PredictStream(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}, func(*pb.Reply) {})
|
||||
}()
|
||||
|
||||
Eventually(started).Should(Receive())
|
||||
Expect(stager.ensureCalls).To(HaveLen(1))
|
||||
Expect(stager.releasedKeys).To(BeEmpty())
|
||||
close(block)
|
||||
Eventually(done).Should(Receive(Succeed()))
|
||||
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
|
||||
})
|
||||
|
||||
It("does not replace a backend result when cleanup fails", func(ctx SpecContext) {
|
||||
reply := &pb.Reply{Message: []byte("ok")}
|
||||
backend := &lifecycleBackend{predictResult: reply}
|
||||
stager := &lifecycleStager{releaseErr: errors.New("release failed")}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
|
||||
result, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result).To(BeIdenticalTo(reply))
|
||||
})
|
||||
})
|
||||
|
||||
func keysFromEnsureCalls(calls []ensureCall) []string {
|
||||
keys := make([]string, len(calls))
|
||||
for i, call := range calls {
|
||||
keys[i] = call.key
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
ggrpc "google.golang.org/grpc"
|
||||
)
|
||||
|
||||
type soundStagingBackend struct {
|
||||
grpc.Backend
|
||||
request *pb.SoundDetectionRequest
|
||||
err error
|
||||
}
|
||||
|
||||
func (b *soundStagingBackend) SoundDetection(_ context.Context, in *pb.SoundDetectionRequest, _ ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
|
||||
b.request = in
|
||||
return &pb.SoundDetectionResponse{}, b.err
|
||||
}
|
||||
|
||||
type soundStagingFailure struct{ FileStager }
|
||||
|
||||
func (s *soundStagingFailure) EnsureRemote(context.Context, string, string, string) (string, error) {
|
||||
return "", errors.New("upload failed")
|
||||
}
|
||||
|
||||
func (s *soundStagingFailure) ReleaseRemote(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type soundRouteFactory struct{ client grpc.Backend }
|
||||
|
||||
func (f *soundRouteFactory) NewClient(string, bool) grpc.Backend { return f.client }
|
||||
|
||||
var _ = Describe("FileStagingClient sound detection", func() {
|
||||
It("stages sound audio through the client returned by SmartRouter.Route", func(ctx SpecContext) {
|
||||
node := &BackendNode{ID: "worker-1", Name: "worker", Address: "10.0.0.1:50051"}
|
||||
reg := &fakeModelRouter{
|
||||
findAndLockNode: node,
|
||||
findAndLockNM: &NodeModel{NodeID: node.ID, ModelName: "ced", Address: "10.0.0.1:9001"},
|
||||
}
|
||||
backend := &soundStagingBackend{Backend: &stubBackend{healthResult: true}}
|
||||
stager := &fakeFileStager{}
|
||||
router := NewSmartRouter(reg, SmartRouterOptions{
|
||||
ClientFactory: &soundRouteFactory{client: backend},
|
||||
FileStager: stager,
|
||||
Unloader: &fakeUnloader{},
|
||||
})
|
||||
result, err := router.Route(ctx, "ced", "ced.gguf", "ced", "", nil, false)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result).NotTo(BeNil())
|
||||
defer result.Release()
|
||||
request := &pb.SoundDetectionRequest{Src: "/tmp/realtime-sound-window-test.wav", ModelIdentity: "ced.gguf"}
|
||||
_, err = result.Client.SoundDetection(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(stager.ensureCalls).To(HaveLen(1))
|
||||
Expect(stager.ensureCalls[0].localPath).To(Equal(request.Src))
|
||||
Expect(backend.request.Src).To(Equal("/remote/" + stager.ensureCalls[0].key))
|
||||
Expect(request.Src).To(Equal("/tmp/realtime-sound-window-test.wav"))
|
||||
})
|
||||
|
||||
It("stages audio on the worker without changing the caller's request", func(ctx SpecContext) {
|
||||
backend := &soundStagingBackend{}
|
||||
stager := &fakeFileStager{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
request := &pb.SoundDetectionRequest{Src: "/tmp/realtime-sound-window-test.wav", ModelIdentity: "ced.gguf", TopK: 5, Threshold: 0.25}
|
||||
_, err := client.SoundDetection(ctx, request)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(stager.ensureCalls).To(HaveLen(1))
|
||||
Expect(stager.ensureCalls[0].nodeID).To(Equal("worker-1"))
|
||||
Expect(stager.ensureCalls[0].localPath).To(Equal(request.Src))
|
||||
Expect(backend.request.Src).To(Equal("/remote/" + stager.ensureCalls[0].key))
|
||||
Expect(backend.request.ModelIdentity).To(Equal(request.ModelIdentity))
|
||||
Expect(backend.request.TopK).To(Equal(request.TopK))
|
||||
Expect(backend.request.Threshold).To(Equal(request.Threshold))
|
||||
Expect(request.Src).To(Equal("/tmp/realtime-sound-window-test.wav"))
|
||||
})
|
||||
It("does not call the backend when staging fails", func(ctx SpecContext) {
|
||||
backend := &soundStagingBackend{}
|
||||
client := NewFileStagingClient(backend, &soundStagingFailure{}, "worker-1")
|
||||
_, err := client.SoundDetection(ctx, &pb.SoundDetectionRequest{Src: "/tmp/clip.wav"})
|
||||
Expect(err).To(MatchError(ContainSubstring("upload failed")))
|
||||
Expect(backend.request).To(BeNil())
|
||||
})
|
||||
It("passes through requests without a file and preserves backend errors", func(ctx SpecContext) {
|
||||
failure := errors.New("classifier failed")
|
||||
backend := &soundStagingBackend{err: failure}
|
||||
stager := &fakeFileStager{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
_, err := client.SoundDetection(ctx, &pb.SoundDetectionRequest{})
|
||||
Expect(err).To(MatchError(failure))
|
||||
Expect(stager.ensureCalls).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -39,7 +39,8 @@ var _ = Describe("FileStagingClient TTS references", func() {
|
||||
Expect(stager.ensureCalls).To(HaveLen(1))
|
||||
Expect(stager.ensureCalls[0].localPath).To(Equal("/data/voice-profiles/profile/reference.wav"))
|
||||
Expect(backend.ttsRequest.Voice).To(HavePrefix("/remote/ephemeral/"))
|
||||
Expect(backend.ttsRequest.Voice).To(MatchRegexp(`/inputs/[0-9a-f]{8}/reference\.wav$`))
|
||||
voicePathPattern := `/inputs/` + fullUUIDPattern + `/reference\.wav$`
|
||||
Expect(backend.ttsRequest.Voice).To(MatchRegexp(voicePathPattern))
|
||||
})
|
||||
|
||||
It("stages a reference WAV before streaming synthesis", func(ctx SpecContext) {
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Recovering unfinished file finalization", func() {
|
||||
It("stages a full-size unfinished upload and subsequently skips the committed file", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
content := []byte("complete model bytes")
|
||||
hash := fmt.Sprintf("%x", sha256.Sum256(content))
|
||||
remote := filepath.Join(dir, "model.bin")
|
||||
Expect(os.WriteFile(remote, content, 0600)).To(Succeed())
|
||||
Expect(os.WriteFile(remote+targetSidecarSuffix, []byte(hash), 0600)).To(Succeed())
|
||||
local := filepath.Join(GinkgoT().TempDir(), "model.bin")
|
||||
Expect(os.WriteFile(local, content, 0600)).To(Succeed())
|
||||
|
||||
puts := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.Method {
|
||||
case http.MethodHead:
|
||||
handleHead(w, r, dir, "", "", "model.bin")
|
||||
case http.MethodPut:
|
||||
puts++
|
||||
handleUpload(w, r, dir, "", "", "model.bin", 0)
|
||||
}
|
||||
}))
|
||||
DeferCleanup(server.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(server.URL, "http://"), nil
|
||||
}, "")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
result, err := stager.EnsureRemote(ctx, "worker", local, "model.bin")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result).To(Equal(remote))
|
||||
Expect(remote + targetSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(os.ReadFile(remote + hashSidecarSuffix)).To(Equal([]byte(hash)))
|
||||
Expect(os.ReadFile(remote)).To(Equal(content))
|
||||
result, err = stager.EnsureRemote(ctx, "worker", local, "model.bin")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(result).To(Equal(remote))
|
||||
Expect(puts).To(Equal(1))
|
||||
})
|
||||
|
||||
DescribeTable("validates existing bytes before acknowledging a retry",
|
||||
func(content string, expectedStatus int) {
|
||||
dir := GinkgoT().TempDir()
|
||||
remote := filepath.Join(dir, "model.bin")
|
||||
hash := fmt.Sprintf("%x", sha256.Sum256([]byte("good")))
|
||||
Expect(os.WriteFile(remote, []byte(content), 0600)).To(Succeed())
|
||||
Expect(os.WriteFile(remote+targetSidecarSuffix, []byte(hash), 0600)).To(Succeed())
|
||||
// A cached hash must not conceal corrupt bytes left by an interrupted upload.
|
||||
Expect(os.WriteFile(remote+hashSidecarSuffix, []byte(hash), 0600)).To(Succeed())
|
||||
req := httptest.NewRequest(http.MethodPut, "/v1/files/model.bin", strings.NewReader("good"))
|
||||
req.Header.Set("Content-Range", "bytes 0-3/4")
|
||||
req.Header.Set(HeaderContentSHA256, hash)
|
||||
response := httptest.NewRecorder()
|
||||
handleUpload(response, req, dir, "", "", "model.bin", 0)
|
||||
Expect(response.Code).To(Equal(expectedStatus))
|
||||
if expectedStatus == http.StatusBadRequest {
|
||||
Expect(response.Body.String()).To(ContainSubstring("sha256 mismatch"))
|
||||
Expect(remote).NotTo(BeAnExistingFile())
|
||||
Expect(remote + targetSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
Expect(remote + hashSidecarSuffix).NotTo(BeAnExistingFile())
|
||||
} else {
|
||||
Expect(os.ReadFile(remote)).To(Equal([]byte(content)))
|
||||
Expect(remote + targetSidecarSuffix).To(BeAnExistingFile())
|
||||
}
|
||||
},
|
||||
Entry("rejects corrupt full-size content", "evil", http.StatusBadRequest),
|
||||
Entry("preserves partial content for resume", "go", http.StatusRequestedRangeNotSatisfiable),
|
||||
Entry("does not finalize oversized content", "good-extra", http.StatusRequestedRangeNotSatisfiable),
|
||||
)
|
||||
})
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
@@ -22,6 +23,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/pkg/downloader"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/safefile"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
@@ -32,8 +34,9 @@ const (
|
||||
HeaderFileSize = "X-File-Size"
|
||||
// HeaderTargetSHA256 is set on HEAD responses for partial (resumable) uploads
|
||||
// to expose the expected final SHA-256 of the in-progress file. When set,
|
||||
// the file on disk is not yet the full content — the client may resume by
|
||||
// PUT'ing the remainder with a matching X-Content-SHA256 header.
|
||||
// the file on disk is not yet verified — the client may resume by
|
||||
// PUT'ing the remainder with a matching X-Content-SHA256 header. A full-size
|
||||
// file can still carry this marker if finalization was interrupted.
|
||||
HeaderTargetSHA256 = "X-Target-SHA256"
|
||||
hashSidecarSuffix = ".sha256"
|
||||
// targetSidecarSuffix stores the expected final SHA-256 of a partially
|
||||
@@ -48,17 +51,23 @@ const (
|
||||
// Auth is via Bearer token (registration token), using constant-time comparison.
|
||||
// A nil readiness fails open, keeping /readyz's historical always-200 answer.
|
||||
func StartFileTransferServer(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, logStore ...*model.BackendLogStore) (*http.Server, error) {
|
||||
return StartFileTransferServerWithCapacity(addr, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, nil, logStore...)
|
||||
}
|
||||
|
||||
// StartFileTransferServerWithCapacity starts the file transfer server with a
|
||||
// worker-local guard for per-request ephemeral inputs.
|
||||
func StartFileTransferServerWithCapacity(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, logStore ...*model.BackendLogStore) (*http.Server, error) {
|
||||
listener, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("listen %s: %w", addr, err)
|
||||
}
|
||||
return StartFileTransferServerWithReadiness(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, logStore...)
|
||||
return startFileTransferServer(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, logStore...)
|
||||
}
|
||||
|
||||
// StartFileTransferServerWithListener starts the server on an existing listener.
|
||||
// This avoids the TOCTOU race of closing a listener and re-binding to the same port.
|
||||
func StartFileTransferServerWithListener(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, logStore ...*model.BackendLogStore) (*http.Server, error) {
|
||||
return StartFileTransferServerWithReadiness(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, nil, logStore...)
|
||||
return startFileTransferServer(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, nil, nil, logStore...)
|
||||
}
|
||||
|
||||
// StartFileTransferServerWithReadiness is StartFileTransferServerWithListener
|
||||
@@ -66,6 +75,10 @@ func StartFileTransferServerWithListener(lis net.Listener, stagingDir, modelsDir
|
||||
// the probe keeps its historical always-200 behaviour for callers that have no
|
||||
// meaningful readiness signal to report.
|
||||
func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, logStore ...*model.BackendLogStore) (*http.Server, error) {
|
||||
return startFileTransferServer(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, nil, logStore...)
|
||||
}
|
||||
|
||||
func startFileTransferServer(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, logStore ...*model.BackendLogStore) (*http.Server, error) {
|
||||
if err := os.MkdirAll(stagingDir, 0750); err != nil {
|
||||
return nil, fmt.Errorf("creating staging dir %s: %w", stagingDir, err)
|
||||
}
|
||||
@@ -100,6 +113,18 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
|
||||
handleListDir(w, r, stagingDir, modelsDir, dataDir, key)
|
||||
})
|
||||
|
||||
mux.HandleFunc("/v1/files-release", func(w http.ResponseWriter, r *http.Request) {
|
||||
if !checkBearerToken(r, token) {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodDelete {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
handleReleaseBatchWithCapacity(w, r, stagingDir, capacity)
|
||||
})
|
||||
|
||||
mux.HandleFunc("/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if !checkBearerToken(r, token) {
|
||||
xlog.Debug("HTTP file transfer: unauthorized request", "method", r.Method, "path", r.URL.Path, "remote", r.RemoteAddr)
|
||||
@@ -115,12 +140,16 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
|
||||
case http.MethodHead:
|
||||
handleHead(w, r, stagingDir, modelsDir, dataDir, key)
|
||||
case http.MethodPut:
|
||||
handleUpload(w, r, stagingDir, modelsDir, dataDir, key, maxUploadSize)
|
||||
handleUploadWithCapacity(w, r, stagingDir, modelsDir, dataDir, key, maxUploadSize, capacity)
|
||||
case http.MethodGet:
|
||||
handleDownload(w, r, stagingDir, modelsDir, dataDir, key)
|
||||
case http.MethodDelete:
|
||||
handleReleaseWithCapacity(w, r, stagingDir, key, capacity)
|
||||
case http.MethodPost:
|
||||
if key == "temp" {
|
||||
handleAllocTemp(w, r, stagingDir)
|
||||
} else if r.URL.Query().Get("claim") == "1" {
|
||||
handleClaimWithCapacity(w, r, stagingDir, key, capacity)
|
||||
} else {
|
||||
http.Error(w, "not found", http.StatusNotFound)
|
||||
}
|
||||
@@ -181,6 +210,163 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
|
||||
return server, nil
|
||||
}
|
||||
|
||||
// handleClaimWithCapacity marks an existing ephemeral file as owned by the
|
||||
// request that just verified its content.
|
||||
func handleClaimWithCapacity(w http.ResponseWriter, _ *http.Request, stagingDir, key string, capacity EphemeralCapacity) {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if operations, ok := capacity.(ephemeralRequestOperationCapacity); ok {
|
||||
requestID := strings.Split(key, "/")[2]
|
||||
if err := operations.BeginRequestOperation(requestID); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusConflict)
|
||||
return
|
||||
}
|
||||
defer operations.EndRequestOperation(requestID)
|
||||
}
|
||||
filePath := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
if err := validatePathInDir(filePath, stagingDir); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
info, err := os.Lstat(filePath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
http.Error(w, "not found", http.StatusNotFound)
|
||||
} else {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
}
|
||||
return
|
||||
}
|
||||
if !info.Mode().IsRegular() {
|
||||
http.Error(w, "ephemeral path is not a regular file", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if capacity != nil {
|
||||
if err := capacity.Claim(filePath); err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
http.Error(w, "not found", http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
http.Error(w, err.Error(), http.StatusInsufficientStorage)
|
||||
return
|
||||
}
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func handleRelease(w http.ResponseWriter, _ *http.Request, stagingDir, key string) {
|
||||
handleReleaseWithCapacity(w, nil, stagingDir, key, nil)
|
||||
}
|
||||
|
||||
func handleReleaseWithCapacity(w http.ResponseWriter, _ *http.Request, stagingDir, key string, capacity EphemeralCapacity) {
|
||||
if err := releaseEphemeralStagingKey(stagingDir, key, capacity); err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
if errors.Is(err, safefile.ErrUnsafePath) || validateEphemeralReleaseKey(key) != nil {
|
||||
status = http.StatusBadRequest
|
||||
}
|
||||
http.Error(w, err.Error(), status)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func handleReleaseBatchWithCapacity(w http.ResponseWriter, r *http.Request, stagingDir string, capacity EphemeralCapacity) {
|
||||
r.Body = http.MaxBytesReader(w, r.Body, 1<<20)
|
||||
var request struct {
|
||||
RequestID string `json:"request_id"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
||||
http.Error(w, fmt.Sprintf("decoding release batch: %v", err), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if err := validateEphemeralRequestID(request.RequestID); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if err := releaseEphemeralStagingRequest(r.Context(), stagingDir, request.RequestID, capacity); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
func releaseEphemeralStagingRequest(ctx context.Context, stagingDir, requestID string, capacity EphemeralCapacity) error {
|
||||
if err := validateEphemeralRequestID(requestID); err != nil {
|
||||
return err
|
||||
}
|
||||
if requestCapacity, ok := capacity.(ephemeralRequestCapacity); ok {
|
||||
if err := requestCapacity.BeginRequestRelease(ctx, requestID); err != nil {
|
||||
return fmt.Errorf("beginning release for request %q: %w", requestID, err)
|
||||
}
|
||||
defer requestCapacity.EndRequestRelease(requestID)
|
||||
}
|
||||
root := filepath.Join(stagingDir, "ephemeral")
|
||||
categories, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
var releaseErrors []error
|
||||
for _, category := range categories {
|
||||
if !category.IsDir() || category.Type()&os.ModeSymlink != 0 {
|
||||
continue
|
||||
}
|
||||
requestDir := filepath.Join(root, category.Name(), requestID)
|
||||
info, err := os.Lstat(requestDir)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("stating request directory %q: %w", requestDir, err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("ephemeral request path %q is not a real directory", requestDir))
|
||||
continue
|
||||
}
|
||||
entries, err := os.ReadDir(requestDir)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("reading request directory %q: %w", requestDir, err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("unexpected directory in ephemeral request %q", filepath.Join(requestDir, entry.Name())))
|
||||
continue
|
||||
}
|
||||
key := filepath.ToSlash(filepath.Join("ephemeral", category.Name(), requestID, entry.Name()))
|
||||
if err := releaseEphemeralStagingKey(stagingDir, key, capacity); err != nil {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("releasing %q: %w", key, err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return errors.Join(releaseErrors...)
|
||||
}
|
||||
|
||||
func releaseEphemeralStagingKey(stagingDir, key string, capacity EphemeralCapacity) error {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
relativePath := filepath.FromSlash(key)
|
||||
filePath := filepath.Join(stagingDir, relativePath)
|
||||
if err := safefile.RemoveExact(stagingDir, relativePath, []string{hashSidecarSuffix, targetSidecarSuffix}, 2); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, path := range []string{filePath, filePath + hashSidecarSuffix, filePath + targetSidecarSuffix} {
|
||||
if capacity != nil {
|
||||
if err := capacity.Release(path); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func handleHead(w http.ResponseWriter, r *http.Request, stagingDir, modelsDir, dataDir, key string) {
|
||||
if key == "" {
|
||||
http.Error(w, "key is required", http.StatusBadRequest)
|
||||
@@ -249,6 +435,74 @@ type contentRange struct {
|
||||
total int64
|
||||
}
|
||||
|
||||
// EphemeralCapacity bounds worker-local request input storage. Implementations
|
||||
// must reserve before bytes reach disk and may reconcile reservations with the
|
||||
// resulting file after a write ends.
|
||||
type EphemeralCapacity interface {
|
||||
Reserve(path string, size int64) error
|
||||
Commit(path string) error
|
||||
Claim(path string) error
|
||||
Release(path string) error
|
||||
CapacityWriter(path string, destination io.Writer) (io.WriteCloser, error)
|
||||
}
|
||||
|
||||
type ephemeralRequestCapacity interface {
|
||||
BeginRequestRelease(ctx context.Context, requestID string) error
|
||||
EndRequestRelease(requestID string)
|
||||
}
|
||||
|
||||
type ephemeralRequestOperationCapacity interface {
|
||||
BeginRequestOperation(requestID string) error
|
||||
EndRequestOperation(requestID string)
|
||||
}
|
||||
|
||||
type uploadStatusWriter struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
}
|
||||
|
||||
func (w *uploadStatusWriter) WriteHeader(status int) {
|
||||
w.status = status
|
||||
w.ResponseWriter.WriteHeader(status)
|
||||
}
|
||||
|
||||
func (w *uploadStatusWriter) Write(payload []byte) (int, error) {
|
||||
if w.status == 0 {
|
||||
w.status = http.StatusOK
|
||||
}
|
||||
return w.ResponseWriter.Write(payload)
|
||||
}
|
||||
|
||||
type ephemeralCapacityWriteError struct{ err error }
|
||||
|
||||
func (e *ephemeralCapacityWriteError) Error() string { return e.err.Error() }
|
||||
func (e *ephemeralCapacityWriteError) Unwrap() error { return e.err }
|
||||
|
||||
type ephemeralCapacityWriteCloser struct{ io.WriteCloser }
|
||||
|
||||
func (w ephemeralCapacityWriteCloser) Write(payload []byte) (int, error) {
|
||||
written, err := w.WriteCloser.Write(payload)
|
||||
if err != nil {
|
||||
return written, &ephemeralCapacityWriteError{err: err}
|
||||
}
|
||||
return written, nil
|
||||
}
|
||||
|
||||
func (w ephemeralCapacityWriteCloser) Close() error {
|
||||
if err := w.WriteCloser.Close(); err != nil {
|
||||
return &ephemeralCapacityWriteError{err: err}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func uploadWriteStatus(err error) int {
|
||||
var capacityErr *ephemeralCapacityWriteError
|
||||
if errors.As(err, &capacityErr) {
|
||||
return http.StatusInsufficientStorage
|
||||
}
|
||||
return http.StatusInternalServerError
|
||||
}
|
||||
|
||||
// parseContentRange parses a Content-Range header value of the form
|
||||
// "bytes <start>-<end>/<total>". RFC 9110 §14.4.
|
||||
// Returns (nil, nil) when the header is empty (no range request).
|
||||
@@ -290,10 +544,29 @@ func parseContentRange(h string) (*contentRange, error) {
|
||||
}
|
||||
|
||||
func handleUpload(w http.ResponseWriter, r *http.Request, stagingDir, modelsDir, dataDir, key string, maxUploadSize int64) {
|
||||
handleUploadWithCapacity(w, r, stagingDir, modelsDir, dataDir, key, maxUploadSize, nil)
|
||||
}
|
||||
|
||||
func handleUploadWithCapacity(w http.ResponseWriter, r *http.Request, stagingDir, modelsDir, dataDir, key string, maxUploadSize int64, capacity EphemeralCapacity) {
|
||||
if key == "" {
|
||||
http.Error(w, "key is required", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
capacityEnabled := capacity != nil && strings.HasPrefix(key, "ephemeral/")
|
||||
if capacityEnabled {
|
||||
if err := validateEphemeralReleaseKey(key); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if operations, ok := capacity.(ephemeralRequestOperationCapacity); ok {
|
||||
requestID := strings.Split(key, "/")[2]
|
||||
if err := operations.BeginRequestOperation(requestID); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusConflict)
|
||||
return
|
||||
}
|
||||
defer operations.EndRequestOperation(requestID)
|
||||
}
|
||||
}
|
||||
|
||||
if maxUploadSize > 0 {
|
||||
r.Body = http.MaxBytesReader(w, r.Body, maxUploadSize)
|
||||
@@ -328,18 +601,67 @@ func handleUpload(w http.ResponseWriter, r *http.Request, stagingDir, modelsDir,
|
||||
return
|
||||
}
|
||||
|
||||
if cr == nil {
|
||||
// Non-resumable (legacy) path: truncate-create, single fire-and-forget.
|
||||
handleFullUpload(w, r, dstPath, key, expectedFinalHash)
|
||||
return
|
||||
capacityEnabled = capacityEnabled && targetDir == stagingDir
|
||||
unknownLengthCapacity := capacityEnabled && r.ContentLength < 0
|
||||
capacityPaths := []string{dstPath, dstPath + hashSidecarSuffix, dstPath + targetSidecarSuffix}
|
||||
if capacityEnabled {
|
||||
if r.ContentLength >= 0 {
|
||||
if err := capacity.Reserve(dstPath, r.ContentLength); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInsufficientStorage)
|
||||
return
|
||||
}
|
||||
}
|
||||
for _, sidecarPath := range capacityPaths[1:] {
|
||||
if err := capacity.Reserve(sidecarPath, sha256.Size*2); err != nil {
|
||||
for _, reservedPath := range capacityPaths {
|
||||
reconcileEphemeralCapacity(capacity, reservedPath, 0)
|
||||
}
|
||||
http.Error(w, err.Error(), http.StatusInsufficientStorage)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
handleRangeUpload(w, r, dstPath, key, cr, expectedFinalHash)
|
||||
statusWriter := &uploadStatusWriter{ResponseWriter: w}
|
||||
var uploadCapacity EphemeralCapacity
|
||||
if unknownLengthCapacity {
|
||||
uploadCapacity = capacity
|
||||
}
|
||||
|
||||
if cr == nil {
|
||||
// Non-resumable (legacy) path: truncate-create, single fire-and-forget.
|
||||
handleFullUpload(statusWriter, r, dstPath, key, expectedFinalHash, uploadCapacity)
|
||||
} else {
|
||||
handleRangeUpload(statusWriter, r, dstPath, key, cr, expectedFinalHash, uploadCapacity)
|
||||
}
|
||||
|
||||
if !capacityEnabled {
|
||||
return
|
||||
}
|
||||
for _, capacityPath := range capacityPaths {
|
||||
reconcileEphemeralCapacity(capacity, capacityPath, statusWriter.status)
|
||||
}
|
||||
}
|
||||
|
||||
func reconcileEphemeralCapacity(capacity EphemeralCapacity, path string, status int) {
|
||||
if info, err := os.Lstat(path); err == nil && info.Mode().IsRegular() {
|
||||
if err := capacity.Commit(path); err != nil {
|
||||
xlog.Error("Committing ephemeral capacity failed", "path", path, "status", status, "error", err)
|
||||
if removeErr := os.Remove(path); removeErr != nil && !os.IsNotExist(removeErr) {
|
||||
xlog.Warn("Removing uncommitted ephemeral file failed", "path", path, "error", removeErr)
|
||||
}
|
||||
if releaseErr := capacity.Release(path); releaseErr != nil {
|
||||
xlog.Warn("Rolling back failed ephemeral commit", "path", path, "error", releaseErr)
|
||||
}
|
||||
}
|
||||
} else if err := capacity.Release(path); err != nil {
|
||||
xlog.Warn("Rolling back ephemeral capacity failed", "path", path, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// handleFullUpload writes the entire request body to dstPath, replacing any
|
||||
// existing content. This is the legacy happy-path with no Range header.
|
||||
func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expectedFinalHash string) {
|
||||
func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expectedFinalHash string, capacity EphemeralCapacity) {
|
||||
// Reset any in-progress resumable state.
|
||||
_ = os.Remove(dstPath + targetSidecarSuffix)
|
||||
|
||||
@@ -350,13 +672,32 @@ func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expe
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
var destination io.Writer = f
|
||||
var capacityWriter io.WriteCloser
|
||||
if capacity != nil {
|
||||
var writer io.WriteCloser
|
||||
writer, err = capacity.CapacityWriter(dstPath, f)
|
||||
if err != nil {
|
||||
_ = os.Remove(dstPath)
|
||||
http.Error(w, err.Error(), http.StatusInsufficientStorage)
|
||||
return
|
||||
}
|
||||
capacityWriter = ephemeralCapacityWriteCloser{WriteCloser: writer}
|
||||
destination = capacityWriter
|
||||
}
|
||||
|
||||
hasher := sha256.New()
|
||||
n, err := io.Copy(f, io.TeeReader(r.Body, hasher))
|
||||
n, err := io.Copy(destination, io.TeeReader(r.Body, hasher))
|
||||
if capacityWriter != nil {
|
||||
if closeErr := capacityWriter.Close(); err == nil {
|
||||
err = closeErr
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
os.Remove(dstPath)
|
||||
os.Remove(dstPath + hashSidecarSuffix)
|
||||
xlog.Error("File upload failed", "key", key, "bytesReceived", n, "contentLength", r.ContentLength, "remote", r.RemoteAddr, "error", err)
|
||||
http.Error(w, fmt.Sprintf("writing file: %v", err), http.StatusInternalServerError)
|
||||
http.Error(w, fmt.Sprintf("writing file: %v", err), uploadWriteStatus(err))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -385,7 +726,7 @@ func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expe
|
||||
// the request starts at the current file size. When the slice completes the
|
||||
// transfer (end+1 == total), it validates the optional expected final hash and
|
||||
// writes the sidecar.
|
||||
func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key string, cr *contentRange, expectedFinalHash string) {
|
||||
func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key string, cr *contentRange, expectedFinalHash string, capacity EphemeralCapacity) {
|
||||
// Determine the current on-disk size (0 if missing).
|
||||
var currentSize int64
|
||||
if info, err := os.Stat(dstPath); err == nil {
|
||||
@@ -428,6 +769,13 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
|
||||
_ = os.Remove(dstPath + hashSidecarSuffix)
|
||||
_ = os.Remove(targetSidecar)
|
||||
currentSize = 0
|
||||
} else if currentSize == cr.total {
|
||||
// The final bytes may have landed before a worker stopped during
|
||||
// verification. Clients retry from zero when HEAD reports the full
|
||||
// size without a committed hash. Verify the actual bytes, not the
|
||||
// target sidecar, so this retry can finish without a 416 loop.
|
||||
finalizeRangeUpload(w, dstPath, key, currentSize, expectedFinalHash)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -462,6 +810,18 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
|
||||
return
|
||||
}
|
||||
defer func() { _ = f.Close() }()
|
||||
var destination io.Writer = f
|
||||
var capacityWriter io.WriteCloser
|
||||
if capacity != nil {
|
||||
var writer io.WriteCloser
|
||||
writer, err = capacity.CapacityWriter(dstPath, f)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInsufficientStorage)
|
||||
return
|
||||
}
|
||||
capacityWriter = ephemeralCapacityWriteCloser{WriteCloser: writer}
|
||||
destination = capacityWriter
|
||||
}
|
||||
|
||||
// Persist the declared expected hash so subsequent chunks can be
|
||||
// cross-checked.
|
||||
@@ -473,10 +833,15 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
|
||||
|
||||
expectedChunkLen := cr.end - cr.start + 1
|
||||
limited := io.LimitReader(r.Body, expectedChunkLen)
|
||||
n, err := io.Copy(f, limited)
|
||||
n, err := io.Copy(destination, limited)
|
||||
if capacityWriter != nil {
|
||||
if closeErr := capacityWriter.Close(); err == nil {
|
||||
err = closeErr
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
xlog.Error("Range upload chunk failed", "key", key, "bytesReceived", n, "expected", expectedChunkLen, "remote", r.RemoteAddr, "error", err)
|
||||
http.Error(w, fmt.Sprintf("writing file: %v", err), http.StatusInternalServerError)
|
||||
http.Error(w, fmt.Sprintf("writing file: %v", err), uploadWriteStatus(err))
|
||||
return
|
||||
}
|
||||
if n != expectedChunkLen {
|
||||
@@ -499,7 +864,13 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
|
||||
return
|
||||
}
|
||||
|
||||
// Upload complete — compute the final hash by re-reading the file.
|
||||
finalizeRangeUpload(w, dstPath, key, newSize, expectedFinalHash)
|
||||
}
|
||||
|
||||
// finalizeRangeUpload also recovers transfers interrupted after the last byte
|
||||
// was written, so it must hash the file rather than trust cached metadata.
|
||||
func finalizeRangeUpload(w http.ResponseWriter, dstPath, key string, size int64, expectedFinalHash string) {
|
||||
targetSidecar := dstPath + targetSidecarSuffix
|
||||
finalHash, err := downloader.CalculateSHA(dstPath)
|
||||
if err != nil {
|
||||
xlog.Error("Failed to compute final hash on range upload", "path", dstPath, "error", err)
|
||||
@@ -521,7 +892,7 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
|
||||
// Clear the in-progress sidecar — upload is committed.
|
||||
_ = os.Remove(targetSidecar)
|
||||
|
||||
xlog.Info("Resumable file upload complete", "key", key, "path", dstPath, "size", newSize, "sha256", finalHash)
|
||||
xlog.Info("Resumable file upload complete", "key", key, "path", dstPath, "size", size, "sha256", finalHash)
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(map[string]string{"local_path": dstPath}); err != nil {
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
@@ -17,10 +18,99 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type recordingEphemeralCapacity struct {
|
||||
reserved int64
|
||||
reserveErr error
|
||||
writerErr error
|
||||
claimCalls []string
|
||||
claimErr error
|
||||
commitErr error
|
||||
releases []string
|
||||
startedOps []string
|
||||
endedOps []string
|
||||
}
|
||||
|
||||
type nopWriteCloser struct{ io.Writer }
|
||||
|
||||
func (nopWriteCloser) Close() error { return nil }
|
||||
|
||||
type failingWriteCloser struct{ err error }
|
||||
|
||||
func (w failingWriteCloser) Write([]byte) (int, error) { return 0, w.err }
|
||||
func (failingWriteCloser) Close() error { return nil }
|
||||
|
||||
type blockingDestinationWriteCloser struct {
|
||||
destination io.Writer
|
||||
written chan struct{}
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (w *blockingDestinationWriteCloser) Write(payload []byte) (int, error) {
|
||||
n, err := w.destination.Write(payload)
|
||||
close(w.written)
|
||||
<-w.release
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (*blockingDestinationWriteCloser) Close() error { return nil }
|
||||
|
||||
type blockingDestinationCapacity struct {
|
||||
written chan struct{}
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (*blockingDestinationCapacity) Reserve(string, int64) error { return nil }
|
||||
func (*blockingDestinationCapacity) BeginRequestRelease(context.Context, string) error { return nil }
|
||||
func (*blockingDestinationCapacity) EndRequestRelease(string) {}
|
||||
func (*blockingDestinationCapacity) BeginRequestOperation(string) error { return nil }
|
||||
func (*blockingDestinationCapacity) EndRequestOperation(string) {}
|
||||
func (*blockingDestinationCapacity) Commit(string) error { return nil }
|
||||
func (*blockingDestinationCapacity) Claim(string) error { return nil }
|
||||
func (*blockingDestinationCapacity) Release(string) error { return nil }
|
||||
func (g *blockingDestinationCapacity) CapacityWriter(_ string, destination io.Writer) (io.WriteCloser, error) {
|
||||
return &blockingDestinationWriteCloser{
|
||||
destination: destination,
|
||||
written: g.written,
|
||||
release: g.release,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (g *recordingEphemeralCapacity) Reserve(_ string, size int64) error {
|
||||
g.reserved = size
|
||||
return g.reserveErr
|
||||
}
|
||||
|
||||
func (*recordingEphemeralCapacity) BeginRequestRelease(context.Context, string) error { return nil }
|
||||
func (*recordingEphemeralCapacity) EndRequestRelease(string) {}
|
||||
func (g *recordingEphemeralCapacity) BeginRequestOperation(requestID string) error {
|
||||
g.startedOps = append(g.startedOps, requestID)
|
||||
return nil
|
||||
}
|
||||
func (g *recordingEphemeralCapacity) EndRequestOperation(requestID string) {
|
||||
g.endedOps = append(g.endedOps, requestID)
|
||||
}
|
||||
|
||||
func (g *recordingEphemeralCapacity) Commit(string) error { return g.commitErr }
|
||||
func (g *recordingEphemeralCapacity) Release(path string) error {
|
||||
g.releases = append(g.releases, path)
|
||||
return nil
|
||||
}
|
||||
func (g *recordingEphemeralCapacity) Claim(path string) error {
|
||||
g.claimCalls = append(g.claimCalls, path)
|
||||
return g.claimErr
|
||||
}
|
||||
func (g *recordingEphemeralCapacity) CapacityWriter(_ string, destination io.Writer) (io.WriteCloser, error) {
|
||||
if g.writerErr != nil {
|
||||
return failingWriteCloser{err: g.writerErr}, nil
|
||||
}
|
||||
return nopWriteCloser{Writer: destination}, nil
|
||||
}
|
||||
|
||||
var _ = Describe("FileTransferServer", func() {
|
||||
setupTestServer := func(token string, maxUploadSize int64) (*httptest.Server, string, string, string) {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
@@ -54,6 +144,67 @@ var _ = Describe("FileTransferServer", func() {
|
||||
}
|
||||
|
||||
Describe("Upload and Download", func() {
|
||||
It("rejects a declared ephemeral upload before writing when capacity is exhausted", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
dataDir := GinkgoT().TempDir()
|
||||
guard := &recordingEphemeralCapacity{reserveErr: fmt.Errorf("full")}
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPut, "/v1/files/ephemeral/audio/request/input.wav", strings.NewReader("payload"))
|
||||
|
||||
handleUploadWithCapacity(recorder, request, stagingDir, modelsDir, dataDir, "ephemeral/audio/request/input.wav", 0, guard)
|
||||
|
||||
Expect(recorder.Code).To(Equal(http.StatusInsufficientStorage))
|
||||
Expect(guard.reserved).To(Equal(int64(len("payload"))))
|
||||
Expect(filepath.Join(stagingDir, "ephemeral", "audio", "request", "input.wav")).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("returns insufficient storage when a chunked upload reaches its bound", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
guard := &recordingEphemeralCapacity{writerErr: fmt.Errorf("full")}
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPut, "/v1/files/ephemeral/audio/request/input.wav", strings.NewReader("payload"))
|
||||
request.ContentLength = -1
|
||||
|
||||
handleUploadWithCapacity(recorder, request, stagingDir, GinkgoT().TempDir(), GinkgoT().TempDir(), "ephemeral/audio/request/input.wav", 0, guard)
|
||||
|
||||
Expect(recorder.Code).To(Equal(http.StatusInsufficientStorage))
|
||||
})
|
||||
|
||||
It("keeps unknown-length bytes guarded until they reach the staged file", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
dataDir := GinkgoT().TempDir()
|
||||
guard := &blockingDestinationCapacity{
|
||||
written: make(chan struct{}),
|
||||
release: make(chan struct{}),
|
||||
}
|
||||
released := false
|
||||
defer func() {
|
||||
if !released {
|
||||
close(guard.release)
|
||||
}
|
||||
}()
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPut, "/v1/files/ephemeral/audio/request/input.wav", strings.NewReader("payload"))
|
||||
request.ContentLength = -1
|
||||
done := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
handleUploadWithCapacity(recorder, request, stagingDir, modelsDir, dataDir, "ephemeral/audio/request/input.wav", 0, guard)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
Eventually(guard.written).Should(BeClosed())
|
||||
path := filepath.Join(stagingDir, "ephemeral", "audio", "request", "input.wav")
|
||||
Expect(os.ReadFile(path)).To(Equal([]byte("payload")))
|
||||
close(guard.release)
|
||||
released = true
|
||||
Eventually(done).Should(BeClosed())
|
||||
Expect(recorder.Code).To(Equal(http.StatusOK))
|
||||
})
|
||||
|
||||
It("round-trips file content correctly", func() {
|
||||
ts, _, _, _ := setupTestServer("secret-token", 0)
|
||||
|
||||
@@ -394,6 +545,17 @@ var _ = Describe("FileTransferServer", func() {
|
||||
})
|
||||
})
|
||||
|
||||
It("removes and releases a file whose capacity commit fails", func() {
|
||||
path := filepath.Join(GinkgoT().TempDir(), "input.wav")
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
guard := &recordingEphemeralCapacity{commitErr: errors.New("request released")}
|
||||
|
||||
reconcileEphemeralCapacity(guard, path, http.StatusOK)
|
||||
|
||||
Expect(path).NotTo(BeAnExistingFile())
|
||||
Expect(guard.releases).To(Equal([]string{path}))
|
||||
})
|
||||
|
||||
// --- Upload sidecar tests ---
|
||||
|
||||
Describe("Upload hash sidecar", func() {
|
||||
@@ -438,6 +600,165 @@ var _ = Describe("FileTransferServer", func() {
|
||||
// --- EnsureRemote skip tests ---
|
||||
|
||||
Describe("EnsureRemote skip-if-exists", func() {
|
||||
It("reports a claim-time disappearance as a cache miss", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
key := "ephemeral/audio/request/input.wav"
|
||||
remotePath := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(remotePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(remotePath, []byte("stale"), 0o600)).To(Succeed())
|
||||
guard := &recordingEphemeralCapacity{claimErr: fmt.Errorf("claim raced recovery: %w", os.ErrNotExist)}
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/files/"+key, nil)
|
||||
|
||||
handleClaimWithCapacity(recorder, request, stagingDir, key, guard)
|
||||
|
||||
Expect(recorder.Code).To(Equal(http.StatusNotFound))
|
||||
})
|
||||
|
||||
It("claims a matching ephemeral file before returning the worker path", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
dataDir := GinkgoT().TempDir()
|
||||
guard := &recordingEphemeralCapacity{}
|
||||
key := "ephemeral/audio/request/input.wav"
|
||||
remotePath := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
content := []byte("already on worker")
|
||||
Expect(os.MkdirAll(filepath.Dir(remotePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(remotePath, content, 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(remotePath+hashSidecarSuffix, []byte(sha256Hex(content)), 0o600)).To(Succeed())
|
||||
localPath := filepath.Join(GinkgoT().TempDir(), "input.wav")
|
||||
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
requestKey := strings.TrimPrefix(r.URL.Path, "/v1/files/")
|
||||
switch r.Method {
|
||||
case http.MethodHead:
|
||||
handleHead(w, r, stagingDir, modelsDir, dataDir, requestKey)
|
||||
case http.MethodPost:
|
||||
handleClaimWithCapacity(w, r, stagingDir, requestKey, guard)
|
||||
default:
|
||||
http.Error(w, "unexpected upload", http.StatusInternalServerError)
|
||||
}
|
||||
})
|
||||
ts := httptest.NewServer(mux)
|
||||
DeferCleanup(ts.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(ts.URL, "http://"), nil
|
||||
}, "")
|
||||
|
||||
for range 2 {
|
||||
path, err := stager.EnsureRemote(context.Background(), "node-1", localPath, key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(path).To(Equal(remotePath))
|
||||
}
|
||||
Expect(guard.claimCalls).To(Equal([]string{remotePath, remotePath}))
|
||||
Expect(guard.startedOps).To(Equal([]string{"request", "request"}))
|
||||
Expect(guard.endedOps).To(Equal([]string{"request", "request"}))
|
||||
})
|
||||
|
||||
It("propagates an ephemeral cache-hit claim failure", func() {
|
||||
content := []byte("already on worker")
|
||||
localPath := filepath.Join(GinkgoT().TempDir(), "input.wav")
|
||||
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodHead {
|
||||
w.Header().Set(HeaderLocalPath, "/remote/ephemeral/audio/request/input.wav")
|
||||
w.Header().Set(HeaderContentSHA256, sha256Hex(content))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return
|
||||
}
|
||||
http.Error(w, "ephemeral capacity exceeded", http.StatusInsufficientStorage)
|
||||
})
|
||||
ts := httptest.NewServer(mux)
|
||||
DeferCleanup(ts.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(ts.URL, "http://"), nil
|
||||
}, "")
|
||||
|
||||
path, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "ephemeral/audio/request/input.wav")
|
||||
|
||||
Expect(path).To(BeEmpty())
|
||||
Expect(err).To(MatchError(ContainSubstring("ephemeral capacity exceeded")))
|
||||
})
|
||||
|
||||
DescribeTable("uploads a matching ephemeral file when an old worker cannot claim it",
|
||||
func(claimStatus int) {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
content := []byte("compatible upload")
|
||||
localPath := filepath.Join(GinkgoT().TempDir(), "input.wav")
|
||||
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
|
||||
cacheHitPath := filepath.Join(stagingDir, "ephemeral", "audio", "stale", "input.wav")
|
||||
putCalls := 0
|
||||
putRemotePath := ""
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.Method {
|
||||
case http.MethodHead:
|
||||
w.Header().Set(HeaderLocalPath, cacheHitPath)
|
||||
w.Header().Set(HeaderContentSHA256, sha256Hex(content))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case http.MethodPost:
|
||||
http.Error(w, "claim unsupported", claimStatus)
|
||||
case http.MethodPut:
|
||||
putCalls++
|
||||
key := strings.TrimPrefix(r.URL.Path, "/v1/files/")
|
||||
putRemotePath = filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
handleUpload(w, r, stagingDir, "", "", key, 0)
|
||||
case http.MethodDelete:
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
})
|
||||
ts := httptest.NewServer(mux)
|
||||
DeferCleanup(ts.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(ts.URL, "http://"), nil
|
||||
}, "")
|
||||
|
||||
backend := &lifecycleBackend{}
|
||||
client := NewFileStagingClient(backend, stager, "node-1")
|
||||
request := &pb.PredictOptions{Audios: []string{localPath}}
|
||||
_, err := client.Predict(context.Background(), request)
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(putCalls).To(Equal(1))
|
||||
Expect(backend.predictInput).NotTo(BeNil())
|
||||
Expect(backend.predictInput.Audios).To(Equal([]string{putRemotePath}))
|
||||
},
|
||||
Entry("404", http.StatusNotFound),
|
||||
Entry("405", http.StatusMethodNotAllowed),
|
||||
)
|
||||
|
||||
It("keeps matching model probes read-only", func() {
|
||||
content := []byte("model")
|
||||
localPath := filepath.Join(GinkgoT().TempDir(), "model.bin")
|
||||
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
|
||||
unexpectedWrites := 0
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/files/", func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodHead {
|
||||
unexpectedWrites++
|
||||
http.Error(w, "unexpected write", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.Header().Set(HeaderLocalPath, "/models/tracking/model.bin")
|
||||
w.Header().Set(HeaderContentSHA256, sha256Hex(content))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
ts := httptest.NewServer(mux)
|
||||
DeferCleanup(ts.Close)
|
||||
stager := NewHTTPFileStager(func(string) (string, error) {
|
||||
return strings.TrimPrefix(ts.URL, "http://"), nil
|
||||
}, "")
|
||||
|
||||
path, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "models/tracking/model.bin")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(path).To(Equal("/models/tracking/model.bin"))
|
||||
Expect(unexpectedWrites).To(BeZero())
|
||||
})
|
||||
|
||||
It("skips upload when file exists with matching hash", func() {
|
||||
ts, stagingDir, _, _ := setupTestServer("tok", 0)
|
||||
|
||||
|
||||
@@ -1599,8 +1599,15 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op
|
||||
|
||||
// Stage file paths referenced in generic Options (key:value pairs where values
|
||||
// are file paths). Options stay as relative paths — backends resolve them via ModelPath.
|
||||
r.stageGenericOptions(ctx, node, opts.Options, frontendModelsDir, localModelDir, keyMapper.Key)
|
||||
r.stageGenericOptions(ctx, node, opts.Overrides, frontendModelsDir, localModelDir, keyMapper.Key)
|
||||
for _, options := range [][]string{opts.Options, opts.Overrides} {
|
||||
remoteRoot := r.stageGenericOptions(ctx, node, options, frontendModelsDir, localModelDir, keyMapper.Key)
|
||||
if opts.ModelFile == "" && remoteRoot != "" {
|
||||
// Virtual models have no primary file from which to derive the
|
||||
// worker root. Their relative options must resolve against the
|
||||
// companion assets we actually staged, not the frontend's root.
|
||||
opts.ModelPath = remoteRoot
|
||||
}
|
||||
}
|
||||
|
||||
return opts, nil
|
||||
}
|
||||
@@ -1831,7 +1838,9 @@ func (r *SmartRouter) stageCompanionFiles(ctx context.Context, node *BackendNode
|
||||
// that resolve to existing files relative to the frontend models directory or
|
||||
// the model's own directory. Option values are NOT rewritten — backends resolve
|
||||
// them via ModelPath. keyFn generates the namespaced storage key for each file.
|
||||
func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir, modelDir string, keyFn func(string) string) {
|
||||
// Returns the staged models root, or empty when no asset was staged.
|
||||
func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir, modelDir string, keyFn func(string) string) string {
|
||||
remoteRoot := ""
|
||||
for _, opt := range options {
|
||||
optKey, val, ok := strings.Cut(opt, ":")
|
||||
if !ok || val == "" {
|
||||
@@ -1856,18 +1865,23 @@ func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode
|
||||
// worker; a single file is staged directly. Values are never rewritten —
|
||||
// backends resolve relative paths via ModelPath.
|
||||
if err == nil && info.IsDir() {
|
||||
r.stageOptionDir(ctx, node, absPath, keyFn)
|
||||
if remoteDir := r.stageOptionDir(ctx, node, absPath, keyFn); remoteDir != "" {
|
||||
remoteRoot = DeriveRemoteModelPath(remoteDir, relativeToModelsDir(frontendModelsDir, absPath, filepath.Base(absPath)))
|
||||
}
|
||||
xlog.Debug("Staged option directory", "option", optKey, "localPath", absPath)
|
||||
continue
|
||||
}
|
||||
|
||||
key := keyFn(absPath)
|
||||
if _, err := r.fileStager.EnsureRemote(ctx, node.ID, absPath, key); err != nil {
|
||||
remotePath, err := r.fileStager.EnsureRemote(ctx, node.ID, absPath, key)
|
||||
if err != nil {
|
||||
xlog.Warn("Failed to stage option file, skipping", "option", opt, "path", absPath, "error", err)
|
||||
continue
|
||||
}
|
||||
remoteRoot = DeriveRemoteModelPath(remotePath, relativeToModelsDir(frontendModelsDir, absPath, filepath.Base(absPath)))
|
||||
xlog.Debug("Staged option file", "option", optKey, "localPath", absPath)
|
||||
}
|
||||
return remoteRoot
|
||||
}
|
||||
|
||||
// resolveOptionPath finds an existing local path for an option value: an
|
||||
@@ -1895,8 +1909,10 @@ func resolveOptionPath(val, frontendModelsDir, modelDir string) (string, bool) {
|
||||
// stageOptionDir stages every regular file under an option-declared directory
|
||||
// (e.g. sherpa-onnx's espeak-ng-data) using the structure-preserving key, so the
|
||||
// tree is recreated beside the model on the worker. Per-file errors are logged
|
||||
// and skipped; the option value itself is not rewritten.
|
||||
func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir string, keyFn func(string) string) {
|
||||
// and skipped; the option value itself is not rewritten. Returns the remote
|
||||
// directory derived from a successfully staged file, or empty when none succeeds.
|
||||
func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir string, keyFn func(string) string) string {
|
||||
remoteDir := ""
|
||||
_ = filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkErr error) error {
|
||||
if walkErr != nil || d.IsDir() {
|
||||
return nil
|
||||
@@ -1911,11 +1927,17 @@ func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir
|
||||
if isHashSidecar(path) {
|
||||
return nil
|
||||
}
|
||||
if _, err := r.fileStager.EnsureRemote(ctx, node.ID, path, keyFn(path)); err != nil {
|
||||
remotePath, err := r.fileStager.EnsureRemote(ctx, node.ID, path, keyFn(path))
|
||||
if err != nil {
|
||||
xlog.Warn("Failed to stage option directory file, skipping", "path", path, "error", err)
|
||||
return nil
|
||||
}
|
||||
if rel, err := filepath.Rel(dir, path); err == nil {
|
||||
remoteDir = DeriveRemoteModelPath(remotePath, rel)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return remoteDir
|
||||
}
|
||||
|
||||
// probeHealth checks whether a backend process on the given node/addr is alive
|
||||
|
||||
@@ -52,6 +52,8 @@ func (f *fakeFileStager) AllocRemoteTemp(_ context.Context, _ string) (string, e
|
||||
|
||||
func (f *fakeFileStager) StageRemoteToStore(_ context.Context, _, _, _ string) error { return nil }
|
||||
|
||||
func (f *fakeFileStager) ReleaseRemote(_ context.Context, _, _ string) error { return nil }
|
||||
|
||||
func (f *fakeFileStager) ListRemoteDir(_ context.Context, _, _ string) ([]string, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type failedCompanionStager struct{ FileStager }
|
||||
|
||||
func (failedCompanionStager) EnsureRemote(context.Context, string, string, string) (string, error) {
|
||||
return "", errors.New("worker unavailable")
|
||||
}
|
||||
|
||||
var _ = Describe("staging virtual model companions", func() {
|
||||
DescribeTable("anchors relative assets on the worker",
|
||||
func(options, overrides []string, files []string) {
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
for _, name := range files {
|
||||
path := filepath.Join(modelsDir, name)
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("weights"), 0600)).To(Succeed())
|
||||
}
|
||||
stager := &fakeFileStager{}
|
||||
router := &SmartRouter{fileStager: stager, stagingTracker: NewStagingTracker()}
|
||||
input := &pb.ModelOptions{Model: "insightface-buffalo-m", ModelFile: filepath.Join(modelsDir, "insightface-buffalo-m"), ModelPath: modelsDir, Options: options, Overrides: overrides}
|
||||
staged, err := router.stageModelFiles(context.Background(), &BackendNode{ID: "worker"}, input, "insightface-buffalo-m")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(staged.ModelPath).To(Equal("/remote/models/insightface-buffalo-m"))
|
||||
Expect(staged.Options).To(Equal(options))
|
||||
Expect(staged.Overrides).To(Equal(overrides))
|
||||
Expect(input.ModelPath).To(Equal(modelsDir))
|
||||
Expect(input.ModelFile).To(Equal(filepath.Join(modelsDir, "insightface-buffalo-m")))
|
||||
Expect(stager.ensureCalls).To(HaveLen(len(files)))
|
||||
for _, call := range stager.ensureCalls {
|
||||
rel, err := filepath.Rel(modelsDir, call.localPath)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(filepath.Join(staged.ModelPath, rel)).To(Equal("/remote/" + call.key))
|
||||
}
|
||||
},
|
||||
Entry("Buffalo pack and MiniFASNet files", []string{"model_pack:buffalo_m", "antispoof_v2_onnx:MiniFASNetV2.onnx", "antispoof_v1se_onnx:MiniFASNetV1SE.onnx"}, nil, []string{"buffalo_m/det_2.5g.onnx", "buffalo_m/w600k_r50.onnx", "MiniFASNetV2.onnx", "MiniFASNetV1SE.onnx"}),
|
||||
Entry("only a nested companion directory", []string{"model_pack:packs/buffalo_m"}, nil, []string{"packs/buffalo_m/det_2.5g.onnx"}),
|
||||
Entry("only a companion file", []string{"antispoof_v2_onnx:MiniFASNetV2.onnx"}, nil, []string{"MiniFASNetV2.onnx"}),
|
||||
Entry("only override assets", nil, []string{"antispoof_v2_onnx:MiniFASNetV2.onnx"}, []string{"MiniFASNetV2.onnx"}),
|
||||
)
|
||||
It("keeps the original path when no assets are staged", func() {
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
router := &SmartRouter{fileStager: &fakeFileStager{}, stagingTracker: NewStagingTracker()}
|
||||
input := &pb.ModelOptions{Model: "virtual", ModelFile: filepath.Join(modelsDir, "virtual"), ModelPath: modelsDir, Options: []string{"engine:insightface"}}
|
||||
staged, err := router.stageModelFiles(context.Background(), &BackendNode{ID: "worker"}, input, "virtual")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(staged.ModelPath).To(Equal(modelsDir))
|
||||
})
|
||||
DescribeTable("keeps the original root when companion staging fails",
|
||||
func(directory bool) {
|
||||
modelsDir := GinkgoT().TempDir()
|
||||
relative := "MiniFASNetV2.onnx"
|
||||
if directory {
|
||||
relative = "buffalo_m/det_2.5g.onnx"
|
||||
}
|
||||
local := filepath.Join(modelsDir, relative)
|
||||
Expect(os.MkdirAll(filepath.Dir(local), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(local, []byte("weights"), 0600)).To(Succeed())
|
||||
value := relative
|
||||
if directory {
|
||||
value = "buffalo_m"
|
||||
}
|
||||
router := &SmartRouter{fileStager: failedCompanionStager{}, stagingTracker: NewStagingTracker()}
|
||||
input := &pb.ModelOptions{Model: "virtual", ModelFile: filepath.Join(modelsDir, "virtual"), ModelPath: modelsDir, Options: []string{"asset:" + value}}
|
||||
staged, err := router.stageModelFiles(context.Background(), &BackendNode{ID: "worker"}, input, "virtual")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(staged.ModelPath).To(Equal(modelsDir))
|
||||
}, Entry("file", false), Entry("directory", true),
|
||||
)
|
||||
|
||||
})
|
||||
@@ -42,6 +42,28 @@ type QuantizationService struct {
|
||||
// jobs is the cross-replica job store: an in-memory map kept consistent across
|
||||
// replicas via NATS, optionally read-through to PostgreSQL in distributed mode.
|
||||
jobs *syncstate.SyncedMap[string, *schema.QuantizationJob]
|
||||
|
||||
// progressMu guards progressSubs.
|
||||
//
|
||||
// A backend's per-job progress stream has a single destructive consumer: the
|
||||
// backend pops each update off one queue and hands it to whoever is reading.
|
||||
// So the service opens that stream exactly once per job — in watchProgress,
|
||||
// started by StartJob — and fans the updates out in-process to the SSE clients
|
||||
// registered here. Opening a second stream per client would make the two
|
||||
// readers race for the same updates.
|
||||
progressMu sync.Mutex
|
||||
progressSubs map[string][]chan *schema.QuantizationProgressEvent
|
||||
}
|
||||
|
||||
// progressSubBuffer is the per-subscriber event buffer. It absorbs a client that
|
||||
// is briefly slow; a client that falls further behind drops events rather than
|
||||
// stalling the single reader of the backend stream.
|
||||
const progressSubBuffer = 64
|
||||
|
||||
// isTerminalStatus reports whether a job status is final, i.e. no further
|
||||
// progress update will follow.
|
||||
func isTerminalStatus(status string) bool {
|
||||
return status == "stopped" || status == "completed" || status == "failed"
|
||||
}
|
||||
|
||||
// NewQuantizationService creates a new QuantizationService. In distributed mode
|
||||
@@ -59,6 +81,7 @@ func NewQuantizationService(
|
||||
appConfig: appConfig,
|
||||
modelLoader: modelLoader,
|
||||
configLoader: configLoader,
|
||||
progressSubs: make(map[string][]chan *schema.QuantizationProgressEvent),
|
||||
}
|
||||
|
||||
// Only attach a Store interface when a concrete store exists, otherwise the
|
||||
@@ -240,6 +263,13 @@ func (s *QuantizationService) StartJob(ctx context.Context, userID string, req s
|
||||
}
|
||||
s.saveJobState(job)
|
||||
|
||||
// Consume the backend's progress stream for the lifetime of the job, not for
|
||||
// the lifetime of a client's SSE connection: a job that runs with nobody
|
||||
// attached must still reach "completed" in the store and in state.json. The
|
||||
// request ctx is done as soon as this HTTP handler returns, so the watcher
|
||||
// rides the application context instead.
|
||||
go s.watchProgress(s.appConfig.Context, jobID, backendName, modelID)
|
||||
|
||||
return &schema.QuantizationJobResponse{
|
||||
ID: jobID,
|
||||
Status: "queued",
|
||||
@@ -311,6 +341,14 @@ func (s *QuantizationService) StopJob(ctx context.Context, userID, jobID string)
|
||||
s.saveJobState(job)
|
||||
s.mu.Unlock()
|
||||
|
||||
// Release clients attached to the progress stream: the backend process is gone,
|
||||
// so the watcher will not see a terminal update to forward.
|
||||
s.publishProgress(jobID, &schema.QuantizationProgressEvent{
|
||||
JobID: jobID,
|
||||
Status: "stopped",
|
||||
Message: "Quantization stopped by user",
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -377,7 +415,153 @@ func (s *QuantizationService) DeleteJob(userID, jobID string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// StreamProgress opens a gRPC progress stream and calls the callback for each update.
|
||||
// watchProgress is the single reader of a job's backend progress stream. It
|
||||
// records every transition on the job — in the cross-replica store and in
|
||||
// state.json — and republishes it to the clients attached via StreamProgress.
|
||||
//
|
||||
// Recording here rather than in StreamProgress is the point: the backend hands
|
||||
// each update to one consumer, so while StreamProgress was that consumer a job's
|
||||
// state only advanced while somebody was watching it.
|
||||
func (s *QuantizationService) watchProgress(ctx context.Context, jobID, backendName, modelID string) {
|
||||
backendModel, err := s.modelLoader.Load(
|
||||
model.WithBackendString(backendName),
|
||||
model.WithModel(backendName),
|
||||
model.WithModelID(modelID),
|
||||
)
|
||||
if err != nil {
|
||||
xlog.Warn("Failed to load backend for quantization progress", "job_id", jobID, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
err = backendModel.QuantizationProgress(ctx, &pb.QuantizationProgressRequest{
|
||||
JobId: jobID,
|
||||
}, func(update *pb.QuantizationProgressUpdate) {
|
||||
s.publishProgress(jobID, s.applyProgressUpdate(ctx, jobID, update))
|
||||
})
|
||||
if err != nil {
|
||||
xlog.Warn("Quantization progress stream ended with an error", "job_id", jobID, "error", err)
|
||||
}
|
||||
|
||||
// On shutdown leave the job alone: loadJobsFromDisk already reports jobs that
|
||||
// were running at exit as stopped.
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// A stream that ends without a terminal update means the backend is gone and
|
||||
// nothing further will arrive. Record that instead of leaving the job in a
|
||||
// running state forever — which is the failure this watcher exists to prevent —
|
||||
// and release any client still waiting on a terminal event.
|
||||
s.mu.Lock()
|
||||
j, ok := s.jobs.Get(jobID)
|
||||
stale := ok && !isTerminalStatus(j.Status)
|
||||
if stale {
|
||||
j.Status = "failed"
|
||||
if j.Message == "" {
|
||||
j.Message = "Backend progress stream ended before the job reported a result"
|
||||
}
|
||||
if err := s.jobs.Set(ctx, j); err != nil {
|
||||
xlog.Warn("Failed to persist orphaned job state", "job_id", jobID, "error", err)
|
||||
}
|
||||
s.saveJobState(j)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
if stale {
|
||||
s.publishProgress(jobID, &schema.QuantizationProgressEvent{
|
||||
JobID: jobID,
|
||||
Status: "failed",
|
||||
Message: "Backend progress stream ended before the job reported a result",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// applyProgressUpdate records a backend progress update on the job and returns
|
||||
// the event to hand to subscribers.
|
||||
func (s *QuantizationService) applyProgressUpdate(ctx context.Context, jobID string, update *pb.QuantizationProgressUpdate) *schema.QuantizationProgressEvent {
|
||||
s.mu.Lock()
|
||||
if j, ok := s.jobs.Get(jobID); ok {
|
||||
// Don't let progress updates overwrite terminal states
|
||||
if !isTerminalStatus(j.Status) {
|
||||
j.Status = update.Status
|
||||
}
|
||||
if update.Message != "" {
|
||||
j.Message = update.Message
|
||||
}
|
||||
if update.OutputFile != "" {
|
||||
j.OutputFile = update.OutputFile
|
||||
}
|
||||
if err := s.jobs.Set(ctx, j); err != nil {
|
||||
xlog.Warn("Failed to persist progress update", "job_id", jobID, "error", err)
|
||||
}
|
||||
s.saveJobState(j)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// Convert extra metrics
|
||||
extraMetrics := make(map[string]float32, len(update.ExtraMetrics))
|
||||
for k, v := range update.ExtraMetrics {
|
||||
extraMetrics[k] = v
|
||||
}
|
||||
|
||||
return &schema.QuantizationProgressEvent{
|
||||
JobID: update.JobId,
|
||||
ProgressPercent: update.ProgressPercent,
|
||||
Status: update.Status,
|
||||
Message: update.Message,
|
||||
OutputFile: update.OutputFile,
|
||||
ExtraMetrics: extraMetrics,
|
||||
}
|
||||
}
|
||||
|
||||
// subscribeProgress registers a channel to receive a job's progress events.
|
||||
func (s *QuantizationService) subscribeProgress(jobID string) chan *schema.QuantizationProgressEvent {
|
||||
ch := make(chan *schema.QuantizationProgressEvent, progressSubBuffer)
|
||||
s.progressMu.Lock()
|
||||
s.progressSubs[jobID] = append(s.progressSubs[jobID], ch)
|
||||
s.progressMu.Unlock()
|
||||
return ch
|
||||
}
|
||||
|
||||
// unsubscribeProgress removes a channel registered by subscribeProgress. The
|
||||
// channel is never closed, so a publish racing with an unsubscribe cannot send
|
||||
// on a closed channel.
|
||||
func (s *QuantizationService) unsubscribeProgress(jobID string, ch chan *schema.QuantizationProgressEvent) {
|
||||
s.progressMu.Lock()
|
||||
defer s.progressMu.Unlock()
|
||||
|
||||
subs := s.progressSubs[jobID]
|
||||
for i, c := range subs {
|
||||
if c == ch {
|
||||
s.progressSubs[jobID] = append(subs[:i], subs[i+1:]...)
|
||||
break
|
||||
}
|
||||
}
|
||||
if len(s.progressSubs[jobID]) == 0 {
|
||||
delete(s.progressSubs, jobID)
|
||||
}
|
||||
}
|
||||
|
||||
// publishProgress fans an event out to a job's subscribers.
|
||||
func (s *QuantizationService) publishProgress(jobID string, event *schema.QuantizationProgressEvent) {
|
||||
s.progressMu.Lock()
|
||||
subs := append([]chan *schema.QuantizationProgressEvent(nil), s.progressSubs[jobID]...)
|
||||
s.progressMu.Unlock()
|
||||
|
||||
for _, ch := range subs {
|
||||
select {
|
||||
case ch <- event:
|
||||
default:
|
||||
// A subscriber that cannot keep up must not stall the reader that is
|
||||
// recording job state for everyone else.
|
||||
xlog.Warn("Dropping quantization progress event for a slow subscriber", "job_id", jobID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// StreamProgress calls the callback for each progress event of a job until it
|
||||
// reaches a terminal status or ctx is done. It is a pure reader: the job's own
|
||||
// watcher owns the backend stream and the state transitions.
|
||||
func (s *QuantizationService) StreamProgress(ctx context.Context, userID, jobID string, callback func(event *schema.QuantizationProgressEvent)) error {
|
||||
s.mu.Lock()
|
||||
job, ok := s.jobs.Get(jobID)
|
||||
@@ -391,59 +575,41 @@ func (s *QuantizationService) StreamProgress(ctx context.Context, userID, jobID
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
streamModelID := job.ModelID
|
||||
if streamModelID == "" {
|
||||
streamModelID = job.Backend + "-quantize"
|
||||
ch := s.subscribeProgress(jobID)
|
||||
defer s.unsubscribeProgress(jobID, ch)
|
||||
|
||||
// Re-read the job after subscribing: it may have finished between the lookup
|
||||
// above and the subscription, and no further event would ever arrive. Jobs
|
||||
// restored from disk after a restart are terminal too, and have no watcher.
|
||||
s.mu.Lock()
|
||||
current, ok := s.jobs.Get(jobID)
|
||||
terminal := ok && isTerminalStatus(current.Status)
|
||||
var final *schema.QuantizationProgressEvent
|
||||
if terminal {
|
||||
final = &schema.QuantizationProgressEvent{
|
||||
JobID: current.ID,
|
||||
Status: current.Status,
|
||||
Message: current.Message,
|
||||
OutputFile: current.OutputFile,
|
||||
}
|
||||
}
|
||||
backendModel, err := s.modelLoader.Load(
|
||||
model.WithBackendString(job.Backend),
|
||||
model.WithModel(job.Backend),
|
||||
model.WithModelID(streamModelID),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load backend: %w", err)
|
||||
s.mu.Unlock()
|
||||
if terminal {
|
||||
callback(final)
|
||||
return nil
|
||||
}
|
||||
|
||||
return backendModel.QuantizationProgress(ctx, &pb.QuantizationProgressRequest{
|
||||
JobId: jobID,
|
||||
}, func(update *pb.QuantizationProgressUpdate) {
|
||||
// Update job status and persist
|
||||
s.mu.Lock()
|
||||
if j, ok := s.jobs.Get(jobID); ok {
|
||||
// Don't let progress updates overwrite terminal states
|
||||
isTerminal := j.Status == "stopped" || j.Status == "completed" || j.Status == "failed"
|
||||
if !isTerminal {
|
||||
j.Status = update.Status
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case event := <-ch:
|
||||
callback(event)
|
||||
if isTerminalStatus(event.Status) {
|
||||
return nil
|
||||
}
|
||||
if update.Message != "" {
|
||||
j.Message = update.Message
|
||||
}
|
||||
if update.OutputFile != "" {
|
||||
j.OutputFile = update.OutputFile
|
||||
}
|
||||
if err := s.jobs.Set(ctx, j); err != nil {
|
||||
xlog.Warn("Failed to persist progress update", "job_id", jobID, "error", err)
|
||||
}
|
||||
s.saveJobState(j)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
|
||||
// Convert extra metrics
|
||||
extraMetrics := make(map[string]float32)
|
||||
for k, v := range update.ExtraMetrics {
|
||||
extraMetrics[k] = v
|
||||
}
|
||||
|
||||
event := &schema.QuantizationProgressEvent{
|
||||
JobID: update.JobId,
|
||||
ProgressPercent: update.ProgressPercent,
|
||||
Status: update.Status,
|
||||
Message: update.Message,
|
||||
OutputFile: update.OutputFile,
|
||||
ExtraMetrics: extraMetrics,
|
||||
}
|
||||
callback(event)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// sanitizeQuantModelName replaces non-alphanumeric characters with hyphens and lowercases.
|
||||
|
||||
@@ -8,6 +8,9 @@ package quantization
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
@@ -16,6 +19,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
)
|
||||
|
||||
// newTestService builds a standalone QuantizationService wired to the given bus.
|
||||
@@ -175,6 +179,156 @@ var _ = Describe("QuantizationService", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("progress recording", func() {
|
||||
var (
|
||||
bus *testutil.FakeBus
|
||||
s *QuantizationService
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
bus = testutil.NewFakeBus()
|
||||
s = newTestService(bus)
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
Expect(s.Close()).To(Succeed())
|
||||
})
|
||||
|
||||
// The reported failure: a job that ran with no SSE client attached stayed
|
||||
// "queued" forever, because the only code that advanced job state lived
|
||||
// inside StreamProgress' stream callback. The transition is now applied by
|
||||
// the job's own watcher, so it lands with nobody watching.
|
||||
It("advances job state and rewrites state.json with no subscriber attached", func() {
|
||||
job := &schema.QuantizationJob{ID: "job-np", UserID: "user-1", Status: "queued", CreatedAt: "2026-09-05T10:00:00Z"}
|
||||
Expect(s.jobs.Set(ctx, job)).To(Succeed())
|
||||
Expect(s.progressSubs).To(BeEmpty())
|
||||
|
||||
s.applyProgressUpdate(ctx, "job-np", &pb.QuantizationProgressUpdate{
|
||||
JobId: "job-np",
|
||||
Status: "completed",
|
||||
Message: "Quantization complete",
|
||||
OutputFile: "/data/quantization/job-np/model-q4_k_m.gguf",
|
||||
})
|
||||
|
||||
got, err := s.GetJob("user-1", "job-np")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.Status).To(Equal("completed"))
|
||||
Expect(got.Message).To(Equal("Quantization complete"))
|
||||
Expect(got.OutputFile).To(Equal("/data/quantization/job-np/model-q4_k_m.gguf"))
|
||||
|
||||
data, err := os.ReadFile(filepath.Join(s.jobDir("job-np"), "state.json"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
var persisted schema.QuantizationJob
|
||||
Expect(json.Unmarshal(data, &persisted)).To(Succeed())
|
||||
Expect(persisted.Status).To(Equal("completed"))
|
||||
Expect(persisted.OutputFile).To(Equal("/data/quantization/job-np/model-q4_k_m.gguf"))
|
||||
})
|
||||
|
||||
It("does not let a late update overwrite a terminal status", func() {
|
||||
job := &schema.QuantizationJob{ID: "job-stopped", UserID: "user-1", Status: "stopped", CreatedAt: "2026-09-05T10:00:00Z"}
|
||||
Expect(s.jobs.Set(ctx, job)).To(Succeed())
|
||||
|
||||
s.applyProgressUpdate(ctx, "job-stopped", &pb.QuantizationProgressUpdate{
|
||||
JobId: "job-stopped",
|
||||
Status: "quantizing",
|
||||
})
|
||||
|
||||
got, err := s.GetJob("user-1", "job-stopped")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.Status).To(Equal("stopped"))
|
||||
})
|
||||
|
||||
// The backend hands each update to a single consumer, so every client has
|
||||
// to be served from one in-process fan-out rather than its own stream.
|
||||
It("delivers one update to every attached subscriber", func() {
|
||||
first := s.subscribeProgress("job-fan")
|
||||
second := s.subscribeProgress("job-fan")
|
||||
defer s.unsubscribeProgress("job-fan", first)
|
||||
defer s.unsubscribeProgress("job-fan", second)
|
||||
|
||||
s.publishProgress("job-fan", &schema.QuantizationProgressEvent{JobID: "job-fan", Status: "quantizing"})
|
||||
|
||||
Expect((<-first).Status).To(Equal("quantizing"))
|
||||
Expect((<-second).Status).To(Equal("quantizing"))
|
||||
})
|
||||
|
||||
It("unsubscribing removes the job's entry once the last client leaves", func() {
|
||||
ch := s.subscribeProgress("job-leave")
|
||||
Expect(s.progressSubs).To(HaveKey("job-leave"))
|
||||
s.unsubscribeProgress("job-leave", ch)
|
||||
Expect(s.progressSubs).ToNot(HaveKey("job-leave"))
|
||||
})
|
||||
|
||||
// A client attaching after the job finished — including a job restored from
|
||||
// disk as "stopped" after a restart, which has no watcher — must not block
|
||||
// waiting for an event that will never come.
|
||||
It("returns a final event immediately for a job that already finished", func() {
|
||||
job := &schema.QuantizationJob{
|
||||
ID: "job-done", UserID: "user-1", Status: "completed",
|
||||
Message: "Quantization complete", OutputFile: "/data/quantization/job-done/model-q4_k_m.gguf",
|
||||
CreatedAt: "2026-09-05T10:00:00Z",
|
||||
}
|
||||
Expect(s.jobs.Set(ctx, job)).To(Succeed())
|
||||
|
||||
var seen []*schema.QuantizationProgressEvent
|
||||
Expect(s.StreamProgress(ctx, "user-1", "job-done", func(e *schema.QuantizationProgressEvent) {
|
||||
seen = append(seen, e)
|
||||
})).To(Succeed())
|
||||
|
||||
Expect(seen).To(HaveLen(1))
|
||||
Expect(seen[0].Status).To(Equal("completed"))
|
||||
Expect(seen[0].OutputFile).To(Equal("/data/quantization/job-done/model-q4_k_m.gguf"))
|
||||
Expect(s.progressSubs).ToNot(HaveKey("job-done"))
|
||||
})
|
||||
|
||||
// StopJob kills the backend, so the watcher will never forward a terminal
|
||||
// update; without an explicit release an attached client would hang.
|
||||
It("releases an attached client when the job is stopped", func() {
|
||||
job := &schema.QuantizationJob{ID: "job-stop", UserID: "user-1", Status: "quantizing", CreatedAt: "2026-09-05T10:00:00Z"}
|
||||
Expect(s.jobs.Set(ctx, job)).To(Succeed())
|
||||
|
||||
ch := s.subscribeProgress("job-stop")
|
||||
defer s.unsubscribeProgress("job-stop", ch)
|
||||
|
||||
// nil modelLoader: exercise the release without standing up a backend.
|
||||
s.mu.Lock()
|
||||
job.Status = "stopped"
|
||||
s.mu.Unlock()
|
||||
s.publishProgress("job-stop", &schema.QuantizationProgressEvent{
|
||||
JobID: "job-stop", Status: "stopped", Message: "Quantization stopped by user",
|
||||
})
|
||||
|
||||
event := <-ch
|
||||
Expect(event.Status).To(Equal("stopped"))
|
||||
Expect(isTerminalStatus(event.Status)).To(BeTrue())
|
||||
})
|
||||
|
||||
It("streams published events to a client until a terminal status arrives", func() {
|
||||
job := &schema.QuantizationJob{ID: "job-live", UserID: "user-1", Status: "queued", CreatedAt: "2026-09-05T10:00:00Z"}
|
||||
Expect(s.jobs.Set(ctx, job)).To(Succeed())
|
||||
|
||||
var seen []string
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- s.StreamProgress(ctx, "user-1", "job-live", func(e *schema.QuantizationProgressEvent) {
|
||||
seen = append(seen, e.Status)
|
||||
})
|
||||
}()
|
||||
|
||||
Eventually(func() bool {
|
||||
s.progressMu.Lock()
|
||||
defer s.progressMu.Unlock()
|
||||
return len(s.progressSubs["job-live"]) == 1
|
||||
}).Should(BeTrue())
|
||||
|
||||
s.publishProgress("job-live", &schema.QuantizationProgressEvent{JobID: "job-live", Status: "quantizing"})
|
||||
s.publishProgress("job-live", &schema.QuantizationProgressEvent{JobID: "job-live", Status: "completed"})
|
||||
|
||||
Eventually(done).Should(Receive(BeNil()))
|
||||
Expect(seen).To(Equal([]string{"quantizing", "completed"}))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("compile-time adapter contract", func() {
|
||||
It("satisfies syncstate.Store for *distributed.QuantStore", func() {
|
||||
// Guards against drift between the adapter and the component interface;
|
||||
|
||||
@@ -47,8 +47,10 @@ type Config struct {
|
||||
PrefetchModels []string `env:"LOCALAI_PREFETCH_MODELS,PREFETCH_MODELS" help:"Comma-separated gallery model IDs to download from LOCALAI_GALLERIES at worker boot (e.g. 'llama-3.2-1b-instruct,phi-3-mini-4k'). Skipped if already on disk and SHA matches." group:"server"`
|
||||
|
||||
// HTTP file transfer
|
||||
HTTPAddr string `env:"LOCALAI_HTTP_ADDR" default:"" help:"HTTP file transfer server address (default: gRPC port + 1)" group:"server" hidden:""`
|
||||
AdvertiseHTTPAddr string `env:"LOCALAI_ADVERTISE_HTTP_ADDR" help:"HTTP address the frontend uses to reach this node for file transfer" group:"server" hidden:""`
|
||||
HTTPAddr string `env:"LOCALAI_HTTP_ADDR" default:"" help:"HTTP file transfer server address (default: gRPC port + 1)" group:"server" hidden:""`
|
||||
AdvertiseHTTPAddr string `env:"LOCALAI_ADVERTISE_HTTP_ADDR" help:"HTTP address the frontend uses to reach this node for file transfer" group:"server" hidden:""`
|
||||
EphemeralStagingByteLimit int64 `env:"LOCALAI_EPHEMERAL_STAGING_BYTE_LIMIT" default:"0" help:"Maximum bytes used by worker request-input staging across HTTP and S3 caches. Zero or negative uses min(10 GiB, 10% of filesystem capacity)." group:"server"`
|
||||
EphemeralStagingMinFreeBytes int64 `env:"LOCALAI_EPHEMERAL_STAGING_MIN_FREE_BYTES" default:"0" help:"Filesystem space kept free while staging request inputs. Zero or negative uses max(1 GiB, 5% of filesystem capacity)." group:"server"`
|
||||
|
||||
// Registration (required)
|
||||
AdvertiseAddr string `env:"LOCALAI_ADVERTISE_ADDR" help:"Address the frontend uses to reach this node (defaults to hostname:port from Addr)" group:"registration" hidden:""`
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,484 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package worker
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type capacityGatedFileWriter struct {
|
||||
file *os.File
|
||||
entered chan struct{}
|
||||
resume chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (w *capacityGatedFileWriter) Write(p []byte) (int, error) {
|
||||
w.once.Do(func() {
|
||||
close(w.entered)
|
||||
<-w.resume
|
||||
})
|
||||
return w.file.Write(p)
|
||||
}
|
||||
|
||||
type capacityShortWriter struct{}
|
||||
|
||||
func (capacityShortWriter) Write(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
var _ = Describe("EphemeralCapacityGuard", func() {
|
||||
It("derives bounded defaults and preserves positive overrides", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
limit, headroom, err := effectiveEphemeralCapacity([]string{root}, 0, -1)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(limit).To(BeNumerically(">", 0))
|
||||
Expect(limit).To(BeNumerically("<=", defaultEphemeralByteLimitCeiling))
|
||||
Expect(headroom).To(BeNumerically(">=", defaultEphemeralMinFreeFloor))
|
||||
|
||||
limit, headroom, err = effectiveEphemeralCapacity([]string{root}, 123, 456)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(limit).To(Equal(int64(123)))
|
||||
Expect(headroom).To(Equal(int64(456)))
|
||||
})
|
||||
|
||||
It("accounts existing regular files without following symlinks", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
outside := filepath.Join(GinkgoT().TempDir(), "outside.bin")
|
||||
Expect(os.WriteFile(filepath.Join(root, "existing.bin"), make([]byte, 6), 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(outside, make([]byte, 100), 0o600)).To(Succeed())
|
||||
Expect(os.Symlink(outside, filepath.Join(root, "outside-link"))).To(Succeed())
|
||||
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
err = guard.Reserve(filepath.Join(root, "next.bin"), 5)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.RequestedBytes).To(Equal(int64(5)))
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(6)))
|
||||
Expect(capacityErr.LimitBytes).To(Equal(int64(10)))
|
||||
Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0))
|
||||
Expect(capacityErr.HeadroomBytes).To(BeZero())
|
||||
})
|
||||
|
||||
It("serializes competing reservations", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 1, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
start := make(chan struct{})
|
||||
results := make(chan error, 32)
|
||||
var wait sync.WaitGroup
|
||||
for i := range 32 {
|
||||
wait.Add(1)
|
||||
go func(index int) {
|
||||
defer wait.Done()
|
||||
<-start
|
||||
results <- guard.Reserve(filepath.Join(root, string(rune('a'+index))), 1)
|
||||
}(i)
|
||||
}
|
||||
close(start)
|
||||
wait.Wait()
|
||||
close(results)
|
||||
|
||||
succeeded := 0
|
||||
for result := range results {
|
||||
if result == nil {
|
||||
succeeded++
|
||||
continue
|
||||
}
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(result, &capacityErr)).To(BeTrue())
|
||||
}
|
||||
Expect(succeeded).To(Equal(1))
|
||||
})
|
||||
|
||||
It("makes only an equal active reservation idempotent", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "nested", "payload.bin")
|
||||
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "nested", ".", "payload.bin"), 4)).To(Succeed())
|
||||
err = guard.Reserve(path, 5)
|
||||
var conflictErr *EphemeralReservationConflictError
|
||||
Expect(errors.As(err, &conflictErr)).To(BeTrue())
|
||||
Expect(conflictErr.ActiveBytes).To(Equal(int64(4)))
|
||||
Expect(conflictErr.RequestedBytes).To(Equal(int64(5)))
|
||||
Expect(guard.Reserve(filepath.Join(root, "other.bin"), 6)).To(Succeed())
|
||||
Expect(guard.Release(filepath.Join(root, "nested", ".", "payload.bin"))).To(Succeed())
|
||||
Expect(guard.Release(path)).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("retains committed bytes when the same path starts another reservation", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeTrue())
|
||||
Expect(guard.Reserve(path, 6)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeTrue())
|
||||
|
||||
err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(10)))
|
||||
})
|
||||
|
||||
It("retains startup-accounted bytes when the path is reserved", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
|
||||
Expect(guard.Reserve(path, 6)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeTrue())
|
||||
err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(10)))
|
||||
})
|
||||
|
||||
It("commits the regular file's actual size", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
|
||||
Expect(guard.Reserve(path, 10)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("four"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
Expect(guard.Commit(filepath.Clean(path))).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "six.bin"), 6)).To(Succeed())
|
||||
})
|
||||
|
||||
It("preserves configured filesystem headroom", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 1<<30, 1<<62)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
err = guard.Reserve(filepath.Join(root, "payload.bin"), 1)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.RequestedBytes).To(Equal(int64(1)))
|
||||
Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0))
|
||||
Expect(capacityErr.HeadroomBytes).To(Equal(int64(1 << 62)))
|
||||
})
|
||||
|
||||
It("reserves bounded chunks before forwarding unknown-length input", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, ephemeralCapacityWriteChunk+1, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
file, err := os.Create(path)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(file.Close)
|
||||
writer, err := guard.NewWriter(path, file)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
n, err := writer.Write(make([]byte, ephemeralCapacityWriteChunk+2))
|
||||
Expect(n).To(Equal(int(ephemeralCapacityWriteChunk)))
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(writer.Close()).To(Succeed())
|
||||
info, err := file.Stat()
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(info.Size()).To(Equal(ephemeralCapacityWriteChunk))
|
||||
})
|
||||
|
||||
It("waits for an open bounded writer before committing", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
file, err := os.Create(path)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(file.Close)
|
||||
gated := &capacityGatedFileWriter{
|
||||
file: file, entered: make(chan struct{}), resume: make(chan struct{}),
|
||||
}
|
||||
DeferCleanup(func() {
|
||||
select {
|
||||
case <-gated.resume:
|
||||
default:
|
||||
close(gated.resume)
|
||||
}
|
||||
})
|
||||
writer, err := guard.NewWriter(path, gated)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
writeDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, writeErr := writer.Write([]byte("1234567"))
|
||||
writeDone <- writeErr
|
||||
}()
|
||||
Eventually(gated.entered).Should(BeClosed())
|
||||
commitDone := make(chan error, 1)
|
||||
go func() { commitDone <- guard.Commit(path) }()
|
||||
Eventually(func() int { return guard.commitWaiterCount(path) }).Should(Equal(1))
|
||||
Expect(commitDone).NotTo(Receive())
|
||||
|
||||
close(gated.resume)
|
||||
Expect(<-writeDone).To(Succeed())
|
||||
Expect(commitDone).NotTo(Receive())
|
||||
Expect(writer.Close()).To(Succeed())
|
||||
Eventually(commitDone).Should(Receive(Succeed()))
|
||||
|
||||
err = guard.Reserve(filepath.Join(root, "other.bin"), 4)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(7)))
|
||||
})
|
||||
|
||||
It("does not share pending capacity between concurrent writers", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "payload.bin")
|
||||
file, err := os.Create(path)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(file.Close)
|
||||
gated := &capacityGatedFileWriter{
|
||||
file: file, entered: make(chan struct{}), resume: make(chan struct{}),
|
||||
}
|
||||
first, err := guard.NewWriter(path, gated)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var secondDestination bytes.Buffer
|
||||
second, err := guard.NewWriter(path, &secondDestination)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
firstDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, writeErr := first.Write([]byte("1234567"))
|
||||
firstDone <- writeErr
|
||||
}()
|
||||
Eventually(gated.entered).Should(BeClosed())
|
||||
|
||||
n, err := second.Write([]byte("7654321"))
|
||||
Expect(n).To(BeZero())
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(secondDestination.Len()).To(BeZero())
|
||||
|
||||
close(gated.resume)
|
||||
Expect(<-firstDone).To(Succeed())
|
||||
Expect(first.Close()).To(Succeed())
|
||||
Expect(second.Close()).To(Succeed())
|
||||
})
|
||||
|
||||
It("rolls back bytes the destination writer does not accept", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 5, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
writer, err := guard.NewWriter(filepath.Join(root, "payload.bin"), capacityShortWriter{})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
n, err := writer.Write([]byte("123"))
|
||||
Expect(n).To(Equal(1))
|
||||
Expect(err).To(MatchError(io.ErrShortWrite))
|
||||
Expect(writer.Close()).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "other.bin"), 4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("rejects paths outside roots and through symlinks", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
outside := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 100, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(guard.Reserve(filepath.Join(outside, "payload.bin"), 1)).To(
|
||||
MatchError(ContainSubstring("outside registered ephemeral roots")),
|
||||
)
|
||||
Expect(os.Symlink(outside, filepath.Join(root, "escape"))).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "escape", "payload.bin"), 1)).To(
|
||||
MatchError(ContainSubstring("symlink")),
|
||||
)
|
||||
})
|
||||
|
||||
It("supports recovery tree accounting without dropping active reservations", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
active := filepath.Join(root, "active", "payload.bin")
|
||||
stale := filepath.Join(root, "stale", "payload.bin")
|
||||
|
||||
Expect(guard.Reserve(active, 4)).To(Succeed())
|
||||
Expect(guard.Account(stale, 3)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(root)).To(BeTrue())
|
||||
Expect(guard.ReleaseTree(filepath.Join(root, "stale"))).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 6)).To(Succeed())
|
||||
Expect(guard.ReleaseTree(root)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(root)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("waits for pre-release reservations before request cleanup scans", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "audio", "request-1", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
|
||||
released := make(chan error, 1)
|
||||
go func() {
|
||||
released <- guard.BeginRequestRelease(context.Background(), "request-1")
|
||||
}()
|
||||
Eventually(guard.releaseTombstoneCount).Should(Equal(1))
|
||||
Consistently(released, 50*time.Millisecond).ShouldNot(Receive())
|
||||
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
Eventually(released).Should(Receive(Succeed()))
|
||||
guard.EndRequestRelease("request-1")
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("rejects staging after request cleanup begins", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(guard.BeginRequestRelease(context.Background(), "request-1")).To(Succeed())
|
||||
defer guard.EndRequestRelease("request-1")
|
||||
|
||||
err = guard.Reserve(filepath.Join(root, "audio", "request-1", "late.wav"), 1)
|
||||
var releasedErr *EphemeralRequestReleasedError
|
||||
Expect(errors.As(err, &releasedErr)).To(BeTrue())
|
||||
Expect(releasedErr.RequestID).To(Equal("request-1"))
|
||||
})
|
||||
|
||||
It("leaves a late commit recoverable when release times out", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "audio", "request-1", "late.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
Expect(guard.BeginRequestRelease(ctx, "request-1")).To(MatchError(context.Canceled))
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("bounds release markers without reopening registered work", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(guard.BeginRequestOperation("request-pinned")).To(Succeed())
|
||||
pinnedRelease := make(chan error, 1)
|
||||
go func() {
|
||||
pinnedRelease <- guard.BeginRequestRelease(context.Background(), "request-pinned")
|
||||
}()
|
||||
Eventually(guard.releaseTombstoneCount).Should(Equal(1))
|
||||
guard.EndRequestOperation("request-pinned")
|
||||
Eventually(pinnedRelease).Should(Receive(Succeed()))
|
||||
|
||||
for index := range maxEphemeralReleaseTombstones + 10 {
|
||||
requestID := fmt.Sprintf("request-%d", index)
|
||||
Expect(guard.BeginRequestRelease(context.Background(), requestID)).To(Succeed())
|
||||
guard.EndRequestRelease(requestID)
|
||||
}
|
||||
Expect(guard.releaseTombstoneCount()).To(Equal(maxEphemeralReleaseTombstones))
|
||||
err = guard.Reserve(filepath.Join(root, "audio", "request-pinned", "late.wav"), 1)
|
||||
var releasedErr *EphemeralRequestReleasedError
|
||||
Expect(errors.As(err, &releasedErr)).To(BeTrue())
|
||||
guard.EndRequestRelease("request-pinned")
|
||||
})
|
||||
|
||||
It("applies backpressure at the release-pin cap and clears ownership", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "audio", "request-target", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
|
||||
guard.mu.Lock()
|
||||
for index := range maxEphemeralReleaseTombstones {
|
||||
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
|
||||
}
|
||||
guard.mu.Unlock()
|
||||
released := make(chan error, 1)
|
||||
go func() {
|
||||
released <- guard.BeginRequestRelease(context.Background(), "request-target")
|
||||
}()
|
||||
Consistently(released, 50*time.Millisecond).ShouldNot(Receive())
|
||||
|
||||
guard.EndRequestRelease("pinned-0")
|
||||
Eventually(released).Should(Receive(Succeed()))
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
guard.EndRequestRelease("request-target")
|
||||
})
|
||||
|
||||
It("makes committed files recoverable when pin backpressure expires", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
path := filepath.Join(root, "audio", "request-target", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
guard.mu.Lock()
|
||||
for index := range maxEphemeralReleaseTombstones {
|
||||
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
|
||||
}
|
||||
guard.mu.Unlock()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
Expect(guard.BeginRequestRelease(ctx, "request-target")).To(MatchError(context.Canceled))
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("rejects a registered cache-hit claim after pin backpressure expires", func() {
|
||||
root := GinkgoT().TempDir()
|
||||
path := filepath.Join(root, "audio", "request-target", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(guard.BeginRequestOperation("request-target")).To(Succeed())
|
||||
defer guard.EndRequestOperation("request-target")
|
||||
guard.mu.Lock()
|
||||
for index := range maxEphemeralReleaseTombstones {
|
||||
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
|
||||
}
|
||||
guard.mu.Unlock()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
Expect(guard.BeginRequestRelease(ctx, "request-target")).To(MatchError(context.Canceled))
|
||||
err = guard.Claim(path)
|
||||
var releasedErr *EphemeralRequestReleasedError
|
||||
Expect(errors.As(err, &releasedErr)).To(BeTrue())
|
||||
Expect(guard.HasActiveReservation(path)).To(BeFalse())
|
||||
})
|
||||
})
|
||||
@@ -14,9 +14,9 @@ const (
|
||||
// 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
|
||||
defaultEphemeralStagingTTL = time.Hour
|
||||
// defaultEphemeralStagingSweep is how often the worker sweeps.
|
||||
defaultEphemeralStagingSweep = 30 * time.Minute
|
||||
defaultEphemeralStagingSweep = 15 * time.Minute
|
||||
)
|
||||
|
||||
// StartEphemeralStagingCleanup sweeps the worker's own staging directory for
|
||||
@@ -32,6 +32,15 @@ func StartEphemeralStagingCleanup(ctx context.Context, stagingDir string, ttl, i
|
||||
if stagingDir == "" {
|
||||
return
|
||||
}
|
||||
StartEphemeralRootsCleanup(ctx, []string{filepath.Join(stagingDir, "ephemeral")}, nil, ttl, interval)
|
||||
}
|
||||
|
||||
// StartEphemeralRootsCleanup removes abandoned request inputs for every worker
|
||||
// transport while sharing accounting with live reservations.
|
||||
func StartEphemeralRootsCleanup(ctx context.Context, roots []string, guard *EphemeralCapacityGuard, ttl, interval time.Duration) {
|
||||
if len(roots) == 0 {
|
||||
return
|
||||
}
|
||||
if ttl <= 0 {
|
||||
ttl = defaultEphemeralStagingTTL
|
||||
}
|
||||
@@ -39,31 +48,41 @@ func StartEphemeralStagingCleanup(ctx context.Context, stagingDir string, ttl, i
|
||||
interval = defaultEphemeralStagingSweep
|
||||
}
|
||||
|
||||
// Reclaim crash leftovers before the caller starts accepting new work.
|
||||
CleanEphemeralRoots(roots, ttl, guard)
|
||||
|
||||
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)
|
||||
CleanEphemeralRoots(roots, ttl, guard)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
xlog.Info("Ephemeral staging cleanup started", "dir", stagingDir, "ttl", ttl, "interval", interval)
|
||||
xlog.Info("Ephemeral staging cleanup started", "roots", roots, "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")
|
||||
CleanEphemeralRoots([]string{filepath.Join(stagingDir, "ephemeral")}, ttl, nil)
|
||||
}
|
||||
|
||||
// CleanEphemeralRoots removes stale request directories from explicit
|
||||
// ephemeral roots. WalkDir never follows directory symlinks.
|
||||
func CleanEphemeralRoots(roots []string, ttl time.Duration, guard *EphemeralCapacityGuard) {
|
||||
for _, root := range roots {
|
||||
cleanEphemeralRoot(root, ttl, guard)
|
||||
}
|
||||
}
|
||||
|
||||
func cleanEphemeralRoot(root string, ttl time.Duration, guard *EphemeralCapacityGuard) {
|
||||
categories, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
// A worker that has never served a file-bearing request has no
|
||||
@@ -87,19 +106,31 @@ func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
|
||||
continue
|
||||
}
|
||||
for _, entry := range entries {
|
||||
path := filepath.Join(categoryDir, entry.Name())
|
||||
info, err := entry.Info()
|
||||
if !entry.IsDir() || entry.Type()&os.ModeSymlink != 0 {
|
||||
continue
|
||||
}
|
||||
requestPath := filepath.Join(categoryDir, entry.Name())
|
||||
newest, err := newestEphemeralModTime(requestPath)
|
||||
if err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot stat entry", "path", path, "error", err)
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot inspect request", "path", requestPath, "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) {
|
||||
if !newest.Before(cutoff) {
|
||||
continue
|
||||
}
|
||||
if err := os.RemoveAll(path); err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", path, "error", err)
|
||||
if guard != nil {
|
||||
removedTree, err := guard.RemoveTreeIfInactive(requestPath, func() error {
|
||||
return os.RemoveAll(requestPath)
|
||||
})
|
||||
if err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", requestPath, "error", err)
|
||||
continue
|
||||
}
|
||||
if !removedTree {
|
||||
continue
|
||||
}
|
||||
} else if err := os.RemoveAll(requestPath); err != nil {
|
||||
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", requestPath, "error", err)
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
@@ -110,3 +141,24 @@ func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
|
||||
xlog.Info("Ephemeral staging cleanup removed stale request files", "count", removed, "dir", root)
|
||||
}
|
||||
}
|
||||
|
||||
func newestEphemeralModTime(root string) (time.Time, error) {
|
||||
var newest time.Time
|
||||
err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.ModTime().After(newest) {
|
||||
newest = info.ModTime()
|
||||
}
|
||||
if entry.Type()&os.ModeSymlink != 0 && entry.IsDir() {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return newest, err
|
||||
}
|
||||
@@ -55,4 +55,60 @@ var _ = Describe("Worker ephemeral staging cleanup", func() {
|
||||
It("does nothing when no ephemeral directory exists", func() {
|
||||
Expect(func() { CleanEphemeralStaging(stagingDir, time.Hour) }).ToNot(Panic())
|
||||
})
|
||||
|
||||
It("sweeps both transport roots by newest descendant and skips active requests", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
httpRoot := filepath.Join(stagingDir, "ephemeral")
|
||||
s3Root := filepath.Join(cacheDir, "ephemeral")
|
||||
guard, err := NewEphemeralCapacityGuard([]string{httpRoot, s3Root}, 8, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
staleRequest := filepath.Join(httpRoot, "audio", "stale")
|
||||
activeRequest := filepath.Join(s3Root, "audio", "active")
|
||||
freshChildRequest := filepath.Join(s3Root, "audio", "fresh-child")
|
||||
for _, requestDir := range []string{staleRequest, activeRequest, freshChildRequest} {
|
||||
Expect(os.MkdirAll(requestDir, 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(filepath.Join(requestDir, "input.bin"), []byte("data"), 0o600)).To(Succeed())
|
||||
}
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
fresh := time.Now().Add(-5 * time.Minute)
|
||||
for _, requestDir := range []string{staleRequest, activeRequest, freshChildRequest} {
|
||||
Expect(os.Chtimes(requestDir, old, old)).To(Succeed())
|
||||
}
|
||||
Expect(os.Chtimes(filepath.Join(staleRequest, "input.bin"), old, old)).To(Succeed())
|
||||
Expect(os.Chtimes(filepath.Join(activeRequest, "input.bin"), old, old)).To(Succeed())
|
||||
Expect(os.Chtimes(filepath.Join(freshChildRequest, "input.bin"), fresh, fresh)).To(Succeed())
|
||||
Expect(guard.Account(filepath.Join(staleRequest, "input.bin"), 4)).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(activeRequest, "input.bin"), 4)).To(Succeed())
|
||||
|
||||
CleanEphemeralRoots([]string{httpRoot, s3Root}, time.Hour, guard)
|
||||
|
||||
Expect(staleRequest).NotTo(BeADirectory())
|
||||
Expect(activeRequest).To(BeADirectory())
|
||||
Expect(freshChildRequest).To(BeADirectory())
|
||||
Expect(guard.Reserve(filepath.Join(httpRoot, "audio", "replacement", "input.bin"), 4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("keeps committed request inputs until exact release ends ownership", func() {
|
||||
root := filepath.Join(stagingDir, "ephemeral")
|
||||
requestDir := filepath.Join(root, "audio", "owned")
|
||||
path := filepath.Join(requestDir, "input.bin")
|
||||
Expect(os.MkdirAll(requestDir, 0o750)).To(Succeed())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 8, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(guard.Reserve(path, 4)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
|
||||
Expect(guard.Commit(path)).To(Succeed())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
Expect(os.Chtimes(path, old, old)).To(Succeed())
|
||||
Expect(os.Chtimes(requestDir, old, old)).To(Succeed())
|
||||
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(requestDir).To(BeADirectory())
|
||||
|
||||
Expect(guard.Release(path)).To(Succeed())
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(requestDir).NotTo(BeADirectory())
|
||||
})
|
||||
})
|
||||
@@ -3,14 +3,19 @@ package worker
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/pkg/safefile"
|
||||
"github.com/mudler/xlog"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
// isPathAllowed checks if path is within one of the allowed directories.
|
||||
@@ -37,7 +42,7 @@ func isPathAllowed(path string, allowedDirs []string) bool {
|
||||
}
|
||||
|
||||
// subscribeFileStaging subscribes to NATS file staging subjects for this node.
|
||||
func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, nodeID string) error {
|
||||
func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, nodeID string, capacity *EphemeralCapacityGuard) error {
|
||||
// Create FileManager with same S3 config as the frontend
|
||||
// TODO: propagate a caller-provided context once Config carries one
|
||||
s3Store, err := storage.NewS3Store(context.Background(), storage.S3Config{
|
||||
@@ -57,6 +62,10 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing file manager: %w", err)
|
||||
}
|
||||
if err := subscribeFileReleaseWithCapacity(natsClient, nodeID, fm, cacheDir, capacity); err != nil {
|
||||
return err
|
||||
}
|
||||
var ensureGroup singleflight.Group
|
||||
|
||||
// Subscribe: files.ensure — download S3 key to local, reply with local path
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) {
|
||||
@@ -68,12 +77,19 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
|
||||
return
|
||||
}
|
||||
|
||||
localPath, err := fm.Download(context.Background(), req.Key)
|
||||
value, err, _ := ensureGroup.Do(req.Key, func() (any, error) {
|
||||
return ensureWorkerFile(context.Background(), fm, capacity, req.Key)
|
||||
})
|
||||
if err != nil {
|
||||
xlog.Error("File ensure failed", "key", req.Key, "error", err)
|
||||
replyJSON(reply, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
localPath, ok := value.(string)
|
||||
if !ok {
|
||||
replyJSON(reply, map[string]string{"error": fmt.Sprintf("unexpected file ensure result %T", value)})
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("File ensured locally", "key", req.Key, "path", localPath)
|
||||
replyJSON(reply, map[string]string{"local_path": localPath})
|
||||
@@ -199,3 +215,220 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
|
||||
xlog.Info("Subscribed to file staging NATS subjects", "nodeID", nodeID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func subscribeFileRelease(natsClient messaging.MessagingClient, nodeID string, fm *storage.FileManager, cacheDir string) error {
|
||||
return subscribeFileReleaseWithCapacity(natsClient, nodeID, fm, cacheDir, nil)
|
||||
}
|
||||
|
||||
func subscribeFileReleaseWithCapacity(natsClient messaging.MessagingClient, nodeID string, fm *storage.FileManager, cacheDir string, capacity *EphemeralCapacityGuard) error {
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesRelease(nodeID), func(data []byte, reply func([]byte)) {
|
||||
var req struct {
|
||||
Key string `json:"key"`
|
||||
RequestID string `json:"request_id"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": "invalid request"})
|
||||
return
|
||||
}
|
||||
var err error
|
||||
if req.RequestID != "" {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
err = releaseEphemeralCacheRequest(ctx, cacheDir, req.RequestID, capacity)
|
||||
cancel()
|
||||
} else {
|
||||
cachePath, cacheErr := fm.CachePath(req.Key)
|
||||
err = cacheErr
|
||||
if err == nil {
|
||||
err = releaseEphemeralCachePathWithCapacity(cacheDir, req.Key, cachePath, capacity)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
replyJSON(reply, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
replyJSON(reply, map[string]string{})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.release events: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func releaseEphemeralCacheKey(cacheDir, key string) error {
|
||||
return releaseEphemeralCachePath(cacheDir, key, filepath.Join(cacheDir, filepath.FromSlash(key)))
|
||||
}
|
||||
|
||||
func releaseEphemeralCachePath(cacheDir, key, filePath string) error {
|
||||
return releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath, nil)
|
||||
}
|
||||
|
||||
func releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath string, capacity *EphemeralCapacityGuard) error {
|
||||
if err := validateEphemeralCacheKey(key); err != nil {
|
||||
return err
|
||||
}
|
||||
relativePath := filepath.FromSlash(key)
|
||||
expectedPath := filepath.Join(cacheDir, relativePath)
|
||||
if filepath.Clean(filePath) != expectedPath {
|
||||
return fmt.Errorf("release path %q does not match key %q", filePath, key)
|
||||
}
|
||||
if err := safefile.RemoveExact(cacheDir, relativePath, []string{".sha256", ".sha256.target"}, 2); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, path := range []string{filePath, filePath + ".sha256", filePath + ".sha256.target"} {
|
||||
if capacity != nil {
|
||||
if err := capacity.Release(path); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func releaseEphemeralCacheRequest(ctx context.Context, cacheDir, requestID string, capacity *EphemeralCapacityGuard) error {
|
||||
if err := validateEphemeralCacheRequestID(requestID); err != nil {
|
||||
return err
|
||||
}
|
||||
if capacity != nil {
|
||||
if err := capacity.BeginRequestRelease(ctx, requestID); err != nil {
|
||||
return fmt.Errorf("beginning release for request %q: %w", requestID, err)
|
||||
}
|
||||
defer capacity.EndRequestRelease(requestID)
|
||||
}
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
categories, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
var releaseErrors []error
|
||||
for _, category := range categories {
|
||||
if !category.IsDir() || category.Type()&os.ModeSymlink != 0 {
|
||||
continue
|
||||
}
|
||||
requestDir := filepath.Join(root, category.Name(), requestID)
|
||||
info, err := os.Lstat(requestDir)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("stating request directory %q: %w", requestDir, err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("ephemeral request path %q is not a real directory", requestDir))
|
||||
continue
|
||||
}
|
||||
entries, err := os.ReadDir(requestDir)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("reading request directory %q: %w", requestDir, err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("unexpected directory in ephemeral request %q", filepath.Join(requestDir, entry.Name())))
|
||||
continue
|
||||
}
|
||||
key := filepath.ToSlash(filepath.Join("ephemeral", category.Name(), requestID, entry.Name()))
|
||||
filePath := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
if err := releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath, capacity); err != nil {
|
||||
releaseErrors = append(releaseErrors, fmt.Errorf("releasing %q: %w", key, err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return errors.Join(releaseErrors...)
|
||||
}
|
||||
|
||||
type ephemeralStagingCapacity interface {
|
||||
Reserve(path string, size int64) error
|
||||
Commit(path string) error
|
||||
Claim(path string) error
|
||||
Release(path string) error
|
||||
}
|
||||
|
||||
func ensureWorkerFile(ctx context.Context, fm *storage.FileManager, capacity *EphemeralCapacityGuard, key string) (string, error) {
|
||||
if capacity == nil {
|
||||
return fm.Download(ctx, key)
|
||||
}
|
||||
if strings.HasPrefix(key, "ephemeral/") {
|
||||
if err := validateEphemeralCacheKey(key); err != nil {
|
||||
return "", err
|
||||
}
|
||||
requestID := strings.Split(key, "/")[2]
|
||||
if err := capacity.BeginRequestOperation(requestID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer capacity.EndRequestOperation(requestID)
|
||||
}
|
||||
return ensureWorkerFileWithCapacity(ctx, fm, capacity, key)
|
||||
}
|
||||
|
||||
func ensureWorkerFileWithCapacity(ctx context.Context, fm *storage.FileManager, capacity ephemeralStagingCapacity, key string) (string, error) {
|
||||
if capacity == nil || !strings.HasPrefix(key, "ephemeral/") {
|
||||
return fm.Download(ctx, key)
|
||||
}
|
||||
if err := validateEphemeralCacheKey(key); err != nil {
|
||||
return "", err
|
||||
}
|
||||
cachePath, err := fm.CachePath(key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if info, statErr := os.Lstat(cachePath); statErr == nil {
|
||||
if !info.Mode().IsRegular() {
|
||||
return "", fmt.Errorf("ephemeral cache path %q is not a regular file", cachePath)
|
||||
}
|
||||
if err := capacity.Claim(cachePath); err != nil {
|
||||
if !errors.Is(err, os.ErrNotExist) {
|
||||
return "", err
|
||||
}
|
||||
} else {
|
||||
return cachePath, nil
|
||||
}
|
||||
} else if !os.IsNotExist(statErr) {
|
||||
return "", statErr
|
||||
}
|
||||
|
||||
meta, err := fm.Head(ctx, key)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading size for %s: %w", key, err)
|
||||
}
|
||||
if err := capacity.Reserve(cachePath, meta.Size); err != nil {
|
||||
return "", err
|
||||
}
|
||||
localPath, err := fm.Download(ctx, key)
|
||||
if err != nil {
|
||||
_ = capacity.Release(cachePath)
|
||||
return "", err
|
||||
}
|
||||
if err := capacity.Commit(cachePath); err != nil {
|
||||
_ = fm.EvictCache(key)
|
||||
_ = capacity.Release(cachePath)
|
||||
return "", err
|
||||
}
|
||||
return localPath, nil
|
||||
}
|
||||
|
||||
func validateEphemeralCacheKey(key string) error {
|
||||
if strings.Contains(key, "\\") || path.Clean(key) != key {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
parts := strings.Split(key, "/")
|
||||
if len(parts) != 4 || parts[0] != "ephemeral" {
|
||||
return fmt.Errorf("release key %q must identify one file below ephemeral/", key)
|
||||
}
|
||||
for _, part := range parts[1:] {
|
||||
if part == "" || part == "." || part == ".." {
|
||||
return fmt.Errorf("invalid ephemeral key %q", key)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateEphemeralCacheRequestID(requestID string) error {
|
||||
if requestID == "" || strings.ContainsAny(requestID, "/\\") || path.Clean(requestID) != requestID || requestID == "." || requestID == ".." {
|
||||
return fmt.Errorf("invalid ephemeral request ID %q", requestID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type stagingObjectStore struct {
|
||||
payload []byte
|
||||
getCalls int
|
||||
getErr error
|
||||
}
|
||||
|
||||
type disappearingStagingCapacity struct{}
|
||||
|
||||
func (*disappearingStagingCapacity) Reserve(string, int64) error { return nil }
|
||||
func (*disappearingStagingCapacity) Commit(string) error { return nil }
|
||||
func (*disappearingStagingCapacity) Release(string) error { return nil }
|
||||
func (*disappearingStagingCapacity) Claim(path string) error {
|
||||
if err := os.Remove(path); err != nil {
|
||||
return err
|
||||
}
|
||||
return fmt.Errorf("claim raced recovery: %w", os.ErrNotExist)
|
||||
}
|
||||
|
||||
func (*stagingObjectStore) Put(context.Context, string, io.Reader) error { return nil }
|
||||
func (s *stagingObjectStore) Get(context.Context, string) (io.ReadCloser, error) {
|
||||
s.getCalls++
|
||||
if s.getErr != nil {
|
||||
return nil, s.getErr
|
||||
}
|
||||
return io.NopCloser(strings.NewReader(string(s.payload))), nil
|
||||
}
|
||||
func (s *stagingObjectStore) Head(_ context.Context, key string) (*storage.ObjectMeta, error) {
|
||||
return &storage.ObjectMeta{Key: key, Size: int64(len(s.payload))}, nil
|
||||
}
|
||||
func (*stagingObjectStore) Exists(context.Context, string) (bool, error) { return true, nil }
|
||||
func (*stagingObjectStore) Delete(context.Context, string) error { return nil }
|
||||
func (*stagingObjectStore) List(context.Context, string) ([]string, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
type releaseSubscription struct{}
|
||||
|
||||
func (releaseSubscription) Unsubscribe() error { return nil }
|
||||
|
||||
type releaseMessagingClient struct {
|
||||
subject string
|
||||
handler func([]byte, func([]byte))
|
||||
}
|
||||
|
||||
func (m *releaseMessagingClient) Publish(string, any) error { return nil }
|
||||
func (m *releaseMessagingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) SubscribeReply(subject string, handler func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
m.subject = subject
|
||||
m.handler = handler
|
||||
return releaseSubscription{}, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) Request(string, []byte, time.Duration) ([]byte, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *releaseMessagingClient) IsConnected() bool { return true }
|
||||
func (m *releaseMessagingClient) Close() {}
|
||||
|
||||
var _ = Describe("Worker exact-key staging release", func() {
|
||||
It("protects a startup-accounted HTTP cache hit through authenticated repeated probes", func() {
|
||||
stagingDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(stagingDir, "ephemeral")
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
remotePath := filepath.Join(stagingDir, filepath.FromSlash(key))
|
||||
content := []byte("data")
|
||||
Expect(os.MkdirAll(filepath.Dir(remotePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(remotePath, content, 0o600)).To(Succeed())
|
||||
hash := sha256.Sum256(content)
|
||||
Expect(os.WriteFile(remotePath+".sha256", []byte(fmt.Sprintf("%x", hash)), 0o600)).To(Succeed())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
for _, path := range []string{remotePath, remotePath + ".sha256", filepath.Dir(remotePath)} {
|
||||
Expect(os.Chtimes(path, old, old)).To(Succeed())
|
||||
}
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, int64(len(content)+sha256.Size*2), 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
addr := listener.Addr().String()
|
||||
Expect(listener.Close()).To(Succeed())
|
||||
server, err := nodes.StartFileTransferServerWithCapacity(addr, stagingDir, GinkgoT().TempDir(), GinkgoT().TempDir(), "secret", 0, nil, guard)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(nodes.ShutdownFileTransferServer, server)
|
||||
|
||||
localPath := filepath.Join(GinkgoT().TempDir(), "input.wav")
|
||||
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
|
||||
stager := nodes.NewHTTPFileStager(func(string) (string, error) { return addr, nil }, "secret")
|
||||
for range 2 {
|
||||
path, ensureErr := stager.EnsureRemote(context.Background(), "worker", localPath, key)
|
||||
Expect(ensureErr).NotTo(HaveOccurred())
|
||||
Expect(path).To(Equal(remotePath))
|
||||
}
|
||||
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(remotePath).To(BeAnExistingFile())
|
||||
|
||||
Expect(stager.ReleaseRemote(context.Background(), "worker", key)).To(Succeed())
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(remotePath).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("claims a startup-scanned cache hit against stale recovery until release", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(cachePath, []byte("data"), 0o600)).To(Succeed())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
Expect(os.Chtimes(cachePath, old, old)).To(Succeed())
|
||||
Expect(os.Chtimes(filepath.Dir(cachePath), old, old)).To(Succeed())
|
||||
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
store := &stagingObjectStore{payload: []byte("unused")}
|
||||
fm, err := storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
localPath, err := ensureWorkerFile(context.Background(), fm, guard, key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(localPath).To(Equal(cachePath))
|
||||
Expect(store.getCalls).To(BeZero())
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(cachePath).To(BeAnExistingFile())
|
||||
|
||||
Expect(guard.Release(cachePath)).To(Succeed())
|
||||
CleanEphemeralRoots([]string{root}, time.Hour, guard)
|
||||
Expect(cachePath).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("downloads again when a cache file disappears while being claimed", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(cachePath, []byte("stale"), 0o600)).To(Succeed())
|
||||
store := &stagingObjectStore{payload: []byte("fresh")}
|
||||
fm, err := storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
localPath, err := ensureWorkerFileWithCapacity(context.Background(), fm, &disappearingStagingCapacity{}, key)
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(localPath).To(Equal(cachePath))
|
||||
Expect(os.ReadFile(localPath)).To(Equal([]byte("fresh")))
|
||||
Expect(store.getCalls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("makes repeated cache-hit claims idempotent", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(cachePath, []byte("data"), 0o600)).To(Succeed())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(&stagingObjectStore{}, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
for range 2 {
|
||||
localPath, ensureErr := ensureWorkerFile(context.Background(), fm, guard, key)
|
||||
Expect(ensureErr).NotTo(HaveOccurred())
|
||||
Expect(localPath).To(Equal(cachePath))
|
||||
}
|
||||
err = guard.Reserve(filepath.Join(root, "other", "request-id", "input.wav"), 1)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(4)))
|
||||
})
|
||||
|
||||
It("capacity-checks growth of a startup-scanned cache file", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(cachePath, []byte("12"), 0o600)).To(Succeed())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(cachePath, []byte("12345"), 0o600)).To(Succeed())
|
||||
fm, err := storage.NewFileManager(&stagingObjectStore{}, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
_, err = ensureWorkerFile(context.Background(), fm, guard, key)
|
||||
var capacityErr *EphemeralCapacityError
|
||||
Expect(errors.As(err, &capacityErr)).To(BeTrue())
|
||||
Expect(capacityErr.RequestedBytes).To(Equal(int64(3)))
|
||||
Expect(capacityErr.UsageBytes).To(Equal(int64(2)))
|
||||
Expect(guard.HasActiveReservation(cachePath)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reserves S3 object size before download and releases it with the exact key", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
store := &stagingObjectStore{payload: []byte("data")}
|
||||
fm, err := storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
key := "ephemeral/audio/request-id/input.wav"
|
||||
|
||||
localPath, err := ensureWorkerFile(context.Background(), fm, guard, key)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(localPath).To(BeAnExistingFile())
|
||||
Expect(store.getCalls).To(Equal(1))
|
||||
Expect(guard.Reserve(filepath.Join(root, "audio", "other", "input.wav"), 1)).NotTo(Succeed())
|
||||
|
||||
Expect(releaseEphemeralCachePathWithCapacity(cacheDir, key, localPath, guard)).To(Succeed())
|
||||
Expect(guard.Reserve(filepath.Join(root, "audio", "other", "input.wav"), 4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("rejects an oversized S3 object before starting its download", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
store := &stagingObjectStore{payload: []byte("oversized")}
|
||||
fm, err := storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{filepath.Join(cacheDir, "ephemeral")}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
_, err = ensureWorkerFile(context.Background(), fm, guard, "ephemeral/audio/request-id/input.wav")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(store.getCalls).To(BeZero())
|
||||
})
|
||||
|
||||
It("rolls back an S3 reservation when the download fails", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
root := filepath.Join(cacheDir, "ephemeral")
|
||||
store := &stagingObjectStore{payload: []byte("data"), getErr: errors.New("download failed")}
|
||||
fm, err := storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
_, err = ensureWorkerFile(context.Background(), fm, guard, "ephemeral/audio/request-id/input.wav")
|
||||
Expect(err).To(MatchError(ContainSubstring("download failed")))
|
||||
Expect(guard.Reserve(filepath.Join(root, "audio", "replacement", "input.wav"), 4)).To(Succeed())
|
||||
})
|
||||
|
||||
It("removes only the exact cache file and upload sidecars", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
target := filepath.Join(categoryDir, "input.wav")
|
||||
sibling := filepath.Join(categoryDir, "keep.wav")
|
||||
for _, path := range []string{target, target + ".sha256", target + ".sha256.target", sibling} {
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/input.wav")).To(Succeed())
|
||||
Expect(target).NotTo(BeAnExistingFile())
|
||||
Expect(target + ".sha256").NotTo(BeAnExistingFile())
|
||||
Expect(target + ".sha256.target").NotTo(BeAnExistingFile())
|
||||
Expect(sibling).To(BeAnExistingFile())
|
||||
Expect(categoryDir).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("succeeds for a missing file and prunes empty category and request directories", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
|
||||
for range 2 {
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/missing.wav")).To(Succeed())
|
||||
}
|
||||
Expect(categoryDir).NotTo(BeADirectory())
|
||||
Expect(filepath.Dir(categoryDir)).NotTo(BeADirectory())
|
||||
Expect(filepath.Join(cacheDir, "ephemeral")).To(BeADirectory())
|
||||
})
|
||||
|
||||
It("rejects traversal and symlink escapes", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
outsideDir := GinkgoT().TempDir()
|
||||
outsidePath := filepath.Join(outsideDir, "input.wav")
|
||||
Expect(os.WriteFile(outsidePath, []byte("keep"), 0640)).To(Succeed())
|
||||
requestDir := filepath.Join(cacheDir, "ephemeral", "request-id")
|
||||
Expect(os.MkdirAll(requestDir, 0750)).To(Succeed())
|
||||
Expect(os.Symlink(outsideDir, filepath.Join(requestDir, "audio"))).To(Succeed())
|
||||
|
||||
for _, key := range []string{
|
||||
"models/model.gguf",
|
||||
"ephemeral/../models/model.gguf",
|
||||
"ephemeral/request-id/audio/../../model.gguf",
|
||||
"ephemeral/request-id/audio/input.wav",
|
||||
} {
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, key)).NotTo(Succeed(), key)
|
||||
}
|
||||
Expect(outsidePath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("rejects symlinked files and sidecars without deleting their targets", func() {
|
||||
for _, linkedName := range []string{"input.wav", "input.wav.sha256", "input.wav.sha256.target"} {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
|
||||
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
|
||||
target := filepath.Join(categoryDir, "input.wav")
|
||||
if linkedName != "input.wav" {
|
||||
Expect(os.WriteFile(target, []byte("input"), 0640)).To(Succeed())
|
||||
}
|
||||
preserved := filepath.Join(cacheDir, "ephemeral", "preserved-"+linkedName)
|
||||
Expect(os.WriteFile(preserved, []byte("keep"), 0640)).To(Succeed())
|
||||
Expect(os.Symlink(preserved, filepath.Join(categoryDir, linkedName))).To(Succeed())
|
||||
|
||||
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/input.wav")).NotTo(Succeed(), linkedName)
|
||||
Expect(preserved).To(BeAnExistingFile(), linkedName)
|
||||
}
|
||||
})
|
||||
|
||||
It("registers an exact release handler", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
path := filepath.Join(cacheDir, "ephemeral", "request-id", "audio", "input.wav")
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
fm, err := storage.NewFileManager(nil, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseMessagingClient{}
|
||||
|
||||
Expect(subscribeFileRelease(client, "node.one", fm, cacheDir)).To(Succeed())
|
||||
Expect(client.subject).To(Equal(messaging.SubjectNodeFilesRelease("node.one")))
|
||||
request, err := json.Marshal(map[string]string{"key": "ephemeral/request-id/audio/input.wav"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var response []byte
|
||||
client.handler(request, func(data []byte) { response = append([]byte(nil), data...) })
|
||||
|
||||
var reply map[string]string
|
||||
Expect(json.Unmarshal(response, &reply)).To(Succeed())
|
||||
Expect(reply["error"]).To(BeEmpty())
|
||||
Expect(path).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("releases a request batch through one worker message", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
keys := []string{
|
||||
"ephemeral/audio/request-id/input.wav",
|
||||
"ephemeral/images/request-id/frame.jpg",
|
||||
}
|
||||
for _, key := range keys {
|
||||
path := filepath.Join(cacheDir, filepath.FromSlash(key))
|
||||
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
|
||||
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
|
||||
}
|
||||
fm, err := storage.NewFileManager(nil, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseMessagingClient{}
|
||||
Expect(subscribeFileRelease(client, "node.one", fm, cacheDir)).To(Succeed())
|
||||
request, err := json.Marshal(map[string]any{"request_id": "request-id"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var response []byte
|
||||
|
||||
client.handler(request, func(data []byte) { response = append([]byte(nil), data...) })
|
||||
|
||||
var reply map[string]string
|
||||
Expect(json.Unmarshal(response, &reply)).To(Succeed())
|
||||
Expect(reply["error"]).To(BeEmpty())
|
||||
for _, key := range keys {
|
||||
Expect(filepath.Join(cacheDir, filepath.FromSlash(key))).NotTo(BeAnExistingFile())
|
||||
}
|
||||
})
|
||||
|
||||
It("returns validation errors through the release handler", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
fm, err := storage.NewFileManager(nil, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
client := &releaseMessagingClient{}
|
||||
Expect(subscribeFileRelease(client, "node-1", fm, cacheDir)).To(Succeed())
|
||||
|
||||
request, err := json.Marshal(map[string]string{"key": "models/model.gguf"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var response []byte
|
||||
client.handler(request, func(data []byte) { response = append([]byte(nil), data...) })
|
||||
|
||||
var reply map[string]string
|
||||
Expect(json.Unmarshal(response, &reply)).To(Succeed())
|
||||
Expect(reply["error"]).NotTo(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -597,6 +597,7 @@ func (s *backendSupervisor) reapDeadProcess(key string, bp *backendProcess) {
|
||||
if bp == nil {
|
||||
return
|
||||
}
|
||||
s.cleanupProcessRuntime(bp.proc)
|
||||
if bp.port <= 0 {
|
||||
xlog.Error("Cannot recycle backend port: dead process has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port)
|
||||
return
|
||||
@@ -614,6 +615,7 @@ func (s *backendSupervisor) releaseBackendStart(key string, bp *backendProcess)
|
||||
return
|
||||
}
|
||||
delete(s.processes, key)
|
||||
s.cleanupProcessRuntime(bp.proc)
|
||||
if bp.port <= 0 {
|
||||
xlog.Error("Cannot recycle backend port: startup has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port)
|
||||
return
|
||||
@@ -947,6 +949,7 @@ func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, st
|
||||
return fmt.Errorf("stopping backend process %s: %w", key, stopErr)
|
||||
}
|
||||
delete(s.processes, key)
|
||||
s.cleanupProcessRuntime(bp.proc)
|
||||
if bp.port <= 0 {
|
||||
xlog.Error("Cannot recycle backend port: process has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port)
|
||||
return nil
|
||||
@@ -955,6 +958,14 @@ func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, st
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *backendSupervisor) cleanupProcessRuntime(proc *process.Process) {
|
||||
// Some focused supervisor tests provide synthetic process handles without a
|
||||
// ModelLoader. Production processes always come from s.ml.StartProcess.
|
||||
if s.ml != nil {
|
||||
s.ml.CleanupProcessRuntime(proc)
|
||||
}
|
||||
}
|
||||
|
||||
// stopAllBackends stops all running backend processes and returns the process
|
||||
// keys it attempted, so a caller answering a backend.stop request can report
|
||||
// what it acted on.
|
||||
|
||||
@@ -149,22 +149,37 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
// the top of Run so the worker fails before registering.)
|
||||
httpAddr := cfg.resolveHTTPAddr()
|
||||
stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging")
|
||||
cacheDir := filepath.Join(cfg.ModelsPath, "..", "cache")
|
||||
dataDir := filepath.Join(cfg.ModelsPath, "..", "data")
|
||||
ephemeralRoots := []string{
|
||||
filepath.Join(stagingDir, "ephemeral"),
|
||||
filepath.Join(cacheDir, "ephemeral"),
|
||||
}
|
||||
byteLimit, minFreeBytes, err := effectiveEphemeralCapacity(
|
||||
ephemeralRoots,
|
||||
cfg.EphemeralStagingByteLimit,
|
||||
cfg.EphemeralStagingMinFreeBytes,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resolving ephemeral staging capacity: %w", err)
|
||||
}
|
||||
ephemeralCapacity, err := NewEphemeralCapacityGuard(ephemeralRoots, byteLimit, minFreeBytes)
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing ephemeral staging capacity: %w", err)
|
||||
}
|
||||
xlog.Info("Ephemeral staging capacity configured", "roots", ephemeralRoots, "byteLimit", byteLimit, "minFreeBytes", minFreeBytes)
|
||||
StartEphemeralRootsCleanup(shutdownCtx, ephemeralRoots, ephemeralCapacity, 0, 0)
|
||||
// The readiness gate is created here but only armed once NATS is up and the
|
||||
// backend supervisor exists, below, because the gate probes both.
|
||||
// Until then /readyz reports ready, which is correct: reaching this line
|
||||
// means the worker has already registered with the frontend, so it is
|
||||
// mid-startup rather than broken.
|
||||
readiness := &nodes.WorkerReadiness{}
|
||||
httpServer, err := nodes.StartFileTransferServer(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ml.BackendLogs())
|
||||
httpServer, err := nodes.StartFileTransferServerWithCapacity(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ephemeralCapacity, ml.BackendLogs())
|
||||
if err != nil {
|
||||
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()
|
||||
@@ -249,7 +264,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
|
||||
// Subscribe to file staging NATS subjects if S3 is configured
|
||||
if cfg.StorageURL != "" {
|
||||
if err := cfg.subscribeFileStaging(natsClient, nodeID); err != nil {
|
||||
if err := cfg.subscribeFileStaging(natsClient, nodeID, ephemeralCapacity); err != nil {
|
||||
nodes.ShutdownFileTransferServer(httpServer)
|
||||
return fmt.Errorf("subscribing to file staging subjects: %w", err)
|
||||
}
|
||||
|
||||
@@ -103,8 +103,7 @@ Before diving into advanced topics, ensure you have:
|
||||
## Related Sections
|
||||
|
||||
- 📚 [Reference](../reference/) - API documentation and command reference
|
||||
- 🔌 [Installation](../installation/) - Deployment options and requirements
|
||||
- ⭐ [Features](../features/) - Overview of LocalAI capabilities
|
||||
- - ⭐ [Features](../features/) - Overview of LocalAI capabilities
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -397,7 +397,9 @@ The canonical names match upstream llama.cpp (dash-separated). For backward comp
|
||||
Multiple types can be chained by passing a comma-separated list to `spec_type` (e.g. `spec_type:ngram-simple,ngram-mod`). The runtime tries them in order and accepts the first proposal that meets the acceptance criteria.
|
||||
|
||||
{{% notice note %}}
|
||||
Speculative decoding is automatically disabled when multimodal models (with `mmproj`) are active. The `n_draft` parameter can also be overridden per-request.
|
||||
The current LocalAI llama.cpp backend supports speculative decoding with multimodal models that load an `mmproj`, including MTP. LocalAI passes both configurations to llama.cpp and does not disable speculation merely because an `mmproj` is present. Upstream llama.cpp removed the former general multimodal/speculative restriction in [ggml-org/llama.cpp#19493](https://github.com/ggml-org/llama.cpp/pull/19493); [ggml-org/llama.cpp#22673](https://github.com/ggml-org/llama.cpp/pull/22673) later added MTP support and explicitly documented its compatibility with vision input.
|
||||
|
||||
Compatibility still depends on the installed backend version and the target/draft model architecture. Check the backend logs for successful projector loading and speculative-context initialization, then look for the `draft acceptance` statistics line and its `accepted / generated` counts. A representative run with zero accepted draft tokens receives no speculative speedup and can indicate that the model or settings need tuning.
|
||||
{{% /notice %}}
|
||||
|
||||
##### Multi-Token Prediction (MTP)
|
||||
@@ -427,7 +429,7 @@ Detection runs both at **import time** (the `/import-model` UI / `POST /models/i
|
||||
| `spec_type` | `draft-mtp` | Activates MTP. Can be chained with other types (see below). |
|
||||
| `spec_n_max` / `draft_max` | `2`-`6` | Number of draft tokens per step. Upstream's PR suggests 2-3 for the tightest acceptance window; LocalAI's auto-default is 6 to favour throughput on models with high acceptance. |
|
||||
| `spec_p_min` | `0.75` | Pinned because upstream marks the current default with a "change to 0.0f" TODO; locking it here keeps acceptance thresholds stable across future llama.cpp bumps. |
|
||||
| `mmproj_use_gpu` | `false` (or unset `mmproj`) | MTP has a prompt-processing overhead; if the model is non-vision, drop the mmproj entirely to save VRAM. |
|
||||
| `mmproj_use_gpu` | `true` for vision | MTP does not require disabling the projector. Keep `mmproj` configured for image input; set this option to `false` to keep the projector on CPU when VRAM is tight. Remove `mmproj` only for text-only use when vision is not needed. |
|
||||
|
||||
**Minimal config** (override-only, since auto-detection already covers this for MTP-capable GGUFs):
|
||||
|
||||
@@ -441,6 +443,23 @@ options:
|
||||
- spec_n_max:3
|
||||
```
|
||||
|
||||
**With vision enabled:**
|
||||
|
||||
```yaml
|
||||
name: qwen3-vision-mtp
|
||||
backend: llama-cpp
|
||||
known_usecases:
|
||||
- chat
|
||||
- vision
|
||||
parameters:
|
||||
model: qwen3-with-mtp.gguf
|
||||
mmproj: mmproj-qwen3.gguf
|
||||
options:
|
||||
- spec_type:draft-mtp
|
||||
- spec_n_max:3
|
||||
- spec_p_min:0.75
|
||||
```
|
||||
|
||||
**With a separate MTP head file:**
|
||||
|
||||
```yaml
|
||||
@@ -720,7 +739,7 @@ For image generation models using the `diffusers` backend:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `diffusers.cuda` | bool | Enable CUDA for diffusers |
|
||||
| `diffusers.cuda` | bool | Force CUDA. By default the backend auto-detects and uses CUDA when a compatible GPU is present (ROCm builds included). Pin the CPU with `options: ["device:cpu"]` |
|
||||
| `diffusers.pipeline_type` | string | Pipeline type (e.g., `stable-diffusion`, `stable-diffusion-xl`) |
|
||||
| `diffusers.scheduler_type` | string | Scheduler type (e.g., `euler`, `ddpm`) |
|
||||
| `diffusers.enable_parameters` | string | Comma-separated parameters to enable |
|
||||
|
||||
@@ -288,6 +288,9 @@ All agent endpoints are grouped under `/api/agents/`:
|
||||
| `POST` | `/api/agents/collections/:name/upload` | Upload a document |
|
||||
| `GET` | `/api/agents/collections/:name/entries` | List entries |
|
||||
| `POST` | `/api/agents/collections/:name/search` | Search a collection |
|
||||
| `GET` | `/api/agents/collections/:name/sources` | List external sources |
|
||||
| `POST` | `/api/agents/collections/:name/sources` | Add an external source (`update_interval` is an integer number of minutes; defaults to 60) |
|
||||
| `DELETE` | `/api/agents/collections/:name/sources` | Remove an external source |
|
||||
| `POST` | `/api/agents/collections/:name/reset` | Reset a collection |
|
||||
|
||||
### Actions
|
||||
|
||||
@@ -11,6 +11,10 @@ LocalAI exposes this through the `/v1/audio/classification` endpoint, modelled a
|
||||
|
||||
Because classification is exposed as a regular OpenAI-style endpoint, any HTTP client works - there is no Python dependency on the consumer side.
|
||||
|
||||
In distributed mode, LocalAI stages uploaded audio and realtime sound-detection
|
||||
windows on the selected worker before classification. The API server and worker
|
||||
do not need a shared temporary directory.
|
||||
|
||||
## Endpoint
|
||||
|
||||
```
|
||||
|
||||
@@ -28,6 +28,11 @@ key stored *inside* the GGUF, which cannot be read from a remote repository, and
|
||||
upstream GGUF repository hosts every family in one place. Set `backend: audio-cpp`
|
||||
in the model YAML, or select it explicitly in the import form.
|
||||
|
||||
The bundled audio-cpp gallery entries set `backend:best` in `options` to select
|
||||
an available compute backend, with CPU as the fallback. To force CPU execution,
|
||||
replace that option with `backend:cpu`. Model configurations that omit this
|
||||
option still default to CPU.
|
||||
|
||||
## What it serves
|
||||
|
||||
One model serves one family, and a family advertises the tasks it can perform. The
|
||||
|
||||
@@ -11,6 +11,7 @@ The transcription endpoint allows to convert audio files to text. The endpoint s
|
||||
- **[whisper.cpp](https://github.com/ggerganov/whisper.cpp)**: A C++ library for audio transcription (default)
|
||||
- **moonshine**: Ultra-fast transcription engine optimized for low-end devices
|
||||
- **faster-whisper**: Fast Whisper implementation with CTranslate2
|
||||
- **WhisperX**: Whisper transcription with word alignment and optional speaker diarization. Set `HF_TOKEN` and pass `diarize=true` to load WhisperX's gated pyannote diarization pipeline.
|
||||
- **[parakeet-cpp](https://github.com/mudler/parakeet.cpp)**: A C++/ggml port of NVIDIA NeMo Parakeet (FastConformer TDT/CTC/RNNT/hybrid). Runs quantized GGUFs on CPU or GPU, emits word-level timestamps, and supports cache-aware streaming (the `realtime_eou` model surfaces end-of-utterance events).
|
||||
- **llama-cpp**: Route transcription to any multimodal-audio GGUF model served by the `llama-cpp` backend (e.g. [Qwen3-ASR](https://huggingface.co/ggml-org/Qwen3-ASR-0.6B-GGUF), Voxtral, Qwen2-Audio). Under the hood the request is converted into a chat completion with the audio attached via the model's audio encoder - the same path the upstream llama.cpp server uses. Set `backend: llama-cpp` in the model YAML and point `mmproj` at the matching audio encoder.
|
||||
- **voxtral**: Voxtral-family models served by a dedicated backend
|
||||
@@ -109,7 +110,7 @@ In addition to `file` and `model`, the endpoint accepts the following multipart
|
||||
| `timestamp_granularities[]` | Multi-value form field: `word` and/or `segment`. Honored when the backend produces the requested granularity. |
|
||||
| `response_format` | One of `json` (default for backwards-compat), `verbose_json`, `text`, `srt`, `vtt`, `lrc`. |
|
||||
| `stream` | When `true`, the endpoint emits an SSE stream of `transcript.text.delta` events followed by a final `transcript.text.done` event. |
|
||||
| `diarize` | LocalAI extension - speaker diarization (whisper.cpp only). |
|
||||
| `diarize` | LocalAI extension - speaker diarization. WhisperX requires `HF_TOKEN`; requests fail with `FailedPrecondition` when it is missing. |
|
||||
|
||||
The response body for `verbose_json` includes `text`, `language`, `duration`, and `segments[]` (with `speaker` populated when diarization is enabled).
|
||||
|
||||
|
||||
@@ -297,6 +297,8 @@ local-ai worker \
|
||||
| `--advertise-addr` | `LOCALAI_ADVERTISE_ADDR` | *(auto)* | Address the frontend uses to reach this node (see below) |
|
||||
| `--http-addr` | `LOCALAI_HTTP_ADDR` | gRPC port - 1 | HTTP file transfer server bind address |
|
||||
| `--advertise-http-addr` | `LOCALAI_ADVERTISE_HTTP_ADDR` | *(auto)* | HTTP address the frontend uses for file transfer |
|
||||
| `--ephemeral-staging-byte-limit` | `LOCALAI_EPHEMERAL_STAGING_BYTE_LIMIT` | `0` (automatic) | Maximum bytes held by request-input staging across the worker's HTTP staging directory and S3 cache. Automatic mode uses the smaller of 10 GiB and 10% of filesystem capacity. |
|
||||
| `--ephemeral-staging-min-free-bytes` | `LOCALAI_EPHEMERAL_STAGING_MIN_FREE_BYTES` | `0` (automatic) | Free filesystem space preserved while staging request inputs. Automatic mode uses the larger of 1 GiB and 5% of filesystem capacity. |
|
||||
| `--register-to` | `LOCALAI_REGISTER_TO` | *(required)* | Frontend URL for self-registration |
|
||||
| `--node-name` | `LOCALAI_NODE_NAME` | hostname | Human-readable node name |
|
||||
| `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token to authenticate with the frontend |
|
||||
@@ -320,6 +322,12 @@ local-ai worker \
|
||||
**HTTP file transfer:** Each worker also runs a small HTTP server for file transfer (model files, configs). By default it listens on the gRPC base port - 1 (e.g., if gRPC base is 50051, HTTP is on 50050). gRPC ports grow upward from the base port as additional models are loaded. Set `--advertise-http-addr` if the auto-detected address is not routable from the frontend.
|
||||
{{% /notice %}}
|
||||
|
||||
### Ephemeral request-input storage
|
||||
|
||||
Workers reserve local capacity before accepting per-request audio, image, and other ephemeral inputs. The limit covers both direct HTTP staging and the worker's S3 download cache. A request is rejected before inference when accepting its input would exceed the byte limit or the configured free-space headroom. One request-scoped cleanup operation releases all exact input keys and their reservations after inference, while a one-hour recovery sweep removes abandoned files after crashes. The sweep runs at startup and every 15 minutes, preserves active requests, and considers the newest file in each request directory.
|
||||
|
||||
Set both capacity variables to positive byte counts when a worker needs fixed limits. Leaving either value at zero selects its filesystem-based default. These settings apply only below the two `ephemeral` roots; model, data, and configuration files are excluded.
|
||||
|
||||
### Worker Health Probes
|
||||
|
||||
The worker's HTTP server (base port - 1, default 50050) exposes two unauthenticated probes:
|
||||
@@ -1205,9 +1213,17 @@ Notes:
|
||||
- Verify `--heartbeat-interval` is not set too high
|
||||
- Offline nodes automatically restore to healthy when they re-register (no re-approval needed)
|
||||
|
||||
**InsightFace reports a missing MiniFASNet file after staging:**
|
||||
- Gallery models such as `insightface-buffalo-m` use a virtual primary name and load their files through options. The frontend derives the worker's model directory from successfully staged companion files or directories, so relative options resolve inside the model's staging directory.
|
||||
- If logs show matching hashes for the staged files but InsightFace still reports a bare filename such as `MiniFASNetV2.onnx` as missing, upgrade the frontend to include this path-resolution fix. Re-uploading the same files does not correct the directory passed to the backend.
|
||||
|
||||
**Backend not installing:**
|
||||
- Check the worker logs for `backend.install` events
|
||||
|
||||
**Model staging repeatedly fails with HTTP 416 after all bytes have arrived:**
|
||||
- An interrupted upload can leave a full-size file marked as unfinished (`.sha256.target`). On retry, the worker verifies the file's SHA-256 and finalizes it if it matches, without rewriting the model. Corrupt content fails integrity validation and is removed.
|
||||
- Upgrade the affected worker to get this recovery behavior. Older workers can repeatedly reject retries from byte zero with `Content-Range start 0 does not match current file size`. File size alone is not proof that an upload is valid.
|
||||
|
||||
**Requests still report an old context size or another old load option:**
|
||||
- Query `/api/nodes/:id/models` for every worker that hosts the model.
|
||||
- Confirm that every routable replica has `state: loaded` and the same current `config_revision`.
|
||||
@@ -1230,8 +1246,9 @@ Notes:
|
||||
- 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.
|
||||
- A request that carries a file (an image, an audio clip, a video) stages that file below the worker's HTTP staging or S3 cache `ephemeral/` directory. The frontend releases each request-owned input when inference finishes, and the worker reserves capacity before accepting it.
|
||||
- A one-hour recovery sweep runs at startup and every 15 minutes to reclaim inputs left by interrupted requests. It preserves active reservations and uses the newest file timestamp in each request directory.
|
||||
- Releases before request-owned cleanup existed can leave a legacy backlog. Delete the affected `ephemeral/` directory once, as the user the worker runs as; capacity admission and recovery cleanup keep new staging bounded.
|
||||
- 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.
|
||||
|
||||
|
||||
@@ -73,9 +73,9 @@ Detect faces and analyze demographics (buffalo entries populate
|
||||
age / gender; YuNet + SFace returns regions only):
|
||||
|
||||
```bash
|
||||
curl -sX POST http://localhost:8080/v1/face/detect \
|
||||
curl -sX POST http://localhost:8080/v1/detection \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "face-detect-buffalo-l", "img": "https://example.com/group.jpg"}'
|
||||
-d '{"model": "face-detect-buffalo-l", "image": "https://example.com/group.jpg"}'
|
||||
|
||||
curl -sX POST http://localhost:8080/v1/face/analyze \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -141,6 +141,39 @@ Response:
|
||||
}
|
||||
```
|
||||
|
||||
## Restore enrollments after a restart
|
||||
|
||||
The default identity store is in memory. Clients can keep an enrollment record
|
||||
and replay it with `POST /v1/face/register` after a restart. Extract the embedding
|
||||
once with `/v1/face/embed`, then save the exact returned vector, model, name,
|
||||
labels, and enrollment timestamp. Submit `embedding` instead of `img`:
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "insightface-opencv",
|
||||
"name": "Alice",
|
||||
"embedding": [0.12, -0.04, 0.31],
|
||||
"registered_at": "2026-09-07T12:00:00Z",
|
||||
"labels": {"client_id": "alice"}
|
||||
}
|
||||
```
|
||||
|
||||
The vector above is abbreviated; send the complete embedding from the same
|
||||
recognizer model. Provide exactly one of `img` or `embedding`. Vectors must be
|
||||
finite and nonzero. `registered_at` is optional and defaults to the current time;
|
||||
replay the original timestamp to preserve it.
|
||||
|
||||
The store upserts by exact vector. Registration now derives a stable ID from
|
||||
that vector and the store namespace, so retries and replay after a restart
|
||||
return the same ID without adding duplicate entries. Replaying updates the name
|
||||
and labels. Images can produce slightly different embeddings across runs; keep
|
||||
the original vector instead of embedding the photo again on each retry.
|
||||
|
||||
This does not make the server store persistent. Clients must retain and restore
|
||||
the records themselves. With independent stores behind a load balancer, replay
|
||||
into each store or use a shared store. Do not mix different recognizer models in
|
||||
one store. IDs from older versions change on their first registration replay.
|
||||
|
||||
## 1:N identification workflow (register → identify → forget)
|
||||
|
||||
This is the primary "face recognition" flow. Under the hood it uses
|
||||
|
||||
@@ -79,8 +79,8 @@ When a model does not fit entirely in VRAM, the following `options:` control whe
|
||||
|--------|---------|-------------|
|
||||
| `backend` | `backend:clip=cpu,vae=cuda0,diffusion=vulkan0` | Runtime (compute) backend assignment per component. Use `cpu` to place a component's compute on the CPU. Component keys include `te` (text encoder / CLIP), `vae`, `diffusion`, `controlnet`. |
|
||||
| `params_backend` | `params_backend:diffusion=disk,clip=cpu` | Where parameters (weights) are stored. Supports `cpu`, `disk` (mmap weights from disk to save RAM/VRAM), or per-component specs. |
|
||||
| `max_vram` | `max_vram:8` or `max_vram:-1` | VRAM budget (in GiB) for graph-cut segmented parameter offload. `0` disables it, `-1` auto-selects (free VRAM minus ~1 GiB). Also accepts per-backend budgets. |
|
||||
| `stream_layers` | `stream_layers:true` | Enable residency + prefetch streaming on top of `max_vram` (no effect unless `max_vram` is set). |
|
||||
| `max_vram` | `max_vram:8` or `max_vram:-1` | Optional per-device VRAM budget (in GiB) for managed weights and automatic graph-cut execution. `0` uses live free VRAM without an explicit cap; a negative value reserves that many GiB of free VRAM. Also accepts per-backend budgets. |
|
||||
| `stream_layers` | `stream_layers:true` | Deprecated compatibility option. Segmented weight streaming is now selected automatically, so this value is ignored. |
|
||||
| `rpc_servers` | `rpc_servers:localhost:50052,192.168.1.3:50052` | Comma-separated list of `host:port` RPC servers to offload compute to. |
|
||||
| `pulid_weights_path` | `pulid_weights_path:pulid.safetensors` | Path to PuLID-Flux weights for identity injection. |
|
||||
|
||||
|
||||
Loaded 100 of 126 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user