|
1 | 1 | package gallery_test |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "bytes" |
4 | 5 | "context" |
| 6 | + "encoding/binary" |
| 7 | + "math" |
| 8 | + "net/http" |
| 9 | + "net/http/httptest" |
5 | 10 | "os" |
| 11 | + "path/filepath" |
| 12 | + "time" |
6 | 13 |
|
| 14 | + gguf "github.com/gpustack/gguf-parser-go" |
7 | 15 | . "github.com/onsi/ginkgo/v2" |
8 | 16 | . "github.com/onsi/gomega" |
| 17 | + "gopkg.in/yaml.v3" |
9 | 18 |
|
10 | 19 | "github.com/mudler/LocalAI/core/config" |
11 | 20 | "github.com/mudler/LocalAI/core/gallery" |
@@ -57,6 +66,46 @@ var _ = Describe("VRAM estimate warm-up", func() { |
57 | 66 | Consistently(func() bool { return true }, "100ms").Should(BeTrue()) |
58 | 67 | }) |
59 | 68 |
|
| 69 | + It("does not crash the server when remote GGUF metadata is malformed", func() { |
| 70 | + payload := warmMalformedGGUF() |
| 71 | + requested := make(chan struct{}) |
| 72 | + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 73 | + select { |
| 74 | + case <-requested: |
| 75 | + default: |
| 76 | + close(requested) |
| 77 | + } |
| 78 | + http.ServeContent(w, r, "model.gguf", time.Time{}, bytes.NewReader(payload)) |
| 79 | + })) |
| 80 | + DeferCleanup(server.Close) |
| 81 | + |
| 82 | + galleryPath := filepath.Join(state.Model.ModelsPath, "malformed-gallery.yaml") |
| 83 | + index, err := yaml.Marshal([]gallery.GalleryModel{{Metadata: gallery.Metadata{ |
| 84 | + Name: "malformed-gguf", |
| 85 | + AdditionalFiles: []gallery.File{{ |
| 86 | + Filename: "model.gguf", |
| 87 | + URI: server.URL + "/model.gguf", |
| 88 | + }}, |
| 89 | + }}}) |
| 90 | + Expect(err).NotTo(HaveOccurred()) |
| 91 | + Expect(os.WriteFile(galleryPath, index, 0600)).To(Succeed()) |
| 92 | + |
| 93 | + cfg := gallery.DefaultEstimateWarmConfig |
| 94 | + cfg.Limit = 1 |
| 95 | + cfg.Concurrency = 1 |
| 96 | + cfg.Contexts = []uint32{8192} |
| 97 | + gallery.WarmEstimateCache(context.Background(), []config.Gallery{{ |
| 98 | + Name: "malformed", |
| 99 | + URL: "file://" + galleryPath, |
| 100 | + }}, state, cfg) |
| 101 | + |
| 102 | + Eventually(requested, "2s").Should(BeClosed()) |
| 103 | + // The warm-up is detached. Give its parser time to consume the response; |
| 104 | + // before the recovery boundary, that goroutine panicked and killed the |
| 105 | + // entire test process (and the LocalAI server in production). |
| 106 | + Consistently(func() bool { return true }, "300ms").Should(BeTrue()) |
| 107 | + }) |
| 108 | + |
60 | 109 | Describe("configuration from the environment", func() { |
61 | 110 | AfterEach(func() { |
62 | 111 | os.Unsetenv("LOCALAI_VRAM_WARM_LIMIT") |
@@ -113,3 +162,19 @@ var _ = Describe("VRAM estimate warm-up", func() { |
113 | 162 | }) |
114 | 163 |
|
115 | 164 | }) |
| 165 | + |
| 166 | +func warmMalformedGGUF() []byte { |
| 167 | + payload := make([]byte, 0, 128) |
| 168 | + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMagicGGUFLe)) |
| 169 | + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFVersionV3)) |
| 170 | + payload = binary.LittleEndian.AppendUint64(payload, 0) |
| 171 | + payload = binary.LittleEndian.AppendUint64(payload, 1) |
| 172 | + key := "tokenizer.ggml.tokens" |
| 173 | + payload = binary.LittleEndian.AppendUint64(payload, uint64(len(key))) |
| 174 | + payload = append(payload, key...) |
| 175 | + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeArray)) |
| 176 | + payload = binary.LittleEndian.AppendUint32(payload, uint32(gguf.GGUFMetadataValueTypeString)) |
| 177 | + payload = binary.LittleEndian.AppendUint64(payload, 1) |
| 178 | + payload = binary.LittleEndian.AppendUint64(payload, math.MaxUint64) |
| 179 | + return payload |
| 180 | +} |
0 commit comments