diff --git a/backend/go/crispasr/gocrispasr.go b/backend/go/crispasr/gocrispasr.go index 1fd64b08e943..be431165dd82 100644 --- a/backend/go/crispasr/gocrispasr.go +++ b/backend/go/crispasr/gocrispasr.go @@ -67,7 +67,16 @@ const defaultTTSSampleRate = 24000 // resampling, so the WAV header must match it. Returns ok=false for non-piper // models (key absent) or an unreadable file, letting the caller fall back to // defaultTTSSampleRate. -func piperSampleRate(modelPath string) (int, bool) { +func piperSampleRate(modelPath string) (rate int, ok bool) { + // A malformed metadata length can make gguf-parser-go panic before it can + // return an error. Keep a bad voice file from crash-looping the backend. + defer func() { + if recover() != nil { + rate = 0 + ok = false + } + }() + // Only scalar architecture keys are read, so skip the large array metadata // (phoneme map) and mmap the header - same rationale as pkg/vram's reader. f, err := gguf.ParseGGUFFile(modelPath, gguf.UseMMap(), gguf.SkipLargeMetadata()) @@ -78,7 +87,7 @@ func piperSampleRate(modelPath string) (int, bool) { if !ok || kv.ValueType != gguf.GGUFMetadataValueTypeUint32 { return 0, false } - rate := int(kv.ValueUint32()) + rate = int(kv.ValueUint32()) if rate <= 0 { return 0, false } diff --git a/backend/go/crispasr/gocrispasr_samplerate_test.go b/backend/go/crispasr/gocrispasr_samplerate_test.go index 6b0cf726b5fd..c36e7a75a284 100644 --- a/backend/go/crispasr/gocrispasr_samplerate_test.go +++ b/backend/go/crispasr/gocrispasr_samplerate_test.go @@ -3,6 +3,7 @@ package main import ( "bytes" "encoding/binary" + "math" "os" "path/filepath" @@ -102,6 +103,24 @@ var _ = Describe("piper sample rate", func() { _, ok := piperSampleRate(p) Expect(ok).To(BeFalse()) }) + + It("returns ok=false instead of panicking on a malformed string length", func() { + p := filepath.Join(GinkgoT().TempDir(), "malformed.gguf") + var b bytes.Buffer + b.WriteString("GGUF") + Expect(binary.Write(&b, binary.LittleEndian, uint32(3))).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(0))).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(1))).To(Succeed()) + key := "general.name" + Expect(binary.Write(&b, binary.LittleEndian, uint64(len(key)))).To(Succeed()) + b.WriteString(key) + Expect(binary.Write(&b, binary.LittleEndian, ggufTypeString)).To(Succeed()) + Expect(binary.Write(&b, binary.LittleEndian, uint64(math.MaxInt64))).To(Succeed()) + Expect(os.WriteFile(p, b.Bytes(), 0o644)).To(Succeed()) + + _, ok := piperSampleRate(p) + Expect(ok).To(BeFalse()) + }) }) // End-to-end through the built .so. Gated on CRISPASR_PIPER_MODEL_PATH (a diff --git a/core/gallery/estimate_warm.go b/core/gallery/estimate_warm.go index 680ebcdbb0ad..4b291cd589d4 100644 --- a/core/gallery/estimate_warm.go +++ b/core/gallery/estimate_warm.go @@ -9,6 +9,7 @@ import ( "time" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/concurrency" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/xlog" @@ -101,7 +102,7 @@ func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemSt return } - go func() { + concurrency.SafeGo(func() { started := time.Now() models, err := AvailableGalleryModelsCached(galleries, systemState) @@ -131,7 +132,7 @@ func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemSt for i := 0; i < cfg.Concurrency; i++ { wg.Add(1) - go func() { + concurrency.SafeGo(func() { defer wg.Done() for m := range cursor { // Per entry, not for the run: one unreachable weight file @@ -164,7 +165,7 @@ func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemSt cancel() } - }() + }) } feed: @@ -183,7 +184,7 @@ func WarmEstimateCache(ctx context.Context, galleries []config.Gallery, systemSt return } xlog.Info("gallery caches warmed", "estimates", warmed, "variants", warmedVariants, "of", len(models), "took", time.Since(started).Round(time.Second)) - }() + }) } // EstimateWarmConfigFromEnv reads the warm-up bounds from the environment, diff --git a/core/gallery/estimate_warm_test.go b/core/gallery/estimate_warm_test.go index 7f75bb412ff1..e247f4bfce36 100644 --- a/core/gallery/estimate_warm_test.go +++ b/core/gallery/estimate_warm_test.go @@ -1,11 +1,20 @@ package gallery_test import ( + "bytes" "context" + "encoding/binary" + "math" + "net/http" + "net/http/httptest" "os" + "path/filepath" + "time" + gguf "github.com/gpustack/gguf-parser-go" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" @@ -57,6 +66,46 @@ var _ = Describe("VRAM estimate warm-up", func() { Consistently(func() bool { return true }, "100ms").Should(BeTrue()) }) + It("does not crash the server when remote GGUF metadata is malformed", func() { + payload := warmMalformedGGUF() + requested := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-requested: + default: + close(requested) + } + http.ServeContent(w, r, "model.gguf", time.Time{}, bytes.NewReader(payload)) + })) + DeferCleanup(server.Close) + + galleryPath := filepath.Join(state.Model.ModelsPath, "malformed-gallery.yaml") + index, err := yaml.Marshal([]gallery.GalleryModel{{Metadata: gallery.Metadata{ + Name: "malformed-gguf", + AdditionalFiles: []gallery.File{{ + Filename: "model.gguf", + URI: server.URL + "/model.gguf", + }}, + }}}) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(galleryPath, index, 0600)).To(Succeed()) + + cfg := gallery.DefaultEstimateWarmConfig + cfg.Limit = 1 + cfg.Concurrency = 1 + cfg.Contexts = []uint32{8192} + gallery.WarmEstimateCache(context.Background(), []config.Gallery{{ + Name: "malformed", + URL: "file://" + galleryPath, + }}, state, cfg) + + Eventually(requested, "2s").Should(BeClosed()) + // The warm-up is detached. Give its parser time to consume the response; + // before the recovery boundary, that goroutine panicked and killed the + // entire test process (and the LocalAI server in production). + Consistently(func() bool { return true }, "300ms").Should(BeTrue()) + }) + Describe("configuration from the environment", func() { AfterEach(func() { os.Unsetenv("LOCALAI_VRAM_WARM_LIMIT") @@ -113,3 +162,19 @@ var _ = Describe("VRAM estimate warm-up", func() { }) }) + +func warmMalformedGGUF() []byte { + payload := make([]byte, 0, 128) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMagicGGUFLe)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFVersionV3)) + payload = binary.LittleEndian.AppendUint64(payload, 0) + payload = binary.LittleEndian.AppendUint64(payload, 1) + key := "tokenizer.ggml.tokens" + payload = binary.LittleEndian.AppendUint64(payload, uint64(len(key))) + payload = append(payload, key...) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeArray)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) + payload = binary.LittleEndian.AppendUint64(payload, 1) + payload = binary.LittleEndian.AppendUint64(payload, math.MaxUint64) + return payload +} diff --git a/core/gallery/importers/llama-cpp.go b/core/gallery/importers/llama-cpp.go index 0804ce34f904..923437c89062 100644 --- a/core/gallery/importers/llama-cpp.go +++ b/core/gallery/importers/llama-cpp.go @@ -401,7 +401,10 @@ func maybeApplyMTPDefaults(modelConfig *config.ModelConfig, details Details, cfg } }() - f, err := gguf.ParseGGUFFileRemote(ctx, probeURL) + // MTP markers are architecture scalars. Avoid allocating tokenizer and + // other large arrays from an untrusted remote header; panic recovery cannot + // contain a fatal out-of-memory condition. + f, err := gguf.ParseGGUFFileRemote(ctx, probeURL, gguf.SkipLargeMetadata()) if err != nil { xlog.Debug("[mtp-importer] failed to read remote GGUF header for MTP detection", "uri", probeURL, "error", err) return diff --git a/docker-compose.distributed.yaml b/docker-compose.distributed.yaml index 9971b57c3fcf..3387e313415b 100644 --- a/docker-compose.distributed.yaml +++ b/docker-compose.distributed.yaml @@ -74,6 +74,9 @@ services: GODEBUG: "netdns=go" # Paths MODELS_PATH: /models + # Avoid probing remote gallery GGUF metadata during container startup. + # Remove this line or set a positive limit to opt back into cache warming. + LOCALAI_VRAM_WARM_LIMIT: "0" volumes: - frontend_models:/models - frontend_data:/data diff --git a/docker-compose.yaml b/docker-compose.yaml index a432a699a3d0..ee137e83c4f9 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -18,6 +18,9 @@ services: - .env environment: - MODELS_PATH=/models + # Avoid probing remote gallery GGUF metadata during container startup. + # Remove this line or set a positive limit to opt back into cache warming. + - LOCALAI_VRAM_WARM_LIMIT=0 # - DEBUG=true ## Agents (LocalAGI) - https://localai.io/features/agents/ # - LOCALAI_DISABLE_AGENTS=false diff --git a/docs/content/advanced/vram-management.md b/docs/content/advanced/vram-management.md index dda39ba07409..229d39f9f4a9 100644 --- a/docs/content/advanced/vram-management.md +++ b/docs/content/advanced/vram-management.md @@ -477,6 +477,11 @@ then on. | `LOCALAI_VRAM_WARM_LIMIT` | `300` | How many gallery entries to warm at startup, estimates and variants alike. Set to `0` to disable the warm-up entirely. | | `LOCALAI_VRAM_WARM_CONCURRENCY` | `4` | How many estimates to run at once. | +The provided Docker Compose configurations set `LOCALAI_VRAM_WARM_LIMIT=0` +as a defensive default, so container startup does not probe remote GGUF files. +Remove that override or set it to a positive number to opt into background +warming. + ```bash # Air-gapped, or you would rather not make the requests at all LOCALAI_VRAM_WARM_LIMIT=0 local-ai run diff --git a/go.mod b/go.mod index eced8c0916cb..0daf7f61201a 100644 --- a/go.mod +++ b/go.mod @@ -24,7 +24,7 @@ require ( github.com/gofrs/flock v0.13.0 github.com/google/go-containerregistry v0.21.6 github.com/google/uuid v1.6.0 - github.com/gpustack/gguf-parser-go v0.24.0 + github.com/gpustack/gguf-parser-go v0.25.0 github.com/hpcloud/tail v1.0.0 github.com/ipfs/go-log v1.0.5 github.com/jaypipes/ghw v0.24.0 diff --git a/go.sum b/go.sum index 61ccede93c39..f582b871d22f 100644 --- a/go.sum +++ b/go.sum @@ -666,8 +666,8 @@ github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 h1:JeSE6pjso5THxAzdVpqr6/geYxZytqFMBCOtn/ujyeo= github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674/go.mod h1:r4w70xmWCQKmi1ONH4KIaBptdivuRPyosB9RmPlGEwA= -github.com/gpustack/gguf-parser-go v0.24.0 h1:tdJceXYp9e5RhE9RwVYIuUpir72Jz2D68NEtDXkKCKc= -github.com/gpustack/gguf-parser-go v0.24.0/go.mod h1:y4TwTtDqFWTK+xvprOjRUh+dowgU2TKCX37vRKvGiZ0= +github.com/gpustack/gguf-parser-go v0.25.0 h1:1AMBhMKtI24nTtn588Bq53FqNiOvEw1x9Nb4HbRrThs= +github.com/gpustack/gguf-parser-go v0.25.0/go.mod h1:y4TwTtDqFWTK+xvprOjRUh+dowgU2TKCX37vRKvGiZ0= github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI= github.com/grpc-ecosystem/go-grpc-middleware v1.4.0/go.mod h1:g5qyo/la0ALbONm6Vbp88Yd8NsDy6rZz+RcrMPxvld8= github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw= diff --git a/pkg/vram/gguf_reader.go b/pkg/vram/gguf_reader.go index f2842ba83c2b..9e224d110cf5 100644 --- a/pkg/vram/gguf_reader.go +++ b/pkg/vram/gguf_reader.go @@ -2,6 +2,7 @@ package vram import ( "context" + "fmt" "strings" gguf "github.com/gpustack/gguf-parser-go" @@ -10,7 +11,18 @@ import ( type defaultGGUFReader struct{} -func (defaultGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFMeta, error) { +func (defaultGGUFReader) ReadMetadata(ctx context.Context, uri string) (meta *GGUFMeta, err error) { + // gguf-parser-go parses lengths supplied by the file and has historically + // panicked on values that cannot fit in a Go slice. Metadata can come from + // an untrusted remote host, and this reader is also used by a background + // gallery worker, where an escaped panic would terminate the whole server. + defer func() { + if recovered := recover(); recovered != nil { + meta = nil + err = fmt.Errorf("read GGUF metadata: parser panic: %v", recovered) + } + }() + u := downloader.URI(uri) urlStr := u.ResolveURL() @@ -28,7 +40,10 @@ func (defaultGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFMet if !u.LooksLikeHTTPURL() { return nil, nil } - f, err := gguf.ParseGGUFFileRemote(ctx, urlStr) + // The estimator only consumes architecture scalars. Tokenizer arrays can + // be very large and are unnecessary here, so avoid downloading or + // allocating them for remote files just as the local path does above. + f, err := gguf.ParseGGUFFileRemote(ctx, urlStr, gguf.SkipLargeMetadata()) if err != nil { return nil, err } diff --git a/pkg/vram/gguf_reader_test.go b/pkg/vram/gguf_reader_test.go new file mode 100644 index 000000000000..f1cd4a21eb14 --- /dev/null +++ b/pkg/vram/gguf_reader_test.go @@ -0,0 +1,115 @@ +package vram_test + +import ( + "bytes" + "context" + "encoding/binary" + "math" + "net/http" + "net/http/httptest" + "time" + + gguf "github.com/gpustack/gguf-parser-go" + "github.com/mudler/LocalAI/pkg/vram" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("DefaultGGUFReader", func() { + It("reads architecture scalars from a valid remote GGUF", func() { + server := serveGGUF(validRemoteGGUF()) + + meta, err := vram.DefaultGGUFReader().ReadMetadata(context.Background(), server.URL+"/model.gguf") + + Expect(err).NotTo(HaveOccurred()) + Expect(meta).To(Equal(&vram.GGUFMeta{ + BlockCount: 32, + EmbeddingLength: 4096, + HeadCount: 32, + HeadCountKV: 8, + MaximumContextLength: 8192, + })) + }) + + It("rejects an overflowing tokenizer array without allocating it", func() { + server := serveGGUF(malformedGGUFArray(math.MaxUint64)) + + _, err := vram.DefaultGGUFReader().ReadMetadata(context.Background(), server.URL+"/model.gguf") + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).NotTo(ContainSubstring("parser panic"), + "large tokenizer metadata should be skipped with a bounds error") + }) + + It("converts a parser panic from malformed string metadata to an error", func() { + server := serveGGUF(malformedGGUFString(uint64(math.MaxInt64))) + + _, err := vram.DefaultGGUFReader().ReadMetadata(context.Background(), server.URL+"/model.gguf") + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("parser panic")) + }) +}) + +func serveGGUF(payload []byte) *httptest.Server { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.ServeContent(w, r, "model.gguf", time.Time{}, bytes.NewReader(payload)) + })) + DeferCleanup(server.Close) + return server +} + +func malformedGGUFString(length uint64) []byte { + payload := ggufHeader(1) + payload = appendGGUFString(payload, "general.name") + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) + payload = binary.LittleEndian.AppendUint64(payload, length) + return payload +} + +func validRemoteGGUF() []byte { + payload := ggufHeader(6) + payload = appendGGUFStringValue(payload, "general.architecture", "llama") + payload = appendGGUFUint32(payload, "llama.block_count", 32) + payload = appendGGUFUint32(payload, "llama.embedding_length", 4096) + payload = appendGGUFUint32(payload, "llama.attention.head_count", 32) + payload = appendGGUFUint32(payload, "llama.attention.head_count_kv", 8) + payload = appendGGUFUint32(payload, "llama.context_length", 8192) + return payload +} + +func malformedGGUFArray(itemLength uint64) []byte { + payload := ggufHeader(1) + payload = appendGGUFString(payload, "tokenizer.ggml.tokens") + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeArray)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) + payload = binary.LittleEndian.AppendUint64(payload, 1) + payload = binary.LittleEndian.AppendUint64(payload, itemLength) + return payload +} + +func ggufHeader(metadataCount uint64) []byte { + payload := make([]byte, 0, 128) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMagicGGUFLe)) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFVersionV3)) + payload = binary.LittleEndian.AppendUint64(payload, 0) + payload = binary.LittleEndian.AppendUint64(payload, metadataCount) + return payload +} + +func appendGGUFString(payload []byte, value string) []byte { + payload = binary.LittleEndian.AppendUint64(payload, uint64(len(value))) + return append(payload, value...) +} + +func appendGGUFStringValue(payload []byte, key, value string) []byte { + payload = appendGGUFString(payload, key) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) + return appendGGUFString(payload, value) +} + +func appendGGUFUint32(payload []byte, key string, value uint32) []byte { + payload = appendGGUFString(payload, key) + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeUint32)) + return binary.LittleEndian.AppendUint32(payload, value) +}