mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 14:22:11 -04:00
feat(distributed): serve file staging over the worker tunnel
The four nodes.<id>.files.* subjects were the last commands a serve-backend worker took off the bus. They are now HTTP routes under workerctl.Prefix, on the same loopback server and behind the same bearer check as the ten lifecycle verbs, so the frontend reaches them through the worker's tunnel. files.listdir is the verb this matters most for. Its reply had to fit a payload the bus would carry, which put a wide model directory close to the limit; a response body has no such ceiling, so nothing truncates the listing at either end. A short listing reads to the frontend as files the worker does not have. S3NATSFileStager becomes S3FileStager and calls ControlClient, which means every failure now lands in the bucket phase 3 exists to keep straight: a route this frontend could not use is unroutable and nothing may act on it, while the worker's own answer, including "that file is not there", is evidence a caller may act on. Each RPC's deadline is DERIVED FROM the caller's context rather than started fresh, at every one of the five call sites, so a caller that gave up stops the RPC too. A worker started without an object store mounts no file verb at all and answers 404, which is the same answer a build too old to know them gives. The subjects and the backend worker's files.> publish grant go with them; a backend worker now publishes nowhere but its own inbox. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
5880f4e6dd
commit
7fc617c8ed
19 files changed
+1258
-378
No files matched your search
@@ -390,11 +390,22 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
}
|
||||
xlog.Info("File manager initialized", "cacheDir", cacheDir)
|
||||
|
||||
// The frontend's control plane client. It reaches every worker over that
|
||||
// worker's own tunnel, on the same `http` stream tag the file stager below
|
||||
// uses, so a control RPC to a worker another replica holds is relayed the
|
||||
// way an inference request is.
|
||||
//
|
||||
// ONE of these for the whole frontend, and the S3 file stager takes this
|
||||
// one rather than minting a second. The client caches an http.Client per
|
||||
// node, which is what keeps a worker's tunnel stream warm between verbs; a
|
||||
// second client would open its own and the two would never share one.
|
||||
controlClient := nodes.NewControlClient(workerHTTPDialer, cfg.Distributed.RegistrationToken)
|
||||
|
||||
// Create FileStager for distributed file transfer
|
||||
var fileStager nodes.FileStager
|
||||
if cfg.Distributed.StorageURL != "" {
|
||||
fileStager = nodes.NewS3NATSFileStager(fileMgr, natsClient)
|
||||
xlog.Info("File stager initialized (S3+NATS)")
|
||||
fileStager = nodes.NewS3FileStager(fileMgr, controlClient)
|
||||
xlog.Info("File stager initialized (object store + worker tunnel)")
|
||||
} else {
|
||||
fileStager = nodes.NewHTTPFileStager(func(nodeID string) (string, error) {
|
||||
node, err := registry.Get(context.Background(), nodeID)
|
||||
@@ -410,12 +421,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
}, cfg.Distributed.RegistrationToken, workerHTTPDialer)
|
||||
xlog.Info("File stager initialized (HTTP direct transfer)")
|
||||
}
|
||||
// The frontend's control plane client. It reaches every worker over that
|
||||
// worker's own tunnel, on the same `http` stream tag the file stager above
|
||||
// uses, so a control RPC to a worker another replica holds is relayed the
|
||||
// way an inference request is.
|
||||
controlClient := nodes.NewControlClient(workerHTTPDialer, cfg.Distributed.RegistrationToken)
|
||||
|
||||
// Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go
|
||||
remoteUnloader := nodes.NewRemoteUnloaderAdapter(
|
||||
registry,
|
||||
|
||||
@@ -376,32 +376,9 @@ type RunningModelInfo struct {
|
||||
Address string `json:"address,omitempty"`
|
||||
}
|
||||
|
||||
// File Staging (Request-Reply — targeted to specific nodes)
|
||||
// These subjects use request-reply for synchronous file operations.
|
||||
|
||||
// SubjectNodeFilesEnsure tells a serve-backend node to download an S3 key to its local cache.
|
||||
// Reply: {local_path, error}
|
||||
func SubjectNodeFilesEnsure(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.ensure"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesStage tells a serve-backend node to upload a local file to S3.
|
||||
// Reply: {key, error}
|
||||
func SubjectNodeFilesStage(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.stage"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesTemp tells a serve-backend node to allocate a temp file.
|
||||
// Reply: {local_path, error}
|
||||
func SubjectNodeFilesTemp(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.temp"
|
||||
}
|
||||
|
||||
// SubjectNodeFilesListDir tells a serve-backend node to list files in a directory.
|
||||
// Reply: {files: [...], error}
|
||||
func SubjectNodeFilesListDir(nodeID string) string {
|
||||
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".files.listdir"
|
||||
}
|
||||
// File staging is no longer carried here. The four nodes.<id>.files.* subjects
|
||||
// are HTTP routes under workerctl.Prefix, served on the worker's own server and
|
||||
// reached through its tunnel, so no subject is minted for them.
|
||||
|
||||
// Cache Invalidation (Pub/Sub — broadcast to all instances)
|
||||
const (
|
||||
|
||||
@@ -5,29 +5,47 @@ import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// S3NATSFileStager implements FileStager using S3 for storage and NATS
|
||||
// request-reply for coordination with backend nodes. Both frontend and
|
||||
// backend nodes share the same S3 bucket. The flow is:
|
||||
// S3FileStager implements FileStager using an object store for the bytes and
|
||||
// the worker's tunnelled control plane for coordination. Both frontend and
|
||||
// worker share the same bucket. The flow is:
|
||||
//
|
||||
// 1. Frontend uploads file to S3
|
||||
// 2. Frontend sends NATS request to nodes.{nodeID}.files.ensure
|
||||
// 3. Backend downloads from S3 to local cache, replies with local path
|
||||
type S3NATSFileStager struct {
|
||||
fm *storage.FileManager
|
||||
nats messaging.MessagingClient
|
||||
// 1. Frontend uploads the file to the object store
|
||||
// 2. Frontend calls POST /v1/control/files/ensure on the worker's tunnel
|
||||
// 3. Worker downloads from the store to its local cache and replies with the
|
||||
// local path
|
||||
type S3FileStager struct {
|
||||
fm *storage.FileManager
|
||||
control *ControlClient
|
||||
}
|
||||
|
||||
// NewS3NATSFileStager creates a new S3+NATS file stager.
|
||||
func NewS3NATSFileStager(fm *storage.FileManager, nats messaging.MessagingClient) *S3NATSFileStager {
|
||||
return &S3NATSFileStager{fm: fm, nats: nats}
|
||||
// NewS3FileStager creates a file stager that moves bytes through fm and
|
||||
// commands workers over control.
|
||||
func NewS3FileStager(fm *storage.FileManager, control *ControlClient) *S3FileStager {
|
||||
return &S3FileStager{fm: fm, control: control}
|
||||
}
|
||||
|
||||
// NATS request/reply message types
|
||||
// The two budgets a file-staging RPC gets. They are the ones the NATS
|
||||
// request-reply timeouts carried, kept verbatim: a transfer verb waits out a
|
||||
// multi-gigabyte copy, and a metadata verb does not.
|
||||
//
|
||||
// They are CEILINGS on the caller's own budget rather than budgets of their
|
||||
// own. Every call below derives its deadline from the caller's context, so a
|
||||
// caller that has already given up stops the RPC too; deriving from a fresh
|
||||
// background context would keep commanding a worker nobody is listening to and
|
||||
// would let a late answer be read as a live one.
|
||||
const (
|
||||
fileTransferRPCTimeout = 10 * time.Minute
|
||||
fileMetadataRPCTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
// Control request/reply message types. Their JSON is the shape the
|
||||
// nodes.<id>.files.* subjects carried, so the worker's handler bodies did not
|
||||
// have to change when the carrier did.
|
||||
|
||||
type fileEnsureRequest struct {
|
||||
Key string `json:"key"`
|
||||
@@ -64,10 +82,24 @@ type fileListDirReply struct {
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// EnsureRemote uploads a local file to S3 (if not already there) and sends
|
||||
// a NATS request-reply to the backend node to download it locally.
|
||||
func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
|
||||
// Upload to S3 if not already present
|
||||
// callWorker issues one file-staging RPC under a deadline derived from the
|
||||
// caller's context.
|
||||
//
|
||||
// It exists so the derivation is written ONCE. Five call sites each repeating
|
||||
// context.WithTimeout is five chances for one of them to start from a
|
||||
// background context instead, and that one site would then keep commanding a
|
||||
// worker after its caller had gone, with nothing else in the suite any redder
|
||||
// for it.
|
||||
func (s *S3FileStager) callWorker(ctx context.Context, nodeID, path string, budget time.Duration, req, reply any) error {
|
||||
rpcCtx, cancel := context.WithTimeout(ctx, budget)
|
||||
defer cancel()
|
||||
return s.control.Call(rpcCtx, nodeID, path, req, reply)
|
||||
}
|
||||
|
||||
// EnsureRemote uploads a local file to the object store (if not already there)
|
||||
// and tells the worker to fetch it.
|
||||
func (s *S3FileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
|
||||
// Upload to the store if not already present
|
||||
exists, _ := s.fm.Exists(ctx, key)
|
||||
if !exists {
|
||||
// Wrap with progress reporting if a staging callback is available
|
||||
@@ -78,14 +110,13 @@ func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath,
|
||||
}
|
||||
}
|
||||
if err := s.fm.UploadWithProgress(ctx, key, localPath, progressFn); err != nil {
|
||||
return "", fmt.Errorf("uploading %s to S3: %w", localPath, err)
|
||||
return "", fmt.Errorf("uploading %s to the object store: %w", localPath, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Send NATS request-reply to backend
|
||||
subject := messaging.SubjectNodeFilesEnsure(nodeID)
|
||||
reply, err := messaging.RequestJSON[fileEnsureRequest, fileEnsureReply](s.nats, subject, fileEnsureRequest{Key: key}, 10*time.Minute)
|
||||
if err != nil {
|
||||
var reply fileEnsureReply
|
||||
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesEnsure, fileTransferRPCTimeout,
|
||||
fileEnsureRequest{Key: key}, &reply); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
@@ -96,36 +127,38 @@ func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath,
|
||||
return reply.LocalPath, nil
|
||||
}
|
||||
|
||||
// FetchRemote tells the backend to upload a file to S3, then downloads it locally.
|
||||
func (s *S3NATSFileStager) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error {
|
||||
// Tell backend to upload to S3
|
||||
// FetchRemote tells the worker to upload a file to the object store, then
|
||||
// downloads it locally.
|
||||
func (s *S3FileStager) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error {
|
||||
key := storage.EphemeralKey(remotePath, "fetch", "output")
|
||||
return s.fetchRemoteWithKey(ctx, nodeID, remotePath, key, localDst, true)
|
||||
}
|
||||
|
||||
// FetchRemoteByKey tells the backend to upload a file (identified by key) to S3,
|
||||
// then downloads it locally. The key is used as-is for S3 routing.
|
||||
func (s *S3NATSFileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error {
|
||||
// For S3 mode, we still need the remote path — derive it from the key.
|
||||
// The backend serves the file from its data dir based on the key prefix.
|
||||
// FetchRemoteByKey tells the worker to upload a file (identified by key) to the
|
||||
// object store, then downloads it locally. The key is used as-is for routing.
|
||||
func (s *S3FileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error {
|
||||
// The remote path is derived from the key: the worker serves the file from
|
||||
// its data dir based on the key prefix.
|
||||
remotePath := "/" + key // e.g. "/data/quantization/{jobID}/model.gguf"
|
||||
return s.fetchRemoteWithKey(ctx, nodeID, remotePath, key, localDst, true)
|
||||
}
|
||||
|
||||
func (s *S3NATSFileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePath, key, localDst string, cleanup bool) error {
|
||||
subject := messaging.SubjectNodeFilesStage(nodeID)
|
||||
reply, err := messaging.RequestJSON[fileStageRequest, fileStageReply](s.nats, subject, fileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute)
|
||||
if err != nil {
|
||||
func (s *S3FileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePath, key, localDst string, cleanup bool) error {
|
||||
var reply fileStageReply
|
||||
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileTransferRPCTimeout,
|
||||
fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
|
||||
return err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
return fmt.Errorf("backend stage failed: %s", reply.Error)
|
||||
}
|
||||
|
||||
// Download from S3 to local cache
|
||||
// Download from the store to the local cache. The CALLER's context bounds
|
||||
// this rather than the RPC's, because it is this frontend's own work and
|
||||
// the RPC it belonged to has already finished.
|
||||
cachedPath, err := s.fm.Download(ctx, key)
|
||||
if err != nil {
|
||||
return fmt.Errorf("downloading %s from S3: %w", key, err)
|
||||
return fmt.Errorf("downloading %s from the object store: %w", key, err)
|
||||
}
|
||||
|
||||
// Copy from cache to destination
|
||||
@@ -141,11 +174,11 @@ func (s *S3NATSFileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remot
|
||||
return nil
|
||||
}
|
||||
|
||||
// AllocRemoteTemp asks the backend to allocate a temp file via NATS request-reply.
|
||||
func (s *S3NATSFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) {
|
||||
subject := messaging.SubjectNodeFilesTemp(nodeID)
|
||||
reply, err := messaging.RequestJSON[fileTempRequest, fileTempReply](s.nats, subject, fileTempRequest{}, 30*time.Second)
|
||||
if err != nil {
|
||||
// AllocRemoteTemp asks the worker to allocate a temp file.
|
||||
func (s *S3FileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) {
|
||||
var reply fileTempReply
|
||||
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesTemp, fileMetadataRPCTimeout,
|
||||
fileTempRequest{}, &reply); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
@@ -155,10 +188,17 @@ func (s *S3NATSFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (
|
||||
return reply.LocalPath, nil
|
||||
}
|
||||
|
||||
func (s *S3NATSFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) {
|
||||
subject := messaging.SubjectNodeFilesListDir(nodeID)
|
||||
reply, err := messaging.RequestJSON[fileListDirRequest, fileListDirReply](s.nats, subject, fileListDirRequest{KeyPrefix: keyPrefix}, 30*time.Second)
|
||||
if err != nil {
|
||||
// ListRemoteDir returns the relative paths of every file under keyPrefix on the
|
||||
// worker.
|
||||
//
|
||||
// Nothing truncates the answer, at either end. The bus this used to ride put a
|
||||
// ceiling on how big a reply could be, and a wide model directory was the case
|
||||
// that pushed against it; a response body has no such ceiling, and a short
|
||||
// listing would read to the caller as files the worker does not have.
|
||||
func (s *S3FileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) {
|
||||
var reply fileListDirReply
|
||||
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesListDir, fileMetadataRPCTimeout,
|
||||
fileListDirRequest{KeyPrefix: keyPrefix}, &reply); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
@@ -168,11 +208,11 @@ func (s *S3NATSFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix
|
||||
return reply.Files, nil
|
||||
}
|
||||
|
||||
// StageRemoteToStore tells the backend to upload a local file to S3.
|
||||
func (s *S3NATSFileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error {
|
||||
subject := messaging.SubjectNodeFilesStage(nodeID)
|
||||
reply, err := messaging.RequestJSON[fileStageRequest, fileStageReply](s.nats, subject, fileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute)
|
||||
if err != nil {
|
||||
// StageRemoteToStore tells the worker to upload a local file to shared storage.
|
||||
func (s *S3FileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error {
|
||||
var reply fileStageReply
|
||||
if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileTransferRPCTimeout,
|
||||
fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil {
|
||||
return err
|
||||
}
|
||||
if reply.Error != "" {
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// The S3 file stager's half of the file-staging contract: which verb each
|
||||
// method addresses, and whose budget bounds it.
|
||||
//
|
||||
// The stager is reached through the real ControlClient over a real HTTP
|
||||
// transport onto a scripted worker, because what these specs are about is
|
||||
// transport behaviour: a double that never dials anything cannot fail the way
|
||||
// a spent budget or a rejected route fails.
|
||||
var _ = Describe("the S3 file stager's control RPCs", func() {
|
||||
const nodeID = "stager-node"
|
||||
|
||||
var (
|
||||
workers *scriptedControlWorkers
|
||||
stager *S3FileStager
|
||||
local string
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
workers = newScriptedControlWorkers()
|
||||
|
||||
dir := GinkgoT().TempDir()
|
||||
store, err := storage.NewFilesystemStore(filepath.Join(dir, "objectstore"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err := storage.NewFileManager(store, filepath.Join(dir, "cache"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
stager = NewS3FileStager(fm, workers.controlClient())
|
||||
|
||||
local = filepath.Join(dir, "model.gguf")
|
||||
Expect(os.WriteFile(local, []byte("weights"), 0o600)).To(Succeed())
|
||||
})
|
||||
|
||||
DescribeTable("addresses the verb that verb's path names",
|
||||
func(path string, reply any, call func(*S3FileStager) error) {
|
||||
workers.scriptReply(controlKey(nodeID, path), reply)
|
||||
Expect(call(stager)).To(Succeed())
|
||||
Expect(workers.callSubjects()).To(ContainElement(controlKey(nodeID, path)))
|
||||
},
|
||||
Entry("ensure", workerctl.PathFilesEnsure, fileEnsureReply{LocalPath: "/w/models/m.gguf"},
|
||||
func(s *S3FileStager) error {
|
||||
_, err := s.EnsureRemote(context.Background(), nodeID, local, storage.ModelKey("m.gguf"))
|
||||
return err
|
||||
}),
|
||||
Entry("temp", workerctl.PathFilesTemp, fileTempReply{LocalPath: "/w/tmp/x"},
|
||||
func(s *S3FileStager) error {
|
||||
_, err := s.AllocRemoteTemp(context.Background(), nodeID)
|
||||
return err
|
||||
}),
|
||||
Entry("listdir", workerctl.PathFilesListDir, fileListDirReply{Files: []string{"a", "b"}},
|
||||
func(s *S3FileStager) error {
|
||||
_, err := s.ListRemoteDir(context.Background(), nodeID, "models/m")
|
||||
return err
|
||||
}),
|
||||
Entry("stage", workerctl.PathFilesStage, fileStageReply{Key: "data/out"},
|
||||
func(s *S3FileStager) error {
|
||||
return s.StageRemoteToStore(context.Background(), nodeID, "/w/models/out", "data/out")
|
||||
}),
|
||||
)
|
||||
|
||||
// THE rule of this change, and it is written out at five separate call
|
||||
// sites: the RPC's budget is DERIVED FROM the caller's context, never
|
||||
// started fresh from a background one. A site that started fresh would keep
|
||||
// commanding a worker after its caller had given up, and would report the
|
||||
// worker's late answer as a live one. Each site is pinned on its own,
|
||||
// because five sites behind one spec is four sites nothing holds.
|
||||
DescribeTable("never reaches the worker once the caller's context is spent",
|
||||
func(call func(context.Context, *S3FileStager) error) {
|
||||
// Every verb is scripted to answer, so the ONLY thing that can stop
|
||||
// the call is the caller's own spent context.
|
||||
for _, p := range []string{
|
||||
workerctl.PathFilesEnsure, workerctl.PathFilesStage,
|
||||
workerctl.PathFilesTemp, workerctl.PathFilesListDir,
|
||||
} {
|
||||
workers.scriptRawReply(controlKey(nodeID, p), []byte(`{}`))
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
err := call(ctx, stager)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "got %v", err)
|
||||
// An expiry is never evidence about a file, so it must arrive
|
||||
// wearing the umbrella that stops a caller acting on it.
|
||||
Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue(), "got %v", err)
|
||||
Expect(cluster.IsWorkerAnswer(err)).To(BeFalse())
|
||||
Expect(workers.callSubjects()).To(BeEmpty())
|
||||
},
|
||||
Entry("ensure", func(ctx context.Context, s *S3FileStager) error {
|
||||
_, err := s.EnsureRemote(ctx, nodeID, local, storage.ModelKey("m.gguf"))
|
||||
return err
|
||||
}),
|
||||
Entry("temp", func(ctx context.Context, s *S3FileStager) error {
|
||||
_, err := s.AllocRemoteTemp(ctx, nodeID)
|
||||
return err
|
||||
}),
|
||||
Entry("listdir", func(ctx context.Context, s *S3FileStager) error {
|
||||
_, err := s.ListRemoteDir(ctx, nodeID, "models/m")
|
||||
return err
|
||||
}),
|
||||
Entry("stage to store", func(ctx context.Context, s *S3FileStager) error {
|
||||
return s.StageRemoteToStore(ctx, nodeID, "/w/models/out", "data/out")
|
||||
}),
|
||||
Entry("fetch", func(ctx context.Context, s *S3FileStager) error {
|
||||
return s.FetchRemote(ctx, nodeID, "/w/models/out", filepath.Join(GinkgoT().TempDir(), "dst"))
|
||||
}),
|
||||
Entry("fetch by key", func(ctx context.Context, s *S3FileStager) error {
|
||||
return s.FetchRemoteByKey(ctx, nodeID, "data/out", filepath.Join(GinkgoT().TempDir(), "dst"))
|
||||
}),
|
||||
)
|
||||
|
||||
// The other half of the same distinction, also at every site: what the
|
||||
// WORKER said comes back as the worker's answer, so a caller may act on it,
|
||||
// and it must not be dressed up as a route failure.
|
||||
DescribeTable("reports the worker's own refusal as the worker's answer",
|
||||
func(path string, call func(*S3FileStager) error) {
|
||||
workers.scriptRawReply(controlKey(nodeID, path), []byte(`{"error":"no space left on device"}`))
|
||||
err := call(stager)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no space left on device"))
|
||||
Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse(), "got %v", err)
|
||||
},
|
||||
Entry("ensure", workerctl.PathFilesEnsure, func(s *S3FileStager) error {
|
||||
_, err := s.EnsureRemote(context.Background(), nodeID, local, storage.ModelKey("m.gguf"))
|
||||
return err
|
||||
}),
|
||||
Entry("temp", workerctl.PathFilesTemp, func(s *S3FileStager) error {
|
||||
_, err := s.AllocRemoteTemp(context.Background(), nodeID)
|
||||
return err
|
||||
}),
|
||||
Entry("listdir", workerctl.PathFilesListDir, func(s *S3FileStager) error {
|
||||
_, err := s.ListRemoteDir(context.Background(), nodeID, "models/m")
|
||||
return err
|
||||
}),
|
||||
Entry("stage to store", workerctl.PathFilesStage, func(s *S3FileStager) error {
|
||||
return s.StageRemoteToStore(context.Background(), nodeID, "/w/models/out", "data/out")
|
||||
}),
|
||||
Entry("fetch", workerctl.PathFilesStage, func(s *S3FileStager) error {
|
||||
return s.FetchRemote(context.Background(), nodeID, "/w/models/out",
|
||||
filepath.Join(GinkgoT().TempDir(), "dst"))
|
||||
}),
|
||||
)
|
||||
|
||||
// A worker too old to serve a file verb answers 404. That is a DEPLOYMENT
|
||||
// fact about the worker's build and says nothing about the file, so it must
|
||||
// not reach a caller as "that file is not there".
|
||||
It("reports a worker that serves no file verbs as unsupported, not as an absent file", func() {
|
||||
workers.scriptUnsupported(controlKey(nodeID, workerctl.PathFilesListDir))
|
||||
_, err := stager.ListRemoteDir(context.Background(), nodeID, "models/m")
|
||||
Expect(err).To(MatchError(ErrWorkerControlUnsupported))
|
||||
Expect(cluster.IsWorkerAnswer(err)).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reports a worker it has no route to as unroutable", func() {
|
||||
workers.scriptUnroutable(nodeID)
|
||||
_, err := stager.AllocRemoteTemp(context.Background(), nodeID)
|
||||
Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue(), "got %v", err)
|
||||
Expect(cluster.IsWorkerAnswer(err)).To(BeFalse())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,268 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// The worker's file-staging control plane: four 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.
|
||||
//
|
||||
// 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
|
||||
// listing is no longer sized against a bus payload: it is a response body the
|
||||
// caller is already reading, which is why nothing here truncates one.
|
||||
//
|
||||
// The 200-with-error-field shape is the same one the lifecycle verbs take, and
|
||||
// for the same reason: "that file is not there" is the worker's own ANSWER, and
|
||||
// a reap guard may act on it, while a non-2xx is this frontend failing to reach
|
||||
// the worker, which nothing may act on. A handler that answered 500 for a
|
||||
// failed upload would move its verdict into the bucket reserved for a broken
|
||||
// link.
|
||||
|
||||
// The file-staging reply bodies. They mirror the frontend's decode structs in
|
||||
// core/services/nodes/file_stager_s3.go field for field; the two are written
|
||||
// separately because neither package may import the other, and the roundtrip
|
||||
// spec is what holds them together.
|
||||
type fileEnsureReply struct {
|
||||
LocalPath string `json:"local_path,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileStageReply struct {
|
||||
Key string `json:"key,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileTempReply struct {
|
||||
LocalPath string `json:"local_path,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type fileListDirReply struct {
|
||||
Files []string `json:"files,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// stagingCacheDir is the one place the worker's staging cache directory is
|
||||
// derived from its configuration.
|
||||
//
|
||||
// One place and not two, because the FileManager caches INTO this directory and
|
||||
// the listdir and temp verbs resolve paths AGAINST it. Two derivations that
|
||||
// drifted would give a worker that downloads a file to one directory and then
|
||||
// reports it missing from another, and both halves would still be self
|
||||
// consistent.
|
||||
func (cfg *Config) stagingCacheDir() string {
|
||||
return filepath.Join(cfg.ModelsPath, "..", "cache")
|
||||
}
|
||||
|
||||
// stagingDataDir is where keys under storage.DataKeyPrefix resolve, and it is
|
||||
// derived from the cache directory for the same single-source reason.
|
||||
func (cfg *Config) stagingDataDir() string {
|
||||
return filepath.Join(cfg.stagingCacheDir(), "..", "data")
|
||||
}
|
||||
|
||||
// NewStagingFileManager builds the FileManager the file-staging verbs serve
|
||||
// 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
|
||||
// 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) {
|
||||
s3Store, err := storage.NewS3Store(ctx, storage.S3Config{
|
||||
Endpoint: cfg.StorageURL,
|
||||
Region: cfg.StorageRegion,
|
||||
Bucket: cfg.StorageBucket,
|
||||
AccessKeyID: cfg.StorageAccessKey,
|
||||
SecretAccessKey: cfg.StorageSecretKey,
|
||||
ForcePathStyle: true,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("initializing S3 store: %w", err)
|
||||
}
|
||||
fm, err := storage.NewFileManager(s3Store, cfg.stagingCacheDir())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("initializing file manager: %w", err)
|
||||
}
|
||||
return fm, nil
|
||||
}
|
||||
|
||||
// RegisterFileControlRoutes mounts the four 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) {
|
||||
cacheDir := cfg.stagingCacheDir()
|
||||
|
||||
// files.ensure: download an object-store key into this worker's cache and
|
||||
// say where it landed.
|
||||
//
|
||||
// It takes the caller's context. Nothing is terminated and no resource is
|
||||
// held if it is abandoned half way: the download simply stops, and the next
|
||||
// attempt starts over. So the caller's budget is the operation's budget.
|
||||
postControlVerb(mux, workerctl.PathFilesEnsure, func(ctx context.Context, body []byte) (any, error) {
|
||||
var req struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
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)
|
||||
if err != nil {
|
||||
xlog.Error("File ensure failed", "key", req.Key, "error", err)
|
||||
return fileEnsureReply{Error: err.Error()}, nil
|
||||
}
|
||||
xlog.Debug("File ensured locally", "key", req.Key, "path", localPath)
|
||||
return fileEnsureReply{LocalPath: localPath}, nil
|
||||
})
|
||||
|
||||
// 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
|
||||
// primitive: the token holder can name any absolute path, so only the
|
||||
// directories this worker stages out of are served.
|
||||
postControlVerb(mux, workerctl.PathFilesStage, func(ctx context.Context, body []byte) (any, error) {
|
||||
var req struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, fmt.Errorf("invalid files.stage request: %w", err)
|
||||
}
|
||||
allowedDirs := []string{cacheDir}
|
||||
if cfg.ModelsPath != "" {
|
||||
allowedDirs = append(allowedDirs, cfg.ModelsPath)
|
||||
}
|
||||
if !isPathAllowed(req.LocalPath, allowedDirs) {
|
||||
return fileStageReply{Error: "path outside allowed directories"}, nil
|
||||
}
|
||||
if err := fm.Upload(ctx, req.Key, req.LocalPath); err != nil {
|
||||
xlog.Error("File stage failed", "path", req.LocalPath, "key", req.Key, "error", err)
|
||||
return fileStageReply{Error: err.Error()}, nil
|
||||
}
|
||||
xlog.Debug("File staged to the object store", "path", req.LocalPath, "key", req.Key)
|
||||
return fileStageReply{Key: req.Key}, nil
|
||||
})
|
||||
|
||||
// files.temp: allocate an empty file the frontend may then upload into.
|
||||
postControlVerb(mux, workerctl.PathFilesTemp, func(context.Context, []byte) (any, error) {
|
||||
tmpDir := filepath.Join(cacheDir, "staging-tmp")
|
||||
if err := os.MkdirAll(tmpDir, 0750); err != nil {
|
||||
return fileTempReply{Error: fmt.Sprintf("creating temp dir: %v", err)}, nil
|
||||
}
|
||||
f, err := os.CreateTemp(tmpDir, "localai-staging-*.tmp")
|
||||
if err != nil {
|
||||
return fileTempReply{Error: fmt.Sprintf("creating temp file: %v", err)}, nil
|
||||
}
|
||||
localPath := f.Name()
|
||||
if err := f.Close(); err != nil {
|
||||
return fileTempReply{Error: fmt.Sprintf("closing temp file: %v", err)}, nil
|
||||
}
|
||||
xlog.Debug("Allocated temp file", "path", localPath)
|
||||
return fileTempReply{LocalPath: localPath}, nil
|
||||
})
|
||||
|
||||
// files.listdir: the relative paths of every file under one key prefix.
|
||||
//
|
||||
// Nothing here caps the answer. Over NATS the reply had to fit a payload
|
||||
// the bus was willing to carry, and a wide model directory was the case
|
||||
// that risked it; over HTTP the listing is written into a body the caller
|
||||
// is already reading, so its size is no longer a property of the carrier. A
|
||||
// cap would silently return a SHORT listing, which the frontend reads as
|
||||
// files that are not there.
|
||||
postControlVerb(mux, workerctl.PathFilesListDir, func(ctx context.Context, body []byte) (any, error) {
|
||||
var req struct {
|
||||
KeyPrefix string `json:"key_prefix"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, fmt.Errorf("invalid files.listdir request: %w", err)
|
||||
}
|
||||
dirPath, ok := cfg.resolveStagingDir(req.KeyPrefix)
|
||||
if !ok {
|
||||
return fileListDirReply{Error: "invalid key prefix"}, nil
|
||||
}
|
||||
files, err := listStagedFiles(ctx, dirPath)
|
||||
if err != nil {
|
||||
xlog.Error("Failed to list staged files", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "error", err)
|
||||
return fileListDirReply{Error: err.Error()}, nil
|
||||
}
|
||||
xlog.Debug("Listed remote dir", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "fileCount", len(files))
|
||||
return fileListDirReply{Files: files}, nil
|
||||
})
|
||||
}
|
||||
|
||||
// resolveStagingDir maps a storage key prefix onto the local directory it names,
|
||||
// and reports whether that directory is one this worker serves.
|
||||
//
|
||||
// The second return is not an error string on purpose: a prefix that climbs out
|
||||
// of the served directories is refused before anything touches the filesystem,
|
||||
// so a crafted key_prefix cannot turn this verb into a directory reader for the
|
||||
// whole host.
|
||||
func (cfg *Config) resolveStagingDir(keyPrefix string) (string, bool) {
|
||||
cacheDir := cfg.stagingCacheDir()
|
||||
dataDir := cfg.stagingDataDir()
|
||||
|
||||
dirPath := filepath.Join(cacheDir, keyPrefix)
|
||||
if rel, ok := strings.CutPrefix(keyPrefix, storage.ModelKeyPrefix); ok && cfg.ModelsPath != "" {
|
||||
dirPath = filepath.Join(cfg.ModelsPath, rel)
|
||||
} else if rel, ok := strings.CutPrefix(keyPrefix, storage.DataKeyPrefix); ok {
|
||||
dirPath = filepath.Join(dataDir, rel)
|
||||
}
|
||||
|
||||
dirPath = filepath.Clean(dirPath)
|
||||
cleanCache := filepath.Clean(cacheDir)
|
||||
cleanModels := filepath.Clean(cfg.ModelsPath)
|
||||
cleanData := filepath.Clean(dataDir)
|
||||
within := func(root string) bool {
|
||||
return dirPath == root || strings.HasPrefix(dirPath, root+string(filepath.Separator))
|
||||
}
|
||||
if within(cleanCache) || (cleanModels != "." && within(cleanModels)) || within(cleanData) {
|
||||
return dirPath, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// listStagedFiles walks dirPath and returns every file's path relative to it.
|
||||
//
|
||||
// The walk honours ctx because a very wide directory is real work, and a caller
|
||||
// that has already given up should not keep this worker stat-ing files. The
|
||||
// context error is returned as the walk's error, so it travels back as a
|
||||
// FAILURE of the listing rather than as an empty listing, which the frontend
|
||||
// would read as a directory with nothing in it.
|
||||
func listStagedFiles(ctx context.Context, dirPath string) ([]string, error) {
|
||||
var files []string
|
||||
err := filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
return ctxErr
|
||||
}
|
||||
if d.IsDir() {
|
||||
return nil
|
||||
}
|
||||
rel, relErr := filepath.Rel(dirPath, path)
|
||||
if relErr != nil {
|
||||
return relErr
|
||||
}
|
||||
files = append(files, rel)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return files, nil
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// The four file-staging verbs on the worker's own control plane.
|
||||
//
|
||||
// They used to be NATS request-reply subjects, and the size of a listdir reply
|
||||
// was a property of the carrier: a directory with enough files in it produced a
|
||||
// payload the bus was not comfortable with. Over HTTP it is a response body,
|
||||
// which is why one of the specs below asserts a listing far past any payload
|
||||
// cap rather than asserting a cap of its own.
|
||||
var _ = Describe("worker file-staging control routes", func() {
|
||||
var (
|
||||
cfg *Config
|
||||
srv *httptest.Server
|
||||
fm *storage.FileManager
|
||||
store *storage.FilesystemStore
|
||||
modelsDir string
|
||||
cacheDir string
|
||||
)
|
||||
|
||||
// post issues one control verb the way the frontend does and hands back the
|
||||
// raw response, so a spec can assert the STATUS as well as the body. The
|
||||
// two carry different meanings and a helper that decoded only the body
|
||||
// would hide the one this file exists to pin.
|
||||
post := func(path string, body any) *http.Response {
|
||||
GinkgoHelper()
|
||||
raw, err := json.Marshal(body)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
resp, err := http.Post(srv.URL+path, "application/json", bytes.NewReader(raw))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
return resp
|
||||
}
|
||||
|
||||
decode := func(resp *http.Response, out any) {
|
||||
GinkgoHelper()
|
||||
Expect(json.NewDecoder(resp.Body).Decode(out)).To(Succeed())
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
modelsDir = filepath.Join(dir, "models")
|
||||
Expect(os.MkdirAll(modelsDir, 0o750)).To(Succeed())
|
||||
cacheDir = filepath.Join(dir, "cache")
|
||||
|
||||
var err error
|
||||
store, err = storage.NewFilesystemStore(filepath.Join(dir, "objectstore"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
fm, err = storage.NewFileManager(store, cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
cfg = &Config{ModelsPath: modelsDir}
|
||||
mux := http.NewServeMux()
|
||||
cfg.RegisterFileControlRoutes(mux, fm)
|
||||
srv = httptest.NewServer(mux)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("allocates a temp path and returns it", func() {
|
||||
resp := post(workerctl.PathFilesTemp, struct{}{})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(BeEmpty())
|
||||
Expect(reply.LocalPath).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("downloads a key the store already holds and reports where it landed", func() {
|
||||
key := storage.ModelKey("ensure-me.gguf")
|
||||
Expect(store.Put(context.Background(), key, strings.NewReader("weights"))).To(Succeed())
|
||||
|
||||
resp := post(workerctl.PathFilesEnsure, map[string]string{"key": key})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(BeEmpty())
|
||||
Expect(reply.LocalPath).To(BeAnExistingFile())
|
||||
Expect(os.ReadFile(reply.LocalPath)).To(Equal([]byte("weights")))
|
||||
})
|
||||
|
||||
It("uploads a file under an allowed directory and answers with its key", func() {
|
||||
local := filepath.Join(modelsDir, "staged.bin")
|
||||
Expect(os.WriteFile(local, []byte("output"), 0o600)).To(Succeed())
|
||||
|
||||
resp := post(workerctl.PathFilesStage, map[string]string{"local_path": local, "key": "data/out.bin"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Key string `json:"key"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(BeEmpty())
|
||||
Expect(reply.Key).To(Equal("data/out.bin"))
|
||||
exists, err := store.Exists(context.Background(), "data/out.bin")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeTrue())
|
||||
})
|
||||
|
||||
It("returns a listing longer than a NATS payload would have carried", func() {
|
||||
// 4000 files at ~40 bytes of name each is ~160 KB, twenty times the
|
||||
// NOTIFY cap and well past what the old carrier was comfortable with.
|
||||
// The point of the spec is that size is no longer a property of the
|
||||
// transport.
|
||||
bigDir := filepath.Join(modelsDir, "big")
|
||||
Expect(os.MkdirAll(bigDir, 0o750)).To(Succeed())
|
||||
for i := range 4000 {
|
||||
name := fmt.Sprintf("shard-%030d.safetensors", i)
|
||||
Expect(os.WriteFile(filepath.Join(bigDir, name), []byte("x"), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
resp := post(workerctl.PathFilesListDir, map[string]string{"key_prefix": "models/big"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Files []string `json:"files"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(BeEmpty())
|
||||
Expect(reply.Files).To(HaveLen(4000))
|
||||
})
|
||||
|
||||
// The rule these four pin is ONE rule stated at four handlers: a verb's own
|
||||
// failure is the worker's ANSWER and travels as a 200 with the error field
|
||||
// set, never as a 5xx. A 5xx is what the frontend maps onto "this frontend
|
||||
// could not reach that worker", which nothing may act on, so a handler that
|
||||
// answered 500 would move its own verdict into the bucket reserved for a
|
||||
// broken link. Each handler is pinned separately because each writes the
|
||||
// rule out for itself.
|
||||
It("reports a staging failure as a 200 with an error field, not as a 5xx", func() {
|
||||
resp := post(workerctl.PathFilesStage, map[string]string{"local_path": "/nope", "key": "k"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).NotTo(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports an upload that failed as a 200 with an error field, not as a 5xx", func() {
|
||||
// The path is INSIDE an allowed directory and the file is simply not
|
||||
// there, so this reaches the upload rather than stopping at the
|
||||
// allow-list. The two are separate statements of the same rule inside
|
||||
// one handler, and a spec that only ever reaches the allow-list leaves
|
||||
// the upload free to answer 500 for the worker's own verdict.
|
||||
missing := filepath.Join(modelsDir, "was-never-written.bin")
|
||||
resp := post(workerctl.PathFilesStage, map[string]string{"local_path": missing, "key": "data/x"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Key string `json:"key"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).NotTo(BeEmpty())
|
||||
Expect(reply.Error).NotTo(ContainSubstring("outside allowed directories"))
|
||||
Expect(reply.Key).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports an ensure of a key the store does not hold as a 200 with an error field", func() {
|
||||
resp := post(workerctl.PathFilesEnsure, map[string]string{"key": "models/absent.gguf"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).NotTo(BeEmpty())
|
||||
Expect(reply.LocalPath).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports a listdir of a directory that is not there as a 200 with an error field", func() {
|
||||
resp := post(workerctl.PathFilesListDir, map[string]string{"key_prefix": "models/never-created"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Files []string `json:"files"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).NotTo(BeEmpty())
|
||||
Expect(reply.Files).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports a temp allocation it cannot make as a 200 with an error field", func() {
|
||||
// A regular file where the staging directory must go: MkdirAll cannot
|
||||
// create through it, for root as much as for anyone else, so the verb
|
||||
// fails for a reason that is entirely this worker's own.
|
||||
Expect(os.MkdirAll(cacheDir, 0o750)).To(Succeed())
|
||||
Expect(os.WriteFile(filepath.Join(cacheDir, "staging-tmp"), []byte("not a dir"), 0o600)).To(Succeed())
|
||||
|
||||
resp := post(workerctl.PathFilesTemp, struct{}{})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).NotTo(BeEmpty())
|
||||
Expect(reply.LocalPath).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses a key prefix that climbs out of the directories it serves", func() {
|
||||
resp := post(workerctl.PathFilesListDir, map[string]string{"key_prefix": "../../../etc"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Files []string `json:"files"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(ContainSubstring("invalid key prefix"))
|
||||
Expect(reply.Files).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses to upload a path outside the directories it serves", func() {
|
||||
outside := filepath.Join(GinkgoT().TempDir(), "secret")
|
||||
Expect(os.WriteFile(outside, []byte("nope"), 0o600)).To(Succeed())
|
||||
|
||||
resp := post(workerctl.PathFilesStage, map[string]string{"local_path": outside, "key": "data/leak"})
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
decode(resp, &reply)
|
||||
Expect(reply.Error).To(ContainSubstring("outside allowed directories"))
|
||||
})
|
||||
|
||||
// The bounded, POST-only entry point, pinned at every one of the four
|
||||
// paths rather than at one of them. A route that skipped it would be a
|
||||
// second, unbounded door onto a boundary this worker serves, and it would
|
||||
// also let a liveness probe or an address bar run a command.
|
||||
DescribeTable("refuses a method that is not POST",
|
||||
func(path string) {
|
||||
resp, err := http.Get(srv.URL + path)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusMethodNotAllowed))
|
||||
},
|
||||
Entry("ensure", workerctl.PathFilesEnsure),
|
||||
Entry("stage", workerctl.PathFilesStage),
|
||||
Entry("temp", workerctl.PathFilesTemp),
|
||||
Entry("listdir", workerctl.PathFilesListDir),
|
||||
)
|
||||
|
||||
DescribeTable("refuses a body past the control cap, and as a rejected request rather than an answer",
|
||||
func(path string) {
|
||||
// VALID JSON past the cap, deliberately. A body of filler is
|
||||
// refused by the decoder whether or not anything bounds it, so a
|
||||
// spec written that way passes for a reason that has nothing to do
|
||||
// with the bound and would keep passing after the bound was
|
||||
// removed. This one can only be refused by the bound.
|
||||
oversized := append([]byte(`{"key":"`), bytes.Repeat([]byte("a"), maxControlRequestBytes+1)...)
|
||||
oversized = append(oversized, []byte(`"}`)...)
|
||||
resp, err := http.Post(srv.URL+path, "application/json", bytes.NewReader(oversized))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
|
||||
},
|
||||
Entry("ensure", workerctl.PathFilesEnsure),
|
||||
Entry("stage", workerctl.PathFilesStage),
|
||||
Entry("temp", workerctl.PathFilesTemp),
|
||||
Entry("listdir", workerctl.PathFilesListDir),
|
||||
)
|
||||
|
||||
It("reports a request body it cannot read as a rejection, never as a file that is not there", func() {
|
||||
resp, err := http.Post(srv.URL+workerctl.PathFilesEnsure, "application/json", strings.NewReader("{not json"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
|
||||
})
|
||||
|
||||
It("mounts every file verb, so none can be dropped from the set the frontend calls", func() {
|
||||
for _, path := range []string{
|
||||
workerctl.PathFilesEnsure, workerctl.PathFilesStage,
|
||||
workerctl.PathFilesTemp, workerctl.PathFilesListDir,
|
||||
} {
|
||||
resp := post(path, struct{}{})
|
||||
Expect(resp.StatusCode).NotTo(Equal(http.StatusNotFound), "%s is not mounted", path)
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -85,34 +85,7 @@ func (s *backendSupervisor) upgrader() upgradeFunc {
|
||||
// plane shares one bearer check with the file routes rather than growing a
|
||||
// second one.
|
||||
func (s *backendSupervisor) RegisterControlRoutes(mux *http.ServeMux) {
|
||||
// post registers one JSON-in, JSON-out verb. reply may be nil, in which
|
||||
// case the verb answers 204: that is the shape the two former
|
||||
// publish-no-reply subjects take.
|
||||
post := func(path string, h func(ctx context.Context, body []byte) (any, error)) {
|
||||
mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
|
||||
body, ok := readControlBody(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
reply, err := h(r.Context(), body)
|
||||
if err != nil {
|
||||
// Reaching here means the request could not be READ or routed,
|
||||
// which is this worker failing rather than answering. A verb's
|
||||
// own failure never reaches here: it comes back as a reply with
|
||||
// Error set, and a 200.
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if reply == nil {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if encErr := json.NewEncoder(w).Encode(reply); encErr != nil {
|
||||
xlog.Debug("worker control reply could not be written", "path", path, "error", encErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
post := func(path string, h controlVerb) { postControlVerb(mux, path, h) }
|
||||
|
||||
post(workerctl.PathModelsRunning, func(context.Context, []byte) (any, error) {
|
||||
return messaging.ModelsRunningReply{Models: s.runningModels()}, nil
|
||||
@@ -200,6 +173,46 @@ func (s *backendSupervisor) RegisterControlRoutes(mux *http.ServeMux) {
|
||||
})
|
||||
}
|
||||
|
||||
// controlVerb is the shape of a JSON-in, JSON-out control handler. A nil reply
|
||||
// means the verb answers 204, which is the shape the former publish-no-reply
|
||||
// subjects take.
|
||||
//
|
||||
// An error returned here means the request could not be READ or routed, which
|
||||
// is this worker FAILING rather than answering, and it becomes a non-2xx. A
|
||||
// verb's own failure is never returned this way: it comes back as a reply with
|
||||
// Error set, on a 200. See the note at the top of this file for why the two
|
||||
// must not be confused.
|
||||
type controlVerb func(ctx context.Context, body []byte) (any, error)
|
||||
|
||||
// postControlVerb registers one control verb on mux.
|
||||
//
|
||||
// It is the single door every control route enters through, so the POST-only
|
||||
// check, the bounded read and the 200-with-error-field shape are written once
|
||||
// rather than once per route set. The file-staging routes mount through it for
|
||||
// exactly that reason: a second registrar with its own copy of the rules is how
|
||||
// one of them ends up unbounded, or answering 500 for a verdict.
|
||||
func postControlVerb(mux *http.ServeMux, path string, h controlVerb) {
|
||||
mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
|
||||
body, ok := readControlBody(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
reply, err := h(r.Context(), body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if reply == nil {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if encErr := json.NewEncoder(w).Encode(reply); encErr != nil {
|
||||
xlog.Debug("worker control reply could not be written", "path", path, "error", encErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// readControlBody enforces the two things every control verb requires of a
|
||||
// request: that it is a POST, and that its body is bounded.
|
||||
//
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
@@ -447,8 +448,16 @@ 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 worker that has one, and the mounting assertion below walks every
|
||||
// path this package names.
|
||||
store, err := storage.NewFilesystemStore(filepath.Join(dir, "objectstore"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
stagingFM, err := storage.NewFileManager(store, filepath.Join(dir, "..", "cache"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
srv, err = startWorkerHTTPServer("127.0.0.1:0", filepath.Join(dir, "staging"), dir,
|
||||
filepath.Join(dir, "data"), token, &nodes.WorkerReadiness{}, sup, nil)
|
||||
filepath.Join(dir, "data"), token, &nodes.WorkerReadiness{}, sup, sup.cfg, stagingFM, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { nodes.ShutdownFileTransferServer(srv) })
|
||||
Expect(srv.Addr).NotTo(BeEmpty(), "the worker HTTP server must report the address it bound")
|
||||
@@ -481,6 +490,31 @@ var _ = Describe("the worker's HTTP server", func() {
|
||||
Expect(postCtl(workerctl.PathModelsRunning, "wrong").StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
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
|
||||
// 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()
|
||||
bare, err := startWorkerHTTPServer("127.0.0.1:0", filepath.Join(dir, "staging"), dir,
|
||||
filepath.Join(dir, "data"), token, &nodes.WorkerReadiness{}, sup, sup.cfg, nil, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { nodes.ShutdownFileTransferServer(bare) })
|
||||
|
||||
for _, p := range []string{
|
||||
workerctl.PathFilesEnsure, workerctl.PathFilesStage,
|
||||
workerctl.PathFilesTemp, workerctl.PathFilesListDir,
|
||||
} {
|
||||
req, reqErr := http.NewRequest(http.MethodPost, "http://"+bare.Addr+p, strings.NewReader("{}"))
|
||||
Expect(reqErr).NotTo(HaveOccurred())
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, doErr := http.DefaultClient.Do(req)
|
||||
Expect(doErr).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusNotFound), "%s answered without an object store", p)
|
||||
}
|
||||
})
|
||||
|
||||
It("mounts every control verb, not just the one this spec reads", func() {
|
||||
for _, p := range workerctl.AllPaths() {
|
||||
if p == workerctl.PathNodeStop {
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// The frontend's S3 file stager against the REAL worker file-staging routes: a
|
||||
// real object store both sides share, the real handlers mounted through the
|
||||
// real nodes.AuthenticatedRoutes so the real bearer check runs first, reached
|
||||
// by a real nodes.ControlClient over a real HTTP transport.
|
||||
//
|
||||
// It lives in this package for the same reason the control-client roundtrip
|
||||
// does: the dependency runs worker -> nodes, so the frontend half cannot be
|
||||
// exercised against the real handler from the other side without an import
|
||||
// cycle. A spec on either side alone proves only that side agrees with itself;
|
||||
// the path literals, the JSON shapes and the status codes are pinned together
|
||||
// only here.
|
||||
var _ = Describe("the frontend's file stager against the real worker", func() {
|
||||
const (
|
||||
token = "s3cret-registration-token"
|
||||
nodeID = "staging-worker"
|
||||
)
|
||||
|
||||
var (
|
||||
stager *nodes.S3FileStager
|
||||
store *storage.FilesystemStore
|
||||
workerFM *storage.FileManager
|
||||
modelsDir string
|
||||
srvAddr string
|
||||
)
|
||||
|
||||
newStager := func(tok string) *nodes.S3FileStager {
|
||||
GinkgoHelper()
|
||||
frontendFM, err := storage.NewFileManager(store, GinkgoT().TempDir())
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
control := nodes.NewControlClient(func(string) func(context.Context, string, string) (net.Conn, error) {
|
||||
return func(ctx context.Context, _, _ string) (net.Conn, error) {
|
||||
var d net.Dialer
|
||||
return d.DialContext(ctx, "tcp", srvAddr)
|
||||
}
|
||||
}, tok)
|
||||
return nodes.NewS3FileStager(frontendFM, control)
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
modelsDir = filepath.Join(dir, "worker", "models")
|
||||
Expect(os.MkdirAll(modelsDir, 0o750)).To(Succeed())
|
||||
|
||||
var err error
|
||||
store, err = storage.NewFilesystemStore(filepath.Join(dir, "objectstore"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
cfg := &Config{ModelsPath: modelsDir}
|
||||
workerFM, err = storage.NewFileManager(store, filepath.Join(dir, "worker", "cache"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
srv, err := nodes.StartFileTransferServerWithRoutes(lis,
|
||||
filepath.Join(dir, "worker", "staging"), modelsDir, filepath.Join(dir, "worker", "data"),
|
||||
token, config.DefaultMaxUploadSize, nil,
|
||||
&nodes.AuthenticatedRoutes{
|
||||
Prefix: workerctl.Prefix,
|
||||
Register: func(mux *http.ServeMux) {
|
||||
cfg.RegisterFileControlRoutes(mux, workerFM)
|
||||
},
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() { _ = srv.Close() })
|
||||
|
||||
srvAddr = lis.Addr().String()
|
||||
stager = newStager(token)
|
||||
})
|
||||
|
||||
It("puts a frontend file in the store and has the worker fetch it", func() {
|
||||
local := filepath.Join(GinkgoT().TempDir(), "model.gguf")
|
||||
Expect(os.WriteFile(local, []byte("checkpoint bytes"), 0o600)).To(Succeed())
|
||||
|
||||
remote, err := stager.EnsureRemote(context.Background(), nodeID, local, storage.ModelKey("rt/model.gguf"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(remote).To(BeAnExistingFile())
|
||||
Expect(os.ReadFile(remote)).To(Equal([]byte("checkpoint bytes")))
|
||||
})
|
||||
|
||||
It("allocates a temp file on the worker", func() {
|
||||
remote, err := stager.AllocRemoteTemp(context.Background(), nodeID)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(remote).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("allocates and downloads into the SAME cache directory", func() {
|
||||
// The staging cache root is derived once and read by two things: the
|
||||
// FileManager caches downloads into it, and the temp verb allocates
|
||||
// inside it. Two derivations that drifted would each stay self
|
||||
// consistent, so nothing else in this suite would notice; what would
|
||||
// notice is an operator whose disk budget covers one of the two.
|
||||
local := filepath.Join(GinkgoT().TempDir(), "same-root.gguf")
|
||||
Expect(os.WriteFile(local, []byte("bytes"), 0o600)).To(Succeed())
|
||||
downloaded, err := stager.EnsureRemote(context.Background(), nodeID, local, storage.ModelKey("root/x.gguf"))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
tmp, err := stager.AllocRemoteTemp(context.Background(), nodeID)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
cacheRoot := filepath.Join(filepath.Dir(modelsDir), "cache")
|
||||
Expect(downloaded).To(HavePrefix(cacheRoot + string(filepath.Separator)))
|
||||
Expect(tmp).To(HavePrefix(cacheRoot + string(filepath.Separator)))
|
||||
})
|
||||
|
||||
It("stages a worker file into the store", func() {
|
||||
remote := filepath.Join(modelsDir, "result.bin")
|
||||
Expect(os.WriteFile(remote, []byte("job output"), 0o600)).To(Succeed())
|
||||
|
||||
Expect(stager.StageRemoteToStore(context.Background(), nodeID, remote, "data/rt/result.bin")).To(Succeed())
|
||||
exists, err := store.Exists(context.Background(), "data/rt/result.bin")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(exists).To(BeTrue())
|
||||
})
|
||||
|
||||
It("fetches a worker file back through the store", func() {
|
||||
remote := filepath.Join(modelsDir, "fetched.bin")
|
||||
Expect(os.WriteFile(remote, []byte("fetch me"), 0o600)).To(Succeed())
|
||||
dst := filepath.Join(GinkgoT().TempDir(), "local.bin")
|
||||
|
||||
Expect(stager.FetchRemote(context.Background(), nodeID, remote, dst)).To(Succeed())
|
||||
Expect(os.ReadFile(dst)).To(Equal([]byte("fetch me")))
|
||||
})
|
||||
|
||||
It("lists a worker directory whose listing outgrows any bus payload", func() {
|
||||
big := filepath.Join(modelsDir, "wide")
|
||||
Expect(os.MkdirAll(big, 0o750)).To(Succeed())
|
||||
for i := range 4000 {
|
||||
name := fmt.Sprintf("shard-%030d.safetensors", i)
|
||||
Expect(os.WriteFile(filepath.Join(big, name), []byte("x"), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
files, err := stager.ListRemoteDir(context.Background(), nodeID, "models/wide")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(files).To(HaveLen(4000))
|
||||
})
|
||||
|
||||
It("reports the worker's own refusal as the worker's answer", func() {
|
||||
_, err := stager.ListRemoteDir(context.Background(), nodeID, "../../../etc")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("backend listdir failed"))
|
||||
// It said something about a file, not about the route, so it must not
|
||||
// be wearing the umbrella that stops a caller acting on it.
|
||||
Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeFalse())
|
||||
})
|
||||
|
||||
// One rule, stated once per verb: every file-staging RPC goes through the
|
||||
// control client, so a 401 is a failure of the ROUTE and never the worker
|
||||
// saying a file is not there. Pinned at each verb because each verb writes
|
||||
// the call out for itself, and a single verb that reached the worker some
|
||||
// other way would still leave the other five green.
|
||||
DescribeTable("reports a rejected token as unroutable, never as a verdict about a file",
|
||||
func(call func(*nodes.S3FileStager) error) {
|
||||
wrong := newStager("not-the-token")
|
||||
err := call(wrong)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeTrue(), "got %v", err)
|
||||
Expect(cluster.IsWorkerAnswer(err)).To(BeFalse())
|
||||
Expect(errors.Is(err, nodes.ErrWorkerControlUnsupported)).To(BeFalse())
|
||||
},
|
||||
Entry("ensure", func(s *nodes.S3FileStager) error {
|
||||
local := filepath.Join(GinkgoT().TempDir(), "m.gguf")
|
||||
Expect(os.WriteFile(local, []byte("x"), 0o600)).To(Succeed())
|
||||
_, err := s.EnsureRemote(context.Background(), nodeID, local, storage.ModelKey("auth/m.gguf"))
|
||||
return err
|
||||
}),
|
||||
Entry("temp", func(s *nodes.S3FileStager) error {
|
||||
_, err := s.AllocRemoteTemp(context.Background(), nodeID)
|
||||
return err
|
||||
}),
|
||||
Entry("listdir", func(s *nodes.S3FileStager) error {
|
||||
_, err := s.ListRemoteDir(context.Background(), nodeID, "models/wide")
|
||||
return err
|
||||
}),
|
||||
Entry("stage to store", func(s *nodes.S3FileStager) error {
|
||||
return s.StageRemoteToStore(context.Background(), nodeID, filepath.Join(modelsDir, "x"), "data/x")
|
||||
}),
|
||||
Entry("fetch", func(s *nodes.S3FileStager) error {
|
||||
return s.FetchRemote(context.Background(), nodeID, filepath.Join(modelsDir, "x"),
|
||||
filepath.Join(GinkgoT().TempDir(), "out"))
|
||||
}),
|
||||
Entry("fetch by key", func(s *nodes.S3FileStager) error {
|
||||
return s.FetchRemoteByKey(context.Background(), nodeID, "data/x",
|
||||
filepath.Join(GinkgoT().TempDir(), "out"))
|
||||
}),
|
||||
)
|
||||
})
|
||||
@@ -1,16 +1,8 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// isPathAllowed checks if path is within one of the allowed directories.
|
||||
@@ -35,167 +27,3 @@ func isPathAllowed(path string, allowedDirs []string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// subscribeFileStaging subscribes to NATS file staging subjects for this node.
|
||||
func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, nodeID string) error {
|
||||
// Create FileManager with same S3 config as the frontend
|
||||
// TODO: propagate a caller-provided context once Config carries one
|
||||
s3Store, err := storage.NewS3Store(context.Background(), storage.S3Config{
|
||||
Endpoint: cfg.StorageURL,
|
||||
Region: cfg.StorageRegion,
|
||||
Bucket: cfg.StorageBucket,
|
||||
AccessKeyID: cfg.StorageAccessKey,
|
||||
SecretAccessKey: cfg.StorageSecretKey,
|
||||
ForcePathStyle: true,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing S3 store: %w", err)
|
||||
}
|
||||
|
||||
cacheDir := filepath.Join(cfg.ModelsPath, "..", "cache")
|
||||
fm, err := storage.NewFileManager(s3Store, cacheDir)
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing file manager: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe: files.ensure — download S3 key to local, reply with local path
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) {
|
||||
var req struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": "invalid request"})
|
||||
return
|
||||
}
|
||||
|
||||
localPath, err := fm.Download(context.Background(), req.Key)
|
||||
if err != nil {
|
||||
xlog.Error("File ensure failed", "key", req.Key, "error", err)
|
||||
replyJSON(reply, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("File ensured locally", "key", req.Key, "path", localPath)
|
||||
replyJSON(reply, map[string]string{"local_path": localPath})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.ensure events: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe: files.stage — upload local path to S3, reply with key
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesStage(nodeID), func(data []byte, reply func([]byte)) {
|
||||
var req struct {
|
||||
LocalPath string `json:"local_path"`
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": "invalid request"})
|
||||
return
|
||||
}
|
||||
|
||||
allowedDirs := []string{cacheDir}
|
||||
if cfg.ModelsPath != "" {
|
||||
allowedDirs = append(allowedDirs, cfg.ModelsPath)
|
||||
}
|
||||
if !isPathAllowed(req.LocalPath, allowedDirs) {
|
||||
replyJSON(reply, map[string]string{"error": "path outside allowed directories"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := fm.Upload(context.Background(), req.Key, req.LocalPath); err != nil {
|
||||
xlog.Error("File stage failed", "path", req.LocalPath, "key", req.Key, "error", err)
|
||||
replyJSON(reply, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("File staged to S3", "path", req.LocalPath, "key", req.Key)
|
||||
replyJSON(reply, map[string]string{"key": req.Key})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.stage events: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe: files.temp — allocate temp file, reply with local path
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesTemp(nodeID), func(data []byte, reply func([]byte)) {
|
||||
tmpDir := filepath.Join(cacheDir, "staging-tmp")
|
||||
if err := os.MkdirAll(tmpDir, 0750); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": fmt.Sprintf("creating temp dir: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
f, err := os.CreateTemp(tmpDir, "localai-staging-*.tmp")
|
||||
if err != nil {
|
||||
replyJSON(reply, map[string]string{"error": fmt.Sprintf("creating temp file: %v", err)})
|
||||
return
|
||||
}
|
||||
localPath := f.Name()
|
||||
if err := f.Close(); err != nil {
|
||||
replyJSON(reply, map[string]string{"error": fmt.Sprintf("closing temp file: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("Allocated temp file", "path", localPath)
|
||||
replyJSON(reply, map[string]string{"local_path": localPath})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.temp events: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe: files.listdir — list files in a local directory, reply with relative paths
|
||||
if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesListDir(nodeID), func(data []byte, reply func([]byte)) {
|
||||
var req struct {
|
||||
KeyPrefix string `json:"key_prefix"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
replyJSON(reply, map[string]any{"error": "invalid request"})
|
||||
return
|
||||
}
|
||||
|
||||
// Resolve key prefix to local directory
|
||||
dirPath := filepath.Join(cacheDir, req.KeyPrefix)
|
||||
if rel, ok := strings.CutPrefix(req.KeyPrefix, storage.ModelKeyPrefix); ok && cfg.ModelsPath != "" {
|
||||
dirPath = filepath.Join(cfg.ModelsPath, rel)
|
||||
} else if rel, ok := strings.CutPrefix(req.KeyPrefix, storage.DataKeyPrefix); ok {
|
||||
dirPath = filepath.Join(cacheDir, "..", "data", rel)
|
||||
}
|
||||
|
||||
// Sanitize to prevent directory traversal via crafted key_prefix
|
||||
dirPath = filepath.Clean(dirPath)
|
||||
cleanCache := filepath.Clean(cacheDir)
|
||||
cleanModels := filepath.Clean(cfg.ModelsPath)
|
||||
cleanData := filepath.Clean(filepath.Join(cacheDir, "..", "data"))
|
||||
if !(strings.HasPrefix(dirPath, cleanCache+string(filepath.Separator)) ||
|
||||
dirPath == cleanCache ||
|
||||
(cleanModels != "." && strings.HasPrefix(dirPath, cleanModels+string(filepath.Separator))) ||
|
||||
dirPath == cleanModels ||
|
||||
strings.HasPrefix(dirPath, cleanData+string(filepath.Separator)) ||
|
||||
dirPath == cleanData) {
|
||||
replyJSON(reply, map[string]any{"error": "invalid key prefix"})
|
||||
return
|
||||
}
|
||||
|
||||
var files []string
|
||||
if err := filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !d.IsDir() {
|
||||
rel, err := filepath.Rel(dirPath, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
files = append(files, rel)
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
xlog.Error("Failed to list staged files", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "error", err)
|
||||
replyJSON(reply, map[string]any{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("Listed remote dir", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "fileCount", len(files))
|
||||
replyJSON(reply, map[string]any{"files": files})
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to files.listdir events: %w", err)
|
||||
}
|
||||
|
||||
xlog.Info("Subscribed to file staging NATS subjects", "nodeID", nodeID)
|
||||
return nil
|
||||
}
|
||||
@@ -1,17 +0,0 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// replyJSON marshals v to JSON and calls the reply function.
|
||||
func replyJSON(reply func([]byte), v any) {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
xlog.Error("Failed to marshal NATS reply", "error", err)
|
||||
data = []byte(`{"error":"internal marshal error"}`)
|
||||
}
|
||||
reply(data)
|
||||
}
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
@@ -225,8 +226,26 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
maxPort: cfg.effectiveMaxPort(basePort),
|
||||
}
|
||||
|
||||
// The file-staging FileManager is built BEFORE the server, for the same
|
||||
// reason the supervisor is: the server is what serves those four verbs, and
|
||||
// a route that is not mounted when the tunnel comes up is a 404 the
|
||||
// frontend reads as a worker that does not implement the verb. It used to
|
||||
// be built after NATS connected, which is a window that no longer exists.
|
||||
//
|
||||
// A worker with no object store configured mounts NO file verbs. That is
|
||||
// the honest answer rather than a degraded one: it has nowhere to fetch
|
||||
// from or stage to, and the frontend that reaches such a deployment uses
|
||||
// the HTTP file stager, which does not call these paths at all.
|
||||
var stagingFM *storage.FileManager
|
||||
if cfg.StorageURL != "" {
|
||||
stagingFM, err = cfg.NewStagingFileManager(shutdownCtx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("initializing file staging: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
httpServer, err := startWorkerHTTPServer(httpAddr, stagingDir, cfg.ModelsPath, dataDir,
|
||||
cfg.RegistrationToken, readiness, supervisor, ml.BackendLogs())
|
||||
cfg.RegistrationToken, readiness, supervisor, cfg, stagingFM, ml.BackendLogs())
|
||||
if err != nil {
|
||||
return fmt.Errorf("starting HTTP file transfer server: %w", err)
|
||||
}
|
||||
@@ -305,14 +324,6 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
}
|
||||
}()
|
||||
|
||||
// Subscribe to file staging NATS subjects if S3 is configured
|
||||
if cfg.StorageURL != "" {
|
||||
if err := cfg.subscribeFileStaging(natsClient, nodeID); err != nil {
|
||||
nodes.ShutdownFileTransferServer(httpServer)
|
||||
return fmt.Errorf("subscribing to file staging subjects: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
xlog.Info("Worker ready, serving its control plane over the tunnel")
|
||||
// Exit on an OS signal or on an internal fatal condition (e.g. NATS
|
||||
// credentials became unrenewable), so the worker restarts and re-acquires
|
||||
@@ -343,10 +354,19 @@ func Run(ctx *cliContext.Context, cfg *Config) error {
|
||||
// the server without the control plane and producing a worker that looks
|
||||
// healthy while answering 404 to every command.
|
||||
func startWorkerHTTPServer(addr, stagingDir, modelsDir, dataDir, token string,
|
||||
readiness *nodes.WorkerReadiness, sup *backendSupervisor, logStore *model.BackendLogStore) (*http.Server, error) {
|
||||
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{
|
||||
Prefix: workerctl.Prefix,
|
||||
Register: sup.RegisterControlRoutes,
|
||||
Prefix: workerctl.Prefix,
|
||||
// One registrar for both route sets, because there is ONE control
|
||||
// prefix and AuthenticatedRoutes mounts one mux behind one bearer
|
||||
// check. A second route set would be a second check to forget.
|
||||
Register: func(mux *http.ServeMux) {
|
||||
sup.RegisterControlRoutes(mux)
|
||||
if stagingFM != nil {
|
||||
cfg.RegisterFileControlRoutes(mux, stagingFM)
|
||||
}
|
||||
},
|
||||
}, logStore)
|
||||
}
|
||||
@@ -29,6 +29,17 @@ const (
|
||||
PathModelDelete = "/v1/control/model/delete"
|
||||
PathModelsRunning = "/v1/control/models/running"
|
||||
PathNodeStop = "/v1/control/node/stop"
|
||||
|
||||
// The file-staging verbs. They are only ever served by a worker whose
|
||||
// deployment configured an object store, because without one there is
|
||||
// nothing for them to move a file to or from; a worker without one mounts
|
||||
// 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.
|
||||
PathFilesEnsure = "/v1/control/files/ensure"
|
||||
PathFilesStage = "/v1/control/files/stage"
|
||||
PathFilesTemp = "/v1/control/files/temp"
|
||||
PathFilesListDir = "/v1/control/files/listdir"
|
||||
)
|
||||
|
||||
// AllPaths returns every control verb's path.
|
||||
@@ -47,6 +58,10 @@ func AllPaths() []string {
|
||||
PathModelDelete,
|
||||
PathModelsRunning,
|
||||
PathNodeStop,
|
||||
PathFilesEnsure,
|
||||
PathFilesStage,
|
||||
PathFilesTemp,
|
||||
PathFilesListDir,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -28,6 +28,10 @@ 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 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"),
|
||||
Entry("files listdir", workerctl.PathFilesListDir, "/v1/control/files/listdir"),
|
||||
)
|
||||
|
||||
It("names the prefix exactly, since the worker mounts its whole control plane behind it", func() {
|
||||
@@ -57,6 +61,10 @@ var _ = Describe("control plane paths on the wire", func() {
|
||||
workerctl.PathModelDelete,
|
||||
workerctl.PathModelsRunning,
|
||||
workerctl.PathNodeStop,
|
||||
workerctl.PathFilesEnsure,
|
||||
workerctl.PathFilesStage,
|
||||
workerctl.PathFilesListDir,
|
||||
workerctl.PathFilesTemp,
|
||||
))
|
||||
})
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ Each model gets its own gRPC backend process, so a single worker can serve multi
|
||||
## Prerequisites
|
||||
|
||||
- **PostgreSQL** (with pgvector extension recommended for RAG) - used for node registry, job store, auth, and shared state
|
||||
- **NATS** server - used for real-time backend lifecycle events and file staging
|
||||
- **NATS** server - used for agent-worker coordination and the frontend's own cross-replica events. Serve-backend workers no longer take any command over it: the backend and model lifecycle verbs and file staging are HTTP routes on the worker's tunnel.
|
||||
- All services must be on the same network (or reachable via configured URLs)
|
||||
|
||||
## Quick Start with Docker Compose
|
||||
@@ -170,17 +170,24 @@ Every connection the frontend makes to a worker now goes through that worker's t
|
||||
| Inference, model load, health checks | gRPC to a backend process | `grpc` |
|
||||
| Model file staging, backend-log listing | HTTP to the worker's own server | `http` |
|
||||
| Live backend-log streaming | WebSocket to the same server | `http` |
|
||||
| The control plane: backend install/upgrade/list/stop/delete, model stop/unload/delete, models running, node stop | HTTP to the worker's own server | `http` |
|
||||
| The control plane: backend install/upgrade/list/stop/delete, model stop/unload/delete, models running, node stop, and the four object-store staging verbs | HTTP to the worker's own server | `http` |
|
||||
|
||||
#### The worker control plane
|
||||
|
||||
A serve-backend worker serves the commands the frontend gives it as ordinary
|
||||
HTTP routes under `/v1/control/`, on the same loopback server that already
|
||||
carries file staging and backend logs, behind the same `LOCALAI_REGISTRATION_TOKEN`
|
||||
bearer check. They replace the ten `nodes.<id>.*` NATS subjects a worker used to
|
||||
subscribe to; the request and reply bodies are unchanged, so nothing an operator
|
||||
inspects on the wire has a new shape. Agent workers still take `nodes.<id>.backend.stop`
|
||||
over NATS.
|
||||
bearer check. They replace the fourteen `nodes.<id>.*` NATS subjects a worker
|
||||
used to subscribe to - the ten backend and model lifecycle verbs, plus the four
|
||||
object-store staging verbs (`POST /v1/control/files/{ensure,stage,temp,listdir}`,
|
||||
mounted only when the deployment configured an object store). The request and
|
||||
reply bodies are unchanged, so nothing an operator inspects on the wire has a new
|
||||
shape. Agent workers still take `nodes.<id>.backend.stop` over NATS.
|
||||
|
||||
`files/listdir` is the verb the change is most visible on. Its reply used to be
|
||||
sized against what the bus would carry, which put a wide model directory close to
|
||||
the limit; it is now a response body the frontend is already reading, so the
|
||||
listing is returned whole and nothing truncates it at either end.
|
||||
|
||||
Two of the routes stream. `POST /v1/control/backend/install` and
|
||||
`/v1/control/backend/upgrade` answer with `application/x-ndjson`: zero or more
|
||||
@@ -399,7 +406,7 @@ A frontend replica that dies mid-load does not wedge the model: the job row carr
|
||||
|
||||
### NATS JWT authentication (recommended for production)
|
||||
|
||||
By default, NATS connections are anonymous: any client that can reach port `4222` may publish the subjects still carried on it. Serve-backend workers no longer subscribe to `nodes.<id>.backend.install` and its nine siblings - those are HTTP routes on the worker's tunnel now, see [The worker control plane](#the-worker-control-plane) - but agent workers, file staging and the frontend's own service credential still use NATS. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential.
|
||||
By default, NATS connections are anonymous: any client that can reach port `4222` may publish the subjects still carried on it. Serve-backend workers no longer subscribe to `nodes.<id>.backend.install` and its nine siblings - those are HTTP routes on the worker's tunnel now, see [The worker control plane](#the-worker-control-plane) - but agent workers and the frontend's own service credential still use NATS. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential.
|
||||
|
||||
| Flag | Env Var | Description |
|
||||
|------|---------|-------------|
|
||||
@@ -443,7 +450,7 @@ Generate operator/account material with [`scripts/nats-auth-setup.sh`](https://g
|
||||
|
||||
### Optional: S3 Object Storage
|
||||
|
||||
For multi-host deployments where workers don't share a filesystem, S3-compatible storage enables distributed file transfer (model files, configs):
|
||||
For multi-host deployments where workers don't share a filesystem, S3-compatible storage enables distributed file transfer (model files, configs). The frontend uploads the file to the bucket and then tells the worker to fetch it, over that worker's tunnel (`POST /v1/control/files/ensure`); the reverse direction (`.../files/stage`) has the worker upload one of its own files for the frontend to pull down. The bytes travel through the bucket, never through the tunnel:
|
||||
|
||||
| Flag | Env Var | Default | Description |
|
||||
|------|---------|---------|-------------|
|
||||
@@ -453,6 +460,8 @@ For multi-host deployments where workers don't share a filesystem, S3-compatible
|
||||
| `--storage-access-key` | `LOCALAI_STORAGE_ACCESS_KEY` | *(empty)* | S3 access key |
|
||||
| `--storage-secret-key` | `LOCALAI_STORAGE_SECRET_KEY` | *(empty)* | S3 secret key |
|
||||
|
||||
A worker started without `LOCALAI_STORAGE_URL` does not serve the four staging verbs at all, and answers `404` for them, which is the same answer a frontend gets from a worker too old to know them.
|
||||
|
||||
When S3 is not configured, model files are transferred directly from the frontend to workers via **HTTP** - no shared filesystem needed. Each worker runs a small HTTP file transfer server alongside the gRPC backend process. This is the default and works out of the box.
|
||||
|
||||
For high-throughput or very large model files, S3 can be more efficient since it avoids streaming through the frontend.
|
||||
@@ -474,7 +483,10 @@ Hugging Face for managed artifacts.
|
||||
|
||||
With `LOCALAI_DISTRIBUTED_SHARED_MODELS` enabled, workers use the shared
|
||||
absolute snapshot path and skip transfer. Otherwise, the controller stages the
|
||||
complete snapshot tree to each worker before loading the backend.
|
||||
complete snapshot tree to each worker before loading the backend. With an object
|
||||
store configured the controller uploads to the bucket and commands the worker to
|
||||
fetch over its tunnel; without one it pushes the files to the worker's HTTP file
|
||||
transfer server directly. Neither path uses NATS.
|
||||
|
||||
{{% notice warning %}}
|
||||
Every controller and worker must have enough disk space for its own snapshot
|
||||
@@ -546,7 +558,7 @@ local-ai worker \
|
||||
| `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | Umbrella switch implying both `--registration-require-auth` and `--nats-require-auth` |
|
||||
| `--heartbeat-interval` | `LOCALAI_HEARTBEAT_INTERVAL` | `10s` | Interval between heartbeat pings |
|
||||
| `--worker-tunnel` | `LOCALAI_WORKER_TUNNEL` | `true` | Hold one outbound multiplexed tunnel to the frontend and serve its requests over it, so this worker needs no inbound port (see [Worker tunnels](#worker-tunnels)). Setting it to `false` is a **fatal startup error**, not a degraded mode: the frontend has no path that dials a worker's advertised address, so a worker without its tunnel is a worker nothing can reach. To run without tunnels, run the pre-tunnel release on both the worker and the frontend. |
|
||||
| `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS URL for backend installation and file staging |
|
||||
| `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS URL. A serve-backend worker takes no command over it - backend installation and file staging are HTTP routes on its tunnel - but it must still connect for the frontend to consider it healthy. |
|
||||
| `--nats-jwt` | `LOCALAI_NATS_JWT` | *(empty)* | Optional override for the `nats_jwt` returned at registration |
|
||||
| `--nats-user-seed` | `LOCALAI_NATS_USER_SEED` | *(empty)* | Optional override for `nats_user_seed` from registration |
|
||||
| `--nats-require-auth` | `LOCALAI_NATS_REQUIRE_AUTH` | `false` | Require NATS JWT+seed (from registration or env) |
|
||||
|
||||
@@ -37,10 +37,11 @@ var _ = Describe("MintWorkerJWT", func() {
|
||||
uc, err := jwt.DecodeUserClaims(token)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(uc.Permissions.Sub.Allow).To(ContainElement("nodes.550e8400-e29b-41d4-a716-446655440000.>"))
|
||||
Expect(uc.Permissions.Pub.Allow).To(ContainElement("nodes.550e8400-e29b-41d4-a716-446655440000.files.>"))
|
||||
// The install-progress subject is gone with the carrier: progress is a
|
||||
// line in the install response now, so a minted worker JWT must not
|
||||
// still be granted a publish right for it.
|
||||
// still be granted a publish right for it. File staging went the same
|
||||
// way, so a minted worker JWT publishes nowhere but its own inbox.
|
||||
Expect(uc.Permissions.Pub.Allow).To(ConsistOf("_INBOX.>"))
|
||||
for _, subj := range uc.Permissions.Pub.Allow {
|
||||
Expect(subj).NotTo(ContainSubstring("backend.install"))
|
||||
}
|
||||
|
||||
@@ -38,22 +38,23 @@ func WorkerPermissions(nodeID, nodeType string) (pubAllow, subAllow []string) {
|
||||
"_INBOX.>",
|
||||
}
|
||||
default:
|
||||
// Backend worker: file staging on this node only.
|
||||
//
|
||||
// The backend and model lifecycle verbs left the bus: they are HTTP
|
||||
// Backend worker. Every verb a frontend gives it left the bus: the
|
||||
// backend and model lifecycle verbs and now file staging too are HTTP
|
||||
// routes under workerctl.Prefix, served on the worker's own server and
|
||||
// reached through its tunnel, so no subject is minted for them and none
|
||||
// is allowed here. The wildcard stays because the file-staging subjects
|
||||
// are still under this node's prefix; narrowing it to nodes.<id>.files.>
|
||||
// is a separate change that would break a worker mid-upgrade.
|
||||
// is allowed here.
|
||||
//
|
||||
// The subscribe wildcard stays for now. A worker subscribes to nothing
|
||||
// under it on this build, but narrowing it is a change a worker
|
||||
// mid-upgrade would feel, and the connection itself is what the next
|
||||
// step of this removal deletes.
|
||||
subAllow = []string{
|
||||
prefix + ".>",
|
||||
"_INBOX.>",
|
||||
}
|
||||
// backend.install.*.progress is gone with the subject: install progress
|
||||
// is written into the install response the frontend is already reading.
|
||||
// Nothing left to publish. backend.install.*.progress went with the
|
||||
// install subject, and the file-staging replies went with theirs.
|
||||
pubAllow = []string{
|
||||
prefix + ".files.>",
|
||||
"_INBOX.>",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,37 +59,32 @@ var _ = Describe("WorkerPermissions subject coverage", func() {
|
||||
Context("backend worker", func() {
|
||||
pub, sub := natsauth.WorkerPermissions(nodeID, "backend")
|
||||
|
||||
// Every subject core/services/worker/file_staging.go subscribes to.
|
||||
// The backend and model lifecycle verbs are NOT here: they left the bus
|
||||
// for the worker's tunnelled control plane, so there is no subject to
|
||||
// cover. See core/services/workerctl.
|
||||
subscribed := []string{
|
||||
messaging.SubjectNodeFilesEnsure(nodeID),
|
||||
messaging.SubjectNodeFilesStage(nodeID),
|
||||
messaging.SubjectNodeFilesTemp(nodeID),
|
||||
messaging.SubjectNodeFilesListDir(nodeID),
|
||||
}
|
||||
for _, subject := range subscribed {
|
||||
It("allows subscribing to "+subject, func() {
|
||||
Expect(anyAllows(sub, subject)).To(BeTrue(),
|
||||
"backend JWT sub allow-list %v does not cover %s", sub, subject)
|
||||
})
|
||||
}
|
||||
|
||||
It("allows publishing file staging replies", func() {
|
||||
subject := messaging.SubjectNodeFilesStage(nodeID)
|
||||
Expect(anyAllows(pub, subject)).To(BeTrue(),
|
||||
"backend JWT pub allow-list %v does not cover %s", pub, subject)
|
||||
// A backend worker subscribes to no subject of its own on this build.
|
||||
// Every verb a frontend gives it — the backend and model lifecycle ten,
|
||||
// and now the four file-staging verbs — is an HTTP route on its
|
||||
// tunnelled control plane, so there is no subject left to cover. See
|
||||
// core/services/workerctl.
|
||||
//
|
||||
// The subscribe wildcard is asserted rather than removed because the
|
||||
// grant is still minted and a worker mid-upgrade still uses it.
|
||||
It("still grants a backend worker its own node subtree to subscribe on", func() {
|
||||
Expect(sub).To(ConsistOf(
|
||||
"nodes."+workerSubjectTokenForTest(nodeID)+".>",
|
||||
"_INBOX.>",
|
||||
))
|
||||
})
|
||||
|
||||
// The negative half, and it is the one that would catch a verb quietly
|
||||
// coming back to the bus: a backend worker is granted nothing to
|
||||
// publish outside its own file-staging subtree and its inbox.
|
||||
It("grants a backend worker no publish rights outside file staging and its inbox", func() {
|
||||
Expect(pub).To(ConsistOf(
|
||||
"nodes."+workerSubjectTokenForTest(nodeID)+".files.>",
|
||||
"_INBOX.>",
|
||||
))
|
||||
// publish at all beyond its own inbox. File staging used to be the one
|
||||
// exception and is not any more.
|
||||
It("grants a backend worker no publish rights outside its inbox", func() {
|
||||
Expect(pub).To(ConsistOf("_INBOX.>"))
|
||||
})
|
||||
|
||||
It("no longer grants a backend worker the file-staging publish subtree", func() {
|
||||
Expect(anyAllows(pub, "nodes."+workerSubjectTokenForTest(nodeID)+".files.stage")).To(BeFalse(),
|
||||
"backend JWT pub allow-list %v still covers file staging", pub)
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -7,9 +7,9 @@ import (
|
||||
"path/filepath"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"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"
|
||||
@@ -41,8 +41,8 @@ var _ = Describe("File Staging", Label("Distributed"), func() {
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
|
||||
Context("S3NATSFileStager", func() {
|
||||
It("should create S3NATSFileStager with valid config", func() {
|
||||
Context("S3FileStager", func() {
|
||||
It("should create S3FileStager with valid config", func() {
|
||||
storeDir := filepath.Join(tmpDir, "objectstore")
|
||||
cacheDir := filepath.Join(tmpDir, "cache")
|
||||
|
||||
@@ -53,7 +53,7 @@ var _ = Describe("File Staging", Label("Distributed"), func() {
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(fm.IsConfigured()).To(BeTrue())
|
||||
|
||||
stager := nodes.NewS3NATSFileStager(fm, infra.NC)
|
||||
stager := nodes.NewS3FileStager(fm, nodes.NewControlClient(directWorkerDialerFor, ""))
|
||||
Expect(stager).ToNot(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -83,8 +83,8 @@ var _ = Describe("File Staging", Label("Distributed"), func() {
|
||||
})
|
||||
})
|
||||
|
||||
Context("S3NATSFileStager with backend node simulation", func() {
|
||||
It("should coordinate file staging via NATS request-reply", func() {
|
||||
Context("S3FileStager with backend node simulation", func() {
|
||||
It("should coordinate file staging over the worker's control plane", func() {
|
||||
storeDir := filepath.Join(tmpDir, "objectstore")
|
||||
cacheDir := filepath.Join(tmpDir, "cache")
|
||||
|
||||
@@ -103,15 +103,14 @@ var _ = Describe("File Staging", Label("Distributed"), func() {
|
||||
}
|
||||
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
||||
|
||||
// Verify NATS file staging subjects are correctly formed
|
||||
ensureSubj := messaging.SubjectNodeFilesEnsure(node.ID)
|
||||
Expect(ensureSubj).To(ContainSubstring("files.ensure"))
|
||||
|
||||
stageSubj := messaging.SubjectNodeFilesStage(node.ID)
|
||||
Expect(stageSubj).To(ContainSubstring("files.stage"))
|
||||
|
||||
tempSubj := messaging.SubjectNodeFilesTemp(node.ID)
|
||||
Expect(tempSubj).To(ContainSubstring("files.temp"))
|
||||
// The staging verbs are HTTP routes on the worker's own control
|
||||
// plane now, so what a registered node is addressed by is its
|
||||
// tunnel host and the path, not a subject.
|
||||
Expect(nodes.WorkerHTTPHost(node.ID, "")).To(ContainSubstring(node.ID))
|
||||
Expect(workerctl.PathFilesEnsure).To(HavePrefix(workerctl.Prefix))
|
||||
Expect(workerctl.PathFilesStage).To(HavePrefix(workerctl.Prefix))
|
||||
Expect(workerctl.PathFilesTemp).To(HavePrefix(workerctl.Prefix))
|
||||
Expect(workerctl.PathFilesListDir).To(HavePrefix(workerctl.Prefix))
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
Reference in new issue
Block a user