From 799cc9f21133433a3b1cba79b768ec21116b4aad Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Tue, 18 Aug 2026 07:56:59 +0100 Subject: [PATCH] feat: bound global admission and expose running backend traces (#11560) feat: bound backend admission and expose running traces Add process-wide backend execution admission without blocking UI or administrative HTTP work. Represent backend operations while they are in flight, surface running traces with immediate log links, and tie streaming admission leases to the gRPC receive lifecycle. Assisted-by: OpenAI Codex: GPT-5 Signed-off-by: Richard Palethorpe --- core/application/application.go | 5 +- core/backend/audio_transform.go | 21 +++- core/backend/depth.go | 9 ++ core/backend/detection.go | 9 ++ core/backend/detokenize.go | 9 ++ core/backend/embeddings.go | 23 +++- core/backend/face_analyze.go | 9 ++ core/backend/face_embed.go | 5 + core/backend/face_verify.go | 9 ++ core/backend/global_admission.go | 72 +++++++++++ core/backend/global_admission_test.go | 39 ++++++ core/backend/image.go | 14 ++- core/backend/llm.go | 20 +++- core/backend/model3d.go | 14 ++- core/backend/model_load_trace_test.go | 17 ++- core/backend/options.go | 46 ++++--- core/backend/rerank.go | 9 ++ core/backend/score.go | 9 ++ core/backend/soundgeneration.go | 13 ++ core/backend/stores.go | 47 ++++++-- core/backend/token_classify.go | 12 +- core/backend/token_metrics.go | 5 + core/backend/tokenize.go | 9 ++ core/backend/transcript.go | 15 +++ core/backend/transcript_live.go | 15 ++- core/backend/transcript_live_internal_test.go | 1 + core/backend/tts.go | 18 +++ core/backend/upscale.go | 7 +- core/backend/vad.go | 5 + core/backend/video.go | 14 ++- core/backend/voice_analyze.go | 9 ++ core/backend/voice_embed.go | 9 ++ core/backend/voice_verify.go | 9 ++ core/cli/run.go | 2 + core/config/application_config.go | 17 ++- core/http/app.go | 12 ++ core/http/react-ui/e2e/traces-live.spec.js | 32 +++++ core/http/react-ui/src/pages/Traces.jsx | 7 +- core/trace/backend_trace.go | 113 ++++++++++++++---- core/trace/backend_trace_live_test.go | 65 ++++++++++ docs/content/reference/cli-reference.md | 1 + docs/content/reference/runtime-errors.md | 1 + pkg/grpc/client.go | 77 ++++++++++-- pkg/grpc/embed.go | 6 + pkg/grpc/stream_cleanup_test.go | 40 +++++++ pkg/model/initializers.go | 13 +- pkg/model/loader.go | 19 +++ 47 files changed, 853 insertions(+), 79 deletions(-) create mode 100644 core/backend/global_admission.go create mode 100644 core/backend/global_admission_test.go create mode 100644 core/trace/backend_trace_live_test.go create mode 100644 pkg/grpc/stream_cleanup_test.go 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) {