diff --git a/core/application/distributed.go b/core/application/distributed.go index af7bd5501..ad822cc1a 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -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, diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index fe9c1691a..3bb8a49eb 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -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..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 ( diff --git a/core/services/nodes/file_stager_s3.go b/core/services/nodes/file_stager_s3.go index 0d3847b7c..19015ed5c 100644 --- a/core/services/nodes/file_stager_s3.go +++ b/core/services/nodes/file_stager_s3.go @@ -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..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 != "" { diff --git a/core/services/nodes/file_stager_s3_test.go b/core/services/nodes/file_stager_s3_test.go new file mode 100644 index 000000000..8fdc4ae01 --- /dev/null +++ b/core/services/nodes/file_stager_s3_test.go @@ -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()) + }) +}) diff --git a/core/services/worker/control_files.go b/core/services/worker/control_files.go new file mode 100644 index 000000000..2b05781a9 --- /dev/null +++ b/core/services/worker/control_files.go @@ -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..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 +} diff --git a/core/services/worker/control_files_test.go b/core/services/worker/control_files_test.go new file mode 100644 index 000000000..42a815c05 --- /dev/null +++ b/core/services/worker/control_files_test.go @@ -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) + } + }) +}) diff --git a/core/services/worker/control_routes.go b/core/services/worker/control_routes.go index d787f23cc..56d251cac 100644 --- a/core/services/worker/control_routes.go +++ b/core/services/worker/control_routes.go @@ -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. // diff --git a/core/services/worker/control_routes_test.go b/core/services/worker/control_routes_test.go index f57793324..584e118c2 100644 --- a/core/services/worker/control_routes_test.go +++ b/core/services/worker/control_routes_test.go @@ -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 { diff --git a/core/services/worker/file_stager_roundtrip_test.go b/core/services/worker/file_stager_roundtrip_test.go new file mode 100644 index 000000000..1c8e45530 --- /dev/null +++ b/core/services/worker/file_stager_roundtrip_test.go @@ -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")) + }), + ) +}) diff --git a/core/services/worker/file_staging.go b/core/services/worker/file_staging.go index 019afcba9..9e65f2ced 100644 --- a/core/services/worker/file_staging.go +++ b/core/services/worker/file_staging.go @@ -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 -} diff --git a/core/services/worker/reply.go b/core/services/worker/reply.go deleted file mode 100644 index 9700f19fa..000000000 --- a/core/services/worker/reply.go +++ /dev/null @@ -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) -} diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 22f116654..6ccf1bbc7 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -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) } diff --git a/core/services/workerctl/paths.go b/core/services/workerctl/paths.go index e76b5566d..4af1164d8 100644 --- a/core/services/workerctl/paths.go +++ b/core/services/workerctl/paths.go @@ -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, } } diff --git a/core/services/workerctl/paths_test.go b/core/services/workerctl/paths_test.go index 8ffabb2e3..4d72124ec 100644 --- a/core/services/workerctl/paths_test.go +++ b/core/services/workerctl/paths_test.go @@ -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, )) }) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 865b2b7af..95febd2e1 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -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..*` 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..backend.stop` -over NATS. +bearer check. They replace the fourteen `nodes..*` 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..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..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..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) | diff --git a/pkg/natsauth/mint_test.go b/pkg/natsauth/mint_test.go index e33992ec1..195f05b2c 100644 --- a/pkg/natsauth/mint_test.go +++ b/pkg/natsauth/mint_test.go @@ -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")) } diff --git a/pkg/natsauth/permissions.go b/pkg/natsauth/permissions.go index a411e05fd..8e0fa985e 100644 --- a/pkg/natsauth/permissions.go +++ b/pkg/natsauth/permissions.go @@ -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..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.>", } } diff --git a/pkg/natsauth/permissions_coverage_test.go b/pkg/natsauth/permissions_coverage_test.go index b02541e9a..bc5cfc693 100644 --- a/pkg/natsauth/permissions_coverage_test.go +++ b/pkg/natsauth/permissions_coverage_test.go @@ -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) }) }) diff --git a/tests/e2e/distributed/file_staging_test.go b/tests/e2e/distributed/file_staging_test.go index e765d87c3..a22a639c4 100644 --- a/tests/e2e/distributed/file_staging_test.go +++ b/tests/e2e/distributed/file_staging_test.go @@ -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)) }) })