diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index d6c14cf2e..57c5bad9b 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -19,6 +19,8 @@ import ( "google.golang.org/protobuf/proto" ) +const stagedInputReleaseTimeout = 30 * time.Second + // FileStagingClient wraps a grpc.Backend to transparently handle file transfer // for distributed mode. Input files are staged on the backend node before the // gRPC call. Output files are retrieved from the backend after the call. @@ -49,21 +51,64 @@ func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string) // requestID generates a unique ID for ephemeral file keys. func requestID() string { - return uuid.New().String()[:8] + return uuid.NewString() +} + +type stagedInputLifecycle struct { + client *FileStagingClient + requestID string + keys []string + seen map[string]struct{} +} + +func (f *FileStagingClient) newStagedInputLifecycle() *stagedInputLifecycle { + return &stagedInputLifecycle{ + client: f, + requestID: requestID(), + keys: []string{}, + seen: map[string]struct{}{}, + } +} + +func (l *stagedInputLifecycle) track(key string) { + if _, ok := l.seen[key]; ok { + return + } + l.seen[key] = struct{}{} + l.keys = append(l.keys, key) +} + +func (l *stagedInputLifecycle) release() { + if len(l.keys) == 0 { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), stagedInputReleaseTimeout) + defer cancel() + for _, key := range l.keys { + if err := l.client.stager.ReleaseRemote(ctx, l.client.nodeID, key); err != nil { + xlog.Warn("Failed to release staged input", "node", l.client.nodeID, "key", key, "error", err) + } + } } // stageInputFile uploads a local file to the remote node via the FileStager. -// Returns the remote-local path and the ephemeral key. -func (f *FileStagingClient) stageInputFile(ctx context.Context, reqID, localPath, category string) (string, string, error) { +func (f *FileStagingClient) stageInputFile( + ctx context.Context, + lifecycle *stagedInputLifecycle, + localPath, + category string, +) (string, error) { basename := filepath.Base(localPath) - key := storage.EphemeralKey(reqID, category, basename) + key := storage.EphemeralKey(lifecycle.requestID, category, basename) + lifecycle.track(key) remotePath, err := f.stager.EnsureRemote(ctx, f.nodeID, localPath, key) if err != nil { - return "", "", fmt.Errorf("staging input file: %w", err) + return "", fmt.Errorf("staging input file: %w", err) } - return remotePath, key, nil + return remotePath, nil } // retrieveOutputFile retrieves an output file from the backend to a local path. @@ -101,23 +146,29 @@ func (f *FileStagingClient) translateModelPath(frontendPath string) string { } func (f *FileStagingClient) Predict(ctx context.Context, in *pb.PredictOptions, opts ...ggrpc.CallOption) (*pb.Reply, error) { - reqID := requestID() - in, _ = f.stageMultimodalInputs(ctx, reqID, in) + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.PredictOptions) + in = f.stageMultimodalInputs(ctx, lifecycle, in) return f.Backend.Predict(ctx, in, opts...) } func (f *FileStagingClient) PredictStream(ctx context.Context, in *pb.PredictOptions, fn func(reply *pb.Reply), opts ...ggrpc.CallOption) error { - reqID := requestID() - in, _ = f.stageMultimodalInputs(ctx, reqID, in) + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.PredictOptions) + in = f.stageMultimodalInputs(ctx, lifecycle, in) return f.Backend.PredictStream(ctx, in, fn, opts...) } func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.GenerateImageRequest) // Stage input source image if present if in.Src != "" && isFilePath(in.Src) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs") if err != nil { return nil, fmt.Errorf("staging image src: %w", err) } @@ -127,7 +178,7 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm // Stage reference images for i, img := range in.RefImages { if isFilePath(img) { - backendPath, _, err := f.stageInputFile(ctx, reqID, img, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, img, "inputs") if err != nil { return nil, fmt.Errorf("staging ref image: %w", err) } @@ -161,25 +212,27 @@ func (f *FileStagingClient) GenerateImage(ctx context.Context, in *pb.GenerateIm } func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.GenerateVideoRequest) // Stage start/end images and optional audio conditioning. if in.StartImage != "" && isFilePath(in.StartImage) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.StartImage, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.StartImage, "inputs") if err != nil { return nil, fmt.Errorf("staging start image: %w", err) } in.StartImage = backendPath } if in.EndImage != "" && isFilePath(in.EndImage) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.EndImage, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.EndImage, "inputs") if err != nil { return nil, fmt.Errorf("staging end image: %w", err) } in.EndImage = backendPath } if in.Audio != "" && isFilePath(in.Audio) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Audio, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Audio, "inputs") if err != nil { return nil, fmt.Errorf("staging video audio: %w", err) } @@ -211,11 +264,13 @@ func (f *FileStagingClient) GenerateVideo(ctx context.Context, in *pb.GenerateVi } func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.Generate3DRequest) // Stage the conditioning image or existing GLB used by 3D post-processing. if in.Src != "" && isFilePath(in.Src) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Src, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs") if err != nil { return nil, fmt.Errorf("staging 3D input asset: %w", err) } @@ -247,7 +302,9 @@ func (f *FileStagingClient) Generate3D(ctx context.Context, in *pb.Generate3DReq } func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.TTSRequest) // Translate model path from frontend to remote worker path. // The model and its companion files (e.g. .onnx.json) were already staged @@ -258,7 +315,7 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ... // Voice may be a named backend speaker or a request-scoped reference WAV. // Only path-shaped values are staged; speaker IDs pass through unchanged. if in.Voice != "" && isFilePath(in.Voice) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs") if err != nil { return nil, fmt.Errorf("staging TTS voice reference: %w", err) } @@ -290,14 +347,16 @@ func (f *FileStagingClient) TTS(ctx context.Context, in *pb.TTSRequest, opts ... } func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn func(*pb.Reply), opts ...ggrpc.CallOption) error { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.TTSRequest) // Translate model path from frontend to remote worker path (same as TTS above) if in.Model != "" && isFilePath(in.Model) { in.Model = f.translateModelPath(in.Model) } if in.Voice != "" && isFilePath(in.Voice) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Voice, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Voice, "inputs") if err != nil { return fmt.Errorf("staging streaming TTS voice reference: %w", err) } @@ -308,11 +367,13 @@ func (f *FileStagingClient) TTSStream(ctx context.Context, in *pb.TTSRequest, fn } func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest, opts ...ggrpc.CallOption) (*pb.Result, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.SoundGenerationRequest) // Stage input source if in.Src != nil && *in.Src != "" && isFilePath(*in.Src) { - backendPath, _, err := f.stageInputFile(ctx, reqID, *in.Src, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, *in.Src, "inputs") if err != nil { return nil, fmt.Errorf("staging sound src: %w", err) } @@ -344,24 +405,27 @@ func (f *FileStagingClient) SoundGeneration(ctx context.Context, in *pb.SoundGen } func (f *FileStagingClient) SoundDetection(ctx context.Context, in *pb.SoundDetectionRequest, opts ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) { + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.SoundDetectionRequest) if in.Src != "" && isFilePath(in.Src) { - backendPath, _, err := f.stageInputFile(ctx, requestID(), in.Src, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Src, "inputs") if err != nil { return nil, fmt.Errorf("staging audio for sound detection: %w", err) } - // Keep the frontend path available if the caller retries on another node. - in = proto.Clone(in).(*pb.SoundDetectionRequest) in.Src = backendPath } return f.Backend.SoundDetection(ctx, in, opts...) } func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.TranscriptRequest, opts ...ggrpc.CallOption) (*pb.TranscriptResult, error) { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.TranscriptRequest) // Stage input audio file if in.Dst != "" && isFilePath(in.Dst) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs") if err != nil { return nil, fmt.Errorf("staging audio for transcription: %w", err) } @@ -372,11 +436,13 @@ func (f *FileStagingClient) AudioTranscription(ctx context.Context, in *pb.Trans } func (f *FileStagingClient) AudioTranscriptionStream(ctx context.Context, in *pb.TranscriptRequest, fn func(chunk *pb.TranscriptStreamResponse), opts ...ggrpc.CallOption) error { - reqID := requestID() + lifecycle := f.newStagedInputLifecycle() + defer lifecycle.release() + in = proto.Clone(in).(*pb.TranscriptRequest) // Stage input audio file if in.Dst != "" && isFilePath(in.Dst) { - backendPath, _, err := f.stageInputFile(ctx, reqID, in.Dst, "inputs") + backendPath, err := f.stageInputFile(ctx, lifecycle, in.Dst, "inputs") if err != nil { return fmt.Errorf("staging audio for transcription stream: %w", err) } @@ -458,26 +524,33 @@ func (f *FileStagingClient) QuantizationProgress(ctx context.Context, in *pb.Qua // stageMultimodalInputs stages Images, Videos, Audios fields in PredictOptions // if they are file paths (not base64 or URLs). -func (f *FileStagingClient) stageMultimodalInputs(ctx context.Context, reqID string, in *pb.PredictOptions) (*pb.PredictOptions, []string) { - var keys []string - in.Images = f.stagePathSlice(ctx, reqID, in.Images, "inputs", &keys) - in.Videos = f.stagePathSlice(ctx, reqID, in.Videos, "inputs", &keys) - in.Audios = f.stagePathSlice(ctx, reqID, in.Audios, "inputs", &keys) - return in, keys +func (f *FileStagingClient) stageMultimodalInputs( + ctx context.Context, + lifecycle *stagedInputLifecycle, + in *pb.PredictOptions, +) *pb.PredictOptions { + in.Images = f.stagePathSlice(ctx, lifecycle, in.Images, "inputs") + in.Videos = f.stagePathSlice(ctx, lifecycle, in.Videos, "inputs") + in.Audios = f.stagePathSlice(ctx, lifecycle, in.Audios, "inputs") + return in } -func (f *FileStagingClient) stagePathSlice(ctx context.Context, reqID string, paths []string, category string, keys *[]string) []string { +func (f *FileStagingClient) stagePathSlice( + ctx context.Context, + lifecycle *stagedInputLifecycle, + paths []string, + category string, +) []string { result := make([]string, len(paths)) for i, p := range paths { if isFilePath(p) { - backendPath, key, err := f.stageInputFile(ctx, reqID, p, category) + backendPath, err := f.stageInputFile(ctx, lifecycle, p, category) if err != nil { xlog.Warn("Failed to stage multimodal file, passing through", "path", p, "error", err) result[i] = p continue } result[i] = backendPath - *keys = append(*keys, key) } else { result[i] = p } diff --git a/core/services/nodes/file_staging_lifecycle_test.go b/core/services/nodes/file_staging_lifecycle_test.go new file mode 100644 index 000000000..25892e998 --- /dev/null +++ b/core/services/nodes/file_staging_lifecycle_test.go @@ -0,0 +1,288 @@ +package nodes + +import ( + "context" + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" + "google.golang.org/protobuf/proto" +) + +const fullUUIDPattern = `[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}` + +type lifecycleStager struct { + fakeFileStager + ensureErr error + releaseErr error + releasedKeys []string + releaseCtxErr []error + releaseHasDeadline []bool + releaseDeadlines []time.Time +} + +func (s *lifecycleStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { + s.fakeFileStager.EnsureRemote(ctx, nodeID, localPath, key) + if s.ensureErr != nil { + return "", s.ensureErr + } + return "/remote/" + key, nil +} + +func (s *lifecycleStager) ReleaseRemote(ctx context.Context, _ string, key string) error { + s.releasedKeys = append(s.releasedKeys, key) + s.releaseCtxErr = append(s.releaseCtxErr, ctx.Err()) + deadline, ok := ctx.Deadline() + s.releaseHasDeadline = append(s.releaseHasDeadline, ok) + s.releaseDeadlines = append(s.releaseDeadlines, deadline) + return s.releaseErr +} + +type lifecycleBackend struct { + grpc.Backend + predictResult *pb.Reply + predictErr error + streamBlock <-chan struct{} + streamStarted chan<- struct{} +} + +func (b *lifecycleBackend) Predict(_ context.Context, _ *pb.PredictOptions, _ ...ggrpc.CallOption) (*pb.Reply, error) { + if b.predictResult == nil { + b.predictResult = &pb.Reply{} + } + return b.predictResult, b.predictErr +} + +func (b *lifecycleBackend) PredictStream(_ context.Context, _ *pb.PredictOptions, _ func(*pb.Reply), _ ...ggrpc.CallOption) error { + if b.streamStarted != nil { + b.streamStarted <- struct{}{} + } + if b.streamBlock != nil { + <-b.streamBlock + } + return nil +} + +func (b *lifecycleBackend) GenerateImage(_ context.Context, _ *pb.GenerateImageRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +func (b *lifecycleBackend) GenerateVideo(_ context.Context, _ *pb.GenerateVideoRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +func (b *lifecycleBackend) Generate3D(_ context.Context, _ *pb.Generate3DRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +func (b *lifecycleBackend) TTS(_ context.Context, _ *pb.TTSRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +func (b *lifecycleBackend) TTSStream(_ context.Context, _ *pb.TTSRequest, _ func(*pb.Reply), _ ...ggrpc.CallOption) error { + return nil +} + +func (b *lifecycleBackend) SoundGeneration(_ context.Context, _ *pb.SoundGenerationRequest, _ ...ggrpc.CallOption) (*pb.Result, error) { + return &pb.Result{Success: true}, nil +} + +func (b *lifecycleBackend) SoundDetection(_ context.Context, _ *pb.SoundDetectionRequest, _ ...ggrpc.CallOption) (*pb.SoundDetectionResponse, error) { + return &pb.SoundDetectionResponse{}, nil +} + +func (b *lifecycleBackend) AudioTranscription(_ context.Context, _ *pb.TranscriptRequest, _ ...ggrpc.CallOption) (*pb.TranscriptResult, error) { + return &pb.TranscriptResult{}, nil +} + +func (b *lifecycleBackend) AudioTranscriptionStream(_ context.Context, _ *pb.TranscriptRequest, _ func(*pb.TranscriptStreamResponse), _ ...ggrpc.CallOption) error { + return nil +} + +var _ = Describe("FileStagingClient request lifecycle", func() { + It("uses a full UUID for ephemeral request keys", func() { + Expect(requestID()).To(MatchRegexp(`^` + fullUUIDPattern + `$`)) + }) + + It("releases every staged key and preserves caller requests", func(ctx SpecContext) { + tests := []struct { + name string + keyCount int + invoke func(*FileStagingClient) proto.Message + }{ + {name: "predict", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}, Videos: []string{"/tmp/video.mp4"}, Audios: []string{"/tmp/audio.wav"}} + original := proto.Clone(request) + _, err := client.Predict(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "predict stream", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.PredictOptions{Images: []string{"/tmp/image.png"}} + original := proto.Clone(request) + Expect(client.PredictStream(ctx, request, func(*pb.Reply) {})).To(Succeed()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "image generation", keyCount: 2, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.GenerateImageRequest{Src: "/tmp/source.png", RefImages: []string{"/tmp/reference.png"}} + original := proto.Clone(request) + _, err := client.GenerateImage(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "video generation", keyCount: 3, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.GenerateVideoRequest{StartImage: "/tmp/start.png", EndImage: "/tmp/end.png", Audio: "/tmp/audio.wav"} + original := proto.Clone(request) + _, err := client.GenerateVideo(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "3D generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.Generate3DRequest{Src: "/tmp/source.glb"} + original := proto.Clone(request) + _, err := client.Generate3D(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.TTSRequest{Voice: "/tmp/voice.wav"} + original := proto.Clone(request) + _, err := client.TTS(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "streaming TTS", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.TTSRequest{Voice: "/tmp/voice.wav"} + original := proto.Clone(request) + Expect(client.TTSStream(ctx, request, func(*pb.Reply) {})).To(Succeed()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "sound generation", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + source := "/tmp/source.wav" + request := &pb.SoundGenerationRequest{Src: &source} + original := proto.Clone(request) + _, err := client.SoundGeneration(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "sound detection", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.SoundDetectionRequest{Src: "/tmp/source.wav"} + original := proto.Clone(request) + _, err := client.SoundDetection(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"} + original := proto.Clone(request) + _, err := client.AudioTranscription(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + {name: "streaming transcription", keyCount: 1, invoke: func(client *FileStagingClient) proto.Message { + request := &pb.TranscriptRequest{Dst: "/tmp/source.wav"} + original := proto.Clone(request) + Expect(client.AudioTranscriptionStream(ctx, request, func(*pb.TranscriptStreamResponse) {})).To(Succeed()) + Expect(proto.Equal(request, original)).To(BeTrue()) + return request + }}, + } + + for _, test := range tests { + By(test.name) + stager := &lifecycleStager{} + client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1") + test.invoke(client) + Expect(stager.ensureCalls).To(HaveLen(test.keyCount)) + Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls))) + Expect(stager.releaseCtxErr).To(ConsistOf(make([]error, test.keyCount))) + Expect(stager.releaseHasDeadline).To(HaveLen(test.keyCount)) + for _, hasDeadline := range stager.releaseHasDeadline { + Expect(hasDeadline).To(BeTrue()) + } + for _, deadline := range stager.releaseDeadlines { + Expect(time.Until(deadline)).To(BeNumerically(">", 0)) + Expect(time.Until(deadline)).To(BeNumerically("<=", time.Minute)) + } + } + }) + + It("tracks a key before staging so a partial upload failure is released", func(ctx SpecContext) { + uploadErr := errors.New("upload failed") + stager := &lifecycleStager{ensureErr: uploadErr, releaseErr: errors.New("release failed")} + client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1") + + _, err := client.GenerateImage(ctx, &pb.GenerateImageRequest{Src: "/tmp/source.png"}) + + Expect(err).To(MatchError(ContainSubstring("upload failed"))) + Expect(stager.ensureCalls).To(HaveLen(1)) + Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls))) + }) + + It("uses an active bounded cleanup context after caller cancellation", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + stager := &lifecycleStager{} + client := NewFileStagingClient(&lifecycleBackend{}, stager, "worker-1") + + _, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}) + + Expect(err).NotTo(HaveOccurred()) + Expect(stager.releaseCtxErr).To(Equal([]error{nil})) + }) + + It("does not release streaming inputs before the backend completes", func(ctx SpecContext) { + block := make(chan struct{}) + started := make(chan struct{}, 1) + backend := &lifecycleBackend{streamBlock: block, streamStarted: started} + stager := &lifecycleStager{} + client := NewFileStagingClient(backend, stager, "worker-1") + done := make(chan error, 1) + + go func() { + done <- client.PredictStream(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}, func(*pb.Reply) {}) + }() + + Eventually(started).Should(Receive()) + Expect(stager.ensureCalls).To(HaveLen(1)) + Expect(stager.releasedKeys).To(BeEmpty()) + close(block) + Eventually(done).Should(Receive(Succeed())) + Expect(stager.releasedKeys).To(Equal(keysFromEnsureCalls(stager.ensureCalls))) + }) + + It("does not replace a backend result when cleanup fails", func(ctx SpecContext) { + reply := &pb.Reply{Message: []byte("ok")} + backend := &lifecycleBackend{predictResult: reply} + stager := &lifecycleStager{releaseErr: errors.New("release failed")} + client := NewFileStagingClient(backend, stager, "worker-1") + + result, err := client.Predict(ctx, &pb.PredictOptions{Images: []string{"/tmp/image.png"}}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(BeIdenticalTo(reply)) + }) +}) + +func keysFromEnsureCalls(calls []ensureCall) []string { + keys := make([]string, len(calls)) + for i, call := range calls { + keys[i] = call.key + } + return keys +} diff --git a/core/services/nodes/file_staging_sound_detection_test.go b/core/services/nodes/file_staging_sound_detection_test.go index e66b18f5f..2fc6fb7e3 100644 --- a/core/services/nodes/file_staging_sound_detection_test.go +++ b/core/services/nodes/file_staging_sound_detection_test.go @@ -28,6 +28,10 @@ func (s *soundStagingFailure) EnsureRemote(context.Context, string, string, stri return "", errors.New("upload failed") } +func (s *soundStagingFailure) ReleaseRemote(context.Context, string, string) error { + return nil +} + type soundRouteFactory struct{ client grpc.Backend } func (f *soundRouteFactory) NewClient(string, bool) grpc.Backend { return f.client } diff --git a/core/services/nodes/file_staging_tts_test.go b/core/services/nodes/file_staging_tts_test.go index a4e4f57f4..f9ff2a5d7 100644 --- a/core/services/nodes/file_staging_tts_test.go +++ b/core/services/nodes/file_staging_tts_test.go @@ -39,7 +39,8 @@ var _ = Describe("FileStagingClient TTS references", func() { Expect(stager.ensureCalls).To(HaveLen(1)) Expect(stager.ensureCalls[0].localPath).To(Equal("/data/voice-profiles/profile/reference.wav")) Expect(backend.ttsRequest.Voice).To(HavePrefix("/remote/ephemeral/")) - Expect(backend.ttsRequest.Voice).To(MatchRegexp(`/inputs/[0-9a-f]{8}/reference\.wav$`)) + voicePathPattern := `/inputs/` + fullUUIDPattern + `/reference\.wav$` + Expect(backend.ttsRequest.Voice).To(MatchRegexp(voicePathPattern)) }) It("stages a reference WAV before streaming synthesis", func(ctx SpecContext) {