Merge master into the distributed test branch

Preserve staging capacity limits and request cleanup over worker tunnels.
Adapt the release handlers and tests to the HTTP control transport.

Verified 33 control-path, 148 node staging, and 123 worker specs.
PostgreSQL integration specs require Docker, which is unavailable here.

Assisted-by: Codex:GPT-6
This commit is contained in:
localai-org-maint-bot committed 2026-09-08 20:11:13 +00:00
commit fbf7634725
43 files changed
+4672 -131

No files matched your search

+1 -1
View File
@@ -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?=f6277c1695a83cf388a8282c1c1a8757cf626f18
AUDIO_CPP_VERSION?=9c6a282337cc83f227cc10428867a478947706ad
AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
+2 -2
View File
@@ -1,10 +1,10 @@
# ds4 backend Makefile.
#
# Upstream pin lives below as DS4_VERSION?=9ab705347c1775e7599ede7eb81a6255ec7dccb5
# Upstream pin lives below as DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
# (.github/bump_deps.sh) can find and update it - matches the
# llama-cpp / ik-llama-cpp / turboquant convention.
DS4_VERSION?=9ab705347c1775e7599ede7eb81a6255ec7dccb5
DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
DS4_REPO?=https://github.com/antirez/ds4
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
+1 -1
View File
@@ -1,5 +1,5 @@
LLAMA_VERSION?=465e49b9cea78a68b9c244ffb48d0ee24a82873d
LLAMA_VERSION?=67672dc5b76f8bc17785a19d3dc6d1463fc2902c
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
CMAKE_ARGS?=
+1 -1
View File
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
# CrispASR version (release tag)
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
CRISPASR_VERSION?=c6d4265ae2ee2b8931b09d7d25d5c65c75c36a41
CRISPASR_VERSION?=301acd87b036764973b8bfba71e0a21818036d33
SO_TARGET?=libgocrispasr.so
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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?=d8fb10c02977c8ca999f3fb4e02df9ecf10f7ba6
STABLEDIFFUSION_GGML_VERSION?=d04e8950c1ec8d30248cbe996682b3182fb1adf6
CMAKE_ARGS+=-DGGML_MAX_NAME=128
+22 -6
View File
@@ -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 )
+20
View File
@@ -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
@@ -425,3 +426,22 @@ class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
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")
+58 -1
View File
@@ -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:
@@ -32,6 +37,9 @@ 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)
@@ -47,3 +55,52 @@ type FileStager interface {
// documented no-op.
ForgetNode(nodeID string)
}
// 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
}
+143 -11
View File
@@ -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"
@@ -145,6 +147,94 @@ func (h *HTTPFileStager) ForgetNode(nodeID string) {
}
}
// 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)
}
client, err := h.clientFor(nodeID)
if err != nil {
return 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 := 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)
}
client, err := h.clientFor(nodeID)
if err != nil {
return 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 := 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)
@@ -158,7 +248,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, client, addr, localPath, key); ok {
if remotePath, ok, probeErr := h.probeExisting(ctx, client, 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
}
@@ -507,14 +599,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, client *http.Client, 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, client *http.Client, 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)
@@ -522,18 +615,18 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, client *http.Client,
resp, err := 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
@@ -543,14 +636,53 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, client *http.Client,
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, client, addr, key)
if err != nil {
return "", false, err
}
if !claimed {
return "", false, nil
}
}
return remotePath, true, nil
}
func (h *HTTPFileStager) claimExisting(ctx context.Context, client *http.Client, 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 := 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,367 @@
package nodes
import (
"context"
"encoding/json"
"fmt"
"net"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"time"
"github.com/mudler/LocalAI/core/services/storage"
"github.com/mudler/LocalAI/core/services/workerctl"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"io"
)
type releaseTestControl struct {
subject string
payload []byte
onRequest func()
requestCalled bool
requestCount int
timeout time.Duration
replies [][]byte
}
func (m *releaseTestControl) RoundTrip(req *http.Request) (*http.Response, error) {
if err := req.Context().Err(); err != nil {
return nil, err
}
m.subject = req.URL.Path
m.payload, _ = io.ReadAll(req.Body)
m.requestCalled = true
m.requestCount++
if deadline, ok := req.Context().Deadline(); ok {
m.timeout = time.Until(deadline)
}
if m.onRequest != nil {
m.onRequest()
}
reply := []byte(`{}`)
if m.requestCount <= len(m.replies) {
reply = m.replies[m.requestCount-1]
}
return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(string(reply))), Request: req}, nil
}
func (m *releaseTestControl) control() *ControlClient {
c := NewControlClient(directNetDialerFor, "")
c.clients["node.one"] = &nodeHTTPClient{client: &http.Client{Transport: m}}
return c
}
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, directNetDialerFor), 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
}, "", directNetDialerFor)
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", directNetDialerFor)
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 := &releaseTestControl{}
client.onRequest = func() {
exists, existsErr := store.Exists(context.Background(), key)
Expect(existsErr).NotTo(HaveOccurred())
Expect(exists).To(BeTrue())
}
stager := NewS3FileStager(fm, client.control())
Expect(stager.ReleaseRemote(context.Background(), "node.one", key)).To(Succeed())
Expect(client.requestCalled).To(BeTrue())
Expect(client.subject).To(Equal(workerctl.PathFilesRelease))
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 control 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 := &releaseTestControl{}
stager := NewS3FileStager(fm, client.control())
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 := &releaseTestControl{}
stager := NewS3FileStager(fm, client.control())
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 control 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 := &releaseTestControl{replies: [][]byte{
[]byte(`{"error":"batch payload unsupported"}`),
[]byte(`{}`),
[]byte(`{}`),
}}
stager := NewS3FileStager(fm, client.control())
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 := &releaseTestControl{}
stager := NewS3FileStager(fm, client.control())
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 control 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 := &releaseTestControl{}
stager := NewS3FileStager(fm, client.control())
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))
})
})
+64
View File
@@ -2,6 +2,7 @@ package nodes
import (
"context"
"errors"
"fmt"
"time"
@@ -100,6 +101,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 {
@@ -251,3 +261,57 @@ func (s *S3FileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePat
return nil
}
// ReleaseRemote evicts one exact ephemeral key from the worker before deleting
// the shared object.
func (s *S3FileStager) 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 control request.
// A worker that only understands the exact-key payload returns an error, so the
// frontend retries each key during a rolling upgrade.
func (s *S3FileStager) 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 *S3FileStager) releaseWorkerKeys(ctx context.Context, nodeID string, request fileReleaseRequest) error {
var reply fileReleaseReply
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesRelease, request, &reply); err != nil {
return err
}
if reply.Error != "" {
return fmt.Errorf("backend release failed: %s", reply.Error)
}
return nil
}
+140 -45
View File
@@ -19,6 +19,8 @@ import (
"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.
@@ -53,21 +55,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.
@@ -105,23 +156,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)
}
@@ -131,7 +196,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)
}
@@ -165,25 +230,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)
}
@@ -215,11 +282,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)
}
@@ -251,7 +320,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
@@ -262,7 +333,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)
}
@@ -294,14 +365,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)
}
@@ -312,11 +385,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)
}
@@ -348,24 +423,27 @@ func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGen
}
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, requestID(), in.Src, "inputs")
backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs")
if err != nil {
return nil, fmt.Errorf("staging audio for sound detection: %w", err)
}
// Keep the frontend path available if the caller retries on another node.
in = proto.Clone(in).(*pb.SoundDetectionRequest)
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)
}
@@ -376,11 +454,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)
}
@@ -462,31 +542,46 @@ func (f *FileStagingClient) QuantizationProgress(ctx context.Context, in *pb.Qua
// stageMultimodalInputs stages Images, Videos, Audios fields in PredictOptions
// if they are file paths (not base64 or URLs).
func (f *FileStagingClient) stageMultimodalInputs(ctx context.Context, reqID string, in *pb.PredictOptions) (*pb.PredictOptions, []string) {
var keys []string
in.Images = f.stagePathSlice(ctx, reqID, in.Images, "inputs", &keys)
in.Videos = f.stagePathSlice(ctx, reqID, in.Videos, "inputs", &keys)
in.Audios = f.stagePathSlice(ctx, reqID, in.Audios, "inputs", &keys)
return in, keys
func (f *FileStagingClient) stageMultimodalInputs(
ctx context.Context,
lifecycle *stagedInputLifecycle,
in *pb.PredictOptions,
) (*pb.PredictOptions, error) {
var err error
in.Images, err = f.stagePathSlice(ctx, lifecycle, in.Images, "inputs")
if err != nil {
return nil, fmt.Errorf("staging predict images: %w", err)
}
in.Videos, err = f.stagePathSlice(ctx, lifecycle, in.Videos, "inputs")
if err != nil {
return nil, fmt.Errorf("staging predict videos: %w", err)
}
in.Audios, err = f.stagePathSlice(ctx, lifecycle, in.Audios, "inputs")
if err != nil {
return nil, fmt.Errorf("staging predict audios: %w", err)
}
return in, nil
}
func (f *FileStagingClient) stagePathSlice(ctx context.Context, reqID string, paths []string, category string, keys *[]string) []string {
func (f *FileStagingClient) stagePathSlice(
ctx context.Context,
lifecycle *stagedInputLifecycle,
paths []string,
category string,
) ([]string, error) {
result := make([]string, len(paths))
for i, p := range paths {
if isFilePath(p) {
backendPath, key, err := f.stageInputFile(ctx, reqID, p, category)
backendPath, err := f.stageInputFile(ctx, lifecycle, p, category)
if err != nil {
xlog.Warn("Failed to stage multimodal file, passing through", "path", p, "error", err)
result[i] = p
continue
return nil, fmt.Errorf("staging %q: %w", p, err)
}
result[i] = backendPath
*keys = append(*keys, key)
} else {
result[i] = p
}
}
return result
return result, nil
}
// isFilePath checks if a string looks like a local file path (not base64 or URL).
@@ -0,0 +1,341 @@
package nodes
import (
"context"
"errors"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
grpc "github.com/mudler/LocalAI/pkg/grpc"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
ggrpc "google.golang.org/grpc"
"google.golang.org/protobuf/proto"
)
const fullUUIDPattern = `[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`
type lifecycleStager struct {
fakeFileStager
ensureErr error
ensureErrAt int
releaseErr error
releasedKeys []string
releaseBatches [][]string
releaseCtxErr []error
releaseHasDeadline []bool
releaseDeadlines []time.Time
}
func (s *lifecycleStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
s.fakeFileStager.EnsureRemote(ctx, nodeID, localPath, key)
if s.ensureErr != nil && (s.ensureErrAt == 0 || len(s.ensureCalls) == s.ensureErrAt) {
return "", s.ensureErr
}
return "/remote/" + key, nil
}
func (s *lifecycleStager) ReleaseRemote(ctx context.Context, _ string, key string) error {
s.releasedKeys = append(s.releasedKeys, key)
s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err())
deadline, ok := ctx.Deadline()
s.releaseHasDeadline = append(s.releaseHasDeadline, ok)
s.releaseDeadlines = append(s.releaseDeadlines, deadline)
return s.releaseErr
}
func (s *lifecycleStager) ReleaseRemoteRequest(ctx context.Context, _, _ string, keys []string) error {
s.releaseBatches = append(s.releaseBatches, append([]string(nil), keys...))
s.releasedKeys = append(s.releasedKeys, keys...)
s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err())
deadline, ok := ctx.Deadline()
s.releaseHasDeadline = append(s.releaseHasDeadline, ok)
s.releaseDeadlines = append(s.releaseDeadlines, deadline)
return s.releaseErr
}
type lifecycleBackend struct {
grpc.Backend
predictResult *pb.Reply
predictErr error
predictCalls int
predictInput *pb.PredictOptions
streamCalls int
streamBlock <-chan struct{}
streamStarted chan<- struct{}
}
func (b *lifecycleBackend) Predict(_ context.Context, in *pb.PredictOptions, _ ...ggrpc.CallOption) (*pb.Reply, error) {
b.predictCalls++
b.predictInput = proto.Clone(in).(*pb.PredictOptions)
if b.predictResult == nil {
b.predictResult = &pb.Reply{}
}
return b.predictResult, b.predictErr
}
func (b *lifecycleBackend) PredictStream(_ context.Context, _ *pb.PredictOptions, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
b.streamCalls++
if b.streamStarted != nil {
b.streamStarted <- struct{}{}
}
if b.streamBlock != nil {
<-b.streamBlock
}
return nil
}
func (b *lifecycleBackend) GenerateImage(_ context.Context, _ *pb.GenerateImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
func (b *lifecycleBackend) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
func (b *lifecycleBackend) Generate3D(_ context.Context, _ *pb.Generate3DRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
func (b *lifecycleBackend) TTS(_ context.Context, _ *pb.TTSRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
func (b *lifecycleBackend) TTSStream(_ context.Context, _ *pb.TTSRequest, _ func(*pb.Reply), _ ...ggrpc.CallOption) error {
return nil
}
func (b *lifecycleBackend) SoundGeneration(_ context.Context, _ *pb.SoundGenerationRequest, _ ...ggrpc.CallOption) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
func (b *lifecycleBackend) SoundDetection(_ context.Context, _ *pb.SoundDetectionRequest, _ ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) {
return &pb.SoundDetectionResponse{}, nil
}
func (b *lifecycleBackend) AudioTranscription(_ context.Context, _ *pb.TranscriptRequest, _ ...ggrpc.CallOption) (*pb.TranscriptResult, error) {
return &pb.TranscriptResult{}, nil
}
func (b *lifecycleBackend) AudioTranscriptionStream(_ context.Context, _ *pb.TranscriptRequest, _ func(*pb.TranscriptStreamResponse), _ ...ggrpc.CallOption) error {
return nil
}
var _ = Describe("FileStagingClient request lifecycle", func() {
It("uses a full UUID for ephemeral request keys", func() {
Expect(requestID()).To(MatchRegexp(`^` + fullUUIDPattern + `$`))
})
It("releases every staged key and preserves caller requests", func(ctx SpecContext) {
tests := []struct {
name string
keyCount int
invoke func(*FileStagingClient) proto.Message
}{
{name: "predict", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}, Videos: []string{"/tmp/video.mp4"}, Audios: []string{"/tmp/audio.wav"}}
original := proto.Clone(request)
_, err := client.Predict(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "predict stream", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}}
original := proto.Clone(request)
Expect(client.PredictStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "image generation", keyCount: 2, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.GenerateImageRequest{Src: "/tmp/source.png", RefImages: []string{"/tmp/reference.png"}}
original := proto.Clone(request)
_, err := client.GenerateImage(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "video generation", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.GenerateVideoRequest{StartImage: "/tmp/start.png", EndImage: "/tmp/end.png", Audio: "/tmp/audio.wav"}
original := proto.Clone(request)
_, err := client.GenerateVideo(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "3D generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.Generate3DRequest{Src: "/tmp/source.glb"}
original := proto.Clone(request)
_, err := client.Generate3D(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
original := proto.Clone(request)
_, err := client.TTS(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "streaming TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.TTSRequest{Voice: "/tmp/voice.wav"}
original := proto.Clone(request)
Expect(client.TTSStream(ctx, request, func(*pb.Reply) {})).To(Succeed())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "sound generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
source := "/tmp/source.wav"
request := &pb.SoundGenerationRequest{Src: &source}
original := proto.Clone(request)
_, err := client.SoundGeneration(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "sound detection", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.SoundDetectionRequest{Src: "/tmp/source.wav"}
original := proto.Clone(request)
_, err := client.SoundDetection(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
original := proto.Clone(request)
_, err := client.AudioTranscription(ctx, request)
Expect(err).NotTo(HaveOccurred())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
{name: "streaming transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message {
request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"}
original := proto.Clone(request)
Expect(client.AudioTranscriptionStream(ctx, request, func(*pb.TranscriptStreamResponse) {})).To(Succeed())
Expect(proto.Equal(request, original)).To(BeTrue())
return request
}},
}
for _, test := range tests {
By(test.name)
stager := &lifecycleStager{}
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
test.invoke(client)
Expect(stager.ensureCalls).To(HaveLen(test.keyCount))
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
Expect(stager.releaseBatches).To(Equal([][]string{keysFromEnsureCalls(stager.ensureCalls)}))
Expect(stager.releaseCtxErr).To(Equal([]error{nil}))
Expect(stager.releaseHasDeadline).To(HaveLen(1))
for _, hasDeadline := range stager.releaseHasDeadline {
Expect(hasDeadline).To(BeTrue())
}
for _, deadline := range stager.releaseDeadlines {
Expect(time.Until(deadline)).To(BeNumerically(">", 0))
Expect(time.Until(deadline)).To(BeNumerically("<=", time.Minute))
}
}
})
It("tracks a key before staging so a partial upload failure is released", func(ctx SpecContext) {
uploadErr := errors.New("upload failed")
stager := &lifecycleStager{ensureErr: uploadErr, releaseErr: errors.New("release failed")}
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
_, err := client.GenerateImage(ctx, &pb.GenerateImageRequest{Src: "/tmp/source.png"})
Expect(err).To(MatchError(ContainSubstring("upload failed")))
Expect(stager.ensureCalls).To(HaveLen(1))
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
})
It("does not invoke predict when multimodal staging fails", func(ctx SpecContext) {
uploadErr := errors.New("ephemeral capacity exceeded")
stager := &lifecycleStager{ensureErr: uploadErr, ensureErrAt: 2}
backend := &lifecycleBackend{}
client := NewFileStagingClient(backend, stager, "worker-1")
request := &pb.PredictOptions{Images: []string{"/tmp/first.png", "/tmp/second.png"}}
original := proto.Clone(request)
result, err := client.Predict(ctx, request)
Expect(result).To(BeNil())
Expect(err).To(MatchError(ContainSubstring("ephemeral capacity exceeded")))
Expect(backend.predictCalls).To(BeZero())
Expect(proto.Equal(request, original)).To(BeTrue())
Expect(stager.ensureCalls).To(HaveLen(2))
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
})
It("does not invoke streaming predict when multimodal staging fails", func(ctx SpecContext) {
uploadErr := errors.New("ephemeral capacity exceeded")
stager := &lifecycleStager{ensureErr: uploadErr}
backend := &lifecycleBackend{}
client := NewFileStagingClient(backend, stager, "worker-1")
request := &pb.PredictOptions{Audios: []string{"/tmp/audio.wav"}}
original := proto.Clone(request)
err := client.PredictStream(ctx, request, func(*pb.Reply) {})
Expect(err).To(MatchError(ContainSubstring("ephemeral capacity exceeded")))
Expect(backend.streamCalls).To(BeZero())
Expect(proto.Equal(request, original)).To(BeTrue())
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
})
It("uses an active bounded cleanup context after caller cancellation", func() {
ctx, cancel := context.WithCancel(context.Background())
cancel()
stager := &lifecycleStager{}
client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1")
_, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
Expect(err).NotTo(HaveOccurred())
Expect(stager.releaseCtxErr).To(Equal([]error{nil}))
})
It("does not release streaming inputs before the backend completes", func(ctx SpecContext) {
block := make(chan struct{})
started := make(chan struct{}, 1)
backend := &lifecycleBackend{streamBlock: block, streamStarted: started}
stager := &lifecycleStager{}
client := NewFileStagingClient(backend, stager, "worker-1")
done := make(chan error, 1)
go func() {
done <- client.PredictStream(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}, func(*pb.Reply) {})
}()
Eventually(started).Should(Receive())
Expect(stager.ensureCalls).To(HaveLen(1))
Expect(stager.releasedKeys).To(BeEmpty())
close(block)
Eventually(done).Should(Receive(Succeed()))
Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls)))
})
It("does not replace a backend result when cleanup fails", func(ctx SpecContext) {
reply := &pb.Reply{Message: []byte("ok")}
backend := &lifecycleBackend{predictResult: reply}
stager := &lifecycleStager{releaseErr: errors.New("release failed")}
client := NewFileStagingClient(backend, stager, "worker-1")
result, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}})
Expect(err).NotTo(HaveOccurred())
Expect(result).To(BeIdenticalTo(reply))
})
})
func keysFromEnsureCalls(calls []ensureCall) []string {
keys := make([]string, len(calls))
for i, call := range calls {
keys[i] = call.key
}
return keys
}
@@ -28,6 +28,10 @@ func (s *soundStagingFailure) EnsureRemote(context.Context, string, string, stri
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) NewClientForNode(string, string, bool) (grpc.Backend, error) {
+2 -1
View File
@@ -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) {
+380 -14
View File
@@ -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"
)
@@ -67,17 +69,28 @@ type AuthenticatedRoutes struct {
// A nil readiness fails open, keeping /readyz's historical always-200 answer.
// A nil extra mounts no additional routes.
func StartFileTransferServer(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, extra *AuthenticatedRoutes, logStore ...*model.BackendLogStore) (*http.Server, error) {
return StartFileTransferServerWithCapacityAndRoutes(addr, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, nil, extra, 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) {
return StartFileTransferServerWithCapacityAndRoutes(addr, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, nil, logStore...)
}
// StartFileTransferServerWithCapacityAndRoutes combines capacity admission with authenticated control routes.
func StartFileTransferServerWithCapacityAndRoutes(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, extra *AuthenticatedRoutes, 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 StartFileTransferServerWithRoutes(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, extra, logStore...)
return startFileTransferServerWithRoutes(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, extra, 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
@@ -91,6 +104,14 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi
// StartFileTransferServerWithRoutes is StartFileTransferServerWithReadiness
// plus an extra authenticated route set. See AuthenticatedRoutes.
func StartFileTransferServerWithRoutes(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, extra *AuthenticatedRoutes, logStore ...*model.BackendLogStore) (*http.Server, error) {
return startFileTransferServerWithRoutes(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, nil, extra, logStore...)
}
func startFileTransferServer(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, logStore ...*model.BackendLogStore) (*http.Server, error) {
return startFileTransferServerWithRoutes(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, nil, logStore...)
}
func startFileTransferServerWithRoutes(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, extra *AuthenticatedRoutes, logStore ...*model.BackendLogStore) (*http.Server, error) {
// Checked before anything is created; see validateAuthenticatedRoutes.
if err := validateAuthenticatedRoutes(extra); err != nil {
return nil, err
@@ -129,6 +150,18 @@ func StartFileTransferServerWithRoutes(lis net.Listener, stagingDir, modelsDir,
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)
@@ -144,12 +177,16 @@ func StartFileTransferServerWithRoutes(lis net.Listener, stagingDir, modelsDir,
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)
}
@@ -314,6 +351,163 @@ func serveWorkerHTTP(lis net.Listener, mux *http.ServeMux, logFields ...any) *ht
return server
}
// 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)
@@ -382,6 +576,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).
@@ -423,10 +685,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)
@@ -461,18 +742,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)
@@ -483,13 +813,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
}
@@ -518,7 +867,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 {
@@ -602,6 +951,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.
@@ -613,10 +974,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 {
@@ -5,6 +5,7 @@ import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"net"
@@ -17,6 +18,7 @@ import (
"sync"
"time"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
@@ -65,6 +67,94 @@ var _ = Describe("The HTTP file stager without a worker dialer", func() {
})
})
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()
@@ -98,6 +188,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)
@@ -438,6 +589,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() {
@@ -482,6 +644,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
}, "", directNetDialerFor)
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
}, "", directNetDialerFor)
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
}, "", directNetDialerFor)
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
}, "", directNetDialerFor)
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)
+2
View File
@@ -51,6 +51,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
}
+3 -1
View File
@@ -50,7 +50,9 @@ type Config struct {
// HTTPAddr binds the HTTP file-transfer server. Default is loopback on
// basePort-1; an explicit value is bound exactly as given.
HTTPAddr string `env:"LOCALAI_HTTP_ADDR" default:"" help:"HTTP file transfer server bind address (default: loopback on the gRPC base port - 1)" group:"server" hidden:""`
HTTPAddr string `env:"LOCALAI_HTTP_ADDR" default:"" help:"HTTP file transfer server bind address (default: loopback on the gRPC base port - 1)" 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)
RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"`
+49 -5
View File
@@ -12,11 +12,12 @@ import (
"github.com/mudler/LocalAI/core/services/storage"
"github.com/mudler/LocalAI/core/services/workerctl"
"github.com/mudler/xlog"
"golang.org/x/sync/singleflight"
)
// The worker's file-staging control plane: four verbs that move model and job
// The worker's file-staging control plane: five verbs that move model and job
// artifacts between the object store both sides share and this worker's disk.
// They replace the four nodes.<id>.files.* NATS subjects.
// They replace the five nodes.<id>.files.* NATS subjects.
//
// The reply shapes are the ones those subjects already carried, unchanged, so
// an operator reading the wire sees the same fields. What DID change is that a
@@ -76,7 +77,7 @@ func (cfg *Config) stagingDataDir() string {
// from, over the same object store the frontend uses.
//
// It returns an error rather than degrading, because a worker whose deployment
// asked for object storage and could not reach it would otherwise mount four
// asked for object storage and could not reach it would otherwise mount five
// verbs that fail every call, which the frontend cannot tell from a worker out
// of disk.
func (cfg *Config) NewStagingFileManager(ctx context.Context) (*storage.FileManager, error) {
@@ -98,13 +99,19 @@ func (cfg *Config) NewStagingFileManager(ctx context.Context) (*storage.FileMana
return fm, nil
}
// RegisterFileControlRoutes mounts the four file-staging verbs on mux.
// RegisterFileControlRoutes mounts the five file-staging verbs on mux.
//
// The caller is responsible for putting mux behind authentication; see
// nodes.AuthenticatedRoutes, which is how the worker mounts this so the file
// verbs share one bearer check with the lifecycle verbs and the file routes
// rather than growing a third one.
func (cfg *Config) RegisterFileControlRoutes(mux *http.ServeMux, fm *storage.FileManager) {
cfg.RegisterFileControlRoutesWithCapacity(mux, fm, nil)
}
// RegisterFileControlRoutesWithCapacity applies worker staging capacity to the tunnel control verbs.
func (cfg *Config) RegisterFileControlRoutesWithCapacity(mux *http.ServeMux, fm *storage.FileManager, capacity *EphemeralCapacityGuard) {
var ensureGroup singleflight.Group
cacheDir := cfg.stagingCacheDir()
// files.ensure: download an object-store key into this worker's cache and
@@ -120,15 +127,23 @@ func (cfg *Config) RegisterFileControlRoutes(mux *http.ServeMux, fm *storage.Fil
if err := json.Unmarshal(body, &req); err != nil {
return nil, fmt.Errorf("invalid files.ensure request: %w", err)
}
localPath, err := fm.Download(ctx, req.Key)
value, err, _ := ensureGroup.Do(req.Key, func() (any, error) {
return ensureWorkerFile(ctx, fm, capacity, req.Key)
})
if err != nil {
xlog.Error("File ensure failed", "key", req.Key, "error", err)
return fileEnsureReply{Error: err.Error()}, nil
}
localPath, ok := value.(string)
if !ok {
return fileEnsureReply{Error: fmt.Sprintf("unexpected file ensure result %T", value)}, nil
}
xlog.Debug("File ensured locally", "key", req.Key, "path", localPath)
return fileEnsureReply{LocalPath: localPath}, nil
})
registerFileReleaseControlRoute(mux, fm, cacheDir, capacity)
// files.stage: upload one of this worker's files to the object store.
//
// The path allow-list is what keeps this verb from being an exfiltration
@@ -266,3 +281,32 @@ func listStagedFiles(ctx context.Context, dirPath string) ([]string, error) {
}
return files, nil
}
func registerFileReleaseControlRoute(mux *http.ServeMux, fm *storage.FileManager, cacheDir string, capacity *EphemeralCapacityGuard) {
postControlVerb(mux, workerctl.PathFilesRelease, func(ctx context.Context, body []byte) (any, error) {
var req struct {
Key string `json:"key"`
RequestID string `json:"request_id"`
}
if err := json.Unmarshal(body, &req); err != nil {
return nil, fmt.Errorf("invalid files.release request: %w", err)
}
var err error
if req.RequestID != "" {
err = releaseEphemeralCacheRequest(ctx, cacheDir, req.RequestID, capacity)
} else {
var cachePath string
cachePath, err = fm.CachePath(req.Key)
if err == nil {
err = releaseEphemeralCachePathWithCapacity(cacheDir, req.Key, cachePath, capacity)
}
}
reply := struct {
Error string `json:"error,omitempty"`
}{}
if err != nil {
reply.Error = err.Error()
}
return reply, nil
})
}
+3 -3
View File
@@ -448,7 +448,7 @@ var _ = Describe("the worker's HTTP server", func() {
return nil, nil
},
}
// A real object store, because the four file verbs are only mounted for
// A real object store, because the five file verbs are only mounted for
// a worker that has one, and the mounting assertion below walks every
// path this package names.
store, err := storage.NewFilesystemStore(filepath.Join(dir, "objectstore"))
@@ -492,7 +492,7 @@ var _ = Describe("the worker's HTTP server", func() {
It("serves no file verb at all when the deployment configured no object store", func() {
// Not a degraded mount: a worker with nowhere to stage to answers the
// four file paths the way a build that never had them does, which is
// five file paths the way a build that never had them does, which is
// the 404 the frontend already reads as "this worker does not serve
// that verb" rather than as a file that is not there.
dir := GinkgoT().TempDir()
@@ -503,7 +503,7 @@ var _ = Describe("the worker's HTTP server", func() {
for _, p := range []string{
workerctl.PathFilesEnsure, workerctl.PathFilesStage,
workerctl.PathFilesTemp, workerctl.PathFilesListDir,
workerctl.PathFilesTemp, workerctl.PathFilesListDir, workerctl.PathFilesRelease,
} {
req, reqErr := http.NewRequest(http.MethodPost, "http://"+bare.Addr+p, strings.NewReader("{}"))
Expect(reqErr).NotTo(HaveOccurred())
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,484 @@
// SPDX-License-Identifier: MIT
package worker
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"sync"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
type capacityGatedFileWriter struct {
file *os.File
entered chan struct{}
resume chan struct{}
once sync.Once
}
func (w *capacityGatedFileWriter) Write(p []byte) (int, error) {
w.once.Do(func() {
close(w.entered)
<-w.resume
})
return w.file.Write(p)
}
type capacityShortWriter struct{}
func (capacityShortWriter) Write(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
return 1, nil
}
var _ = Describe("EphemeralCapacityGuard", func() {
It("derives bounded defaults and preserves positive overrides", func() {
root := GinkgoT().TempDir()
limit, headroom, err := effectiveEphemeralCapacity([]string{root}, 0, -1)
Expect(err).NotTo(HaveOccurred())
Expect(limit).To(BeNumerically(">", 0))
Expect(limit).To(BeNumerically("<=", defaultEphemeralByteLimitCeiling))
Expect(headroom).To(BeNumerically(">=", defaultEphemeralMinFreeFloor))
limit, headroom, err = effectiveEphemeralCapacity([]string{root}, 123, 456)
Expect(err).NotTo(HaveOccurred())
Expect(limit).To(Equal(int64(123)))
Expect(headroom).To(Equal(int64(456)))
})
It("accounts existing regular files without following symlinks", func() {
root := GinkgoT().TempDir()
outside := filepath.Join(GinkgoT().TempDir(), "outside.bin")
Expect(os.WriteFile(filepath.Join(root, "existing.bin"), make([]byte, 6), 0o600)).To(Succeed())
Expect(os.WriteFile(outside, make([]byte, 100), 0o600)).To(Succeed())
Expect(os.Symlink(outside, filepath.Join(root, "outside-link"))).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
err = guard.Reserve(filepath.Join(root, "next.bin"), 5)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.RequestedBytes).To(Equal(int64(5)))
Expect(capacityErr.UsageBytes).To(Equal(int64(6)))
Expect(capacityErr.LimitBytes).To(Equal(int64(10)))
Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0))
Expect(capacityErr.HeadroomBytes).To(BeZero())
})
It("serializes competing reservations", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 1, 0)
Expect(err).NotTo(HaveOccurred())
start := make(chan struct{})
results := make(chan error, 32)
var wait sync.WaitGroup
for i := range 32 {
wait.Add(1)
go func(index int) {
defer wait.Done()
<-start
results <- guard.Reserve(filepath.Join(root, string(rune('a'+index))), 1)
}(i)
}
close(start)
wait.Wait()
close(results)
succeeded := 0
for result := range results {
if result == nil {
succeeded++
continue
}
var capacityErr *EphemeralCapacityError
Expect(errors.As(result, &capacityErr)).To(BeTrue())
}
Expect(succeeded).To(Equal(1))
})
It("makes only an equal active reservation idempotent", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "nested", "payload.bin")
Expect(guard.Reserve(path, 4)).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "nested", ".", "payload.bin"), 4)).To(Succeed())
err = guard.Reserve(path, 5)
var conflictErr *EphemeralReservationConflictError
Expect(errors.As(err, &conflictErr)).To(BeTrue())
Expect(conflictErr.ActiveBytes).To(Equal(int64(4)))
Expect(conflictErr.RequestedBytes).To(Equal(int64(5)))
Expect(guard.Reserve(filepath.Join(root, "other.bin"), 6)).To(Succeed())
Expect(guard.Release(filepath.Join(root, "nested", ".", "payload.bin"))).To(Succeed())
Expect(guard.Release(path)).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 4)).To(Succeed())
})
It("retains committed bytes when the same path starts another reservation", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "payload.bin")
Expect(guard.Reserve(path, 4)).To(Succeed())
Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
Expect(guard.HasActiveReservation(path)).To(BeTrue())
Expect(guard.Reserve(path, 6)).To(Succeed())
Expect(guard.HasActiveReservation(path)).To(BeTrue())
err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.UsageBytes).To(Equal(int64(10)))
})
It("retains startup-accounted bytes when the path is reserved", func() {
root := GinkgoT().TempDir()
path := filepath.Join(root, "payload.bin")
Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.HasActiveReservation(path)).To(BeFalse())
Expect(guard.Reserve(path, 6)).To(Succeed())
Expect(guard.HasActiveReservation(path)).To(BeTrue())
err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.UsageBytes).To(Equal(int64(10)))
})
It("commits the regular file's actual size", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "payload.bin")
Expect(guard.Reserve(path, 10)).To(Succeed())
Expect(os.WriteFile(path, []byte("four"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
Expect(guard.Commit(filepath.Clean(path))).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "six.bin"), 6)).To(Succeed())
})
It("preserves configured filesystem headroom", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 1<<30, 1<<62)
Expect(err).NotTo(HaveOccurred())
err = guard.Reserve(filepath.Join(root, "payload.bin"), 1)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.RequestedBytes).To(Equal(int64(1)))
Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0))
Expect(capacityErr.HeadroomBytes).To(Equal(int64(1 << 62)))
})
It("reserves bounded chunks before forwarding unknown-length input", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, ephemeralCapacityWriteChunk+1, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "payload.bin")
file, err := os.Create(path)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(file.Close)
writer, err := guard.NewWriter(path, file)
Expect(err).NotTo(HaveOccurred())
n, err := writer.Write(make([]byte, ephemeralCapacityWriteChunk+2))
Expect(n).To(Equal(int(ephemeralCapacityWriteChunk)))
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(writer.Close()).To(Succeed())
info, err := file.Stat()
Expect(err).NotTo(HaveOccurred())
Expect(info.Size()).To(Equal(ephemeralCapacityWriteChunk))
})
It("waits for an open bounded writer before committing", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "payload.bin")
file, err := os.Create(path)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(file.Close)
gated := &capacityGatedFileWriter{
file: file, entered: make(chan struct{}), resume: make(chan struct{}),
}
DeferCleanup(func() {
select {
case <-gated.resume:
default:
close(gated.resume)
}
})
writer, err := guard.NewWriter(path, gated)
Expect(err).NotTo(HaveOccurred())
writeDone := make(chan error, 1)
go func() {
_, writeErr := writer.Write([]byte("1234567"))
writeDone <- writeErr
}()
Eventually(gated.entered).Should(BeClosed())
commitDone := make(chan error, 1)
go func() { commitDone <- guard.Commit(path) }()
Eventually(func() int { return guard.commitWaiterCount(path) }).Should(Equal(1))
Expect(commitDone).NotTo(Receive())
close(gated.resume)
Expect(<-writeDone).To(Succeed())
Expect(commitDone).NotTo(Receive())
Expect(writer.Close()).To(Succeed())
Eventually(commitDone).Should(Receive(Succeed()))
err = guard.Reserve(filepath.Join(root, "other.bin"), 4)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.UsageBytes).To(Equal(int64(7)))
})
It("does not share pending capacity between concurrent writers", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "payload.bin")
file, err := os.Create(path)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(file.Close)
gated := &capacityGatedFileWriter{
file: file, entered: make(chan struct{}), resume: make(chan struct{}),
}
first, err := guard.NewWriter(path, gated)
Expect(err).NotTo(HaveOccurred())
var secondDestination bytes.Buffer
second, err := guard.NewWriter(path, &secondDestination)
Expect(err).NotTo(HaveOccurred())
firstDone := make(chan error, 1)
go func() {
_, writeErr := first.Write([]byte("1234567"))
firstDone <- writeErr
}()
Eventually(gated.entered).Should(BeClosed())
n, err := second.Write([]byte("7654321"))
Expect(n).To(BeZero())
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(secondDestination.Len()).To(BeZero())
close(gated.resume)
Expect(<-firstDone).To(Succeed())
Expect(first.Close()).To(Succeed())
Expect(second.Close()).To(Succeed())
})
It("rolls back bytes the destination writer does not accept", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 5, 0)
Expect(err).NotTo(HaveOccurred())
writer, err := guard.NewWriter(filepath.Join(root, "payload.bin"), capacityShortWriter{})
Expect(err).NotTo(HaveOccurred())
n, err := writer.Write([]byte("123"))
Expect(n).To(Equal(1))
Expect(err).To(MatchError(io.ErrShortWrite))
Expect(writer.Close()).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "other.bin"), 4)).To(Succeed())
})
It("rejects paths outside roots and through symlinks", func() {
root := GinkgoT().TempDir()
outside := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 100, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.Reserve(filepath.Join(outside, "payload.bin"), 1)).To(
MatchError(ContainSubstring("outside registered ephemeral roots")),
)
Expect(os.Symlink(outside, filepath.Join(root, "escape"))).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "escape", "payload.bin"), 1)).To(
MatchError(ContainSubstring("symlink")),
)
})
It("supports recovery tree accounting without dropping active reservations", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
active := filepath.Join(root, "active", "payload.bin")
stale := filepath.Join(root, "stale", "payload.bin")
Expect(guard.Reserve(active, 4)).To(Succeed())
Expect(guard.Account(stale, 3)).To(Succeed())
Expect(guard.HasActiveReservation(root)).To(BeTrue())
Expect(guard.ReleaseTree(filepath.Join(root, "stale"))).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 6)).To(Succeed())
Expect(guard.ReleaseTree(root)).To(Succeed())
Expect(guard.HasActiveReservation(root)).To(BeFalse())
})
It("waits for pre-release reservations before request cleanup scans", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "audio", "request-1", "input.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
Expect(guard.Reserve(path, 4)).To(Succeed())
released := make(chan error, 1)
go func() {
released <- guard.BeginRequestRelease(context.Background(), "request-1")
}()
Eventually(guard.releaseTombstoneCount).Should(Equal(1))
Consistently(released, 50*time.Millisecond).ShouldNot(Receive())
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
Eventually(released).Should(Receive(Succeed()))
guard.EndRequestRelease("request-1")
Expect(guard.HasActiveReservation(path)).To(BeFalse())
})
It("rejects staging after request cleanup begins", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.BeginRequestRelease(context.Background(), "request-1")).To(Succeed())
defer guard.EndRequestRelease("request-1")
err = guard.Reserve(filepath.Join(root, "audio", "request-1", "late.wav"), 1)
var releasedErr *EphemeralRequestReleasedError
Expect(errors.As(err, &releasedErr)).To(BeTrue())
Expect(releasedErr.RequestID).To(Equal("request-1"))
})
It("leaves a late commit recoverable when release times out", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "audio", "request-1", "late.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
Expect(guard.Reserve(path, 4)).To(Succeed())
ctx, cancel := context.WithCancel(context.Background())
cancel()
Expect(guard.BeginRequestRelease(ctx, "request-1")).To(MatchError(context.Canceled))
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
Expect(guard.HasActiveReservation(path)).To(BeFalse())
})
It("bounds release markers without reopening registered work", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.BeginRequestOperation("request-pinned")).To(Succeed())
pinnedRelease := make(chan error, 1)
go func() {
pinnedRelease <- guard.BeginRequestRelease(context.Background(), "request-pinned")
}()
Eventually(guard.releaseTombstoneCount).Should(Equal(1))
guard.EndRequestOperation("request-pinned")
Eventually(pinnedRelease).Should(Receive(Succeed()))
for index := range maxEphemeralReleaseTombstones + 10 {
requestID := fmt.Sprintf("request-%d", index)
Expect(guard.BeginRequestRelease(context.Background(), requestID)).To(Succeed())
guard.EndRequestRelease(requestID)
}
Expect(guard.releaseTombstoneCount()).To(Equal(maxEphemeralReleaseTombstones))
err = guard.Reserve(filepath.Join(root, "audio", "request-pinned", "late.wav"), 1)
var releasedErr *EphemeralRequestReleasedError
Expect(errors.As(err, &releasedErr)).To(BeTrue())
guard.EndRequestRelease("request-pinned")
})
It("applies backpressure at the release-pin cap and clears ownership", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "audio", "request-target", "input.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
Expect(guard.Reserve(path, 4)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
guard.mu.Lock()
for index := range maxEphemeralReleaseTombstones {
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
}
guard.mu.Unlock()
released := make(chan error, 1)
go func() {
released <- guard.BeginRequestRelease(context.Background(), "request-target")
}()
Consistently(released, 50*time.Millisecond).ShouldNot(Receive())
guard.EndRequestRelease("pinned-0")
Eventually(released).Should(Receive(Succeed()))
Expect(guard.HasActiveReservation(path)).To(BeFalse())
guard.EndRequestRelease("request-target")
})
It("makes committed files recoverable when pin backpressure expires", func() {
root := GinkgoT().TempDir()
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
path := filepath.Join(root, "audio", "request-target", "input.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
Expect(guard.Reserve(path, 4)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
guard.mu.Lock()
for index := range maxEphemeralReleaseTombstones {
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
}
guard.mu.Unlock()
ctx, cancel := context.WithCancel(context.Background())
cancel()
Expect(guard.BeginRequestRelease(ctx, "request-target")).To(MatchError(context.Canceled))
Expect(guard.HasActiveReservation(path)).To(BeFalse())
})
It("rejects a registered cache-hit claim after pin backpressure expires", func() {
root := GinkgoT().TempDir()
path := filepath.Join(root, "audio", "request-target", "input.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.BeginRequestOperation("request-target")).To(Succeed())
defer guard.EndRequestOperation("request-target")
guard.mu.Lock()
for index := range maxEphemeralReleaseTombstones {
guard.releasePins[fmt.Sprintf("pinned-%d", index)] = 1
}
guard.mu.Unlock()
ctx, cancel := context.WithCancel(context.Background())
cancel()
Expect(guard.BeginRequestRelease(ctx, "request-target")).To(MatchError(context.Canceled))
err = guard.Claim(path)
var releasedErr *EphemeralRequestReleasedError
Expect(errors.As(err, &releasedErr)).To(BeTrue())
Expect(guard.HasActiveReservation(path)).To(BeFalse())
})
})
+69 -17
View File
@@ -14,9 +14,9 @@ const (
// outlive the request that needed it. Inference reads these files while the
// request runs, so the window has to cover a slow multimodal request; it
// does not have to cover anything longer.
defaultEphemeralStagingTTL = 6 * time.Hour
defaultEphemeralStagingTTL = time.Hour
// defaultEphemeralStagingSweep is how often the worker sweeps.
defaultEphemeralStagingSweep = 30 * time.Minute
defaultEphemeralStagingSweep = 15 * time.Minute
)
// StartEphemeralStagingCleanup sweeps the worker's own staging directory for
@@ -32,6 +32,15 @@ func StartEphemeralStagingCleanup(ctx context.Context, stagingDir string, ttl, i
if stagingDir == "" {
return
}
StartEphemeralRootsCleanup(ctx, []string{filepath.Join(stagingDir, "ephemeral")}, nil, ttl, interval)
}
// StartEphemeralRootsCleanup removes abandoned request inputs for every worker
// transport while sharing accounting with live reservations.
func StartEphemeralRootsCleanup(ctx context.Context, roots []string, guard *EphemeralCapacityGuard, ttl, interval time.Duration) {
if len(roots) == 0 {
return
}
if ttl <= 0 {
ttl = defaultEphemeralStagingTTL
}
@@ -39,31 +48,41 @@ func StartEphemeralStagingCleanup(ctx context.Context, stagingDir string, ttl, i
interval = defaultEphemeralStagingSweep
}
// Reclaim crash leftovers before the caller starts accepting new work.
CleanEphemeralRoots(roots, ttl, guard)
go func() {
ticker := time.NewTicker(interval)
defer ticker.Stop()
// Sweep once at startup: a worker that crashed with staged files leaves
// them behind, and waiting a full interval to reclaim that space is the
// case that hurts on a volume that is already close to full.
CleanEphemeralStaging(stagingDir, ttl)
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
CleanEphemeralStaging(stagingDir, ttl)
CleanEphemeralRoots(roots, ttl, guard)
}
}
}()
xlog.Info("Ephemeral staging cleanup started", "dir", stagingDir, "ttl", ttl, "interval", interval)
xlog.Info("Ephemeral staging cleanup started", "roots", roots, "ttl", ttl, "interval", interval)
}
// CleanEphemeralStaging removes staged per-request directories older than ttl.
// It only ever descends into <stagingDir>/ephemeral, so staged model weights,
// which live alongside it and are not scratch, are never considered.
func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
root := filepath.Join(stagingDir, "ephemeral")
CleanEphemeralRoots([]string{filepath.Join(stagingDir, "ephemeral")}, ttl, nil)
}
// CleanEphemeralRoots removes stale request directories from explicit
// ephemeral roots. WalkDir never follows directory symlinks.
func CleanEphemeralRoots(roots []string, ttl time.Duration, guard *EphemeralCapacityGuard) {
for _, root := range roots {
cleanEphemeralRoot(root, ttl, guard)
}
}
func cleanEphemeralRoot(root string, ttl time.Duration, guard *EphemeralCapacityGuard) {
categories, err := os.ReadDir(root)
if err != nil {
// A worker that has never served a file-bearing request has no
@@ -87,19 +106,31 @@ func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
continue
}
for _, entry := range entries {
path := filepath.Join(categoryDir, entry.Name())
info, err := entry.Info()
if !entry.IsDir() || entry.Type()&os.ModeSymlink != 0 {
continue
}
requestPath := filepath.Join(categoryDir, entry.Name())
newest, err := newestEphemeralModTime(requestPath)
if err != nil {
xlog.Warn("Ephemeral staging cleanup: cannot stat entry", "path", path, "error", err)
xlog.Warn("Ephemeral staging cleanup: cannot inspect request", "path", requestPath, "error", err)
continue
}
// A request rewrites nothing after staging, so the entry's own
// modification time is when its request was served.
if !info.ModTime().Before(cutoff) {
if !newest.Before(cutoff) {
continue
}
if err := os.RemoveAll(path); err != nil {
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", path, "error", err)
if guard != nil {
removedTree, err := guard.RemoveTreeIfInactive(requestPath, func() error {
return os.RemoveAll(requestPath)
})
if err != nil {
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", requestPath, "error", err)
continue
}
if !removedTree {
continue
}
} else if err := os.RemoveAll(requestPath); err != nil {
xlog.Warn("Ephemeral staging cleanup: cannot remove", "path", requestPath, "error", err)
continue
}
removed++
@@ -110,3 +141,24 @@ func CleanEphemeralStaging(stagingDir string, ttl time.Duration) {
xlog.Info("Ephemeral staging cleanup removed stale request files", "count", removed, "dir", root)
}
}
func newestEphemeralModTime(root string) (time.Time, error) {
var newest time.Time
err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
info, err := entry.Info()
if err != nil {
return err
}
if info.ModTime().After(newest) {
newest = info.ModTime()
}
if entry.Type()&os.ModeSymlink != 0 && entry.IsDir() {
return filepath.SkipDir
}
return nil
})
return newest, err
}
@@ -55,4 +55,60 @@ var _ = Describe("Worker ephemeral staging cleanup", func() {
It("does nothing when no ephemeral directory exists", func() {
Expect(func() { CleanEphemeralStaging(stagingDir, time.Hour) }).ToNot(Panic())
})
It("sweeps both transport roots by newest descendant and skips active requests", func() {
cacheDir := GinkgoT().TempDir()
httpRoot := filepath.Join(stagingDir, "ephemeral")
s3Root := filepath.Join(cacheDir, "ephemeral")
guard, err := NewEphemeralCapacityGuard([]string{httpRoot, s3Root}, 8, 0)
Expect(err).NotTo(HaveOccurred())
staleRequest := filepath.Join(httpRoot, "audio", "stale")
activeRequest := filepath.Join(s3Root, "audio", "active")
freshChildRequest := filepath.Join(s3Root, "audio", "fresh-child")
for _, requestDir := range []string{staleRequest, activeRequest, freshChildRequest} {
Expect(os.MkdirAll(requestDir, 0o750)).To(Succeed())
Expect(os.WriteFile(filepath.Join(requestDir, "input.bin"), []byte("data"), 0o600)).To(Succeed())
}
old := time.Now().Add(-2 * time.Hour)
fresh := time.Now().Add(-5 * time.Minute)
for _, requestDir := range []string{staleRequest, activeRequest, freshChildRequest} {
Expect(os.Chtimes(requestDir, old, old)).To(Succeed())
}
Expect(os.Chtimes(filepath.Join(staleRequest, "input.bin"), old, old)).To(Succeed())
Expect(os.Chtimes(filepath.Join(activeRequest, "input.bin"), old, old)).To(Succeed())
Expect(os.Chtimes(filepath.Join(freshChildRequest, "input.bin"), fresh, fresh)).To(Succeed())
Expect(guard.Account(filepath.Join(staleRequest, "input.bin"), 4)).To(Succeed())
Expect(guard.Reserve(filepath.Join(activeRequest, "input.bin"), 4)).To(Succeed())
CleanEphemeralRoots([]string{httpRoot, s3Root}, time.Hour, guard)
Expect(staleRequest).NotTo(BeADirectory())
Expect(activeRequest).To(BeADirectory())
Expect(freshChildRequest).To(BeADirectory())
Expect(guard.Reserve(filepath.Join(httpRoot, "audio", "replacement", "input.bin"), 4)).To(Succeed())
})
It("keeps committed request inputs until exact release ends ownership", func() {
root := filepath.Join(stagingDir, "ephemeral")
requestDir := filepath.Join(root, "audio", "owned")
path := filepath.Join(requestDir, "input.bin")
Expect(os.MkdirAll(requestDir, 0o750)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 8, 0)
Expect(err).NotTo(HaveOccurred())
Expect(guard.Reserve(path, 4)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0o600)).To(Succeed())
Expect(guard.Commit(path)).To(Succeed())
old := time.Now().Add(-2 * time.Hour)
Expect(os.Chtimes(path, old, old)).To(Succeed())
Expect(os.Chtimes(requestDir, old, old)).To(Succeed())
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(requestDir).To(BeADirectory())
Expect(guard.Release(path)).To(Succeed())
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(requestDir).NotTo(BeADirectory())
})
})
+188
View File
@@ -1,8 +1,16 @@
package worker
import (
"context"
"errors"
"fmt"
"os"
"path"
"path/filepath"
"strings"
"github.com/mudler/LocalAI/core/services/storage"
"github.com/mudler/LocalAI/pkg/safefile"
)
// isPathAllowed checks if path is within one of the allowed directories.
@@ -31,3 +39,183 @@ func isPathAllowed(path string, allowedDirs []string) bool {
}
return false
}
func releaseEphemeralCacheKey(cacheDir, key string) error {
return releaseEphemeralCachePath(cacheDir, key, filepath.Join(cacheDir, filepath.FromSlash(key)))
}
func releaseEphemeralCachePath(cacheDir, key, filePath string) error {
return releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath, nil)
}
func releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath string, capacity *EphemeralCapacityGuard) error {
if err := validateEphemeralCacheKey(key); err != nil {
return err
}
relativePath := filepath.FromSlash(key)
expectedPath := filepath.Join(cacheDir, relativePath)
if filepath.Clean(filePath) != expectedPath {
return fmt.Errorf("release path %q does not match key %q", filePath, key)
}
if err := safefile.RemoveExact(cacheDir, relativePath, []string{".sha256", ".sha256.target"}, 2); err != nil {
return err
}
for _, path := range []string{filePath, filePath + ".sha256", filePath + ".sha256.target"} {
if capacity != nil {
if err := capacity.Release(path); err != nil {
return err
}
}
}
return nil
}
func releaseEphemeralCacheRequest(ctx context.Context, cacheDir, requestID string, capacity *EphemeralCapacityGuard) error {
if err := validateEphemeralCacheRequestID(requestID); err != nil {
return err
}
if capacity != nil {
if err := capacity.BeginRequestRelease(ctx, requestID); err != nil {
return fmt.Errorf("beginning release for request %q: %w", requestID, err)
}
defer capacity.EndRequestRelease(requestID)
}
root := filepath.Join(cacheDir, "ephemeral")
categories, err := os.ReadDir(root)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
var releaseErrors []error
for _, category := range categories {
if !category.IsDir() || category.Type()&os.ModeSymlink != 0 {
continue
}
requestDir := filepath.Join(root, category.Name(), requestID)
info, err := os.Lstat(requestDir)
if err != nil {
if !os.IsNotExist(err) {
releaseErrors = append(releaseErrors, fmt.Errorf("stating request directory %q: %w", requestDir, err))
}
continue
}
if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
releaseErrors = append(releaseErrors, fmt.Errorf("ephemeral request path %q is not a real directory", requestDir))
continue
}
entries, err := os.ReadDir(requestDir)
if err != nil {
if !os.IsNotExist(err) {
releaseErrors = append(releaseErrors, fmt.Errorf("reading request directory %q: %w", requestDir, err))
}
continue
}
for _, entry := range entries {
if entry.IsDir() {
releaseErrors = append(releaseErrors, fmt.Errorf("unexpected directory in ephemeral request %q", filepath.Join(requestDir, entry.Name())))
continue
}
key := filepath.ToSlash(filepath.Join("ephemeral", category.Name(), requestID, entry.Name()))
filePath := filepath.Join(cacheDir, filepath.FromSlash(key))
if err := releaseEphemeralCachePathWithCapacity(cacheDir, key, filePath, capacity); err != nil {
releaseErrors = append(releaseErrors, fmt.Errorf("releasing %q: %w", key, err))
}
}
}
return errors.Join(releaseErrors...)
}
type ephemeralStagingCapacity interface {
Reserve(path string, size int64) error
Commit(path string) error
Claim(path string) error
Release(path string) error
}
func ensureWorkerFile(ctx context.Context, fm *storage.FileManager, capacity *EphemeralCapacityGuard, key string) (string, error) {
if capacity == nil {
return fm.Download(ctx, key)
}
if strings.HasPrefix(key, "ephemeral/") {
if err := validateEphemeralCacheKey(key); err != nil {
return "", err
}
requestID := strings.Split(key, "/")[2]
if err := capacity.BeginRequestOperation(requestID); err != nil {
return "", err
}
defer capacity.EndRequestOperation(requestID)
}
return ensureWorkerFileWithCapacity(ctx, fm, capacity, key)
}
func ensureWorkerFileWithCapacity(ctx context.Context, fm *storage.FileManager, capacity ephemeralStagingCapacity, key string) (string, error) {
if capacity == nil || !strings.HasPrefix(key, "ephemeral/") {
return fm.Download(ctx, key)
}
if err := validateEphemeralCacheKey(key); err != nil {
return "", err
}
cachePath, err := fm.CachePath(key)
if err != nil {
return "", err
}
if info, statErr := os.Lstat(cachePath); statErr == nil {
if !info.Mode().IsRegular() {
return "", fmt.Errorf("ephemeral cache path %q is not a regular file", cachePath)
}
if err := capacity.Claim(cachePath); err != nil {
if !errors.Is(err, os.ErrNotExist) {
return "", err
}
} else {
return cachePath, nil
}
} else if !os.IsNotExist(statErr) {
return "", statErr
}
meta, err := fm.Head(ctx, key)
if err != nil {
return "", fmt.Errorf("reading size for %s: %w", key, err)
}
if err := capacity.Reserve(cachePath, meta.Size); err != nil {
return "", err
}
localPath, err := fm.Download(ctx, key)
if err != nil {
_ = capacity.Release(cachePath)
return "", err
}
if err := capacity.Commit(cachePath); err != nil {
_ = fm.EvictCache(key)
_ = capacity.Release(cachePath)
return "", err
}
return localPath, nil
}
func validateEphemeralCacheKey(key string) error {
if strings.Contains(key, "\\") || path.Clean(key) != key {
return fmt.Errorf("invalid ephemeral key %q", key)
}
parts := strings.Split(key, "/")
if len(parts) != 4 || parts[0] != "ephemeral" {
return fmt.Errorf("release key %q must identify one file below ephemeral/", key)
}
for _, part := range parts[1:] {
if part == "" || part == "." || part == ".." {
return fmt.Errorf("invalid ephemeral key %q", key)
}
}
return nil
}
func validateEphemeralCacheRequestID(requestID string) error {
if requestID == "" || strings.ContainsAny(requestID, "/\\") || path.Clean(requestID) != requestID || requestID == "." || requestID == ".." {
return fmt.Errorf("invalid ephemeral request ID %q", requestID)
}
return nil
}
@@ -0,0 +1,430 @@
package worker
import (
"context"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"os"
"path/filepath"
"strings"
"time"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/storage"
"github.com/mudler/LocalAI/core/services/workerctl"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"net/http"
"net/http/httptest"
)
type stagingObjectStore struct {
payload []byte
getCalls int
getErr error
}
type disappearingStagingCapacity struct{}
func (*disappearingStagingCapacity) Reserve(string, int64) error { return nil }
func (*disappearingStagingCapacity) Commit(string) error { return nil }
func (*disappearingStagingCapacity) Release(string) error { return nil }
func (*disappearingStagingCapacity) Claim(path string) error {
if err := os.Remove(path); err != nil {
return err
}
return fmt.Errorf("claim raced recovery: %w", os.ErrNotExist)
}
func (*stagingObjectStore) Put(context.Context, string, io.Reader) error { return nil }
func (s *stagingObjectStore) Get(context.Context, string) (io.ReadCloser, error) {
s.getCalls++
if s.getErr != nil {
return nil, s.getErr
}
return io.NopCloser(strings.NewReader(string(s.payload))), nil
}
func (s *stagingObjectStore) Head(_ context.Context, key string) (*storage.ObjectMeta, error) {
return &storage.ObjectMeta{Key: key, Size: int64(len(s.payload))}, nil
}
func (*stagingObjectStore) Exists(context.Context, string) (bool, error) { return true, nil }
func (*stagingObjectStore) Delete(context.Context, string) error { return nil }
func (*stagingObjectStore) List(context.Context, string) ([]string, error) {
return nil, nil
}
var _ = Describe("Worker exact-key staging release", func() {
It("shares capacity between HTTP uploads and S3 control requests and releases both over the worker transport", func() {
ctx := context.Background()
dir := GinkgoT().TempDir()
cfg := &Config{ModelsPath: filepath.Join(dir, "models")}
stagingDir := filepath.Join(dir, "staging")
cacheDir := cfg.stagingCacheDir()
guard, err := NewEphemeralCapacityGuard([]string{filepath.Join(stagingDir, "ephemeral"), filepath.Join(cacheDir, "ephemeral")}, 4+sha256.Size*4, 0)
Expect(err).NotTo(HaveOccurred())
store, err := storage.NewFilesystemStore(filepath.Join(dir, "objects"))
Expect(err).NotTo(HaveOccurred())
workerFM, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
frontendFM, err := storage.NewFileManager(store, filepath.Join(dir, "frontend-cache"))
Expect(err).NotTo(HaveOccurred())
sup := &backendSupervisor{cfg: cfg, processes: map[string]*backendProcess{}}
server, err := startWorkerHTTPServer("127.0.0.1:0", stagingDir, cfg.ModelsPath, cfg.stagingDataDir(), "secret", nil, sup, cfg, workerFM, nil, guard)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(func() { nodes.ShutdownFileTransferServer(server) })
dialFor := func(string) func(context.Context, string, string) (net.Conn, error) {
return func(ctx context.Context, _, _ string) (net.Conn, error) {
return (&net.Dialer{}).DialContext(ctx, "tcp", server.Addr)
}
}
httpStager := nodes.NewHTTPFileStager(func(string) (string, error) { return server.Addr, nil }, "secret", dialFor)
s3Stager := nodes.NewS3FileStager(frontendFM, nodes.NewControlClient(dialFor, "secret"))
input := filepath.Join(dir, "input.wav")
Expect(os.WriteFile(input, []byte("data"), 0o600)).To(Succeed())
s3Input := filepath.Join(dir, "s3-input.wav")
Expect(os.WriteFile(s3Input, []byte(strings.Repeat("x", 100)), 0o600)).To(Succeed())
httpKey := "ephemeral/audio/http-request/input.wav"
s3Key := "ephemeral/audio/s3-request/input.wav"
httpPath, err := httpStager.EnsureRemote(ctx, "worker", input, httpKey)
Expect(err).NotTo(HaveOccurred())
_, err = s3Stager.EnsureRemote(ctx, "worker", s3Input, s3Key)
Expect(err).To(HaveOccurred())
Expect(filepath.Join(cacheDir, filepath.FromSlash(s3Key))).NotTo(BeAnExistingFile())
Expect(httpStager.ReleaseRemoteRequest(ctx, "worker", "http-request", []string{httpKey})).To(Succeed())
Expect(httpPath).NotTo(BeAnExistingFile())
s3Path, err := s3Stager.EnsureRemote(ctx, "worker", s3Input, s3Key)
Expect(err).NotTo(HaveOccurred())
Expect(guard.HasActiveReservation(s3Path)).To(BeTrue())
Expect(s3Stager.ReleaseRemoteRequest(ctx, "worker", "s3-request", []string{s3Key})).To(Succeed())
Expect(s3Path).NotTo(BeAnExistingFile())
Expect(guard.HasActiveReservation(s3Path)).To(BeFalse())
exists, err := store.Exists(ctx, s3Key)
Expect(err).NotTo(HaveOccurred())
Expect(exists).To(BeFalse())
})
It("protects a startup-accounted HTTP cache hit through authenticated repeated probes", func() {
stagingDir := GinkgoT().TempDir()
root := filepath.Join(stagingDir, "ephemeral")
key := "ephemeral/audio/request-id/input.wav"
remotePath := filepath.Join(stagingDir, filepath.FromSlash(key))
content := []byte("data")
Expect(os.MkdirAll(filepath.Dir(remotePath), 0o750)).To(Succeed())
Expect(os.WriteFile(remotePath, content, 0o600)).To(Succeed())
hash := sha256.Sum256(content)
Expect(os.WriteFile(remotePath+".sha256", []byte(fmt.Sprintf("%x", hash)), 0o600)).To(Succeed())
old := time.Now().Add(-2 * time.Hour)
for _, path := range []string{remotePath, remotePath + ".sha256", filepath.Dir(remotePath)} {
Expect(os.Chtimes(path, old, old)).To(Succeed())
}
guard, err := NewEphemeralCapacityGuard([]string{root}, int64(len(content)+sha256.Size*2), 0)
Expect(err).NotTo(HaveOccurred())
listener, err := net.Listen("tcp", "127.0.0.1:0")
Expect(err).NotTo(HaveOccurred())
addr := listener.Addr().String()
Expect(listener.Close()).To(Succeed())
server, err := nodes.StartFileTransferServerWithCapacity(addr, stagingDir, GinkgoT().TempDir(), GinkgoT().TempDir(), "secret", 0, nil, guard)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(nodes.ShutdownFileTransferServer, server)
localPath := filepath.Join(GinkgoT().TempDir(), "input.wav")
Expect(os.WriteFile(localPath, content, 0o600)).To(Succeed())
stager := nodes.NewHTTPFileStager(func(string) (string, error) { return addr, nil }, "secret", func(string) func(context.Context, string, string) (net.Conn, error) {
return (&net.Dialer{}).DialContext
})
for range 2 {
path, ensureErr := stager.EnsureRemote(context.Background(), "worker", localPath, key)
Expect(ensureErr).NotTo(HaveOccurred())
Expect(path).To(Equal(remotePath))
}
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(remotePath).To(BeAnExistingFile())
Expect(stager.ReleaseRemote(context.Background(), "worker", key)).To(Succeed())
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(remotePath).NotTo(BeAnExistingFile())
})
It("claims a startup-scanned cache hit against stale recovery until release", func() {
cacheDir := GinkgoT().TempDir()
root := filepath.Join(cacheDir, "ephemeral")
key := "ephemeral/audio/request-id/input.wav"
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
Expect(os.WriteFile(cachePath, []byte("data"), 0o600)).To(Succeed())
old := time.Now().Add(-2 * time.Hour)
Expect(os.Chtimes(cachePath, old, old)).To(Succeed())
Expect(os.Chtimes(filepath.Dir(cachePath), old, old)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
Expect(err).NotTo(HaveOccurred())
store := &stagingObjectStore{payload: []byte("unused")}
fm, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
localPath, err := ensureWorkerFile(context.Background(), fm, guard, key)
Expect(err).NotTo(HaveOccurred())
Expect(localPath).To(Equal(cachePath))
Expect(store.getCalls).To(BeZero())
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(cachePath).To(BeAnExistingFile())
Expect(guard.Release(cachePath)).To(Succeed())
CleanEphemeralRoots([]string{root}, time.Hour, guard)
Expect(cachePath).NotTo(BeAnExistingFile())
})
It("downloads again when a cache file disappears while being claimed", func() {
cacheDir := GinkgoT().TempDir()
key := "ephemeral/audio/request-id/input.wav"
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
Expect(os.WriteFile(cachePath, []byte("stale"), 0o600)).To(Succeed())
store := &stagingObjectStore{payload: []byte("fresh")}
fm, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
localPath, err := ensureWorkerFileWithCapacity(context.Background(), fm, &disappearingStagingCapacity{}, key)
Expect(err).NotTo(HaveOccurred())
Expect(localPath).To(Equal(cachePath))
Expect(os.ReadFile(localPath)).To(Equal([]byte("fresh")))
Expect(store.getCalls).To(Equal(1))
})
It("makes repeated cache-hit claims idempotent", func() {
cacheDir := GinkgoT().TempDir()
root := filepath.Join(cacheDir, "ephemeral")
key := "ephemeral/audio/request-id/input.wav"
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
Expect(os.WriteFile(cachePath, []byte("data"), 0o600)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
Expect(err).NotTo(HaveOccurred())
fm, err := storage.NewFileManager(&stagingObjectStore{}, cacheDir)
Expect(err).NotTo(HaveOccurred())
for range 2 {
localPath, ensureErr := ensureWorkerFile(context.Background(), fm, guard, key)
Expect(ensureErr).NotTo(HaveOccurred())
Expect(localPath).To(Equal(cachePath))
}
err = guard.Reserve(filepath.Join(root, "other", "request-id", "input.wav"), 1)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.UsageBytes).To(Equal(int64(4)))
})
It("capacity-checks growth of a startup-scanned cache file", func() {
cacheDir := GinkgoT().TempDir()
root := filepath.Join(cacheDir, "ephemeral")
key := "ephemeral/audio/request-id/input.wav"
cachePath := filepath.Join(cacheDir, filepath.FromSlash(key))
Expect(os.MkdirAll(filepath.Dir(cachePath), 0o750)).To(Succeed())
Expect(os.WriteFile(cachePath, []byte("12"), 0o600)).To(Succeed())
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
Expect(err).NotTo(HaveOccurred())
Expect(os.WriteFile(cachePath, []byte("12345"), 0o600)).To(Succeed())
fm, err := storage.NewFileManager(&stagingObjectStore{}, cacheDir)
Expect(err).NotTo(HaveOccurred())
_, err = ensureWorkerFile(context.Background(), fm, guard, key)
var capacityErr *EphemeralCapacityError
Expect(errors.As(err, &capacityErr)).To(BeTrue())
Expect(capacityErr.RequestedBytes).To(Equal(int64(3)))
Expect(capacityErr.UsageBytes).To(Equal(int64(2)))
Expect(guard.HasActiveReservation(cachePath)).To(BeFalse())
})
It("reserves S3 object size before download and releases it with the exact key", func() {
cacheDir := GinkgoT().TempDir()
root := filepath.Join(cacheDir, "ephemeral")
store := &stagingObjectStore{payload: []byte("data")}
fm, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
Expect(err).NotTo(HaveOccurred())
key := "ephemeral/audio/request-id/input.wav"
localPath, err := ensureWorkerFile(context.Background(), fm, guard, key)
Expect(err).NotTo(HaveOccurred())
Expect(localPath).To(BeAnExistingFile())
Expect(store.getCalls).To(Equal(1))
Expect(guard.Reserve(filepath.Join(root, "audio", "other", "input.wav"), 1)).NotTo(Succeed())
Expect(releaseEphemeralCachePathWithCapacity(cacheDir, key, localPath, guard)).To(Succeed())
Expect(guard.Reserve(filepath.Join(root, "audio", "other", "input.wav"), 4)).To(Succeed())
})
It("rejects an oversized S3 object before starting its download", func() {
cacheDir := GinkgoT().TempDir()
store := &stagingObjectStore{payload: []byte("oversized")}
fm, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
guard, err := NewEphemeralCapacityGuard([]string{filepath.Join(cacheDir, "ephemeral")}, 4, 0)
Expect(err).NotTo(HaveOccurred())
_, err = ensureWorkerFile(context.Background(), fm, guard, "ephemeral/audio/request-id/input.wav")
Expect(err).To(HaveOccurred())
Expect(store.getCalls).To(BeZero())
})
It("rolls back an S3 reservation when the download fails", func() {
cacheDir := GinkgoT().TempDir()
root := filepath.Join(cacheDir, "ephemeral")
store := &stagingObjectStore{payload: []byte("data"), getErr: errors.New("download failed")}
fm, err := storage.NewFileManager(store, cacheDir)
Expect(err).NotTo(HaveOccurred())
guard, err := NewEphemeralCapacityGuard([]string{root}, 4, 0)
Expect(err).NotTo(HaveOccurred())
_, err = ensureWorkerFile(context.Background(), fm, guard, "ephemeral/audio/request-id/input.wav")
Expect(err).To(MatchError(ContainSubstring("download failed")))
Expect(guard.Reserve(filepath.Join(root, "audio", "replacement", "input.wav"), 4)).To(Succeed())
})
It("removes only the exact cache file and upload sidecars", func() {
cacheDir := GinkgoT().TempDir()
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
target := filepath.Join(categoryDir, "input.wav")
sibling := filepath.Join(categoryDir, "keep.wav")
for _, path := range []string{target, target + ".sha256", target + ".sha256.target", sibling} {
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
}
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/input.wav")).To(Succeed())
Expect(target).NotTo(BeAnExistingFile())
Expect(target + ".sha256").NotTo(BeAnExistingFile())
Expect(target + ".sha256.target").NotTo(BeAnExistingFile())
Expect(sibling).To(BeAnExistingFile())
Expect(categoryDir).To(BeADirectory())
})
It("succeeds for a missing file and prunes empty category and request directories", func() {
cacheDir := GinkgoT().TempDir()
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
for range 2 {
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/missing.wav")).To(Succeed())
}
Expect(categoryDir).NotTo(BeADirectory())
Expect(filepath.Dir(categoryDir)).NotTo(BeADirectory())
Expect(filepath.Join(cacheDir, "ephemeral")).To(BeADirectory())
})
It("rejects traversal and symlink escapes", func() {
cacheDir := GinkgoT().TempDir()
outsideDir := GinkgoT().TempDir()
outsidePath := filepath.Join(outsideDir, "input.wav")
Expect(os.WriteFile(outsidePath, []byte("keep"), 0640)).To(Succeed())
requestDir := filepath.Join(cacheDir, "ephemeral", "request-id")
Expect(os.MkdirAll(requestDir, 0750)).To(Succeed())
Expect(os.Symlink(outsideDir, filepath.Join(requestDir, "audio"))).To(Succeed())
for _, key := range []string{
"models/model.gguf",
"ephemeral/../models/model.gguf",
"ephemeral/request-id/audio/../../model.gguf",
"ephemeral/request-id/audio/input.wav",
} {
Expect(releaseEphemeralCacheKey(cacheDir, key)).NotTo(Succeed(), key)
}
Expect(outsidePath).To(BeAnExistingFile())
})
It("rejects symlinked files and sidecars without deleting their targets", func() {
for _, linkedName := range []string{"input.wav", "input.wav.sha256", "input.wav.sha256.target"} {
cacheDir := GinkgoT().TempDir()
categoryDir := filepath.Join(cacheDir, "ephemeral", "request-id", "audio")
Expect(os.MkdirAll(categoryDir, 0750)).To(Succeed())
target := filepath.Join(categoryDir, "input.wav")
if linkedName != "input.wav" {
Expect(os.WriteFile(target, []byte("input"), 0640)).To(Succeed())
}
preserved := filepath.Join(cacheDir, "ephemeral", "preserved-"+linkedName)
Expect(os.WriteFile(preserved, []byte("keep"), 0640)).To(Succeed())
Expect(os.Symlink(preserved, filepath.Join(categoryDir, linkedName))).To(Succeed())
Expect(releaseEphemeralCacheKey(cacheDir, "ephemeral/request-id/audio/input.wav")).NotTo(Succeed(), linkedName)
Expect(preserved).To(BeAnExistingFile(), linkedName)
}
})
It("registers an exact release handler", func() {
cacheDir := GinkgoT().TempDir()
path := filepath.Join(cacheDir, "ephemeral", "request-id", "audio", "input.wav")
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
fm, err := storage.NewFileManager(nil, cacheDir)
Expect(err).NotTo(HaveOccurred())
mux := http.NewServeMux()
registerFileReleaseControlRoute(mux, fm, cacheDir, nil)
request, err := json.Marshal(map[string]string{"key": "ephemeral/request-id/audio/input.wav"})
Expect(err).NotTo(HaveOccurred())
response := httptest.NewRecorder()
mux.ServeHTTP(response, httptest.NewRequest(http.MethodPost, workerctl.PathFilesRelease, strings.NewReader(string(request))))
Expect(response.Code).To(Equal(http.StatusOK))
var reply map[string]string
Expect(json.Unmarshal(response.Body.Bytes(), &reply)).To(Succeed())
Expect(reply["error"]).To(BeEmpty())
Expect(path).NotTo(BeAnExistingFile())
})
It("releases a request batch through one worker control request", func() {
cacheDir := GinkgoT().TempDir()
keys := []string{
"ephemeral/audio/request-id/input.wav",
"ephemeral/images/request-id/frame.jpg",
}
for _, key := range keys {
path := filepath.Join(cacheDir, filepath.FromSlash(key))
Expect(os.MkdirAll(filepath.Dir(path), 0750)).To(Succeed())
Expect(os.WriteFile(path, []byte("data"), 0640)).To(Succeed())
}
fm, err := storage.NewFileManager(nil, cacheDir)
Expect(err).NotTo(HaveOccurred())
mux := http.NewServeMux()
registerFileReleaseControlRoute(mux, fm, cacheDir, nil)
request, err := json.Marshal(map[string]any{"request_id": "request-id"})
Expect(err).NotTo(HaveOccurred())
response := httptest.NewRecorder()
mux.ServeHTTP(response, httptest.NewRequest(http.MethodPost, workerctl.PathFilesRelease, strings.NewReader(string(request))))
Expect(response.Code).To(Equal(http.StatusOK))
var reply map[string]string
Expect(json.Unmarshal(response.Body.Bytes(), &reply)).To(Succeed())
Expect(reply["error"]).To(BeEmpty())
for _, key := range keys {
Expect(filepath.Join(cacheDir, filepath.FromSlash(key))).NotTo(BeAnExistingFile())
}
})
It("returns validation errors through the release handler", func() {
cacheDir := GinkgoT().TempDir()
fm, err := storage.NewFileManager(nil, cacheDir)
Expect(err).NotTo(HaveOccurred())
mux := http.NewServeMux()
registerFileReleaseControlRoute(mux, fm, cacheDir, nil)
request, err := json.Marshal(map[string]string{"key": "models/model.gguf"})
Expect(err).NotTo(HaveOccurred())
response := httptest.NewRecorder()
mux.ServeHTTP(response, httptest.NewRequest(http.MethodPost, workerctl.PathFilesRelease, strings.NewReader(string(request))))
Expect(response.Code).To(Equal(http.StatusOK))
var reply map[string]string
Expect(json.Unmarshal(response.Body.Bytes(), &reply)).To(Succeed())
Expect(reply["error"]).NotTo(BeEmpty())
})
})
+32 -13
View File
@@ -138,11 +138,26 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
// the top of Run so the worker fails before registering.)
httpAddr := cfg.resolveHTTPAddr()
stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging")
// Derived through the same helper the listdir verb resolves `data/` keys
// against, and not a second time here. Two independent joins that agreed
// today would each stay self consistent if one moved, and the symptom would
// be a verb that lists files the file server does not serve.
cacheDir := cfg.stagingCacheDir()
dataDir := cfg.stagingDataDir()
ephemeralRoots := []string{
filepath.Join(stagingDir, "ephemeral"),
filepath.Join(cacheDir, "ephemeral"),
}
byteLimit, minFreeBytes, err := effectiveEphemeralCapacity(
ephemeralRoots,
cfg.EphemeralStagingByteLimit,
cfg.EphemeralStagingMinFreeBytes,
)
if err != nil {
return fmt.Errorf("resolving ephemeral staging capacity: %w", err)
}
ephemeralCapacity, err := NewEphemeralCapacityGuard(ephemeralRoots, byteLimit, minFreeBytes)
if err != nil {
return fmt.Errorf("initializing ephemeral staging capacity: %w", err)
}
xlog.Info("Ephemeral staging capacity configured", "roots", ephemeralRoots, "byteLimit", byteLimit, "minFreeBytes", minFreeBytes)
StartEphemeralRootsCleanup(shutdownCtx, ephemeralRoots, ephemeralCapacity, 0, 0)
// The readiness gate is created here but only armed once the tunnel exists,
// below. Until then /readyz reports ready, which is correct: reaching this
// line means the worker has already registered with the frontend, so it is
@@ -201,15 +216,11 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
}
httpServer, err := startWorkerHTTPServer(httpAddr, stagingDir, cfg.ModelsPath, dataDir,
cfg.RegistrationToken, readiness, supervisor, cfg, stagingFM, ml.BackendLogs())
cfg.RegistrationToken, readiness, supervisor, cfg, stagingFM, ml.BackendLogs(), ephemeralCapacity)
if err != nil {
return fmt.Errorf("starting HTTP file transfer server: %w", err)
}
// Per-request input files land in stagingDir over that server and nothing
// used to remove them, so a long-lived worker filled its own disk.
StartEphemeralStagingCleanup(shutdownCtx, stagingDir, 0, 0)
// The tunnel is started here, after the HTTP server it fronts is listening
// and before any backend process exists. Both orders are deliberate: a
// stream tagged for HTTP that arrived before the server bound would be
@@ -336,9 +347,17 @@ func startTunnelAndArmReadiness(ctx context.Context, readiness *nodes.WorkerRead
// healthy while answering 404 to every command.
func startWorkerHTTPServer(addr, stagingDir, modelsDir, dataDir, token string,
readiness *nodes.WorkerReadiness, sup *backendSupervisor, cfg *Config,
stagingFM *storage.FileManager, logStore *model.BackendLogStore) (*http.Server, error) {
return nodes.StartFileTransferServer(addr, stagingDir, modelsDir, dataDir, token,
config.DefaultMaxUploadSize, readiness, &nodes.AuthenticatedRoutes{
stagingFM *storage.FileManager, logStore *model.BackendLogStore, capacities ...*EphemeralCapacityGuard) (*http.Server, error) {
var capacity *EphemeralCapacityGuard
if len(capacities) > 0 {
capacity = capacities[0]
}
var httpCapacity nodes.EphemeralCapacity
if capacity != nil {
httpCapacity = capacity
}
return nodes.StartFileTransferServerWithCapacityAndRoutes(addr, stagingDir, modelsDir, dataDir, token,
config.DefaultMaxUploadSize, readiness, httpCapacity, &nodes.AuthenticatedRoutes{
Prefix: workerctl.Prefix,
// One registrar for both route sets, because there is ONE control
// prefix and AuthenticatedRoutes mounts one mux behind one bearer
@@ -346,7 +365,7 @@ func startWorkerHTTPServer(addr, stagingDir, modelsDir, dataDir, token string,
Register: func(mux *http.ServeMux) {
sup.RegisterControlRoutes(mux)
if stagingFM != nil {
cfg.RegisterFileControlRoutes(mux, stagingFM)
cfg.RegisterFileControlRoutesWithCapacity(mux, stagingFM, capacity)
}
},
}, logStore)
+2
View File
@@ -36,6 +36,7 @@ const (
// them not at all, and the catch-all under Prefix answers for them. That is
// the same 404 an older build gives, which is what the frontend already
// reads as "this worker does not serve that verb" rather than as absence.
PathFilesRelease = "/v1/control/files/release"
PathFilesEnsure = "/v1/control/files/ensure"
PathFilesStage = "/v1/control/files/stage"
PathFilesTemp = "/v1/control/files/temp"
@@ -73,6 +74,7 @@ func BackendPaths() []string {
PathModelDelete,
PathModelsRunning,
PathNodeStop,
PathFilesRelease,
PathFilesEnsure,
PathFilesStage,
PathFilesTemp,
+2
View File
@@ -28,6 +28,7 @@ var _ = Describe("control plane paths on the wire", func() {
Entry("model delete", workerctl.PathModelDelete, "/v1/control/model/delete"),
Entry("models running", workerctl.PathModelsRunning, "/v1/control/models/running"),
Entry("node stop", workerctl.PathNodeStop, "/v1/control/node/stop"),
Entry("files release", workerctl.PathFilesRelease, "/v1/control/files/release"),
Entry("files ensure", workerctl.PathFilesEnsure, "/v1/control/files/ensure"),
Entry("files stage", workerctl.PathFilesStage, "/v1/control/files/stage"),
Entry("files temp", workerctl.PathFilesTemp, "/v1/control/files/temp"),
@@ -66,6 +67,7 @@ var _ = Describe("control plane paths on the wire", func() {
workerctl.PathModelDelete,
workerctl.PathModelsRunning,
workerctl.PathNodeStop,
workerctl.PathFilesRelease,
workerctl.PathFilesEnsure,
workerctl.PathFilesStage,
workerctl.PathFilesListDir,
+1 -1
View File
@@ -739,7 +739,7 @@ For image generation models using the `diffusers` backend:
| Field | Type | Description |
|-------|------|-------------|
| `diffusers.cuda` | bool | Enable CUDA for diffusers |
| `diffusers.cuda` | bool | Force CUDA. By default the backend auto-detects and uses CUDA when a compatible GPU is present (ROCm builds included). Pin the CPU with `options: ["device:cpu"]` |
| `diffusers.pipeline_type` | string | Pipeline type (e.g., `stable-diffusion`, `stable-diffusion-xl`) |
| `diffusers.scheduler_type` | string | Scheduler type (e.g., `euler`, `ddpm`) |
| `diffusers.enable_parameters` | string | Comma-separated parameters to enable |
+5
View File
@@ -28,6 +28,11 @@ key stored *inside* the GGUF, which cannot be read from a remote repository, and
upstream GGUF repository hosts every family in one place. Set `backend: audio-cpp`
in the model YAML, or select it explicitly in the import form.
The bundled audio-cpp gallery entries set `backend:best` in `options` to select
an available compute backend, with CPU as the fallback. To force CPU execution,
replace that option with `backend:cpu`. Model configurations that omit this
option still default to CPU.
## What it serves
One model serves one family, and a family advertises the tasks it can perform. The
+13 -2
View File
@@ -818,6 +818,8 @@ There is no broker flag here. A serve-backend worker dials one outbound tunnel t
| `--serve-addr` | `LOCALAI_SERVE_ADDR` | `0.0.0.0:50051` | Same, used when `--addr` is unset |
| `--grpc-max-port` | `LOCALAI_GRPC_MAX_PORT` | `65535` | Highest port the worker may assign to a backend gRPC process. Each backend gets its own port, allocated upward from the base port, so the width of `[base port, this]` caps how many backends this worker can run at once (see [Backend gRPC port range](#backend-grpc-port-range)) |
| `--http-addr` | `LOCALAI_HTTP_ADDR` | `127.0.0.1:{gRPC port - 1}` | HTTP file transfer server bind address |
| `--ephemeral-staging-byte-limit` | `LOCALAI_EPHEMERAL_STAGING_BYTE_LIMIT` | `0` (automatic) | Maximum bytes held by request-input staging across the worker's HTTP staging directory and S3 cache. Automatic mode uses the smaller of 10 GiB and 10% of filesystem capacity. |
| `--ephemeral-staging-min-free-bytes` | `LOCALAI_EPHEMERAL_STAGING_MIN_FREE_BYTES` | `0` (automatic) | Free filesystem space preserved while staging request inputs. Automatic mode uses the larger of 1 GiB and 5% of filesystem capacity. |
| `--register-to` | `LOCALAI_REGISTER_TO` | *(required)* | Frontend URL for self-registration |
| `--node-name` | `LOCALAI_NODE_NAME` | hostname | Human-readable node name |
| `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token to authenticate with the frontend |
@@ -835,6 +837,14 @@ There is no broker flag here. A serve-backend worker dials one outbound tunnel t
**HTTP file transfer:** Each worker also runs a small HTTP server for file transfer (model files, configs). It listens on loopback at the gRPC base port - 1 (e.g., if gRPC base is 50051, HTTP is on 50050). gRPC ports grow upward from the base port as additional models are loaded.
{{% /notice %}}
### Ephemeral request-input storage
Workers reserve local capacity before accepting per-request audio, image, and other ephemeral inputs. The limit covers both direct HTTP staging and the worker's S3 download cache. A request is rejected before inference when accepting its input would exceed the byte limit or the configured free-space headroom. One request-scoped cleanup operation releases all exact input keys and their reservations after inference, while a one-hour recovery sweep removes abandoned files after crashes. The sweep runs at startup and every 15 minutes, preserves active requests, and considers the newest file in each request directory.
For S3 staging, the frontend releases request inputs with `POST /v1/control/files/release` through the worker tunnel. The worker applies the same capacity limits to `files/ensure` downloads and direct HTTP uploads.
Set both capacity variables to positive byte counts when a worker needs fixed limits. Leaving either value at zero selects its filesystem-based default. These settings apply only below the two `ephemeral` roots; model, data, and configuration files are excluded.
### Worker Health Probes
The worker's HTTP server (loopback, base port - 1, default 50050) exposes two unauthenticated probes. They are reachable from the worker host - which is where a container healthcheck runs - and not from the network:
@@ -1777,8 +1787,9 @@ Notes:
- Older releases decided absence from `nats: no responders available for request`, which was one frontend's observation that nobody answered *it* within a request budget. Two replicas asking at the same moment could disagree and demote each other's workers. That signal is gone from the scheduler, and no component opens a bus connection to produce it.
**A worker fills its own disk over time:**
- A request that carries a file (an image, an audio clip, a video) stages that file to the worker under `<models>/../staging/ephemeral/`. The worker deletes these 6 hours after the request that needed them, and sweeps every 30 minutes plus once at startup, so a worker that crashed mid-request still reclaims the space.
- Releases before this sweep existed kept every staged input for the lifetime of the worker. Delete `<models>/../staging/ephemeral/` on an affected worker once, as the user the worker runs as; the sweep keeps it bounded from then on.
- A request that carries a file (an image, an audio clip, a video) stages that file below the worker's HTTP staging or S3 cache `ephemeral/` directory. The frontend releases each request-owned input when inference finishes, and the worker reserves capacity before accepting it.
- A one-hour recovery sweep runs at startup and every 15 minutes to reclaim inputs left by interrupted requests. It preserves active reservations and uses the newest file timestamp in each request directory.
- Releases before request-owned cleanup existed can leave a legacy backlog. Delete the affected `ephemeral/` directory once, as the user the worker runs as; capacity admission and recovery cleanup keep new staging bounded.
- Staged **model** files are not touched by this. They live beside the ephemeral directory and are not per-request scratch.
- A worker whose volume is genuinely full reports `creating backend process state directory under ...: no space left on device` when a backend starts.
@@ -0,0 +1,158 @@
# Request-owned ephemeral staging
## Problem
Distributed requests copy transient inputs below
`<staging>/ephemeral/<category>/<request-id>`. The worker currently removes
these files only when a periodic age sweep considers them stale. A Reachy Mini
sending camera and sound data about once per second created more than 21,000
request directories and filled its Mac worker before the six-hour retention
window elapsed.
Reducing the retention window is insufficient. A time limit bounds residence
time, but the retained bytes still scale with request rate and input size. A
quota sweeper would also have to infer whether an old file is still in use.
Neither rule prevents concurrent uploads from consuming the worker's last free
space.
## Goals
- Give every ephemeral input an explicit owner and release it when that request
finishes, fails, or is cancelled.
- Keep cleanup transport-independent for HTTP and S3/NATS workers.
- Reserve capacity before accepting ephemeral bytes so concurrent requests
cannot consume configured disk headroom.
- Reject a request cleanly when its input does not fit; never evict an input
that a running request may still be reading.
- Recover abandoned files after frontend or worker crashes.
- Never inspect or remove models, data, configuration, or paths outside the
worker's ephemeral staging tree.
## Non-goals
- Retaining request inputs as a cache.
- Evicting persistent model or data files to make an inference request fit.
- Treating modification timestamps as proof that a request is active.
## Request ownership
The `FileStagingClient` already creates one request ID before staging inputs and
waits for synchronous and streaming backend calls to finish. It will track each
ephemeral key before attempting to stage it and defer one request-scoped
release identified by the request ID. Release runs after the backend call
returns, including error and cancellation paths, using a short background
timeout so cancellation of the request does not cancel its cleanup.
Request IDs will use the full UUID rather than the current eight-character
prefix. The worker enumerates only category directories for that validated
request ID and removes each entry with exact, symlink-safe deletion.
`FileStager` will expose an idempotent exact-key `ReleaseRemote` operation and
an optional request-scoped operation. The client uses one fixed-size request
message for the normal path and retains exact-key calls as a rolling-upgrade
fallback:
- HTTP sends one authenticated request containing the fixed-size request ID.
The worker derives and removes that request's exact files, then prunes empty
request and category directories without following symlinks.
- S3/NATS sends one request-reply containing the request ID so the selected
worker evicts the request's local cached files. The frontend then deletes the
matching objects from its tracked exact-key list.
Either deletion may already have happened and still counts as success.
If staging fails partway through a request, the deferred release still includes
the planned key, allowing it to remove a partial file when the transport can
identify one. Cleanup errors are logged and do not replace the inference result.
HTTP and S3 ingress register request operations before any pre-reservation
work. Before enumerating files, the capacity guard marks the request released
and waits for registered operations and admitted writes to finish. Later
operations, reservations, and cache claims for that request are rejected.
Markers expire after one hour and are capped at 16,384 entries, but a marker is
never evicted while its registered operation or cleanup scan is active.
Concurrent operation and cleanup state have the same hard cap. Disk bytes
remain independently bounded by capacity admission.
Cleanup waits within its deadline when all cleanup-pin slots are occupied.
If that deadline expires, existing entries lose active ownership so recovery
can reclaim them; registered ingress for the request remains closed until it
exits.
## Capacity admission
A worker-local ephemeral capacity guard is shared by its HTTP and S3/NATS input
paths. It accounts for both `<staging>/ephemeral`, used by HTTP, and
`<cache>/ephemeral`, used by S3 downloads. Before writing an ephemeral object,
the transport reserves its declared size. HTTP obtains the size from the upload
metadata; S3/NATS obtains it from object metadata. Reservations are serialized
in memory, cover both committed ephemeral bytes and concurrent writes, and are
returned on release or failed transfer.
Admission succeeds only when both conditions remain true after the reservation:
1. Total ephemeral bytes remain below the configured ephemeral staging limit.
2. The filesystem retains the configured minimum free-space headroom.
The guard rejects the transfer before inference when either condition fails.
An input with unknown size is written through a bounded accounting writer that
reserves fixed-size chunks before writing each chunk and stops before crossing
the limit. The existing maximum-upload-size check remains the per-file ceiling.
The limit and headroom are worker settings. By default, ephemeral data may use
the smaller of 10 GiB or 10 percent of filesystem capacity, while the worker
preserves the larger of 1 GiB or 5 percent as free-space headroom. The worker
logs the effective values at startup. A zero or negative operator value selects
the default rather than disabling protection. The guard scans the ephemeral
tree at startup to account for abandoned committed bytes. Filesystem free-space
checks are repeated at reservation time because other processes may share the
volume.
## Crash recovery
The existing periodic cleanup remains as a fallback for ownership messages lost
when a frontend or worker process dies. It uses a one-hour recovery TTL,
performs one startup sweep, and repeats every 15 minutes. It skips every key
held by an active reservation, considers the newest modification time in each
remaining request tree, and does not follow directory symlinks. It removes only
request directories below the registered `<staging>/ephemeral` and
`<cache>/ephemeral` roots.
The recovery window does not control normal storage growth. Request completion
and capacity reservations do. A recovery deletion updates the capacity guard's
accounted bytes.
## Error handling and observability
Admission failures report the requested bytes, current ephemeral usage, limit,
available bytes, and required headroom. Successful release and recovery update
usage counters. Read, stat, and remove failures include the affected path and
allow unrelated cleanup to continue. Missing ephemeral files and directories
are normal for idempotent release.
## Testing
Regression tests will establish the following behavior:
1. Successful, failed, cancelled, and streaming calls issue one request-scoped
worker cleanup only after the backend has returned.
2. Partial staging failures release the planned key without changing the main
error returned to the caller.
3. HTTP and S3/NATS release remove local files; S3/NATS also removes the object.
4. Release rejects persistent keys and path traversal, does not follow
symlinks, and leaves paths outside `ephemeral` untouched.
5. Concurrent reservations cannot exceed the byte limit or free-space
headroom, and failed transfers return their reservations.
6. Unknown-length writes stop at the capacity boundary.
7. Startup accounting includes abandoned ephemeral files, and the recovery
sweep removes only stale, inactive leftovers and updates accounting.
Focused package tests will run with race detection, followed by the relevant
repository lint and vet checks.
## Rollout
The change requires a new LocalAI worker and frontend build because both sides
participate in release. The Mac worker starts by accounting for its existing
backlog and removing recovery-expired files. The deployment check will verify
available space, admission and release logs, stable ephemeral usage under
continuous camera and audio traffic, and successful vision, sound detection,
and transcription requests.
+26
View File
@@ -56579,6 +56579,8 @@
known_usecases:
- tts
name: audio-cpp-indextts-2.5
options:
- backend:best
parameters:
model: audio-cpp/index-tts2_5-orig.gguf
files:
@@ -56616,6 +56618,8 @@
known_usecases:
- tts
name: audio-cpp-supertonic
options:
- backend:best
parameters:
model: audio-cpp/supertonic-3-orig.gguf
files:
@@ -56658,6 +56662,8 @@
known_usecases:
- tts
name: audio-cpp-higgs-audio-v3
options:
- backend:best
parameters:
model: audio-cpp/higgs-audio-v3-tts-4b-q8_0.gguf
files:
@@ -56707,6 +56713,8 @@
- tts
- audio_transform
name: audio-cpp-chatterbox
options:
- backend:best
parameters:
model: audio-cpp/chatterbox-q8_0.gguf
files:
@@ -56743,6 +56751,8 @@
known_usecases:
- transcript
name: audio-cpp-citrinet-asr
options:
- backend:best
parameters:
model: audio-cpp/citrinet-asr-q8_0.gguf
files:
@@ -56779,6 +56789,8 @@
known_usecases:
- transcript
name: audio-cpp-nemotron-asr
options:
- backend:best
parameters:
model: audio-cpp/nemotron-3.5-asr-streaming-0.6b-f16.gguf
files:
@@ -56816,6 +56828,8 @@
known_usecases:
- diarization
name: audio-cpp-sortformer-diarization
options:
- backend:best
parameters:
model: audio-cpp/sortformer-diar-4spk-v1-q8_0.gguf
files:
@@ -56853,6 +56867,8 @@
known_usecases:
- audio_transform
name: audio-cpp-htdemucs
options:
- backend:best
parameters:
model: audio-cpp/htdemucs-f16.gguf
files:
@@ -56892,6 +56908,8 @@
known_usecases:
- sound_generation
name: audio-cpp-stable-audio-sfx
options:
- backend:best
parameters:
model: audio-cpp/stable-audio-3-small-sfx-q8_0.gguf
files:
@@ -56930,6 +56948,8 @@
known_usecases:
- transcript
name: audio-cpp-forced-aligner
options:
- backend:best
parameters:
language: en
model: audio-cpp/qwen3-forced-aligner-0.6b-q8_0.gguf
@@ -56959,6 +56979,7 @@
- vad
name: audio-cpp-silero-vad
options:
- backend:best
- family:silero_vad
parameters:
model: bundled:silero_vad
@@ -56986,6 +57007,7 @@
- vad
name: audio-cpp-marblenet-vad
options:
- backend:best
- family:marblenet_vad
parameters:
model: bundled:marblenet_vad
@@ -57025,6 +57047,8 @@
known_usecases:
- tts
name: audio-cpp-irodori-voicedesign
options:
- backend:best
parameters:
model: audio-cpp/irodori-tts-600m-v3-voicedesign-q8_0.gguf
files:
@@ -57069,6 +57093,7 @@
- audio_transform
name: audio-cpp-seedvc-singing
options:
- backend:best
- task:svc
parameters:
model: audio-cpp/seed-vc-mlx-q8_0.gguf
@@ -57121,6 +57146,7 @@
- audio_transform
name: audio-cpp-vevo2-speech-to-speech
options:
- backend:best
- task:s2s
parameters:
model: audio-cpp/vevo2-q8_0.gguf
+17
View File
@@ -0,0 +1,17 @@
//go:build !aix && !darwin && !dragonfly && !freebsd && !linux && !netbsd && !openbsd && !solaris
package safefile
import (
"errors"
"fmt"
)
// ErrUnsafePath reports a path shape or file type that exact removal refuses.
var ErrUnsafePath = errors.New("unsafe removal path")
// RemoveExact fails closed on platforms without component-relative no-follow
// filesystem operations.
func RemoveExact(root, relativePath string, sidecarSuffixes []string, pruneParents int) error {
return fmt.Errorf("secure exact removal is unsupported on this platform")
}
+144
View File
@@ -0,0 +1,144 @@
//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
package safefile
import (
"errors"
"fmt"
"path/filepath"
"strings"
"golang.org/x/sys/unix"
)
// ErrUnsafePath reports a path shape or file type that exact removal refuses.
var ErrUnsafePath = errors.New("unsafe removal path")
// RemoveExact removes a file and its named sidecars below root without
// following symbolic links. It prunes up to pruneParents empty parent
// directories, but never removes root itself.
func RemoveExact(root, relativePath string, sidecarSuffixes []string, pruneParents int) error {
return removeExact(root, relativePath, sidecarSuffixes, pruneParents, nil)
}
func removeExact(root, relativePath string, sidecarSuffixes []string, pruneParents int, parentsOpened func()) error {
parts, err := cleanRelativeParts(relativePath)
if err != nil {
return err
}
if pruneParents < 0 {
return fmt.Errorf("%w: prune parent count must not be negative", ErrUnsafePath)
}
rootFD, err := unix.Open(root, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
if err != nil {
return fmt.Errorf("opening removal root %q: %w", root, err)
}
handles := []int{rootFD}
defer func() {
for i := len(handles) - 1; i >= 0; i-- {
_ = unix.Close(handles[i])
}
}()
for _, component := range parts[:len(parts)-1] {
fd, openErr := unix.Openat(handles[len(handles)-1], component, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
if errors.Is(openErr, unix.ENOENT) {
return nil
}
if openErr != nil {
if errors.Is(openErr, unix.ELOOP) || errors.Is(openErr, unix.ENOTDIR) {
return fmt.Errorf("%w: path component %q is not a directory: %v", ErrUnsafePath, component, openErr)
}
return fmt.Errorf("opening removal path component %q: %w", component, openErr)
}
handles = append(handles, fd)
}
if parentsOpened != nil {
parentsOpened()
}
parentFD := handles[len(handles)-1]
leaf := parts[len(parts)-1]
names := make([]string, 0, len(sidecarSuffixes)+1)
names = append(names, leaf)
for _, suffix := range sidecarSuffixes {
if suffix == "" || strings.ContainsAny(suffix, `/\\`) {
return fmt.Errorf("%w: invalid sidecar suffix %q", ErrUnsafePath, suffix)
}
names = append(names, leaf+suffix)
}
for _, name := range names {
var stat unix.Stat_t
statErr := unix.Fstatat(parentFD, name, &stat, unix.AT_SYMLINK_NOFOLLOW)
if errors.Is(statErr, unix.ENOENT) {
continue
}
if statErr != nil {
return fmt.Errorf("stating removal entry %q: %w", name, statErr)
}
switch stat.Mode & unix.S_IFMT {
case unix.S_IFLNK:
return fmt.Errorf("%w: refusing to remove symbolic link %q", ErrUnsafePath, name)
case unix.S_IFDIR:
return fmt.Errorf("%w: release key identifies a directory", ErrUnsafePath)
}
}
for _, name := range names {
if unlinkErr := unix.Unlinkat(parentFD, name, 0); unlinkErr != nil && !errors.Is(unlinkErr, unix.ENOENT) {
return fmt.Errorf("removing entry %q: %w", name, unlinkErr)
}
}
maxPrune := min(pruneParents, len(handles)-1)
for childIndex := len(handles) - 1; childIndex >= len(handles)-maxPrune; childIndex-- {
parentIndex := childIndex - 1
name := parts[childIndex-1]
same, identityErr := sameDirectoryEntry(handles[parentIndex], name, handles[childIndex])
if identityErr != nil {
if errors.Is(identityErr, unix.ENOENT) {
break
}
return fmt.Errorf("checking directory %q before pruning: %w", name, identityErr)
}
if !same {
break
}
removeErr := unix.Unlinkat(handles[parentIndex], name, unix.AT_REMOVEDIR)
if removeErr == nil {
continue
}
if errors.Is(removeErr, unix.ENOENT) || errors.Is(removeErr, unix.ENOTEMPTY) || errors.Is(removeErr, unix.EEXIST) {
break
}
return fmt.Errorf("pruning directory %q: %w", name, removeErr)
}
return nil
}
func cleanRelativeParts(relativePath string) ([]string, error) {
if relativePath == "" || filepath.IsAbs(relativePath) || filepath.Clean(relativePath) != relativePath {
return nil, fmt.Errorf("%w: %q is not a clean relative path", ErrUnsafePath, relativePath)
}
parts := strings.Split(relativePath, string(filepath.Separator))
for _, part := range parts {
if part == "" || part == "." || part == ".." {
return nil, fmt.Errorf("%w: %q is not a clean relative path", ErrUnsafePath, relativePath)
}
}
return parts, nil
}
func sameDirectoryEntry(parentFD int, name string, openedFD int) (bool, error) {
var opened unix.Stat_t
if err := unix.Fstat(openedFD, &opened); err != nil {
return false, err
}
var current unix.Stat_t
if err := unix.Fstatat(parentFD, name, &current, unix.AT_SYMLINK_NOFOLLOW); err != nil {
return false, err
}
return current.Mode&unix.S_IFMT == unix.S_IFDIR && current.Dev == opened.Dev && current.Ino == opened.Ino, nil
}
+57
View File
@@ -0,0 +1,57 @@
//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris
package safefile
import (
"os"
"path/filepath"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ = Describe("Race-safe exact removal", func() {
DescribeTable("keeps replacement targets untouched after an opened parent is exchanged",
func(targetInsideRoot bool) {
root := GinkgoT().TempDir()
originalRequest := filepath.Join(root, "ephemeral", "request-id")
originalCategory := filepath.Join(originalRequest, "audio")
Expect(os.MkdirAll(originalCategory, 0750)).To(Succeed())
Expect(os.WriteFile(filepath.Join(originalCategory, "input.wav"), []byte("original"), 0640)).To(Succeed())
replacement := GinkgoT().TempDir()
if targetInsideRoot {
replacement = filepath.Join(root, "replacement")
Expect(os.MkdirAll(replacement, 0750)).To(Succeed())
}
Expect(os.MkdirAll(filepath.Join(replacement, "audio"), 0750)).To(Succeed())
replacementFile := filepath.Join(replacement, "audio", "input.wav")
Expect(os.WriteFile(replacementFile, []byte("replacement"), 0640)).To(Succeed())
renamedRequest := filepath.Join(root, "ephemeral", "opened-request")
opened := make(chan struct{})
swapped := make(chan struct{})
done := make(chan error, 1)
go func() {
done <- removeExact(root, filepath.Join("ephemeral", "request-id", "audio", "input.wav"), []string{".sha256", ".sha256.target"}, 2, func() {
close(opened)
<-swapped
})
}()
<-opened
Expect(os.Rename(originalRequest, renamedRequest)).To(Succeed())
Expect(os.Symlink(replacement, originalRequest)).To(Succeed())
close(swapped)
Expect(<-done).To(Succeed())
data, err := os.ReadFile(replacementFile)
Expect(err).NotTo(HaveOccurred())
Expect(data).To(Equal([]byte("replacement")))
Expect(originalRequest).To(BeAnExistingFile())
Expect(filepath.Join(renamedRequest, "audio", "input.wav")).NotTo(BeAnExistingFile())
},
Entry("for an internal replacement", true),
Entry("for an external replacement", false),
)
})
+10
View File
@@ -5652,6 +5652,12 @@ const docTemplate = `{
"schema.FaceRegisterRequest": {
"type": "object",
"properties": {
"embedding": {
"type": "array",
"items": {
"type": "number"
}
},
"img": {
"type": "string"
},
@@ -5667,6 +5673,10 @@ const docTemplate = `{
"name": {
"type": "string"
},
"registered_at": {
"description": "original enrollment time when replaying a saved embedding",
"type": "string"
},
"store": {
"description": "vector store model; empty = local-store default",
"type": "string"
+10
View File
@@ -5649,6 +5649,12 @@
"schema.FaceRegisterRequest": {
"type": "object",
"properties": {
"embedding": {
"type": "array",
"items": {
"type": "number"
}
},
"img": {
"type": "string"
},
@@ -5664,6 +5670,10 @@
"name": {
"type": "string"
},
"registered_at": {
"description": "original enrollment time when replaying a saved embedding",
"type": "string"
},
"store": {
"description": "vector store model; empty = local-store default",
"type": "string"
+7
View File
@@ -1138,6 +1138,10 @@ definitions:
type: object
schema.FaceRegisterRequest:
properties:
embedding:
items:
type: number
type: array
img:
type: string
labels:
@@ -1148,6 +1152,9 @@ definitions:
type: string
name:
type: string
registered_at:
description: original enrollment time when replaying a saved embedding
type: string
store:
description: vector store model; empty = local-store default
type: string
+4 -4
View File
@@ -3,10 +3,10 @@
# The four GitHub fields are rewritten by .github/ci/refresh-site-counters.sh,
# which runs weekly from .github/workflows/refresh-site-counters.yml. Editing
# them by hand works but will be overwritten on the next run.
stars: 48646
forks: 4377
contributors: 230
releases: 136
stars: 48949
forks: 4430
contributors: 237
releases: 135
# The GitHub API cannot answer for this one, so it is maintained by hand and
# the refresh script carries it through untouched.