diff --git a/core/application/application.go b/core/application/application.go index 7633c8f81..fd2f1e404 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -25,6 +25,7 @@ import ( "github.com/mudler/LocalAI/core/services/voiceprofile" "github.com/mudler/LocalAI/core/services/voicerecognition" "github.com/mudler/LocalAI/core/templates" + "github.com/mudler/LocalAI/core/trace" pkggrpc "github.com/mudler/LocalAI/pkg/grpc" localaitools "github.com/mudler/LocalAI/pkg/mcp/localaitools" localaiInproc "github.com/mudler/LocalAI/pkg/mcp/localaitools/inproc" @@ -119,6 +120,8 @@ func (a *Application) Ready() bool { return a.startupComplete.Load() } func (a *Application) markStartupComplete() { a.startupComplete.Store(true) } func newApplication(appConfig *config.ApplicationConfig) *Application { + corebackend.ConfigureGlobalBackendAdmission(appConfig.MaxConcurrentBackendRequests) + trace.ConfigureBackendTraceMaxInFlight(appConfig.MaxConcurrentBackendRequests) ml := model.NewModelLoader(appConfig.SystemState) // Apply the per-model load-failure cooldown (0 disables). Set here rather @@ -134,7 +137,7 @@ func newApplication(appConfig *config.ApplicationConfig) *Application { // Record a model_load backend trace for every real backend load, so the // Traces UI shows which backend runtime served each model and how long // the load took. Load failures are traced by the modality wrappers. - ml.SetLoadObserver(corebackend.ModelLoadTraceObserver(appConfig)) + ml.SetLoadLifecycleObserver(corebackend.ModelLoadTraceObserver(appConfig)) app := &Application{ backendLoader: config.NewModelConfigLoader( diff --git a/core/backend/audio_transform.go b/core/backend/audio_transform.go index 99c46a4ff..81dcace04 100644 --- a/core/backend/audio_transform.go +++ b/core/backend/audio_transform.go @@ -91,12 +91,20 @@ func ModelAudioTransform( return AudioTransformOutputs{}, nil, fmt.Errorf("persist reference: %w", err) } } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return AudioTransformOutputs{}, nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceAudioTransform, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(filepath.Base(audioPath), 200)}) } + defer trace.CancelBackendTrace(traceID) res, err := transformModel.AudioTransform(ctx, &proto.AudioTransformRequest{ ModelIdentity: modelConfig.Model, @@ -126,6 +134,7 @@ func ModelAudioTransform( } } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceAudioTransform, @@ -198,7 +207,17 @@ func ModelAudioTransformStream( if transformModel == nil { return nil, fmt.Errorf("could not load audio-transform model %q", modelConfig.Model) } - return transformModel.AudioTransformStream(ctx) + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + stream, err := transformModel.AudioTransformStream(ctx) + if err != nil { + release() + return nil, err + } + stream.AddCleanup(release) + return stream, nil } // persistAudioInput copies a transient input file (typically a multipart diff --git a/core/backend/depth.go b/core/backend/depth.go index 154cf3335..df32b46de 100644 --- a/core/backend/depth.go +++ b/core/backend/depth.go @@ -33,12 +33,20 @@ func Depth( if depthModel == nil { return nil, fmt.Errorf("could not load depth model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceDepth, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(in.GetSrc(), 200)}) } + defer trace.CancelBackendTrace(traceID) // Stamped here for the same reason as in rerank.go: the caller builds the // request without a ModelConfig, this function has the one that loaded. @@ -53,6 +61,7 @@ func Depth( } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceDepth, diff --git a/core/backend/detection.go b/core/backend/detection.go index 4d0ddae33..3b14f4c73 100644 --- a/core/backend/detection.go +++ b/core/backend/detection.go @@ -33,11 +33,19 @@ func Detection( return nil, fmt.Errorf("could not load detection model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceDetection, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(sourceFile, 200)}) } + defer trace.CancelBackendTrace(traceID) res, err := detectionModel.Detect(ctx, &proto.DetectOptions{ ModelIdentity: modelConfig.Model, @@ -55,6 +63,7 @@ func Detection( } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceDetection, diff --git a/core/backend/detokenize.go b/core/backend/detokenize.go index 05c16c4a5..c26df9a19 100644 --- a/core/backend/detokenize.go +++ b/core/backend/detokenize.go @@ -23,11 +23,19 @@ func ModelDetokenize(tokens []int32, loader *model.ModelLoader, modelConfig conf return schema.DetokenizeResponse{}, err } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return schema.DetokenizeResponse{}, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTokenize, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "detokenize"}) } + defer trace.CancelBackendTrace(traceID) resp, err := inferenceModel.Detokenize(appConfig.Context, &pb.DetokenizeRequest{Tokens: tokens}) @@ -43,6 +51,7 @@ func ModelDetokenize(tokens []int32, loader *model.ModelLoader, modelConfig conf } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTokenize, diff --git a/core/backend/embeddings.go b/core/backend/embeddings.go index eff88ef04..d68a334fb 100644 --- a/core/backend/embeddings.go +++ b/core/backend/embeddings.go @@ -109,9 +109,15 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M traceData["input_tokens_count"] = len(tokens) } - startTime := time.Now() + summary := trace.TruncateString(s, 200) + if summary == "" { + summary = fmt.Sprintf("tokens[%d]", len(tokens)) + } originalFn := wrappedFn wrappedFn = func() ([]float32, error) { + startTime := time.Now() + traceID := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceEmbedding, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: summary}) + defer trace.CancelBackendTrace(traceID) result, err := originalFn() duration := time.Since(startTime) @@ -122,12 +128,8 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M errStr = err.Error() } - summary := trace.TruncateString(s, 200) - if summary == "" { - summary = fmt.Sprintf("tokens[%d]", len(tokens)) - } - trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: duration, Type: trace.BackendTraceEmbedding, @@ -141,6 +143,15 @@ func ModelEmbedding(ctx context.Context, s string, tokens []int, loader *model.M return result, err } } + originalFn := wrappedFn + wrappedFn = func() ([]float32, error) { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() + return originalFn() + } return wrappedFn, nil } diff --git a/core/backend/face_analyze.go b/core/backend/face_analyze.go index c729506fe..0d734361d 100644 --- a/core/backend/face_analyze.go +++ b/core/backend/face_analyze.go @@ -30,11 +30,19 @@ func FaceAnalyze( return nil, fmt.Errorf("could not load face recognition model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceFaceAnalyze, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "face analysis"}) } + defer trace.CancelBackendTrace(traceID) res, err := faceModel.FaceAnalyze(ctx, &proto.FaceAnalyzeRequest{ ModelIdentity: modelConfig.Model, @@ -49,6 +57,7 @@ func FaceAnalyze( errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceFaceAnalyze, diff --git a/core/backend/face_embed.go b/core/backend/face_embed.go index dc9fecad0..6062c96ea 100644 --- a/core/backend/face_embed.go +++ b/core/backend/face_embed.go @@ -32,6 +32,11 @@ func FaceEmbed( predictOpts := gRPCPredictOpts(modelConfig, loader.ModelPath) predictOpts.Images = []string{imgBase64} + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() res, err := faceModel.Embeddings(ctx, predictOpts) if err != nil { diff --git a/core/backend/face_verify.go b/core/backend/face_verify.go index 0b51b6942..e5d3bb400 100644 --- a/core/backend/face_verify.go +++ b/core/backend/face_verify.go @@ -30,11 +30,19 @@ func FaceVerify( return nil, fmt.Errorf("could not load face recognition model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceFaceVerify, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "face verification"}) } + defer trace.CancelBackendTrace(traceID) res, err := faceModel.FaceVerify(ctx, &proto.FaceVerifyRequest{ ModelIdentity: modelConfig.Model, @@ -50,6 +58,7 @@ func FaceVerify( errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceFaceVerify, diff --git a/core/backend/global_admission.go b/core/backend/global_admission.go new file mode 100644 index 000000000..9b216ea83 --- /dev/null +++ b/core/backend/global_admission.go @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: MIT + +package backend + +import ( + "fmt" + "sync" + "time" + + "github.com/mudler/LocalAI/core/config" +) + +// BackendAdmissionError reports that the process-wide backend execution +// ceiling is full. HTTP callers map it to 503; internal callers receive the +// same typed error instead of silently queueing and growing in-flight state. +type BackendAdmissionError struct { + Limit int + RetryAfter time.Duration +} + +func (e *BackendAdmissionError) Error() string { + return fmt.Sprintf("backend inference capacity reached (max_concurrent=%d); retry after %s", e.Limit, e.RetryAfter) +} + +var backendAdmission = struct { + sync.RWMutex + limit int + slots chan struct{} +}{} + +// ConfigureGlobalBackendAdmission sets the process-wide ceiling. It is called +// during application construction, before backend work can begin. +func ConfigureGlobalBackendAdmission(limit int) { + if limit <= 0 { + limit = config.DefaultMaxConcurrentBackendRequests + } + backendAdmission.Lock() + backendAdmission.limit = limit + backendAdmission.slots = make(chan struct{}, limit) + backendAdmission.Unlock() +} + +// AcquireGlobalBackendSlot admits one backend operation without queueing. +// Callers must invoke release on every completion path. +func AcquireGlobalBackendSlot() (release func(), err error) { + backendAdmission.RLock() + limit, slots := backendAdmission.limit, backendAdmission.slots + backendAdmission.RUnlock() + if slots == nil { + backendAdmission.Lock() + if backendAdmission.slots == nil { + backendAdmission.limit = config.DefaultMaxConcurrentBackendRequests + backendAdmission.slots = make(chan struct{}, backendAdmission.limit) + } + limit, slots = backendAdmission.limit, backendAdmission.slots + backendAdmission.Unlock() + } + select { + case slots <- struct{}{}: + var once sync.Once + return func() { once.Do(func() { <-slots }) }, nil + default: + return nil, &BackendAdmissionError{Limit: limit, RetryAfter: time.Second} + } +} + +// GlobalBackendInFlight is the current number of admitted backend operations. +func GlobalBackendInFlight() int { + backendAdmission.RLock() + defer backendAdmission.RUnlock() + return len(backendAdmission.slots) +} diff --git a/core/backend/global_admission_test.go b/core/backend/global_admission_test.go new file mode 100644 index 000000000..d735e9e1e --- /dev/null +++ b/core/backend/global_admission_test.go @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT + +package backend_test + +import ( + "errors" + + "github.com/mudler/LocalAI/core/backend" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("global backend admission", func() { + BeforeEach(func() { + backend.ConfigureGlobalBackendAdmission(1) + }) + + It("rejects excess backend work without queueing", func() { + release, err := backend.AcquireGlobalBackendSlot() + Expect(err).NotTo(HaveOccurred()) + Expect(backend.GlobalBackendInFlight()).To(Equal(1)) + + _, err = backend.AcquireGlobalBackendSlot() + var capacityErr *backend.BackendAdmissionError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.Limit).To(Equal(1)) + + release() + Expect(backend.GlobalBackendInFlight()).To(BeZero()) + }) + + It("makes release idempotent", func() { + release, err := backend.AcquireGlobalBackendSlot() + Expect(err).NotTo(HaveOccurred()) + release() + release() + Expect(backend.GlobalBackendInFlight()).To(BeZero()) + }) +}) diff --git a/core/backend/image.go b/core/backend/image.go index 95f4dc122..75c2a9faa 100644 --- a/core/backend/image.go +++ b/core/backend/image.go @@ -63,9 +63,11 @@ func ImageGeneration(ctx context.Context, height, width, step, seed int, positiv "destination": dst, } - startTime := time.Now() originalFn := fn fn = func() error { + startTime := time.Now() + traceID := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceImageGeneration, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(positive_prompt, 200)}) + defer trace.CancelBackendTrace(traceID) err := originalFn() duration := time.Since(startTime) @@ -75,6 +77,7 @@ func ImageGeneration(ctx context.Context, height, width, step, seed int, positiv } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: duration, Type: trace.BackendTraceImageGeneration, @@ -88,6 +91,15 @@ func ImageGeneration(ctx context.Context, height, width, step, seed int, positiv return err } } + originalFn := fn + fn = func() error { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + return originalFn() + } return fn, nil } diff --git a/core/backend/llm.go b/core/backend/llm.go index caf03f74d..2abd4618e 100644 --- a/core/backend/llm.go +++ b/core/backend/llm.go @@ -378,9 +378,17 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima "xml_format_preset": c.FunctionsConfig.XMLFormatPreset, } - startTime := time.Now() originalFn := fn fn = func() (LLMResponse, error) { + startTime := time.Now() + traceID := trace.BeginBackendTrace(trace.BackendTrace{ + Timestamp: startTime, + Type: trace.BackendTraceLLM, + ModelName: c.Name, + Backend: c.Backend, + Summary: trace.GenerateLLMSummary(messages, s), + }) + defer trace.CancelBackendTrace(traceID) resp, err := originalFn() duration := time.Since(startTime) @@ -432,6 +440,7 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: duration, Type: trace.BackendTraceLLM, @@ -445,6 +454,15 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima return resp, err } } + originalFn := fn + fn = func() (LLMResponse, error) { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return LLMResponse{}, err + } + defer release() + return originalFn() + } return fn, nil } diff --git a/core/backend/model3d.go b/core/backend/model3d.go index 1952501d9..7cccbc7db 100644 --- a/core/backend/model3d.go +++ b/core/backend/model3d.go @@ -74,9 +74,11 @@ func Model3DGeneration(options Model3DGenerationOptions, loader *model.ModelLoad } } - startTime := time.Now() originalFn := fn fn = func() error { + startTime := time.Now() + traceID := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: traceType, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(traceSummary, 200)}) + defer trace.CancelBackendTrace(traceID) err := originalFn() duration := time.Since(startTime) @@ -86,6 +88,7 @@ func Model3DGeneration(options Model3DGenerationOptions, loader *model.ModelLoad } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: duration, Type: traceType, @@ -99,6 +102,15 @@ func Model3DGeneration(options Model3DGenerationOptions, loader *model.ModelLoad return err } } + originalFn := fn + fn = func() error { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + return originalFn() + } return fn, nil } diff --git a/core/backend/model_load_trace_test.go b/core/backend/model_load_trace_test.go index 1cce5da26..273963f0f 100644 --- a/core/backend/model_load_trace_test.go +++ b/core/backend/model_load_trace_test.go @@ -30,6 +30,7 @@ var _ = Describe("ModelLoadTraceObserver", func() { } BeforeEach(func() { + backend.ConfigureGlobalBackendAdmission(64) appConfig = &config.ApplicationConfig{ EnableTracing: true, TracingMaxItems: 64, @@ -39,7 +40,11 @@ var _ = Describe("ModelLoadTraceObserver", func() { }) It("records a model_load trace with the backend runtime on success", func() { - backend.ModelLoadTraceObserver(appConfig)(successEvent) + finish, err := backend.ModelLoadTraceObserver(appConfig)(successEvent) + Expect(err).NotTo(HaveOccurred()) + Expect(finish).NotTo(BeNil()) + Expect(trace.GetBackendTraces()).To(ContainElement(HaveField("Status", trace.BackendTraceRunning))) + finish(successEvent) Eventually(trace.GetBackendTraces).Should(HaveLen(1)) got := trace.GetBackendTraces()[0] @@ -57,7 +62,10 @@ var _ = Describe("ModelLoadTraceObserver", func() { failed := successEvent failed.Err = errors.New("grpc service not ready") - backend.ModelLoadTraceObserver(appConfig)(failed) + finish, err := backend.ModelLoadTraceObserver(appConfig)(failed) + Expect(err).NotTo(HaveOccurred()) + Expect(finish).NotTo(BeNil()) + finish(failed) Consistently(trace.GetBackendTraces, "100ms", "20ms").Should(BeEmpty()) }) @@ -65,7 +73,10 @@ var _ = Describe("ModelLoadTraceObserver", func() { It("records nothing when tracing is disabled", func() { appConfig.EnableTracing = false - backend.ModelLoadTraceObserver(appConfig)(successEvent) + finish, err := backend.ModelLoadTraceObserver(appConfig)(successEvent) + Expect(err).NotTo(HaveOccurred()) + Expect(finish).NotTo(BeNil()) + finish(successEvent) Consistently(trace.GetBackendTraces, "100ms", "20ms").Should(BeEmpty()) }) diff --git a/core/backend/options.go b/core/backend/options.go index c7525fee8..f8056d964 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -33,24 +33,38 @@ import ( // backend's launcher path, which names the variant directory) — that is what // identifies WHICH build served the load. A stale installed backend is // invisible in the model config but obvious here. -func ModelLoadTraceObserver(appConfig *config.ApplicationConfig) func(model.BackendLoadEvent) { - return func(ev model.BackendLoadEvent) { - if ev.Err != nil || !appConfig.EnableTracing { - return +func ModelLoadTraceObserver(appConfig *config.ApplicationConfig) func(model.BackendLoadEvent) (func(model.BackendLoadEvent), error) { + return func(ev model.BackendLoadEvent) (func(model.BackendLoadEvent), error) { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + if !appConfig.EnableTracing { + return func(model.BackendLoadEvent) { release() }, nil } trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) - trace.RecordBackendTrace(trace.BackendTrace{ - Timestamp: time.Now(), - Duration: ev.Duration, - Type: trace.BackendTraceModelLoad, - ModelName: ev.ModelID, - Backend: ev.Backend, - Summary: "Model loaded", - Data: map[string]any{ - "model_file": ev.ModelName, - "backend_runtime": ev.BackendURI, - }, - }) + started := time.Now() + id := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: started, Type: trace.BackendTraceModelLoad, ModelName: ev.ModelID, Backend: ev.Backend, Summary: "Loading model"}) + return func(done model.BackendLoadEvent) { + defer release() + if done.Err != nil { + trace.CancelBackendTrace(id) + return + } + trace.RecordBackendTrace(trace.BackendTrace{ + ID: id, + Timestamp: started, + Duration: done.Duration, + Type: trace.BackendTraceModelLoad, + ModelName: done.ModelID, + Backend: done.Backend, + Summary: "Model loaded", + Data: map[string]any{ + "model_file": done.ModelName, + "backend_runtime": done.BackendURI, + }, + }) + }, nil } } diff --git a/core/backend/rerank.go b/core/backend/rerank.go index 94314276d..af58a43c1 100644 --- a/core/backend/rerank.go +++ b/core/backend/rerank.go @@ -71,11 +71,19 @@ func Rerank(ctx context.Context, request *proto.RerankRequest, loader *model.Mod return nil, fmt.Errorf("could not load rerank model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceRerank, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(request.Query, 200)}) } + defer trace.CancelBackendTrace(traceID) // Stamped here, not at the HTTP handler: this is the function that also // builds ModelOptions from the same config, so the two values are equal by @@ -91,6 +99,7 @@ func Rerank(ctx context.Context, request *proto.RerankRequest, loader *model.Mod } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceRerank, diff --git a/core/backend/score.go b/core/backend/score.go index 174e86500..cf37f848d 100644 --- a/core/backend/score.go +++ b/core/backend/score.go @@ -96,15 +96,23 @@ func ModelScore(prompt string, candidates []string, opts ScoreOptions, loader *m return nil, fmt.Errorf("Score: candidates must be non-empty") } return func(ctx context.Context) ([]CandidateScore, error) { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() // Surface score calls in the Traces UI alongside the LLM calls // they typically gate (router classifier, eval scoring). Without // this, a router-classified request shows only the downstream LLM // trace with no record of the classification that picked it. var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceScore, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(prompt, 200)}) } + defer trace.CancelBackendTrace(traceID) resp, err := b.Score(ctx, &pb.ScoreRequest{ ModelIdentity: modelConfig.Model, Prompt: prompt, @@ -120,6 +128,7 @@ func ModelScore(prompt string, candidates []string, opts ScoreOptions, loader *m errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceScore, diff --git a/core/backend/soundgeneration.go b/core/backend/soundgeneration.go index e01e4b4a4..82cbe8492 100644 --- a/core/backend/soundgeneration.go +++ b/core/backend/soundgeneration.go @@ -101,11 +101,23 @@ func SoundGeneration( req.Instrumental = instrumental } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return "", nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + summary := trace.TruncateString(text, 200) + if summary == "" && caption != "" { + summary = trace.TruncateString(caption, 200) + } + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceSoundGeneration, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: summary}) } + defer trace.CancelBackendTrace(traceID) res, err := soundGenModel.SoundGeneration(ctx, req) @@ -135,6 +147,7 @@ func SoundGeneration( } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceSoundGeneration, diff --git a/core/backend/stores.go b/core/backend/stores.go index 1c324ab1a..5c97e9792 100644 --- a/core/backend/stores.go +++ b/core/backend/stores.go @@ -46,16 +46,24 @@ func (s *localVectorStore) backend(_ context.Context) (grpc.Backend, error) { } func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float64, payload []byte, ok bool, err error) { - start := time.Now() outcome := "hit" - defer func() { - s.recordTrace(start, "search", len(vec), sim, outcome, err) - }() be, berr := s.backend(ctx) if berr != nil { outcome = "backend_load_error" - return 0, nil, false, fmt.Errorf("vector store load: %w", berr) + err = fmt.Errorf("vector store load: %w", berr) + s.recordTrace("", time.Now(), "search", len(vec), 0, outcome, err) + return 0, nil, false, err } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return 0, nil, false, err + } + defer release() + start := time.Now() + traceID := s.beginTrace(start, "search") + defer func() { + s.recordTrace(traceID, start, "search", len(vec), sim, outcome, err) + }() _, values, similarities, ferr := store.Find(ctx, be, vec, 1) if ferr != nil { outcome = "find_error" @@ -69,16 +77,24 @@ func (s *localVectorStore) Search(ctx context.Context, vec []float32) (sim float } func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload []byte) (err error) { - start := time.Now() outcome := "ok" - defer func() { - s.recordTrace(start, "insert", len(vec), 0, outcome, err) - }() be, berr := s.backend(ctx) if berr != nil { outcome = "backend_load_error" - return fmt.Errorf("vector store load: %w", berr) + err = fmt.Errorf("vector store load: %w", berr) + s.recordTrace("", time.Now(), "insert", len(vec), 0, outcome, err) + return err } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + start := time.Now() + traceID := s.beginTrace(start, "insert") + defer func() { + s.recordTrace(traceID, start, "insert", len(vec), 0, outcome, err) + }() if serr := store.SetSingle(ctx, be, vec, payload); serr != nil { outcome = "insert_error" return serr @@ -91,7 +107,15 @@ func (s *localVectorStore) Insert(ctx context.Context, vec []float32, payload [] // modelName uses the store namespace (e.g. "router-cache-smart-router") so // admins can tell which router's cache misbehaved; the backend is always // "local-store" and can't disambiguate. -func (s *localVectorStore) recordTrace(start time.Time, op string, vecDim int, sim float64, outcome string, err error) { +func (s *localVectorStore) beginTrace(start time.Time, op string) string { + if s.appConfig == nil || !s.appConfig.EnableTracing { + return "" + } + trace.InitBackendTracingIfEnabled(s.appConfig.TracingMaxItems, s.appConfig.TracingMaxBodyBytes) + return trace.BeginBackendTrace(trace.BackendTrace{Timestamp: start, Type: trace.BackendTraceVectorStore, ModelName: s.storeName, Backend: model.LocalStoreBackend, Summary: op}) +} + +func (s *localVectorStore) recordTrace(traceID string, start time.Time, op string, vecDim int, sim float64, outcome string, err error) { if s.appConfig == nil || !s.appConfig.EnableTracing { return } @@ -115,6 +139,7 @@ func (s *localVectorStore) recordTrace(start time.Time, op string, vecDim int, s data["similarity"] = sim } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: start, Duration: time.Since(start), Type: trace.BackendTraceVectorStore, diff --git a/core/backend/token_classify.go b/core/backend/token_classify.go index 254b765cf..e316f2253 100644 --- a/core/backend/token_classify.go +++ b/core/backend/token_classify.go @@ -81,11 +81,19 @@ func ModelTokenClassify(text string, opts TokenClassifyOptions, loader *model.Mo return nil, err } return func(ctx context.Context) ([]TokenEntity, error) { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTokenClassify, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(text, 200)}) } + defer trace.CancelBackendTrace(traceID) resp, err := inferenceModel.TokenClassify(ctx, &pb.TokenClassifyRequest{ ModelIdentity: modelConfig.Model, Text: text, @@ -93,7 +101,9 @@ func ModelTokenClassify(text string, opts TokenClassifyOptions, loader *model.Mo }) entities := tokenClassifyResponseToEntities(resp) if appConfig.EnableTracing { - trace.RecordBackendTrace(tokenClassifyTrace(modelConfig, text, opts.Threshold, entities, startTime, err)) + bt := tokenClassifyTrace(modelConfig, text, opts.Threshold, entities, startTime, err) + bt.ID = traceID + trace.RecordBackendTrace(bt) } if err != nil { return nil, err diff --git a/core/backend/token_metrics.go b/core/backend/token_metrics.go index 45d60a406..fa564169d 100644 --- a/core/backend/token_metrics.go +++ b/core/backend/token_metrics.go @@ -26,6 +26,11 @@ func TokenMetrics( if model == nil { return nil, fmt.Errorf("could not loadmodel model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() res, err := model.GetTokenMetrics(ctx, &proto.MetricsRequest{}) diff --git a/core/backend/tokenize.go b/core/backend/tokenize.go index 6b926b179..3bc4cabf8 100644 --- a/core/backend/tokenize.go +++ b/core/backend/tokenize.go @@ -39,11 +39,19 @@ func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.Model predictOptions := gRPCPredictOpts(modelConfig, loader.ModelPath) predictOptions.Prompt = s + release, err := AcquireGlobalBackendSlot() + if err != nil { + return schema.TokenizeResponse{}, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTokenize, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(s, 200)}) } + defer trace.CancelBackendTrace(traceID) // tokenize the string resp, err := inferenceModel.TokenizeString(appConfig.Context, predictOptions) @@ -57,6 +65,7 @@ func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.Model tokenCount := tokenizeTokenCount(resp) trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTokenize, diff --git a/core/backend/transcript.go b/core/backend/transcript.go index 73826f3ae..0ddee3add 100644 --- a/core/backend/transcript.go +++ b/core/backend/transcript.go @@ -81,15 +81,23 @@ func ModelTranscriptionWithOptions(ctx context.Context, req TranscriptionRequest if err != nil { return nil, err } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string var audioSnippet map[string]any if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTranscription, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(req.Audio, 200)}) // Capture audio before the backend call — the backend may delete the file. audioSnippet = trace.AudioSnippet(req.Audio, appConfig.TracingMaxBodyBytes) } + defer trace.CancelBackendTrace(traceID) r, err := transcriptionModel.AudioTranscription(ctx, req.toProto(uint32(*modelConfig.Threads), modelConfig.Model)) if err != nil { @@ -105,6 +113,7 @@ func ModelTranscriptionWithOptions(ctx context.Context, req TranscriptionRequest maps.Copy(errData, audioSnippet) } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTranscription, @@ -133,6 +142,7 @@ func ModelTranscriptionWithOptions(ctx context.Context, req TranscriptionRequest maps.Copy(data, audioSnippet) } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTranscription, @@ -166,6 +176,11 @@ func ModelTranscriptionStream(ctx context.Context, req TranscriptionRequest, ml pbReq := req.toProto(uint32(*modelConfig.Threads), modelConfig.Model) pbReq.Stream = true + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() return transcriptionModel.AudioTranscriptionStream(ctx, pbReq, func(chunk *proto.TranscriptStreamResponse) { if chunk == nil { diff --git a/core/backend/transcript_live.go b/core/backend/transcript_live.go index 956e7e717..0ad6d72e1 100644 --- a/core/backend/transcript_live.go +++ b/core/backend/transcript_live.go @@ -55,6 +55,7 @@ type liveTranscriptionSession struct { closeOnce sync.Once closeErr error trace *liveTraceState // nil when tracing was disabled at open + release func() } func (s *liveTranscriptionSession) Feed(pcm []float32) error { @@ -80,6 +81,7 @@ func (s *liveTranscriptionSession) Close() error { } s.closeErr = err s.trace.record(err) + s.release() }) return s.closeErr } @@ -100,6 +102,7 @@ type liveTraceState struct { backend string language string started time.Time + traceID string mu sync.Mutex pcm []byte // first trace.MaxSnippetSeconds of fed audio, int16 LE @@ -115,12 +118,14 @@ func newLiveTraceState(modelConfig config.ModelConfig, appConfig *config.Applica return nil } trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) + started := time.Now() return &liveTraceState{ appConfig: appConfig, modelName: modelConfig.Name, backend: modelConfig.Backend, language: language, - started: time.Now(), + started: started, + traceID: trace.BeginBackendTrace(trace.BackendTrace{Timestamp: started, Type: trace.BackendTraceTranscription, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "live transcription"}), } } @@ -180,6 +185,7 @@ func (ts *liveTraceState) record(closeErr error) { ts.mu.Unlock() bt := trace.BackendTrace{ + ID: ts.traceID, Timestamp: ts.started, Duration: time.Since(ts.started), Type: trace.BackendTraceTranscription, @@ -210,6 +216,10 @@ func ModelTranscriptionLive(ctx context.Context, language string, if err != nil { return nil, err } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } // The derived cancel out-lives this call inside the session: Close uses // it to unwind the stream (and, in embed mode, the server-side recv @@ -218,12 +228,14 @@ func ModelTranscriptionLive(ctx context.Context, language string, stream, err := transcriptionModel.AudioTranscriptionLive(streamCtx) if err != nil { cancel() + release() return nil, err } fail := func(err error) (LiveTranscriptionSession, error) { _ = stream.CloseSend() cancel() + release() return nil, err } @@ -252,6 +264,7 @@ func ModelTranscriptionLive(ctx context.Context, language string, cancel: cancel, recvDone: make(chan struct{}), trace: newLiveTraceState(modelConfig, appConfig, language), + release: release, } go func() { diff --git a/core/backend/transcript_live_internal_test.go b/core/backend/transcript_live_internal_test.go index cbd7fac54..6f6bed6a4 100644 --- a/core/backend/transcript_live_internal_test.go +++ b/core/backend/transcript_live_internal_test.go @@ -148,6 +148,7 @@ var _ = Describe("liveTraceState", func() { It("clamps out-of-range float samples instead of wrapping", func() { ts := newLiveTraceState(modelCfg(), appConfig, "") + defer trace.CancelBackendTrace(ts.traceID) ts.addPCM([]float32{2.0, -2.0}) Expect(ts.pcm).To(Equal([]byte{0xff, 0x7f, 0x00, 0x80})) // 32767, -32768 }) diff --git a/core/backend/tts.go b/core/backend/tts.go index 11ca0aabd..3451eb471 100644 --- a/core/backend/tts.go +++ b/core/backend/tts.go @@ -96,11 +96,19 @@ func ModelTTS( modelPath = modelConfig.Model // skip this step if it fails????? } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return "", nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTTS, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(text, 200)}) } + defer trace.CancelBackendTrace(traceID) ttsRequest := newTTSRequest(text, modelPath, voice, filePath, language, instructions, params, modelConfig.Model) @@ -125,6 +133,7 @@ func ModelTTS( } } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTTS, @@ -188,11 +197,19 @@ func ModelTTSStream( modelPath = modelConfig.Model // skip this step if it fails????? } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceTTS, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(text, 200)}) } + defer trace.CancelBackendTrace(traceID) var sampleRate uint32 = 16000 // default headerSent := false @@ -292,6 +309,7 @@ func ModelTTSStream( } } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceTTS, diff --git a/core/backend/upscale.go b/core/backend/upscale.go index c821f56cd..ad3885174 100644 --- a/core/backend/upscale.go +++ b/core/backend/upscale.go @@ -19,7 +19,12 @@ func ImageUpscale(ctx context.Context, src, dst string, scale int, loader *model } fn := func() error { - _, err := inferenceModel.UpscaleImage( + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + _, err = inferenceModel.UpscaleImage( ctx, &proto.UpscaleImageRequest{ Src: src, diff --git a/core/backend/vad.go b/core/backend/vad.go index c74055b76..c6fefa345 100644 --- a/core/backend/vad.go +++ b/core/backend/vad.go @@ -28,6 +28,11 @@ func VAD(request *schema.VADRequest, ModelIdentity: modelConfig.Model, Audio: request.Audio, } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() resp, err := vadModel.VAD(ctx, &req) if err != nil { return nil, err diff --git a/core/backend/video.go b/core/backend/video.go index db56e363c..d991bb895 100644 --- a/core/backend/video.go +++ b/core/backend/video.go @@ -81,9 +81,11 @@ func VideoGeneration(options VideoGenerationOptions, loader *model.ModelLoader, "has_audio": options.Audio != "", } - startTime := time.Now() originalFn := fn fn = func() error { + startTime := time.Now() + traceID := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceVideoGeneration, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: trace.TruncateString(options.Prompt, 200)}) + defer trace.CancelBackendTrace(traceID) err := originalFn() duration := time.Since(startTime) @@ -93,6 +95,7 @@ func VideoGeneration(options VideoGenerationOptions, loader *model.ModelLoader, } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: duration, Type: trace.BackendTraceVideoGeneration, @@ -106,6 +109,15 @@ func VideoGeneration(options VideoGenerationOptions, loader *model.ModelLoader, return err } } + originalFn := fn + fn = func() error { + release, err := AcquireGlobalBackendSlot() + if err != nil { + return err + } + defer release() + return originalFn() + } return fn, nil } diff --git a/core/backend/voice_analyze.go b/core/backend/voice_analyze.go index ddb20a6a9..3b90e63ed 100644 --- a/core/backend/voice_analyze.go +++ b/core/backend/voice_analyze.go @@ -29,11 +29,19 @@ func VoiceAnalyze( return nil, fmt.Errorf("could not load voice recognition model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceVoiceAnalyze, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "voice analysis"}) } + defer trace.CancelBackendTrace(traceID) res, err := voiceModel.VoiceAnalyze(ctx, &proto.VoiceAnalyzeRequest{ ModelIdentity: modelConfig.Model, @@ -47,6 +55,7 @@ func VoiceAnalyze( errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceVoiceAnalyze, diff --git a/core/backend/voice_embed.go b/core/backend/voice_embed.go index 6608d4842..bfbe43e75 100644 --- a/core/backend/voice_embed.go +++ b/core/backend/voice_embed.go @@ -32,11 +32,19 @@ func VoiceEmbed( return nil, fmt.Errorf("could not load voice recognition model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceVoiceEmbed, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "voice embedding"}) } + defer trace.CancelBackendTrace(traceID) res, err := voiceModel.VoiceEmbed(ctx, &proto.VoiceEmbedRequest{ ModelIdentity: modelConfig.Model, @@ -49,6 +57,7 @@ func VoiceEmbed( errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceVoiceEmbed, diff --git a/core/backend/voice_verify.go b/core/backend/voice_verify.go index e05380687..ec8baea8d 100644 --- a/core/backend/voice_verify.go +++ b/core/backend/voice_verify.go @@ -30,11 +30,19 @@ func VoiceVerify( return nil, fmt.Errorf("could not load voice recognition model") } + release, err := AcquireGlobalBackendSlot() + if err != nil { + return nil, err + } + defer release() var startTime time.Time + var traceID string if appConfig.EnableTracing { trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems, appConfig.TracingMaxBodyBytes) startTime = time.Now() + traceID = trace.BeginBackendTrace(trace.BackendTrace{Timestamp: startTime, Type: trace.BackendTraceVoiceVerify, ModelName: modelConfig.Name, Backend: modelConfig.Backend, Summary: "voice verification"}) } + defer trace.CancelBackendTrace(traceID) res, err := voiceModel.VoiceVerify(ctx, &proto.VoiceVerifyRequest{ ModelIdentity: modelConfig.Model, @@ -50,6 +58,7 @@ func VoiceVerify( errStr = err.Error() } trace.RecordBackendTrace(trace.BackendTrace{ + ID: traceID, Timestamp: startTime, Duration: time.Since(startTime), Type: trace.BackendTraceVoiceVerify, diff --git a/core/cli/run.go b/core/cli/run.go index 53c8cbf47..4bcd80f2c 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -114,6 +114,7 @@ type RunCMD struct { MachineTag string `env:"LOCALAI_MACHINE_TAG,MACHINE_TAG" help:"Add Machine-Tag header to each response which is useful to track the machine in the P2P network" group:"api"` LoadToMemory []string `env:"LOCALAI_LOAD_TO_MEMORY,LOAD_TO_MEMORY" help:"A list of models to load into memory at startup" group:"models"` EnableTracing bool `env:"LOCALAI_ENABLE_TRACING,ENABLE_TRACING" help:"Enable API tracing" group:"api"` + MaxConcurrentBackendRequests int `env:"LOCALAI_MAX_CONCURRENT_BACKEND_REQUESTS,MAX_CONCURRENT_BACKEND_REQUESTS" default:"1024" help:"Maximum number of concurrent backend inference operations. Excess inference requests receive HTTP 503; UI and administrative endpoints remain available." group:"backends"` TracingMaxItems int `env:"LOCALAI_TRACING_MAX_ITEMS" default:"1024" help:"Maximum number of traces to keep" group:"api"` TracingMaxBodyBytes int `env:"LOCALAI_TRACING_MAX_BODY_BYTES" default:"65536" help:"Maximum bytes captured per request/response body in the trace buffer (0 = uncapped). Caps memory growth from chatty endpoints like /embeddings." group:"api"` AgentJobRetentionDays int `env:"LOCALAI_AGENT_JOB_RETENTION_DAYS,AGENT_JOB_RETENTION_DAYS" default:"30" help:"Number of days to keep agent job history (default: 30)" group:"api"` @@ -474,6 +475,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.EnableTracing { opts = append(opts, config.EnableTracing) } + opts = append(opts, config.WithMaxConcurrentBackendRequests(r.MaxConcurrentBackendRequests)) opts = append(opts, config.WithTracingMaxItems(r.TracingMaxItems)) opts = append(opts, config.WithTracingMaxBodyBytes(r.TracingMaxBodyBytes)) diff --git a/core/config/application_config.go b/core/config/application_config.go index 009fc1cea..aecb81fd3 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -42,6 +42,7 @@ type ApplicationConfig struct { F16 bool Debug bool EnableTracing bool + MaxConcurrentBackendRequests int TracingMaxItems int TracingMaxBodyBytes int // Per-body cap for captured request/response bodies; 0 disables the cap EnableBackendLogging bool @@ -212,6 +213,8 @@ type ApplicationConfig struct { Branding BrandingConfig } +const DefaultMaxConcurrentBackendRequests = 1024 + // BrandingConfig holds the whitelabel/branding configuration of the instance. // Text fields are exposed via the public GET /api/branding endpoint so the // login page can read them before authentication. Binary asset filenames @@ -294,8 +297,9 @@ func NewApplicationConfig(o ...AppOption) *ApplicationConfig { // Interval to the default on every restart (#10601). The effective // 500ms default is supplied at the watchdog layer (DefaultWatchdogInterval) // when the value is still 0. - TracingMaxItems: 1024, - TracingMaxBodyBytes: 64 * 1024, // 64 KiB - caps each request/response body in the trace buffer + MaxConcurrentBackendRequests: DefaultMaxConcurrentBackendRequests, + TracingMaxItems: 1024, + TracingMaxBodyBytes: 64 * 1024, // 64 KiB - caps each request/response body in the trace buffer AgentPool: AgentPoolConfig{ Enabled: true, Timeout: "5m", @@ -785,6 +789,15 @@ func WithTracingMaxItems(items int) AppOption { } } +func WithMaxConcurrentBackendRequests(requests int) AppOption { + return func(o *ApplicationConfig) { + if requests <= 0 { + requests = DefaultMaxConcurrentBackendRequests + } + o.MaxConcurrentBackendRequests = requests + } +} + func WithTracingMaxBodyBytes(bytes int) AppOption { return func(o *ApplicationConfig) { o.TracingMaxBodyBytes = bytes diff --git a/core/http/app.go b/core/http/app.go index 65e8ee917..2e1453ac0 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -19,6 +19,7 @@ import ( "github.com/mudler/LocalAI/pkg/model" + corebackend "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/http/endpoints/localai" @@ -78,6 +79,15 @@ func applyModelLoadCooldown(err error, code int, c echo.Context) int { return http.StatusServiceUnavailable } +func applyBackendAdmission(err error, code int, c echo.Context) int { + var capacityErr *corebackend.BackendAdmissionError + if !errors.As(err, &capacityErr) { + return code + } + c.Response().Header().Set("Retry-After", strconv.Itoa(int(capacityErr.RetryAfter.Seconds()))) + return http.StatusServiceUnavailable +} + // respondModelLoading answers a request whose model is still cold-loading with // 503, a Retry-After header and the live `loading` object, reporting true when // it handled the error. @@ -197,6 +207,7 @@ func API(application *application.Application) (*echo.Echo, error) { code = he.Code } code = applyModelLoadCooldown(err, code, c) + code = applyBackendAdmission(err, code, c) // Handle 404 errors: serve React SPA for HTML requests, JSON otherwise if code == http.StatusNotFound { @@ -225,6 +236,7 @@ func API(application *application.Application) (*echo.Echo, error) { code = he.Code } code = applyModelLoadCooldown(err, code, c) + code = applyBackendAdmission(err, code, c) // Opaque errors deliberately withhold the body, so a still-loading // model gets the status and Retry-After but no progress detail. code = applyModelLoading(err, code, c) diff --git a/core/http/react-ui/e2e/traces-live.spec.js b/core/http/react-ui/e2e/traces-live.spec.js index 1dfa6e5ec..08661c8a2 100644 --- a/core/http/react-ui/e2e/traces-live.spec.js +++ b/core/http/react-ui/e2e/traces-live.spec.js @@ -21,6 +21,38 @@ test('marks an API trace with no response status as in progress', async ({ page await expect(row.locator('.fa-check-circle')).toHaveCount(0) }) +test('shows a running backend trace and exposes its logs immediately', async ({ page }) => { + await page.route('**/api/traces?*', route => route.fulfill({ json: [] })) + await page.route('**/api/backend-traces?*', route => route.fulfill({ + json: [{ + id: 'backend-running-1', + timestamp: '2026-08-05T02:00:00Z', + duration: 2_000_000_000, + status: 'running', + type: 'llm', + model_name: 'slow-model', + backend: 'llama-cpp', + summary: 'generating a reply', + }], + headers: { 'X-Total-Count': '1' }, + })) + await page.route('**/api/backend-traces/backend-running-1', route => route.fulfill({ + json: { + id: 'backend-running-1', status: 'running', type: 'llm', + timestamp: '2026-08-05T02:00:00Z', duration: 2_000_000_000, + model_name: 'slow-model', backend: 'llama-cpp', summary: 'generating a reply', + }, + })) + + await page.goto('/app/traces?tab=backend') + await page.getByRole('button', { name: /Backend Traces/ }).click() + const row = page.locator('tbody tr').filter({ hasText: 'generating a reply' }) + await expect(row.locator('[title="Running"]')).toBeVisible() + await row.click() + await expect(page.getByRole('link', { name: 'View backend logs' })).toHaveAttribute('href', /\/app\/backend-logs\/slow-model/) + await expect(page.getByText(/running/, { exact: false })).toBeVisible() +}) + // Regression for #11376: switching from Backend Traces back to API Traces // used to crash the page. `traces` holds whichever list was fetched last, so // right after `setActiveTab('api')` — before the refetch effect lands — the diff --git a/core/http/react-ui/src/pages/Traces.jsx b/core/http/react-ui/src/pages/Traces.jsx index c53aadce9..8f55742d0 100644 --- a/core/http/react-ui/src/pages/Traces.jsx +++ b/core/http/react-ui/src/pages/Traces.jsx @@ -246,11 +246,12 @@ function DataFields({ data, nested }) { // Expanded detail for a backend trace row function BackendTraceDetail({ trace }) { + const running = trace.status === 'running' const infoItems = [ { label: 'Type', value: trace.type }, { label: 'Model', value: trace.model_name || '-' }, { label: 'Backend', value: trace.backend || '-' }, - { label: 'Duration', value: formatDuration(trace.duration) }, + { label: 'Duration', value: running ? `${formatDuration(trace.duration)} (running)` : formatDuration(trace.duration) }, ] return ( @@ -717,7 +718,9 @@ export default function Traces() { {formatDuration(trace.duration)} - {trace.error + {trace.status === 'running' + ? + : trace.error ? : } diff --git a/core/trace/backend_trace.go b/core/trace/backend_trace.go index 6dbf03ab6..e66b5268b 100644 --- a/core/trace/backend_trace.go +++ b/core/trace/backend_trace.go @@ -20,6 +20,8 @@ import ( type BackendTraceType string +type BackendTraceStatus string + const ( BackendTraceLLM BackendTraceType = "llm" BackendTraceEmbedding BackendTraceType = "embedding" @@ -47,17 +49,24 @@ const ( BackendTraceVectorStore BackendTraceType = "vector_store" ) +const ( + BackendTraceRunning BackendTraceStatus = "running" + BackendTraceCompleted BackendTraceStatus = "completed" + BackendTraceFailed BackendTraceStatus = "failed" +) + type BackendTrace struct { // ID identifies this trace for the lifetime of the process, so the list // endpoint can return trimmed entries and clients can fetch the full // record back from /api/backend-traces/:id. - ID string `json:"id"` - Timestamp time.Time `json:"timestamp"` - Duration time.Duration `json:"duration"` - Type BackendTraceType `json:"type"` - ModelName string `json:"model_name"` - Backend string `json:"backend"` - Summary string `json:"summary"` + ID string `json:"id"` + Timestamp time.Time `json:"timestamp"` + Duration time.Duration `json:"duration"` + Status BackendTraceStatus `json:"status,omitempty"` + Type BackendTraceType `json:"type"` + ModelName string `json:"model_name"` + Backend string `json:"backend"` + Summary string `json:"summary"` // Body is the full request payload sent to the backend, when one // applies (currently: cloud-proxy passthrough forwards). Summary // is a short preview for the trace list; Body is the full @@ -82,6 +91,8 @@ var ( backendTraceDataPath string backendTraceStore *tracepersist.Store[BackendTrace] backendTraceStoreKey string + backendInFlight = make(map[string]BackendTrace) + backendMaxInFlight = 1024 ) type backendTraceCommand struct { @@ -169,6 +180,17 @@ func ConfigureBackendTracePersistence(dataPath string) { backendMu.Unlock() } +// ConfigureBackendTraceMaxInFlight applies the process-wide admission ceiling +// to backend work that can outlive the HTTP request which started it. +func ConfigureBackendTraceMaxInFlight(max int) { + if max <= 0 { + max = 1024 + } + backendMu.Lock() + backendMaxInFlight = max + backendMu.Unlock() +} + func advanceBackendTraceID(id string) { n, err := strconv.ParseUint(id, 10, 64) if err != nil { @@ -178,7 +200,7 @@ func advanceBackendTraceID(id string) { } } -func RecordBackendTrace(t BackendTrace) { +func RecordBackendTrace(trace BackendTrace) { backendMu.Lock() maxBody := backendMaxBodyBytes store := backendTraceStore @@ -188,17 +210,61 @@ func RecordBackendTrace(t BackendTrace) { // (Inf/NaN) that encoding/json cannot marshal. A single such value — e.g. a // -Inf dBFS audio metric from a silent clip — would otherwise fail the whole // /api/backend-traces response and blank the Traces UI. - if t.Data != nil { - t.Data = sanitizeData(t.Data, maxBody) + if trace.Data != nil { + trace.Data = sanitizeData(trace.Data, maxBody) + } + if trace.ID == "" { + trace.ID = strconv.FormatUint(backendTraceIDSeq.Add(1), 10) + } + if trace.Error != "" { + trace.Status = BackendTraceFailed + } else { + trace.Status = BackendTraceCompleted + } + backendMu.Lock() + delete(backendInFlight, trace.ID) + backendMu.Unlock() + select { + case backendLogChan <- backendTraceCommand{trace: &trace, store: store}: + default: + xlog.Warn("Backend trace channel full, dropping trace") + } +} + +// BeginBackendTrace makes an operation visible before its backend call starts. +// RecordBackendTrace completes the same record when passed the trace again. +func BeginBackendTrace(t BackendTrace) string { + if t.Timestamp.IsZero() { + t.Timestamp = time.Now() } if t.ID == "" { t.ID = strconv.FormatUint(backendTraceIDSeq.Add(1), 10) } - select { - case backendLogChan <- backendTraceCommand{trace: &t, store: store}: - default: - xlog.Warn("Backend trace channel full, dropping trace") + t.Status = BackendTraceRunning + + backendMu.Lock() + defer backendMu.Unlock() + if backendMaxInFlight > 0 && len(backendInFlight) >= backendMaxInFlight { + return "" } + running := t + // Completion-only payloads may be mutated by producers while the request is + // running. The list and backend-log link only need this stable metadata. + running.Body = "" + running.Data = nil + backendInFlight[running.ID] = running + return running.ID +} + +// CancelBackendTrace removes an operation that is represented by a different +// completion trace, such as a model-load failure enriched by its caller. +func CancelBackendTrace(id string) { + if id == "" { + return + } + backendMu.Lock() + delete(backendInFlight, id) + backendMu.Unlock() } // sanitizeData walks a trace Data map (recursing into nested maps and slices) @@ -290,17 +356,22 @@ func sanitizeValue(v any, maxBytes int) (any, bool) { func GetBackendTraces() []BackendTrace { backendMu.Lock() - if backendTraceBuffer == nil { - backendMu.Unlock() - return []BackendTrace{} + var ptrs []*BackendTrace + if backendTraceBuffer != nil { + ptrs = backendTraceBuffer.Values() + } + running := make([]BackendTrace, 0, len(backendInFlight)) + for _, t := range backendInFlight { + t.Duration = time.Since(t.Timestamp) + running = append(running, t) } - ptrs := backendTraceBuffer.Values() backendMu.Unlock() - traces := make([]BackendTrace, len(ptrs)) - for i, p := range ptrs { - traces[i] = *p + traces := make([]BackendTrace, 0, len(ptrs)+len(running)) + for _, p := range ptrs { + traces = append(traces, *p) } + traces = append(traces, running...) slices.SortFunc(traces, func(a, b BackendTrace) int { return b.Timestamp.Compare(a.Timestamp) diff --git a/core/trace/backend_trace_live_test.go b/core/trace/backend_trace_live_test.go new file mode 100644 index 000000000..33a912234 --- /dev/null +++ b/core/trace/backend_trace_live_test.go @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: MIT + +package trace_test + +import ( + "time" + + "github.com/mudler/LocalAI/core/trace" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("live backend traces", func() { + BeforeEach(func() { + trace.InitBackendTracingIfEnabled(64, 1024) + trace.ConfigureBackendTraceMaxInFlight(64) + trace.ClearBackendTraces() + }) + + It("lists an operation while it is running and replaces it on completion", func() { + started := time.Now().Add(-time.Millisecond) + id := trace.BeginBackendTrace(trace.BackendTrace{ + Timestamp: started, + Type: trace.BackendTraceLLM, + ModelName: "slow-model", + Backend: "test-backend", + Summary: "hello", + }) + Expect(id).NotTo(BeEmpty()) + + running, ok := trace.GetBackendTrace(id) + Expect(ok).To(BeTrue()) + Expect(running.Status).To(Equal(trace.BackendTraceRunning)) + Expect(running.Duration).To(BeNumerically(">", 0)) + + trace.RecordBackendTrace(trace.BackendTrace{ + ID: id, + Timestamp: started, + Duration: time.Since(started), + Type: trace.BackendTraceLLM, + ModelName: "slow-model", + Backend: "test-backend", + Summary: "hello -> world", + }) + + Eventually(func() trace.BackendTrace { + completed, _ := trace.GetBackendTrace(id) + return completed + }).Should(HaveField("Status", trace.BackendTraceCompleted)) + Expect(trace.GetBackendTraces()).To(HaveLen(1)) + }) + + It("bounds running traces without dropping their eventual completion", func() { + trace.ConfigureBackendTraceMaxInFlight(1) + first := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: time.Now(), Type: trace.BackendTraceLLM}) + second := trace.BeginBackendTrace(trace.BackendTrace{Timestamp: time.Now(), Type: trace.BackendTraceEmbedding}) + Expect(first).NotTo(BeEmpty()) + Expect(second).To(BeEmpty()) + Expect(trace.GetBackendTraces()).To(HaveLen(1)) + + trace.RecordBackendTrace(trace.BackendTrace{Timestamp: time.Now(), Type: trace.BackendTraceEmbedding}) + trace.RecordBackendTrace(trace.BackendTrace{ID: first, Timestamp: time.Now(), Type: trace.BackendTraceLLM}) + Eventually(trace.GetBackendTraces).Should(HaveLen(2)) + }) +}) diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 4007ae106..820bc876a 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -81,6 +81,7 @@ For more information on VRAM management, see [VRAM and Memory Management]({{%rel | Parameter | Default | Description | Environment Variable | |-----------|---------|-------------|----------------------| | `--address` | `:8080` | Bind address for the API server | `$LOCALAI_ADDRESS`, `$ADDRESS` | +| `--max-concurrent-backend-requests` | `1024` | Process-wide ceiling for concurrent backend inference operations. Excess inference receives HTTP 503 with `Retry-After`; UI and administrative endpoints remain available | `$LOCALAI_MAX_CONCURRENT_BACKEND_REQUESTS`, `$MAX_CONCURRENT_BACKEND_REQUESTS` | | `--cors` | `false` | Enable CORS (Cross-Origin Resource Sharing) | `$LOCALAI_CORS`, `$CORS` | | `--cors-allow-origins` | | Comma-separated list of allowed CORS origins | `$LOCALAI_CORS_ALLOW_ORIGINS`, `$CORS_ALLOW_ORIGINS` | | `--csrf` | `false` | Enable Fiber CSRF middleware | `$LOCALAI_CSRF` | diff --git a/docs/content/reference/runtime-errors.md b/docs/content/reference/runtime-errors.md index 9122a1356..1e14dc0b5 100644 --- a/docs/content/reference/runtime-errors.md +++ b/docs/content/reference/runtime-errors.md @@ -22,6 +22,7 @@ The left column is the literal string as it appears in the LocalAI server log (o | `failed to load model: ...` | Returned by the load endpoints and several feature paths (voice, realtime, audio transform) when the model config could not be resolved or the backend load failed. | Confirm the model name exists (`local-ai models list`) and its YAML is valid. The trailing text carries the specific reason. | | HTTP `503` with a `Retry-After` header, after a load failed | Model-load failure cooldown. After a model fails to load, LocalAI refuses new load attempts for that model for a short window so a client that keeps polling a broken model does not respawn a crashing backend on every request. The window starts at `--model-load-failure-cooldown` (default `10s`) and doubles per consecutive failure up to 5m; it resets on the first success. | Fix the underlying load failure (see the rows above), then wait out the `Retry-After` seconds before retrying, or restart LocalAI to clear the cooldown. Set `--model-load-failure-cooldown 0` (or `LOCALAI_MODEL_LOAD_FAILURE_COOLDOWN=0`) to disable the cooldown entirely. See {{% relref "reference/cli-reference" %}}. | | HTTP `503` with a `Retry-After` header, under load | Per-model concurrency limit reached. When a model config sets a `MaxConcurrent` limit, extra requests are rejected with `503` and a `Retry-After` (whole seconds, floor 1) instead of queueing. | Retry after the advised delay, raise the model's concurrency limit, or run more replicas. | +| HTTP `503` when backend inference is saturated | The process-wide `--max-concurrent-backend-requests` backend-execution ceiling is full. This protects inference and in-flight backend-trace memory without blocking UI or administrative endpoints. | Retry after the advised delay, reduce inference concurrency, raise the limit if the host has capacity, or add replicas. | | `invalid pitch` (with CUDA) | The prompt exceeded the model's context size. | Reduce the prompt length, or raise the model's context size (`context_size:` in the model YAML). | | `SIGILL` (illegal instruction) on startup | The prebuilt backend binary uses CPU instructions your CPU does not have (for example AVX512, AVX2, F16C, FMA). | Rebuild the backend for your CPU. In a container, set `REBUILD=true` and disable the unsupported instructions, for example `CMAKE_ARGS="-DGGML_F16C=OFF -DGGML_AVX512=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF" make build`. | | CUDA / VRAM out of memory (backend log shows `out of memory`, `CUDA error: out of memory`, or the process is killed loading) | The model plus its KV cache does not fit in GPU memory. | Use a smaller quantization, reduce `context_size:`, offload fewer layers to the GPU (lower `gpu_layers:`), or free VRAM held by other processes. On multi-GPU hosts, confirm the model is not trying to load entirely onto one device. | diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index 506cee725..a6f8947eb 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -913,14 +913,76 @@ func (c *Client) Forward(ctx context.Context, opts ...grpc.CallOption) (ForwardC }, nil } +// LifecycleStream lets callers attach resource ownership to a stream's true +// terminal state without wrapping Recv themselves. +type LifecycleStream interface { + Context() context.Context + AddCleanup(func()) +} + // AudioTransformStreamClient is the duplex interface returned by // (*Client).AudioTransformStream. Wraps the generated bidi client without // leaking the proto package across the public boundary. type AudioTransformStreamClient interface { + LifecycleStream Send(*pb.AudioTransformFrameRequest) error Recv() (*pb.AudioTransformFrameResponse, error) CloseSend() error - Context() context.Context +} + +// streamCleanup owns callbacks whose lifetime follows the receive side of a +// stream. CloseSend is intentionally not terminal: bidi backends may continue +// producing their response tail after the caller finishes sending. +type streamCleanup struct { + mu sync.Mutex + done bool + doneCh chan struct{} + callbacks []func() +} + +func newStreamCleanup(ctx context.Context, initial func()) *streamCleanup { + l := &streamCleanup{doneCh: make(chan struct{})} + if initial != nil { + l.callbacks = append(l.callbacks, initial) + } + go func() { + select { + case <-ctx.Done(): + l.finish() + case <-l.doneCh: + } + }() + return l +} + +func (l *streamCleanup) add(fn func()) { + if fn == nil { + return + } + l.mu.Lock() + if !l.done { + l.callbacks = append(l.callbacks, fn) + l.mu.Unlock() + return + } + l.mu.Unlock() + fn() +} + +func (l *streamCleanup) finish() { + l.mu.Lock() + if l.done { + l.mu.Unlock() + return + } + l.done = true + callbacks := l.callbacks + l.callbacks = nil + close(l.doneCh) + l.mu.Unlock() + for _, fn := range callbacks { + fn() + } } // audioTransformStreamClient is the concrete wrapper. It also owns the @@ -930,14 +992,15 @@ type AudioTransformStreamClient interface { // forwardClient. type audioTransformStreamClient struct { pb.Backend_AudioTransformStreamClient - closeOnce sync.Once - closer func() + lifecycle *streamCleanup } +func (s *audioTransformStreamClient) AddCleanup(fn func()) { s.lifecycle.add(fn) } + func (s *audioTransformStreamClient) Recv() (*pb.AudioTransformFrameResponse, error) { resp, err := s.Backend_AudioTransformStreamClient.Recv() - if err != nil && s.closer != nil { - s.closeOnce.Do(s.closer) + if err != nil { + s.lifecycle.finish() } return resp, err } @@ -971,10 +1034,10 @@ func (c *Client) AudioTransformStream(ctx context.Context, opts ...grpc.CallOpti } return &audioTransformStreamClient{ Backend_AudioTransformStreamClient: stream, - closer: func() { + lifecycle: newStreamCleanup(ctx, func() { _ = conn.Close() cleanup() - }, + }), }, nil } diff --git a/pkg/grpc/embed.go b/pkg/grpc/embed.go index 6b616dfa0..e962c74c5 100644 --- a/pkg/grpc/embed.go +++ b/pkg/grpc/embed.go @@ -207,6 +207,7 @@ func (e *embedBackend) AudioTransformStream(ctx context.Context, opts ...grpc.Ca reqs: reqs, resps: resps, srvDone: srvDone, + cleanup: newStreamCleanup(ctx, nil), }, nil } @@ -397,8 +398,11 @@ type embedBackendAudioTransformStreamClient struct { resps <-chan *pb.AudioTransformFrameResponse srvDone <-chan error closeOnce bool + cleanup *streamCleanup } +func (e *embedBackendAudioTransformStreamClient) AddCleanup(fn func()) { e.cleanup.add(fn) } + func (e *embedBackendAudioTransformStreamClient) Send(req *pb.AudioTransformFrameRequest) error { select { case e.reqs <- req: @@ -412,6 +416,7 @@ func (e *embedBackendAudioTransformStreamClient) Recv() (*pb.AudioTransformFrame select { case resp, ok := <-e.resps: if !ok { + e.cleanup.finish() // Server-side finished. Surface its terminal error if any. select { case err := <-e.srvDone: @@ -424,6 +429,7 @@ func (e *embedBackendAudioTransformStreamClient) Recv() (*pb.AudioTransformFrame } return resp, nil case <-e.ctx.Done(): + e.cleanup.finish() return nil, e.ctx.Err() } } diff --git a/pkg/grpc/stream_cleanup_test.go b/pkg/grpc/stream_cleanup_test.go new file mode 100644 index 000000000..8b7fc20f8 --- /dev/null +++ b/pkg/grpc/stream_cleanup_test.go @@ -0,0 +1,40 @@ +package grpc + +import ( + "context" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("stream cleanup", func() { + It("runs registered cleanup exactly once at stream termination", func() { + var calls atomic.Int32 + lifecycle := newStreamCleanup(context.Background(), func() { calls.Add(1) }) + lifecycle.add(func() { calls.Add(1) }) + + lifecycle.finish() + lifecycle.finish() + + Expect(calls.Load()).To(Equal(int32(2))) + }) + + It("runs callbacks registered after termination immediately", func() { + var calls atomic.Int32 + lifecycle := newStreamCleanup(context.Background(), nil) + lifecycle.finish() + lifecycle.add(func() { calls.Add(1) }) + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("terminates when the stream context is cancelled", func() { + ctx, cancel := context.WithCancel(context.Background()) + var calls atomic.Int32 + lifecycle := newStreamCleanup(ctx, func() { calls.Add(1) }) + cancel() + Eventually(calls.Load).Should(Equal(int32(1))) + lifecycle.finish() + Expect(calls.Load()).To(Equal(int32(1))) + }) +}) diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index 8d36f1e7a..2462973e6 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -73,15 +73,24 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string uri := ml.GetAllExternalBackends(o)[backend] start := time.Now() + started := BackendLoadEvent{ModelID: modelID, ModelName: modelName, Backend: backend, BackendURI: uri} + finish, admissionErr := ml.notifyLoadStarted(started) + if admissionErr != nil { + return nil, admissionErr + } m, err := ml.spawnGRPCModel(backend, uri, o, modelID, modelName, modelFile) - ml.notifyLoadObserver(BackendLoadEvent{ + completed := BackendLoadEvent{ ModelID: modelID, ModelName: modelName, Backend: backend, BackendURI: uri, Duration: time.Since(start), Err: err, - }) + } + ml.notifyLoadObserver(completed) + if finish != nil { + finish(completed) + } return m, err } } diff --git a/pkg/model/loader.go b/pkg/model/loader.go index fd54d358f..77b7a0ed6 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -94,6 +94,7 @@ type ModelLoader struct { lruEvictionRetryInterval time.Duration // Interval between retries when waiting for busy models onUnloadHooks []ModelUnloadHook loadObserver func(BackendLoadEvent) + loadLifecycleObserver func(BackendLoadEvent) (func(BackendLoadEvent), error) remoteUnloader RemoteModelUnloader modelRouter ModelRouter // distributed mode: route to remote node backendLogs *BackendLogStore @@ -256,6 +257,24 @@ func (ml *ModelLoader) notifyLoadObserver(ev BackendLoadEvent) { } } +// SetLoadLifecycleObserver registers a begin callback whose returned function +// is invoked when the same real backend load attempt finishes. +func (ml *ModelLoader) SetLoadLifecycleObserver(obs func(BackendLoadEvent) (func(BackendLoadEvent), error)) { + ml.mu.Lock() + defer ml.mu.Unlock() + ml.loadLifecycleObserver = obs +} + +func (ml *ModelLoader) notifyLoadStarted(ev BackendLoadEvent) (func(BackendLoadEvent), error) { + ml.mu.Lock() + obs := ml.loadLifecycleObserver + ml.mu.Unlock() + if obs == nil { + return nil, nil + } + return obs(ev) +} + // SetRemoteUnloader sets the handler for unloading models on remote nodes. // In distributed mode, this should be set to the SmartRouter adapter. func (ml *ModelLoader) SetRemoteUnloader(u RemoteModelUnloader) {