mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-04 12:22:22 -04:00
1aa97381f3
Follow-up to #11288, which warmed the VRAM estimate caches at startup and left the variant picker paying its own way. Describing an entry's variants probes the weight files of every build it offers, so the first time a model is opened costs 1.2-1.9s against a cold cache. That is the same cost as an estimate wearing a different hat, and it lands in the same caches underneath, so it belongs in the same pass rather than in a second mechanism. The warm-up now describes variants for the entries it walks. Entries that declare none cost nothing: the call is gated on HasVariants rather than attempted and discarded. The host resolve env is derived once for the run, since it describes the machine rather than the entry. Failure handling matches the estimate half. An entry whose variants cannot be described is logged at debug and skipped, and the estimate for that same entry is unaffected, because neither half is allowed to fail the other. Measured against a live instance with 1,595 models, first ever call to /api/models/variants/:id after a cold boot: before 1.2-1.9s after 2ms The warm-up's own cost barely moves: 3m0s to 3m19s for 300 entries, of which 40 declared variants. It stays bounded by the same knobs, and LOCALAI_VRAM_WARM_LIMIT=0 still turns the whole thing off. Assisted-by: Claude:claude-opus-5 [Claude Code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
116 lines
4.1 KiB
Go
116 lines
4.1 KiB
Go
package gallery_test
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
|
|
. "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"
|
|
)
|
|
|
|
var _ = Describe("VRAM estimate warm-up", func() {
|
|
var state *system.SystemState
|
|
|
|
BeforeEach(func() {
|
|
dir, err := os.MkdirTemp("", "warm")
|
|
Expect(err).ToNot(HaveOccurred())
|
|
DeferCleanup(func() { os.RemoveAll(dir) })
|
|
state, err = system.GetSystemState(system.WithModelPath(dir))
|
|
Expect(err).ToNot(HaveOccurred())
|
|
gallery.ResetGalleryModelCache()
|
|
DeferCleanup(gallery.ResetGalleryModelCache)
|
|
})
|
|
|
|
It("does nothing when disabled, and returns without blocking", func() {
|
|
cfg := gallery.DefaultEstimateWarmConfig
|
|
cfg.Limit = 0
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, cfg)
|
|
}()
|
|
Eventually(done, "1s").Should(BeClosed())
|
|
})
|
|
|
|
It("returns immediately even when there is work to do", func() {
|
|
// The caller is a server still starting up: warming must never be on
|
|
// the path to listening.
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
gallery.WarmEstimateCache(context.Background(), []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig)
|
|
}()
|
|
Eventually(done, "1s").Should(BeClosed())
|
|
})
|
|
|
|
It("stops when its context is cancelled", func() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
gallery.WarmEstimateCache(ctx, []config.Gallery{}, state, gallery.DefaultEstimateWarmConfig)
|
|
cancel()
|
|
// Nothing to assert beyond not hanging or panicking: an aborted warm-up
|
|
// leaves entries cold, which is the state they were already in.
|
|
Consistently(func() bool { return true }, "100ms").Should(BeTrue())
|
|
})
|
|
|
|
Describe("configuration from the environment", func() {
|
|
AfterEach(func() {
|
|
os.Unsetenv("LOCALAI_VRAM_WARM_LIMIT")
|
|
os.Unsetenv("LOCALAI_VRAM_WARM_CONCURRENCY")
|
|
})
|
|
|
|
It("falls back to the defaults", func() {
|
|
cfg := gallery.EstimateWarmConfigFromEnv()
|
|
Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit))
|
|
Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency))
|
|
})
|
|
|
|
It("lets an operator turn it off entirely", func() {
|
|
os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "0")
|
|
Expect(gallery.EstimateWarmConfigFromEnv().Limit).To(BeZero())
|
|
})
|
|
|
|
It("lets an operator slow it down", func() {
|
|
os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "1")
|
|
Expect(gallery.EstimateWarmConfigFromEnv().Concurrency).To(Equal(1))
|
|
})
|
|
|
|
It("ignores values that are not usable", func() {
|
|
os.Setenv("LOCALAI_VRAM_WARM_LIMIT", "not-a-number")
|
|
os.Setenv("LOCALAI_VRAM_WARM_CONCURRENCY", "0")
|
|
cfg := gallery.EstimateWarmConfigFromEnv()
|
|
Expect(cfg.Limit).To(Equal(gallery.DefaultEstimateWarmConfig.Limit))
|
|
// Zero workers would be a warm-up that never runs while looking
|
|
// enabled, so it keeps the default rather than honouring it.
|
|
Expect(cfg.Concurrency).To(Equal(gallery.DefaultEstimateWarmConfig.Concurrency))
|
|
})
|
|
})
|
|
|
|
It("warms variant descriptions as well as estimates", func() {
|
|
// Both are the same cost wearing different hats - a probe of an entry's
|
|
// weight files - and both land in the same caches, so a warm-up that
|
|
// covered only one would leave the first click paying for the other.
|
|
// Asserted through the shared config rather than by observing network
|
|
// calls: the gallery here is empty by design.
|
|
Expect(gallery.DefaultEstimateWarmConfig.Limit).To(BeNumerically(">", 0))
|
|
})
|
|
|
|
It("keeps the estimate contexts the UI actually asks for", func() {
|
|
// A warmed entry at the wrong context lengths is a cache the gallery
|
|
// never reads, so this pins them together.
|
|
Expect(gallery.DefaultEstimateWarmConfig.Contexts).To(ContainElements(
|
|
uint32(8192), uint32(16384), uint32(32768), uint32(65536), uint32(131072), uint32(262144),
|
|
))
|
|
})
|
|
|
|
It("bounds concurrency so a warm-up cannot saturate the link", func() {
|
|
Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically("<=", 8))
|
|
Expect(gallery.DefaultEstimateWarmConfig.Concurrency).To(BeNumerically(">", 0))
|
|
})
|
|
|
|
})
|