mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-10 13:08:55 -04:00
Compare commits
104
Commits
No files matched your search
@@ -49,6 +49,40 @@ AI agents MUST NOT add `Co-Authored-By` trailers for themselves either.
|
||||
A human reviewer owns the contribution; the AI's involvement is recorded
|
||||
via `Assisted-by` (see below).
|
||||
|
||||
### Exception: automation operated by a maintainer
|
||||
|
||||
The rule above addresses the common case, an AI assistant helping a human
|
||||
contributor who then signs off. It does not fit automation that a
|
||||
maintainer runs themselves, which opens pull requests with no human
|
||||
submitter to sign. Applied literally there, nothing ever signs and the
|
||||
DCO check blocks the pull request permanently.
|
||||
|
||||
A maintainer-operated bot MUST therefore add a `Signed-off-by` trailer
|
||||
naming **the maintainer who operates it**, not the bot and not the model:
|
||||
|
||||
```
|
||||
Assisted-by: Codex:gpt-5
|
||||
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
|
||||
```
|
||||
|
||||
This is not the AI certifying the DCO. The maintainer is, exactly as they
|
||||
do for a commit they typed by hand: they configured the automation, they
|
||||
own its output, and they take responsibility for it when they merge it.
|
||||
The `Assisted-by` trailer still records that a model produced the code, so
|
||||
the provenance trail is unchanged.
|
||||
|
||||
The exception is narrow and does not widen the rule for anyone else:
|
||||
|
||||
- It applies only to automation a LocalAI maintainer operates and whose
|
||||
output that maintainer reviews before merge.
|
||||
- The sign-off names a real person who accepts DCO responsibility.
|
||||
- An AI assistant helping an outside contributor still MUST NOT sign off.
|
||||
That contributor adds their own trailer.
|
||||
- A bot MUST NOT sign off on behalf of anyone other than its operator, and
|
||||
MUST NOT add a trailer for a contributor whose branch it pushes to. If
|
||||
automation contributes to someone else's branch, it leaves the sign-off
|
||||
to that contributor.
|
||||
|
||||
## Attribution
|
||||
|
||||
When AI tools contribute to LocalAI development, proper attribution helps
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ Human contributors: see [CONTRIBUTING.md](CONTRIBUTING.md) for the development w
|
||||
|
||||
LocalAI follows the Linux kernel project's [guidelines for AI coding assistants](https://docs.kernel.org/process/coding-assistants.html). Before submitting AI-assisted code, read [.agents/ai-coding-assistants.md](.agents/ai-coding-assistants.md). Key rules:
|
||||
|
||||
- **No `Signed-off-by` from AI.** Only the human submitter may sign off on the Developer Certificate of Origin.
|
||||
- **No `Signed-off-by` from AI.** Only the human submitter may sign off on the Developer Certificate of Origin. One exception: automation a maintainer operates signs off with *that maintainer's* identity, since no other human submitter exists to certify it. See [.agents/ai-coding-assistants.md](.agents/ai-coding-assistants.md).
|
||||
- **No `Co-Authored-By: <AI>` trailers.** The human contributor owns the change.
|
||||
- **Use an `Assisted-by:` trailer** to attribute AI involvement. Format: `Assisted-by: AGENT_NAME:MODEL_VERSION [TOOL1] [TOOL2]`.
|
||||
- **The human submitter is responsible** for reviewing, testing, and understanding every line of generated code.
|
||||
|
||||
+1
-1
@@ -218,7 +218,7 @@ LocalAI follows the **same guidelines as the Linux kernel project** for AI-assis
|
||||
|
||||
The full policy for this repository lives in [`.agents/ai-coding-assistants.md`](.agents/ai-coding-assistants.md). Summary:
|
||||
|
||||
- **AI agents MUST NOT add `Signed-off-by` tags.** Only humans can certify the Developer Certificate of Origin.
|
||||
- **AI agents MUST NOT add `Signed-off-by` tags.** Only humans can certify the Developer Certificate of Origin. Automation operated by a maintainer is the one exception: it signs off with that maintainer's identity, because there is no other human submitter to certify it.
|
||||
- **AI agents MUST NOT add `Co-Authored-By` trailers** attributing themselves as co-authors.
|
||||
- **Attribute AI involvement with an `Assisted-by` trailer** in the commit message:
|
||||
|
||||
|
||||
@@ -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?=fa5aaac9266a98c68f8a5c9fcd1ba6ff65875416
|
||||
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?=6289c516273979173abbc062209a81dd3706b804
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=8db89fe083ae4d17c9a2428ccd29803d3ae8f577
|
||||
DS4_VERSION?=6289c516273979173abbc062209a81dd3706b804
|
||||
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?=3e416d7f5a9d4cc3195e8171dbf891541ca59c6a
|
||||
LLAMA_REPO?=https://github.com/ikawrakow/ik_llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=3466812d1f06728effe7c0f3c0671117f461672d
|
||||
LLAMA_VERSION?=434ddbbc0e30522e897670681e503b797c12b7c1
|
||||
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?=a5b6953c4a579a2bbd1c0913ad8a85c2a4d99953
|
||||
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?=c44b60b8053bbf2a5c1e014f11323fb3f2485177
|
||||
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
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
# on a cu130 host. Pull the cu130-flavoured wheel from vLLM's per-tag index
|
||||
# instead — the cublas13 case in install.sh adds --index-strategy=unsafe-best-match
|
||||
# so uv consults this index alongside PyPI.
|
||||
--extra-index-url https://wheels.vllm.ai/0.28.0/cu130
|
||||
--extra-index-url https://wheels.vllm.ai/0.29.0/cu130
|
||||
# VERSION COUPLING: darwin/Apple-Silicon builds use vllm-metal (see install.sh),
|
||||
# which pins this exact vLLM version. Bumping vllm here means coordinating with a
|
||||
# vllm-metal release that supports the new version, or macOS/Metal builds break.
|
||||
vllm==0.28.0
|
||||
vllm==0.29.0
|
||||
@@ -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)
|
||||
@@ -99,3 +99,12 @@ var DiffusersSchedulerOptions = []FieldOption{
|
||||
{Value: "heun", Label: "Heun"},
|
||||
{Value: "unipc", Label: "UniPC"},
|
||||
}
|
||||
|
||||
// SystemMessagesAfterFirstOptions are the values of template.system_messages_after_first:
|
||||
// how system messages that appear after the first turn are handled before the chat
|
||||
// template runs (empty = pass through unchanged, which strict Jinja templates reject).
|
||||
var SystemMessagesAfterFirstOptions = []FieldOption{
|
||||
{Value: "", Label: "Pass through (default)"},
|
||||
{Value: "merge", Label: "Merge into the first system message"},
|
||||
{Value: "user", Label: "Forward as user messages"},
|
||||
}
|
||||
@@ -382,6 +382,14 @@ func DefaultRegistry() map[string]FieldMetaOverride {
|
||||
Description: "Use the chat template from the model's tokenizer config",
|
||||
Order: 44,
|
||||
},
|
||||
"template.system_messages_after_first": {
|
||||
Section: "templates",
|
||||
Label: "System Messages After First",
|
||||
Description: "How system messages that appear after the first turn are handled before templating: merge into the first system message, or forward as user messages. Empty passes them through unchanged, which strict Jinja templates reject.",
|
||||
Component: "select",
|
||||
Options: SystemMessagesAfterFirstOptions,
|
||||
Order: 45,
|
||||
},
|
||||
// Router section template — kept in the templates UI section
|
||||
// (rather than the router section under "other") so operators
|
||||
// editing prompt shapes find all template-typed fields in one
|
||||
|
||||
@@ -1351,6 +1351,16 @@ type TemplateConfig struct {
|
||||
// that can use the tokenizers specified in the JSON config files of the models
|
||||
UseTokenizerTemplate bool `yaml:"use_tokenizer_template,omitempty" json:"use_tokenizer_template,omitempty"`
|
||||
|
||||
// SystemMessagesAfterFirst controls what happens to system-role messages that
|
||||
// appear after the leading system block. Some tokenizer chat templates (e.g.
|
||||
// Qwen3.8 / Flash-Next) raise "System message must be at the beginning" for
|
||||
// them, while agent frameworks (cogito tool selection, adjustment prompts)
|
||||
// legitimately append system instructions mid-conversation.
|
||||
// ""/"error": pass through unchanged (template decides)
|
||||
// "merge": fold them into the leading system message
|
||||
// "user": forward them as user-role instructions (keeps their position)
|
||||
SystemMessagesAfterFirst string `yaml:"system_messages_after_first,omitempty" json:"system_messages_after_first,omitempty"`
|
||||
|
||||
// JoinChatMessagesByCharacter is a string that will be used to join chat messages together.
|
||||
// It defaults to \n
|
||||
JoinChatMessagesByCharacter *string `yaml:"join_chat_messages_by_character,omitempty" json:"join_chat_messages_by_character,omitempty"`
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
})
|
||||
@@ -66,6 +66,65 @@ func stripEmptySystemMessages(messages []schema.Message) []schema.Message {
|
||||
return out
|
||||
}
|
||||
|
||||
// normalizeLateSystemMessages handles system-role messages that appear after the
|
||||
// leading system block, according to template.system_messages_after_first:
|
||||
// "merge" folds them into the first system message (created if absent), "user"
|
||||
// forwards them as user-role turns at their original position. Any other value
|
||||
// returns the messages unchanged. Needed for tokenizer templates that reject
|
||||
// late system turns (Qwen3.8: "System message must be at the beginning") while
|
||||
// agent frameworks append instructions mid-conversation.
|
||||
func normalizeLateSystemMessages(messages []schema.Message, mode string) []schema.Message {
|
||||
if mode != "merge" && mode != "user" {
|
||||
return messages
|
||||
}
|
||||
lead := 0
|
||||
for lead < len(messages) && messages[lead].Role == "system" {
|
||||
lead++
|
||||
}
|
||||
late := false
|
||||
for _, m := range messages[lead:] {
|
||||
if m.Role == "system" {
|
||||
late = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !late {
|
||||
return messages
|
||||
}
|
||||
out := make([]schema.Message, 0, len(messages)+1)
|
||||
out = append(out, messages[:lead]...)
|
||||
if mode == "merge" && lead == 0 {
|
||||
out = append(out, schema.Message{Role: "system"})
|
||||
}
|
||||
for _, m := range messages[lead:] {
|
||||
if m.Role != "system" {
|
||||
out = append(out, m)
|
||||
continue
|
||||
}
|
||||
text := strings.TrimSpace(messageText(m))
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
switch mode {
|
||||
case "merge":
|
||||
first := &out[0]
|
||||
joined := strings.TrimSpace(messageText(*first))
|
||||
if joined != "" {
|
||||
joined += "\n\n"
|
||||
}
|
||||
joined += text
|
||||
first.Content = joined
|
||||
first.StringContent = joined
|
||||
case "user":
|
||||
m.Role = "user"
|
||||
m.Content = text
|
||||
m.StringContent = text
|
||||
out = append(out, m)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// mergeToolCallDeltas merges streaming tool call deltas into complete tool calls.
|
||||
// In SSE streaming, a single tool call arrives as multiple chunks sharing the same Index:
|
||||
// the first chunk carries the ID, Type, and Name; subsequent chunks append to Arguments.
|
||||
@@ -182,6 +241,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
|
||||
// Drop blank system turns from the web UI (and similar clients) so they
|
||||
// cannot suppress the model YAML system_prompt / tokenizer defaults.
|
||||
input.Messages = stripEmptySystemMessages(input.Messages)
|
||||
input.Messages = normalizeLateSystemMessages(input.Messages, config.TemplateConfig.SystemMessagesAfterFirst)
|
||||
|
||||
// Tokenizer-template models pass messages through to the backend as-is,
|
||||
// so apply the configured system_prompt when the request did not supply
|
||||
|
||||
@@ -378,6 +378,50 @@ var _ = Describe("system message helpers", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("normalizeLateSystemMessages", func() {
|
||||
msgs := func() []schema.Message {
|
||||
return []schema.Message{
|
||||
{Role: "system", Content: "lead", StringContent: "lead"},
|
||||
{Role: "user", Content: "q", StringContent: "q"},
|
||||
{Role: "assistant", Content: "a", StringContent: "a"},
|
||||
{Role: "system", Content: "late", StringContent: "late"},
|
||||
{Role: "user", Content: "q2", StringContent: "q2"},
|
||||
}
|
||||
}
|
||||
It("leaves messages untouched by default", func() {
|
||||
out := normalizeLateSystemMessages(msgs(), "")
|
||||
Expect(out).To(HaveLen(5))
|
||||
Expect(out[3].Role).To(Equal("system"))
|
||||
})
|
||||
It("merge folds late system turns into the leading one", func() {
|
||||
out := normalizeLateSystemMessages(msgs(), "merge")
|
||||
Expect(out).To(HaveLen(4))
|
||||
Expect(out[0].Role).To(Equal("system"))
|
||||
Expect(out[0].StringContent).To(Equal("lead\n\nlate"))
|
||||
for _, m := range out[1:] {
|
||||
Expect(m.Role).NotTo(Equal("system"))
|
||||
}
|
||||
})
|
||||
It("merge creates a leading system message when none exists", func() {
|
||||
in := msgs()[1:]
|
||||
out := normalizeLateSystemMessages(in, "merge")
|
||||
Expect(out[0].Role).To(Equal("system"))
|
||||
Expect(out[0].StringContent).To(Equal("late"))
|
||||
Expect(out).To(HaveLen(4))
|
||||
})
|
||||
It("user forwards late system turns as user turns in place", func() {
|
||||
out := normalizeLateSystemMessages(msgs(), "user")
|
||||
Expect(out).To(HaveLen(5))
|
||||
Expect(out[3].Role).To(Equal("user"))
|
||||
Expect(out[3].StringContent).To(Equal("late"))
|
||||
Expect(out[0].Role).To(Equal("system"))
|
||||
})
|
||||
It("does nothing when no late system turn exists", func() {
|
||||
in := msgs()[:3]
|
||||
Expect(normalizeLateSystemMessages(in, "user")).To(HaveLen(3))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("stripEmptySystemMessages", func() {
|
||||
It("removes blank system turns and keeps the rest", func() {
|
||||
in := []schema.Message{
|
||||
|
||||
@@ -9,9 +9,11 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -31,6 +33,7 @@ import (
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openai/types"
|
||||
"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"
|
||||
laudio "github.com/mudler/LocalAI/pkg/audio"
|
||||
"github.com/mudler/LocalAI/pkg/functions"
|
||||
@@ -135,6 +138,8 @@ type Session struct {
|
||||
Instructions string
|
||||
DefaultConversationID string
|
||||
ModelInterface Model
|
||||
ttsParams map[string]string
|
||||
voiceRelease func()
|
||||
// The pipeline model config or the config for an any-to-any model
|
||||
ModelConfig *config.ModelConfig
|
||||
InputSampleRate int
|
||||
@@ -198,6 +203,22 @@ type Session struct {
|
||||
respSink *responseSink
|
||||
}
|
||||
|
||||
func (s *Session) installVoiceBinding(voice string, params map[string]string, release func()) {
|
||||
if release == nil {
|
||||
release = func() {}
|
||||
}
|
||||
var once sync.Once
|
||||
s.Voice = voice
|
||||
s.ttsParams = maps.Clone(params)
|
||||
s.voiceRelease = func() { once.Do(release) }
|
||||
}
|
||||
|
||||
func (s *Session) releaseVoiceBinding() {
|
||||
if s.voiceRelease != nil {
|
||||
s.voiceRelease()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) FromClient(session *types.SessionUnion) {
|
||||
}
|
||||
|
||||
@@ -633,6 +654,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(), wrapped.TTSConfig.TTSConfig.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
|
||||
}
|
||||
session.installVoiceBinding(resolvedVoice, params, release)
|
||||
defer session.releaseVoiceBinding()
|
||||
wrapped.setTTSParams(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
|
||||
@@ -826,6 +858,7 @@ func runRealtimeSession(application *application.Application, t Transport, model
|
||||
application.ApplicationConfig(),
|
||||
evaluator,
|
||||
buildRealtimeRoutingContext(application, session.ID),
|
||||
application.VoiceProfileStore(),
|
||||
); err != nil {
|
||||
xlog.Error("failed to update session", "error", err)
|
||||
sendError(t, "session_update_error", fmt.Sprintf("Failed to update session: %v", err), "", "")
|
||||
@@ -1164,7 +1197,7 @@ func updateTransSession(session *Session, update *types.SessionUnion, cl *config
|
||||
return nil
|
||||
}
|
||||
|
||||
func updateSession(session *Session, update *types.SessionUnion, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig, evaluator *templates.Evaluator, routing *RealtimeRoutingContext) error {
|
||||
func updateSession(session *Session, update *types.SessionUnion, cl *config.ModelConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig, evaluator *templates.Evaluator, routing *RealtimeRoutingContext, profiles *voiceprofile.Store) error {
|
||||
sessionLock.Lock()
|
||||
defer sessionLock.Unlock()
|
||||
|
||||
@@ -1172,8 +1205,13 @@ func updateSession(session *Session, update *types.SessionUnion, cl *config.Mode
|
||||
return nil
|
||||
}
|
||||
|
||||
session.TranscriptionOnly = false
|
||||
rt := update.Realtime
|
||||
explicitVoice := rt.Audio != nil && rt.Audio.Output != nil && rt.Audio.Output.Voice != ""
|
||||
rebuild := rt.Model != "" || explicitVoice || (rt.Audio != nil && rt.Audio.Input != nil && rt.Audio.Input.Transcription != nil)
|
||||
|
||||
candidateModelName := session.Model
|
||||
candidateConfig := session.ModelConfig
|
||||
candidateTranscription := session.InputAudioTranscription
|
||||
|
||||
if rt.Model != "" {
|
||||
cfg, err := cl.LoadModelConfigFileByNameDefaultOptions(rt.Model, appConfig)
|
||||
@@ -1184,40 +1222,78 @@ func updateSession(session *Session, update *types.SessionUnion, cl *config.Mode
|
||||
return fmt.Errorf("model is not a valid pipeline model: %s", rt.Model)
|
||||
}
|
||||
|
||||
if session.InputAudioTranscription == nil {
|
||||
session.InputAudioTranscription = &types.AudioTranscription{}
|
||||
}
|
||||
session.InputAudioTranscription.Model = cfg.Pipeline.Transcription
|
||||
session.Voice = cfg.TTSConfig.Voice
|
||||
session.Model = rt.Model
|
||||
session.ModelConfig = cfg
|
||||
}
|
||||
|
||||
if rt.Audio != nil && rt.Audio.Output != nil && rt.Audio.Output.Voice != "" {
|
||||
session.Voice = string(rt.Audio.Output.Voice)
|
||||
candidateModelName = rt.Model
|
||||
candidateConfig = cfg
|
||||
candidateTranscription = &types.AudioTranscription{Model: cfg.Pipeline.Transcription}
|
||||
}
|
||||
|
||||
if rt.Audio != nil && rt.Audio.Input != nil && rt.Audio.Input.Transcription != nil {
|
||||
trUpd := rt.Audio.Input.Transcription
|
||||
trUpd := *rt.Audio.Input.Transcription
|
||||
// A language-only update (e.g. a client forcing the STT language) carries
|
||||
// an empty Model. Preserve the pipeline's configured transcription backend
|
||||
// instead of blanking it — otherwise the next utterance transcribes against
|
||||
// an empty model and the backend RPC fails with "unimplemented".
|
||||
if trUpd.Model == "" && session.InputAudioTranscription != nil {
|
||||
trUpd.Model = session.InputAudioTranscription.Model
|
||||
if trUpd.Model == "" && candidateTranscription != nil {
|
||||
trUpd.Model = candidateTranscription.Model
|
||||
}
|
||||
session.InputAudioTranscription = trUpd
|
||||
candidateTranscription = &trUpd
|
||||
if trUpd.Model != "" {
|
||||
session.ModelConfig.Pipeline.Transcription = trUpd.Model
|
||||
cfgCopy := *candidateConfig
|
||||
candidateConfig = &cfgCopy
|
||||
candidateConfig.Pipeline.Transcription = trUpd.Model
|
||||
}
|
||||
}
|
||||
|
||||
if rt.Model != "" || (rt.Audio != nil && rt.Audio.Output != nil && rt.Audio.Output.Voice != "") || (rt.Audio != nil && rt.Audio.Input != nil && rt.Audio.Input.Transcription != nil) {
|
||||
m, err := newModel(&session.ModelConfig.Pipeline, cl, ml, appConfig, evaluator, routing)
|
||||
candidateModel := session.ModelInterface
|
||||
candidateVoice := session.Voice
|
||||
candidateParams := maps.Clone(session.ttsParams)
|
||||
var candidateRelease func()
|
||||
selectVoice := rt.Model != "" || explicitVoice
|
||||
if rebuild {
|
||||
m, err := newModel(&candidateConfig.Pipeline, cl, ml, appConfig, evaluator, routing)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
session.ModelInterface = m
|
||||
candidateModel = m
|
||||
wrapped := m.(*wrappedModel)
|
||||
if selectVoice {
|
||||
configuredVoice := wrapped.TTSConfig.TTSConfig.Voice
|
||||
if explicitVoice {
|
||||
configuredVoice = string(rt.Audio.Output.Voice)
|
||||
}
|
||||
candidateVoice, candidateParams, candidateRelease, err = resolveRealtimeVoice(
|
||||
context.Background(), configuredVoice, wrapped.TTSConfig, profiles,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
wrapped.setTTSParams(candidateParams)
|
||||
}
|
||||
|
||||
if rt.LocalAIClassifier != nil {
|
||||
if err := validateClassifierActivation(candidateModel, rt.LocalAIClassifier); err != nil {
|
||||
if candidateRelease != nil {
|
||||
candidateRelease()
|
||||
}
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
oldRelease := session.voiceRelease
|
||||
session.TranscriptionOnly = false
|
||||
session.Model = candidateModelName
|
||||
session.ModelConfig = candidateConfig
|
||||
session.ModelInterface = candidateModel
|
||||
session.InputAudioTranscription = candidateTranscription
|
||||
if selectVoice {
|
||||
session.installVoiceBinding(candidateVoice, candidateParams, candidateRelease)
|
||||
if oldRelease != nil {
|
||||
oldRelease()
|
||||
}
|
||||
}
|
||||
|
||||
if rebuild {
|
||||
// A session.update that swaps the model/voice rebuilds the pipeline, so
|
||||
// warm the new backends too (unless opted out) — otherwise the next turn
|
||||
// pays the cold-start load the original session warm-up already avoided.
|
||||
@@ -1226,9 +1302,9 @@ func updateSession(session *Session, update *types.SessionUnion, cl *config.Mode
|
||||
// stall every other session. Load errors are logged (and still surface on
|
||||
// first use); per-stage failures are already warned inside
|
||||
// backend.PreloadStages.
|
||||
if !session.ModelConfig.Pipeline.DisableWarmup {
|
||||
if !candidateConfig.Pipeline.DisableWarmup {
|
||||
go func() {
|
||||
if err := m.Warmup(context.Background()); err != nil {
|
||||
if err := candidateModel.Warmup(context.Background()); err != nil {
|
||||
xlog.Error("realtime warmup failed after session.update", "error", err)
|
||||
}
|
||||
}()
|
||||
@@ -1287,9 +1363,6 @@ func updateSession(session *Session, update *types.SessionUnion, cl *config.Mode
|
||||
// Replace-not-merge, like tools: the client owns the whole option
|
||||
// list. Invalid configs reject the update without touching the
|
||||
// session's current classifier.
|
||||
if err := validateClassifierActivation(session.ModelInterface, rt.LocalAIClassifier); err != nil {
|
||||
return err
|
||||
}
|
||||
session.Classifier = rt.LocalAIClassifier
|
||||
prewarmClassifier(session)
|
||||
}
|
||||
@@ -1923,7 +1996,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,39 @@ 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) setTTSParams(params map[string]string) {
|
||||
m.ttsParams = maps.Clone(params)
|
||||
}
|
||||
|
||||
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 +706,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,354 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openai/types"
|
||||
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"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("realtime session voice switching", func() {
|
||||
type fixture struct {
|
||||
store *voiceprofile.Store
|
||||
voiceDir string
|
||||
loader *config.ModelConfigLoader
|
||||
models *model.ModelLoader
|
||||
appConfig *config.ApplicationConfig
|
||||
profileA voiceprofile.Profile
|
||||
profileB voiceprofile.Profile
|
||||
}
|
||||
|
||||
newFixture := func(ctx SpecContext) *fixture {
|
||||
modelDir := GinkgoT().TempDir()
|
||||
voiceDir := GinkgoT().TempDir()
|
||||
store := voiceprofile.NewStore(voiceDir)
|
||||
DeferCleanup(func() { Expect(store.Close()).To(Succeed()) })
|
||||
profileA, err := store.Create(ctx, voiceprofile.CreateInput{
|
||||
Name: "Alpha", Language: "en", Transcript: "Alpha transcript", ConsentConfirmed: true,
|
||||
}, bytes.NewReader(realtimeProfileWAV(time.Second)))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
profileB, err := store.Create(ctx, voiceprofile.CreateInput{
|
||||
Name: "Beta", Language: "it", Transcript: "Beta transcript", ConsentConfirmed: true,
|
||||
}, bytes.NewReader(realtimeProfileWAV(time.Second)))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
configs := map[string]string{
|
||||
"vad": "name: vad\nbackend: test\nparameters:\n model: vad.bin\n",
|
||||
"stt": "name: stt\nbackend: test\nparameters:\n model: stt.bin\n",
|
||||
"llm": "name: llm\nbackend: test\nparameters:\n model: llm.bin\n",
|
||||
"tts-a": fmt.Sprintf("name: tts-a\nbackend: qwen3-tts-cpp\nparameters:\n model: tts-a.bin\ntts:\n voice: %s\n voice_cloning: true\n", profileA.Voice),
|
||||
"tts-b": fmt.Sprintf("name: tts-b\nbackend: qwen3-tts-cpp\nparameters:\n model: tts-b.bin\ntts:\n voice: %s\n voice_cloning: true\n", profileB.Voice),
|
||||
"pipe-a": "name: pipe-a\npipeline:\n vad: vad\n transcription: stt\n llm: llm\n tts: tts-a\n disable_warmup: true\n",
|
||||
"pipe-b": "name: pipe-b\npipeline:\n vad: vad\n transcription: stt\n llm: llm\n tts: tts-b\n disable_warmup: true\n",
|
||||
}
|
||||
for name, body := range configs {
|
||||
Expect(os.WriteFile(filepath.Join(modelDir, name+".yaml"), []byte(body), 0o644)).To(Succeed())
|
||||
}
|
||||
loader := config.NewModelConfigLoader(modelDir)
|
||||
Expect(loader.LoadModelConfigsFromPath(modelDir)).To(Succeed())
|
||||
state, err := system.GetSystemState(system.WithModelPath(modelDir))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return &fixture{
|
||||
store: store, voiceDir: voiceDir, loader: loader, models: model.NewModelLoader(state),
|
||||
appConfig: config.NewApplicationConfig(config.WithSystemState(state)),
|
||||
profileA: profileA, profileB: profileB,
|
||||
}
|
||||
}
|
||||
|
||||
newSession := func(f *fixture, voice string) *Session {
|
||||
cfg, err := f.loader.LoadModelConfigFileByNameDefaultOptions("pipe-a", f.appConfig)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
m, err := newModel(&cfg.Pipeline, f.loader, f.models, f.appConfig, nil, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
session := &Session{
|
||||
Model: "pipe-a", Voice: voice, ModelConfig: cfg, ModelInterface: m,
|
||||
InputAudioTranscription: &types.AudioTranscription{Model: "stt"},
|
||||
}
|
||||
resolved, params, release, err := resolveRealtimeVoice(context.Background(), voice, m.(*wrappedModel).TTSConfig, f.store)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
session.installVoiceBinding(resolved, params, release)
|
||||
return session
|
||||
}
|
||||
|
||||
update := func(f *fixture, session *Session, rt *types.RealtimeSession) error {
|
||||
return updateSession(session, &types.SessionUnion{Realtime: rt}, f.loader, f.models, f.appConfig, nil, nil, f.store)
|
||||
}
|
||||
|
||||
It("switches ordinary voices to profiles and clears the lease at final cleanup", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, "speaker-1")
|
||||
Expect(update(f, session, &types.RealtimeSession{Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: types.Voice(f.profileA.Voice)}}})).To(Succeed())
|
||||
|
||||
Expect(session.Voice).To(BeAnExistingFile())
|
||||
Expect(session.ttsParams).To(Equal(map[string]string{"ref_text": "Alpha transcript"}))
|
||||
leased := session.Voice
|
||||
session.releaseVoiceBinding()
|
||||
session.releaseVoiceBinding()
|
||||
Expect(leased).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("replaces one profile lease with another and then an ordinary voice", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, f.profileA.Voice)
|
||||
firstLease := session.Voice
|
||||
|
||||
Expect(update(f, session, &types.RealtimeSession{Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: types.Voice(f.profileB.Voice)}}})).To(Succeed())
|
||||
Expect(firstLease).NotTo(BeAnExistingFile())
|
||||
secondLease := session.Voice
|
||||
Expect(secondLease).To(BeAnExistingFile())
|
||||
Expect(session.ttsParams).To(HaveKeyWithValue("ref_text", "Beta transcript"))
|
||||
|
||||
Expect(update(f, session, &types.RealtimeSession{Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: "speaker-2"}}})).To(Succeed())
|
||||
Expect(secondLease).NotTo(BeAnExistingFile())
|
||||
Expect(session.Voice).To(Equal("speaker-2"))
|
||||
Expect(session.ttsParams).To(BeNil())
|
||||
Expect(session.ModelInterface.(*wrappedModel).ttsParams).To(BeNil())
|
||||
})
|
||||
|
||||
It("uses a new model default profile unless an explicit voice takes precedence", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, "speaker-1")
|
||||
Expect(update(f, session, &types.RealtimeSession{Model: "pipe-b"})).To(Succeed())
|
||||
Expect(session.ttsParams).To(HaveKeyWithValue("ref_text", "Beta transcript"))
|
||||
defaultLease := session.Voice
|
||||
|
||||
Expect(update(f, session, &types.RealtimeSession{
|
||||
Model: "pipe-a",
|
||||
Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: "speaker-explicit"}},
|
||||
})).To(Succeed())
|
||||
Expect(defaultLease).NotTo(BeAnExistingFile())
|
||||
Expect(session.Voice).To(Equal("speaker-explicit"))
|
||||
Expect(session.ttsParams).To(BeNil())
|
||||
})
|
||||
|
||||
It("preserves a profile binding across a language-only rebuild", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, f.profileA.Voice)
|
||||
lease := session.Voice
|
||||
Expect(update(f, session, &types.RealtimeSession{Audio: &types.RealtimeSessionAudio{Input: &types.SessionAudioInput{
|
||||
Transcription: &types.AudioTranscription{Language: "fr"},
|
||||
}}})).To(Succeed())
|
||||
|
||||
Expect(session.Voice).To(Equal(lease))
|
||||
Expect(session.InputAudioTranscription.Model).To(Equal("stt"))
|
||||
Expect(session.InputAudioTranscription.Language).To(Equal("fr"))
|
||||
wrapped := session.ModelInterface.(*wrappedModel)
|
||||
Expect(wrapped.ttsParams).To(Equal(map[string]string{"ref_text": "Alpha transcript"}))
|
||||
wrapped.ttsParams["ref_text"] = "wrapper mutation"
|
||||
Expect(session.ttsParams).To(HaveKeyWithValue("ref_text", "Alpha transcript"))
|
||||
})
|
||||
|
||||
It("rolls back the model, wrapper, voice, and lease when preparation fails", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, f.profileA.Voice)
|
||||
oldModel, oldConfig, oldVoice := session.ModelInterface, session.ModelConfig, session.Voice
|
||||
err := update(f, session, &types.RealtimeSession{
|
||||
Model: "pipe-b",
|
||||
Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: "localai://voice-profiles/00000000-0000-0000-0000-000000000001"}},
|
||||
})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(session.Model).To(Equal("pipe-a"))
|
||||
Expect(session.ModelConfig).To(BeIdenticalTo(oldConfig))
|
||||
Expect(session.ModelInterface).To(BeIdenticalTo(oldModel))
|
||||
Expect(session.Voice).To(Equal(oldVoice))
|
||||
Expect(oldVoice).To(BeAnExistingFile())
|
||||
Expect(session.ttsParams).To(HaveKeyWithValue("ref_text", "Alpha transcript"))
|
||||
})
|
||||
|
||||
It("releases a candidate profile lease when later validation fails", func(ctx SpecContext) {
|
||||
f := newFixture(ctx)
|
||||
session := newSession(f, "speaker-1")
|
||||
leases := func() []string {
|
||||
matches, err := filepath.Glob(filepath.Join(f.voiceDir, voiceprofile.DirectoryName, ".leases", "*", "*.wav"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
return matches
|
||||
}
|
||||
Expect(leases()).To(BeEmpty())
|
||||
|
||||
err := update(f, session, &types.RealtimeSession{
|
||||
Audio: &types.RealtimeSessionAudio{Output: &types.SessionAudioOutput{Voice: types.Voice(f.profileA.Voice)}},
|
||||
LocalAIClassifier: classifierTestConfig(0, nil),
|
||||
})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(session.Voice).To(Equal("speaker-1"))
|
||||
Expect(leases()).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
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')
|
||||
})
|
||||
})
|
||||
@@ -103,3 +103,48 @@ test.describe("Models gallery - recommended panel prominence", () => {
|
||||
await expect(grid(page).locator(".lane__tag--evidence")).toHaveCount(1);
|
||||
});
|
||||
});
|
||||
|
||||
// Start with a fitting model so absence assertions cannot pass during loading.
|
||||
// Then change the polled hardware budget while keeping the same gallery.
|
||||
for (const view of ["models", "home"]) {
|
||||
test(`${view} removes GPU recommendations when no candidate fits`, async ({ page }) => {
|
||||
await mockGallery(page, 0);
|
||||
await page.route("**/v1/models", (route) =>
|
||||
route.fulfill({ json: { data: [] } }),
|
||||
);
|
||||
const gib = 1024 ** 3;
|
||||
let budget = 24 * gib;
|
||||
await page.route("**/api/resources", (route) =>
|
||||
route.fulfill({ json: {
|
||||
type: "gpu",
|
||||
aggregate: { total_memory: budget, gpu_count: 1 },
|
||||
gpus: [{ vendor: "nvidia", total_memory: budget }],
|
||||
} }),
|
||||
);
|
||||
await page.route("**/api/models/estimate/*", (route) =>
|
||||
route.fulfill({ json: {
|
||||
sizeBytes: 17.4 * gib,
|
||||
sizeDisplay: "17.4 GB",
|
||||
estimates: { 4096: { vramBytes: 18.4 * gib, vramDisplay: "18.4 GB" } },
|
||||
} }),
|
||||
);
|
||||
await page.goto(view === "models" ? "/app/models" : "/app/");
|
||||
const section = view === "models" ? panel(page) : page.locator(".home-starters");
|
||||
await expect(section).toBeVisible();
|
||||
await expect(section).toContainText("tiny-chat");
|
||||
|
||||
// Wait for BOTH recommendation estimates, not the hook's loading render
|
||||
// or the gallery rail's separate context-size requests.
|
||||
const estimatesFinished = REC_MODELS.map(model => page.waitForResponse(response => {
|
||||
const url = new URL(response.url());
|
||||
return url.pathname.endsWith('/api/models/estimate/' + model.name) &&
|
||||
url.searchParams.get('contexts') === '4096' && response.status() === 200;
|
||||
}).then(response => response.finished()));
|
||||
budget = 12 * gib;
|
||||
await Promise.all(estimatesFinished);
|
||||
await page.evaluate(() => new Promise(resolve =>
|
||||
requestAnimationFrame(() => requestAnimationFrame(resolve)),
|
||||
));
|
||||
await expect(section).toHaveCount(0, { timeout: 15_000 });
|
||||
});
|
||||
}
|
||||
@@ -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
+139
-82
@@ -24,7 +24,7 @@
|
||||
"@modelcontextprotocol/sdk": "^1.30.0",
|
||||
"dompurify": "^3.4.13",
|
||||
"highlight.js": "^11.11.1",
|
||||
"hono": "4.12.34",
|
||||
"hono": "4.13.5",
|
||||
"i18next": "^26.0.8",
|
||||
"i18next-browser-languagedetector": "^8.2.1",
|
||||
"i18next-http-backend": "^3.0.6",
|
||||
@@ -32,7 +32,8 @@
|
||||
"react": "^19.1.0",
|
||||
"react-dom": "^19.1.0",
|
||||
"react-i18next": "^17.0.6",
|
||||
"react-router-dom": "^7.18.1",
|
||||
"react-router": "8.3.1",
|
||||
"react-router-dom": "7.18.2",
|
||||
"yaml": "^2.8.3"
|
||||
},
|
||||
"devDependencies": {
|
||||
@@ -648,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",
|
||||
@@ -754,9 +771,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/@istanbuljs/load-nyc-config/node_modules/js-yaml": {
|
||||
"version": "3.14.2",
|
||||
"resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-3.14.2.tgz",
|
||||
"integrity": "sha512-PMSmkqxr106Xa156c2M265Z+FTrPl+oxd/rgOQy2tijQeK5TxQ43psO1ZCwhVOSdnn+RzkzlRz/eY4BgJBYVpg==",
|
||||
"version": "3.15.2",
|
||||
"resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-3.15.2.tgz",
|
||||
"integrity": "sha512-6EuL879VkRA+1Cz578mKMiKvjPNEuk6+r1JaFzoSWejZmtf7xWbIyw1e3KkxlkzTIt9Taw6JBhEppG7utc1P+w==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
@@ -1639,9 +1656,9 @@
|
||||
]
|
||||
},
|
||||
"node_modules/baseline-browser-mapping": {
|
||||
"version": "2.10.38",
|
||||
"resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.10.38.tgz",
|
||||
"integrity": "sha512-31/02mVB4yuQU6adKk5SlY6m+mxDwUq5KZkyYgnLrrKl7TEm1+3PyDtDBz2kOv/wxZz41GHsvV1A/u6RmiyBvw==",
|
||||
"version": "2.11.20",
|
||||
"resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz",
|
||||
"integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"bin": {
|
||||
@@ -1810,9 +1827,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/browserslist": {
|
||||
"version": "4.28.2",
|
||||
"resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.2.tgz",
|
||||
"integrity": "sha512-48xSriZYYg+8qXna9kwqjIVzuQxi+KYWp2+5nCYnYKPTr0LvD89Jqk2Or5ogxz0NUMfIjhh2lIUX/LyX9B4oIg==",
|
||||
"version": "4.28.8",
|
||||
"resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz",
|
||||
"integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
@@ -1830,11 +1847,11 @@
|
||||
],
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"baseline-browser-mapping": "^2.10.12",
|
||||
"caniuse-lite": "^1.0.30001782",
|
||||
"electron-to-chromium": "^1.5.328",
|
||||
"node-releases": "^2.0.36",
|
||||
"update-browserslist-db": "^1.2.3"
|
||||
"baseline-browser-mapping": "^2.11.12",
|
||||
"caniuse-lite": "^1.0.30001809",
|
||||
"electron-to-chromium": "^1.5.402",
|
||||
"node-releases": "^2.0.53",
|
||||
"update-browserslist-db": "^1.3.0"
|
||||
},
|
||||
"bin": {
|
||||
"browserslist": "cli.js"
|
||||
@@ -1940,9 +1957,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/caniuse-lite": {
|
||||
"version": "1.0.30001799",
|
||||
"resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001799.tgz",
|
||||
"integrity": "sha512-hG1bReV+OUU+MOqK4t/ZWI0tZOyz3rqS9XuhOUz1cIcbwBKjOyJEJuw9ER5JuNyqxNk8u/JUVbGibBOL1yrjFw==",
|
||||
"version": "1.0.30001810",
|
||||
"resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz",
|
||||
"integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
@@ -2197,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",
|
||||
@@ -2437,9 +2460,9 @@
|
||||
"integrity": "sha512-WMwm9LhRUo+WUaRN+vRuETqG89IgZphVSNkdFgeb6sS/E4OrDIN7t48CAewSHXc6C8lefD8KKfr5vY61brQlow=="
|
||||
},
|
||||
"node_modules/electron-to-chromium": {
|
||||
"version": "1.5.375",
|
||||
"resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.375.tgz",
|
||||
"integrity": "sha512-ZWP5eB4BVPW/ZYo9252hQZHZ5XavtsTgpbhcmMmRwymavC5AsLWQWBPaKMeNd2LW0KGby5HPXvj7+sr4ta5j/Q==",
|
||||
"version": "1.5.420",
|
||||
"resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.420.tgz",
|
||||
"integrity": "sha512-2yD6XreGusOfNV+dUcvipJEXc3n/n7fgr7996aszTG+YY5E4mqM4tOq/3uhP129cazL9YHbVWSpc79ePotWtPA==",
|
||||
"dev": true,
|
||||
"license": "ISC"
|
||||
},
|
||||
@@ -2847,6 +2870,15 @@
|
||||
"express": ">= 4.11"
|
||||
}
|
||||
},
|
||||
"node_modules/express-rate-limit/node_modules/ip-address": {
|
||||
"version": "10.3.1",
|
||||
"resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.3.1.tgz",
|
||||
"integrity": "sha512-1e9d3kb97NHJTIJDZW9rKqW2h6+dFa50Dy0fpPSMQp2ADje5gvKsXmdiK6dwY5t76TaTt5+P5N1Y/LoToIxP6g==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">= 12"
|
||||
}
|
||||
},
|
||||
"node_modules/express/node_modules/cookie": {
|
||||
"version": "0.7.2",
|
||||
"resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.2.tgz",
|
||||
@@ -2879,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",
|
||||
@@ -3435,9 +3467,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/hono": {
|
||||
"version": "4.12.34",
|
||||
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.34.tgz",
|
||||
"integrity": "sha512-GqXJqY/xJkJmuloTrnV1ZEXG3fqte+VjkUqoRNZXcrUidiUOP4fMSIHHY4tsqZBK++kVyWmt/AAfSUuy57/eSA==",
|
||||
"version": "4.13.5",
|
||||
"resolved": "https://registry.npmjs.org/hono/-/hono-4.13.5.tgz",
|
||||
"integrity": "sha512-O6+/eCYRkzzzy0rPWwKLiGBR1nFuUPZynnwjxN1MBA62NNqbT0wQEzQyK2gSO5yDIDB336sXQleAhOHrzlYyKw==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">=16.9.0"
|
||||
@@ -4195,15 +4227,6 @@
|
||||
"resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz",
|
||||
"integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ=="
|
||||
},
|
||||
"node_modules/ip-address": {
|
||||
"version": "10.4.0",
|
||||
"resolved": "https://registry.npmjs.org/ip-address/-/ip-address-10.4.0.tgz",
|
||||
"integrity": "sha512-oSK96Grm3aP6OrS263xVxbNDGVL7rzBtYdpGqlDG8iQdoenDoTs/nkki+DflYbAEE8Xl6o5YxhxlrKvI3nqKXQ==",
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
"node": ">= 12"
|
||||
}
|
||||
},
|
||||
"node_modules/ipaddr.js": {
|
||||
"version": "1.9.1",
|
||||
"resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz",
|
||||
@@ -4570,10 +4593,21 @@
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/js-yaml": {
|
||||
"version": "4.1.1",
|
||||
"resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.1.1.tgz",
|
||||
"integrity": "sha512-qQKT4zQxXl8lLwBtHMWwaTcGfFOZviOJet3Oy/xmGk2gZH677CJM9EvtfdSkgWcATZhj/55JZ0rmy3myCT5lsA==",
|
||||
"version": "4.3.2",
|
||||
"resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.2.tgz",
|
||||
"integrity": "sha512-SFNOvSJ+Dgf/9An904Yx+CgSlIPCkIpao4qo51lpee25TIRejdH3rhR4EZMGoNx3/TP3O+wzWuiTFl4sqbltzA==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
"type": "github",
|
||||
"url": "https://github.com/sponsors/puzrin"
|
||||
},
|
||||
{
|
||||
"type": "github",
|
||||
"url": "https://github.com/sponsors/nodeca"
|
||||
}
|
||||
],
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"argparse": "^2.0.1"
|
||||
},
|
||||
@@ -5186,9 +5220,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/node-releases": {
|
||||
"version": "2.0.48",
|
||||
"resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.48.tgz",
|
||||
"integrity": "sha512-1uz8041X6LoI6ZSdZacM9lVY28vuzDlSKitnpbSNK0RfKoIJkX29NBPVEFXhnuSuEOA9Ww0xnPJ+ILWbGAv8DA==",
|
||||
"version": "2.0.54",
|
||||
"resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz",
|
||||
"integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"engines": {
|
||||
@@ -5950,9 +5984,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",
|
||||
@@ -6060,22 +6094,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": {
|
||||
@@ -6105,9 +6141,46 @@
|
||||
}
|
||||
},
|
||||
"node_modules/react-router": {
|
||||
"version": "7.18.1",
|
||||
"resolved": "https://registry.npmjs.org/react-router/-/react-router-7.18.1.tgz",
|
||||
"integrity": "sha512-GDLgg3i3uM0aeJO3Fm+TCS+sDQ7gu12T6x0qdTEzcwqEfleci7JwugVNIF3U//0FWKnJT7ptG+20B2jfDqnZAg==",
|
||||
"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-es": "^3.1.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=22.22.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": ">=19.2.7",
|
||||
"react-dom": ">=19.2.7"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
"react-dom": {
|
||||
"optional": true
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/react-router-dom": {
|
||||
"version": "7.18.2",
|
||||
"resolved": "https://registry.npmjs.org/react-router-dom/-/react-router-dom-7.18.2.tgz",
|
||||
"integrity": "sha512-AIKJ/jgGlFb3EbfCXk5Gzshiwt+l3mqbCrNjmEWMMjqQxNJ3svBa6bgzFyCC2Sw3RA0VWF1kg3uQf2OFhxb8hw==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"react-router": "7.18.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": ">=18",
|
||||
"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",
|
||||
@@ -6126,22 +6199,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/react-router-dom": {
|
||||
"version": "7.18.1",
|
||||
"resolved": "https://registry.npmjs.org/react-router-dom/-/react-router-dom-7.18.1.tgz",
|
||||
"integrity": "sha512-KaZh+X/6UtEp28x51AUYZDMg9NGoz2ja3dNHa+ta/tk40vCzKhQ/RypCWBMLbmDr6//E24Vv5uPsrqXFozdkAg==",
|
||||
"license": "MIT",
|
||||
"dependencies": {
|
||||
"react-router": "7.18.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"react": ">=18",
|
||||
"react-dom": ">=18"
|
||||
}
|
||||
},
|
||||
"node_modules/readable-stream": {
|
||||
"version": "2.3.8",
|
||||
"resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-2.3.8.tgz",
|
||||
@@ -7165,9 +7222,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/update-browserslist-db": {
|
||||
"version": "1.2.3",
|
||||
"resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz",
|
||||
"integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==",
|
||||
"version": "1.3.2",
|
||||
"resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz",
|
||||
"integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==",
|
||||
"dev": true,
|
||||
"funding": [
|
||||
{
|
||||
|
||||
@@ -19,8 +19,9 @@
|
||||
"coverage:report": "nyc report"
|
||||
},
|
||||
"overrides": {
|
||||
"hono": "4.12.34",
|
||||
"ip-address": "10.3.1"
|
||||
"hono": "4.13.5",
|
||||
"ip-address": "10.3.1",
|
||||
"path-to-regexp": "^8.4.0"
|
||||
},
|
||||
"dependencies": {
|
||||
"@codemirror/autocomplete": "^6.18.6",
|
||||
@@ -39,7 +40,7 @@
|
||||
"@modelcontextprotocol/sdk": "^1.30.0",
|
||||
"dompurify": "^3.4.13",
|
||||
"highlight.js": "^11.11.1",
|
||||
"hono": "4.12.34",
|
||||
"hono": "4.13.5",
|
||||
"i18next": "^26.0.8",
|
||||
"i18next-browser-languagedetector": "^8.2.1",
|
||||
"i18next-http-backend": "^3.0.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"
|
||||
},
|
||||
|
||||
@@ -3,63 +3,17 @@ import { useTranslation } from 'react-i18next'
|
||||
import { modelsApi } from '../utils/api'
|
||||
import { useRecommendedModels, isNvfp4Name } from '../hooks/useRecommendedModels'
|
||||
|
||||
// Static fallback used only when the live gallery / estimates can't be reached
|
||||
// (offline, trimmed gallery). The hook is the primary, data-driven path; these
|
||||
// are real gallery names kept as a safety net so onboarding never shows nothing.
|
||||
// Gemma picks use the QAT (quantization-aware-trained) Q4 builds. NVIDIA boxes
|
||||
// get NVFP4 + MTP variants at the mid/large tiers (see NVIDIA below).
|
||||
const BASE = {
|
||||
cpu: [
|
||||
{ name: 'gemma-4-e2b-it-qat-q4_0', size: '~1.5 GB' },
|
||||
{ name: 'qwen3.5-4b-claude-4.6-opus-reasoning-distilled', size: '~2.5 GB' },
|
||||
{ name: 'gemma-4-e4b-it-qat-q4_0', size: '~3 GB' },
|
||||
{ name: 'lfm2.5-1.2b-instruct', size: '~0.8 GB' },
|
||||
],
|
||||
'gpu-small': [
|
||||
{ name: 'gemma-4-e4b-it-qat-q4_0', size: '~3 GB' },
|
||||
{ name: 'lfm2.5-8b-a1b', size: '~5 GB' },
|
||||
{ name: 'qwen3.5-9b', size: '~5.5 GB' },
|
||||
{ name: 'gemma-4-12b-it-qat-q4_0', size: '~7 GB' },
|
||||
],
|
||||
'gpu-mid': [
|
||||
{ name: 'qwen3.6-27b', size: '~16 GB' },
|
||||
{ name: 'qwen3.6-27b-mtp-pi-tune', size: '~16 GB' },
|
||||
{ name: 'gemma-4-26b-a4b-it-qat-q4_0', size: '~16 GB' },
|
||||
{ name: 'qwen3.5-27b', size: '~16 GB' },
|
||||
],
|
||||
'gpu-large': [
|
||||
{ name: 'qwen3.6-35b-a3b-apex', size: '~20 GB' },
|
||||
{ name: 'qwen3.6-35b-a3b-claude-4.6-opus-reasoning-distilled', size: '~20 GB' },
|
||||
{ name: 'gemma-4-31b-it-qat-q4_0', size: '~18 GB' },
|
||||
{ name: 'qwen3.5-35b-a3b-apex', size: '~20 GB' },
|
||||
],
|
||||
}
|
||||
|
||||
// NVIDIA-only overrides: NVFP4 is a Blackwell-optimised 4-bit format paired with
|
||||
// MTP (multi-token prediction) for speed. Only the mid/large tiers have these.
|
||||
const NVIDIA = {
|
||||
'gpu-mid': [
|
||||
{ name: 'qwen3.6-27b-nvfp4-mtp', size: '~14 GB' },
|
||||
{ name: 'qwen3.6-27b-mtp-pi-tune', size: '~16 GB' },
|
||||
{ name: 'gemma-4-26b-a4b-it-qat-q4_0', size: '~16 GB' },
|
||||
{ name: 'qwen3.6-27b', size: '~16 GB' },
|
||||
],
|
||||
'gpu-large': [
|
||||
{ name: 'qwen3.6-35b-a3b-nvfp4-mtp', size: '~18 GB' },
|
||||
{ name: 'qwen3.6-27b-nvfp4-mtp', size: '~14 GB' },
|
||||
{ name: 'qwen3.6-35b-a3b-apex', size: '~20 GB' },
|
||||
{ name: 'gemma-4-31b-it-qat-q4_0', size: '~18 GB' },
|
||||
],
|
||||
}
|
||||
|
||||
function fallbackFor(tierId, isNvidia) {
|
||||
if (isNvidia && NVIDIA[tierId]) return NVIDIA[tierId]
|
||||
return BASE[tierId] || BASE.cpu
|
||||
}
|
||||
// Offline CPU suggestions do not claim a measured GPU fit.
|
||||
const CPU_FALLBACK = [
|
||||
{ name: 'gemma-4-e2b-it-qat-q4_0', size: '~1.5 GB' },
|
||||
{ name: 'qwen3.5-4b-claude-4.6-opus-reasoning-distilled', size: '~2.5 GB' },
|
||||
{ name: 'gemma-4-e4b-it-qat-q4_0', size: '~3 GB' },
|
||||
{ name: 'lfm2.5-1.2b-instruct', size: '~0.8 GB' },
|
||||
]
|
||||
|
||||
export default function StarterModels({ addToast, onInstallStarted }) {
|
||||
const { t } = useTranslation('home')
|
||||
const { recommended, tier, isNvidia, loading } = useRecommendedModels({ count: 4 })
|
||||
const { recommended, tier, loading } = useRecommendedModels({ count: 4 })
|
||||
const [installing, setInstalling] = useState(() => new Set())
|
||||
|
||||
// While the hardware probe + gallery query are in flight, render nothing
|
||||
@@ -67,10 +21,11 @@ export default function StarterModels({ addToast, onInstallStarted }) {
|
||||
if (loading) return null
|
||||
|
||||
// Prefer live recommendations; fall back to the static list only when the
|
||||
// gallery yielded nothing.
|
||||
// gallery yielded nothing on a CPU host. Static GPU picks have no measured
|
||||
// fit and must not replace an empty set of fitting recommendations.
|
||||
const items = (recommended && recommended.length > 0)
|
||||
? recommended.map(r => ({ name: r.name, size: r.sizeDisplay }))
|
||||
: fallbackFor(tier.id, isNvidia)
|
||||
: tier.id === 'cpu' ? CPU_FALLBACK : []
|
||||
|
||||
if (items.length === 0) return null
|
||||
|
||||
|
||||
+3
-3
@@ -53,16 +53,16 @@ function rank(candidates, tier, count, isNvidia) {
|
||||
}
|
||||
const limit = tier.vram * 0.95
|
||||
const fits = pool.filter(c => c.vramBytes != null && c.vramBytes <= limit)
|
||||
const base = fits.length > 0 ? fits : pool // tiny GPU where nothing fits → fall through to smallest
|
||||
const byPreference = (a, b) => {
|
||||
// On NVIDIA, surface NVFP4 first; then largest-that-fits (best quality).
|
||||
if (isNvidia) {
|
||||
const an = isNvfp4Name(a.name), bn = isNvfp4Name(b.name)
|
||||
if (an !== bn) return an ? -1 : 1
|
||||
}
|
||||
return fits.length > 0 ? b.sizeBytes - a.sizeBytes : a.sizeBytes - b.sizeBytes
|
||||
return b.sizeBytes - a.sizeBytes
|
||||
}
|
||||
return [...base].sort(byPreference).slice(0, count)
|
||||
// An oversized or unestimated model cannot be labelled a hardware fit.
|
||||
return [...fits].sort(byPreference).slice(0, count)
|
||||
}
|
||||
|
||||
export function useRecommendedModels({ count = 4, candidatePool = 10 } = {}) {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -439,6 +439,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
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
@@ -16,8 +17,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 +52,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 +153,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 +193,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 +227,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 +279,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 +317,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 +330,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 +362,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 +382,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 +419,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 +451,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 +539,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).
|
||||
@@ -484,10 +594,27 @@ func isFilePath(s string) bool {
|
||||
if strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://") {
|
||||
return false
|
||||
}
|
||||
// Raw JPEG base64 begins with /9j because every JPEG starts with the
|
||||
// FF D8 FF marker. Do not mistake that leading slash for an absolute path.
|
||||
if isRawJPEGBase64(s) {
|
||||
return false
|
||||
}
|
||||
// Starts with / (absolute path) or contains path separator
|
||||
return s[0] == '/' || filepath.IsAbs(s)
|
||||
}
|
||||
|
||||
func isRawJPEGBase64(s string) bool {
|
||||
if len(s) < 4 {
|
||||
return false
|
||||
}
|
||||
prefix, err := base64.StdEncoding.DecodeString(s[:4])
|
||||
if err != nil || len(prefix) != 3 || prefix[0] != 0xff || prefix[1] != 0xd8 || prefix[2] != 0xff {
|
||||
return false
|
||||
}
|
||||
_, err = io.Copy(io.Discard, base64.NewDecoder(base64.StdEncoding, strings.NewReader(s)))
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// copyFile copies src to dst.
|
||||
func copyFile(src, dst string) error {
|
||||
if err := os.MkdirAll(filepath.Dir(dst), 0750); err != nil {
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
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("passes raw JPEG base64 to predict without staging it as a path", func(ctx SpecContext) {
|
||||
const jpegBase64 = "/9j/2Q=="
|
||||
backend := &lifecycleBackend{}
|
||||
stager := &lifecycleStager{}
|
||||
client := NewFileStagingClient(backend, stager, "worker-1")
|
||||
|
||||
_, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{jpegBase64}})
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(stager.ensureCalls).To(BeEmpty())
|
||||
Expect(backend.predictInput.Images).To(Equal([]string{jpegBase64}))
|
||||
})
|
||||
|
||||
It("still recognizes an invalid JPEG-like base64 string as a path", func() {
|
||||
Expect(isFilePath("/9j/not-a-jpeg")).To(BeTrue())
|
||||
})
|
||||
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
outside := filepath.Join(canonicalWorkerTempDir(), "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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
outside := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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 := canonicalWorkerTempDir()
|
||||
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
|
||||
}
|
||||
@@ -24,7 +24,7 @@ var _ = Describe("Worker ephemeral staging cleanup", func() {
|
||||
return dir
|
||||
}
|
||||
|
||||
BeforeEach(func() { stagingDir = GinkgoT().TempDir() })
|
||||
BeforeEach(func() { stagingDir = canonicalWorkerTempDir() })
|
||||
|
||||
It("removes staged request directories older than the TTL", func() {
|
||||
old := mkEphemeral("aaaa1111", 48*time.Hour)
|
||||
@@ -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 := canonicalWorkerTempDir()
|
||||
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
|
||||
}
|
||||
Loaded 100 of 141 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user