diff --git a/core/gallery/gallery.go b/core/gallery/gallery.go index bd054b0ff..68d6de9e8 100644 --- a/core/gallery/gallery.go +++ b/core/gallery/gallery.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "os" - "path/filepath" "slices" "strings" "sync" @@ -19,6 +18,7 @@ import ( "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/LocalAI/pkg/xsync" "github.com/mudler/xlog" + "golang.org/x/sync/singleflight" "gopkg.in/yaml.v3" ) @@ -276,13 +276,12 @@ func FindGalleryElement[T GalleryElement](models []T, name string) T { func AvailableGalleryModels(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryModel], error) { var models []*GalleryModel + isInstalled := installedConfigs(systemState.Model.ModelsPath) + // Get models from galleries for _, gallery := range galleries { galleryModels, err := getGalleryElements(gallery, systemState.Model.ModelsPath, systemState.RequireBackendIntegrity, func(model *GalleryModel) bool { - if _, err := os.Stat(filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", model.GetName()))); err == nil { - return true - } - return false + return isInstalled(model.GetName()) }) if err != nil { return nil, err @@ -351,6 +350,7 @@ var ( // same cache-defeating loop the refresh interval exists to stop. availableModelsLoaded bool refreshing atomic.Bool + coldLoad singleflight.Group galleryGeneration atomic.Uint64 lastRefreshUnixNano atomic.Int64 ) @@ -429,12 +429,15 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste availableModelsMu.RUnlock() if loaded { + // The directory is read before taking the lock. Held across the + // filesystem work, the lock serialized every caller behind it, and a + // page view is dozens of concurrent callers. + isInstalled := installedConfigs(systemState.Model.ModelsPath) // Refresh installed status under write lock to avoid races with // concurrent readers and the background refresh goroutine. availableModelsMu.Lock() for _, m := range cached { - _, err := os.Stat(filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", m.GetName()))) - m.SetInstalled(err == nil) + m.SetInstalled(isInstalled(m.GetName())) } availableModelsMu.Unlock() // Trigger a background refresh if one is not already running. @@ -442,20 +445,29 @@ func AvailableGalleryModelsCached(galleries []config.Gallery, systemState *syste return cached, nil } - // No cache yet — must do a blocking load. - models, err := AvailableGalleryModels(galleries, systemState) + // No cache yet, so the load blocks. Callers arriving while it runs wait + // for it instead of each starting their own: a page view on a fresh + // server is the listing plus one estimate per row at once, and each load + // fetches the gallery index and every config it references. + v, err, _ := coldLoad.Do("gallery", func() (any, error) { + models, err := AvailableGalleryModels(galleries, systemState) + if err != nil { + return nil, err + } + + availableModelsMu.Lock() + availableModelsCache = models + availableModelsLoaded = true + galleryGeneration.Add(1) + availableModelsMu.Unlock() + lastRefreshUnixNano.Store(time.Now().UnixNano()) + + return models, nil + }) if err != nil { return nil, err } - - availableModelsMu.Lock() - availableModelsCache = models - availableModelsLoaded = true - galleryGeneration.Add(1) - availableModelsMu.Unlock() - lastRefreshUnixNano.Store(time.Now().UnixNano()) - - return models, nil + return v.(GalleryElements[*GalleryModel]), nil } // triggerGalleryRefresh starts a background goroutine that refreshes the diff --git a/core/gallery/gallery_installed_scan_test.go b/core/gallery/gallery_installed_scan_test.go new file mode 100644 index 000000000..7bcf22a1f --- /dev/null +++ b/core/gallery/gallery_installed_scan_test.go @@ -0,0 +1,138 @@ +package gallery_test + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/system" +) + +// The models directory is often network storage (SMB, NFS), where every +// filesystem call is a round trip. The cached listing is read by the gallery +// page and by one VRAM estimate per row, so whatever it costs is paid dozens +// of times per page view. +var _ = Describe("Gallery cache installed status", func() { + const index = ` +- name: plain + backend: llama-cpp +- name: linked + backend: llama-cpp +- name: dangling + backend: llama-cpp +- name: later + backend: llama-cpp +- name: absent + backend: llama-cpp +` + + var ( + modelsDir string + state *system.SystemState + galleries []config.Gallery + hits atomic.Int32 + delay time.Duration + ) + + BeforeEach(func() { + var err error + modelsDir, err = os.MkdirTemp("", "gallery-installed") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = os.RemoveAll(modelsDir) }) + state, err = system.GetSystemState(system.WithModelPath(modelsDir)) + Expect(err).ToNot(HaveOccurred()) + + hits.Store(0) + delay = 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + time.Sleep(delay) + _, _ = w.Write([]byte(index)) + })) + DeferCleanup(server.Close) + galleries = []config.Gallery{{Name: "test", URL: server.URL + "/index.yaml"}} + + gallery.ResetGalleryModelCache() + DeferCleanup(gallery.ResetGalleryModelCache) + }) + + installed := func(models gallery.GalleryElements[*gallery.GalleryModel]) map[string]bool { + out := map[string]bool{} + for _, m := range models { + out[m.Name] = m.Installed + } + return out + } + + It("reports what os.Stat would, for files, symlinks and dangling symlinks", func() { + Expect(os.WriteFile(filepath.Join(modelsDir, "plain.yaml"), []byte("name: plain\n"), 0o644)).To(Succeed()) + target := filepath.Join(modelsDir, "target.txt") + Expect(os.WriteFile(target, []byte("name: linked\n"), 0o644)).To(Succeed()) + Expect(os.Symlink(target, filepath.Join(modelsDir, "linked.yaml"))).To(Succeed()) + Expect(os.Symlink(filepath.Join(modelsDir, "missing"), filepath.Join(modelsDir, "dangling.yaml"))).To(Succeed()) + + // Both the blocking first load and the cached path set the flag, and + // they must agree. + for range 2 { + models, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + Expect(installed(models)).To(Equal(map[string]bool{ + "plain": true, + "linked": true, + "dangling": false, + "later": false, + "absent": false, + })) + } + }) + + It("picks up a config written after the gallery was cached", func() { + _, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + + Expect(os.WriteFile(filepath.Join(modelsDir, "later.yaml"), []byte("name: later\n"), 0o644)).To(Succeed()) + + models, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + Expect(installed(models)).To(HaveKeyWithValue("later", true)) + }) + + It("reports nothing installed when the models directory is gone", func() { + _, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + Expect(os.RemoveAll(modelsDir)).To(Succeed()) + + models, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + Expect(installed(models)).To(HaveEach(BeFalse())) + }) + + It("shares one upstream load between concurrent callers on a cold cache", func() { + // Slow enough that every caller arrives while the first load is still + // in flight, which is what a page view does to a freshly started + // server: the listing and every row's estimate at once. + delay = 300 * time.Millisecond + + var wg sync.WaitGroup + for range 8 { + wg.Go(func() { + defer GinkgoRecover() + models, err := gallery.AvailableGalleryModelsCached(galleries, state) + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(HaveLen(5)) + }) + } + wg.Wait() + + Expect(hits.Load()).To(Equal(int32(1))) + }) +}) diff --git a/core/gallery/installed_configs.go b/core/gallery/installed_configs.go new file mode 100644 index 000000000..fb36e5d29 --- /dev/null +++ b/core/gallery/installed_configs.go @@ -0,0 +1,61 @@ +package gallery + +import ( + "errors" + "io/fs" + "os" + "path/filepath" + "strings" +) + +const modelConfigExt = ".yaml" + +// installedConfigs answers "does /.yaml exist?" for every +// entry of a gallery from a single read of the models directory. +// +// The question used to be asked with one os.Stat per gallery entry. The gallery +// holds thousands of entries and the models directory is often network storage +// (SMB, NFS), where each Stat is a round trip, so one listing cost seconds. The +// listing is read by the gallery page and by one VRAM estimate per row, which +// turned a page view into minutes. +// +// Answers match os.Stat on the same path: a symlink counts only when its target +// exists, and anything else carrying the name counts, directories included. +// Names that are not a plain file name are checked with os.Stat directly, since +// they point outside the listed directory. +func installedConfigs(modelsPath string) func(name string) bool { + statInstalled := func(name string) bool { + _, err := os.Stat(filepath.Join(modelsPath, name+modelConfigExt)) + return err == nil + } + + entries, err := os.ReadDir(modelsPath) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + return func(string) bool { return false } + } + // A directory that exists but cannot be listed may still answer a + // Stat, so fall back rather than report everything as not installed. + return statInstalled + } + + present := make(map[string]struct{}, len(entries)) + for _, e := range entries { + base, ok := strings.CutSuffix(e.Name(), modelConfigExt) + if !ok { + continue + } + if e.Type()&fs.ModeSymlink != 0 && !statInstalled(base) { + continue + } + present[base] = struct{}{} + } + + return func(name string) bool { + if strings.ContainsRune(name, '/') || strings.ContainsRune(name, filepath.Separator) { + return statInstalled(name) + } + _, ok := present[name] + return ok + } +}