mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-28 09:08:54 -04:00
fix: stop listing the models directory on request paths and cache GGUF header reads (#12683)
* fix(http): stop listing the models directory on embeddings and bearer requests Every embeddings request read the whole models directory before the body was parsed, even when the body named its model. On a network-backed models directory this cost seconds per request. Embeddings now binds the body first and looks up a default model only when the body has none, like chat. The pre-bind default lookup and the gpt-4o constant default are gone from that route. The first-available default lookup skips the directory when the filter cannot accept an unconfigured file. Every usecase filter rejects files without a config, so the scan never added results for them. Match-all, nil and name filters keep the scan. The bearer token check is now an exact test: a configured name, or one stat of that name in the models directory. The value is a literal, not a regular expression, and a name that is not a local relative path never matches a file. Assisted-by: Claude Code:claude-haiku-5-5 * perf(config): cache GGUF header reads for model defaults Model defaults re-parsed the GGUF header on every SetDefaults call, which runs per request and at startup reload. Put a bounded LRU in front of the parser, keyed by absolute path and validated by size and modification time. The cached value is a snapshot of the fields the defaults use, not the parser output. A changed size or mtime is a miss. Missing files and parse errors are not cached, so the existing fallback is unchanged. LOCALAI_DISABLE_GUESSING, the empty-path guard, the alias early return, the panic recovery and the context-size default are all kept in front of the cache. Assisted-by: Claude:claude-haiku-5-5 * test(http): cover listing-free default lookup on every usecase route Assisted-by: Claude Code:claude-haiku-5-5 * feat(config): persist the GGUF header cache per instance The header cache kept parsed GGUF results only in memory, so every restart parsed each model file again before its defaults were known. Write the cache to a per-instance file in the data path and load it at startup. A restarted frontend with a fixed instance ID starts warm. The file is versioned and records the estimator inputs (gguf-parser version, candidate contexts, auto context size, offload count). A mismatch on any of them makes the file ignored. Each entry keeps path, size and mtime, so a changed model file is a miss on load. Writes run on one background saver, debounced, through a temp file and rename. The entry bound applies to the file as well. Without a data path the cache stays in memory only. Documented as a per-instance exception in the model configuration docs. Assisted-by: Claude:claude-haiku-5-5 * docs(embeddings): document model selection and the no-model error Assisted-by: Claude Code:claude-haiku-5-5 * refactor(config): split the header cache into a generic cache The GGUF header cache mixed two things: the mechanics (bounded LRU keyed by path and identity, per-instance file, debounced atomic save, tolerant load) and the GGUF fields. Move the mechanics into headerCache[S], driven by a headerCacheSpec[S] that names the file prefix, schema, fingerprint, parser and the snapshot encode/decode pair. The GGUF cache becomes one instance of it. File names, the identity check, LRU size, debounce, atomic write and the corruption handling are unchanged. The on-disk JSON keeps the "estimator" key so existing files still load. Assisted-by: Claude Code:claude-haiku-5-5 * feat(config): share one GGUF header cache file between frontends Every frontend used to write its own header cache file, named after its instance ID, so a header read by one frontend was never reused by another. The cache now uses one file, header-cache.json, in the data path, for all frontends and for standalone mode. The instance ID no longer changes the file name. Frontends share the file in three ways. A save first merges the file's entries in, so one frontend's write does not drop another's entries. A miss reloads the file when its modification time has changed, and a hit never reads it. Writes go to a temp file that is renamed into place, so a reader never sees a partial file. Two frontends saving at the same moment can still drop each other's last entries. That costs one re-parse, never a wrong value, because every entry is checked against its model file before use. The shared file is valid only when every frontend mounts the same data volume. The distributed-mode and model-configuration docs state this mode and its volume limit. Assisted-by: Claude:claude-haiku-5-5 * feat(config): give an instance ID its own header cache file Assisted-by: Claude Code:claude-haiku-5-5 * style(config): gofmt the distributed config struct Assisted-by: Claude Code:claude-haiku-5-5 --------- Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
80bf7e5a28
commit
c7ce0c602c
27 files changed
+2283
-54
No files matched your search
@@ -109,6 +109,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
// Generate instance ID if not set
|
||||
if cfg.Distributed.InstanceID == "" {
|
||||
cfg.Distributed.InstanceID = uuid.New().String()
|
||||
cfg.Distributed.InstanceIDGenerated = true
|
||||
}
|
||||
xlog.Info("Distributed instance", "id", cfg.Distributed.InstanceID)
|
||||
|
||||
|
||||
@@ -444,6 +444,16 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// Attach the GGUF header cache now, before any model config is loaded, so the
|
||||
// first load already starts from the shared file.
|
||||
// Only an ID the operator set keys a per-instance file; a generated ID would
|
||||
// give every restart its own file.
|
||||
cacheInstanceID := ""
|
||||
if !options.Distributed.InstanceIDGenerated {
|
||||
cacheInstanceID = options.Distributed.InstanceID
|
||||
}
|
||||
config.ConfigureGGUFHeaderCache(options.Context, options.DataPath, cacheInstanceID)
|
||||
|
||||
configLoaderOpts := options.ToConfigLoaderOptions()
|
||||
|
||||
if err := application.ModelConfigLoader().LoadModelConfigsFromPath(options.SystemState.Model.ModelsPath, configLoaderOpts...); err != nil {
|
||||
|
||||
@@ -86,7 +86,7 @@ func contextFitsVRAM(footprint, vram uint64) bool {
|
||||
return vram >= footprint+footprint/contextFitHeadroomDivisor
|
||||
}
|
||||
|
||||
// autoContextSize picks the default context to use for f when the user did not
|
||||
// autoContextSize picks the default context to use when the user did not
|
||||
// set context_size. The choice is deliberately conservative, NOT
|
||||
// VRAM-maximizing:
|
||||
//
|
||||
@@ -100,10 +100,13 @@ func contextFitsVRAM(footprint, vram uint64) bool {
|
||||
// DefaultContextSize. When VRAM is unknown we skip this — the base cap is
|
||||
// already safe and we must not regress CPU / detection-gap hosts.
|
||||
//
|
||||
// footprint reports the per-device VRAM footprint of the model at a context; it
|
||||
// is a parameter so the caller can supply a snapshot instead of the parsed file.
|
||||
//
|
||||
// trainedMax <= 0 means the estimate yielded nothing usable; the caller keeps
|
||||
// its existing DefaultContextSize fallback in that case, so this is only called
|
||||
// with a positive trainedMax.
|
||||
func autoContextSize(f *gguf.GGUFFile, trainedMax int) int {
|
||||
func autoContextSize(footprint func(ctx int) uint64, trainedMax int) int {
|
||||
chosen := trainedMax
|
||||
if chosen > DefaultAutoContextSize {
|
||||
chosen = DefaultAutoContextSize
|
||||
@@ -118,7 +121,7 @@ func autoContextSize(f *gguf.GGUFFile, trainedMax int) int {
|
||||
return chosen
|
||||
}
|
||||
|
||||
if contextFitsVRAM(estimateContextVRAM(f, chosen), vram) {
|
||||
if contextFitsVRAM(footprint(chosen), vram) {
|
||||
return chosen
|
||||
}
|
||||
|
||||
@@ -128,7 +131,7 @@ func autoContextSize(f *gguf.GGUFFile, trainedMax int) int {
|
||||
if cand > chosen || cand < DefaultContextSize {
|
||||
continue
|
||||
}
|
||||
if contextFitsVRAM(estimateContextVRAM(f, cand), vram) {
|
||||
if contextFitsVRAM(footprint(cand), vram) {
|
||||
xlog.Debug("[context_fit] capped auto context to fit per-device VRAM",
|
||||
"context", cand, "base_cap", chosen, "vram_gib", vram>>30)
|
||||
return cand
|
||||
|
||||
@@ -10,6 +10,10 @@ import (
|
||||
// (perDeviceVRAM, estimateContextVRAM) are package vars so a deterministic VRAM
|
||||
// ceiling and footprint can be injected without a real GPU or model file — the
|
||||
// same pattern hardware_defaults_internal_test.go uses for localGPU.
|
||||
// nilFileFootprint asks the package estimator for a footprint without a parsed
|
||||
// file, which is what the test stubs need.
|
||||
func nilFileFootprint(ctx int) uint64 { return estimateContextVRAM(nil, ctx) }
|
||||
|
||||
var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
const gib = uint64(1) << 30
|
||||
|
||||
@@ -33,14 +37,14 @@ var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
// we keep the conservative 8k cap (users opt into more via context_size).
|
||||
perDeviceVRAM = func() uint64 { return 80 * gib }
|
||||
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return gib } // trivially fits
|
||||
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultAutoContextSize))
|
||||
Expect(autoContextSize(nilFileFootprint, 1048576)).To(Equal(DefaultAutoContextSize))
|
||||
})
|
||||
|
||||
It("keeps a small model's trained window instead of inflating it", func() {
|
||||
// trained 4096 < 8192: min() keeps 4096, it is not raised to the cap.
|
||||
perDeviceVRAM = func() uint64 { return 80 * gib }
|
||||
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return gib }
|
||||
Expect(autoContextSize(nil, 4096)).To(Equal(4096))
|
||||
Expect(autoContextSize(nilFileFootprint, 4096)).To(Equal(4096))
|
||||
})
|
||||
|
||||
It("steps below the cap when even 8k would not fit a tiny card", func() {
|
||||
@@ -51,7 +55,7 @@ var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
estimateContextVRAM = func(_ *gguf.GGUFFile, ctx int) uint64 {
|
||||
return gib + uint64(ctx)*100000
|
||||
}
|
||||
chosen := autoContextSize(nil, 1048576)
|
||||
chosen := autoContextSize(nilFileFootprint, 1048576)
|
||||
Expect(chosen).To(BeNumerically("<", DefaultAutoContextSize))
|
||||
Expect(chosen).To(BeNumerically(">=", DefaultContextSize))
|
||||
// The chosen context's footprint must actually fit the card with headroom.
|
||||
@@ -64,7 +68,7 @@ var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
// than defaulting to a window guaranteed to abort.
|
||||
perDeviceVRAM = func() uint64 { return 1 * gib }
|
||||
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return 100 * gib }
|
||||
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultContextSize))
|
||||
Expect(autoContextSize(nilFileFootprint, 1048576)).To(Equal(DefaultContextSize))
|
||||
})
|
||||
|
||||
It("does not clamp when per-device VRAM is unknown", func() {
|
||||
@@ -72,7 +76,7 @@ var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
// not regress — keep the conservative base cap regardless of estimate.
|
||||
perDeviceVRAM = func() uint64 { return 0 }
|
||||
estimateContextVRAM = func(_ *gguf.GGUFFile, _ int) uint64 { return 999 * gib }
|
||||
Expect(autoContextSize(nil, 1048576)).To(Equal(DefaultAutoContextSize))
|
||||
Expect(autoContextSize(nilFileFootprint, 1048576)).To(Equal(DefaultAutoContextSize))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -93,7 +97,7 @@ var _ = Describe("Auto-derived default context (VRAM-aware cap)", func() {
|
||||
},
|
||||
},
|
||||
}
|
||||
guessGGUFFromFile(cfg, f, 0)
|
||||
guessGGUFFromFile(cfg, readGGUFHeader(f), 0)
|
||||
Expect(cfg.ContextSize).ToNot(BeNil())
|
||||
Expect(*cfg.ContextSize).To(Equal(262144))
|
||||
})
|
||||
|
||||
@@ -13,11 +13,14 @@ import (
|
||||
// DistributedConfig holds configuration for horizontal scaling mode.
|
||||
// When Enabled is true, PostgreSQL and NATS are required.
|
||||
type DistributedConfig struct {
|
||||
Enabled bool // --distributed / LOCALAI_DISTRIBUTED
|
||||
InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty)
|
||||
NatsURL string // --nats-url / LOCALAI_NATS_URL
|
||||
StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint)
|
||||
RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration)
|
||||
Enabled bool // --distributed / LOCALAI_DISTRIBUTED
|
||||
InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty)
|
||||
// InstanceIDGenerated is set when the ID above was generated at startup. A generated
|
||||
// ID changes on every start, so on-disk state must not be keyed by it.
|
||||
InstanceIDGenerated bool
|
||||
NatsURL string // --nats-url / LOCALAI_NATS_URL
|
||||
StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint)
|
||||
RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration)
|
||||
// RegistrationRequireAuth fails startup when distributed mode is enabled but
|
||||
// RegistrationToken is empty. The default (false) keeps the historical
|
||||
// fail-open behavior with a loud warning; production should set it so the
|
||||
|
||||
+21
-26
@@ -47,22 +47,22 @@ func HasGenAudioProjector(f *gguf.GGUFFile) bool {
|
||||
return kv.ValueBool()
|
||||
}
|
||||
|
||||
func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
|
||||
func guessGGUFFromFile(cfg *ModelConfig, f ggufHeader, defaultCtx int) {
|
||||
// Explicit opt-in: a negative context_size (canonically -1) means "use the
|
||||
// model's full trained context (n_ctx_train) from GGUF metadata". Unlike the
|
||||
// silent unset path below, this overrides an already-present value and warns
|
||||
// when the resolved window will not fit detected VRAM.
|
||||
if cfg.ContextSize != nil && *cfg.ContextSize < 0 {
|
||||
if maxCtx := int(f.Architecture().MaximumContextLength); maxCtx > 0 {
|
||||
if maxCtx := f.MaxContextLength; maxCtx > 0 {
|
||||
cfg.ContextSize = &maxCtx
|
||||
warnIfContextExceedsVRAM(f, maxCtx, f.Metadata().Name)
|
||||
warnIfContextExceedsVRAM(f.maxCtxFootprint, maxCtx, f.Name)
|
||||
} else {
|
||||
// No usable trained max in metadata: degrade to the safe default
|
||||
// rather than leak a negative n_ctx downstream.
|
||||
d := DefaultContextSize
|
||||
cfg.ContextSize = &d
|
||||
xlog.Warn("[gguf] context_size=-1 requested but GGUF exposes no trained max; using default",
|
||||
"default", d, "model", f.Metadata().Name)
|
||||
"default", d, "model", f.Name)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,9 +72,13 @@ func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
|
||||
// load: a 128k / 256k / 1M KV cache cannot fit a consumer GPU and the
|
||||
// backend aborts (exitCode=-1). autoContextSize instead caps to a modest
|
||||
// default and only steps below it when detected per-device VRAM demands.
|
||||
trainedMax := int(f.EstimateLLaMACppRun().ContextSize)
|
||||
if f.trainedPanic != nil {
|
||||
// Replay the estimator panic where the defaults path used to hit it.
|
||||
panic(f.trainedPanic)
|
||||
}
|
||||
trainedMax := f.TrainedContext
|
||||
if trainedMax > 0 {
|
||||
cSize := autoContextSize(f, trainedMax)
|
||||
cSize := autoContextSize(f.footprintAt, trainedMax)
|
||||
cfg.ContextSize = &cSize
|
||||
} else {
|
||||
defaultCtx = DefaultContextSize
|
||||
@@ -95,20 +99,19 @@ func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
|
||||
cfg.NGPULayers = &defaultHigh
|
||||
}
|
||||
|
||||
xlog.Debug("[gguf] guessDefaultsFromFile: NGPULayers set", "NGPULayers", cfg.NGPULayers, "modelName", f.Metadata().Name)
|
||||
xlog.Debug("[gguf] guessDefaultsFromFile: NGPULayers set", "NGPULayers", cfg.NGPULayers, "modelName", f.Name)
|
||||
|
||||
// identify from well known templates first, otherwise use the raw jinja template
|
||||
chatTemplate, found := f.Header.MetadataKV.Get("tokenizer.chat_template")
|
||||
if found {
|
||||
if f.HasChatTemplate {
|
||||
// fill jinja template
|
||||
cfg.modelTemplate = chatTemplate.ValueString()
|
||||
cfg.modelTemplate = f.ChatTemplate
|
||||
}
|
||||
|
||||
// Auto-enable Multi-Token Prediction (ggml-org/llama.cpp#22673) when the
|
||||
// GGUF carries an embedded MTP head. Skipped silently for non-MTP models
|
||||
// and when the user already configured a spec_type.
|
||||
if n, ok := HasEmbeddedMTPHead(f); ok {
|
||||
ApplyMTPDefaults(cfg, n)
|
||||
if f.HasMTP {
|
||||
ApplyMTPDefaults(cfg, f.MTPLayers)
|
||||
}
|
||||
|
||||
// Thinking support detection is done after model load via DetectThinkingSupportFromBackend
|
||||
@@ -116,15 +119,15 @@ func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
|
||||
// template estimations
|
||||
if cfg.HasTemplate() {
|
||||
// nothing to guess here
|
||||
xlog.Debug("[gguf] guessDefaultsFromFile: template already set", "name", cfg.Name, "modelName", f.Metadata().Name)
|
||||
xlog.Debug("[gguf] guessDefaultsFromFile: template already set", "name", cfg.Name, "modelName", f.Name)
|
||||
return
|
||||
}
|
||||
|
||||
xlog.Debug("[gguf] Model file loaded", "file", cfg.ModelFileName(), "eosTokenID", f.Tokenizer().EOSTokenID, "bosTokenID", f.Tokenizer().BOSTokenID, "modelName", f.Metadata().Name, "architecture", f.Architecture().Architecture)
|
||||
xlog.Debug("[gguf] Model file loaded", "file", cfg.ModelFileName(), "eosTokenID", f.EOSTokenID, "bosTokenID", f.BOSTokenID, "modelName", f.Name, "architecture", f.Architecture)
|
||||
|
||||
// guess the name
|
||||
if cfg.Name == "" {
|
||||
cfg.Name = f.Metadata().Name
|
||||
cfg.Name = f.Name
|
||||
}
|
||||
|
||||
// A model the operator reserved for an internal primitive (the router
|
||||
@@ -142,7 +145,7 @@ func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
|
||||
}
|
||||
|
||||
// Apply per-model-family inference parameter defaults (temperature, top_p, etc.)
|
||||
ApplyInferenceDefaults(cfg, f.Metadata().Name)
|
||||
ApplyInferenceDefaults(cfg, f.Name)
|
||||
}
|
||||
|
||||
// DetectThinkingSupportFromBackend calls the ModelMetadata gRPC method to detect
|
||||
@@ -273,8 +276,8 @@ func applyDetectedThinkingConfig(cfg *ModelConfig, metadata *pb.ModelMetadataRes
|
||||
// detection or estimation gap (no GPU, unknown VRAM, estimate failure) silently
|
||||
// skips the warning. Used by the context_size=-1 auto-max path, where the raw
|
||||
// trained max can be far larger than a consumer card holds.
|
||||
func warnIfContextExceedsVRAM(f *gguf.GGUFFile, ctx int, name string) {
|
||||
defer func() { _ = recover() }() // the run estimate can panic on unusual headers
|
||||
func warnIfContextExceedsVRAM(used uint64, ctx int, name string) {
|
||||
defer func() { _ = recover() }()
|
||||
|
||||
if !xsysinfo.HasGPU("nvidia") && !xsysinfo.HasGPU("amd") {
|
||||
return // no VRAM to compare against
|
||||
@@ -284,14 +287,6 @@ func warnIfContextExceedsVRAM(f *gguf.GGUFFile, ctx int, name string) {
|
||||
return
|
||||
}
|
||||
|
||||
sum := f.EstimateLLaMACppRun(gguf.WithLLaMACppContextSize(int32(ctx))).Summarize(true, 0, 0)
|
||||
if len(sum.Items) == 0 {
|
||||
return
|
||||
}
|
||||
var used uint64
|
||||
for _, v := range sum.Items[0].VRAMs {
|
||||
used += uint64(v.NonUMA)
|
||||
}
|
||||
if used == 0 || used <= vram {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
)
|
||||
|
||||
// ggufHeaderCacheSize bounds the in-memory header cache. Each entry is a few
|
||||
// hundred bytes, so this covers a large model directory while keeping the
|
||||
// memory footprint trivial. Entries past the bound are evicted least recently
|
||||
// used and simply re-parsed on their next use.
|
||||
const ggufHeaderCacheSize = 256
|
||||
|
||||
// ggufParser matches gguf.ParseGGUFFile and is the seam tests use to count
|
||||
// header reads.
|
||||
type ggufParser func(string, ...gguf.GGUFReadOption) (*gguf.GGUFFile, error)
|
||||
|
||||
// ggufHeader is the set of values model defaults read from a parsed GGUF. It
|
||||
// is the only thing kept between requests: the parser output (tensors, the
|
||||
// full metadata map) is not. Every field is a property of the file bytes, so a
|
||||
// snapshot stays valid for as long as the file identity does not change.
|
||||
type ggufHeader struct {
|
||||
Name string // general.name, falling back as the parser does
|
||||
Architecture string // general.architecture
|
||||
MaxContextLength int // architecture's trained maximum, 0 when absent
|
||||
BOSTokenID int64 // tokenizer, for diagnostic logging
|
||||
EOSTokenID int64 // tokenizer, for diagnostic logging
|
||||
|
||||
ChatTemplate string // tokenizer.chat_template
|
||||
HasChatTemplate bool
|
||||
|
||||
MTPLayers uint32 // <arch>.nextn_predict_layers, when an embedded MTP head exists
|
||||
HasMTP bool
|
||||
|
||||
// TrainedContext is the estimator's default context size (0 when the
|
||||
// estimator yields nothing usable).
|
||||
TrainedContext int
|
||||
// trainedPanic is the value the estimator panicked with, if it did. The
|
||||
// defaults path has never recovered from that panic locally; it is replayed
|
||||
// at the same point so the outcome for such a file does not change.
|
||||
trainedPanic any
|
||||
|
||||
// footprints holds the per-device VRAM footprint at each context that the
|
||||
// auto-context logic may ask about. It is filled only when TrainedContext
|
||||
// is positive, the same condition under which that logic runs.
|
||||
footprints map[int]uint64
|
||||
// maxCtxFootprint is the VRAM estimate at MaxContextLength used by the
|
||||
// context_size=-1 warning.
|
||||
maxCtxFootprint uint64
|
||||
}
|
||||
|
||||
// footprintAt returns the precomputed footprint at ctx. A context that was not
|
||||
// precomputed reports 0, the same "cannot confirm" answer as an estimate gap.
|
||||
func (h ggufHeader) footprintAt(ctx int) uint64 {
|
||||
return h.footprints[ctx]
|
||||
}
|
||||
|
||||
// readGGUFHeader converts a parsed file into a snapshot. It is the only place
|
||||
// the parser output is consumed on the defaults path.
|
||||
func readGGUFHeader(f *gguf.GGUFFile) ggufHeader {
|
||||
var h ggufHeader
|
||||
arch := f.Architecture()
|
||||
h.Name = f.Metadata().Name
|
||||
h.Architecture = arch.Architecture
|
||||
h.MaxContextLength = int(arch.MaximumContextLength)
|
||||
tok := f.Tokenizer()
|
||||
h.BOSTokenID = tok.BOSTokenID
|
||||
h.EOSTokenID = tok.EOSTokenID
|
||||
|
||||
if t, ok := f.Header.MetadataKV.Get("tokenizer.chat_template"); ok {
|
||||
h.ChatTemplate = t.ValueString()
|
||||
h.HasChatTemplate = true
|
||||
}
|
||||
|
||||
h.MTPLayers, h.HasMTP = HasEmbeddedMTPHead(f)
|
||||
|
||||
h.TrainedContext, h.trainedPanic = safeTrainedContext(f)
|
||||
if h.TrainedContext > 0 {
|
||||
chosen := h.TrainedContext
|
||||
if chosen > DefaultAutoContextSize {
|
||||
chosen = DefaultAutoContextSize
|
||||
}
|
||||
h.footprints = map[int]uint64{chosen: estimateContextVRAM(f, chosen)}
|
||||
for _, cand := range contextFitCandidates {
|
||||
h.footprints[cand] = estimateContextVRAM(f, cand)
|
||||
}
|
||||
}
|
||||
|
||||
if h.MaxContextLength > 0 {
|
||||
h.maxCtxFootprint = estimateNonUMAAt(f, h.MaxContextLength)
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
// safeTrainedContext returns the estimator's default context size and, when the
|
||||
// estimator panics, the panic value instead of letting it escape.
|
||||
func safeTrainedContext(f *gguf.GGUFFile) (n int, panicked any) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
n, panicked = 0, r
|
||||
}
|
||||
}()
|
||||
return int(f.EstimateLLaMACppRun().ContextSize), nil
|
||||
}
|
||||
|
||||
// estimateNonUMAAt returns the discrete-GPU footprint at ctx for the -1 warning.
|
||||
// It returns 0 on any estimator failure, which suppresses the warning as before.
|
||||
func estimateNonUMAAt(f *gguf.GGUFFile, ctx int) (used uint64) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
used = 0
|
||||
}
|
||||
}()
|
||||
sum := f.EstimateLLaMACppRun(gguf.WithLLaMACppContextSize(int32(ctx))).Summarize(true, 0, 0)
|
||||
if len(sum.Items) == 0 {
|
||||
return 0
|
||||
}
|
||||
for _, v := range sum.Items[0].VRAMs {
|
||||
used += uint64(v.NonUMA)
|
||||
}
|
||||
return used
|
||||
}
|
||||
|
||||
// ggufIdentity is what a cached header is validated against.
|
||||
type ggufIdentity = headerIdentity
|
||||
|
||||
// ggufHeaderCache is the GGUF instance of the generic header cache.
|
||||
type ggufHeaderCache = headerCache[ggufHeader]
|
||||
|
||||
// persistedGGUFCache is the GGUF cache file as read back from disk.
|
||||
type persistedGGUFCache = headerFile
|
||||
|
||||
func newGGUFHeaderCache(max int) *ggufHeaderCache {
|
||||
return newHeaderCache(ggufHeaderSpec, max)
|
||||
}
|
||||
|
||||
// ggufHeaders is the process-wide header cache used by llamaCppDefaults.
|
||||
var ggufHeaders = newGGUFHeaderCache(ggufHeaderCacheSize)
|
||||
|
||||
// parseGGUFHeader parses modelFile with parse and reduces the result to a
|
||||
// snapshot. Only the scalar metadata and array lengths are read, see
|
||||
// llamaCppDefaultsWithParser for why.
|
||||
func parseGGUFHeader(modelFile string, parse ggufParser) (ggufHeader, error) {
|
||||
f, err := parse(modelFile, gguf.UseMMap(), gguf.SkipLargeMetadata())
|
||||
if err != nil {
|
||||
return ggufHeader{}, err
|
||||
}
|
||||
return readGGUFHeader(f), nil
|
||||
}
|
||||
|
||||
// loadGGUFHeader returns the header snapshot for modelFile, parsing it only when
|
||||
// no cached snapshot matches the file's current identity. A file that cannot be
|
||||
// stat'ed or parsed is never cached, so the caller sees the same error path as
|
||||
// before the cache existed.
|
||||
func loadGGUFHeader(modelFile string, parse ggufParser) (ggufHeader, error) {
|
||||
return ggufHeaders.loadVia(modelFile, func(p string) (ggufHeader, error) {
|
||||
return parseGGUFHeader(p, parse)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// newPersistentCache returns a fresh cache attached to dataPath for instanceID,
|
||||
// the way startup attaches the process cache. The saver is stopped in cleanup.
|
||||
func newPersistentCache(dataPath string) *ggufHeaderCache {
|
||||
c := newGGUFHeaderCache(ggufHeaderCacheSize)
|
||||
c.attachDisk(context.Background(), dataPath, "")
|
||||
DeferCleanup(c.stopPersist)
|
||||
return c
|
||||
}
|
||||
|
||||
// useHeaderCache swaps the process cache that llamaCppDefaultsWithParser reads,
|
||||
// so a test can drive a new instance through the existing seam.
|
||||
func useHeaderCache(c *ggufHeaderCache) {
|
||||
prev := ggufHeaders
|
||||
ggufHeaders = c
|
||||
DeferCleanup(func() { ggufHeaders = prev })
|
||||
}
|
||||
|
||||
// persistFirstEntry parses weights.gguf through a persistent cache and writes
|
||||
// the file, so later tests start from a real cache file.
|
||||
func persistFirstEntry(dir string) {
|
||||
c := newPersistentCache(dir)
|
||||
useHeaderCache(c)
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
ExpectWithOffset(1, p.calls.Load()).To(Equal(int32(1)))
|
||||
ExpectWithOffset(1, c.flush()).To(Succeed())
|
||||
c.stopPersist()
|
||||
}
|
||||
|
||||
var _ = Describe("GGUF header cache persistence", func() {
|
||||
BeforeEach(func() {
|
||||
GinkgoT().Setenv("LOCALAI_DISABLE_GUESSING", "false")
|
||||
// Keep the background saver out of the way: each test flushes explicitly,
|
||||
// so a timer-driven write cannot race the rename override or the assertions.
|
||||
prev := headerCacheSaveDebounce
|
||||
headerCacheSaveDebounce = time.Hour
|
||||
DeferCleanup(func() { headerCacheSaveDebounce = prev })
|
||||
})
|
||||
|
||||
It("serves entries written by an earlier instance without parsing", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
|
||||
restarted := newPersistentCache(dir)
|
||||
Expect(restarted.len()).To(Equal(1))
|
||||
useHeaderCache(restarted)
|
||||
p := &countingParser{}
|
||||
cfg := presetConfig("weights.gguf")
|
||||
llamaCppDefaultsWithParser(cfg, dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(0)), "a warm restart must not parse the header")
|
||||
Expect(cfg.Name).To(Equal("header-cache-fixture"))
|
||||
})
|
||||
|
||||
It("misses on a restart when the recorded size or mtime no longer matches the file", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
path := filepath.Join(dir, "weights.gguf")
|
||||
writeHeaderCacheFile(path, 128)
|
||||
persistFirstEntry(dir)
|
||||
|
||||
writeHeaderCacheFile(path, 256)
|
||||
resized := newPersistentCache(dir)
|
||||
useHeaderCache(resized)
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)), "a changed size must be a miss")
|
||||
|
||||
persistFirstEntry(dir)
|
||||
later := time.Now().Add(time.Hour)
|
||||
ExpectWithOffset(1, os.Chtimes(path, later, later)).To(Succeed())
|
||||
touched := newPersistentCache(dir)
|
||||
useHeaderCache(touched)
|
||||
p = &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)), "a changed mtime must be a miss")
|
||||
})
|
||||
|
||||
It("ignores a file whose schema version or estimator fingerprint differs", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
path := ggufHeaderCachePath(dir, "")
|
||||
|
||||
raw, err := os.ReadFile(path)
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
var doc map[string]any
|
||||
ExpectWithOffset(1, json.Unmarshal(raw, &doc)).To(Succeed())
|
||||
|
||||
for _, tamper := range []map[string]any{{"schema": 999}, {"estimator": "other"}} {
|
||||
doc2 := map[string]any{}
|
||||
for k, v := range doc {
|
||||
doc2[k] = v
|
||||
}
|
||||
for k, v := range tamper {
|
||||
doc2[k] = v
|
||||
}
|
||||
out, err := json.Marshal(doc2)
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
ExpectWithOffset(1, os.WriteFile(path, out, 0o600)).To(Succeed())
|
||||
|
||||
c := newPersistentCache(dir)
|
||||
Expect(c.len()).To(Equal(0), "tampered file %v must be ignored", tamper)
|
||||
c.stopPersist()
|
||||
}
|
||||
})
|
||||
|
||||
It("ignores corrupt and truncated files without panicking", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
path := ggufHeaderCachePath(dir, "")
|
||||
valid, err := os.ReadFile(path)
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
|
||||
for _, content := range [][]byte{
|
||||
[]byte("{not json"),
|
||||
valid[:len(valid)/2],
|
||||
[]byte(`{"schema":1,"estimator":"x","entries":[{"path":42}]}`),
|
||||
{},
|
||||
} {
|
||||
ExpectWithOffset(1, os.WriteFile(path, content, 0o600)).To(Succeed())
|
||||
var c *ggufHeaderCache
|
||||
Expect(func() { c = newPersistentCache(dir) }).NotTo(Panic())
|
||||
Expect(c.len()).To(Equal(0))
|
||||
c.stopPersist()
|
||||
}
|
||||
})
|
||||
|
||||
It("keeps every frontend on the one shared file", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
|
||||
Expect(ggufHeaderCachePath(dir, "")).To(Equal(filepath.Join(dir, "header-cache.json")))
|
||||
Expect(ggufHeaderCachePath(dir, "")).To(BeAnExistingFile())
|
||||
|
||||
other := newPersistentCache(dir)
|
||||
Expect(other.len()).To(Equal(1), "a second frontend starts from the shared file")
|
||||
})
|
||||
|
||||
It("leaves no temp file behind after a successful write", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
|
||||
Expect(tmpFilesIn(dir)).To(BeEmpty())
|
||||
Expect(ggufHeaderCachePath(dir, "")).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("keeps the previous file intact when a write fails and removes its temp file", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
persistFirstEntry(dir)
|
||||
path := ggufHeaderCachePath(dir, "")
|
||||
before, err := os.ReadFile(path)
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
|
||||
c := newPersistentCache(dir)
|
||||
c.put(filepath.Join(dir, "other.gguf"), ggufIdentity{size: 1, mtime: 1}, ggufHeader{Name: "other"})
|
||||
|
||||
prevRename := headerCacheRename
|
||||
headerCacheRename = func(string, string) error { return errors.New("simulated rename failure") }
|
||||
DeferCleanup(func() { headerCacheRename = prevRename })
|
||||
|
||||
Expect(c.flush()).To(HaveOccurred())
|
||||
after, err := os.ReadFile(path)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(after).To(Equal(before), "a failed write must not touch the previous file")
|
||||
Expect(tmpFilesIn(dir)).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("writes at most as many entries as the in-memory bound", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
c := newPersistentCache(dir)
|
||||
for i := 0; i < ggufHeaderCacheSize+50; i++ {
|
||||
c.put(filepath.Join(dir, fmt.Sprintf("m-%03d.gguf", i)), ggufIdentity{size: 1, mtime: 1}, ggufHeader{Name: "m"})
|
||||
}
|
||||
Expect(c.flush()).To(Succeed())
|
||||
|
||||
raw, err := os.ReadFile(ggufHeaderCachePath(dir, ""))
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
var file persistedGGUFCache
|
||||
Expect(json.Unmarshal(raw, &file)).To(Succeed())
|
||||
Expect(file.Entries).To(HaveLen(ggufHeaderCacheSize))
|
||||
})
|
||||
|
||||
It("works without a data path: memory only, nothing is written", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
c := newPersistentCache("")
|
||||
useHeaderCache(c)
|
||||
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)), "the in-memory cache still works")
|
||||
|
||||
Expect(c.flush()).To(Succeed())
|
||||
entries, err := os.ReadDir(dir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(entries).To(HaveLen(1), "only the model file is present")
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,238 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// presetContext is set on fixture-backed configs. EstimateLLaMACppRun needs
|
||||
// tensors, which a fabricated file lacks, so the trained-context estimate
|
||||
// panics on it; a preset context keeps guessGGUFFromFile on the non-estimate
|
||||
// path while the cache is exercised.
|
||||
var presetContext = 8192
|
||||
|
||||
// headerCacheFixture fabricates a minimal in-memory GGUF. It carries the keys
|
||||
// guessGGUFFromFile reads (architecture, name, chat template, MTP head) so a
|
||||
// cached result can be compared field by field with a fresh parse.
|
||||
func headerCacheFixture() *gguf.GGUFFile {
|
||||
return &gguf.GGUFFile{
|
||||
Header: gguf.GGUFHeader{
|
||||
MetadataKV: gguf.GGUFMetadataKVs{
|
||||
{Key: "general.architecture", ValueType: gguf.GGUFMetadataValueTypeString, Value: "llama"},
|
||||
{Key: "general.name", ValueType: gguf.GGUFMetadataValueTypeString, Value: "header-cache-fixture"},
|
||||
{Key: "tokenizer.chat_template", ValueType: gguf.GGUFMetadataValueTypeString, Value: "{{ messages }}"},
|
||||
{Key: "llama.nextn_predict_layers", ValueType: gguf.GGUFMetadataValueTypeUint32, Value: uint32(1)},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// countingParser stands in for gguf.ParseGGUFFile. It counts calls and either
|
||||
// returns the fixture or the configured error.
|
||||
type countingParser struct {
|
||||
calls atomic.Int32
|
||||
fail error
|
||||
}
|
||||
|
||||
func (p *countingParser) parse(string, ...gguf.GGUFReadOption) (*gguf.GGUFFile, error) {
|
||||
p.calls.Add(1)
|
||||
if p.fail != nil {
|
||||
return nil, p.fail
|
||||
}
|
||||
return headerCacheFixture(), nil
|
||||
}
|
||||
|
||||
func writeHeaderCacheFile(path string, size int) {
|
||||
ExpectWithOffset(1, os.WriteFile(path, make([]byte, size), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
var _ = Describe("GGUF header cache in front of model defaults", func() {
|
||||
BeforeEach(func() { GinkgoT().Setenv("LOCALAI_DISABLE_GUESSING", "false") })
|
||||
|
||||
It("parses an unchanged file once and serves later calls from the cache", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
p := &countingParser{}
|
||||
|
||||
first := presetConfig("weights.gguf")
|
||||
llamaCppDefaultsWithParser(first, dir, p.parse)
|
||||
second := presetConfig("weights.gguf")
|
||||
llamaCppDefaultsWithParser(second, dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)), "an unchanged file must not be parsed again")
|
||||
})
|
||||
|
||||
It("re-parses when the file size changes", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
path := filepath.Join(dir, "weights.gguf")
|
||||
writeHeaderCacheFile(path, 128)
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
|
||||
writeHeaderCacheFile(path, 256)
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("re-parses when only the modification time changes", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
path := filepath.Join(dir, "weights.gguf")
|
||||
writeHeaderCacheFile(path, 128)
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
|
||||
later := time.Now().Add(time.Hour)
|
||||
ExpectWithOffset(1, os.Chtimes(path, later, later)).To(Succeed())
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("produces the same ModelConfig from a cache hit as from a fresh parse", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
p := &countingParser{}
|
||||
|
||||
fresh := presetConfig("weights.gguf")
|
||||
llamaCppDefaultsWithParser(fresh, dir, p.parse)
|
||||
hit := presetConfig("weights.gguf")
|
||||
llamaCppDefaultsWithParser(hit, dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)))
|
||||
Expect(hit.Name).To(Equal("header-cache-fixture"))
|
||||
Expect(hit.Name).To(Equal(fresh.Name))
|
||||
Expect(hit.Options).To(Equal(fresh.Options))
|
||||
Expect(hit.ContextSize).To(Equal(fresh.ContextSize))
|
||||
Expect(hit.NGPULayers).To(Equal(fresh.NGPULayers))
|
||||
Expect(hit.KnownUsecaseStrings).To(Equal(fresh.KnownUsecaseStrings))
|
||||
Expect(hit.TemplateConfig.UseTokenizerTemplate).To(Equal(fresh.TemplateConfig.UseTokenizerTemplate))
|
||||
Expect(hit).To(Equal(fresh))
|
||||
})
|
||||
|
||||
It("does not cache a missing file and keeps the existing fallback", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
p := &countingParser{fail: errors.New("no such file")}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
cfg := &ModelConfig{Backend: "llama-cpp", PredictionOptions: predictionModel("missing.gguf")}
|
||||
Expect(func() { llamaCppDefaultsWithParser(cfg, dir, p.parse) }).NotTo(Panic())
|
||||
Expect(cfg.ContextSize).To(HaveValue(Equal(DefaultContextSize)))
|
||||
Expect(cfg.NGPULayers).To(BeNil())
|
||||
}
|
||||
Expect(p.calls.Load()).To(Equal(int32(2)), "a missing file must reach the parser on every call")
|
||||
|
||||
// Once the file exists it must be parsed, not served from a stale error.
|
||||
writeHeaderCacheFile(filepath.Join(dir, "missing.gguf"), 64)
|
||||
p.fail = nil
|
||||
cfg := presetConfig("missing.gguf")
|
||||
llamaCppDefaultsWithParser(cfg, dir, p.parse)
|
||||
Expect(cfg.Name).To(Equal("header-cache-fixture"))
|
||||
})
|
||||
|
||||
It("does not cache a parse error", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
p := &countingParser{fail: errors.New("truncated header")}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
cfg := &ModelConfig{Backend: "llama-cpp", PredictionOptions: predictionModel("weights.gguf")}
|
||||
Expect(func() { llamaCppDefaultsWithParser(cfg, dir, p.parse) }).NotTo(Panic())
|
||||
Expect(cfg.Name).To(BeEmpty())
|
||||
}
|
||||
Expect(p.calls.Load()).To(Equal(int32(2)))
|
||||
})
|
||||
|
||||
It("still short-circuits on LOCALAI_DISABLE_GUESSING before consulting the cache", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "weights.gguf"), 128)
|
||||
p := &countingParser{}
|
||||
llamaCppDefaultsWithParser(presetConfig("weights.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)))
|
||||
|
||||
GinkgoT().Setenv("LOCALAI_DISABLE_GUESSING", "true")
|
||||
cfg := presetConfig("weights.gguf")
|
||||
cfg.ContextSize = nil
|
||||
llamaCppDefaultsWithParser(cfg, dir, p.parse)
|
||||
|
||||
Expect(p.calls.Load()).To(Equal(int32(1)))
|
||||
Expect(cfg.ContextSize).To(BeNil())
|
||||
Expect(cfg.Name).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("evicts entries beyond its bound so a long-unused file is parsed again", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
p := &countingParser{}
|
||||
// More distinct files than the cache holds: the first one must be evicted.
|
||||
for i := 0; i < 300; i++ {
|
||||
name := fmt.Sprintf("model-%03d.gguf", i)
|
||||
writeHeaderCacheFile(filepath.Join(dir, name), 32)
|
||||
llamaCppDefaultsWithParser(presetConfig(name), dir, p.parse)
|
||||
}
|
||||
before := p.calls.Load()
|
||||
llamaCppDefaultsWithParser(presetConfig("model-000.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(before + 1))
|
||||
|
||||
Expect(ggufHeaders.len()).To(BeNumerically("<=", ggufHeaderCacheSize))
|
||||
|
||||
// The most recent file is still cached.
|
||||
before = p.calls.Load()
|
||||
llamaCppDefaultsWithParser(presetConfig("model-299.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(before))
|
||||
})
|
||||
|
||||
It("is safe for concurrent callers on the same and different files", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
writeHeaderCacheFile(filepath.Join(dir, "a.gguf"), 128)
|
||||
writeHeaderCacheFile(filepath.Join(dir, "b.gguf"), 128)
|
||||
p := &countingParser{}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
names := make([]string, 32)
|
||||
for i := 0; i < 32; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
name := "a.gguf"
|
||||
if i%2 == 1 {
|
||||
name = "b.gguf"
|
||||
}
|
||||
cfg := presetConfig(name)
|
||||
llamaCppDefaultsWithParser(cfg, dir, p.parse)
|
||||
names[i] = cfg.Name
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
for _, n := range names {
|
||||
Expect(n).To(Equal("header-cache-fixture"))
|
||||
}
|
||||
|
||||
// Two files, so at most one parse per file in the steady state; concurrent
|
||||
// first callers may each parse before the entry lands.
|
||||
Expect(p.calls.Load()).To(BeNumerically("<=", 32))
|
||||
before := p.calls.Load()
|
||||
llamaCppDefaultsWithParser(presetConfig("a.gguf"), dir, p.parse)
|
||||
llamaCppDefaultsWithParser(presetConfig("b.gguf"), dir, p.parse)
|
||||
Expect(p.calls.Load()).To(Equal(before))
|
||||
})
|
||||
})
|
||||
|
||||
func presetConfig(name string) *ModelConfig {
|
||||
ctx := presetContext
|
||||
return &ModelConfig{Backend: "llama-cpp", PredictionOptions: predictionModel(name), LLMConfig: LLMConfig{ContextSize: &ctx}}
|
||||
}
|
||||
|
||||
func predictionModel(name string) schema.PredictionOptions {
|
||||
return schema.PredictionOptions{BasicModelRequest: schema.BasicModelRequest{Model: name}}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
|
||||
gguf "github.com/gpustack/gguf-parser-go"
|
||||
)
|
||||
|
||||
// ggufHeaderCacheSchema is the on-disk format version. A file written with any
|
||||
// other value is ignored, so bump it when persistedGGUFHeader changes shape.
|
||||
const ggufHeaderCacheSchema = 1
|
||||
|
||||
// ggufHeaderSpec makes the header cache a GGUF cache: the schema, the estimator
|
||||
// fingerprint and the GGUF field mapping. The file name is the generic one.
|
||||
var ggufHeaderSpec = headerCacheSpec[ggufHeader]{
|
||||
name: "GGUF header cache",
|
||||
schema: ggufHeaderCacheSchema,
|
||||
fingerprint: ggufEstimatorFingerprint(),
|
||||
parse: func(path string) (ggufHeader, error) {
|
||||
return parseGGUFHeader(path, gguf.ParseGGUFFile)
|
||||
},
|
||||
encode: func(h ggufHeader) any { return toPersistedHeader(h) },
|
||||
decode: decodeGGUFHeader,
|
||||
}
|
||||
|
||||
// persistedGGUFHeader mirrors ggufHeader. Every field is either a property of
|
||||
// the file bytes or an estimator output whose inputs are fixed in code (no
|
||||
// device metrics, fixed candidate contexts and offload), keyed by the estimator
|
||||
// fingerprint. Nothing hardware-dependent is stored: the per-device budget is
|
||||
// read at decision time from perDeviceVRAM, not from this file.
|
||||
type persistedGGUFHeader struct {
|
||||
Name string `json:"name"`
|
||||
Architecture string `json:"architecture"`
|
||||
MaxContextLength int `json:"max_context_length"`
|
||||
BOSTokenID int64 `json:"bos_token_id"`
|
||||
EOSTokenID int64 `json:"eos_token_id"`
|
||||
ChatTemplate string `json:"chat_template"`
|
||||
HasChatTemplate bool `json:"has_chat_template"`
|
||||
MTPLayers uint32 `json:"mtp_layers"`
|
||||
HasMTP bool `json:"has_mtp"`
|
||||
TrainedContext int `json:"trained_context"`
|
||||
Footprints map[int]uint64 `json:"footprints,omitempty"`
|
||||
MaxCtxFootprint uint64 `json:"max_ctx_footprint"`
|
||||
// TrainedPanic is the estimator panic message, kept so the replay still
|
||||
// happens for this file. The replay panics with the message as a string,
|
||||
// which the defaults path recovers the same way.
|
||||
TrainedPanic *string `json:"trained_panic,omitempty"`
|
||||
}
|
||||
|
||||
// ggufHeaderCachePath returns the GGUF header cache file for instanceID in dataPath.
|
||||
func ggufHeaderCachePath(dataPath, instanceID string) string {
|
||||
return headerCachePath(dataPath, instanceID)
|
||||
}
|
||||
|
||||
// ggufEstimatorFingerprint names the inputs that decide the estimator-derived
|
||||
// fields: the gguf-parser version, the candidate contexts, the auto context
|
||||
// size and the offload count. Changing any of them changes the file's key.
|
||||
func ggufEstimatorFingerprint() string {
|
||||
parser := "unknown"
|
||||
if bi, ok := debug.ReadBuildInfo(); ok {
|
||||
for _, d := range bi.Deps {
|
||||
if d.Path == "github.com/gpustack/gguf-parser-go" {
|
||||
parser = d.Version
|
||||
if d.Replace != nil {
|
||||
parser = d.Replace.Path + "@" + d.Replace.Version
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("gguf-parser=%s;candidates=%v;auto=%d;ngl=%d",
|
||||
parser, contextFitCandidates, DefaultAutoContextSize, DefaultNGPULayers)
|
||||
}
|
||||
|
||||
// decodeGGUFHeader rebuilds a snapshot from its file form. An absent header
|
||||
// decodes to the zero snapshot, as a missing JSON value always did.
|
||||
func decodeGGUFHeader(raw json.RawMessage) (ggufHeader, error) {
|
||||
var p persistedGGUFHeader
|
||||
if len(raw) > 0 {
|
||||
if err := json.Unmarshal(raw, &p); err != nil {
|
||||
return ggufHeader{}, err
|
||||
}
|
||||
}
|
||||
return fromPersistedHeader(p), nil
|
||||
}
|
||||
|
||||
func toPersistedHeader(h ggufHeader) persistedGGUFHeader {
|
||||
var panicked *string
|
||||
if h.trainedPanic != nil {
|
||||
msg := fmt.Sprint(h.trainedPanic)
|
||||
panicked = &msg
|
||||
}
|
||||
return persistedGGUFHeader{
|
||||
Name: h.Name,
|
||||
Architecture: h.Architecture,
|
||||
MaxContextLength: h.MaxContextLength,
|
||||
BOSTokenID: h.BOSTokenID,
|
||||
EOSTokenID: h.EOSTokenID,
|
||||
ChatTemplate: h.ChatTemplate,
|
||||
HasChatTemplate: h.HasChatTemplate,
|
||||
MTPLayers: h.MTPLayers,
|
||||
HasMTP: h.HasMTP,
|
||||
TrainedContext: h.TrainedContext,
|
||||
Footprints: h.footprints,
|
||||
MaxCtxFootprint: h.maxCtxFootprint,
|
||||
TrainedPanic: panicked,
|
||||
}
|
||||
}
|
||||
|
||||
func fromPersistedHeader(p persistedGGUFHeader) ggufHeader {
|
||||
var panicked any
|
||||
if p.TrainedPanic != nil {
|
||||
panicked = *p.TrainedPanic
|
||||
}
|
||||
return ggufHeader{
|
||||
Name: p.Name,
|
||||
Architecture: p.Architecture,
|
||||
MaxContextLength: p.MaxContextLength,
|
||||
BOSTokenID: p.BOSTokenID,
|
||||
EOSTokenID: p.EOSTokenID,
|
||||
ChatTemplate: p.ChatTemplate,
|
||||
HasChatTemplate: p.HasChatTemplate,
|
||||
MTPLayers: p.MTPLayers,
|
||||
HasMTP: p.HasMTP,
|
||||
TrainedContext: p.TrainedContext,
|
||||
footprints: p.Footprints,
|
||||
maxCtxFootprint: p.MaxCtxFootprint,
|
||||
trainedPanic: panicked,
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureGGUFHeaderCache attaches the process-wide header cache to the shared
|
||||
// header-cache.json in dataPath. Call it once at startup, before model configs
|
||||
// are loaded. An empty dataPath (standalone without data) keeps the cache in
|
||||
// memory only. Frontends that share the data volume share the file; see the
|
||||
// shared-by-file mode in the model configuration docs.
|
||||
func ConfigureGGUFHeaderCache(ctx context.Context, dataPath, instanceID string) {
|
||||
ggufHeaders.attachDisk(ctx, dataPath, instanceID)
|
||||
}
|
||||
@@ -0,0 +1,589 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// headerCacheFileName is the one cache file every frontend uses, in the data
|
||||
// path. Frontends that share the data volume share this file: see the
|
||||
// shared-by-file mode in docs/content/advanced/model-configuration.md.
|
||||
const headerCacheFileName = "header-cache.json"
|
||||
|
||||
// headerCacheMaxFileBytes refuses to read a cache file larger than this. The
|
||||
// entry bound already caps the size a LocalAI build writes, so anything bigger
|
||||
// was written by something else and is not worth decoding.
|
||||
const headerCacheMaxFileBytes = 64 << 20
|
||||
|
||||
// headerCacheSaveDebounce coalesces the inserts of a burst (a fresh start
|
||||
// touching many files) into one write.
|
||||
var headerCacheSaveDebounce = 2 * time.Second
|
||||
|
||||
// headerCacheRename is os.Rename behind a variable so a test can simulate a
|
||||
// failed commit and check that the previous file survives.
|
||||
var headerCacheRename = os.Rename
|
||||
|
||||
// headerCacheSpec describes one file format's use of the header cache. S is the
|
||||
// snapshot kept per file path. The cache never looks inside S: everything that
|
||||
// is format-specific comes from the spec.
|
||||
type headerCacheSpec[S any] struct {
|
||||
// name prefixes log messages, for example "GGUF header cache".
|
||||
name string
|
||||
// schema is the on-disk format version. Bump it when encode changes shape.
|
||||
schema int
|
||||
// fingerprint names the inputs that decide the snapshot values. A file
|
||||
// written with another fingerprint may hold different values, so it is
|
||||
// ignored.
|
||||
fingerprint string
|
||||
// parse builds a snapshot from a file path. Its errors are never cached.
|
||||
parse func(path string) (S, error)
|
||||
// encode returns the JSON-serialisable form of a snapshot.
|
||||
encode func(S) any
|
||||
// decode rebuilds a snapshot from what encode wrote. An error means the
|
||||
// file is corrupt, so the whole file is ignored rather than partly loaded.
|
||||
decode func(json.RawMessage) (S, error)
|
||||
}
|
||||
|
||||
// headerIdentity is what a cached snapshot is validated against. A changed size
|
||||
// or modification time means the file was rewritten, so the entry is not used.
|
||||
type headerIdentity struct {
|
||||
size int64
|
||||
mtime int64 // UnixNano
|
||||
}
|
||||
|
||||
type headerEntry[S any] struct {
|
||||
path string
|
||||
identity headerIdentity
|
||||
snap S
|
||||
}
|
||||
|
||||
// headerCache is a bounded LRU keyed by absolute path, optionally persisted to
|
||||
// the shared header-cache.json. It is safe for concurrent use. Lookups hold the
|
||||
// lock only for the map and list operations, never while a file is parsed or
|
||||
// read.
|
||||
type headerCache[S any] struct {
|
||||
spec headerCacheSpec[S]
|
||||
mu sync.Mutex
|
||||
max int
|
||||
order *list.List // front is most recently used; values are *headerEntry[S]
|
||||
index map[string]*list.Element
|
||||
|
||||
// Disk persistence, set by attachDisk. Empty persistPath means memory only.
|
||||
persistPath string
|
||||
dirty bool
|
||||
saveCh chan struct{}
|
||||
stopSaver context.CancelFunc
|
||||
saverDone chan struct{}
|
||||
// loadedMtime is the modification time of the cache file the last time it
|
||||
// was read, UnixNano. Zero means it has not been read. A miss compares it
|
||||
// with the file on disk to decide whether another frontend has written.
|
||||
loadedMtime int64
|
||||
// persistMu serializes file writes so the saver and an explicit flush never
|
||||
// race on the same temp name or rename order.
|
||||
persistMu sync.Mutex
|
||||
// reloadMu keeps concurrent misses from each reading the same file.
|
||||
reloadMu sync.Mutex
|
||||
}
|
||||
|
||||
func newHeaderCache[S any](spec headerCacheSpec[S], max int) *headerCache[S] {
|
||||
return &headerCache[S]{spec: spec, max: max, order: list.New(), index: make(map[string]*list.Element)}
|
||||
}
|
||||
|
||||
// get returns the cached snapshot for path when its identity matches id. It
|
||||
// never reads the disk: a hit is served from memory alone.
|
||||
func (c *headerCache[S]) get(path string, id headerIdentity) (S, bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
el, ok := c.index[path]
|
||||
if !ok {
|
||||
var zero S
|
||||
return zero, false
|
||||
}
|
||||
e := el.Value.(*headerEntry[S])
|
||||
if e.identity != id {
|
||||
// The file changed since it was cached. Drop the stale entry now so the
|
||||
// fresh parse replaces it rather than sitting beside it.
|
||||
c.order.Remove(el)
|
||||
delete(c.index, path)
|
||||
c.markDirtyLocked()
|
||||
var zero S
|
||||
return zero, false
|
||||
}
|
||||
c.order.MoveToFront(el)
|
||||
return e.snap, true
|
||||
}
|
||||
|
||||
// put stores a snapshot and schedules a background save when the cache is
|
||||
// persisted. It never touches the disk itself, so it is safe on the request path.
|
||||
func (c *headerCache[S]) put(path string, id headerIdentity, snap S) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.insertLocked(path, id, snap)
|
||||
c.markDirtyLocked()
|
||||
}
|
||||
|
||||
// insertLocked adds or replaces an entry and enforces the bound. It does not
|
||||
// mark the cache dirty, so loading from disk does not schedule a rewrite.
|
||||
func (c *headerCache[S]) insertLocked(path string, id headerIdentity, snap S) {
|
||||
if el, ok := c.index[path]; ok {
|
||||
e := el.Value.(*headerEntry[S])
|
||||
e.identity, e.snap = id, snap
|
||||
c.order.MoveToFront(el)
|
||||
return
|
||||
}
|
||||
c.index[path] = c.order.PushFront(&headerEntry[S]{path: path, identity: id, snap: snap})
|
||||
c.evictLocked()
|
||||
}
|
||||
|
||||
// insertBackLocked adds an entry as the least recently used one. Entries
|
||||
// merged from the file use it, so that what this frontend used itself keeps
|
||||
// its place ahead of what only another frontend used.
|
||||
func (c *headerCache[S]) insertBackLocked(path string, id headerIdentity, snap S) {
|
||||
c.index[path] = c.order.PushBack(&headerEntry[S]{path: path, identity: id, snap: snap})
|
||||
c.evictLocked()
|
||||
}
|
||||
|
||||
// evictLocked drops the least recently used entries until the bound holds.
|
||||
func (c *headerCache[S]) evictLocked() {
|
||||
for c.order.Len() > c.max {
|
||||
oldest := c.order.Back()
|
||||
c.order.Remove(oldest)
|
||||
delete(c.index, oldest.Value.(*headerEntry[S]).path)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *headerCache[S]) len() int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.order.Len()
|
||||
}
|
||||
|
||||
// load returns the snapshot for path using the spec's parser.
|
||||
func (c *headerCache[S]) load(path string) (S, error) {
|
||||
return c.loadVia(path, c.spec.parse)
|
||||
}
|
||||
|
||||
// loadVia returns the snapshot for path, calling parse only when no cached
|
||||
// snapshot matches the file's current identity. A file that cannot be stat'ed
|
||||
// or parsed is never cached, so the caller sees the same error path as before
|
||||
// the cache existed. Callers pass parse to swap in a test seam per call.
|
||||
//
|
||||
// On a miss the shared cache file is checked once: if another frontend has
|
||||
// written it since this one last read it, its entries are merged in before the
|
||||
// parser runs. A hit never touches the cache file.
|
||||
func (c *headerCache[S]) loadVia(path string, parse func(string) (S, error)) (S, error) {
|
||||
abs, absErr := filepath.Abs(path)
|
||||
info, statErr := os.Stat(path)
|
||||
if absErr != nil || statErr != nil {
|
||||
return parse(path)
|
||||
}
|
||||
|
||||
// The identity is read before the parse. If the file is rewritten during the
|
||||
// parse, the entry is keyed by the older identity and the next call misses.
|
||||
id := headerIdentity{size: info.Size(), mtime: info.ModTime().UnixNano()}
|
||||
if snap, ok := c.get(abs, id); ok {
|
||||
return snap, nil
|
||||
}
|
||||
c.reloadIfChanged()
|
||||
if snap, ok := c.get(abs, id); ok {
|
||||
return snap, nil
|
||||
}
|
||||
|
||||
snap, err := parse(path)
|
||||
if err != nil {
|
||||
var zero S
|
||||
return zero, err
|
||||
}
|
||||
c.put(abs, id, snap)
|
||||
return snap, nil
|
||||
}
|
||||
|
||||
// reloadIfChanged merges the cache file into memory when its modification time
|
||||
// differs from the one recorded at the last read. It does nothing when the
|
||||
// cache is not persistent or the file is absent.
|
||||
func (c *headerCache[S]) reloadIfChanged() {
|
||||
c.reloadMu.Lock()
|
||||
defer c.reloadMu.Unlock()
|
||||
|
||||
c.mu.Lock()
|
||||
path, loaded := c.persistPath, c.loadedMtime
|
||||
c.mu.Unlock()
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || info.ModTime().UnixNano() == loaded {
|
||||
return
|
||||
}
|
||||
c.readAndMerge(path)
|
||||
}
|
||||
|
||||
// headerFile is the whole cache file, shared by every frontend on the data path.
|
||||
type headerFile struct {
|
||||
Schema int `json:"schema"`
|
||||
// Fingerprint is stored under the key "estimator", the name the GGUF cache
|
||||
// wrote before the cache was generic. The key is kept so existing files stay
|
||||
// valid.
|
||||
Fingerprint string `json:"estimator"`
|
||||
Entries []headerFileEntry `json:"entries"`
|
||||
}
|
||||
|
||||
type headerFileEntry struct {
|
||||
Path string `json:"path"`
|
||||
Size int64 `json:"size"`
|
||||
MtimeNS int64 `json:"mtime_ns"`
|
||||
Header json.RawMessage `json:"header"`
|
||||
}
|
||||
|
||||
// attachDisk makes the cache persistent under dataPath, in the shared file. It
|
||||
// loads the existing file once, then starts the background saver. An empty
|
||||
// dataPath keeps the cache in memory only. Calling it again replaces the
|
||||
// previous attach.
|
||||
func (c *headerCache[S]) attachDisk(parent context.Context, dataPath, instanceID string) {
|
||||
c.stopPersist()
|
||||
if dataPath == "" {
|
||||
return
|
||||
}
|
||||
path := headerCachePath(dataPath, instanceID)
|
||||
c.readAndMerge(path)
|
||||
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
done := make(chan struct{})
|
||||
saveCh := make(chan struct{}, 1)
|
||||
c.mu.Lock()
|
||||
c.persistPath = path
|
||||
c.saveCh = saveCh
|
||||
c.stopSaver = cancel
|
||||
c.saverDone = done
|
||||
c.mu.Unlock()
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
c.runSaver(ctx, saveCh)
|
||||
}()
|
||||
}
|
||||
|
||||
// headerCachePath returns the cache file in dataPath. An empty instanceID
|
||||
// selects the shared header-cache.json. A non-empty instanceID selects a file
|
||||
// of its own, so one frontend's entries never mix with another's. Unusual IDs
|
||||
// are hashed so the file name stays a plain token.
|
||||
func headerCachePath(dataPath, instanceID string) string {
|
||||
if instanceID == "" {
|
||||
return filepath.Join(dataPath, headerCacheFileName)
|
||||
}
|
||||
name := instanceID
|
||||
if !headerInstanceIDPattern.MatchString(instanceID) {
|
||||
sum := sha256.Sum256([]byte(instanceID))
|
||||
name = "x" + hex.EncodeToString(sum[:16])
|
||||
}
|
||||
return filepath.Join(dataPath, "header-cache-"+name+".json")
|
||||
}
|
||||
|
||||
var headerInstanceIDPattern = regexp.MustCompile(`^[A-Za-z0-9_-]{1,64}$`)
|
||||
|
||||
// stopPersist stops the saver after a final write of any pending change and
|
||||
// disables persistence. It is safe to call when nothing is attached.
|
||||
func (c *headerCache[S]) stopPersist() {
|
||||
c.mu.Lock()
|
||||
cancel, done := c.stopSaver, c.saverDone
|
||||
c.stopSaver, c.saverDone = nil, nil
|
||||
c.mu.Unlock()
|
||||
if cancel == nil {
|
||||
return
|
||||
}
|
||||
cancel()
|
||||
<-done
|
||||
|
||||
c.mu.Lock()
|
||||
c.persistPath = ""
|
||||
c.saveCh = nil
|
||||
c.dirty = false
|
||||
c.loadedMtime = 0
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// markDirtyLocked records a change and wakes the saver. The send never blocks:
|
||||
// one pending signal is enough because the saver writes the whole cache.
|
||||
func (c *headerCache[S]) markDirtyLocked() {
|
||||
if c.persistPath == "" {
|
||||
return
|
||||
}
|
||||
c.dirty = true
|
||||
select {
|
||||
case c.saveCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// runSaver writes the cache after changes, debounced, and once more on shutdown.
|
||||
// It is the only background writer; request paths never touch the disk.
|
||||
func (c *headerCache[S]) runSaver(ctx context.Context, saveCh <-chan struct{}) {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
c.flushIfDirty()
|
||||
return
|
||||
case <-saveCh:
|
||||
timer := time.NewTimer(headerCacheSaveDebounce)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
c.flushIfDirty()
|
||||
return
|
||||
case <-timer.C:
|
||||
}
|
||||
c.flushIfDirty()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *headerCache[S]) flushIfDirty() {
|
||||
c.mu.Lock()
|
||||
dirty := c.dirty
|
||||
c.mu.Unlock()
|
||||
if !dirty {
|
||||
return
|
||||
}
|
||||
if err := c.flush(); err != nil {
|
||||
xlog.Warn(c.spec.name+": could not save, it will retry on the next change", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// flush merges the shared file into memory, drops entries whose model file has
|
||||
// changed, and writes the result back. It returns nil and writes nothing when
|
||||
// the cache is not persistent.
|
||||
//
|
||||
// Several frontends may write the file at once. Each write is a whole-file
|
||||
// rename, so a reader never sees a partial file. Two writers can still race: a
|
||||
// writer that read the file before the other renamed can replace the other's
|
||||
// new entries with its own. The loss is at worst a dropped entry, which costs
|
||||
// one re-parse on its next use. A wrong value cannot result, because every
|
||||
// entry is validated against its model file's identity before it is used.
|
||||
func (c *headerCache[S]) flush() error {
|
||||
c.persistMu.Lock()
|
||||
defer c.persistMu.Unlock()
|
||||
|
||||
c.mu.Lock()
|
||||
path := c.persistPath
|
||||
c.mu.Unlock()
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fold in what other frontends saved since this one last read the file, so
|
||||
// this write does not drop their entries.
|
||||
c.readAndMerge(path)
|
||||
|
||||
type pending struct {
|
||||
path string
|
||||
id headerIdentity
|
||||
snap S
|
||||
}
|
||||
c.mu.Lock()
|
||||
// Copy under the lock and encode after it: encoding is the slow part and
|
||||
// must not hold up request lookups.
|
||||
snaps := make([]pending, 0, c.order.Len())
|
||||
for el := c.order.Front(); el != nil; el = el.Next() {
|
||||
e := el.Value.(*headerEntry[S])
|
||||
snaps = append(snaps, pending{path: e.path, id: e.identity, snap: e.snap})
|
||||
}
|
||||
c.dirty = false
|
||||
c.mu.Unlock()
|
||||
|
||||
// An entry whose model file has been rewritten since it was cached can
|
||||
// never be used, so it is not written. A file that cannot be stat'ed is
|
||||
// kept: the identity check at lookup still decides whether it is used.
|
||||
stale := make(map[string]headerIdentity)
|
||||
for _, p := range snaps {
|
||||
info, err := os.Stat(p.path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if now := (headerIdentity{size: info.Size(), mtime: info.ModTime().UnixNano()}); now != p.id {
|
||||
stale[p.path] = p.id
|
||||
}
|
||||
}
|
||||
if len(stale) > 0 {
|
||||
c.mu.Lock()
|
||||
for p, id := range stale {
|
||||
if el, ok := c.index[p]; ok && el.Value.(*headerEntry[S]).identity == id {
|
||||
c.order.Remove(el)
|
||||
delete(c.index, p)
|
||||
}
|
||||
}
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
file := headerFile{
|
||||
Schema: c.spec.schema,
|
||||
Fingerprint: c.spec.fingerprint,
|
||||
Entries: make([]headerFileEntry, 0, len(snaps)),
|
||||
}
|
||||
data, err := func() ([]byte, error) {
|
||||
for _, p := range snaps {
|
||||
if _, drop := stale[p.path]; drop {
|
||||
continue
|
||||
}
|
||||
raw, err := json.Marshal(c.spec.encode(p.snap))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
file.Entries = append(file.Entries, headerFileEntry{
|
||||
Path: p.path,
|
||||
Size: p.id.size,
|
||||
MtimeNS: p.id.mtime,
|
||||
Header: raw,
|
||||
})
|
||||
}
|
||||
return json.Marshal(file)
|
||||
}()
|
||||
if err == nil {
|
||||
err = writeFileAtomic(path, data)
|
||||
}
|
||||
if err != nil {
|
||||
c.mu.Lock()
|
||||
c.dirty = true
|
||||
c.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
xlog.Debug(c.spec.name+" saved", "path", path, "entries", len(file.Entries))
|
||||
return nil
|
||||
}
|
||||
|
||||
// readAndMerge reads the cache file at path and merges its entries into memory.
|
||||
// A missing, unreadable, oversized, corrupt or mismatched file is ignored and
|
||||
// logged once. The modification time recorded here is taken before the read, so
|
||||
// a write that lands during the read makes the next miss read again.
|
||||
func (c *headerCache[S]) readAndMerge(path string) {
|
||||
name := c.spec.name
|
||||
info, err := os.Stat(path)
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
xlog.Debug(name+": no file yet", "path", path)
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
xlog.Warn(name+": cannot stat file, starting empty", "path", path, "error", err)
|
||||
return
|
||||
}
|
||||
// Recorded even when the file is then ignored: a corrupt file would
|
||||
// otherwise be re-read and re-logged on every miss until it changes.
|
||||
mtime := info.ModTime().UnixNano()
|
||||
c.mu.Lock()
|
||||
c.loadedMtime = mtime
|
||||
c.mu.Unlock()
|
||||
if info.Size() > headerCacheMaxFileBytes {
|
||||
xlog.Warn(name+": file too large, ignoring it", "path", path, "bytes", info.Size())
|
||||
return
|
||||
}
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
xlog.Warn(name+": cannot read file, starting empty", "path", path, "error", err)
|
||||
return
|
||||
}
|
||||
var file headerFile
|
||||
if err := json.Unmarshal(raw, &file); err != nil {
|
||||
xlog.Warn(name+": file is corrupt, ignoring it", "path", path, "error", err)
|
||||
return
|
||||
}
|
||||
if file.Schema != c.spec.schema {
|
||||
xlog.Warn(name+": unsupported schema, ignoring it", "path", path, "schema", file.Schema)
|
||||
return
|
||||
}
|
||||
if file.Fingerprint != c.spec.fingerprint {
|
||||
xlog.Debug(name+": format fingerprint changed, ignoring file", "path", path)
|
||||
return
|
||||
}
|
||||
|
||||
type loaded struct {
|
||||
path string
|
||||
id headerIdentity
|
||||
snap S
|
||||
}
|
||||
decoded := make([]loaded, 0, len(file.Entries))
|
||||
for _, e := range file.Entries {
|
||||
snap, err := c.spec.decode(e.Header)
|
||||
if err != nil {
|
||||
xlog.Warn(name+": file is corrupt, ignoring it", "path", path, "error", err)
|
||||
return
|
||||
}
|
||||
if e.Path == "" || !filepath.IsAbs(e.Path) {
|
||||
continue
|
||||
}
|
||||
decoded = append(decoded, loaded{path: e.Path, id: headerIdentity{size: e.Size, mtime: e.MtimeNS}, snap: snap})
|
||||
}
|
||||
|
||||
// Identities are read from the model files before the lock is taken: a
|
||||
// stat per entry is cheap, but holding the lock across them is not.
|
||||
current := make(map[string]headerIdentity, len(decoded))
|
||||
for _, d := range decoded {
|
||||
if fi, err := os.Stat(d.path); err == nil {
|
||||
current[d.path] = headerIdentity{size: fi.Size(), mtime: fi.ModTime().UnixNano()}
|
||||
}
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
// The file lists the most recently used entry first, so entries are merged
|
||||
// from the back: the most recent one of the file ends up nearest the front
|
||||
// of the entries this frontend does not already hold.
|
||||
for i := len(decoded) - 1; i >= 0; i-- {
|
||||
d := decoded[i]
|
||||
if el, exists := c.index[d.path]; exists {
|
||||
// Both frontends hold an entry for this model. Keep the one that
|
||||
// matches the model file on disk. If neither does, drop both. If the
|
||||
// file cannot be stat'ed, keep the one already in memory.
|
||||
e := el.Value.(*headerEntry[S])
|
||||
now, ok := current[d.path]
|
||||
switch {
|
||||
case !ok:
|
||||
case e.identity == now:
|
||||
case d.id == now:
|
||||
e.identity, e.snap = d.id, d.snap
|
||||
default:
|
||||
c.order.Remove(el)
|
||||
delete(c.index, d.path)
|
||||
}
|
||||
continue
|
||||
}
|
||||
c.insertBackLocked(d.path, d.id, d.snap)
|
||||
}
|
||||
}
|
||||
|
||||
// writeFileAtomic writes data to a temp file in path's directory and renames it
|
||||
// into place, so a reader sees either the old file or the new one. On any
|
||||
// failure the temp file is removed and the previous file is left as it was.
|
||||
func writeFileAtomic(path string, data []byte) (err error) {
|
||||
tmp, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path)+".tmp-*")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
if err != nil {
|
||||
_ = tmp.Close()
|
||||
_ = os.Remove(tmp.Name())
|
||||
}
|
||||
}()
|
||||
if _, err = tmp.Write(data); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = tmp.Sync(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err = tmp.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
return headerCacheRename(tmp.Name(), path)
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("header cache file naming", func() {
|
||||
It("uses the shared file when no instance ID is set", func() {
|
||||
Expect(headerCachePath("/data", "")).To(Equal(filepath.Join("/data", "header-cache.json")))
|
||||
})
|
||||
|
||||
It("uses a file of its own for a plain instance ID", func() {
|
||||
Expect(headerCachePath("/data", "local-ai-frontend-0")).To(Equal(filepath.Join("/data", "header-cache-local-ai-frontend-0.json")))
|
||||
})
|
||||
|
||||
It("hashes an instance ID that is not a plain token", func() {
|
||||
path := headerCachePath("/data", "a/b c")
|
||||
Expect(filepath.Base(path)).To(HavePrefix("header-cache-x"))
|
||||
Expect(strings.ContainsAny(filepath.Base(path), "/ ")).To(BeFalse())
|
||||
Expect(headerCachePath("/data", "a/b c")).To(Equal(path))
|
||||
Expect(headerCachePath("/data", "a/b d")).NotTo(Equal(path))
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,346 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// fakeSnap is a snapshot of a made-up format. The generic cache must work with
|
||||
// it exactly as it does with the GGUF snapshot, without knowing its fields.
|
||||
type fakeSnap struct {
|
||||
Name string `json:"name"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
}
|
||||
|
||||
// fakeSpecWith returns a spec whose parser counts its calls through parses.
|
||||
func fakeSpecWith(parses *int32, fingerprint string, schema int) headerCacheSpec[fakeSnap] {
|
||||
return headerCacheSpec[fakeSnap]{
|
||||
name: "fake header cache",
|
||||
schema: schema,
|
||||
fingerprint: fingerprint,
|
||||
parse: func(path string) (fakeSnap, error) {
|
||||
*parses++
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return fakeSnap{}, err
|
||||
}
|
||||
return fakeSnap{Name: filepath.Base(path), Bytes: info.Size()}, nil
|
||||
},
|
||||
encode: func(s fakeSnap) any { return s },
|
||||
decode: func(raw json.RawMessage) (fakeSnap, error) {
|
||||
var s fakeSnap
|
||||
err := json.Unmarshal(raw, &s)
|
||||
return s, err
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// newFakeCache returns a cache for the fake format attached to dataPath, the
|
||||
// way startup attaches the GGUF cache. The saver is stopped in cleanup.
|
||||
func newFakeCache(dataPath string, spec headerCacheSpec[fakeSnap]) *headerCache[fakeSnap] {
|
||||
c := newHeaderCache(spec, 8)
|
||||
c.attachDisk(context.Background(), dataPath, "")
|
||||
DeferCleanup(c.stopPersist)
|
||||
return c
|
||||
}
|
||||
|
||||
func writeFakeModel(path string, size int) {
|
||||
ExpectWithOffset(1, os.WriteFile(path, make([]byte, size), 0o600)).To(Succeed())
|
||||
}
|
||||
|
||||
func fakeCacheFile(dataPath string) string {
|
||||
return headerCachePath(dataPath, "")
|
||||
}
|
||||
|
||||
var _ = Describe("Generic header cache", func() {
|
||||
var parses int32
|
||||
|
||||
BeforeEach(func() {
|
||||
parses = 0
|
||||
// Each test flushes explicitly, so the timer-driven saver stays out of the way.
|
||||
prev := headerCacheSaveDebounce
|
||||
headerCacheSaveDebounce = time.Hour
|
||||
DeferCleanup(func() { headerCacheSaveDebounce = prev })
|
||||
})
|
||||
|
||||
It("serves a snapshot from the file after a restart without parsing", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 42)
|
||||
|
||||
first := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
snap, err := first.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(snap).To(Equal(fakeSnap{Name: "model.bin", Bytes: 42}))
|
||||
Expect(parses).To(Equal(int32(1)))
|
||||
Expect(first.flush()).To(Succeed())
|
||||
first.stopPersist()
|
||||
|
||||
restarted := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
Expect(restarted.len()).To(Equal(1))
|
||||
again, err := restarted.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(again).To(Equal(snap))
|
||||
Expect(parses).To(Equal(int32(1)), "the restarted cache must not parse again")
|
||||
})
|
||||
|
||||
It("ignores a file whose schema version differs", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
c := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = c.load(model)
|
||||
Expect(c.flush()).To(Succeed())
|
||||
c.stopPersist()
|
||||
|
||||
other := newFakeCache(dir, fakeSpecWith(&parses, "v1", 2))
|
||||
Expect(other.len()).To(Equal(0), "a file with another schema must be ignored")
|
||||
})
|
||||
|
||||
It("ignores a file whose format fingerprint differs", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
c := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = c.load(model)
|
||||
Expect(c.flush()).To(Succeed())
|
||||
c.stopPersist()
|
||||
|
||||
other := newFakeCache(dir, fakeSpecWith(&parses, "v2", 1))
|
||||
Expect(other.len()).To(Equal(0), "a file with another fingerprint must be ignored")
|
||||
})
|
||||
|
||||
It("misses when the file identity no longer matches", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
c := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = c.load(model)
|
||||
Expect(parses).To(Equal(int32(1)))
|
||||
|
||||
writeFakeModel(model, 16)
|
||||
snap, err := c.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(parses).To(Equal(int32(2)), "a changed size must be a miss")
|
||||
Expect(snap.Bytes).To(Equal(int64(16)))
|
||||
})
|
||||
|
||||
It("ignores a corrupt file without panicking", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
path := fakeCacheFile(dir)
|
||||
|
||||
for _, content := range []string{
|
||||
"{not json",
|
||||
`{"schema":1,"estimator":"v1","entries":[{"path":42}]}`,
|
||||
`{"schema":1,"estimator":"v1","entries":[{"path":"/x","header":{"name":5}}]}`,
|
||||
} {
|
||||
ExpectWithOffset(1, os.WriteFile(path, []byte(content), 0o600)).To(Succeed())
|
||||
var c *headerCache[fakeSnap]
|
||||
Expect(func() { c = newFakeCache(dir, fakeSpecWith(&parses, "v1", 1)) }).NotTo(Panic())
|
||||
Expect(c.len()).To(Equal(0), "content %q must be ignored", content)
|
||||
c.stopPersist()
|
||||
}
|
||||
})
|
||||
|
||||
It("leaves no temp file behind after a write", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
c := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = c.load(model)
|
||||
Expect(c.flush()).To(Succeed())
|
||||
|
||||
entries, err := os.ReadDir(dir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
for _, e := range entries {
|
||||
Expect(strings.Contains(e.Name(), ".tmp")).To(BeFalse(), "temp file %s left behind", e.Name())
|
||||
}
|
||||
Expect(fakeCacheFile(dir)).To(BeAnExistingFile())
|
||||
})
|
||||
|
||||
It("works without a data path: memory only, nothing is written", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
c := newFakeCache("", fakeSpecWith(&parses, "v1", 1))
|
||||
|
||||
_, _ = c.load(model)
|
||||
_, _ = c.load(model)
|
||||
Expect(parses).To(Equal(int32(1)), "the in-memory cache still works")
|
||||
Expect(c.flush()).To(Succeed())
|
||||
entries, err := os.ReadDir(dir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(entries).To(HaveLen(1), "only the model file is present")
|
||||
})
|
||||
|
||||
It("uses header-cache.json for every instance", func() {
|
||||
Expect(headerCachePath("/data", "")).To(Equal(filepath.Join("/data", "header-cache.json")))
|
||||
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
a := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
b := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = a.load(model)
|
||||
Expect(a.flush()).To(Succeed())
|
||||
Expect(b.flush()).To(Succeed())
|
||||
|
||||
entries, err := os.ReadDir(dir)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
names := []string{}
|
||||
for _, e := range entries {
|
||||
names = append(names, e.Name())
|
||||
}
|
||||
Expect(names).To(ConsistOf("model.bin", "header-cache.json"))
|
||||
})
|
||||
|
||||
It("serves an entry written by another instance after a reload on miss, without parsing", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 42)
|
||||
|
||||
reader := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
writer := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, err := writer.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(parses).To(Equal(int32(1)))
|
||||
Expect(writer.flush()).To(Succeed())
|
||||
|
||||
snap, err := reader.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(snap).To(Equal(fakeSnap{Name: "model.bin", Bytes: 42}))
|
||||
Expect(parses).To(Equal(int32(1)), "the reader must reload the file on its miss instead of parsing")
|
||||
})
|
||||
|
||||
It("merges on save so that the file holds the union of two instances", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
m1 := filepath.Join(dir, "one.bin")
|
||||
m2 := filepath.Join(dir, "two.bin")
|
||||
writeFakeModel(m1, 1)
|
||||
writeFakeModel(m2, 2)
|
||||
|
||||
a := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
b := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
_, _ = a.load(m1)
|
||||
_, _ = b.load(m2)
|
||||
Expect(a.flush()).To(Succeed())
|
||||
Expect(b.flush()).To(Succeed())
|
||||
|
||||
file := readFakeFile(dir)
|
||||
paths := []string{}
|
||||
for _, e := range file.Entries {
|
||||
paths = append(paths, e.Path)
|
||||
}
|
||||
Expect(paths).To(ConsistOf(m1, m2))
|
||||
})
|
||||
|
||||
It("produces a valid file under concurrent saves from two instances", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
var models [2][3]string
|
||||
for i := range models {
|
||||
for j := range models[i] {
|
||||
models[i][j] = filepath.Join(dir, fmt.Sprintf("m-%d-%d.bin", i, j))
|
||||
writeFakeModel(models[i][j], 4+i*3+j)
|
||||
}
|
||||
}
|
||||
// Each instance counts its own parses: the counter is not synchronized, and
|
||||
// the two goroutines below run in parallel.
|
||||
var parsesA, parsesB int32
|
||||
insts := [2]*headerCache[fakeSnap]{
|
||||
newFakeCache(dir, fakeSpecWith(&parsesA, "v1", 1)),
|
||||
newFakeCache(dir, fakeSpecWith(&parsesB, "v1", 1)),
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 16)
|
||||
for i := range insts {
|
||||
wg.Add(1)
|
||||
go func(c *headerCache[fakeSnap], files [3]string) {
|
||||
defer wg.Done()
|
||||
for _, f := range files {
|
||||
if _, err := c.load(f); err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
if err := c.flush(); err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}(insts[i], models[i])
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
}
|
||||
|
||||
// Two racing writers may drop each other's last entries. A final pair of
|
||||
// sequential saves merges everything back, because each save folds in what
|
||||
// is on disk. The bound (8) is larger than the six entries, so none is evicted.
|
||||
Expect(insts[0].flush()).To(Succeed())
|
||||
Expect(insts[1].flush()).To(Succeed())
|
||||
|
||||
file := readFakeFile(dir)
|
||||
Expect(file.Entries).To(HaveLen(6))
|
||||
Expect(tmpFilesIn(dir)).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("ignores a corrupt or wrong-schema file and overwrites it with a valid save", func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
model := filepath.Join(dir, "model.bin")
|
||||
writeFakeModel(model, 8)
|
||||
path := fakeCacheFile(dir)
|
||||
|
||||
for _, content := range []string{
|
||||
"{not json",
|
||||
`{"schema":99,"estimator":"v1","entries":[]}`,
|
||||
} {
|
||||
ExpectWithOffset(1, os.WriteFile(path, []byte(content), 0o600)).To(Succeed())
|
||||
c := newFakeCache(dir, fakeSpecWith(&parses, "v1", 1))
|
||||
Expect(c.len()).To(Equal(0))
|
||||
_, err := c.load(model)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(c.flush()).To(Succeed())
|
||||
c.stopPersist()
|
||||
|
||||
file := readFakeFile(dir)
|
||||
Expect(file.Schema).To(Equal(1))
|
||||
Expect(file.Entries).To(HaveLen(1), "content %q must be replaced by a valid save", content)
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
// readFakeFile decodes the shared cache file in dir, failing the spec if it is not valid.
|
||||
func readFakeFile(dir string) headerFile {
|
||||
raw, err := os.ReadFile(fakeCacheFile(dir))
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
var file headerFile
|
||||
ExpectWithOffset(1, json.Unmarshal(raw, &file)).To(Succeed())
|
||||
return file
|
||||
}
|
||||
|
||||
// tmpFilesIn lists temp files left in dir.
|
||||
func tmpFilesIn(dir string) []string {
|
||||
entries, err := os.ReadDir(dir)
|
||||
ExpectWithOffset(1, err).NotTo(HaveOccurred())
|
||||
var out []string
|
||||
for _, e := range entries {
|
||||
if strings.Contains(e.Name(), ".tmp") {
|
||||
out = append(out, e.Name())
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -64,8 +64,10 @@ func llamaCppDefaultsWithParser(cfg *ModelConfig, modelPath string, parse func(s
|
||||
// -> 8 for a 256k-token vocab), the dominant cost on slow filesystems.
|
||||
//
|
||||
// The mapping is released when ParseGGUFFile returns.
|
||||
f, err := parse(guessPath, gguf.UseMMap(), gguf.SkipLargeMetadata())
|
||||
// The header snapshot comes from the cache when the file identity matches;
|
||||
// otherwise the parser runs once and its result is cached.
|
||||
header, err := loadGGUFHeader(guessPath, parse)
|
||||
if err == nil {
|
||||
guessGGUFFromFile(cfg, f, 0)
|
||||
guessGGUFFromFile(cfg, header, 0)
|
||||
}
|
||||
}
|
||||
@@ -33,3 +33,21 @@ func BuildUsecaseFilterFn(usecases ModelConfigUsecase) ModelConfigFilterFn {
|
||||
return config.HasUsecases(usecases)
|
||||
}
|
||||
}
|
||||
|
||||
// FilterMayMatchLooseFiles reports whether filter can accept a model file that
|
||||
// has no configuration. A nil filter means "no filtering" and accepts them.
|
||||
// Usecase filters need a configuration, so they never match loose files and a
|
||||
// directory scan cannot add results for them. Arbitrary filters cannot be
|
||||
// inspected, so the answer comes from probing with typical file names; any
|
||||
// accepted probe keeps the directory scan.
|
||||
func FilterMayMatchLooseFiles(filter ModelConfigFilterFn) bool {
|
||||
if filter == nil {
|
||||
return true
|
||||
}
|
||||
for _, name := range []string{"", "model", "model.gguf", "model.bin", "model.safetensors", "a/b"} {
|
||||
if filter(name, nil) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package config_test
|
||||
|
||||
import (
|
||||
. "github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("FilterMayMatchLooseFiles", func() {
|
||||
It("is false for every usecase filter, which rejects files without a config", func() {
|
||||
for _, flag := range []ModelConfigUsecase{
|
||||
FLAG_CHAT, FLAG_COMPLETION, FLAG_EDIT, FLAG_EMBEDDINGS, FLAG_RERANK, FLAG_IMAGE,
|
||||
FLAG_TRANSCRIPT, FLAG_TTS, FLAG_SOUND_GENERATION, FLAG_TOKENIZE, FLAG_VAD, FLAG_VIDEO,
|
||||
FLAG_DETECTION, FLAG_VISION, FLAG_FACE_RECOGNITION, FLAG_SPEAKER_RECOGNITION,
|
||||
FLAG_AUDIO_TRANSFORM, FLAG_DIARIZATION, FLAG_DEPTH, FLAG_3D, FLAG_UPSCALE, FLAG_LLM,
|
||||
} {
|
||||
Expect(FilterMayMatchLooseFiles(BuildUsecaseFilterFn(flag))).To(BeFalse(), "flag %d", flag)
|
||||
}
|
||||
})
|
||||
It("is true for nil, the match-all filter and name filters", func() {
|
||||
Expect(FilterMayMatchLooseFiles(nil)).To(BeTrue())
|
||||
Expect(FilterMayMatchLooseFiles(NoFilterFn)).To(BeTrue())
|
||||
Expect(FilterMayMatchLooseFiles(BuildUsecaseFilterFn(FLAG_ANY))).To(BeTrue())
|
||||
name, err := BuildNameFilterFn("model")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(FilterMayMatchLooseFiles(name)).To(BeTrue())
|
||||
})
|
||||
})
|
||||
@@ -226,10 +226,10 @@ var _ = Describe("Diagnostics extraction", func() {
|
||||
Expect(ends(diagnostics.PhaseBearerLookup)).To(HaveLen(tc.lookup))
|
||||
Expect(ends(diagnostics.PhaseBodyLookup)).To(HaveLen(1))
|
||||
Expect(ends(diagnostics.PhaseConfigLoadDefaults)).To(HaveLen(1))
|
||||
// Exact body lookup does not filter configs or enumerate files.
|
||||
// Only the unchanged default and bearer discovery paths do.
|
||||
Expect(ends(diagnostics.PhaseConfigFilter)).To(HaveLen(tc.listing + tc.lookup))
|
||||
Expect(ends(diagnostics.PhaseFSEnumeration)).To(HaveLen(tc.listing + tc.lookup))
|
||||
// Exact body and bearer lookups do not filter configs or enumerate
|
||||
// files. Only the default listing (nil filter here) does.
|
||||
Expect(ends(diagnostics.PhaseConfigFilter)).To(HaveLen(tc.listing))
|
||||
Expect(ends(diagnostics.PhaseFSEnumeration)).To(HaveLen(tc.listing))
|
||||
id := events[0].ID
|
||||
_, err := uuid.Parse(id)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
@@ -85,9 +85,11 @@ func (re *RequestExtractor) setModelNameFromRequest(c echo.Context) {
|
||||
bearer := strings.TrimPrefix(auth, "Bearer ")
|
||||
if bearer != "" && bearer != auth {
|
||||
end := diagnostics.Begin(ctx, diagnostics.PhaseBearerLookup)
|
||||
exists, err := galleryop.CheckIfModelExistsContext(ctx, re.modelConfigLoader, re.modelLoader, bearer, galleryop.ALWAYS_INCLUDE)
|
||||
end(diagnosticOutcome(ctx, err != nil), 0)
|
||||
if err == nil && exists {
|
||||
// Exact-name check: a configured name or one file in the models
|
||||
// directory. It never enumerates the directory.
|
||||
exists := galleryop.CheckExactModelExists(re.modelConfigLoader, re.modelLoader, bearer)
|
||||
end(diagnosticOutcome(ctx, false), 0)
|
||||
if exists {
|
||||
model = bearer
|
||||
}
|
||||
}
|
||||
@@ -127,7 +129,13 @@ func (re *RequestExtractor) BuildFilteredFirstAvailableDefaultModel(filterFn con
|
||||
|
||||
ctx := c.Request().Context()
|
||||
end := diagnostics.Begin(ctx, diagnostics.PhaseDefaultListing)
|
||||
modelNames, err := galleryop.ListModelsContext(ctx, re.modelConfigLoader, re.modelLoader, filterFn, galleryop.SKIP_IF_CONFIGURED)
|
||||
// A filter that rejects unconfigured files (every usecase filter)
|
||||
// gains nothing from reading the models directory, so skip it.
|
||||
policy := galleryop.SKIP_IF_CONFIGURED
|
||||
if !config.FilterMayMatchLooseFiles(filterFn) {
|
||||
policy = galleryop.SKIP_ALWAYS
|
||||
}
|
||||
modelNames, err := galleryop.ListModelsContext(ctx, re.modelConfigLoader, re.modelLoader, filterFn, policy)
|
||||
end(diagnosticOutcome(ctx, err != nil), len(modelNames))
|
||||
if err != nil {
|
||||
xlog.Error("non-fatal error calling ListModels during SetDefaultModelNameToFirstAvailable()", "error", err)
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package middleware_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/pkg/diagnostics"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Request paths that must not list the models directory", func() {
|
||||
var (
|
||||
re *RequestExtractor
|
||||
bcl *config.ModelConfigLoader
|
||||
ml *model.ModelLoader
|
||||
events []diagnostics.Event
|
||||
)
|
||||
enumerations := func() int {
|
||||
n := 0
|
||||
for _, e := range events {
|
||||
if e.Phase == diagnostics.PhaseFSEnumeration && e.State == diagnostics.StateEnd {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
// resolve runs mw and returns the model name left in the request context.
|
||||
resolve := func(mw echo.MiddlewareFunc, auth string) string {
|
||||
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
if auth != "" {
|
||||
req.Header.Set("Authorization", auth)
|
||||
}
|
||||
c := echo.New().NewContext(req, httptest.NewRecorder())
|
||||
events = nil
|
||||
got := ""
|
||||
Expect(mw(func(c echo.Context) error {
|
||||
got, _ = c.Get(CONTEXT_LOCALS_KEY_MODEL_NAME).(string)
|
||||
return nil
|
||||
})(c)).To(Succeed())
|
||||
return got
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
ac := config.NewApplicationConfig()
|
||||
ac.Context = context.Background()
|
||||
ac.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) })
|
||||
yes := true
|
||||
emb := config.ModelConfig{Name: "emb", Backend: "test", Embeddings: &yes}
|
||||
emb.Model = "emb.gguf"
|
||||
other := config.ModelConfig{Name: "other", Backend: "test"}
|
||||
other.Model = "other.gguf"
|
||||
bcl = config.NewModelConfigLoader(dir)
|
||||
bcl.ReplaceModelConfigs([]config.ModelConfig{other, emb})
|
||||
for i := 0; i < 64; i++ {
|
||||
Expect(os.WriteFile(filepath.Join(dir, fmt.Sprintf("unrelated-%03d", i)), nil, 0600)).To(Succeed())
|
||||
}
|
||||
Expect(os.WriteFile(filepath.Join(dir, "raw-weights.gguf"), nil, 0600)).To(Succeed())
|
||||
ml = model.NewModelLoader(ac.SystemState)
|
||||
re = NewRequestExtractor(bcl, ml, ac)
|
||||
})
|
||||
|
||||
Describe("bearer token naming a model", func() {
|
||||
fallback := func() echo.MiddlewareFunc { return re.BuildConstantDefaultModelNameMiddleware("fallback") }
|
||||
|
||||
It("resolves an exact configured name without a listing", func() {
|
||||
Expect(resolve(fallback(), "Bearer emb")).To(Equal("emb"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
It("resolves an exact file in the models directory without a listing", func() {
|
||||
Expect(resolve(fallback(), "Bearer raw-weights.gguf")).To(Equal("raw-weights.gguf"))
|
||||
Expect(resolve(fallback(), "Bearer unrelated-001")).To(Equal("unrelated-001"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
It("does not resolve an unknown value and does not list", func() {
|
||||
Expect(resolve(fallback(), "Bearer sk-not-a-model")).To(Equal("fallback"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
It("treats the value as a literal, not a regular expression", func() {
|
||||
Expect(resolve(fallback(), "Bearer em.*")).To(Equal("fallback"))
|
||||
Expect(resolve(fallback(), "Bearer unrelated-0")).To(Equal("fallback"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
It("never probes paths outside the models directory", func() {
|
||||
Expect(resolve(fallback(), "Bearer ../"+filepath.Base(ml.ModelPath))).To(Equal("fallback"))
|
||||
Expect(resolve(fallback(), "Bearer /etc/hostname")).To(Equal("fallback"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("first available default model", func() {
|
||||
It("reads no directory for any usecase filter and returns what the old listing returned", func() {
|
||||
for _, flag := range []config.ModelConfigUsecase{
|
||||
config.FLAG_CHAT, config.FLAG_COMPLETION, config.FLAG_EMBEDDINGS, config.FLAG_TTS,
|
||||
config.FLAG_TRANSCRIPT, config.FLAG_IMAGE, config.FLAG_RERANK, config.FLAG_VAD,
|
||||
} {
|
||||
filter := config.BuildUsecaseFilterFn(flag)
|
||||
old, err := galleryop.ListModels(bcl, ml, filter, galleryop.SKIP_IF_CONFIGURED)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
skipped, err := galleryop.ListModels(bcl, ml, filter, galleryop.SKIP_ALWAYS)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(skipped).To(Equal(old), "flag %d", flag)
|
||||
|
||||
want := ""
|
||||
if len(old) > 0 {
|
||||
want = old[0]
|
||||
}
|
||||
Expect(resolve(re.BuildFilteredFirstAvailableDefaultModel(filter), "")).To(Equal(want), "flag %d", flag)
|
||||
Expect(enumerations()).To(BeZero(), "flag %d", flag)
|
||||
}
|
||||
})
|
||||
It("keeps the directory scan for a filter that accepts loose files", func() {
|
||||
got := resolve(re.BuildFilteredFirstAvailableDefaultModel(config.NoFilterFn), "")
|
||||
Expect(got).NotTo(BeEmpty())
|
||||
Expect(enumerations()).To(Equal(1))
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,98 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/pkg/diagnostics"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("Embeddings model resolution", func() {
|
||||
var (
|
||||
app *echo.Echo
|
||||
events []diagnostics.Event
|
||||
served string
|
||||
)
|
||||
|
||||
post := func(body string, headers ...string) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/embeddings", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
for i := 0; i+1 < len(headers); i += 2 {
|
||||
req.Header.Set(headers[i], headers[i+1])
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
events, served = nil, ""
|
||||
app.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
enumerations := func() int {
|
||||
n := 0
|
||||
for _, e := range events {
|
||||
if e.Phase == diagnostics.PhaseFSEnumeration && e.State == diagnostics.StateEnd {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
dir := GinkgoT().TempDir()
|
||||
ac := config.NewApplicationConfig()
|
||||
ac.Context = context.Background()
|
||||
ac.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) })
|
||||
yes := true
|
||||
emb := config.ModelConfig{Name: "emb", Backend: "test", Embeddings: &yes}
|
||||
emb.Model = "emb.gguf"
|
||||
chat := config.ModelConfig{Name: "chat", Backend: "test"}
|
||||
chat.Model = "chat.gguf"
|
||||
bcl := config.NewModelConfigLoader(dir)
|
||||
bcl.ReplaceModelConfigs([]config.ModelConfig{chat, emb})
|
||||
for i := 0; i < 64; i++ {
|
||||
Expect(os.WriteFile(filepath.Join(dir, fmt.Sprintf("unrelated-%03d", i)), nil, 0600)).To(Succeed())
|
||||
}
|
||||
re := middleware.NewRequestExtractor(bcl, model.NewModelLoader(ac.SystemState), ac)
|
||||
app = echo.New()
|
||||
app.POST("/v1/embeddings", func(c echo.Context) error {
|
||||
served = c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig).Name
|
||||
return c.NoContent(http.StatusNoContent)
|
||||
}, embeddingsModelResolution(re))
|
||||
})
|
||||
|
||||
It("does not list the models directory when the body names the model", func() {
|
||||
Expect(post(`{"model":"emb","input":"x"}`).Code).To(Equal(http.StatusNoContent))
|
||||
Expect(served).To(Equal("emb"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
|
||||
It("does not list the directory when the body names the model and a key is sent", func() {
|
||||
Expect(post(`{"model":"emb","input":"x"}`, "Authorization", "Bearer some-api-key").Code).To(Equal(http.StatusNoContent))
|
||||
Expect(served).To(Equal("emb"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
|
||||
It("picks the configured embeddings default when the body names no model", func() {
|
||||
Expect(post(`{"input":"x"}`).Code).To(Equal(http.StatusNoContent))
|
||||
Expect(served).To(Equal("emb"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
|
||||
It("still honours a model name carried by the bearer token", func() {
|
||||
Expect(post(`{"input":"x"}`, "Authorization", "Bearer chat").Code).To(Equal(http.StatusNoContent))
|
||||
Expect(served).To(Equal("chat"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,127 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/pkg/diagnostics"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// Every usecase filter used by a route's default-model lookup. Each route
|
||||
// passes one of these to BuildFilteredFirstAvailableDefaultModel, so a
|
||||
// listing-free result for each filter covers all of them.
|
||||
var routeUsecaseFilters = []struct {
|
||||
name string
|
||||
flag config.ModelConfigUsecase
|
||||
}{
|
||||
{"chat", config.FLAG_CHAT},
|
||||
{"completion", config.FLAG_COMPLETION},
|
||||
{"edit", config.FLAG_EDIT},
|
||||
{"embeddings", config.FLAG_EMBEDDINGS},
|
||||
{"rerank", config.FLAG_RERANK},
|
||||
{"transcript", config.FLAG_TRANSCRIPT},
|
||||
{"diarization", config.FLAG_DIARIZATION},
|
||||
{"sound classification", config.FLAG_SOUND_CLASSIFICATION},
|
||||
{"sound generation", config.FLAG_SOUND_GENERATION},
|
||||
{"tts", config.FLAG_TTS},
|
||||
{"image", config.FLAG_IMAGE},
|
||||
{"upscale", config.FLAG_UPSCALE},
|
||||
{"detection", config.FLAG_DETECTION},
|
||||
{"depth", config.FLAG_DEPTH},
|
||||
{"face recognition", config.FLAG_FACE_RECOGNITION},
|
||||
{"speaker recognition", config.FLAG_SPEAKER_RECOGNITION},
|
||||
{"audio transform", config.FLAG_AUDIO_TRANSFORM},
|
||||
{"vad", config.FLAG_VAD},
|
||||
{"video", config.FLAG_VIDEO},
|
||||
{"3d animation", config.FLAG_3D_ANIMATION},
|
||||
{"3d", config.FLAG_3D},
|
||||
{"tokenize", config.FLAG_TOKENIZE},
|
||||
}
|
||||
|
||||
var _ = Describe("Default-model lookup on every usecase route", func() {
|
||||
var (
|
||||
events []diagnostics.Event
|
||||
served string
|
||||
)
|
||||
|
||||
// serve builds a one-route app with the same pre-binding lookup the routes
|
||||
// use, posts the body with optional headers, and returns the response code.
|
||||
serve := func(flag config.ModelConfigUsecase, body string, headers ...string) int {
|
||||
dir := GinkgoT().TempDir()
|
||||
ac := config.NewApplicationConfig()
|
||||
ac.Context = context.Background()
|
||||
ac.SystemState = &system.SystemState{Model: system.Model{ModelsPath: dir}}
|
||||
ac.DiagnosticsRecorder = diagnostics.NewRecorder(func(e diagnostics.Event) { events = append(events, e) })
|
||||
chat := config.ModelConfig{Name: "chat", Backend: "test"}
|
||||
chat.Model = "chat.gguf"
|
||||
bcl := config.NewModelConfigLoader(dir)
|
||||
bcl.ReplaceModelConfigs([]config.ModelConfig{chat})
|
||||
for i := 0; i < 64; i++ {
|
||||
Expect(os.WriteFile(filepath.Join(dir, fmt.Sprintf("unrelated-%03d", i)), nil, 0600)).To(Succeed())
|
||||
}
|
||||
re := middleware.NewRequestExtractor(bcl, model.NewModelLoader(ac.SystemState), ac)
|
||||
app := echo.New()
|
||||
app.POST("/route", func(c echo.Context) error {
|
||||
served, _ = c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_NAME).(string)
|
||||
return c.NoContent(http.StatusNoContent)
|
||||
}, re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(flag)))
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/route", strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
for i := 0; i+1 < len(headers); i += 2 {
|
||||
req.Header.Set(headers[i], headers[i+1])
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
events, served = nil, ""
|
||||
app.ServeHTTP(rec, req)
|
||||
return rec.Code
|
||||
}
|
||||
enumerations := func() int {
|
||||
n := 0
|
||||
for _, e := range events {
|
||||
if e.Phase == diagnostics.PhaseFSEnumeration && e.State == diagnostics.StateEnd {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
for _, f := range routeUsecaseFilters {
|
||||
flag := f.flag
|
||||
|
||||
It(fmt.Sprintf("%s: does not list the models directory when the body names a model", f.name), func() {
|
||||
serve(flag, `{"model":"chat","input":"x"}`)
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
|
||||
It(fmt.Sprintf("%s: does not list the models directory when the body names no model", f.name), func() {
|
||||
serve(flag, `{"input":"x"}`)
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
|
||||
It(fmt.Sprintf("%s: does not list the models directory when a bearer key is sent", f.name), func() {
|
||||
serve(flag, `{"input":"x"}`, "Authorization", "Bearer some-api-key")
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
}
|
||||
|
||||
It("still resolves a model named by the bearer token on the chat lookup", func() {
|
||||
code := serve(config.FLAG_CHAT, `{"input":"x"}`, "Authorization", "Bearer chat")
|
||||
Expect(code).To(Equal(http.StatusNoContent))
|
||||
Expect(served).To(Equal("chat"))
|
||||
Expect(enumerations()).To(BeZero())
|
||||
})
|
||||
})
|
||||
@@ -159,9 +159,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
|
||||
nodeHeaderMiddleware,
|
||||
usageMiddleware,
|
||||
traceMiddleware,
|
||||
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_EMBEDDINGS)),
|
||||
re.BuildConstantDefaultModelNameMiddleware("gpt-4o"),
|
||||
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
|
||||
embeddingsModelResolution(re),
|
||||
func(next echo.HandlerFunc) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
if err := re.SetOpenAIRequest(c); err != nil {
|
||||
@@ -305,3 +303,9 @@ func RegisterOpenAIRoutes(app *echo.Echo,
|
||||
app.GET("/v1/models/capabilities", capabilitiesHandler)
|
||||
app.GET("/models/capabilities", capabilitiesHandler)
|
||||
}
|
||||
|
||||
// embeddingsModelResolution binds the request and resolves its model. The
|
||||
// embeddings default is looked up only when the body names no model.
|
||||
func embeddingsModelResolution(re *middleware.RequestExtractor) echo.MiddlewareFunc {
|
||||
return re.SetModelAndConfigWithDefault(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }, config.BuildUsecaseFilterFn(config.FLAG_EMBEDDINGS))
|
||||
}
|
||||
@@ -2,6 +2,8 @@ package galleryop
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/pkg/diagnostics"
|
||||
"github.com/mudler/LocalAI/pkg/model"
|
||||
@@ -132,3 +134,21 @@ func checkRequestModelExistsContext(ctx context.Context, bcl *config.ModelConfig
|
||||
}
|
||||
return checkIfModelExistsContext(ctx, bcl, ml, modelName, ALWAYS_INCLUDE)
|
||||
}
|
||||
|
||||
// CheckExactModelExists reports whether modelName is a configured model name or
|
||||
// a single entry in the models directory. It never lists the directory. Unlike
|
||||
// CheckIfModelExists, modelName is a literal, not a regular expression, and a
|
||||
// name that is not a local relative path never matches a file.
|
||||
func CheckExactModelExists(bcl *config.ModelConfigLoader, ml *model.ModelLoader, modelName string) bool {
|
||||
return checkExactModelExists(bcl, ml, modelName)
|
||||
}
|
||||
|
||||
func checkExactModelExists(bcl *config.ModelConfigLoader, ml discoveryFiles, modelName string) bool {
|
||||
if _, exists := bcl.GetModelConfig(modelName); exists {
|
||||
return true
|
||||
}
|
||||
if !filepath.IsLocal(modelName) {
|
||||
return false
|
||||
}
|
||||
return ml.ExistsInModelPath(modelName)
|
||||
}
|
||||
@@ -161,6 +161,49 @@ The initially migrated backend families are `transformers` and its aliases,
|
||||
artifact declarations only for this set. Compatible external backends may opt
|
||||
in by declaring the artifact explicitly.
|
||||
|
||||
## GGUF header cache
|
||||
|
||||
When a model config does not set a value itself, LocalAI reads the header of its GGUF file to choose defaults. The header gives the layer count, the trained context length, the chat template, the token IDs, and whether the file has an embedded multi-token-prediction head. Reading a large header takes seconds, so LocalAI keeps the result in a cache.
|
||||
|
||||
**Mode: shared by file.** Every frontend that uses the same data path reads and writes one cache file, `header-cache.json`. A frontend can therefore use the headers another frontend has already read. This mode is valid only while all frontends mount the same data volume.
|
||||
|
||||
Set an instance ID (`--instance-id` or `LOCALAI_INSTANCE_ID`) to give a frontend its own file, `header-cache-<id>.json`, instead of the shared one. It then keeps its own entries across restarts and does not share them with other frontends. An instance ID that is not a plain token (letters, digits, `-` and `_`, up to 64 characters) is hashed into the file name. In distributed mode, an ID that LocalAI generates at startup does not select a file of its own, because it changes on every start; such a frontend uses the shared file.
|
||||
|
||||
{{% notice warning %}}
|
||||
All frontends that share this cache must mount the same data volume. With a ReadWriteOnce volume, the volume attaches to one node at a time, so all frontends that share the cache must run on that one node. A frontend on another node cannot mount the volume and will not start. Sharing this cache across nodes needs a ReadWriteMany volume or a shared store, and this mode does not support it.
|
||||
{{% /notice %}}
|
||||
|
||||
### What is stored
|
||||
|
||||
Each entry is keyed by the absolute path of the file, its size and its modification time. If any of the three changes, the entry is not used and the header is read again.
|
||||
|
||||
The file keeps the values that come from the file: the architecture and name, the trained context length, the chat template, the token IDs, the multi-token-prediction head, and the estimates that depend only on the file and on LocalAI's fixed defaults. The file also records the gguf-parser version and those defaults, so a LocalAI upgrade that changes them ignores the old file.
|
||||
|
||||
The file does not store anything about the machine: GPU and memory sizes, and your settings. Those are read again when a model loads, so the same file gives the same defaults on a different host.
|
||||
|
||||
The cache holds at most 256 entries. The least recently used entry is dropped first. Changes are written in the background after a short pause, never while a request waits.
|
||||
|
||||
### How frontends share the file
|
||||
|
||||
- **Merge on save.** Before writing, a frontend reads the file and merges its entries into its own. When both hold an entry for the same model file, the one that matches the model file on disk is kept. Entries that match no current file are dropped. The 256-entry bound applies to the merged set.
|
||||
- **Atomic write.** The new file is written to a temporary file in the same directory and renamed into place. A reader sees either the old file or the new one, never a partial file.
|
||||
- **Reload on miss.** A frontend reads the file when it starts. Later, it checks the file only when a header lookup misses, and only when the file's modification time differs from the one it last read. It reads the file at most once per miss. A lookup that hits the cache never reads the file.
|
||||
- **Concurrent saves.** Two frontends can save at the same moment. The later rename wins, so an entry that the other frontend saved in that moment can be lost. A lost entry costs one header parse the next time that model file is used. It never gives a wrong value, because every entry is checked against its model file before it is used.
|
||||
|
||||
### Where the file lives
|
||||
|
||||
The file is in the data path (`--data-path`, default `data/` under the LocalAI base path):
|
||||
|
||||
- `header-cache.json`, in standalone mode and for every frontend that does not set an instance ID. Frontends that mount the same data path share it.
|
||||
- `header-cache-<id>.json`, for a frontend that sets an instance ID (`--instance-id` or `LOCALAI_INSTANCE_ID`).
|
||||
- No data path: the cache is kept in memory only. Each start parses the headers again.
|
||||
|
||||
Files named `gguf-header-cache*.json` come from earlier versions. LocalAI does not read them, so you can delete them.
|
||||
|
||||
### Clearing the cache
|
||||
|
||||
Stop the frontends that share the file and delete `header-cache.json`. It is rebuilt as models are used, one parse per model file. Deleting the file while a frontend runs is not reliable: that frontend writes its in-memory entries back after the next change.
|
||||
|
||||
## Parameters Section
|
||||
|
||||
The `parameters` section contains all OpenAI-compatible request parameters and model-specific options.
|
||||
|
||||
@@ -672,6 +672,14 @@ Two limits apply:
|
||||
- Models loaded with `--config-file` exist only on the frontend that loaded them. A reload of the models directory keeps them, and a config-file model wins over a directory file with the same name, as it does at startup.
|
||||
- The reload is strict: if any config file in the directory does not parse, the frontend keeps its current configs and logs the error once. It retries on each pass until the file is fixed. This keeps a half-written file from looking like a deleted model.
|
||||
|
||||
### GGUF header cache shared by file
|
||||
|
||||
Frontends that mount the same data volume share one [GGUF header cache]({{% relref "advanced/model-configuration" %}}#gguf-header-cache) file, `header-cache.json`, in the data path. A header that one frontend has read is used by the others. This is the shared-by-file mode, and it requires that every frontend mounts the same data volume.
|
||||
|
||||
A frontend that sets `--instance-id` keeps its own file, `header-cache-<id>.json`, and does not share it. Use this when a frontend should start warm after a restart without sharing its entries.
|
||||
|
||||
With a ReadWriteOnce volume, all frontends that share the cache must run on one node. Sharing across nodes needs a ReadWriteMany volume or a shared store, and this mode does not support it.
|
||||
|
||||
### Model configuration revisions
|
||||
|
||||
Distributed mode assigns a `config_revision` to each validated model configuration. It hashes the persisted semantic configuration, including fields such as `context_size` and parallel settings. YAML formatting, comments, and map order do not change it.
|
||||
|
||||
@@ -17,6 +17,10 @@ For the API documentation you can refer to the OpenAI docs: https://platform.ope
|
||||
|
||||
The embedding endpoint is compatible with `llama.cpp` models, `bert.cpp` models and sentence-transformers models available in huggingface.
|
||||
|
||||
## Choosing the model
|
||||
|
||||
Send the model name in the request body. If the body has no `model`, LocalAI uses the configured embeddings model whose name comes first alphabetically. If no embeddings model is configured, the request returns `400 Bad Request`. A model name is never taken from the bearer token unless it matches a configured model name exactly.
|
||||
|
||||
## Using Gallery Models
|
||||
|
||||
LocalAI provides a model gallery with pre-configured embedding models. To use a gallery model:
|
||||
|
||||
@@ -192,6 +192,7 @@ See also: [Text Generation]({{%relref "features/text-generation" %}})
|
||||
|--------|--------------------------------|---------------------------------------|
|
||||
| 400 | Invalid or malformed request body | `Bad Request` |
|
||||
| 400 | Model not found in configuration | `Bad Request` |
|
||||
| 400 | No `model` in the body and no embeddings model configured | `Bad Request` |
|
||||
| 500 | Backend inference failure | `Internal Server Error` |
|
||||
|
||||
See also: [Embeddings]({{%relref "features/embeddings" %}})
|
||||
|
||||
Reference in new issue
Block a user