diff --git a/core/application/application.go b/core/application/application.go index 326885b31..174e290c6 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -165,6 +165,7 @@ func newApplication(appConfig *config.ApplicationConfig) *Application { backendLoader: config.NewModelConfigLoader( appConfig.SystemState.Model.ModelsPath, config.WithArtifactMaterializer(appConfig.ModelArtifactMaterializer), + config.WithReloadDiagnostics(appConfig.DiagnosticsRecorder), config.WithPreloadDisplay(appConfig.ModelPreloadRenderMode, appConfig.DisableModelPreloadColor), ), modelLoader: ml, diff --git a/core/cli/run.go b/core/cli/run.go index 522cc0499..905586967 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -3,8 +3,10 @@ package cli import ( "context" "encoding/json" + "errors" "fmt" "net" + nethttp "net/http" "os" "path/filepath" "strings" @@ -16,6 +18,7 @@ import ( "github.com/mudler/LocalAI/core/http" "github.com/mudler/LocalAI/core/p2p" "github.com/mudler/LocalAI/internal" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/signals" "github.com/mudler/LocalAI/pkg/system" @@ -30,6 +33,12 @@ import ( // and document the deprecation in the help text. type RunCMD struct { + Pprof bool `name:"pprof" env:"LOCALAI_PPROF" default:"false" help:"Enable the private loopback profiling listener" group:"diagnostics"` + PprofAddress string `name:"pprof-address" env:"LOCALAI_PPROF_ADDRESS" default:"127.0.0.1:6060" help:"Numeric loopback address and port for profiling" group:"diagnostics"` + PprofMutexProfileFraction int `name:"pprof-mutex-profile-fraction" env:"LOCALAI_PPROF_MUTEX_PROFILE_FRACTION" default:"0" help:"Mutex sampling fraction (0 disables sampling); requires pprof" group:"diagnostics"` + PprofBlockProfileRate int `name:"pprof-block-profile-rate" env:"LOCALAI_PPROF_BLOCK_PROFILE_RATE" default:"0" help:"Block sampling rate in nanoseconds (0 disables sampling); requires pprof" group:"diagnostics"` + RequestPhaseTiming bool `name:"request-phase-timing" env:"LOCALAI_REQUEST_PHASE_TIMING" default:"false" help:"Log sanitized request and config-reload phase timings at Info" group:"diagnostics"` + ModelArgs []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"` Color string `env:"COLOR" hidden:""` NoColor string `env:"NO_COLOR" hidden:""` @@ -244,7 +253,7 @@ func parseDistributedDuration(envVar, raw string) (time.Duration, error) { return d, nil } -func (r *RunCMD) Run(ctx *cliContext.Context) error { +func (r *RunCMD) Run(ctx *cliContext.Context) (result error) { warnDeprecatedFlags() if r.Version { @@ -252,6 +261,11 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { return nil } + diagnosticOpts := r.diagnosticsOptions() + if err := diagnosticOpts.Validate(); err != nil { + return err + } + if ctx.CredentialsFile == "" { if err := LoadCredentials(resolveCredentialsFile("", r.DataPath)); err != nil { return err @@ -293,9 +307,14 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { return err } + // CLI context holds configuration, not a cancellation context. + runCtx, cancelRun := context.WithCancelCause(context.Background()) + if !r.PreloadBackendOnly { + defer cancelRun(nil) + } opts := []config.AppOption{ config.WithProxyAPIKeyEnvLookup(os.Getenv), - config.WithContext(context.Background()), + config.WithContext(runCtx), config.WithArtifactDownloadConcurrency(r.ArtifactDownloadConcurrency), config.WithModelArtifactMaterializer(modelartifacts.NewDefaultManager( modelartifacts.WithHuggingFaceToken(r.HFToken), @@ -789,11 +808,61 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { return err } + var stopProfiler func() error + if diagnosticOpts.Pprof { + profiler, err := diagnostics.Start(diagnosticOpts) + if err != nil { + return err + } + monitorDone := make(chan struct{}) + var profilerErr error + go func() { + defer close(monitorDone) + if err := <-profiler.Errors(); err != nil { + profilerErr = err + cancelRun(err) + } + }() + stopProfiler = func() error { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := profiler.Shutdown(shutdownCtx) + <-monitorDone + return errors.Join(err, profilerErr) + } + defer func() { + result = errors.Join(result, stopProfiler()) + }() + // Signals exit the process without running defers. Register private + // cleanup before application cleanup, which has no bounded shutdown. + signals.RegisterGracefulTerminationHandler(func() { + cancelRun(nil) + if err := stopProfiler(); err != nil { + xlog.Error("error while shutting down profiler", "error", err) + } + }) + } + if diagnosticOpts.RequestPhaseTiming { + opts = append(opts, config.WithDiagnosticsRecorder(diagnostics.NewRecorder(diagnostics.LogEvent))) + } + app, err := application.New(opts...) if err != nil { return fmt.Errorf("LocalAI failed to start: %w.\nTroubleshooting steps:\n 1. Check that your models directory exists and is accessible: %s\n 2. Verify model config files are valid YAML: 'local-ai util usecase-heuristic '\n 3. Check available disk space and file permissions\n 4. Run with --log-level=debug for more details\nSee https://localai.io/basics/troubleshooting/ for more help", err, r.ModelsPath) } + defer func() { + cancelRun(nil) + if stopProfiler != nil { + // Release the bounded private resource even if app.Shutdown blocks. + _ = stopProfiler() + } + result = errors.Join(result, app.Shutdown()) + }() + if err := runCtx.Err(); err != nil { + return context.Cause(runCtx) + } + // Refuse to bind a public-internet address without authentication unless // the operator has explicitly opted in. The auth middleware degrades to // pass-through when there is no auth DB and no legacy keys; on a loopback, @@ -830,6 +899,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } signals.RegisterGracefulTerminationHandler(func() { + cancelRun(nil) if err := app.Shutdown(); err != nil { xlog.Error("error while shutting down application", "error", err) } @@ -840,10 +910,33 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { // collection initialization. go func() { waitForServerReady(listenAddress, app.ApplicationConfig().Context) - app.StartAgentPool() + if runCtx.Err() == nil { + app.StartAgentPool() + } }() - return appHTTP.Start(listenAddress) + if !diagnosticOpts.Pprof { + return appHTTP.Start(listenAddress) + } + served := make(chan error, 1) + go func() { served <- appHTTP.Start(listenAddress) }() + select { + case err := <-served: + return errors.Join(err, context.Cause(runCtx)) + case <-runCtx.Done(): + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + shutdownErr := appHTTP.Shutdown(shutdownCtx) + if shutdownErr != nil { + shutdownErr = errors.Join(shutdownErr, appHTTP.Close()) + } + serveErr := <-served + // Only intentional shutdown makes ErrServerClosed an expected result. + if errors.Is(serveErr, nethttp.ErrServerClosed) { + serveErr = nil + } + return errors.Join(serveErr, context.Cause(runCtx), shutdownErr) + } } // waitForServerReady polls the given address until the HTTP server is diff --git a/core/cli/run_diagnostics.go b/core/cli/run_diagnostics.go new file mode 100644 index 000000000..d58701123 --- /dev/null +++ b/core/cli/run_diagnostics.go @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: MIT +package cli + +import "github.com/mudler/LocalAI/pkg/diagnostics" + +func (r *RunCMD) diagnosticsOptions() diagnostics.Options { + return diagnostics.Options{Pprof: r.Pprof, Address: r.PprofAddress, MutexProfileFraction: r.PprofMutexProfileFraction, BlockProfileRate: r.PprofBlockProfileRate, RequestPhaseTiming: r.RequestPhaseTiming} +} diff --git a/core/cli/run_diagnostics_test.go b/core/cli/run_diagnostics_test.go new file mode 100644 index 000000000..a3dddb4e2 --- /dev/null +++ b/core/cli/run_diagnostics_test.go @@ -0,0 +1,256 @@ +// SPDX-License-Identifier: MIT +package cli + +import ( + "context" + "fmt" + "io" + "log/slog" + "net" + "os" + "os/exec" + "path/filepath" + "runtime" + "strconv" + "sync" + "syscall" + "time" + + "github.com/alecthomas/kong" + cliContext "github.com/mudler/LocalAI/core/cli/context" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/xlog" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics CLI", Serial, func() { + parse := func(args ...string) (*RunCMD, error) { + command := struct { + cliContext.Context `embed:""` + Run RunCMD `cmd:"" default:"withargs"` + }{} + r := &command.Run + // Match cmd/local-ai/main.go, including gallery substitutions. Parsing + // must fail for the supplied value, not an unrelated missing variable. + p, err := kong.New(&command, kong.Vars{ + "basepath": kong.ExpandPath("."), + "generatedcontentpath": DefaultGeneratedContentPath(), + "uploadpath": DefaultUploadPath(), + "galleries": config.DefaultGalleriesJSON, + "backends": config.DefaultBackendGalleriesJSON, + }) + if err != nil { + return r, err + } + _, err = p.Parse(args) + return r, err + } + BeforeEach(func() { + for _, key := range []string{"LOCALAI_PPROF", "LOCALAI_PPROF_ADDRESS", "LOCALAI_PPROF_MUTEX_PROFILE_FRACTION", "LOCALAI_PPROF_BLOCK_PROFILE_RATE", "LOCALAI_REQUEST_PHASE_TIMING"} { + value, present := os.LookupEnv(key) + DeferCleanup(func() { + if present { + Expect(os.Setenv(key, value)).To(Succeed()) + } else { + Expect(os.Unsetenv(key)).To(Succeed()) + } + }) + Expect(os.Unsetenv(key)).To(Succeed()) + } + }) + It("defaults to inactive independently of debug", func() { + for _, args := range [][]string{nil, {"--debug"}, {"--log-level=debug"}} { + r, err := parse(args...) + Expect(err).NotTo(HaveOccurred()) + Expect(r.diagnosticsOptions()).To(Equal(diagnostics.DefaultOptions())) + } + }) + It("maps all five flags", func() { + r, err := parse("--pprof", "--pprof-address=[::1]:6061", "--pprof-mutex-profile-fraction=7", "--pprof-block-profile-rate=123", "--request-phase-timing") + Expect(err).NotTo(HaveOccurred()) + Expect(r.diagnosticsOptions()).To(Equal(diagnostics.Options{Pprof: true, Address: "[::1]:6061", MutexProfileFraction: 7, BlockProfileRate: 123, RequestPhaseTiming: true})) + }) + It("maps all environment settings and explicit false", func() { + GinkgoT().Setenv("LOCALAI_PPROF", "true") + GinkgoT().Setenv("LOCALAI_PPROF_ADDRESS", "127.0.0.2:6062") + GinkgoT().Setenv("LOCALAI_PPROF_MUTEX_PROFILE_FRACTION", "3") + GinkgoT().Setenv("LOCALAI_PPROF_BLOCK_PROFILE_RATE", "42") + GinkgoT().Setenv("LOCALAI_REQUEST_PHASE_TIMING", "true") + r, err := parse() + Expect(err).NotTo(HaveOccurred()) + Expect(r.diagnosticsOptions()).To(Equal(diagnostics.Options{Pprof: true, Address: "127.0.0.2:6062", MutexProfileFraction: 3, BlockProfileRate: 42, RequestPhaseTiming: true})) + GinkgoT().Setenv("LOCALAI_PPROF", "false") + GinkgoT().Setenv("LOCALAI_REQUEST_PHASE_TIMING", "false") + r, err = parse() + Expect(err).NotTo(HaveOccurred()) + Expect(r.Pprof).To(BeFalse()) + Expect(r.RequestPhaseTiming).To(BeFalse()) + }) + It("rejects malformed and overflowing rates", func() { + for _, arg := range []string{"--pprof-block-profile-rate=no", "--pprof-mutex-profile-fraction=999999999999999999999999"} { + _, err := parse(arg) + Expect(err).To(HaveOccurred()) + } + GinkgoT().Setenv("LOCALAI_PPROF", "not-bool") + _, err := parse() + Expect(err).To(HaveOccurred()) + }) + It("validates disabled addresses and conflicting rates", func() { + for _, arg := range []string{"--pprof-address=0.0.0.0:6060", "--pprof-block-profile-rate=-1", "--pprof-mutex-profile-fraction=1"} { + r, err := parse(arg) + Expect(err).NotTo(HaveOccurred()) + Expect(r.diagnosticsOptions().Validate()).To(HaveOccurred()) + } + }) + It("stores one optional startup recorder", func() { + Expect(config.NewApplicationConfig().DiagnosticsRecorder).To(BeNil()) + recorder := diagnostics.NewRecorder(nil) + Expect(config.NewApplicationConfig(config.WithDiagnosticsRecorder(recorder)).DiagnosticsRecorder).To(BeIdenticalTo(recorder)) + }) +}) + +var _ = Describe("Diagnostics inline startup", Serial, func() { + fixture := func() *RunCMD { + dir := GinkgoT().TempDir() + return &RunCMD{ + PprofAddress: diagnostics.DefaultOptions().Address, + ModelsPath: filepath.Join(dir, "models"), BackendsPath: filepath.Join(dir, "backends"), + BackendsSystemPath: filepath.Join(dir, "system"), DataPath: filepath.Join(dir, "data"), + Galleries: "[]", BackendGalleries: "[]", + // Invalid JSON fails application configuration without starting services. + PreloadModels: "not-json", + } + } + It("validates before startup side effects even in preload mode", func() { + r := &RunCMD{PreloadBackendOnly: true, PprofAddress: "0.0.0.0:6060"} + Expect(r.Run(nil)).To(MatchError(r.diagnosticsOptions().Validate())) + }) + It("reports a busy profiler before application initialization", func() { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + r := fixture() + r.Pprof, r.PprofAddress = true, l.Addr().String() + Expect(r.Run(&cliContext.Context{CredentialsFile: "unused"})).To(MatchError(ContainSubstring("bind diagnostics listener"))) + }) + It("does not bind profiling for preload or independent timing", func() { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + for _, preload := range []bool{true, false} { + r := fixture() + r.PreloadBackendOnly, r.Pprof = preload, preload + r.RequestPhaseTiming, r.PprofAddress = true, l.Addr().String() + err := r.Run(&cliContext.Context{CredentialsFile: "unused"}) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).NotTo(ContainSubstring("bind diagnostics listener")) + } + }) + It("releases profiling after real application configuration failure", func() { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + address := l.Addr().String() + Expect(l.Close()).To(Succeed()) + r := fixture() + r.Pprof, r.PprofAddress = true, address + Expect(r.Run(&cliContext.Context{CredentialsFile: "unused"})).To(MatchError(ContainSubstring("LocalAI failed to start"))) + l, err = net.Listen("tcp", address) + Expect(err).NotTo(HaveOccurred()) + Expect(l.Close()).To(Succeed()) + }) +}) + +// Fail the real private listener only once Run has constructed the public API. +// The subprocess isolates the logger, signal registrations and socket operation. +type diagnosticsListenerFailure struct { + slog.Handler + port int + failed chan error + canceled chan struct{} + cancelOnce sync.Once +} + +func (h *diagnosticsListenerFailure) Handle(ctx context.Context, record slog.Record) error { + if record.Message == "Context canceled, shutting down" { + h.cancelOnce.Do(func() { close(h.canceled) }) + } + if record.Message == "LocalAI is started and running" { + entries, err := os.ReadDir("/proc/self/fd") + if err != nil { + h.failed <- err + return nil + } + for _, entry := range entries { + fd, err := strconv.Atoi(entry.Name()) + if err != nil { + continue + } + address, err := syscall.Getsockname(fd) + if err != nil { + continue + } + if address, ok := address.(*syscall.SockaddrInet4); ok && address.Port == h.port { + if err := syscall.Shutdown(fd, syscall.SHUT_RDWR); err != nil { + h.failed <- err + return nil + } + // Application's cancellation log acknowledges the real run context, + // without sleeps or a production lifecycle injection point. + select { + case <-h.canceled: + h.failed <- nil + case <-time.After(5 * time.Second): + h.failed <- fmt.Errorf("run context was not canceled") + } + return nil + } + } + h.failed <- fmt.Errorf("private listener not found") + } + return nil +} + +var _ = Describe("Diagnostics concurrent public bind failure", Serial, func() { + It("retains the occupied API port error alongside cancellation", func() { + if runtime.GOOS != "linux" { + Skip("socket failure fixture uses /proc/self/fd") + } + if os.Getenv("LOCALAI_DIAGNOSTICS_BIND_CHILD") != "1" { + executable, err := os.Executable() + Expect(err).NotTo(HaveOccurred()) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, executable, "-test.run=TestCLI", "-ginkgo.focus=Diagnostics concurrent public bind failure", "-ginkgo.fail-on-empty") + cmd.Env = append(os.Environ(), "LOCALAI_DIAGNOSTICS_BIND_CHILD=1") + output, err := cmd.CombinedOutput() + Expect(err).NotTo(HaveOccurred(), string(output)) + return + } + occupied, err := net.Listen("tcp4", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer occupied.Close() + private, err := net.Listen("tcp4", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + address := private.Addr().String() + port := private.Addr().(*net.TCPAddr).Port + Expect(private.Close()).To(Succeed()) + failure := &diagnosticsListenerFailure{Handler: slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelDebug}), port: port, failed: make(chan error, 1), canceled: make(chan struct{})} + xlog.SetLogger(xlog.NewLoggerWithHandler(failure, xlog.LogLevelDebug)) + dir := GinkgoT().TempDir() + r := &RunCMD{ + Pprof: true, PprofAddress: address, Address: occupied.Addr().String(), + ModelsPath: filepath.Join(dir, "models"), BackendsPath: filepath.Join(dir, "backends"), + BackendsSystemPath: filepath.Join(dir, "system"), DataPath: filepath.Join(dir, "data"), + GeneratedContentPath: filepath.Join(dir, "generated"), UploadPath: filepath.Join(dir, "uploads"), + Galleries: "[]", BackendGalleries: "[]", DisableAgents: true, DisableLocalAIAssistant: true, + DisableMCP: true, DisableMetricsEndpoint: true, DisableWebUI: true, + } + err = r.Run(&cliContext.Context{CredentialsFile: "unused"}) + Expect(failure.failed).To(Receive(BeNil())) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("accept")) + Expect(err.Error()).To(ContainSubstring("address already in use")) + }) +}) diff --git a/core/config/application_config.go b/core/config/application_config.go index 8de3ef103..1a6042fb5 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -6,6 +6,7 @@ import ( "regexp" "time" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/vrambudget" @@ -14,6 +15,9 @@ import ( ) type ApplicationConfig struct { + // DiagnosticsRecorder is startup-only and local to this frontend process. + DiagnosticsRecorder *diagnostics.Recorder `json:"-" yaml:"-"` + // ProxyAPIKeyEnvLookup resolves upstream credentials at the CLI boundary. ProxyAPIKeyEnvLookup func(string) string `json:"-" yaml:"-"` @@ -1158,6 +1162,9 @@ func (o *ApplicationConfig) ToConfigLoaderOptions() []ConfigLoaderOption { LoadOptionThreads(o.Threads), ModelPath(o.SystemState.Model.ModelsPath), LoadOptionGalleryFiles(o.Galleries...), + // Select later YAML/default phases; early phases use WithReloadDiagnostics + // at loader construction so callbacks retain their original execution point. + LoadOptionDiagnostics(o.DiagnosticsRecorder), } } @@ -1233,3 +1240,8 @@ func (o *ApplicationConfig) ApplyRuntimeSettings(settings *RuntimeSettings) (req func WithRouterDecisionLog(enabled bool) AppOption { return func(o *ApplicationConfig) { o.RouterDecisionLog = enabled } } + +// WithDiagnosticsRecorder installs the optional startup-only timing recorder. +func WithDiagnosticsRecorder(r *diagnostics.Recorder) AppOption { + return func(o *ApplicationConfig) { o.DiagnosticsRecorder = r } +} diff --git a/core/config/model_config_diagnostics_test.go b/core/config/model_config_diagnostics_test.go new file mode 100644 index 000000000..bddf7ecf3 --- /dev/null +++ b/core/config/model_config_diagnostics_test.go @@ -0,0 +1,405 @@ +// SPDX-License-Identifier: MIT +package config + +import ( + "context" + "errors" + "io/fs" + "os" + "path/filepath" + "sync" + "time" + + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics config loader", func() { + DescribeTable("preserves stateful caller options for each model's defaults", + func(enabled bool) { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(filepath.Join(dir, "models.yaml"), []byte("- name: a\n- name: b\n"), 0600)).To(Succeed()) + var events []diagnostics.Event + var recorder *diagnostics.Recorder + if enabled { + recorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + } + b := NewModelConfigLoader(dir, WithReloadDiagnostics(recorder)) + calls := 0 + Expect(b.LoadModelConfigsFromPath(dir, func(o *LoadOptions) { + calls++ + LoadOptionThreads(calls)(o) + }, LoadOptionDiagnostics(recorder))).To(Succeed()) + Expect(calls).To(Equal(3)) + Expect(*b.configs["a"].Threads).To(Equal(2)) + Expect(*b.configs["b"].Threads).To(Equal(3)) + if enabled { + Expect(events).NotTo(BeEmpty()) + } else { + Expect(events).To(BeEmpty()) + } + }, + Entry("with diagnostics enabled", true), + Entry("with diagnostics disabled", false), + ) + + It("uses explicit startup timing for early failures and preserves identity across late selection", func() { + for _, mode := range []string{"same", "nil", "different", "failure"} { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(filepath.Join(dir, "a.yaml"), []byte("name: a\n"), 0600)).To(Succeed()) + var early, late []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { early = append(early, e) }) + other := diagnostics.NewRecorder(func(e diagnostics.Event) { late = append(late, e) }) + shared := NewModelConfigLoader(dir, WithReloadDiagnostics(r)) + b := NewModelConfigLoader(dir, shared.ReloadDiagnosticsOption()) + selected := r + if mode == "nil" { + selected = nil + } + if mode == "different" { + selected = other + } + calls := 0 + if mode == "failure" { + b.diagnosticsReadDir = func(string) ([]os.DirEntry, error) { return nil, errors.New("enumeration") } + } + err := b.LoadModelConfigsFromPath(dir, func(o *LoadOptions) { calls++; LoadOptionDiagnostics(selected)(o) }) + if mode == "failure" { + Expect(err).To(HaveOccurred()) + Expect(calls).To(BeZero()) + Expect(early).To(HaveLen(3)) + Expect(early[1].Phase).To(Equal(diagnostics.PhaseReloadEnumeration)) + Expect(early[1].Outcome).To(Equal(diagnostics.OutcomeError)) + } else { + Expect(err).NotTo(HaveOccurred()) + Expect(calls).To(Equal(2)) + if mode == "same" { + Expect(early).To(HaveLen(6)) + } else { + Expect(early).To(HaveLen(4)) + } + } + for _, e := range early { + Expect(e.ID).To(Equal(early[0].ID)) + } + if mode == "different" { + Expect(late).To(HaveLen(2)) + Expect(late[0].ID).NotTo(Equal(early[0].ID)) + Expect(late[1].ID).To(Equal(late[0].ID)) + } + } + }) + + It("honors composed diagnostics options and last writes at each original application site", func() { + for _, enabled := range []bool{true, false} { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(filepath.Join(dir, "a.yaml"), []byte("name: a\nbackend: llama-cpp\n"), 0600)).To(Succeed()) + b := NewModelConfigLoader(dir) + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + var order []string + b.diagnosticsReadDir = func(path string) ([]os.DirEntry, error) { + order = append(order, "enumerate") + entries, err := os.ReadDir(path) + return []os.DirEntry{diagnosticsTrackedEntry{DirEntry: entries[0], visit: func() { order = append(order, "metadata") }}}, err + } + wrap := func(label string, option ConfigLoaderOption) ConfigLoaderOption { + return func(o *LoadOptions) { + unlocked := b.TryLock() + if unlocked { + b.Unlock() + } + Expect(unlocked).To(BeFalse()) + order = append(order, label) + option(o) + } + } + var last *diagnostics.Recorder + if enabled { + b.diagnosticsNow = func() time.Time { + Expect(order).To(Or( + Equal([]string{"enumerate", "metadata", "first", "last"}), + Equal([]string{"enumerate", "metadata", "first", "last", "first", "last"}), + )) + return time.Now() + } + last = r + } else { + b.diagnosticsNow = func() time.Time { Fail("disabled clock read"); return time.Time{} } + } + Expect(b.LoadModelConfigsFromPath(dir, LoadOptionDiagnostics(r), wrap("first", LoadOptionDiagnostics(nil)), wrap("last", LoadOptionDiagnostics(last)))).To(Succeed()) + Expect(order).To(Equal([]string{"enumerate", "metadata", "first", "last", "first", "last"})) + if enabled { + Expect(events).To(HaveLen(2)) + Expect(events[0].Phase).To(Equal(diagnostics.PhaseReloadYAMLRead)) + Expect(events[1].Phase).To(Equal(diagnostics.PhaseReloadParseDefaults)) + } else { + Expect(events).To(BeEmpty()) + } + } + }) + + It("records malformed configured gallery parse errors without changing strict behavior", func() { + for _, strict := range []bool{false, true} { + dir := GinkgoT().TempDir() + path := filepath.Join(dir, "gallery.yaml") + Expect(os.WriteFile(path, []byte("[invalid"), 0600)).To(Succeed()) + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + gallery := LoadOptionGalleryFiles(Gallery{URL: "file://" + path}) + b := NewModelConfigLoader(dir) + err := b.loadModelConfigsFromPath(dir, strict, gallery, LoadOptionDiagnostics(r)) + off := NewModelConfigLoader(dir) + offErr := off.loadModelConfigsFromPath(dir, strict, gallery) + if strict { + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(Equal(offErr.Error())) + } else { + Expect(err).NotTo(HaveOccurred()) + Expect(offErr).NotTo(HaveOccurred()) + } + Expect(b.GetAllModelsConfigs()).To(Equal(off.GetAllModelsConfigs())) + var reads, parses int + for _, e := range events { + if e.Phase == diagnostics.PhaseReloadYAMLRead { + reads += e.Count + } + if e.Phase == diagnostics.PhaseReloadParseDefaults { + parses += e.Count + Expect(e.Outcome).To(Equal(diagnostics.OutcomeError)) + } + } + Expect(reads).To(Equal(1)) + Expect(parses).To(Equal(1)) + ok, classifyErr := classifyGalleryDocument(path) + Expect(ok).To(BeFalse()) + Expect(classifyErr).NotTo(HaveOccurred()) + } + }) + + It("applies caller options under the lock after metadata, never on enumeration failure", func() { + b := NewModelConfigLoader("") + enumerated := false + calls := 0 + b.diagnosticsReadDir = func(string) ([]os.DirEntry, error) { enumerated = true; return nil, nil } + option := func(*LoadOptions) { + calls++ + unlocked := b.TryLock() + if unlocked { + b.Unlock() + } + Expect(unlocked).To(BeFalse()) + Expect(enumerated).To(BeTrue()) + } + Expect(b.LoadModelConfigsFromPath("", option)).To(Succeed()) + Expect(calls).To(Equal(1)) + b.diagnosticsReadDir = func(string) ([]os.DirEntry, error) { return nil, errors.New("enumeration") } + Expect(b.LoadModelConfigsFromPath("", option)).To(HaveOccurred()) + Expect(calls).To(Equal(1)) + }) + + It("preserves filtering and panic propagation and emits only after unlock", func() { + b := NewModelConfigLoader("") + b.configs["a"] = ModelConfig{Name: "a"} + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { + Expect(b.TryLock()).To(BeTrue()) + b.Unlock() + events = append(events, e) + }) + ctx, cancel := context.WithCancel(r.Request(context.Background())) + cancel() + calls := 0 + got := b.GetModelConfigsByFilterContext(ctx, func(_ string, _ *ModelConfig) bool { calls++; return true }) + Expect(got).To(Equal(b.GetModelConfigsByFilter(nil))) + Expect(calls).To(Equal(1)) + Expect(events).To(HaveLen(3)) + for _, e := range events { + Expect(e.Kind).To(Equal(diagnostics.KindRequest)) + Expect(e.Count).To(Equal(1)) + } + Expect(func() { b.GetModelConfigsByFilterContext(ctx, func(string, *ModelConfig) bool { panic("sentinel") }) }).To(PanicWith("sentinel")) + Expect(b.TryLock()).To(BeTrue()) + b.Unlock() + Expect(events[len(events)-1].Outcome).To(Equal(diagnostics.OutcomeError)) + }) + It("separates a real gated mutex wait from held work without time thresholds", func() { + b := NewModelConfigLoader("") + b.configs["a"] = ModelConfig{Name: "a"} + held, release, waiting := make(chan struct{}), make(chan struct{}), make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + b.GetModelConfigsByFilter(func(string, *ModelConfig) bool { close(held); <-release; return true }) + }() + <-held + tick := 0 + b.diagnosticsNow = func() time.Time { + tick++ + if tick == 1 { + close(waiting) + } + return time.Unix(0, int64(tick)) + } + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + wg.Add(1) + go func() { defer wg.Done(); b.GetModelConfigsByFilterContext(r.Request(context.Background()), nil) }() + <-waiting + Expect(b.TryLock()).To(BeFalse()) + close(release) + wg.Wait() + Expect(events).To(HaveLen(3)) + Expect(events[0].Elapsed).To(Equal(time.Nanosecond)) + Expect(events[2].Elapsed).To(BeNumerically(">", events[0].Elapsed)) + }) + It("propagates startup options and aggregates actual reads including gallery classification", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(filepath.Join(dir, "a.yaml"), []byte("name: a\nbackend: llama-cpp\nparameters:\n model: a.bin\n"), 0600)).To(Succeed()) + gallery := filepath.Join(dir, "gallery.yaml") + Expect(os.WriteFile(gallery, []byte("- name: g\n url: https://example.invalid/g\n"), 0600)).To(Succeed()) + b := NewModelConfigLoader(dir) + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { Expect(b.TryLock()).To(BeTrue()); b.Unlock(); events = append(events, e) }) + WithReloadDiagnostics(r)(b) + opts := []ConfigLoaderOption{LoadOptionDiagnostics(r), LoadOptionGalleryFiles(Gallery{URL: "file://" + gallery})} + Expect(b.LoadModelConfigsFromPathStrict(dir, opts...)).To(Succeed()) + Expect(events).To(HaveLen(6)) + counts := map[diagnostics.Phase]int{} + for _, e := range events { + counts[e.Phase] = e.Count + Expect(e.Kind).To(Equal(diagnostics.KindConfigReload)) + } + Expect(counts[diagnostics.PhaseReloadMetadata]).To(Equal(2)) + Expect(counts[diagnostics.PhaseReloadYAMLRead]).To(Equal(2)) + Expect(counts[diagnostics.PhaseReloadParseDefaults]).To(Equal(2)) + first := events[0].ID + Expect(b.LoadModelConfigsFromPathStrict(dir, opts...)).To(Succeed()) + Expect(events[6].ID).NotTo(Equal(first)) + off := NewModelConfigLoader(dir) + Expect(off.LoadModelConfigsFromPathStrict(dir, opts[1:]...)).To(Succeed()) + Expect(off.GetModelConfigsByFilter(nil)).To(Equal(b.GetModelConfigsByFilter(nil))) + app := &ApplicationConfig{DiagnosticsRecorder: r, SystemState: &system.SystemState{}} + lo := &LoadOptions{} + lo.Apply(app.ToConfigLoaderOptions()...) + Expect(lo.diagnosticsRecorder).To(BeIdenticalTo(r)) + }) + It("retains strict and non-strict parse/read failures and records failed internal phases", func() { + for _, readFailure := range []bool{false, true} { + for _, strict := range []bool{false, true} { + dir := GinkgoT().TempDir() + if readFailure { + Expect(os.Mkdir(filepath.Join(dir, "bad.yaml"), 0700)).To(Succeed()) + } else { + Expect(os.WriteFile(filepath.Join(dir, "bad.yaml"), []byte("[invalid"), 0600)).To(Succeed()) + } + b := NewModelConfigLoader(dir) + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { Expect(b.TryLock()).To(BeTrue()); b.Unlock(); events = append(events, e) }) + err := b.loadModelConfigsFromPath(dir, strict, LoadOptionDiagnostics(r)) + off := NewModelConfigLoader(dir) + off.diagnosticsNow = func() time.Time { Fail("disabled clock read"); return time.Time{} } + offErr := off.loadModelConfigsFromPath(dir, strict) + if strict { + Expect(err).To(HaveOccurred()) + Expect(offErr.Error()).To(Equal(err.Error())) + } else { + Expect(err).NotTo(HaveOccurred()) + Expect(offErr).NotTo(HaveOccurred()) + } + phase := diagnostics.PhaseReloadParseDefaults + if readFailure { + phase = diagnostics.PhaseReloadYAMLRead + } + found := false + for _, e := range events { + if e.Phase == phase { + found = true + Expect(e.Outcome).To(Equal(diagnostics.OutcomeError)) + Expect(e.Count).To(Equal(1)) + } + } + Expect(found).To(BeTrue()) + Expect(off.GetModelConfigsByFilter(nil)).To(Equal(b.GetModelConfigsByFilter(nil))) + } + } + }) + It("records enumeration and metadata failures without changing filesystem operations", func() { + dir := GinkgoT().TempDir() + for _, metadata := range []bool{false, true} { + b := NewModelConfigLoader(dir) + calls := 0 + if metadata { + b.diagnosticsReadDir = func(string) ([]os.DirEntry, error) { calls++; return []os.DirEntry{diagnosticsBadEntry{}}, nil } + } else { + dir = filepath.Join(dir, "missing") + } + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { Expect(b.TryLock()).To(BeTrue()); b.Unlock(); events = append(events, e) }) + WithReloadDiagnostics(r)(b) + Expect(b.LoadModelConfigsFromPath(dir, LoadOptionDiagnostics(r))).To(HaveOccurred()) + phase := diagnostics.PhaseReloadEnumeration + if metadata { + phase = diagnostics.PhaseReloadMetadata + Expect(calls).To(Equal(1)) + } + found := false + for _, e := range events { + if e.Phase == phase { + found = true + Expect(e.Outcome).To(Equal(diagnostics.OutcomeError)) + Expect(e.Count).To(Equal(1)) + } + } + Expect(found).To(BeTrue()) + } + }) + It("keeps concurrent reload and request identities separate and direct callers disabled", func() { + dir := GinkgoT().TempDir() + b := NewModelConfigLoader(dir) + var mu sync.Mutex + var events []diagnostics.Event + r := diagnostics.NewRecorder(func(e diagnostics.Event) { mu.Lock(); defer mu.Unlock(); events = append(events, e) }) + WithReloadDiagnostics(r)(b) + var wg sync.WaitGroup + for range 2 { + wg.Add(1) + go func() { + defer GinkgoRecover() + defer wg.Done() + Expect(b.LoadModelConfigsFromPath(dir, LoadOptionDiagnostics(r))).To(Succeed()) + }() + } + wg.Wait() + b.GetModelConfigsByFilterContext(r.Request(context.Background()), nil) + ids := map[string]diagnostics.Kind{} + for _, e := range events { + ids[e.ID] = e.Kind + } + Expect(ids).To(HaveLen(3)) + n := len(events) + b = NewModelConfigLoader(dir) + b.diagnosticsNow = func() time.Time { Fail("disabled clock read"); return time.Time{} } + Expect(b.LoadModelConfigsFromPath(dir)).To(Succeed()) + b.GetModelConfigsByFilter(nil) + Expect(events).To(HaveLen(n)) + }) + +}) + +type diagnosticsBadEntry struct{} + +func (diagnosticsBadEntry) Name() string { return "bad.yaml" } +func (diagnosticsBadEntry) IsDir() bool { return false } +func (diagnosticsBadEntry) Type() fs.FileMode { return 0 } +func (diagnosticsBadEntry) Info() (fs.FileInfo, error) { return nil, errors.New("metadata failed") } + +type diagnosticsTrackedEntry struct { + os.DirEntry + visit func() +} + +func (e diagnosticsTrackedEntry) Info() (fs.FileInfo, error) { e.visit(); return e.DirEntry.Info() } diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index 7c1157ace..b6554f937 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -13,9 +13,11 @@ import ( "slices" "strings" "sync" + "time" "github.com/charmbracelet/glamour" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/modelartifacts" @@ -53,6 +55,10 @@ func WithPreloadDisplay(renderMode string, disableColor bool) ModelConfigLoaderO } type ModelConfigLoader struct { + reloadDiagnosticsRecorder *diagnostics.Recorder + // Optional per-instance test seams, configured before use. + diagnosticsNow func() time.Time + diagnosticsReadDir func(string) ([]os.DirEntry, error) configs map[string]ModelConfig modelPath string artifactMaterializer ArtifactMaterializer @@ -86,11 +92,90 @@ func NewModelConfigLoader(modelPath string, options ...ModelConfigLoaderOption) } type LoadOptions struct { - modelPath string - debug bool - threads, ctxSize int - f16 bool - galleryFiles map[string]struct{} + diagnosticsRecorder *diagnostics.Recorder + modelPath string + debug bool + threads, ctxSize int + f16 bool + galleryFiles map[string]struct{} +} + +// WithReloadDiagnostics sets the immutable startup recorder for lock wait/hold, +// directory enumeration and metadata. Use only during loader construction. +// Nil disables these early measurements, including all per-entry clocks. +// Later YAML/read/default phases are independently selected by LoadOptionDiagnostics. +func WithReloadDiagnostics(r *diagnostics.Recorder) ModelConfigLoaderOption { + return func(bcl *ModelConfigLoader) { bcl.reloadDiagnosticsRecorder = r } +} + +// ReloadDiagnosticsOption copies startup observation configuration to another +// loader (for example an authoritative snapshot). Each reload still creates its +// own operation ID; this does not attribute the other loader's lock to this one. +func (bcl *ModelConfigLoader) ReloadDiagnosticsOption() ModelConfigLoaderOption { + return WithReloadDiagnostics(bcl.reloadDiagnosticsRecorder) +} + +// LoadOptionDiagnostics selects observation of YAML reads and parse/default work +// after options execute under the lock, following enumeration and metadata. +// Options compose normally and the last write wins. Nil disables only subsequent +// phases, not early measurements selected by WithReloadDiagnostics. Without a +// constructor recorder no early spans are collected or retrospectively invented. +// Matching startup and late recorders share an operation ID; a different late +// recorder creates a separate observation ID covering only these later phases. +func LoadOptionDiagnostics(r *diagnostics.Recorder) ConfigLoaderOption { + return func(o *LoadOptions) { o.diagnosticsRecorder = r } +} + +// loaderTimings buffers observations until the model-config mutex is released. +// A nil accumulator performs no clock reads, allocation, or synchronization. +type loaderTimings struct { + ctx context.Context + now func() time.Time + phases map[diagnostics.Phase]*diagnostics.Event + order []diagnostics.Phase +} + +func (bcl *ModelConfigLoader) timings(ctx context.Context) *loaderTimings { + if !diagnostics.Enabled(ctx) { + return nil + } + now := bcl.diagnosticsNow + if now == nil { + now = time.Now + } + return &loaderTimings{ctx: ctx, now: now, phases: make(map[diagnostics.Phase]*diagnostics.Event)} +} +func (t *loaderTimings) start() time.Time { + if t == nil { + return time.Time{} + } + return t.now() +} +func (t *loaderTimings) add(phase diagnostics.Phase, start time.Time, err error, count int) { + if t == nil { + return + } + elapsed := t.now().Sub(start) + e := t.phases[phase] + if e == nil { + e = &diagnostics.Event{Phase: phase, Outcome: diagnostics.OutcomeOK} + t.phases[phase] = e + t.order = append(t.order, phase) + } + e.Elapsed += elapsed + e.Count += count + if err != nil { + e.Outcome = diagnostics.OutcomeError + } +} +func (t *loaderTimings) emit() { + if t == nil { + return + } + for _, phase := range t.order { + e := t.phases[phase] + diagnostics.Record(t.ctx, phase, e.Elapsed, e.Outcome, e.Count) + } } func LoadOptionDebug(debug bool) ConfigLoaderOption { @@ -159,11 +244,19 @@ func (lo *LoadOptions) Apply(options ...ConfigLoaderOption) { // ModelConfig or an array of ModelConfigs. It tries to unmarshal as an array first, // then falls back to a single config if that fails. func readModelConfigsFromFile(file string, opts ...ConfigLoaderOption) ([]*ModelConfig, error) { + return readModelConfigsFromFileMeasured(file, nil, opts...) +} + +func readModelConfigsFromFileMeasured(file string, timing *loaderTimings, opts ...ConfigLoaderOption) (_ []*ModelConfig, resultErr error) { + start := timing.start() f, err := os.ReadFile(file) + timing.add(diagnostics.PhaseReloadYAMLRead, start, err, 1) if err != nil { return nil, fmt.Errorf("readModelConfigsFromFile cannot read config file %q: %w", file, err) } + start = timing.start() + defer func() { timing.add(diagnostics.PhaseReloadParseDefaults, start, resultErr, 1) }() // Try to unmarshal as array first var configs []*ModelConfig if err := yaml.Unmarshal(f, &configs); err == nil && len(configs) > 0 { @@ -342,8 +435,38 @@ func (bcl *ModelConfigLoader) GetAllModelsConfigs() []ModelConfig { } func (bcl *ModelConfigLoader) GetModelConfigsByFilter(filter ModelConfigFilterFn) []ModelConfig { + return bcl.GetModelConfigsByFilterContext(context.Background(), filter) +} + +// GetModelConfigsByFilterContext measures request-owned lock and filter work. +func (bcl *ModelConfigLoader) GetModelConfigsByFilterContext(ctx context.Context, filter ModelConfigFilterFn) []ModelConfig { + timing := bcl.timings(ctx) + wait := timing.start() bcl.Lock() - defer bcl.Unlock() + hold := timing.start() + count := 0 + completed := false + defer func() { + var err error + if !completed { + err = errors.New("filter panic") + } + timing.add(diagnostics.PhaseConfigLockHold, hold, err, count) + bcl.Unlock() + timing.emit() + }() + if timing != nil { + timing.phases[diagnostics.PhaseConfigLockWait] = &diagnostics.Event{Elapsed: hold.Sub(wait), Outcome: diagnostics.OutcomeOK, Count: 1} + timing.order = append(timing.order, diagnostics.PhaseConfigLockWait) + } + start := timing.start() + defer func() { + var err error + if !completed { + err = errors.New("filter panic") + } + timing.add(diagnostics.PhaseConfigFilter, start, err, count) + }() var res []ModelConfig if filter == nil { @@ -351,11 +474,13 @@ func (bcl *ModelConfigLoader) GetModelConfigsByFilter(filter ModelConfigFilterFn } for n, v := range bcl.configs { + count++ if filter(n, &v) { res = append(res, v) } } + completed = true // TODO: I don't think this one needs to Sort on name... but we'll see what breaks. return res @@ -846,17 +971,39 @@ func (bcl *ModelConfigLoader) LoadModelConfigsFromPathStrict(path string, opts . return bcl.loadModelConfigsFromPath(path, true, opts...) } -func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, opts ...ConfigLoaderOption) error { +func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, opts ...ConfigLoaderOption) (resultErr error) { + timing := bcl.timings(bcl.reloadDiagnosticsRecorder.Reload(context.Background())) + var lateTiming *loaderTimings + wait := timing.start() bcl.Lock() - defer bcl.Unlock() - - entries, err := os.ReadDir(path) + hold := timing.start() + defer func() { + timing.add(diagnostics.PhaseReloadLockHold, hold, resultErr, 1) + bcl.Unlock() + timing.emit() + if lateTiming != timing { + lateTiming.emit() + } + }() + if timing != nil { + timing.phases[diagnostics.PhaseReloadLockWait] = &diagnostics.Event{Elapsed: hold.Sub(wait), Outcome: diagnostics.OutcomeOK, Count: 1} + timing.order = append(timing.order, diagnostics.PhaseReloadLockWait) + } + start := timing.start() + readDir := os.ReadDir + if bcl.diagnosticsReadDir != nil { + readDir = bcl.diagnosticsReadDir + } + entries, err := readDir(path) + timing.add(diagnostics.PhaseReloadEnumeration, start, err, 1) if err != nil { return fmt.Errorf("LoadModelConfigsFromPath cannot read directory '%s': %w", path, err) } files := make([]fs.FileInfo, 0, len(entries)) for _, entry := range entries { + start := timing.start() info, err := entry.Info() + timing.add(diagnostics.PhaseReloadMetadata, start, err, 1) if err != nil { return err } @@ -864,6 +1011,11 @@ func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, } loadOptions := &LoadOptions{} loadOptions.Apply(opts...) + if loadOptions.diagnosticsRecorder == bcl.reloadDiagnosticsRecorder { + lateTiming = timing + } else { + lateTiming = bcl.timings(loadOptions.diagnosticsRecorder.Reload(context.Background())) + } for _, file := range files { // Only load real YAML config files and ignore dotfiles or backup variants ext := strings.ToLower(filepath.Ext(file.Name())) @@ -877,7 +1029,7 @@ func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, return absErr } if _, gallerySource := loadOptions.galleryFiles[absolutePath]; gallerySource { - galleryDocument, err := classifyGalleryDocument(filePath) + galleryDocument, err := classifyGalleryDocumentMeasured(filePath, lateTiming) if err != nil { if strict { return err @@ -896,7 +1048,7 @@ func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, } // Read config(s) - handles both single and array formats - configs, err := readModelConfigsFromFile(filePath, opts...) + configs, err := readModelConfigsFromFileMeasured(filePath, lateTiming, opts...) if err != nil { if strict { return err @@ -989,13 +1141,28 @@ var galleryMetadataKeys = map[string]struct{}{ // document subject to the complete shape check; malformed or mixed documents // are errors rather than silently disappearing from an authoritative snapshot. func classifyGalleryDocument(path string) (bool, error) { + return classifyGalleryDocumentMeasured(path, nil) +} + +func classifyGalleryDocumentMeasured(path string, timing *loaderTimings) (_ bool, resultErr error) { + start := timing.start() data, _, err := safefile.ReadRegularAt(filepath.Dir(path), filepath.Base(path)) + timing.add(diagnostics.PhaseReloadYAMLRead, start, err, 1) if err != nil { return false, fmt.Errorf("read YAML file %q for classification: %w", path, err) } + start = timing.start() + var parseErr error + defer func() { + if parseErr == nil { + parseErr = resultErr + } + timing.add(diagnostics.PhaseReloadParseDefaults, start, parseErr, 1) + }() var document yaml.Node - if err := yaml.Unmarshal(data, &document); err != nil || len(document.Content) != 1 { - return false, nil // The model-config parser supplies the syntax error. + parseErr = yaml.Unmarshal(data, &document) + if parseErr != nil || len(document.Content) != 1 { + return false, nil // Preserve classification semantics; diagnostics retain the syntax error. } root := document.Content[0] switch root.Kind { diff --git a/core/http/diagnostics_test.go b/core/http/diagnostics_test.go new file mode 100644 index 000000000..f71c2c715 --- /dev/null +++ b/core/http/diagnostics_test.go @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: MIT +package http_test + +import ( + "context" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "time" + + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/config" + api "github.com/mudler/LocalAI/core/http" + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/LocalAI/pkg/httpclient" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics public API isolation", Serial, func() { + for _, debug := range []bool{false, true} { + for _, profiling := range []bool{false, true} { + for _, timing := range []bool{false, true} { + It(fmt.Sprintf("keeps pprof private with debug=%t profiling=%t timing=%t", debug, profiling, timing), func() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + dir := GinkgoT().TempDir() + modelDir, backendDir := filepath.Join(dir, "models"), filepath.Join(dir, "backends") + Expect(os.Mkdir(modelDir, 0700)).To(Succeed()) + Expect(os.Mkdir(backendDir, 0700)).To(Succeed()) + state, err := system.GetSystemState(system.WithModelPath(modelDir), system.WithBackendPath(backendDir)) + Expect(err).NotTo(HaveOccurred()) + // Reserve the public address independently; changing the private option must not change it. + publicListener, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer publicListener.Close() + privateReservation, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + opts := diagnostics.DefaultOptions() + opts.Pprof = profiling + opts.RequestPhaseTiming = timing + opts.Address = privateReservation.Addr().String() + Expect(privateReservation.Close()).To(Succeed()) + private, err := diagnostics.Start(opts) + Expect(err).NotTo(HaveOccurred()) + defer func() { + c, done := context.WithTimeout(context.Background(), 3*time.Second) + defer done() + Expect(private.Shutdown(c)).To(Succeed()) + }() + var recorder *diagnostics.Recorder + if timing { + recorder = diagnostics.NewRecorder(func(diagnostics.Event) {}) + } + app, err := application.New(config.WithContext(ctx), config.WithSystemState(state), config.WithDebug(debug), config.WithDataPath(filepath.Join(dir, "data")), config.DisableRuntimeSettings, config.DisableAgentPool, config.DisableMCP, config.DisableMetricsEndpoint, config.WithDisableStats(true), config.WithAPIAddress(publicListener.Addr().String()), func(c *config.ApplicationConfig) { + c.DiagnosticsRecorder = recorder + c.DisableLocalAIAssistant = true + }) + Expect(err).NotTo(HaveOccurred()) + defer func() { cancel(); Expect(app.Shutdown()).To(Succeed()) }() + router, err := api.API(app) + Expect(err).NotTo(HaveOccurred()) + defer router.Close() + Expect(app.ApplicationConfig().APIAddress).To(Equal(publicListener.Addr().String())) + Expect(app.ApplicationConfig().APIAddress).NotTo(Equal(opts.Address)) + for _, route := range router.Routes() { + Expect(route.Path).NotTo(ContainSubstring("/debug/pprof")) + Expect(strings.ToLower(route.Name)).NotTo(ContainSubstring("pprof")) + } + for _, path := range []string{"/debug/pprof/", "/debug/pprof/profile?seconds=1", "/debug/pprof/heap", "/debug/pprof/cmdline"} { + for _, accept := range []string{"application/json", "text/html"} { + req := httptest.NewRequest(http.MethodGet, path, nil) + req.Header.Set("Accept", accept) + response := httptest.NewRecorder() + router.ServeHTTP(response, req) + Expect(response.Header().Get("Content-Type")).NotTo(ContainSubstring("application/octet-stream")) + Expect(response.Header().Get("Content-Disposition")).To(BeEmpty()) + Expect(response.Body.String()).NotTo(ContainSubstring("Types of profiles available:")) + Expect(strings.HasPrefix(response.Body.String(), "\x1f\x8b")).To(BeFalse()) + } + } + if !profiling { + Expect(private.Addr()).To(BeNil()) + return + } + Expect(private.Addr().String()).To(Equal(opts.Address)) + client := httpclient.NewWithTimeout(5 * time.Second) + defer client.CloseIdleConnections() + for _, path := range []string{"/debug/pprof/", "/debug/pprof/profile?seconds=1", "/debug/pprof/heap", "/debug/pprof/cmdline", "/debug/pprof/cmdline?x=1", "/debug/pprof/cmdline/"} { + response, err := client.Get("http://" + private.Addr().String() + path) + Expect(err).NotTo(HaveOccurred()) + body, err := io.ReadAll(response.Body) + response.Body.Close() + Expect(err).NotTo(HaveOccurred()) + if strings.Contains(path, "cmdline") { + Expect(response.StatusCode).To(Equal(http.StatusNotFound)) + continue + } + Expect(response.StatusCode).To(Equal(http.StatusOK)) + if path == "/debug/pprof/" { + Expect(string(body)).To(ContainSubstring("Types of profiles available:")) + } else { + Expect(response.Header.Get("Content-Type")).To(Equal("application/octet-stream")) + Expect(strings.HasPrefix(string(body), "\x1f\x8b")).To(BeTrue()) + } + } + }) + } + } + } +}) diff --git a/core/http/middleware/diagnostics.go b/core/http/middleware/diagnostics.go new file mode 100644 index 000000000..84e28a1ad --- /dev/null +++ b/core/http/middleware/diagnostics.go @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: MIT +package middleware + +import ( + "context" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/pkg/diagnostics" +) + +func diagnosticOutcome(ctx context.Context, failed bool) diagnostics.Outcome { + if ctx.Err() != nil { + return diagnostics.OutcomeCanceled + } + if failed { + return diagnostics.OutcomeError + } + return diagnostics.OutcomeOK +} + +// Extraction timings are individual middleware/retry segments, not a single +// socket-to-inference duration. Finish before next; its latency and errors do +// not belong to extraction. The deferred call covers early extraction exits. +func noopExtractionEnd(error) {} + +func (re *RequestExtractor) beginExtraction(c echo.Context) func(error) { + ctx := re.getRequestContext(c) + if !diagnostics.Enabled(ctx) { + return noopExtractionEnd + } + end := diagnostics.Begin(ctx, diagnostics.PhaseExtraction) + finished := false + return func(err error) { + if finished { + return + } + finished = true + end(diagnosticOutcome(ctx, err != nil || c.Response().Status >= 400), 0) + } +} diff --git a/core/http/middleware/diagnostics_test.go b/core/http/middleware/diagnostics_test.go new file mode 100644 index 000000000..ea921a11c --- /dev/null +++ b/core/http/middleware/diagnostics_test.go @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: MIT +package middleware_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + + "github.com/google/uuid" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/failover" + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics extraction", func() { + var re *RequestExtractor + var ac *config.ApplicationConfig + var bcl *config.ModelConfigLoader + var events []diagnostics.Event + BeforeEach(func() { + dir := GinkgoT().TempDir() + ac = config.NewApplicationConfig() + ac.Context = context.Background() + ac.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}} + bcl = config.NewModelConfigLoader(dir) + a := config.ModelConfig{Name: "sentinel-model"} + a.Model = "sentinel-model" + bcl.ReplaceModelConfigs([]config.ModelConfig{a}) + events = nil + ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + re = NewRequestExtractor(bcl, model.NewModelLoader(ac.SystemState), ac) + }) + ends := func(p diagnostics.Phase) []diagnostics.Event { + var out []diagnostics.Event + for _, e := range events { + if e.Phase == p && e.State == diagnostics.StateEnd { + out = append(out, e) + } + } + return out + } + It("preserves anonymous default, JSON, bearer and query selection with timing off and on", func() { + for _, tc := range []struct { + body, query, bearer string + listing, lookup int + }{ + {`{}`, "", "", 1, 0}, {`{"model":"sentinel-model","prompt":"sentinel-prompt"}`, "", "", 1, 0}, + {`{}`, "", "sentinel-model", 0, 1}, {`{}`, "?model=sentinel-model", "sentinel-credential", 0, 0}, + } { + var previous string + recorder := ac.DiagnosticsRecorder + for _, enabled := range []bool{false, true} { + events = nil + ac.DiagnosticsRecorder = nil + if enabled { + ac.DiagnosticsRecorder = recorder + } + e := echo.New() + e.POST("/sentinel-path", func(c echo.Context) error { + Expect(len(ends(diagnostics.PhaseExtraction))).To(Equal(map[bool]int{false: 0, true: 2}[enabled])) + input := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) + Expect(re.SetOpenAIRequest(c)).To(Succeed()) + defer input.Cancel() + Expect(diagnostics.Enabled(input.Context)).To(Equal(enabled)) + diagnostics.Mark(input.Context, diagnostics.PhaseModelInit) + return c.String(200, input.Model) + }, re.BuildFilteredFirstAvailableDefaultModel(nil), re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) })) + req := httptest.NewRequest(http.MethodPost, "/sentinel-path"+tc.query, strings.NewReader(tc.body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Correlation-ID", "sentinel-correlation") + if tc.bearer != "" { + req.Header.Set("Authorization", "Bearer "+tc.bearer) + } + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(200)) + Expect(rec.Body.String()).To(Equal("sentinel-model")) + if !enabled { + previous = rec.Body.String() + Expect(events).To(BeEmpty()) + continue + } + Expect(rec.Body.String()).To(Equal(previous)) + Expect(ends(diagnostics.PhaseDefaultListing)).To(HaveLen(tc.listing)) + Expect(ends(diagnostics.PhaseBearerLookup)).To(HaveLen(tc.lookup)) + Expect(ends(diagnostics.PhaseBodyLookup)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseConfigLoadDefaults)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseConfigFilter)).NotTo(BeEmpty()) + id := events[0].ID + _, err := uuid.Parse(id) + Expect(err).NotTo(HaveOccurred()) + for _, event := range events { + Expect(event.ID).To(Equal(id)) + Expect(event.Kind).To(Equal(diagnostics.KindRequest)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("sentinel")) + } + ac.DiagnosticsRecorder = recorder + } + }) + It("records bind, nil initializer result, missing, disabled and alias failures even when JSON returns nil", func() { + disabled := true + d := config.ModelConfig{Name: "disabled", Disabled: &disabled} + d.Model = "disabled" + a := config.ModelConfig{Name: "alias", Alias: "absent"} + bcl.ReplaceModelConfigs([]config.ModelConfig{d, a}) + for _, tc := range []struct { + body string + nilInput bool + code int + }{ + {`{`, false, 400}, {`{}`, true, 400}, {`{"model":"missing"}`, false, 404}, {`{"model":"disabled"}`, false, 403}, {`{"model":"alias"}`, false, 400}, + } { + events = nil + e := echo.New() + e.POST("/", func(c echo.Context) error { Fail("unexpected inference"); return nil }, re.SetModelAndConfig(func() schema.LocalAIRequest { + if tc.nilInput { + return nil + } + return new(schema.OpenAIRequest) + })) + rec := postJSON(e, "/", tc.body) + Expect(rec.Code).To(Equal(tc.code)) + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseExtraction)[0].Outcome).To(Equal(diagnostics.OutcomeError)) + if strings.Contains(tc.body, "alias") { + Expect(ends(diagnostics.PhaseAliasResolution)[0].Outcome).To(Equal(diagnostics.OutcomeError)) + } + } + }) + It("retains one identity across real failover attempts without default middleware", func() { + primary := config.ModelConfig{Name: "sentinel-primary"} + primary.Model = primary.Name + fallback := config.ModelConfig{Name: "sentinel-fallback"} + fallback.Model = fallback.Name + chain := config.ModelConfig{Name: "sentinel-chain", Failover: &config.FailoverConfig{ + Targets: []config.FailoverTarget{{Model: primary.Name}, {Model: fallback.Name}}, + }} + bcl.ReplaceModelConfigs([]config.ModelConfig{primary, fallback, chain}) + recorder := ac.DiagnosticsRecorder + for _, enabled := range []bool{false, true} { + events = nil + ac.DiagnosticsRecorder = nil + if enabled { + ac.DiagnosticsRecorder = recorder + } + manager := failover.New(bcl) + manager.Sync() + re.SetFailoverManager(manager) + var attempts []string + e := echo.New() + e.POST("/", func(c echo.Context) error { + cfg := c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig) + attempts = append(attempts, cfg.Name) + input := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) + Expect(input.Model).To(Equal(chain.Name)) + Expect(input.Prompt).To(Equal("sentinel-prompt")) + if enabled { + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(len(attempts))) + last := events[len(events)-1] + Expect(last.Phase).To(Equal(diagnostics.PhaseExtraction)) + Expect(last.State).To(Equal(diagnostics.StateEnd)) + Expect(last.Outcome).To(Equal(diagnostics.OutcomeOK)) + } + diagnostics.Mark(c.Request().Context(), diagnostics.PhaseModelInit) + if cfg.Name == primary.Name { + return echo.ErrBadGateway + } + return c.String(http.StatusOK, cfg.Name) + }, re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) })) + response := postJSON(e, "/", `{"model":"sentinel-chain","prompt":"sentinel-prompt"}`) + Expect(attempts).To(Equal([]string{primary.Name, fallback.Name})) + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(response.Body.String()).To(Equal(fallback.Name)) + Expect(response.Header().Get(HeaderServedModel)).To(Equal(fallback.Name)) + Expect(response.Header().Get(HeaderFailover)).To(Equal("fallback")) + if !enabled { + Expect(events).To(BeEmpty()) + continue + } + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(2)) + id := events[0].ID + _, err := uuid.Parse(id) + Expect(err).NotTo(HaveOccurred()) + for _, event := range events { + Expect(event.ID).To(Equal(id)) + Expect(event.Kind).To(Equal(diagnostics.KindRequest)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("sentinel")) + } + }) + It("retains identity across repeated segments and records cancellation before downstream", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(ctx), httptest.NewRecorder()) + h := re.BuildConstantDefaultModelNameMiddleware("sentinel-model")(func(c echo.Context) error { + Expect(ends(diagnostics.PhaseExtraction)).NotTo(BeEmpty()) + return echo.ErrBadGateway + }) + Expect(h(c)).To(Equal(echo.ErrBadGateway)) + Expect(h(c)).To(Equal(echo.ErrBadGateway)) + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(2)) + for _, e := range events { + Expect(e.ID).To(Equal(events[0].ID)) + if e.State == diagnostics.StateEnd { + Expect(e.Outcome).To(Equal(diagnostics.OutcomeCanceled)) + } + } + Expect(c.Request().Context().Err()).To(Equal(context.Canceled)) + }) + It("inherits observation and request cancellation in Responses context rebuilding", func() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(ctx), httptest.NewRecorder()) + Expect(re.BuildConstantDefaultModelNameMiddleware("sentinel-model")(func(echo.Context) error { return nil })(c)).To(Succeed()) + input := &schema.OpenResponsesRequest{Model: "sentinel-model"} + c.Set(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) + cfg := config.ModelConfig{Name: "sentinel-model"} + cfg.Model = "sentinel-model" + c.Set(CONTEXT_LOCALS_KEY_MODEL_CONFIG, &cfg) + Expect(re.SetOpenResponsesRequest(c)).To(Succeed()) + defer input.Cancel() + Expect(diagnostics.Enabled(input.Context)).To(BeTrue()) + diagnostics.Mark(input.Context, diagnostics.PhaseModelInit) + Expect(events[len(events)-1].ID).To(Equal(events[0].ID)) + cancel() + Eventually(input.Context.Done()).Should(BeClosed()) + }) + It("isolates concurrent requests and a background operation on the same recorder", func() { + var mu sync.Mutex + ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { mu.Lock(); defer mu.Unlock(); events = append(events, e) }) + e := echo.New() + e.POST("/", func(c echo.Context) error { + diagnostics.Mark(c.Request().Context(), diagnostics.PhaseModelInit) + return c.NoContent(204) + }, re.BuildFilteredFirstAvailableDefaultModel(nil)) + const n = 12 + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { defer GinkgoRecover(); defer wg.Done(); Expect(postJSON(e, "/", `{}`).Code).To(Equal(204)) }() + } + reload := ac.DiagnosticsRecorder.Reload(context.Background()) + diagnostics.Mark(reload, diagnostics.PhaseReloadEnumeration) + wg.Wait() + groups := map[string][]diagnostics.Event{} + for _, event := range events { + groups[event.ID] = append(groups[event.ID], event) + } + Expect(groups).To(HaveLen(n + 1)) + requests := 0 + for _, group := range groups { + if group[0].Kind == diagnostics.KindConfigReload { + Expect(group).To(HaveLen(1)) + continue + } + requests++ + Expect(group[0].Phase).To(Equal(diagnostics.PhaseExtraction)) + Expect(group[0].State).To(Equal(diagnostics.StateStart)) + Expect(group[len(group)-2].Phase).To(Equal(diagnostics.PhaseExtraction)) + Expect(group[len(group)-2].State).To(Equal(diagnostics.StateEnd)) + Expect(group[len(group)-1].Phase).To(Equal(diagnostics.PhaseModelInit)) + } + Expect(requests).To(Equal(n)) + }) + It("ends an empty default listing before downstream and preserves successful alias selection", func() { + bcl.ReplaceModelConfigs(nil) + c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil), httptest.NewRecorder()) + Expect(re.BuildFilteredFirstAvailableDefaultModel(nil)(func(c echo.Context) error { + Expect(ends(diagnostics.PhaseDefaultListing)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1)) + return nil + })(c)).To(Succeed()) + target := config.ModelConfig{Name: "target"} + target.Model = "target" + alias := config.ModelConfig{Name: "alias", Alias: "target"} + bcl.ReplaceModelConfigs([]config.ModelConfig{target, alias}) + events = nil + e := echo.New() + e.POST("/", func(c echo.Context) error { + Expect(c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest).Model).To(Equal("alias")) + Expect(c.Get(CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig).Name).To(Equal("target")) + Expect(ends(diagnostics.PhaseAliasResolution)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseAliasResolution)[0].Outcome).To(Equal(diagnostics.OutcomeOK)) + return c.NoContent(204) + }, re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) })) + Expect(postJSON(e, "/", `{"model":"alias"}`).Code).To(Equal(204)) + }) + + It("records direct OpenAI and Responses extraction errors", func() { + for _, setter := range []func(echo.Context) error{re.SetOpenAIRequest, re.SetOpenResponsesRequest} { + events = nil + c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil), httptest.NewRecorder()) + Expect(setter(c)).To(Equal(echo.ErrBadRequest)) + Expect(ends(diagnostics.PhaseExtraction)).To(HaveLen(1)) + Expect(ends(diagnostics.PhaseExtraction)[0].Outcome).To(Equal(diagnostics.OutcomeError)) + } + }) + +}) diff --git a/core/http/middleware/request.go b/core/http/middleware/request.go index 4518f446d..e545d4808 100644 --- a/core/http/middleware/request.go +++ b/core/http/middleware/request.go @@ -15,6 +15,7 @@ import ( "github.com/mudler/LocalAI/core/services/failover" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/templates" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/utils" @@ -41,12 +42,27 @@ func NewRequestExtractor(modelConfigLoader *config.ModelConfigLoader, modelLoade } } +// getRequestContext returns the request context, adding optional diagnostics +// once and retaining them on the request for subsequent middleware and retries. +func (re *RequestExtractor) getRequestContext(c echo.Context) context.Context { + ctx := c.Request().Context() + if re.applicationConfig == nil || re.applicationConfig.DiagnosticsRecorder == nil { + return ctx + } + observed := re.applicationConfig.DiagnosticsRecorder.Request(ctx) + if observed != ctx { + c.SetRequest(c.Request().WithContext(observed)) + } + return observed +} + const CONTEXT_LOCALS_KEY_MODEL_NAME = "MODEL_NAME" const CONTEXT_LOCALS_KEY_LOCALAI_REQUEST = "LOCALAI_REQUEST" const CONTEXT_LOCALS_KEY_MODEL_CONFIG = "MODEL_CONFIG" // TODO: Refactor to not return error if unchanged func (re *RequestExtractor) setModelNameFromRequest(c echo.Context) { + ctx := re.getRequestContext(c) model, ok := c.Get(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if ok && model != "" { return @@ -67,7 +83,9 @@ func (re *RequestExtractor) setModelNameFromRequest(c echo.Context) { auth := c.Request().Header.Get("Authorization") bearer := strings.TrimPrefix(auth, "Bearer ") if bearer != "" && bearer != auth { - exists, err := galleryop.CheckIfModelExists(re.modelConfigLoader, re.modelLoader, bearer, galleryop.ALWAYS_INCLUDE) + end := diagnostics.Begin(ctx, diagnostics.PhaseBearerLookup) + exists, err := galleryop.CheckIfModelExistsContext(ctx, re.modelConfigLoader, re.modelLoader, bearer, galleryop.ALWAYS_INCLUDE) + end(diagnosticOutcome(ctx, err != nil), 0) if err == nil && exists { model = bearer } @@ -79,13 +97,16 @@ func (re *RequestExtractor) setModelNameFromRequest(c echo.Context) { func (re *RequestExtractor) BuildConstantDefaultModelNameMiddleware(defaultModelName string) echo.MiddlewareFunc { return func(next echo.HandlerFunc) echo.HandlerFunc { - return func(c echo.Context) error { + return func(c echo.Context) (err error) { + finish := re.beginExtraction(c) + defer func() { finish(err) }() re.setModelNameFromRequest(c) localModelName, ok := c.Get(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if !ok || localModelName == "" { c.Set(CONTEXT_LOCALS_KEY_MODEL_NAME, defaultModelName) xlog.Debug("context local model name not found, setting to default", "defaultModelName", defaultModelName) } + finish(nil) return next(c) } } @@ -93,16 +114,23 @@ func (re *RequestExtractor) BuildConstantDefaultModelNameMiddleware(defaultModel func (re *RequestExtractor) BuildFilteredFirstAvailableDefaultModel(filterFn config.ModelConfigFilterFn) echo.MiddlewareFunc { return func(next echo.HandlerFunc) echo.HandlerFunc { - return func(c echo.Context) error { + return func(c echo.Context) (err error) { + finish := re.beginExtraction(c) + defer func() { finish(err) }() re.setModelNameFromRequest(c) localModelName := c.Get(CONTEXT_LOCALS_KEY_MODEL_NAME).(string) if localModelName != "" { // Don't overwrite existing values + finish(nil) return next(c) } - modelNames, err := galleryop.ListModels(re.modelConfigLoader, re.modelLoader, filterFn, galleryop.SKIP_IF_CONFIGURED) + ctx := c.Request().Context() + end := diagnostics.Begin(ctx, diagnostics.PhaseDefaultListing) + modelNames, err := galleryop.ListModelsContext(ctx, re.modelConfigLoader, re.modelLoader, filterFn, galleryop.SKIP_IF_CONFIGURED) + end(diagnosticOutcome(ctx, err != nil), len(modelNames)) if err != nil { xlog.Error("non-fatal error calling ListModels during SetDefaultModelNameToFirstAvailable()", "error", err) + finish(err) return next(c) } @@ -110,11 +138,13 @@ func (re *RequestExtractor) BuildFilteredFirstAvailableDefaultModel(filterFn con xlog.Warn("SetDefaultModelNameToFirstAvailable used with no matching models installed") // This is non-fatal - making it so was breaking the case of direct installation of raw models // return errors.New("this endpoint requires at least one model to be installed") + finish(nil) return next(c) } c.Set(CONTEXT_LOCALS_KEY_MODEL_NAME, modelNames[0]) xlog.Debug("context local model name not found, setting to the first model", "first model name", modelNames[0]) + finish(nil) return next(c) } } @@ -124,7 +154,9 @@ func (re *RequestExtractor) BuildFilteredFirstAvailableDefaultModel(filterFn con // Otherwise, it's in its own method below for now func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIRequest) echo.MiddlewareFunc { return func(next echo.HandlerFunc) echo.HandlerFunc { - return re.failoverRetry(func(c echo.Context) error { + retry := re.failoverRetry(func(c echo.Context) (err error) { + finish := re.beginExtraction(c) + defer func() { finish(err) }() input := initializer() if input == nil { return echo.NewHTTPError(http.StatusBadRequest, "unable to initialize body") @@ -149,7 +181,10 @@ func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIR if strings.HasSuffix(modelName, ":latest") { modelName = strings.TrimSuffix(modelName, ":latest") } + ctx := c.Request().Context() + endConfig := diagnostics.Begin(ctx, diagnostics.PhaseConfigLoadDefaults) cfg, err := re.modelConfigLoader.LoadModelConfigFileByNameDefaultOptions(modelName, re.applicationConfig) + endConfig(diagnosticOutcome(ctx, err != nil), 0) if err != nil { xlog.Warn("Model Configuration File not found", "model", modelName, "error", err) @@ -163,7 +198,9 @@ func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIR // Skip the check for HuggingFace model IDs (contain "/") since backends // like diffusers may download these on the fly. if modelName != "" && !strings.Contains(modelName, "/") { - exists, existsErr := galleryop.CheckIfModelExists(re.modelConfigLoader, re.modelLoader, modelName, galleryop.ALWAYS_INCLUDE) + endLookup := diagnostics.Begin(ctx, diagnostics.PhaseBodyLookup) + exists, existsErr := galleryop.CheckIfModelExistsContext(ctx, re.modelConfigLoader, re.modelLoader, modelName, galleryop.ALWAYS_INCLUDE) + endLookup(diagnosticOutcome(ctx, existsErr != nil || !exists), 0) if existsErr == nil && !exists { return c.JSON(http.StatusNotFound, schema.ErrorResponse{ Error: &schema.APIError{ @@ -181,7 +218,9 @@ func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIR // the alias name (input.ModelName is left unchanged); usage accounting // records requested=alias / served=target. if cfg != nil && cfg.IsAlias() { + endAlias := diagnostics.Begin(ctx, diagnostics.PhaseAliasResolution) resolved, _, aliasErr := re.modelConfigLoader.ResolveAlias(cfg) + endAlias(diagnosticOutcome(ctx, aliasErr != nil), 0) if aliasErr != nil { return c.JSON(http.StatusBadRequest, schema.ErrorResponse{ Error: &schema.APIError{ @@ -228,12 +267,21 @@ func (re *RequestExtractor) SetModelAndConfig(initializer func() schema.LocalAIR c.Set(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) c.Set(CONTEXT_LOCALS_KEY_MODEL_CONFIG, cfg) + finish(nil) return next(c) }) + return func(c echo.Context) error { + // Retries restore the entry request, so attach its observation before + // failover captures it. Extraction spans remain per attempt above. + re.getRequestContext(c) + return retry(c) + } } } -func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) error { +func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) (err error) { + finish := re.beginExtraction(c) + defer func() { finish(err) }() input, ok := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest) if !ok || input.Model == "" { return echo.ErrBadRequest @@ -253,7 +301,7 @@ func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) error { // Use the request context directly - Echo properly supports context cancellation! // No need for workarounds like handleConnectionCancellation - reqCtx := c.Request().Context() + reqCtx := re.getRequestContext(c) c1, cancel := context.WithCancel(re.applicationConfig.Context) // Cancel when request context is cancelled (client disconnects) @@ -269,6 +317,7 @@ func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) error { // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID) ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) + ctxWithCorrelationID = diagnostics.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel @@ -278,7 +327,7 @@ func (re *RequestExtractor) SetOpenAIRequest(c echo.Context) error { return nil } - err := mergeOpenAIRequestAndModelConfig(cfg, input) + err = mergeOpenAIRequestAndModelConfig(cfg, input) if err != nil { return err } @@ -698,7 +747,9 @@ func mergeOpenAIRequestAndModelConfig(config *config.ModelConfig, input *schema. return nil } -func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { +func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) (err error) { + finish := re.beginExtraction(c) + defer func() { finish(err) }() input, ok := c.Get(CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenResponsesRequest) if !ok || input.Model == "" { return echo.ErrBadRequest @@ -719,7 +770,7 @@ func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { c.Response().Header().Set("x-request-id", correlationID) // Use the request context directly - Echo properly supports context cancellation! - reqCtx := c.Request().Context() + reqCtx := re.getRequestContext(c) c1, cancel := context.WithCancel(re.applicationConfig.Context) // Cancel when request context is cancelled (client disconnects) @@ -735,11 +786,12 @@ func (re *RequestExtractor) SetOpenResponsesRequest(c echo.Context) error { // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, CorrelationIDKey, correlationID) ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) + ctxWithCorrelationID = diagnostics.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel - err := MergeOpenResponsesConfig(cfg, input) + err = MergeOpenResponsesConfig(cfg, input) if err != nil { return err } diff --git a/core/http/routes/anthropic.go b/core/http/routes/anthropic.go index a9e6a48ef..dd5916e89 100644 --- a/core/http/routes/anthropic.go +++ b/core/http/routes/anthropic.go @@ -16,6 +16,7 @@ import ( "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/piiadapter" "github.com/mudler/LocalAI/core/services/routing/router" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/xlog" ) @@ -104,6 +105,7 @@ func setAnthropicRequestContext(appConfig *config.ApplicationConfig) echo.Middle // Add the correlation ID to the new context ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) + ctxWithCorrelationID = diagnostics.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel diff --git a/core/http/routes/diagnostics_context_test.go b/core/http/routes/diagnostics_context_test.go new file mode 100644 index 000000000..3a2ce0d37 --- /dev/null +++ b/core/http/routes/diagnostics_context_test.go @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: MIT +package routes + +import ( + "context" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/middleware" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/diagnostics" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "net/http/httptest" +) + +var _ = Describe("Diagnostics route context reconstruction", func() { + It("retains identity and cancellation through Ollama chat/generate and Anthropic middleware", func() { + for _, kind := range []string{"chat", "generate", "anthropic"} { + ac := config.NewApplicationConfig() + ac.Context = context.Background() + var events []diagnostics.Event + recorder := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ctx = recorder.Request(ctx) + diagnostics.Mark(ctx, diagnostics.PhaseExtraction) + c := echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(ctx), httptest.NewRecorder()) + c.Request().Header.Set("x-request-id", "sentinel-client-id") + cfg := config.ModelConfig{Name: "model"} + cfg.Model = "model" + c.Set(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG, &cfg) + var mw echo.MiddlewareFunc + var getContext func() context.Context + switch kind { + case "chat": + input := &schema.OllamaChatRequest{Model: "model"} + c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) + getContext = func() context.Context { return input.Context } + mw = setOllamaChatRequestContext(ac) + case "generate": + input := &schema.OllamaGenerateRequest{Model: "model"} + c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) + getContext = func() context.Context { return input.Ctx } + mw = setOllamaGenerateRequestContext(ac) + default: + input := &schema.AnthropicRequest{Model: "model", MaxTokens: 1} + c.Set(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST, input) + getContext = func() context.Context { return input.Context } + mw = setAnthropicRequestContext(ac) + } + Expect(mw(func(echo.Context) error { + out := getContext() + Expect(diagnostics.Enabled(out)).To(BeTrue()) + Expect(out.Err()).NotTo(HaveOccurred()) + diagnostics.Mark(out, diagnostics.PhaseModelInit) + cancel() + Eventually(out.Done()).Should(BeClosed()) + return nil + })(c)).To(Succeed()) + Expect(events).To(HaveLen(2)) + Expect(events[1].ID).To(Equal(events[0].ID)) + Expect(events[0].ID).NotTo(Equal("sentinel-client-id")) + } + }) +}) diff --git a/core/http/routes/ollama.go b/core/http/routes/ollama.go index 6e8d97612..8295ee38a 100644 --- a/core/http/routes/ollama.go +++ b/core/http/routes/ollama.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/piiadapter" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/distributedhdr" ) @@ -119,6 +120,7 @@ func setOllamaChatRequestContext(appConfig *config.ApplicationConfig) echo.Middl ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) + ctxWithCorrelationID = diagnostics.Inherit(ctxWithCorrelationID, reqCtx) input.Context = ctxWithCorrelationID input.Cancel = cancel @@ -161,6 +163,7 @@ func setOllamaGenerateRequestContext(appConfig *config.ApplicationConfig) echo.M ctxWithCorrelationID := context.WithValue(c1, middleware.CorrelationIDKey, correlationID) ctxWithCorrelationID = distributedhdr.Inherit(ctxWithCorrelationID, reqCtx) + ctxWithCorrelationID = diagnostics.Inherit(ctxWithCorrelationID, reqCtx) input.Ctx = ctxWithCorrelationID input.Cancel = cancel diff --git a/core/services/galleryop/list_models.go b/core/services/galleryop/list_models.go index 47c2d91e5..a9b1a56fe 100644 --- a/core/services/galleryop/list_models.go +++ b/core/services/galleryop/list_models.go @@ -1,7 +1,9 @@ package galleryop import ( + "context" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/diagnostics" "github.com/mudler/LocalAI/pkg/model" ) @@ -14,7 +16,22 @@ const ( ALWAYS_INCLUDE ) +// discoveryFiles is a call-local seam for verifying discovery operation counts. +type discoveryFiles interface { + ListFilesInModelPathContext(context.Context) ([]string, error) + ExistsInModelPath(string) bool +} + func ListModels(bcl *config.ModelConfigLoader, ml *model.ModelLoader, filter config.ModelConfigFilterFn, looseFilePolicy LooseFilePolicy) ([]string, error) { + return ListModelsContext(context.Background(), bcl, ml, filter, looseFilePolicy) +} + +// ListModelsContext carries diagnostics through config and loose-file discovery. +func ListModelsContext(ctx context.Context, bcl *config.ModelConfigLoader, ml *model.ModelLoader, filter config.ModelConfigFilterFn, looseFilePolicy LooseFilePolicy) ([]string, error) { + return listModelsContext(ctx, bcl, ml, filter, looseFilePolicy) +} + +func listModelsContext(ctx context.Context, bcl *config.ModelConfigLoader, ml discoveryFiles, filter config.ModelConfigFilterFn, looseFilePolicy LooseFilePolicy) ([]string, error) { // Callers (e.g. the Ollama /api/tags handler) pass nil to mean "no // filtering". Without this guard the loose-file loop below dereferences @@ -30,7 +47,9 @@ func ListModels(bcl *config.ModelConfigLoader, ml *model.ModelLoader, filter con // Start with known configurations - for _, c := range bcl.GetModelConfigsByFilter(filter) { + configs := bcl.GetModelConfigsByFilterContext(ctx, filter) + end := diagnostics.Begin(ctx, diagnostics.PhaseLooseFilter) + for _, c := range configs { // Is this better than looseFilePolicy <= SKIP_IF_CONFIGURED ? less performant but more readable? if (looseFilePolicy == SKIP_IF_CONFIGURED) || (looseFilePolicy == LOOSE_ONLY) { skipMap[c.Model] = struct{}{} @@ -40,30 +59,43 @@ func ListModels(bcl *config.ModelConfigLoader, ml *model.ModelLoader, filter con } } + end(diagnostics.OutcomeOK, len(configs)) + // Then iterate through the loose files if requested. if looseFilePolicy != SKIP_ALWAYS { - models, err := ml.ListFilesInModelPath() + models, err := ml.ListFilesInModelPathContext(ctx) if err != nil { return nil, err } + end = diagnostics.Begin(ctx, diagnostics.PhaseLooseFilter) for _, m := range models { // And only adds them if they shouldn't be skipped. if _, exists := skipMap[m]; !exists && filter(m, nil) { dataModels = append(dataModels, m) } } + end(diagnostics.OutcomeOK, len(models)) } return dataModels, nil } func CheckIfModelExists(bcl *config.ModelConfigLoader, ml *model.ModelLoader, modelName string, looseFilePolicy LooseFilePolicy) (bool, error) { + return CheckIfModelExistsContext(context.Background(), bcl, ml, modelName, looseFilePolicy) +} + +// CheckIfModelExistsContext observes the existing lookup and weight-file fallback. +func CheckIfModelExistsContext(ctx context.Context, bcl *config.ModelConfigLoader, ml *model.ModelLoader, modelName string, looseFilePolicy LooseFilePolicy) (bool, error) { + return checkIfModelExistsContext(ctx, bcl, ml, modelName, looseFilePolicy) +} + +func checkIfModelExistsContext(ctx context.Context, bcl *config.ModelConfigLoader, ml discoveryFiles, modelName string, looseFilePolicy LooseFilePolicy) (bool, error) { filter, err := config.BuildNameFilterFn(modelName) if err != nil { return false, err } - models, err := ListModels(bcl, ml, filter, looseFilePolicy) + models, err := listModelsContext(ctx, bcl, ml, filter, looseFilePolicy) if err != nil { return false, err } @@ -74,7 +106,11 @@ func CheckIfModelExists(bcl *config.ModelConfigLoader, ml *model.ModelLoader, mo // ListModels may not find raw model weight files (e.g. .ggml, .gguf) // because ListFilesInModelPath skips known weight-file extensions. // Fall back to checking if the file exists directly in the model path. - if ml.ExistsInModelPath(modelName) { + end := diagnostics.Begin(ctx, diagnostics.PhaseExistenceFallback) + exists := ml.ExistsInModelPath(modelName) + // ExistsInModelPath exposes only a boolean, not an underlying filesystem error. + end(diagnostics.OutcomeOK, 1) + if exists { return true, nil } diff --git a/core/services/galleryop/list_models_diagnostics_test.go b/core/services/galleryop/list_models_diagnostics_test.go new file mode 100644 index 000000000..f987d7320 --- /dev/null +++ b/core/services/galleryop/list_models_diagnostics_test.go @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: MIT +package galleryop + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/LocalAI/pkg/model" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type diagnosticDiscovery struct { + ml *model.ModelLoader + lists, exists int +} + +func (d *diagnosticDiscovery) ListFilesInModelPathContext(ctx context.Context) ([]string, error) { + d.lists++ + return d.ml.ListFilesInModelPathContext(ctx) +} +func (d *diagnosticDiscovery) ExistsInModelPath(name string) bool { + d.exists++ + return d.ml.ExistsInModelPath(name) +} + +var _ = Describe("Diagnostics discovery helpers", func() { + var bcl *config.ModelConfigLoader + var ml *model.ModelLoader + var ctx context.Context + var events []diagnostics.Event + BeforeEach(func() { + dir := GinkgoT().TempDir() + ml = &model.ModelLoader{ModelPath: dir} + bcl = config.NewModelConfigLoader(dir) + for _, name := range []string{"secret-candidate", "z-loose", "secret-weight.gguf"} { + Expect(os.WriteFile(filepath.Join(dir, name), nil, 0600)).To(Succeed()) + } + c := config.ModelConfig{Name: "secret-candidate"} + c.Model = "secret-candidate" + bcl.ReplaceModelConfigs([]config.ModelConfig{c}) + events = nil + ctx = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + }) + It("preserves every policy, nil filters, duplicates, ordering and operation counts", func() { + expected := [][]string{{"z-loose"}, {"secret-candidate", "z-loose"}, {"secret-candidate"}, {"secret-candidate", "secret-candidate", "z-loose"}} + for policy := LOOSE_ONLY; policy <= ALWAYS_INCLUDE; policy++ { + events = nil + plainD, observedD := &diagnosticDiscovery{ml: ml}, &diagnosticDiscovery{ml: ml} + plain, err := listModelsContext(context.Background(), bcl, plainD, nil, policy) + Expect(err).NotTo(HaveOccurred()) + got, err := listModelsContext(ctx, bcl, observedD, nil, policy) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal(plain)) + Expect(got).To(Equal(expected[policy])) + Expect(observedD.lists).To(Equal(plainD.lists)) + Expect(observedD.exists).To(BeZero()) + if policy == SKIP_ALWAYS { + Expect(observedD.lists).To(BeZero()) + } else { + Expect(observedD.lists).To(Equal(1)) + } + var looseCounts []int + fsCount := 0 + configCount := 0 + for _, e := range events { + if e.State == diagnostics.StateEnd { + switch e.Phase { + case diagnostics.PhaseLooseFilter: + looseCounts = append(looseCounts, e.Count) + case diagnostics.PhaseFSEnumeration: + fsCount += e.Count + case diagnostics.PhaseConfigFilter: + configCount++ + } + } + Expect(e.ID).To(Equal(events[0].ID)) + Expect(e.Kind).To(Equal(diagnostics.KindRequest)) + } + Expect(configCount).To(Equal(1)) + if policy == SKIP_ALWAYS { + Expect(looseCounts).To(Equal([]int{1})) + Expect(fsCount).To(BeZero()) + } else { + Expect(looseCounts).To(Equal([]int{1, 3, 2})) + Expect(fsCount).To(Equal(1)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("secret")) + Expect(fmt.Sprint(events)).NotTo(ContainSubstring(ml.ModelPath)) + n := len(events) + wrapped, err := ListModels(bcl, ml, nil, policy) + Expect(err).NotTo(HaveOccurred()) + Expect(wrapped).To(Equal(got)) + Expect(events).To(HaveLen(n)) + companion, err := ListModelsContext(ctx, bcl, ml, nil, policy) + Expect(err).NotTo(HaveOccurred()) + Expect(companion).To(Equal(got)) + } + }) + It("preserves invalid regex, early matches, weight fallback and negative fallback for all policies", func() { + for policy := LOOSE_ONLY; policy <= ALWAYS_INCLUDE; policy++ { + for _, name := range []string{"[secret-invalid", "z-loose", "secret-weight.gguf", "secret-absent"} { + events = nil + a, b := &diagnosticDiscovery{ml: ml}, &diagnosticDiscovery{ml: ml} + plain, plainErr := checkIfModelExistsContext(context.Background(), bcl, a, name, policy) + got, err := checkIfModelExistsContext(ctx, bcl, b, name, policy) + Expect(got).To(Equal(plain)) + if plainErr != nil { + Expect(err).To(Equal(plainErr)) + } else { + Expect(err).NotTo(HaveOccurred()) + } + Expect(b.lists).To(Equal(a.lists)) + Expect(b.exists).To(Equal(a.exists)) + if name == "[secret-invalid" { + Expect(err).To(HaveOccurred()) + Expect(b.lists).To(BeZero()) + Expect(b.exists).To(BeZero()) + Expect(events).To(BeEmpty()) + continue + } + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal(name != "secret-absent")) + if name == "secret-weight.gguf" || name == "secret-absent" || policy == SKIP_ALWAYS { + Expect(b.exists).To(Equal(1)) + } else { + Expect(b.exists).To(BeZero()) + } + fallbackEnds := 0 + for _, e := range events { + Expect(e.ID).To(Equal(events[0].ID)) + if e.Phase == diagnostics.PhaseExistenceFallback && e.State == diagnostics.StateEnd { + fallbackEnds++ + Expect(e.Count).To(Equal(1)) + Expect(e.Outcome).To(Equal(diagnostics.OutcomeOK)) + } + } + Expect(fallbackEnds).To(Equal(b.exists)) + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("secret")) + n := len(events) + wrapped, err := CheckIfModelExists(bcl, ml, name, policy) + Expect(err).NotTo(HaveOccurred()) + Expect(wrapped).To(Equal(got)) + Expect(events).To(HaveLen(n)) + companion, err := CheckIfModelExistsContext(ctx, bcl, ml, name, policy) + Expect(err).NotTo(HaveOccurred()) + Expect(companion).To(Equal(got)) + } + } + }) + It("propagates enumeration errors without fallback and skips enumeration only when requested", func() { + ml.ModelPath = filepath.Join(ml.ModelPath, "secret-missing") + d := &diagnosticDiscovery{ml: ml} + found, err := checkIfModelExistsContext(ctx, bcl, d, "secret-absent", ALWAYS_INCLUDE) + Expect(found).To(BeFalse()) + Expect(os.IsNotExist(err)).To(BeTrue()) + Expect(d.lists).To(Equal(1)) + Expect(d.exists).To(BeZero()) + Expect(events[len(events)-1].Outcome).To(Equal(diagnostics.OutcomeError)) + got, err := ListModelsContext(ctx, bcl, ml, nil, SKIP_ALWAYS) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal([]string{"secret-candidate"})) + }) + It("retains canceled-context behavior and filter invocation counts", func() { + canceled, cancel := context.WithCancel(ctx) + cancel() + calls := 0 + filter := func(_ string, _ *config.ModelConfig) bool { calls++; return true } + plain, err := ListModels(bcl, ml, filter, ALWAYS_INCLUDE) + Expect(err).NotTo(HaveOccurred()) + before := calls + got, err := ListModelsContext(canceled, bcl, ml, filter, ALWAYS_INCLUDE) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal(plain)) + Expect(calls).To(Equal(before * 2)) + }) +}) diff --git a/core/services/modeladmin/remote_sync.go b/core/services/modeladmin/remote_sync.go index 9a51d21b3..8dd98706d 100644 --- a/core/services/modeladmin/remote_sync.go +++ b/core/services/modeladmin/remote_sync.go @@ -28,7 +28,7 @@ func ApplyRemoteChange(ctx context.Context, cl *config.ModelConfigLoader, models } func applyRemoteChange(ctx context.Context, cl *config.ModelConfigLoader, modelsPath string, evt messaging.CacheInvalidateEvent, lifecycle ModelRevisionLifecycle, opts ...config.ConfigLoaderOption) error { - authoritative := config.NewModelConfigLoader(modelsPath) + authoritative := config.NewModelConfigLoader(modelsPath, cl.ReloadDiagnosticsOption()) if err := authoritative.LoadModelConfigsFromPathStrict(modelsPath, opts...); err != nil { return err } diff --git a/docs/content/features/diagnostics.md b/docs/content/features/diagnostics.md new file mode 100644 index 000000000..05630e8ac --- /dev/null +++ b/docs/content/features/diagnostics.md @@ -0,0 +1,185 @@ ++++ +disableToc = false +title = "Profiling and request diagnostics" +weight = 30 ++++ + +LocalAI provides two independently opt-in diagnostics: a private Go profiling listener and structured phase timing events. +Both default to disabled. Debug logging alone enables neither feature. + +## Scope and configuration + +These options apply to `local-ai run`, including a distributed frontend. +They do not enable diagnostics in separately launched workers or provide backend-native llama.cpp profiling. +Diagnostics has no MCP, admin, or public API surface and no UI controls. +State is per-instance, not shared between replicas. Profiles describe the frontend Go process, not the whole cluster. + +| Parameter | Default | Description | Environment Variable | +|-----------|---------|-------------|----------------------| +| `--pprof` | `false` | Enable the private loopback profiling listener | `$LOCALAI_PPROF` | +| `--pprof-address` | `127.0.0.1:6060` | Numeric loopback IP with an explicit port | `$LOCALAI_PPROF_ADDRESS` | +| `--pprof-mutex-profile-fraction` | `0` | Nonnegative mutex sampling fraction; zero disables sampling | `$LOCALAI_PPROF_MUTEX_PROFILE_FRACTION` | +| `--pprof-block-profile-rate` | `0` | Nonnegative block sampling rate in nanoseconds; zero disables sampling | `$LOCALAI_PPROF_BLOCK_PROFILE_RATE` | +| `--request-phase-timing` | `false` | Emit sanitized phase events at Info, independently of profiling | `$LOCALAI_REQUEST_PHASE_TIMING` | + +The address accepts numeric IPv4 loopback or bracketed IPv6 loopback, such as `127.0.0.1:6060` or `[::1]:6060`. +Ports must be decimal values from 1–65535. +Validation rejects hostnames, including `localhost`, wildcard addresses, empty hosts, LAN addresses, public addresses, and Unix sockets. +Validation applies even when profiling is disabled. Negative sampling settings fail validation; nonzero sampling settings require `--pprof`. +`--preload-backend-only` validates the options but does not start profiling or sampling. + +For mutex sampling, a value of N records approximately one in N contention events; 1 records every event. +Block sampling targets one blocking event per configured number of nanoseconds spent blocked; 1 records every blocking event. +These settings are not percentages or capture durations. Sampling adds process-wide overhead, especially at 1. + +One diagnostics server owns process-global runtime sampling in each run process. +It applies both configured rates once after binding successfully and resets both rates to zero on shutdown or fatal serving failure. +It does not restore previous rates; Go provides no getter for the previous block rate. +Do not combine competing sampling owners in the same process. +There is no dynamic toggle, continuous capture, or automatic profile storage. Configuration changes require a separately approved restart or rollout. + +The listener binds before normal request serving. Validation or bind errors fail startup; subsequent fatal serving errors propagate through the run lifecycle. +The private listener gets five seconds to drain before profiling connections are force-closed, which can interrupt an active capture. +Its cleanup runs on startup failure, normal return, and through the existing signal handler (which exits without running deferred cleanup). +A fatal profiler serving error cancels the run context and is returned after application cleanup. +Application initialization or shutdown work that ignores cancellation can delay that return; the private-listener budget is not a deadline for the entire application. +With profiling disabled, diagnostics opens no socket, starts no monitoring goroutine, and changes no runtime sampling settings. +Disabled phase timing allocates no diagnostic identities or events and adds no diagnostic filesystem operations. + +## Access and data sensitivity + +{{% notice warning %}} +The profiling listener is unauthenticated. Loopback reachability is its authorization boundary, not the public API's authentication policy. +Other processes in the same network namespace can access it, including containers sharing a pod's namespace. +Profiles can contain sensitive application data. Heap data and goroutine output can expose secrets. +Protect captures, review them before sharing, and delete them when the investigation ends. +{{% /notice %}} + +The separate listener uses a private mux, not the public API router or the default Go mux. +It provides the profiling index, CPU, trace, symbol, and named runtime profiles, including heap, goroutine, mutex, and block. +`/debug/pprof/cmdline` is unavailable, including query and subtree variants, because command-line arguments can contain credentials. +The standard index can still mention that endpoint. + +Timing events omit credentials, bearer candidates, model names, prompts, bodies, raw headers, filenames, paths, and raw errors. +This sanitation applies only to diagnostic timing payloads. It does not redact unrelated logs or profiles. +Events use the constant message `diagnostic_phase` at Info. +Operator-selected warn/error filtering suppresses these events, even when phase timing is enabled. + +## Bounded Kubernetes investigation + +A rollout or build configuration change is a separate, explicit, owner-approved action. +Select one frontend and a defined observation window. Approve profiling, timing, and any nonzero sampling settings before the window. +Do not change replicas, clocks, scheduling, or storage mounts to force a result. +The examples below capture data only; they do not mutate a live deployment. + +Your cluster runtime must support port-forwarding to the pod's loopback listener. +Do not add a Service, ingress, or NodePort for profiling. +Forward directly to the selected frontend pod, not a load-balanced Service. +The local port must be free. Keep kubectl's default local-only binding; do not add a wildcard `--address`. + +1. Start the direct pod port-forward in a separate terminal, replacing both placeholders: + + ```sh + kubectl -n port-forward pod/ 6060:6060 + ``` + +2. Protect newly created files in your capture terminal: + + ```sh + umask 077 + ``` + +3. Capture ten seconds of CPU activity during the approved window: + + ```sh + curl --fail --output cpu.pprof 'http://127.0.0.1:6060/debug/pprof/profile?seconds=10' + go tool pprof cpu.pprof + ``` + +CPU profiles show execution, not time blocked on I/O or locks. +Capture naturally slow requests and matched control requests against the same frontend. +Record the observation window and generated diagnostic IDs alongside protected captures. +Profiles are process-wide; they do not automatically attribute samples to a request ID. + +Mutex and block profiles require their respective nonzero sampling settings before capture. +The following ten-second delta captures bound collection time, not the duration of process-global sampling. +Sequential commands cover different windows; compare them accordingly. + +```sh +curl --fail --max-time 20 --output mutex.pprof 'http://127.0.0.1:6060/debug/pprof/mutex?seconds=10' +curl --fail --max-time 20 --output block.pprof 'http://127.0.0.1:6060/debug/pprof/block?seconds=10' +go tool pprof mutex.pprof +go tool pprof block.pprof +``` + +Mutex profiles locate contention attributed to lock holders; block profiles locate sampled blocking operations. +Neither substitutes for the request-owned lock and filesystem timings below. +A goroutine snapshot can show I/O or lock waits that CPU profiles miss. +It describes one instant, not a measured duration. Its text can be large and sensitive; the client timeout bounds waiting, not file size. + +```sh +curl --fail --max-time 10 --output goroutines.txt 'http://127.0.0.1:6060/debug/pprof/goroutine?debug=2' +``` + +Stop the port-forward after capture. Disable diagnostics and sampling afterward through the separately approved rollout process. +Protect, review, and delete the capture files under your data-handling policy. +If latency does not reproduce, report non-reproduction. Do not infer a predetermined filesystem, lock, or routing root cause. + +## Event identity and phase boundaries + +Each event contains `id`, `kind`, `phase`, `state`, `outcome`, `elapsed`, and `count`. +The server generates diagnostic UUIDs; it does not trust or reuse client `X-Correlation-ID` headers as diagnostic identities. +Requests use kind `request`; config reload operations use kind `config_reload` with fresh operation IDs. +Request identity survives the instrumented internal context reconstructions. Background callers remain unobserved unless explicitly enabled. + +States are `start`, `end`, and `mark`; outcomes are `ok`, `error`, `canceled`, and `skipped`. +Durations use monotonic elapsed time, not cross-host clock subtraction. +Counts describe aggregate work, not filenames or model labels. +Lock-protected helpers emit buffered aggregates after unlocking, so log order alone does not establish execution order. + +Nested timings overlap and are not additive. Repeated extraction segments share one request ID; they are not duplicate whole-request measurements. +Extraction ends before downstream inference. It is not socket arrival or full request latency. +The existing trace start follows global authentication and body reading; it is not socket arrival either. + +| Phase | Measured boundary | +|-------|-------------------| +| `extraction` | Each extractor-owned entry/exit segment, including repeated extraction or retries, before downstream inference. | +| `bearer_lookup` | Existing bearer-derived model existence lookup, without recording the candidate. | +| `default_listing` | Existing default-model candidate listing within extraction. | +| `body_lookup` | Existing body-derived model existence lookup, without recording body data. | +| `config_load_defaults` | Existing config loading and default application within extraction. | +| `alias_resolution` | Existing model alias resolution within extraction. | +| `config_lock_wait` | Time waiting to acquire the config loader mutex for filtering. | +| `config_lock_hold` | Time holding that mutex for the config listing/filter operation. Includes nested filter work. | +| `config_filter` | The in-memory config filter pass under that lock, not filesystem latency. | +| `fs_enumeration` | The existing `os.ReadDir` call for loose model discovery only. | +| `loose_filter` | Existing filename and in-memory loose-file filtering passes, not directory enumeration. | +| `existence_fallback` | The existing `ExistsInModelPath` fallback when discovery requires it. No extra diagnostic filesystem calls. | +| `model_init` | Entry mark immediately after `ModelLoader.Load` constructs options, before admission, cache, or load logic. Not backend-native execution. | +| `model_router_callback` | Only the existing `ModelRouter` callback invocation. Not transport publish or worker receive latency. | +| `reload_lock_wait` | Time waiting for the mutex of the loader performing the reload. | +| `reload_lock_hold` | Time holding that loader's mutex during reload, including nested enumeration, metadata, and later parsing work. | +| `reload_enumeration` | The reload's existing `os.ReadDir` call. | +| `reload_metadata` | Aggregate existing `DirEntry.Info` calls, separate from enumeration. | +| `reload_yaml_read` | Aggregate YAML file reads, including configured-gallery YAML inspection where applicable. | +| `reload_parse_defaults` | Aggregate YAML parsing/default application, including configured-gallery YAML classification where applicable. Excludes the measured file read. | + +### Reload observer boundaries + +Reload timing measures `ModelConfigLoader.loadModelConfigsFromPath`, not the unrelated dynamic JSON watcher. +Distributed resync constructs a temporary authoritative loader with its own mutex. +That loader's reload events do not prove contention on the shared loader lock or ownership of another request's wait. +No event establishes cross-host timing or worker execution. + +The immutable constructor option `WithReloadDiagnostics` controls early lock, enumeration, and metadata timing. +Ordinary `LoadOptionDiagnostics` callbacks run at their original position after enumeration and metadata, under the loader lock. +Their resolved recorder controls later YAML-read and parse/default timing. +An option-only caller cannot observe earlier phases. A late nil recorder disables subsequent instrumented work, not constructor-enabled early measurements. + +Unchanged recorder identity preserves the reload operation ID across early and later phases. +A different non-nil late recorder receives a separate reload operation ID; these are separate observations, not one correlated span. +Callbacks retain ordering, composition, and last-write semantics. +They retain every preexisting invocation, including each parsed model's `ModelConfig.SetDefaults`; diagnostics does not add a replay or snapshot functional options. + +See the [CLI reference]({{% relref "reference/cli-reference" %}}) for other server options. +The [troubleshooting guide]({{% relref "getting-started/troubleshooting" %}}) covers broader operational checks. diff --git a/docs/content/getting-started/troubleshooting.md b/docs/content/getting-started/troubleshooting.md index 3527ac9d9..999a8e23c 100644 --- a/docs/content/getting-started/troubleshooting.md +++ b/docs/content/getting-started/troubleshooting.md @@ -41,6 +41,15 @@ docker ps -a | grep local-ai docker run --rm --gpus all nvidia/cuda:12.8.0-base-ubuntu24.04 nvidia-smi ``` +## Request latency diagnostics + +Debug logging alone does not enable profiling or request phase timing. +Use [Profiling and request diagnostics]({{% relref "features/diagnostics" %}}) for an opt-in, bounded investigation on one frontend. +The guide distinguishes CPU execution, lock waits, discovery filesystem work, and the model router callback. +It includes direct pod port-forward instructions and warnings about unauthenticated, sensitive profiles. +Enable diagnostics only through a separately approved rollout; do not expose profiling through the public API. +Report non-reproduction if matched requests do not reproduce the delay. + ## Installation Issues ### Binary Won't Execute on Linux diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 31bb79013..5d2cc669b 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -23,6 +23,33 @@ model, see [Model Configuration]({{% relref "advanced/model-configuration" %}}). | `--credentials-file` | | YAML file with credentials for private registries, galleries and download hosts. See [Private Registries and Galleries]({{% relref "advanced/private-sources" %}}) | `$LOCALAI_CREDENTIALS_FILE` | | `--debug` | `false` | **DEPRECATED** - Use `--log-level=debug` instead. Enable debug logging | `$LOCALAI_DEBUG`, `$DEBUG` | +## Diagnostics flags + +See [Profiling and request diagnostics]({{% relref "features/diagnostics" %}}) for sampling semantics, phase boundaries, and protected pod captures. + +These options apply only to `local-ai run`, including distributed frontends. +Diagnostics state is per-instance, not shared between replicas. Debug logging enables neither feature. +The profiling listener is separate from the public API and accepts only numeric loopback addresses with explicit ports. +`--preload-backend-only` validates these options but does not start profiling or sampling. + +| Parameter | Default | Description | Environment Variable | +|-----------|---------|-------------|----------------------| +| `--pprof` | `false` | Enable the private loopback profiling listener | `$LOCALAI_PPROF` | +| `--pprof-address` | `127.0.0.1:6060` | Numeric loopback IP and port (1–65535), including bracketed IPv6; validated even when disabled | `$LOCALAI_PPROF_ADDRESS` | +| `--pprof-mutex-profile-fraction` | `0` | Nonnegative mutex sampling fraction; nonzero requires pprof | `$LOCALAI_PPROF_MUTEX_PROFILE_FRACTION` | +| `--pprof-block-profile-rate` | `0` | Nonnegative block sampling rate in nanoseconds; nonzero requires pprof | `$LOCALAI_PPROF_BLOCK_PROFILE_RATE` | +| `--request-phase-timing` | `false` | Emit sanitized phase timing at Info, independently of pprof; warn/error filtering suppresses events | `$LOCALAI_REQUEST_PHASE_TIMING` | + +Sampling adds overhead. The process applies configured rates once and resets them to zero on shutdown. +There is no runtime toggle, continuous capture, or public API route for diagnostics. +The private listener has no authentication: other processes in the same network namespace can access profiles. +Profiles can contain secrets; protect captured files. The `/debug/pprof/cmdline` endpoint is unavailable. +Shutdown allows five seconds for cleanup, then force-closes profiling connections. +Existing application initialization operations may not promptly observe cancellation. +The signal callback still returns after its cleanup budget expires. +If initialization fails without returning an application handle, this lifecycle cannot shut down resources created inside that failed initialization. +Application shutdown has no context parameter; work that ignores cancellation can outlive the cleanup budget. + ## Storage Flags | Parameter | Default | Description | Environment Variable | diff --git a/pkg/diagnostics/config.go b/pkg/diagnostics/config.go new file mode 100644 index 000000000..c2d2cb1a7 --- /dev/null +++ b/pkg/diagnostics/config.go @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "fmt" + "net" + "strconv" +) + +// Options configures independently opt-in profiling and request timing. +type Options struct { + Pprof bool + Address string + MutexProfileFraction int + BlockProfileRate int + RequestPhaseTiming bool +} + +func DefaultOptions() Options { return Options{Address: "127.0.0.1:6060"} } + +// Validate checks even disabled settings so invalid configuration fails startup. +func (o Options) Validate() error { + host, port, err := net.SplitHostPort(o.Address) + if err != nil { + return fmt.Errorf("invalid diagnostics address: %w", err) + } + ip := net.ParseIP(host) + if ip == nil || !ip.IsLoopback() { + return fmt.Errorf("diagnostics address must use a numeric loopback IP") + } + if port == "" { + return fmt.Errorf("diagnostics address requires a port") + } + for _, c := range port { + if c < '0' || c > '9' { + return fmt.Errorf("diagnostics port must be decimal") + } + } + p, err := strconv.ParseUint(port, 10, 16) + if err != nil || p == 0 { + return fmt.Errorf("diagnostics port must be in 1–65535") + } + if o.MutexProfileFraction < 0 || o.BlockProfileRate < 0 { + return fmt.Errorf("diagnostics sampling rates must be nonnegative") + } + if !o.Pprof && (o.MutexProfileFraction != 0 || o.BlockProfileRate != 0) { + return fmt.Errorf("diagnostics sampling requires pprof") + } + return nil +} diff --git a/pkg/diagnostics/config_test.go b/pkg/diagnostics/config_test.go new file mode 100644 index 000000000..ce383edb3 --- /dev/null +++ b/pkg/diagnostics/config_test.go @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics configuration", func() { + It("defaults to inert diagnostics", func() { + Expect(DefaultOptions()).To(Equal(Options{Address: "127.0.0.1:6060"})) + Expect(DefaultOptions().Validate()).To(Succeed()) + }) + It("accepts numeric loopback forms and decimal boundary ports", func() { + for _, address := range []string{"127.0.0.1:1", "127.255.255.254:65535", "[::1]:6060", "[0:0:0:0:0:0:0:1]:6060", "[::ffff:127.0.0.1]:6060", "127.0.0.1:06060"} { + o := DefaultOptions() + o.Address = address + Expect(o.Validate()).To(Succeed(), address) + } + }) + It("rejects invalid addresses even when disabled", func() { + for _, address := range []string{"", ":6060", "localhost:6060", "0.0.0.0:6060", "[::]:6060", "192.168.1.2:6060", "8.8.8.8:6060", "[::1%lo]:6060", "unix:/tmp/profile", "127.0.0.1", "::1:6060", "127.0.0.1:", "127.0.0.1:0", "127.0.0.1:65536", "127.0.0.1:999999999999999999999999", "127.0.0.1:-1", "127.0.0.1:+1", "127.0.0.1:1.0", "127.0.0.1:http", "127.0.0.1: 1", "127.1:6060", "127.000.0.1:6060", "[::ffff:192.168.1.2]:6060"} { + o := DefaultOptions() + o.Address = address + Expect(o.Validate()).NotTo(Succeed(), address) + } + }) + It("validates sampling independently from timing", func() { + for _, rates := range [][2]int{{-1, 0}, {0, -1}, {1, 0}, {0, 1}} { + o := DefaultOptions() + o.RequestPhaseTiming = true + o.MutexProfileFraction = rates[0] + o.BlockProfileRate = rates[1] + Expect(o.Validate()).NotTo(Succeed()) + o.Pprof = true + if rates[0] >= 0 && rates[1] >= 0 { + Expect(o.Validate()).To(Succeed()) + } else { + Expect(o.Validate()).NotTo(Succeed()) + } + } + }) +}) diff --git a/pkg/diagnostics/diagnostics_suite_test.go b/pkg/diagnostics/diagnostics_suite_test.go new file mode 100644 index 000000000..ee49de9c2 --- /dev/null +++ b/pkg/diagnostics/diagnostics_suite_test.go @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "testing" +) + +func TestDiagnostics(t *testing.T) { RegisterFailHandler(Fail); RunSpecs(t, "Diagnostics Suite") } diff --git a/pkg/diagnostics/docs_test.go b/pkg/diagnostics/docs_test.go new file mode 100644 index 000000000..57db18d9e --- /dev/null +++ b/pkg/diagnostics/docs_test.go @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "fmt" + "os" + "path/filepath" + "regexp" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// These checks catch documentation drift, not the correctness of operator prose. +var _ = Describe("Diagnostics documentation", func() { + readDoc := func(path string) string { + data, err := os.ReadFile(filepath.Join("..", "..", "docs", "content", path)) + Expect(err).NotTo(HaveOccurred()) + return string(data) + } + It("lists all five exact flags, environments, and defaults in both references", func() { + defaults := DefaultOptions() + rows := []struct{ flag, env, value string }{ + {"--pprof", "LOCALAI_PPROF", fmt.Sprint(defaults.Pprof)}, + {"--pprof-address", "LOCALAI_PPROF_ADDRESS", defaults.Address}, + {"--pprof-mutex-profile-fraction", "LOCALAI_PPROF_MUTEX_PROFILE_FRACTION", fmt.Sprint(defaults.MutexProfileFraction)}, + {"--pprof-block-profile-rate", "LOCALAI_PPROF_BLOCK_PROFILE_RATE", fmt.Sprint(defaults.BlockProfileRate)}, + {"--request-phase-timing", "LOCALAI_REQUEST_PHASE_TIMING", fmt.Sprint(defaults.RequestPhaseTiming)}, + } + for _, path := range []string{"features/diagnostics.md", "reference/cli-reference.md"} { + doc := readDoc(path) + for _, row := range rows { + Expect(doc).To(MatchRegexp(`(?m)^\| `+regexp.QuoteMeta("`"+row.flag+"`")+` \| `+regexp.QuoteMeta("`"+row.value+"`")+` \|[^\n]*`+regexp.QuoteMeta("`$"+row.env+"`")+` \|$`), path) + } + } + }) + It("documents every phase wire constant in the boundary table", func() { + doc := readDoc("features/diagnostics.md") + phases := []Phase{PhaseExtraction, PhaseBearerLookup, PhaseDefaultListing, PhaseBodyLookup, + PhaseConfigLoadDefaults, PhaseAliasResolution, PhaseConfigLockWait, PhaseConfigLockHold, + PhaseConfigFilter, PhaseFSEnumeration, PhaseLooseFilter, PhaseExistenceFallback, PhaseModelInit, + PhaseModelRouterCallback, PhaseReloadLockWait, PhaseReloadLockHold, PhaseReloadEnumeration, + PhaseReloadMetadata, PhaseReloadYAMLRead, PhaseReloadParseDefaults} + for _, phase := range phases { + Expect(doc).To(ContainSubstring("| `" + string(phase) + "` |")) + } + source, err := os.ReadFile("timing.go") + Expect(err).NotTo(HaveOccurred()) + // Also catch newly introduced constants omitted from the explicit list. + wires := regexp.MustCompile(`(?m)^\s*Phase\w+\s+Phase\s*=\s*"([^"]+)"`).FindAllStringSubmatch(string(source), -1) + Expect(wires).To(HaveLen(len(phases))) + for _, wire := range wires { + Expect(doc).To(ContainSubstring("| `" + wire[1] + "` |")) + } + }) + It("links both entry points to the feature page", func() { + readDoc("features/diagnostics.md") + for _, path := range []string{"reference/cli-reference.md", "getting-started/troubleshooting.md"} { + Expect(readDoc(path)).To(ContainSubstring(`{{% relref "features/diagnostics" %}}`)) + } + }) + It("retains sensitive access warnings and direct bounded capture examples", func() { + doc := readDoc("features/diagnostics.md") + for _, text := range []string{"`/debug/pprof/cmdline` is unavailable", "same network namespace", "unauthenticated", "sensitive", "umask 077", + "kubectl -n port-forward pod/ 6060:6060", + "curl --fail --output cpu.pprof 'http://127.0.0.1:6060/debug/pprof/profile?seconds=10'", "go tool pprof cpu.pprof", + "/debug/pprof/mutex?seconds=10", "/debug/pprof/block?seconds=10", "/debug/pprof/goroutine?debug=2"} { + Expect(doc).To(ContainSubstring(text)) + } + Expect(doc).NotTo(MatchRegexp(`(?i)port-forward\s+(service/|svc/)`)) + Expect(doc).NotTo(MatchRegexp(`(?im)^\s*(kind:\s*(Service|Ingress)|type:\s*NodePort)`)) + }) +}) diff --git a/pkg/diagnostics/log.go b/pkg/diagnostics/log.go new file mode 100644 index 000000000..fbede0aa8 --- /dev/null +++ b/pkg/diagnostics/log.go @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import "github.com/mudler/xlog" + +// LogEvent emits only the closed diagnostic schema at the normal Info level. +// Operator-selected warn/error filtering suppresses these events. +func LogEvent(event Event) { + if !validEvent(event) { + return + } + xlog.Info("diagnostic_phase", "id", event.ID, "kind", event.Kind, + "phase", event.Phase, "state", event.State, "outcome", event.Outcome, + "elapsed", event.Elapsed, "count", event.Count) +} diff --git a/pkg/diagnostics/server.go b/pkg/diagnostics/server.go new file mode 100644 index 000000000..0f9924dc6 --- /dev/null +++ b/pkg/diagnostics/server.go @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "net/http/pprof" + "runtime" + runtimepprof "runtime/pprof" + "sync" + "time" +) + +type samplingControl interface { + SetMutexProfileFraction(int) int + SetBlockProfileRate(int) +} +type runtimeSampling struct{} + +func (runtimeSampling) SetMutexProfileFraction(v int) int { return runtime.SetMutexProfileFraction(v) } +func (runtimeSampling) SetBlockProfileRate(v int) { runtime.SetBlockProfileRate(v) } + +// Server owns a private listener and process-global sampling. The run process +// must have only one sampling owner; prior runtime rates are not restored. +type Server struct { + httpServer *http.Server + addr net.Addr + errors chan error + done chan struct{} + sampling samplingControl + resetOnce sync.Once + shutdownOnce sync.Once + shutdownErr error +} + +// Start validates and binds before returning. Disabled profiling is inert. +func Start(opts Options) (*Server, error) { return start(opts, net.Listen, runtimeSampling{}) } + +func start(opts Options, listen func(string, string) (net.Listener, error), sampling samplingControl) (*Server, error) { + if err := opts.Validate(); err != nil { + return nil, err + } + if !opts.Pprof { + return &Server{}, nil + } + listener, err := listen("tcp", opts.Address) + if err != nil { + return nil, fmt.Errorf("bind diagnostics listener: %w", err) + } + mux := http.NewServeMux() + // Index dispatches subtree requests, so deny cmdline explicitly at both paths. + mux.HandleFunc("/debug/pprof/cmdline", http.NotFound) + mux.HandleFunc("/debug/pprof/cmdline/", http.NotFound) + mux.HandleFunc("/debug/pprof/", pprof.Index) + mux.HandleFunc("/debug/pprof/profile", pprof.Profile) + mux.HandleFunc("/debug/pprof/trace", pprof.Trace) + mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol) + for _, profile := range runtimepprof.Profiles() { + mux.Handle("/debug/pprof/"+profile.Name(), pprof.Handler(profile.Name())) + } + s := &Server{ + httpServer: &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 60 * time.Second, MaxHeaderBytes: 1 << 20}, + addr: listener.Addr(), errors: make(chan error, 1), done: make(chan struct{}), sampling: sampling, + } + sampling.SetMutexProfileFraction(opts.MutexProfileFraction) + sampling.SetBlockProfileRate(opts.BlockProfileRate) + go func() { + defer close(s.done) + defer close(s.errors) + if err := s.httpServer.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + // Stop active connections and sampling even if the caller has not yet + // consumed the buffered fatal error. + _ = s.httpServer.Close() + s.resetSampling() + s.errors <- err + } + }() + return s, nil +} + +func (s *Server) resetSampling() { + s.resetOnce.Do(func() { + s.sampling.SetMutexProfileFraction(0) + s.sampling.SetBlockProfileRate(0) + }) +} + +// Errors delivers one unexpected serving failure, then closes. Normal shutdown +// closes it without a value; disabled servers return nil. +func (s *Server) Errors() <-chan error { return s.errors } +func (s *Server) Addr() net.Addr { return s.addr } + +// Shutdown drains requests, force-closing on expiry. Repeated calls return the +// first result and wait for the serving goroutine to finish. +func (s *Server) Shutdown(ctx context.Context) error { + if s.httpServer == nil { + return nil + } + s.shutdownOnce.Do(func() { + s.shutdownErr = s.httpServer.Shutdown(ctx) + if s.shutdownErr != nil { + _ = s.httpServer.Close() + } + <-s.done + s.resetSampling() + }) + return s.shutdownErr +} diff --git a/pkg/diagnostics/server_test.go b/pkg/diagnostics/server_test.go new file mode 100644 index 000000000..7b770bc98 --- /dev/null +++ b/pkg/diagnostics/server_test.go @@ -0,0 +1,291 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "runtime" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeSampling struct { + mu sync.Mutex + mutex, block []int +} + +func (f *fakeSampling) SetMutexProfileFraction(v int) int { + f.mu.Lock() + defer f.mu.Unlock() + f.mutex = append(f.mutex, v) + return 0 +} +func (f *fakeSampling) SetBlockProfileRate(v int) { + f.mu.Lock() + defer f.mu.Unlock() + f.block = append(f.block, v) +} +func (f *fakeSampling) calls() [2][]int { + f.mu.Lock() + defer f.mu.Unlock() + return [2][]int{append([]int(nil), f.mutex...), append([]int(nil), f.block...)} +} + +type failingListener struct { + net.Listener + err error +} + +func (l failingListener) Accept() (net.Conn, error) { return nil, l.err } + +// A read signals that net/http has registered the connection before shutdown. +type readingListener struct { + net.Listener + reading chan struct{} +} + +func (l readingListener) Accept() (net.Conn, error) { + c, err := l.Listener.Accept() + if err != nil { + return nil, err + } + return &readingConn{Conn: c, reading: l.reading}, nil +} + +type readingConn struct { + net.Conn + reading chan struct{} + once sync.Once +} + +func (c *readingConn) Read(p []byte) (int, error) { + c.once.Do(func() { close(c.reading) }) + return c.Conn.Read(p) +} + +func enabledOptions() Options { o := DefaultOptions(); o.Pprof = true; return o } +func launch(sampling samplingControl) *Server { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = l.Close() }) + s, err := start(enabledOptions(), func(network, address string) (net.Listener, error) { + Expect(network).To(Equal("tcp")) + Expect(address).To(Equal("127.0.0.1:6060")) + return l, nil + }, sampling) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = s.Shutdown(ctx) + }) + return s +} + +var _ = Describe("Diagnostics server", Serial, func() { + It("is inert even with debug logging or request timing enabled", func() { + old := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelDebug}))) + defer slog.SetDefault(old) + for _, timing := range []bool{false, true} { + f := &fakeSampling{} + o := DefaultOptions() + o.RequestPhaseTiming = timing + s, err := start(o, func(string, string) (net.Listener, error) { Fail("disabled listener called"); return nil, nil }, f) + Expect(err).NotTo(HaveOccurred()) + Expect(s).NotTo(BeNil()) + Expect(s.Addr()).To(BeNil()) + Expect(s.Errors()).To(BeNil()) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(f.calls()).To(Equal([2][]int{})) + } + }) + It("validates before binding and reports bind failures without sampling", func() { + f := &fakeSampling{} + o := enabledOptions() + o.Address = "localhost:6060" + s, err := start(o, func(string, string) (net.Listener, error) { Fail("invalid address bound"); return nil, nil }, f) + Expect(err).To(HaveOccurred()) + Expect(s).To(BeNil()) + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + o.Address = l.Addr().String() + s, err = start(o, net.Listen, f) + Expect(err).To(HaveOccurred()) + Expect(s).To(BeNil()) + Expect(f.calls()).To(Equal([2][]int{})) + }) + It("serves explicit standard profiles without leaking default mux or cmdline", func() { + old := http.DefaultServeMux + http.DefaultServeMux = http.NewServeMux() + defer func() { http.DefaultServeMux = old }() + http.DefaultServeMux.HandleFunc("/private-sentinel", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("leaked")) }) + s := launch(&fakeSampling{}) + client := httpclient.NewWithTimeout(5 * time.Second) + defer client.CloseIdleConnections() + for _, path := range []string{"/debug/pprof/", "/debug/pprof/profile?seconds=1", "/debug/pprof/trace?seconds=0.01", "/debug/pprof/symbol", "/debug/pprof/heap", "/debug/pprof/goroutine?debug=1", "/debug/pprof/mutex", "/debug/pprof/block", "/debug/pprof/allocs", "/debug/pprof/threadcreate"} { + res, err := client.Get("http://" + s.Addr().String() + path) + Expect(err).NotTo(HaveOccurred(), path) + body, err := io.ReadAll(res.Body) + res.Body.Close() + Expect(err).NotTo(HaveOccurred()) + Expect(res.StatusCode).To(Equal(200), path) + Expect(body).NotTo(BeEmpty(), path) + } + for _, path := range []string{"/debug/pprof/cmdline", "/debug/pprof/cmdline?debug=1", "/debug/pprof/cmdline/", "/debug/pprof/cmdline/anything?x=1", "/private-sentinel", "/debug/pprof/unknown"} { + res, err := client.Get("http://" + s.Addr().String() + path) + Expect(err).NotTo(HaveOccurred()) + res.Body.Close() + Expect(res.StatusCode).To(Equal(404), path) + } + Expect(s.httpServer.ReadHeaderTimeout).To(Equal(5 * time.Second)) + Expect(s.httpServer.IdleTimeout).To(Equal(60 * time.Second)) + Expect(s.httpServer.WriteTimeout).To(BeZero()) + Expect(s.httpServer.MaxHeaderBytes).To(Equal(1 << 20)) + }) + It("applies rates after binding and resets once on repeated shutdown", func() { + f := &fakeSampling{} + o := enabledOptions() + o.MutexProfileFraction = 7 + o.BlockProfileRate = 123 + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + s, err := start(o, func(string, string) (net.Listener, error) { Expect(f.calls()).To(Equal([2][]int{})); return l, nil }, f) + Expect(err).NotTo(HaveOccurred()) + Expect(f.calls()).To(Equal([2][]int{{7}, {123}})) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(s.Errors()).To(BeClosed()) + Expect(f.calls()).To(Equal([2][]int{{7, 0}, {123, 0}})) + rebound, err := net.Listen("tcp", s.Addr().String()) + Expect(err).NotTo(HaveOccurred()) + rebound.Close() + }) + It("buffers an immediate unexpected Serve failure exactly once", func() { + f := &fakeSampling{} + boom := errors.New("accept failed") + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + s, err := start(enabledOptions(), func(string, string) (net.Listener, error) { return failingListener{l, boom}, nil }, f) + Expect(err).NotTo(HaveOccurred()) + Eventually(s.Errors()).Should(Receive(MatchError(boom))) + Eventually(s.Errors()).Should(BeClosed()) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(f.calls()).To(Equal([2][]int{{0, 0}, {0, 0}})) + }) + It("force closes active connections on graceful expiry and waits for Serve", func() { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + reading := make(chan struct{}) + s, err := start(enabledOptions(), func(string, string) (net.Listener, error) { + return readingListener{Listener: l, reading: reading}, nil + }, &fakeSampling{}) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _ = s.Shutdown(ctx) + }) + conn, err := net.Dial("tcp", s.Addr().String()) + Expect(err).NotTo(HaveOccurred()) + defer conn.Close() + _, err = conn.Write([]byte("GET /debug/pprof/ HTTP/1.1\r\nHost: localhost\r\n")) + Expect(err).NotTo(HaveOccurred()) + Eventually(reading).Should(BeClosed()) + // An incomplete header stays StateNew until header timeout; Shutdown must force-close it. + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + Expect(s.Shutdown(ctx)).To(MatchError(context.DeadlineExceeded)) + Expect(s.Shutdown(context.Background())).To(MatchError(context.DeadlineExceeded)) + Expect(s.Errors()).To(BeClosed()) + Expect(conn.SetReadDeadline(time.Now().Add(time.Second))).To(Succeed()) + _, err = conn.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred()) + }) + It("supports serial runtime sampling and leaves it disabled", func() { + previous := runtime.SetMutexProfileFraction(0) + defer runtime.SetMutexProfileFraction(previous) + defer runtime.SetBlockProfileRate(0) + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + defer l.Close() + o := enabledOptions() + o.Address = l.Addr().String() + l.Close() + o.MutexProfileFraction = 2 + o.BlockProfileRate = 1 + s, err := Start(o) + Expect(err).NotTo(HaveOccurred()) + defer s.Shutdown(context.Background()) + Expect(runtime.SetMutexProfileFraction(-1)).To(Equal(2)) + Expect(s.Shutdown(context.Background())).To(Succeed()) + Expect(runtime.SetMutexProfileFraction(-1)).To(BeZero()) + }) +}) + +var _ = Describe("Diagnostics real capture shutdown", Serial, func() { + for _, seconds := range []int{2, 30} { + It(fmt.Sprintf("bounds shutdown during a %d second capture", seconds), func() { + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + opts := enabledOptions() + opts.Address = l.Addr().String() + Expect(l.Close()).To(Succeed()) + s, err := Start(opts) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = s.Shutdown(context.Background()) }) + client := httpclient.NewWithTimeout(12 * time.Second) + defer client.CloseIdleConnections() + capture := make(chan error, 1) + go func() { + response, err := client.Get(fmt.Sprintf("http://%s/debug/pprof/profile?seconds=%d", s.Addr(), seconds)) + if err == nil { + _, err = io.Copy(io.Discard, response.Body) + response.Body.Close() + if response.StatusCode != http.StatusOK { + err = fmt.Errorf("profile status %d", response.StatusCode) + } + } + capture <- err + }() + Eventually(func() bool { + response, err := client.Get("http://" + s.Addr().String() + "/debug/pprof/goroutine?debug=2") + if err != nil { + return false + } + defer response.Body.Close() + body, _ := io.ReadAll(response.Body) + return strings.Contains(string(body), "net/http/pprof.Profile") + }, 3*time.Second).Should(BeTrue()) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err = s.Shutdown(ctx) + if seconds == 2 { + Expect(err).NotTo(HaveOccurred()) + Eventually(capture).Should(Receive(BeNil())) + } else { + Expect(err).To(MatchError(context.DeadlineExceeded)) + Eventually(capture).Should(Receive(Not(BeNil()))) + } + Expect(s.Errors()).To(BeClosed()) + l, err = net.Listen("tcp", s.Addr().String()) + Expect(err).NotTo(HaveOccurred()) + Expect(l.Close()).To(Succeed()) + }) + } +}) diff --git a/pkg/diagnostics/timing.go b/pkg/diagnostics/timing.go new file mode 100644 index 000000000..b73d5a740 --- /dev/null +++ b/pkg/diagnostics/timing.go @@ -0,0 +1,191 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "context" + "time" + + "github.com/google/uuid" +) + +type Kind string +type Phase string +type State string +type Outcome string + +const ( + KindRequest Kind = "request" + KindConfigReload Kind = "config_reload" + StateStart State = "start" + StateEnd State = "end" + StateMark State = "mark" + OutcomeOK Outcome = "ok" + OutcomeError Outcome = "error" + OutcomeCanceled Outcome = "canceled" + OutcomeSkipped Outcome = "skipped" + + PhaseExtraction Phase = "extraction" + PhaseBearerLookup Phase = "bearer_lookup" + PhaseDefaultListing Phase = "default_listing" + PhaseBodyLookup Phase = "body_lookup" + PhaseConfigLoadDefaults Phase = "config_load_defaults" + PhaseAliasResolution Phase = "alias_resolution" + PhaseConfigLockWait Phase = "config_lock_wait" + PhaseConfigLockHold Phase = "config_lock_hold" + PhaseConfigFilter Phase = "config_filter" + PhaseFSEnumeration Phase = "fs_enumeration" + PhaseLooseFilter Phase = "loose_filter" + PhaseExistenceFallback Phase = "existence_fallback" + PhaseModelInit Phase = "model_init" + PhaseModelRouterCallback Phase = "model_router_callback" + PhaseReloadLockWait Phase = "reload_lock_wait" + PhaseReloadLockHold Phase = "reload_lock_hold" + PhaseReloadEnumeration Phase = "reload_enumeration" + PhaseReloadMetadata Phase = "reload_metadata" + PhaseReloadYAMLRead Phase = "reload_yaml_read" + PhaseReloadParseDefaults Phase = "reload_parse_defaults" +) + +// Event deliberately has no free-form payload or error field. +type Event struct { + ID string + Kind Kind + Phase Phase + State State + Outcome Outcome + Elapsed time.Duration + Count int +} + +// Sink must be safe for concurrent calls. Calls are synchronous; callers should +// use Record to emit aggregates after releasing application locks. +type Sink func(Event) + +// Recorder is immutable after construction and may be shared between requests. +// A nil recorder disables timing. +type Recorder struct { + sink Sink + now func() time.Time +} + +// NewRecorder defaults to Info logging when sink is nil. +func NewRecorder(sink Sink) *Recorder { + if sink == nil { + sink = LogEvent + } + return &Recorder{sink: sink, now: time.Now} +} + +type observationKey struct{} +type observation struct { + recorder *Recorder + id string + kind Kind +} + +func observed(ctx context.Context) *observation { + o, _ := ctx.Value(observationKey{}).(*observation) + return o +} + +func (r *Recorder) Request(ctx context.Context) context.Context { + if r == nil { + return ctx + } + if o := observed(ctx); o != nil && o.kind == KindRequest { + return ctx + } + return r.attach(ctx, KindRequest) +} + +func (r *Recorder) Reload(ctx context.Context) context.Context { + if r == nil { + return ctx + } + return r.attach(ctx, KindConfigReload) +} + +func (r *Recorder) attach(ctx context.Context, kind Kind) context.Context { + return context.WithValue(ctx, observationKey{}, &observation{recorder: r, id: uuid.NewString(), kind: kind}) +} + +// Inherit transfers only diagnostic identity, leaving destination lifetime and +// other values intact. A source without observation leaves dst unchanged. +func Inherit(dst, src context.Context) context.Context { + if o := observed(src); o != nil { + return context.WithValue(dst, observationKey{}, o) + } + return dst +} + +func Enabled(ctx context.Context) bool { return observed(ctx) != nil } + +func noopEnd(Outcome, int) {} + +// Begin emits a start immediately. The returned function should be called once, +// with the actual outcome; cancellation is not inferred or propagated here. +// Nested durations overlap and must not be summed as disjoint work. +func Begin(ctx context.Context, phase Phase) func(outcome Outcome, count int) { + o := observed(ctx) + if o == nil || !validPhase(phase) { + return noopEnd + } + start := o.recorder.now() + o.emit(phase, StateStart, OutcomeOK, 0, 0) + return func(outcome Outcome, count int) { o.emit(phase, StateEnd, outcome, o.recorder.now().Sub(start), count) } +} + +func Mark(ctx context.Context, phase Phase) { + if o := observed(ctx); o != nil { + o.emit(phase, StateMark, OutcomeOK, 0, 0) + } +} + +func Record(ctx context.Context, phase Phase, elapsed time.Duration, outcome Outcome, count int) { + if o := observed(ctx); o != nil { + o.emit(phase, StateEnd, outcome, elapsed, count) + } +} + +func (o *observation) emit(phase Phase, state State, outcome Outcome, elapsed time.Duration, count int) { + e := Event{ID: o.id, Kind: o.kind, Phase: phase, State: state, Outcome: outcome, Elapsed: elapsed, Count: count} + if validEvent(e) { + o.recorder.sink(e) + } +} + +func validPhase(p Phase) bool { + switch p { + case PhaseExtraction, PhaseBearerLookup, PhaseDefaultListing, PhaseBodyLookup, + PhaseConfigLoadDefaults, PhaseAliasResolution, PhaseConfigLockWait, PhaseConfigLockHold, + PhaseConfigFilter, PhaseFSEnumeration, PhaseLooseFilter, PhaseExistenceFallback, + PhaseModelInit, PhaseModelRouterCallback, PhaseReloadLockWait, PhaseReloadLockHold, + PhaseReloadEnumeration, PhaseReloadMetadata, PhaseReloadYAMLRead, PhaseReloadParseDefaults: + return true + } + return false +} + +func validEvent(e Event) bool { + if !validPhase(e.Phase) { + return false + } + switch e.Kind { + case KindRequest, KindConfigReload: + default: + return false + } + switch e.State { + case StateStart, StateEnd, StateMark: + default: + return false + } + switch e.Outcome { + case OutcomeOK, OutcomeError, OutcomeCanceled, OutcomeSkipped: + default: + return false + } + // LogEvent is exported: reject free-form identities even for direct calls. + id, err := uuid.Parse(e.ID) + return err == nil && id.Version() == 4 && id.Variant() == uuid.RFC4122 && id.String() == e.ID && e.Elapsed >= 0 && e.Count >= 0 +} diff --git a/pkg/diagnostics/timing_test.go b/pkg/diagnostics/timing_test.go new file mode 100644 index 000000000..f2a1c8ce0 --- /dev/null +++ b/pkg/diagnostics/timing_test.go @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: MIT +package diagnostics + +import ( + "bytes" + "context" + "encoding/json" + "log/slog" + "os" + "reflect" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/mudler/xlog" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics timing", func() { + It("keeps disabled calls inert and allocation free", func() { + ctx := context.Background() + var r *Recorder + Expect(r.Request(ctx)).To(BeIdenticalTo(ctx)) + Expect(r.Reload(ctx)).To(BeIdenticalTo(ctx)) + Expect(Inherit(ctx, ctx)).To(BeIdenticalTo(ctx)) + Expect(Enabled(ctx)).To(BeFalse()) + Expect(reflect.ValueOf(Begin(ctx, PhaseExtraction)).Pointer()).To(Equal(reflect.ValueOf(Begin(ctx, PhaseBodyLookup)).Pointer())) + Expect(testing.AllocsPerRun(100, func() { + r.Request(ctx) + r.Reload(ctx) + Begin(ctx, PhaseExtraction)(OutcomeOK, 0) + Mark(ctx, PhaseModelInit) + Record(ctx, PhaseConfigFilter, 0, OutcomeOK, 0) + })).To(BeZero()) + }) + It("records exact nested durations and aggregate outcomes without summing", func() { + var events []Event + r := NewRecorder(func(e Event) { events = append(events, e) }) + now := time.Unix(100, 0) + r.now = func() time.Time { return now } + ctx := r.Request(context.Background()) + outer := Begin(ctx, PhaseExtraction) + now = now.Add(2 * time.Second) + inner := Begin(ctx, PhaseBodyLookup) + now = now.Add(3 * time.Second) + inner(OutcomeError, 1) + now = now.Add(time.Second) + outer(OutcomeCanceled, 2) + Mark(ctx, PhaseModelInit) + Record(ctx, PhaseReloadMetadata, 9*time.Second, OutcomeSkipped, 4) + Expect(events).To(HaveLen(6)) + Expect(events[0].State).To(Equal(StateStart)) + Expect(events[2].Elapsed).To(Equal(3 * time.Second)) + Expect(events[2].Outcome).To(Equal(OutcomeError)) + Expect(events[3].Elapsed).To(Equal(6 * time.Second)) + Expect(events[3].Outcome).To(Equal(OutcomeCanceled)) + Expect(events[4].State).To(Equal(StateMark)) + Expect(events[4].Elapsed).To(BeZero()) + Expect(events[5].Count).To(Equal(4)) + Expect(events[5].Elapsed).To(Equal(9 * time.Second)) + for _, e := range events { + Expect(e.ID).To(Equal(events[0].ID)) + Expect(e.Kind).To(Equal(KindRequest)) + } + }) + It("inherits only private observation and preserves destination cancellation and values", func() { + type key struct{} + var events []Event + r := NewRecorder(func(e Event) { events = append(events, e) }) + src, cancel := context.WithCancel(context.WithValue(context.Background(), key{}, "secret-header")) + src = r.Request(src) + cancel() + dst := context.WithValue(context.Background(), key{}, "destination") + inherited := Inherit(dst, src) + Expect(inherited.Err()).To(BeNil()) + _, deadline := inherited.Deadline() + Expect(deadline).To(BeFalse()) + Expect(inherited.Value(key{})).To(Equal("destination")) + Mark(src, PhaseExtraction) + Mark(inherited, PhaseExtraction) + Expect(events[0].ID).To(Equal(events[1].ID)) + Expect(r.Request(inherited)).To(BeIdenticalTo(inherited)) + dstCanceled, cancelDst := context.WithCancel(dst) + cancelDst() + Expect(Inherit(dstCanceled, src).Err()).To(Equal(context.Canceled)) + }) + It("reuses another recorder's request identity and original observer", func() { + var originalEvents, otherEvents []Event + original := NewRecorder(func(e Event) { originalEvents = append(originalEvents, e) }) + other := NewRecorder(func(e Event) { otherEvents = append(otherEvents, e) }) + req := original.Request(context.Background()) + Mark(req, PhaseExtraction) + + reused := other.Request(req) + Expect(reused).To(BeIdenticalTo(req)) + Expect(observed(reused)).To(BeIdenticalTo(observed(req))) + Expect(observed(reused).recorder).To(BeIdenticalTo(original)) + Mark(reused, PhaseModelInit) + Expect(originalEvents).To(HaveLen(2)) + Expect(originalEvents[1].ID).To(Equal(originalEvents[0].ID)) + Expect(otherEvents).To(BeEmpty()) + Expect(testing.AllocsPerRun(100, func() { other.Request(req) })).To(BeZero()) + + var disabled *Recorder + Expect(disabled.Request(req)).To(BeIdenticalTo(req)) + Expect(disabled.Reload(req)).To(BeIdenticalTo(req)) + }) + It("generates fresh reload identities and a fresh request after reload", func() { + var events []Event + sink := func(e Event) { events = append(events, e) } + r := NewRecorder(sink) + req := r.Request(context.Background()) + reload := r.Reload(req) + for _, ctx := range []context.Context{req, reload, r.Reload(reload), r.Request(reload)} { + Mark(ctx, PhaseExtraction) + } + ids := map[string]bool{} + for _, e := range events { + _, err := uuid.Parse(e.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(ids[e.ID]).To(BeFalse()) + ids[e.ID] = true + } + Expect(events[1].Kind).To(Equal(KindConfigReload)) + Expect(events[2].Kind).To(Equal(KindConfigReload)) + Expect(events[3].Kind).To(Equal(KindRequest)) + }) + It("isolates concurrent requests while permitting concurrent phases", func() { + var mu sync.Mutex + var events []Event + r := NewRecorder(func(e Event) { mu.Lock(); defer mu.Unlock(); events = append(events, e) }) + var wg sync.WaitGroup + for i := 0; i < 32; i++ { + wg.Add(1) + go func() { + defer wg.Done() + ctx := r.Request(context.Background()) + end := Begin(ctx, PhaseExtraction) + Mark(ctx, PhaseModelRouterCallback) + end(OutcomeOK, 0) + }() + } + wg.Wait() + grouped := map[string][]Event{} + for _, e := range events { + grouped[e.ID] = append(grouped[e.ID], e) + } + Expect(grouped).To(HaveLen(32)) + for _, group := range grouped { + Expect(group).To(HaveLen(3)) + Expect(group[0].State).To(Equal(StateStart)) + Expect(group[2].State).To(Equal(StateEnd)) + } + }) + It("accepts exactly the documented phase vocabulary", func() { + phases := []Phase{PhaseExtraction, PhaseBearerLookup, PhaseDefaultListing, PhaseBodyLookup, PhaseConfigLoadDefaults, PhaseAliasResolution, PhaseConfigLockWait, PhaseConfigLockHold, PhaseConfigFilter, PhaseFSEnumeration, PhaseLooseFilter, PhaseExistenceFallback, PhaseModelInit, PhaseModelRouterCallback, PhaseReloadLockWait, PhaseReloadLockHold, PhaseReloadEnumeration, PhaseReloadMetadata, PhaseReloadYAMLRead, PhaseReloadParseDefaults} + expected := []string{"extraction", "bearer_lookup", "default_listing", "body_lookup", "config_load_defaults", "alias_resolution", "config_lock_wait", "config_lock_hold", "config_filter", "fs_enumeration", "loose_filter", "existence_fallback", "model_init", "model_router_callback", "reload_lock_wait", "reload_lock_hold", "reload_enumeration", "reload_metadata", "reload_yaml_read", "reload_parse_defaults"} + var events []Event + ctx := NewRecorder(func(e Event) { events = append(events, e) }).Request(context.Background()) + for i, phase := range phases { + Expect(string(phase)).To(Equal(expected[i])) + Mark(ctx, phase) + } + Expect(events).To(HaveLen(len(phases))) + }) + It("rejects arbitrary enums before sink emission", func() { + var events []Event + r := NewRecorder(func(e Event) { events = append(events, e) }) + ctx := r.Request(context.Background()) + Begin(ctx, Phase("secret"))(OutcomeOK, 0) + Mark(ctx, Phase("secret")) + Record(ctx, PhaseExtraction, 0, Outcome("secret"), 0) + Expect(events).To(BeEmpty()) + }) + It("logs only sanitized fixed fields at Info and rejects forged payloads", func() { + var buf bytes.Buffer + xlog.SetLogger(xlog.NewLoggerWithHandler(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo}), xlog.LogLevel("info"))) + defer xlog.SetLogger(xlog.NewLogger(xlog.LogLevel(os.Getenv(xlog.EnvLogLevel)), os.Getenv(xlog.EnvLogFormat))) + ctx := NewRecorder(nil).Request(context.WithValue(context.Background(), "X-Correlation-ID", "secret")) + Mark(ctx, PhaseModelInit) + var logged map[string]any + Expect(json.Unmarshal(buf.Bytes(), &logged)).To(Succeed()) + Expect(logged["msg"]).To(Equal("diagnostic_phase")) + Expect(logged["level"]).To(Equal("INFO")) + Expect(logged).To(HaveLen(10)) + Expect(buf.String()).NotTo(ContainSubstring("secret")) + valid := Event{ID: uuid.NewString(), Kind: KindRequest, Phase: PhaseExtraction, State: StateEnd, Outcome: OutcomeOK} + for _, mutate := range []func(*Event){func(e *Event) { e.ID = "secret" }, func(e *Event) { e.Kind = "secret" }, func(e *Event) { e.Phase = "secret" }, func(e *Event) { e.State = "secret" }, func(e *Event) { e.Outcome = "secret" }} { + e := valid + mutate(&e) + buf.Reset() + LogEvent(e) + Expect(buf.Len()).To(BeZero()) + } + }) +}) diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index cf5003912..88a35160f 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -9,6 +9,7 @@ import ( "strings" "time" + "github.com/mudler/LocalAI/pkg/diagnostics" grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" processManager "github.com/mudler/go-processmanager" @@ -74,7 +75,17 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string ml.mu.Unlock() if router != nil { xlog.Info("Routing model to remote node via ModelRouter", "modelID", modelID, "backend", backend) - return router(o.context, backend, modelID, modelName, modelFile, o.configRevision, o.gRPCOptions, o.parallelRequests) + // This measures the callback, not transport or worker execution. + end := diagnostics.Begin(o.context, diagnostics.PhaseModelRouterCallback) + model, err := router(o.context, backend, modelID, modelName, modelFile, o.configRevision, o.gRPCOptions, o.parallelRequests) + outcome := diagnostics.OutcomeOK + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + outcome = diagnostics.OutcomeCanceled + } else if err != nil { + outcome = diagnostics.OutcomeError + } + end(outcome, 0) + return model, err } uri := ml.GetAllExternalBackends(o)[backend] @@ -444,6 +455,7 @@ func (ml *ModelLoader) updateModelLastUsed(m *Model) { func (ml *ModelLoader) Load(opts ...Option) (grpc.Backend, error) { o := NewOptions(opts...) + diagnostics.Mark(o.context, diagnostics.PhaseModelInit) ml.mu.Lock() distributed := ml.modelRouter != nil diff --git a/pkg/model/initializers_diagnostics_test.go b/pkg/model/initializers_diagnostics_test.go new file mode 100644 index 000000000..250f2e95c --- /dev/null +++ b/pkg/model/initializers_diagnostics_test.go @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: MIT +package model + +import ( + "context" + "errors" + "fmt" + + "github.com/mudler/LocalAI/pkg/diagnostics" + "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Embedding leaves unexpected backend work as a panic rather than a network call. +type diagnosticsHealthyBackend struct{ grpc.Backend } + +func (*diagnosticsHealthyBackend) HealthCheck(context.Context) (bool, error) { return true, nil } + +var _ = Describe("Diagnostics model initialization", func() { + for _, enabled := range []bool{false, true} { + for _, outcome := range []string{"success", "error", "cancel"} { + It(fmt.Sprintf("times only the callback: enabled=%t outcome=%s", enabled, outcome), func() { + var events []diagnostics.Event + source := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + diagnostics.Mark(source, diagnostics.PhaseExtraction) + id := events[0].ID + events = nil + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if enabled { + ctx = diagnostics.Inherit(ctx, source) + } + ml := NewModelLoader(&system.SystemState{}) + want := NewModel("secret-id", "unused", nil) + cause := errors.New("secret-error") + if outcome == "success" { + cause = nil + } + if outcome == "cancel" { + cancel() + cause = context.Canceled + } + opts := NewOptions(WithContext(ctx), WithConfigRevision("secret-revision"), EnableParallelRequests) + calls := 0 + ml.SetModelRouter(func(got context.Context, backend, id, name, file, revision string, options *pb.ModelOptions, parallel bool) (*Model, error) { + calls++ + Expect(got).To(BeIdenticalTo(ctx)) + Expect([]string{backend, id, name, file, revision}).To(Equal([]string{"secret-backend", "secret-id", "secret-name", "secret-file", "secret-revision"})) + Expect(options).To(BeIdenticalTo(opts.gRPCOptions)) + Expect(parallel).To(BeTrue()) + if enabled { + Expect(events).To(HaveLen(1)) + Expect(events[0].State).To(Equal(diagnostics.StateStart)) + } + return want, cause + }) + got, err := ml.grpcModel("secret-backend", opts)("secret-id", "secret-name", "secret-file") + Expect(got).To(BeIdenticalTo(want)) + if cause == nil { + Expect(err).To(BeNil()) + } else { + Expect(err).To(Equal(cause)) + } + Expect(calls).To(Equal(1)) + if !enabled { + Expect(events).To(BeEmpty()) + return + } + Expect(events).To(HaveLen(2)) + expected := diagnostics.OutcomeOK + if outcome == "error" { + expected = diagnostics.OutcomeError + } + if outcome == "cancel" { + expected = diagnostics.OutcomeCanceled + } + Expect(events[1].State).To(Equal(diagnostics.StateEnd)) + Expect(events[1].Outcome).To(Equal(expected)) + for _, e := range events { + Expect(e.ID).To(Equal(id)) + Expect(e.Phase).To(Equal(diagnostics.PhaseModelRouterCallback)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("secret")) + }) + } + } + It("carries one inherited identity from Load into the router", func() { + var events []diagnostics.Event + source := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + diagnostics.Mark(source, diagnostics.PhaseExtraction) + id := events[0].ID + events = nil + ctx := diagnostics.Inherit(context.Background(), source) + ml := NewModelLoader(&system.SystemState{}) + ml.SetModelRouter(func(got context.Context, _, _, _, _, _ string, _ *pb.ModelOptions, _ bool) (*Model, error) { + Expect(got).To(BeIdenticalTo(ctx)) + Expect(events).To(HaveLen(2)) + Expect(events[0].Phase).To(Equal(diagnostics.PhaseModelInit)) + Expect(events[1].Phase).To(Equal(diagnostics.PhaseModelRouterCallback)) + return nil, errors.New("secret-router-error") + }) + got, err := ml.Load(WithContext(ctx), WithModelID("secret-id"), WithBackendString("secret-backend")) + Expect(got).To(BeNil()) + Expect(err).To(MatchError("failed to route model with internal loader: secret-router-error")) + Expect(events).To(HaveLen(3)) + Expect(events[2].Outcome).To(Equal(diagnostics.OutcomeError)) + for _, e := range events { + Expect(e.ID).To(Equal(id)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("secret")) + }) + It("marks Load before admission and does not time the local branch", func() { + var events []diagnostics.Event + ctx := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + ml := NewModelLoader(&system.SystemState{}) + cause := errors.New("admission rejected without spawning") + ml.SetLoadLifecycleObserver(func(BackendLoadEvent) (func(BackendLoadEvent), error) { + Expect(events).To(HaveLen(1)) + Expect(events[0].Phase).To(Equal(diagnostics.PhaseModelInit)) + Expect(events[0].State).To(Equal(diagnostics.StateMark)) + return nil, cause + }) + got, err := ml.Load(WithContext(ctx), WithModelID("secret-id"), WithBackendString("secret-backend")) + Expect(got).To(BeNil()) + Expect(err).To(MatchError("failed to load model with internal loader: " + cause.Error())) + Expect(events).To(HaveLen(1)) + }) + It("marks cache hits and leaves unobserved loads inert", func() { + var events []diagnostics.Event + ctx := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + ml := NewModelLoader(&system.SystemState{}) + cached := NewModelWithClient("cached", "unused", &diagnosticsHealthyBackend{}) + cached.MarkHealthy() + ml.store.Set("cached", cached) + for _, c := range []context.Context{context.Background(), ctx} { + got, err := ml.Load(WithContext(c), WithModelID("cached")) + Expect(err).NotTo(HaveOccurred()) + Expect(got).NotTo(BeNil()) + } + Expect(events).To(HaveLen(1)) + Expect(events[0].Phase).To(Equal(diagnostics.PhaseModelInit)) + }) +}) diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 85e39f3b5..1db9a9cd7 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -12,6 +12,7 @@ import ( "sync/atomic" "time" + "github.com/mudler/LocalAI/pkg/diagnostics" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/utils" @@ -395,11 +396,35 @@ var knownModelsNameSuffixToSkip []string = []string{ } func (ml *ModelLoader) ListFilesInModelPath() ([]string, error) { - files, err := os.ReadDir(ml.ModelPath) + return ml.ListFilesInModelPathContext(context.Background()) +} + +// ListFilesInModelPathContext preserves discovery behavior while observing the caller's phases. +func (ml *ModelLoader) ListFilesInModelPathContext(ctx context.Context) ([]string, error) { + return ml.listFilesInModelPathContext(ctx, os.ReadDir) +} + +// The operation seam is call-local: no request or filesystem hooks live on shared state. +func (ml *ModelLoader) listFilesInModelPathContext(ctx context.Context, readDir func(string) ([]os.DirEntry, error)) ([]string, error) { + observed := diagnostics.Enabled(ctx) + var start time.Time + if observed { + start = time.Now() + } + files, err := readDir(ml.ModelPath) + if observed { + elapsed := time.Since(start) + outcome := diagnostics.OutcomeOK + if err != nil { + outcome = diagnostics.OutcomeError + } + diagnostics.Record(ctx, diagnostics.PhaseFSEnumeration, elapsed, outcome, 1) + } if err != nil { return []string{}, err } + end := diagnostics.Begin(ctx, diagnostics.PhaseLooseFilter) models := []string{} FILE: for _, file := range files { @@ -431,6 +456,7 @@ FILE: models = append(models, file.Name()) } + end(diagnostics.OutcomeOK, len(files)) return models, nil } diff --git a/pkg/model/loader_diagnostics_test.go b/pkg/model/loader_diagnostics_test.go new file mode 100644 index 000000000..3e532548b --- /dev/null +++ b/pkg/model/loader_diagnostics_test.go @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: MIT +package model + +import ( + "context" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/mudler/LocalAI/pkg/diagnostics" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Diagnostics model directory enumeration", func() { + It("counts only the existing ReadDir and separates filename work without leaking paths", func() { + dir := GinkgoT().TempDir() + for _, name := range []string{"secret-candidate", "z-model", "weights.GGUF", "config.yaml", "CACHEDIR.TAG", "x.bak-save"} { + Expect(os.WriteFile(filepath.Join(dir, name), nil, 0600)).To(Succeed()) + } + Expect(os.Mkdir(filepath.Join(dir, "secret-directory"), 0700)).To(Succeed()) + ml := &ModelLoader{ModelPath: dir} + var events []diagnostics.Event + ctx := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + calls := 0 + readDir := func(path string) ([]os.DirEntry, error) { + calls++ + Expect(events).To(BeEmpty()) // Emission happens after ReadDir, outside the measured operation. + return os.ReadDir(path) + } + plain, err := ml.listFilesInModelPathContext(context.Background(), readDir) + Expect(err).NotTo(HaveOccurred()) + got, err := ml.listFilesInModelPathContext(ctx, readDir) + Expect(err).NotTo(HaveOccurred()) + Expect(got).To(Equal(plain)) + Expect(got).To(Equal([]string{"secret-candidate", "z-model"})) + Expect(calls).To(Equal(2)) + Expect(events).To(HaveLen(3)) + Expect(events[0].Count).To(Equal(1)) + Expect(events[2].Count).To(Equal(7)) + Expect(events[2].Phase).To(Equal(diagnostics.PhaseLooseFilter)) + for _, e := range events { + Expect(e.ID).To(Equal(events[0].ID)) + Expect(e.Outcome).To(Equal(diagnostics.OutcomeOK)) + } + Expect(fmt.Sprint(events)).NotTo(ContainSubstring("secret")) + Expect(fmt.Sprint(events)).NotTo(ContainSubstring(dir)) + wrapped, err := ml.ListFilesInModelPath() + Expect(err).NotTo(HaveOccurred()) + Expect(wrapped).To(Equal(got)) + Expect(events).To(HaveLen(3)) + observed, err := ml.ListFilesInModelPathContext(ctx) + Expect(err).NotTo(HaveOccurred()) + Expect(observed).To(Equal(got)) + }) + It("preserves missing-directory errors and does not run filename filtering", func() { + ml := &ModelLoader{ModelPath: filepath.Join(GinkgoT().TempDir(), "secret-missing")} + var events []diagnostics.Event + ctx := diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) }).Request(context.Background()) + plain, plainErr := ml.ListFilesInModelPath() + got, err := ml.ListFilesInModelPathContext(ctx) + Expect(os.IsNotExist(err)).To(BeTrue()) + Expect(err.Error()).To(Equal(plainErr.Error())) + Expect(got).To(Equal(plain)) + Expect(events).To(HaveLen(1)) + Expect(events[0].Outcome).To(Equal(diagnostics.OutcomeError)) + Expect(events[0].Count).To(Equal(1)) + }) + It("adds no allocations for disabled timing on empty enumeration", func() { + ml := &ModelLoader{} + readDir := func(string) ([]os.DirEntry, error) { return nil, nil } + Expect(testing.AllocsPerRun(100, func() { _, _ = ml.listFilesInModelPathContext(context.Background(), readDir) })).To(BeZero()) + }) +})