mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
fix(vram): persist remote probe metadata (#11487)
* fix(vram): persist remote probe metadata The startup warmer repeated remote size and GGUF metadata probes after every restart because both caches lived only in memory. Store successful HTTP probes for 24 hours so frequent restarts reuse the prior results. Bound the cache, reject invalid records, and purge it when gallery data changes. Local model files continue to bypass persistence. Assisted-by: Codex:gpt-5 * fix(vram): check temporary file cleanup The lint gate rejects the unchecked cleanup call in the persistent cache writer. Assisted-by: Codex:gpt-5.6 [golangci-lint] * fix(vram): make persistent cache optional Remote metadata probes can transfer enough data that operators need control over disk reuse and startup warming. Gallery autoload now gates both behaviors, and the runtime setting applies changes immediately. Assisted-by: Codex:gpt-5 * fix(ui): expose gallery startup pre-warm The existing gallery autoload setting also gates the startup metadata warmer. Name both effects in Settings so operators can find the requested boot control. Assisted-by: Codex:gpt-5 --------- Co-authored-by: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com>
This commit is contained in:
1 parent
ffef539866
commit
8f56e4e042
16 files changed
+656
-12
No files matched your search
+333
-2
@@ -2,9 +2,24 @@ package vram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const persistentCacheEntryLimit = 4096
|
||||
const persistentCacheVersion = 1
|
||||
|
||||
var defaultPersistentGuard = &persistentGenerationGuard{}
|
||||
var defaultCacheMu sync.RWMutex
|
||||
|
||||
// galleryGenFunc returns the current gallery generation counter.
|
||||
// When set, cache entries are invalidated when the generation changes.
|
||||
// When nil (e.g., in tests or non-gallery contexts), entries never expire.
|
||||
@@ -23,6 +38,149 @@ func currentGeneration() uint64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// ConfigurePersistentCache replaces the process-wide estimator caches with
|
||||
// instances that reuse successful remote probes across server restarts.
|
||||
func ConfigurePersistentCache(dir string, ttl time.Duration) {
|
||||
defaultCacheMu.Lock()
|
||||
defer defaultCacheMu.Unlock()
|
||||
removeAbandonedPersistentTemps(dir)
|
||||
prunePersistentEntries(dir, ttl, persistentCacheEntryLimit)
|
||||
guard := &persistentGenerationGuard{dir: dir}
|
||||
defaultPersistentGuard = guard
|
||||
defaultCachedSizeResolver = newCachedSizeResolverWithGuard(defaultSizeResolver{}, dir, ttl, guard)
|
||||
defaultCachedGGUFReader = newCachedGGUFReaderWithGuard(defaultGGUFReader{}, dir, ttl, guard)
|
||||
}
|
||||
|
||||
// DisablePersistentCache keeps process-local caching but stops disk reads and writes.
|
||||
func DisablePersistentCache() {
|
||||
defaultCacheMu.Lock()
|
||||
defer defaultCacheMu.Unlock()
|
||||
defaultPersistentGuard = &persistentGenerationGuard{}
|
||||
defaultCachedSizeResolver = newCachedSizeResolver(defaultSizeResolver{}, "", 0)
|
||||
defaultCachedGGUFReader = newCachedGGUFReader(defaultGGUFReader{}, "", 0)
|
||||
}
|
||||
|
||||
// InvalidatePersistentCache removes remote probe results after the gallery
|
||||
// changes, including when no estimate is requested before the next restart.
|
||||
func InvalidatePersistentCache() {
|
||||
defaultCacheMu.RLock()
|
||||
guard := defaultPersistentGuard
|
||||
defaultCacheMu.RUnlock()
|
||||
guard.invalidate(currentGeneration())
|
||||
}
|
||||
|
||||
func removeAbandonedPersistentTemps(dir string) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.Type().IsRegular() && strings.HasPrefix(entry.Name(), ".vram-") {
|
||||
_ = os.Remove(filepath.Join(dir, entry.Name()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func removePersistentEntries(dir string) {
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.Type().IsRegular() && (strings.HasPrefix(name, "size-") || strings.HasPrefix(name, "gguf-")) {
|
||||
_ = os.Remove(filepath.Join(dir, name))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type persistentGenerationGuard struct {
|
||||
mu sync.Mutex
|
||||
dir string
|
||||
generation uint64
|
||||
set bool
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) invalidate(generation uint64) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
removePersistentEntries(g.dir)
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) canRead(generation uint64) bool {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if !g.set {
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
return true
|
||||
}
|
||||
if g.generation == generation {
|
||||
return true
|
||||
}
|
||||
removePersistentEntries(g.dir)
|
||||
g.generation = generation
|
||||
return false
|
||||
}
|
||||
|
||||
func (g *persistentGenerationGuard) persist(generation uint64, write func()) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if !g.set {
|
||||
g.generation = generation
|
||||
g.set = true
|
||||
}
|
||||
if g.generation == generation {
|
||||
write()
|
||||
}
|
||||
}
|
||||
|
||||
func prunePersistentEntries(dir string, ttl time.Duration, limit int) {
|
||||
if dir == "" || ttl <= 0 {
|
||||
return
|
||||
}
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
type cacheFile struct {
|
||||
path string
|
||||
modTime time.Time
|
||||
}
|
||||
files := make([]cacheFile, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.Type().IsRegular() && (strings.HasPrefix(name, "size-") || strings.HasPrefix(name, "gguf-")) {
|
||||
if info, err := entry.Info(); err == nil {
|
||||
path := filepath.Join(dir, name)
|
||||
if time.Since(info.ModTime()) > ttl {
|
||||
_ = os.Remove(path)
|
||||
continue
|
||||
}
|
||||
files = append(files, cacheFile{path: path, modTime: info.ModTime()})
|
||||
}
|
||||
}
|
||||
}
|
||||
if limit <= 0 || len(files) <= limit {
|
||||
return
|
||||
}
|
||||
sort.Slice(files, func(i, j int) bool { return files[i].modTime.Before(files[j].modTime) })
|
||||
for _, file := range files[:len(files)-limit] {
|
||||
_ = os.Remove(file.path)
|
||||
}
|
||||
}
|
||||
|
||||
func persistentRemoteURI(uri string) bool {
|
||||
parsed, err := url.Parse(uri)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
scheme := strings.ToLower(parsed.Scheme)
|
||||
return (scheme == "http" || scheme == "https") && parsed.Host != ""
|
||||
}
|
||||
|
||||
type sizeCacheEntry struct {
|
||||
size int64
|
||||
err error
|
||||
@@ -33,6 +191,28 @@ type cachedSizeResolver struct {
|
||||
underlying SizeResolver
|
||||
mu sync.Mutex
|
||||
cache map[string]sizeCacheEntry
|
||||
diskDir string
|
||||
diskTTL time.Duration
|
||||
diskGuard *persistentGenerationGuard
|
||||
}
|
||||
|
||||
type persistentSizeEntry struct {
|
||||
Version int `json:"version"`
|
||||
Size int64 `json:"size"`
|
||||
}
|
||||
|
||||
func newCachedSizeResolver(underlying SizeResolver, diskDir string, diskTTL time.Duration) *cachedSizeResolver {
|
||||
return newCachedSizeResolverWithGuard(underlying, diskDir, diskTTL, &persistentGenerationGuard{dir: diskDir})
|
||||
}
|
||||
|
||||
func newCachedSizeResolverWithGuard(underlying SizeResolver, diskDir string, diskTTL time.Duration, guard *persistentGenerationGuard) *cachedSizeResolver {
|
||||
return &cachedSizeResolver{
|
||||
underlying: underlying,
|
||||
cache: make(map[string]sizeCacheEntry),
|
||||
diskDir: diskDir,
|
||||
diskTTL: diskTTL,
|
||||
diskGuard: guard,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) ContentLength(ctx context.Context, uri string) (int64, error) {
|
||||
@@ -43,13 +223,62 @@ func (c *cachedSizeResolver) ContentLength(ctx context.Context, uri string) (int
|
||||
if ok && e.generation == gen {
|
||||
return e.size, e.err
|
||||
}
|
||||
if persistentRemoteURI(uri) && c.canReadPersistent(gen) {
|
||||
if size, ok := c.readPersistent(uri); ok {
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = sizeCacheEntry{size: size, generation: gen}
|
||||
c.mu.Unlock()
|
||||
return size, nil
|
||||
}
|
||||
}
|
||||
size, err := c.underlying.ContentLength(ctx, uri)
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = sizeCacheEntry{size: size, err: err, generation: gen}
|
||||
c.mu.Unlock()
|
||||
if err == nil && persistentRemoteURI(uri) {
|
||||
c.writePersistent(uri, size, gen)
|
||||
}
|
||||
return size, err
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) canReadPersistent(generation uint64) bool {
|
||||
return c.diskGuard.canRead(generation)
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) persistentPath(uri string) string {
|
||||
digest := sha256.Sum256([]byte(uri))
|
||||
return filepath.Join(c.diskDir, "size-"+hex.EncodeToString(digest[:])+".json")
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) readPersistent(uri string) (int64, bool) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
path := c.persistentPath(uri)
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || time.Since(info.ModTime()) > c.diskTTL {
|
||||
return 0, false
|
||||
}
|
||||
data, err := os.ReadFile(path) // #nosec G304 -- path is a hash under the configured cache directory.
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
var entry persistentSizeEntry
|
||||
if json.Unmarshal(data, &entry) != nil || entry.Version != persistentCacheVersion || entry.Size <= 0 {
|
||||
return 0, false
|
||||
}
|
||||
return entry.Size, true
|
||||
}
|
||||
|
||||
func (c *cachedSizeResolver) writePersistent(uri string, size int64, generation uint64) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 || os.MkdirAll(c.diskDir, 0o750) != nil {
|
||||
return
|
||||
}
|
||||
c.diskGuard.persist(generation, func() {
|
||||
writePersistentJSON(c.persistentPath(uri), persistentSizeEntry{Version: persistentCacheVersion, Size: size}, c.diskTTL)
|
||||
})
|
||||
}
|
||||
|
||||
type ggufCacheEntry struct {
|
||||
meta *GGUFMeta
|
||||
err error
|
||||
@@ -60,6 +289,28 @@ type cachedGGUFReader struct {
|
||||
underlying GGUFMetadataReader
|
||||
mu sync.Mutex
|
||||
cache map[string]ggufCacheEntry
|
||||
diskDir string
|
||||
diskTTL time.Duration
|
||||
diskGuard *persistentGenerationGuard
|
||||
}
|
||||
|
||||
type persistentGGUFEntry struct {
|
||||
Version int `json:"version"`
|
||||
Meta *GGUFMeta `json:"meta"`
|
||||
}
|
||||
|
||||
func newCachedGGUFReader(underlying GGUFMetadataReader, diskDir string, diskTTL time.Duration) *cachedGGUFReader {
|
||||
return newCachedGGUFReaderWithGuard(underlying, diskDir, diskTTL, &persistentGenerationGuard{dir: diskDir})
|
||||
}
|
||||
|
||||
func newCachedGGUFReaderWithGuard(underlying GGUFMetadataReader, diskDir string, diskTTL time.Duration, guard *persistentGenerationGuard) *cachedGGUFReader {
|
||||
return &cachedGGUFReader{
|
||||
underlying: underlying,
|
||||
cache: make(map[string]ggufCacheEntry),
|
||||
diskDir: diskDir,
|
||||
diskTTL: diskTTL,
|
||||
diskGuard: guard,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFMeta, error) {
|
||||
@@ -70,26 +321,106 @@ func (c *cachedGGUFReader) ReadMetadata(ctx context.Context, uri string) (*GGUFM
|
||||
if ok && e.generation == gen {
|
||||
return e.meta, e.err
|
||||
}
|
||||
if persistentRemoteURI(uri) && c.canReadPersistent(gen) {
|
||||
if meta, ok := c.readPersistent(uri); ok {
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = ggufCacheEntry{meta: meta, generation: gen}
|
||||
c.mu.Unlock()
|
||||
return meta, nil
|
||||
}
|
||||
}
|
||||
meta, err := c.underlying.ReadMetadata(ctx, uri)
|
||||
c.mu.Lock()
|
||||
c.cache[uri] = ggufCacheEntry{meta: meta, err: err, generation: gen}
|
||||
c.mu.Unlock()
|
||||
if err == nil && meta != nil && persistentRemoteURI(uri) {
|
||||
c.writePersistent(uri, meta, gen)
|
||||
}
|
||||
return meta, err
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) canReadPersistent(generation uint64) bool {
|
||||
return c.diskGuard.canRead(generation)
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) persistentPath(uri string) string {
|
||||
digest := sha256.Sum256([]byte(uri))
|
||||
return filepath.Join(c.diskDir, "gguf-"+hex.EncodeToString(digest[:])+".json")
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) readPersistent(uri string) (*GGUFMeta, bool) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 {
|
||||
return nil, false
|
||||
}
|
||||
path := c.persistentPath(uri)
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || time.Since(info.ModTime()) > c.diskTTL {
|
||||
return nil, false
|
||||
}
|
||||
data, err := os.ReadFile(path) // #nosec G304 -- path is a hash under the configured cache directory.
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
var entry persistentGGUFEntry
|
||||
if json.Unmarshal(data, &entry) != nil || entry.Version != persistentCacheVersion || !validPersistentGGUFMeta(entry.Meta) {
|
||||
return nil, false
|
||||
}
|
||||
return entry.Meta, true
|
||||
}
|
||||
|
||||
func (c *cachedGGUFReader) writePersistent(uri string, meta *GGUFMeta, generation uint64) {
|
||||
if c.diskDir == "" || c.diskTTL <= 0 || os.MkdirAll(c.diskDir, 0o750) != nil {
|
||||
return
|
||||
}
|
||||
c.diskGuard.persist(generation, func() {
|
||||
writePersistentJSON(c.persistentPath(uri), persistentGGUFEntry{Version: persistentCacheVersion, Meta: meta}, c.diskTTL)
|
||||
})
|
||||
}
|
||||
|
||||
func validPersistentGGUFMeta(meta *GGUFMeta) bool {
|
||||
return meta != nil && meta.BlockCount > 0 && meta.EmbeddingLength > 0 && meta.HeadCount > 0 && meta.HeadCountKV > 0
|
||||
}
|
||||
|
||||
func writePersistentJSON(path string, value any, ttl time.Duration) {
|
||||
data, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
tmp, err := os.CreateTemp(filepath.Dir(path), ".vram-*.tmp")
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
tmpPath := tmp.Name()
|
||||
defer func() { _ = os.Remove(tmpPath) }()
|
||||
if _, err = tmp.Write(data); err != nil {
|
||||
_ = tmp.Close()
|
||||
return
|
||||
}
|
||||
if err = tmp.Close(); err != nil {
|
||||
return
|
||||
}
|
||||
if os.Rename(tmpPath, path) == nil {
|
||||
prunePersistentEntries(filepath.Dir(path), ttl, persistentCacheEntryLimit)
|
||||
}
|
||||
}
|
||||
|
||||
// DefaultCachedSizeResolver returns a cached SizeResolver using the default implementation.
|
||||
// Entries are invalidated when the gallery generation changes.
|
||||
func DefaultCachedSizeResolver() SizeResolver {
|
||||
defaultCacheMu.RLock()
|
||||
defer defaultCacheMu.RUnlock()
|
||||
return defaultCachedSizeResolver
|
||||
}
|
||||
|
||||
// DefaultCachedGGUFReader returns a cached GGUFMetadataReader using the default implementation.
|
||||
// Entries are invalidated when the gallery generation changes.
|
||||
func DefaultCachedGGUFReader() GGUFMetadataReader {
|
||||
defaultCacheMu.RLock()
|
||||
defer defaultCacheMu.RUnlock()
|
||||
return defaultCachedGGUFReader
|
||||
}
|
||||
|
||||
var (
|
||||
defaultCachedSizeResolver = &cachedSizeResolver{underlying: defaultSizeResolver{}, cache: make(map[string]sizeCacheEntry)}
|
||||
defaultCachedGGUFReader = &cachedGGUFReader{underlying: defaultGGUFReader{}, cache: make(map[string]ggufCacheEntry)}
|
||||
defaultCachedSizeResolver = newCachedSizeResolver(defaultSizeResolver{}, "", 0)
|
||||
defaultCachedGGUFReader = newCachedGGUFReader(defaultGGUFReader{}, "", 0)
|
||||
)
|
||||
@@ -0,0 +1,228 @@
|
||||
package vram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
type countingSizeResolver struct {
|
||||
size int64
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
type countingGGUFReader struct {
|
||||
meta *GGUFMeta
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
type blockingSizeResolver struct {
|
||||
started chan struct{}
|
||||
release chan struct{}
|
||||
}
|
||||
|
||||
func (r *blockingSizeResolver) ContentLength(context.Context, string) (int64, error) {
|
||||
close(r.started)
|
||||
<-r.release
|
||||
return 42, nil
|
||||
}
|
||||
|
||||
func (r *countingGGUFReader) ReadMetadata(context.Context, string) (*GGUFMeta, error) {
|
||||
r.calls++
|
||||
return r.meta, r.err
|
||||
}
|
||||
|
||||
func (r *countingSizeResolver) ContentLength(context.Context, string) (int64, error) {
|
||||
r.calls++
|
||||
return r.size, r.err
|
||||
}
|
||||
|
||||
var _ = Describe("persistent VRAM metadata cache", func() {
|
||||
AfterEach(func() {
|
||||
ConfigurePersistentCache("", 0)
|
||||
SetGalleryGenerationFunc(nil)
|
||||
})
|
||||
|
||||
It("reuses a successful size probe after the in-memory cache is replaced", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
firstSource := &countingSizeResolver{size: 42}
|
||||
first := newCachedSizeResolver(firstSource, cacheDir, time.Hour)
|
||||
|
||||
size, err := first.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
Expect(firstSource.calls).To(Equal(1))
|
||||
|
||||
secondSource := &countingSizeResolver{err: errors.New("unexpected remote probe")}
|
||||
second := newCachedSizeResolver(secondSource, cacheDir, time.Hour)
|
||||
size, err = second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
Expect(secondSource.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("reuses successful GGUF metadata after the in-memory cache is replaced", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
want := &GGUFMeta{BlockCount: 32, EmbeddingLength: 4096, HeadCount: 32, HeadCountKV: 8, MaximumContextLength: 131072}
|
||||
firstSource := &countingGGUFReader{meta: want}
|
||||
first := newCachedGGUFReader(firstSource, cacheDir, time.Hour)
|
||||
|
||||
meta, err := first.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
Expect(firstSource.calls).To(Equal(1))
|
||||
|
||||
secondSource := &countingGGUFReader{err: errors.New("unexpected remote probe")}
|
||||
second := newCachedGGUFReader(secondSource, cacheDir, time.Hour)
|
||||
meta, err = second.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
Expect(secondSource.calls).To(BeZero())
|
||||
})
|
||||
|
||||
It("configures the default caches used by model estimates", func() {
|
||||
cacheDir := filepath.Join(GinkgoT().TempDir(), "vram")
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
|
||||
Expect(defaultCachedSizeResolver.diskDir).To(Equal(cacheDir))
|
||||
Expect(defaultCachedGGUFReader.diskDir).To(Equal(cacheDir))
|
||||
Expect(defaultCachedSizeResolver.diskTTL).To(Equal(time.Hour))
|
||||
Expect(defaultCachedGGUFReader.diskTTL).To(Equal(time.Hour))
|
||||
Expect(defaultCachedSizeResolver.diskGuard).To(BeIdenticalTo(defaultCachedGGUFReader.diskGuard))
|
||||
})
|
||||
|
||||
It("removes expired VRAM entries when the persistent cache is configured", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
stale := filepath.Join(cacheDir, "size-stale.json")
|
||||
abandoned := filepath.Join(cacheDir, ".vram-abandoned.tmp")
|
||||
unrelated := filepath.Join(cacheDir, "keep.txt")
|
||||
Expect(os.WriteFile(stale, []byte("{}"), 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(abandoned, []byte("partial"), 0o600)).To(Succeed())
|
||||
Expect(os.WriteFile(unrelated, []byte("keep"), 0o600)).To(Succeed())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
Expect(os.Chtimes(stale, old, old)).To(Succeed())
|
||||
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
|
||||
Expect(stale).NotTo(BeAnExistingFile())
|
||||
Expect(abandoned).NotTo(BeAnExistingFile())
|
||||
Expect(unrelated).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("does not reuse a persistent entry after the gallery generation changes", func() {
|
||||
var generation uint64 = 1
|
||||
SetGalleryGenerationFunc(func() uint64 { return generation })
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
first := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
_, err := first.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
freshSource := &countingSizeResolver{size: 84}
|
||||
second := newCachedSizeResolver(freshSource, cacheDir, time.Hour)
|
||||
size, err := second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
|
||||
generation = 2
|
||||
size, err = second.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(84)))
|
||||
Expect(freshSource.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("does not persist probes for local model files", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
first := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
_, err := first.ContentLength(context.Background(), "file:///models/model.gguf")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
freshSource := &countingSizeResolver{size: 84}
|
||||
second := newCachedSizeResolver(freshSource, cacheDir, time.Hour)
|
||||
size, err := second.ContentLength(context.Background(), "file:///models/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(84)))
|
||||
Expect(freshSource.calls).To(Equal(1))
|
||||
})
|
||||
|
||||
It("falls back to the remote probe for an invalid persisted size", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
resolver := newCachedSizeResolver(&countingSizeResolver{size: 42}, cacheDir, time.Hour)
|
||||
Expect(os.WriteFile(resolver.persistentPath("https://example.com/model.gguf"), []byte(`{"size":-1}`), 0o600)).To(Succeed())
|
||||
|
||||
size, err := resolver.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(size).To(Equal(int64(42)))
|
||||
})
|
||||
|
||||
It("falls back to the remote probe for empty persisted GGUF metadata", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
want := &GGUFMeta{BlockCount: 32, EmbeddingLength: 4096, HeadCount: 32, HeadCountKV: 8}
|
||||
reader := newCachedGGUFReader(&countingGGUFReader{meta: want}, cacheDir, time.Hour)
|
||||
Expect(os.WriteFile(reader.persistentPath("https://example.com/model.gguf"), []byte(`{"meta":{}}`), 0o600)).To(Succeed())
|
||||
|
||||
meta, err := reader.ReadMetadata(context.Background(), "https://example.com/model.gguf")
|
||||
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(meta).To(Equal(want))
|
||||
})
|
||||
|
||||
It("keeps the persistent cache within its entry limit", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
for _, name := range []string{"size-a.json", "size-b.json", "gguf-c.json"} {
|
||||
Expect(os.WriteFile(filepath.Join(cacheDir, name), []byte("{}"), 0o600)).To(Succeed())
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
prunePersistentEntries(cacheDir, time.Hour, 2)
|
||||
|
||||
entries, err := os.ReadDir(cacheDir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(entries).To(HaveLen(2))
|
||||
Expect(filepath.Join(cacheDir, "size-a.json")).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("removes persistent entries when gallery data is invalidated", func() {
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
ConfigurePersistentCache(cacheDir, time.Hour)
|
||||
stale := filepath.Join(cacheDir, "size-stale.json")
|
||||
Expect(os.WriteFile(stale, []byte(`{"version":1,"size":42}`), 0o600)).To(Succeed())
|
||||
|
||||
InvalidatePersistentCache()
|
||||
|
||||
Expect(stale).NotTo(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("does not persist a probe that finishes after invalidation", func() {
|
||||
var generation uint64 = 1
|
||||
SetGalleryGenerationFunc(func() uint64 { return generation })
|
||||
cacheDir := GinkgoT().TempDir()
|
||||
guard := &persistentGenerationGuard{dir: cacheDir}
|
||||
source := &blockingSizeResolver{started: make(chan struct{}), release: make(chan struct{})}
|
||||
resolver := newCachedSizeResolverWithGuard(source, cacheDir, time.Hour, guard)
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := resolver.ContentLength(context.Background(), "https://example.com/model.gguf")
|
||||
done <- err
|
||||
}()
|
||||
Eventually(source.started).Should(BeClosed())
|
||||
|
||||
generation = 2
|
||||
guard.invalidate(generation)
|
||||
close(source.release)
|
||||
|
||||
Eventually(done).Should(Receive(BeNil()))
|
||||
Expect(resolver.persistentPath("https://example.com/model.gguf")).NotTo(BeAnExistingFile())
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user