diff --git a/core/services/nodes/file_stager.go b/core/services/nodes/file_stager.go index c5ee38556..0c4556ce0 100644 --- a/core/services/nodes/file_stager.go +++ b/core/services/nodes/file_stager.go @@ -5,11 +5,14 @@ import "context" // FileStager abstracts file transfer between frontend and backend nodes // in distributed mode. Two implementations exist: // -// 1. S3NATSFileStager (primary): Both sides have FileManager with same S3. -// Frontend uploads to S3, sends NATS request-reply to backend to download locally. +// 1. S3FileStager (primary): Both sides have FileManager with same S3. +// Frontend uploads to S3, then calls the worker's control plane over its +// tunnel to have it download locally. // // 2. HTTPFileStager (fallback): Frontend pushes/pulls files directly over // HTTP to a small file transfer server on the backend node (no S3 needed). +// +// Both reach the worker through its tunnel and neither uses NATS. type FileStager interface { // EnsureRemote ensures a local file is available on the remote node. // Returns the remote-local path. diff --git a/core/services/nodes/file_stager_s3.go b/core/services/nodes/file_stager_s3.go index 19015ed5c..ddef12d6f 100644 --- a/core/services/nodes/file_stager_s3.go +++ b/core/services/nodes/file_stager_s3.go @@ -43,6 +43,29 @@ const ( fileMetadataRPCTimeout = 30 * time.Second ) +// fileRPCBudget is the ceiling one file-staging verb's RPC gets. +// +// The mapping lives here rather than at the call sites, and that is the same +// argument callWorker is written under: five sites each naming their own +// constant is five chances to name the wrong one, and a stage verb given the +// metadata ceiling would abandon a multi-gigabyte copy after thirty seconds +// while the worker went on making it. +// +// A path this function does not know gets the SHORTER ceiling. That is the safe +// direction: a verb wrongly given 30 seconds fails visibly and is retried, +// while one wrongly given ten minutes parks a caller on a verb that was never +// meant to be slow. +func fileRPCBudget(path string) time.Duration { + switch path { + case workerctl.PathFilesEnsure, workerctl.PathFilesStage: + return fileTransferRPCTimeout + case workerctl.PathFilesTemp, workerctl.PathFilesListDir: + return fileMetadataRPCTimeout + default: + return fileMetadataRPCTimeout + } +} + // 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. @@ -83,15 +106,16 @@ type fileListDirReply struct { } // callWorker issues one file-staging RPC under a deadline derived from the -// caller's context. +// caller's context and bounded by the verb's own ceiling. // // 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) +// for it. The budget comes from the path rather than from an argument for the +// same reason; see fileRPCBudget. +func (s *S3FileStager) callWorker(ctx context.Context, nodeID, path string, req, reply any) error { + rpcCtx, cancel := context.WithTimeout(ctx, fileRPCBudget(path)) defer cancel() return s.control.Call(rpcCtx, nodeID, path, req, reply) } @@ -115,8 +139,7 @@ func (s *S3FileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key } var reply fileEnsureReply - if err := s.callWorker(ctx, nodeID, workerctl.PathFilesEnsure, fileTransferRPCTimeout, - fileEnsureRequest{Key: key}, &reply); err != nil { + if err := s.callWorker(ctx, nodeID, workerctl.PathFilesEnsure, fileEnsureRequest{Key: key}, &reply); err != nil { return "", err } if reply.Error != "" { @@ -145,8 +168,7 @@ func (s *S3FileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localD 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 { + if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, fileStageRequest{LocalPath: remotePath, Key: key}, &reply); err != nil { return err } if reply.Error != "" { @@ -177,8 +199,7 @@ func (s *S3FileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePat // 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 { + if err := s.callWorker(ctx, nodeID, workerctl.PathFilesTemp, fileTempRequest{}, &reply); err != nil { return "", err } if reply.Error != "" { @@ -197,8 +218,7 @@ func (s *S3FileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (stri // 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 { + if err := s.callWorker(ctx, nodeID, workerctl.PathFilesListDir, fileListDirRequest{KeyPrefix: keyPrefix}, &reply); err != nil { return nil, err } if reply.Error != "" { @@ -211,8 +231,7 @@ func (s *S3FileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix stri // 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 { + if err := s.callWorker(ctx, nodeID, workerctl.PathFilesStage, 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 index 8fdc4ae01..beceeb2fa 100644 --- a/core/services/nodes/file_stager_s3_test.go +++ b/core/services/nodes/file_stager_s3_test.go @@ -5,6 +5,7 @@ import ( "errors" "os" "path/filepath" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -44,6 +45,26 @@ var _ = Describe("the S3 file stager's control RPCs", func() { Expect(os.WriteFile(local, []byte("weights"), 0o600)).To(Succeed()) }) + // The budgets are written out BY HAND and deliberately not derived from the + // constants under test. They are the timeouts the NATS request-reply calls + // carried, and keeping them is a compatibility fact rather than a taste: + // an operator whose staging of a 35 GB checkpoint fits inside ten minutes + // today must not find it cut to thirty seconds by the carrier changing. + DescribeTable("gives each verb the ceiling its NATS timeout carried", + func(path string, want time.Duration) { Expect(fileRPCBudget(path)).To(Equal(want)) }, + Entry("ensure moves bytes", workerctl.PathFilesEnsure, 10*time.Minute), + Entry("stage moves bytes", workerctl.PathFilesStage, 10*time.Minute), + Entry("temp is metadata", workerctl.PathFilesTemp, 30*time.Second), + Entry("listdir is metadata", workerctl.PathFilesListDir, 30*time.Second), + ) + + It("gives a verb it does not know the shorter ceiling", func() { + // The safe direction: a verb wrongly given thirty seconds fails visibly + // and is retried, while one wrongly given ten minutes parks a caller on + // a verb that was never meant to be slow. + Expect(fileRPCBudget(workerctl.Prefix + "invented")).To(Equal(30 * time.Second)) + }) + DescribeTable("addresses the verb that verb's path names", func(path string, reply any, call func(*S3FileStager) error) { workers.scriptReply(controlKey(nodeID, path), reply) diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index 2fdedc3d2..45dbe3a57 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -22,7 +22,9 @@ import ( // for distributed mode. Input files are staged on the backend node before the // gRPC call. Output files are retrieved from the backend after the call. // -// Uses the FileStager interface — agnostic to transport (S3+NATS or gRPC). +// Uses the FileStager interface — agnostic to transport (an object store, or +// direct HTTP to the worker), and in both cases reached over the worker's +// tunnel. // The caller gets a grpc.Backend that behaves identically to a local one — // no changes needed in core/backend/*.go. // diff --git a/core/services/worker/control_files_test.go b/core/services/worker/control_files_test.go index 42a815c05..0668e2de6 100644 --- a/core/services/worker/control_files_test.go +++ b/core/services/worker/control_files_test.go @@ -73,6 +73,17 @@ var _ = Describe("worker file-staging control routes", func() { DeferCleanup(srv.Close) }) + // The on-disk layout, written out BY HAND rather than derived from the + // helpers under test. What these directories ARE is an operator-facing + // contract: a deployment mounts a volume per directory and sizes it, so a + // join that moved would put staged bytes on a volume nobody provisioned. + // Deriving the expectation from the helper would pin nothing. + DescribeTable("derives its directories from the models path, once each", + func(got, want string) { Expect(got).To(Equal(want)) }, + Entry("cache", (&Config{ModelsPath: "/srv/localai/models"}).stagingCacheDir(), "/srv/localai/cache"), + Entry("data", (&Config{ModelsPath: "/srv/localai/models"}).stagingDataDir(), "/srv/localai/data"), + ) + It("allocates a temp path and returns it", func() { resp := post(workerctl.PathFilesTemp, struct{}{}) Expect(resp.StatusCode).To(Equal(http.StatusOK)) @@ -220,6 +231,25 @@ var _ = Describe("worker file-staging control routes", func() { Expect(reply.LocalPath).To(BeEmpty()) }) + It("reports a temp file it cannot create as a 200 with an error field", func() { + // The staging directory ALREADY EXISTS and is unwritable, so MkdirAll + // succeeds and CreateTemp is the branch that fails. It is a second + // statement of the same rule inside the same handler as the MkdirAll + // spec above, and a spec that only ever reaches MkdirAll leaves this + // one free to answer a non-2xx for the worker's own verdict. + Expect(os.MkdirAll(filepath.Join(cacheDir, "staging-tmp"), 0o500)).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)) @@ -282,11 +312,87 @@ var _ = Describe("worker file-staging control routes", func() { 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)) + // The other direction of the same bound, kept beside the refusals so the + // pair is local rather than spread across the suite. + // + // The two entries hold DIFFERENT halves and both are needed. The + // cap-relative body holds that the boundary is inclusive, so the bound is a + // ceiling and not an off-by-one; it moves with the cap and therefore says + // nothing about the cap's size. The absolute body holds that the cap is + // large enough for real traffic: BackendInstallRequest.BackendGalleries is + // a serialized gallery list of a few hundred kilobytes, so a cap tightened + // below a megabyte would start refusing ordinary requests, and only an + // entry written in absolute bytes notices that. + DescribeTable("serves a body that fits inside the control cap", + func(path string, size int) { + body := append([]byte(`{"key":"`), bytes.Repeat([]byte("a"), size-len(`{"key":""}`))...) + body = append(body, []byte(`"}`)...) + Expect(body).To(HaveLen(size)) + + resp, err := http.Post(srv.URL+path, "application/json", bytes.NewReader(body)) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = resp.Body.Close() }) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + }, + Entry("ensure, exactly at the cap", workerctl.PathFilesEnsure, maxControlRequestBytes), + Entry("stage, exactly at the cap", workerctl.PathFilesStage, maxControlRequestBytes), + Entry("temp, exactly at the cap", workerctl.PathFilesTemp, maxControlRequestBytes), + Entry("listdir, exactly at the cap", workerctl.PathFilesListDir, maxControlRequestBytes), + Entry("ensure, a megabyte of real traffic", workerctl.PathFilesEnsure, 1<<20), + Entry("stage, a megabyte of real traffic", workerctl.PathFilesStage, 1<<20), + Entry("temp, a megabyte of real traffic", workerctl.PathFilesTemp, 1<<20), + Entry("listdir, a megabyte of real traffic", workerctl.PathFilesListDir, 1<<20), + ) + + // A body this worker could not PARSE is the frontend's request being wrong, + // not this worker's verdict about a file. The two live in opposite buckets: + // a non-2xx is mapped onto ErrWorkerUnroutable, which nothing may act on, + // while a 200 with the error field set passes through unwrapped so + // cluster.IsWorkerAnswer sees it and a reap guard MAY act on it. A handler + // that answered `{"error":"invalid request"}` for an unparseable body would + // hand a malformed request to a reap guard as evidence about a file. + // + // Pinned at every verb that decodes a body, because each one writes the + // exit out for itself. The base's NATS handlers shipped the wrong answer + // here; adopting the shared door fixed it, and this is what keeps it fixed. + DescribeTable("reports a request body it cannot read as a rejection, never as a file that is not there", + func(path string) { + resp, err := http.Post(srv.URL+path, "application/json", strings.NewReader("{not json")) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = resp.Body.Close() }) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + + // The status alone is not the whole rule: a 400 whose body still + // decodes as a reply with an error field would be read correctly by + // today's client and wrongly by anything that reads the body first. + var reply struct { + Error string `json:"error"` + } + Expect(json.NewDecoder(resp.Body).Decode(&reply)).NotTo(Succeed()) + }, + Entry("ensure", workerctl.PathFilesEnsure), + Entry("stage", workerctl.PathFilesStage), + Entry("listdir", workerctl.PathFilesListDir), + // temp decodes no body, so it has no such exit to state. + ) + + It("fails a listing the caller abandoned rather than answering a short one", func() { + // A caller that gave up must not be answered with the files walked so + // far. A partial listing is the one shape the frontend cannot tell from + // a directory that really is that size, and it would read as files the + // worker does not have. So the walk returns the context error and the + // verb reports a FAILED listing. + dir := filepath.Join(modelsDir, "abandoned") + Expect(os.MkdirAll(dir, 0o750)).To(Succeed()) + for i := range 32 { + Expect(os.WriteFile(filepath.Join(dir, fmt.Sprintf("f-%02d", i)), []byte("x"), 0o600)).To(Succeed()) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + files, err := listStagedFiles(ctx, dir) + Expect(err).To(MatchError(context.Canceled)) + Expect(files).To(BeEmpty()) }) It("mounts every file verb, so none can be dropped from the set the frontend calls", func() { diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 6ccf1bbc7..a5f220eb1 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -186,7 +186,11 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // the top of Run so the worker fails before registering.) httpAddr := cfg.resolveHTTPAddr() stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging") - dataDir := filepath.Join(cfg.ModelsPath, "..", "data") + // Derived through the same helper the listdir verb resolves `data/` keys + // against, and not a second time here. Two independent joins that agreed + // today would each stay self consistent if one moved, and the symptom would + // be a verb that lists files the file server does not serve. + dataDir := cfg.stagingDataDir() // The readiness gate is created here but only armed once NATS is up, below. // Until then /readyz reports ready, which is correct: reaching this line // means the worker has already registered with the frontend, so it is diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 95febd2e1..8dd3e55d7 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -180,9 +180,11 @@ carries file staging and backend logs, behind the same `LOCALAI_REGISTRATION_TOK 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. +mounted only when the deployment configured an object store). The request bodies +are unchanged and the reply fields keep their names and types, so nothing an +operator inspects on the wire has a new shape. The one difference is that a +worker now OMITS an empty reply field where the NATS handlers always emitted it, +which a client reading a missing field as the zero value cannot tell apart. 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 diff --git a/tests/e2e/distributed/nats_jwt_test.go b/tests/e2e/distributed/nats_jwt_test.go index bf947e472..40008c824 100644 --- a/tests/e2e/distributed/nats_jwt_test.go +++ b/tests/e2e/distributed/nats_jwt_test.go @@ -17,14 +17,32 @@ var _ = Describe("NATS JWT Auth", Label("Distributed", "NatsJWT"), func() { infra = SetupJWTInfra() }) - It("connects with a minted backend worker JWT and publishes on allowed subjects", func() { - // Backend workers may publish under nodes..files.> (see pkg/natsauth permissions). - subject := nodeSubjectPrefix(infra.NodeID) + ".files.in" - Expect(infra.NC.Publish(subject, map[string]string{"path": "/tmp/model"})).To(Succeed()) + It("connects with a minted backend worker JWT and publishes on its one remaining allowed subject", func() { + // A backend worker's publish grant is `_INBOX.>` and nothing else now. + // Every verb a frontend gives it, file staging included, is an HTTP + // route on its tunnel, so the `nodes..files.>` grant went with the + // subjects. See pkg/natsauth.WorkerPermissions. + Expect(infra.NC.Publish("_INBOX.probe", map[string]string{"path": "/tmp/model"})).To(Succeed()) Expect(infra.NC.Conn().FlushTimeout(2 * time.Second)).To(Succeed()) + Expect(infra.NC.Conn().LastError()).ToNot(HaveOccurred()) Expect(infra.NC.Conn().IsConnected()).To(BeTrue()) }) + It("denies a backend worker the file-staging subjects it no longer serves", func() { + // This spec used to assert the OPPOSITE, and kept passing after the + // grant was deleted. A NATS permission violation does not close the + // connection, so a spec that checks only FlushTimeout and IsConnected + // cannot tell an allowed publish from a denied one; LastError is what + // actually reads the server's verdict, which is why the sibling below + // has always used it. + subject := nodeSubjectPrefix(infra.NodeID) + ".files.stage" + Expect(infra.NC.Publish(subject, map[string]string{"path": "/tmp/model"})).To(Succeed()) + Eventually(func() error { + _ = infra.NC.Conn().FlushTimeout(500 * time.Millisecond) + return infra.NC.Conn().LastError() + }, "3s", "50ms").Should(HaveOccurred()) + }) + It("allows backend subscribe on the node prefix", func() { wild := nodeSubjectPrefix(infra.NodeID) + ".>" sub, err := infra.NC.Subscribe(wild, func(_ []byte) {})