diff --git a/internal/loki/backend_gguf.go b/internal/loki/backend_gguf.go index 9e20a77..9f643ef 100644 --- a/internal/loki/backend_gguf.go +++ b/internal/loki/backend_gguf.go @@ -48,6 +48,15 @@ type GGUFInfo struct { HeadCountKVLayers []int FullAttnInterval int // .full_attention_interval, si présent Hybrid bool // au moins une clé .ssm.* : couches à état récurrent + // Repli de llama.cpp quand key_length / value_length manquent : n_embd / n_head. + EmbeddingLength int // .embedding_length + HeadCount int // .attention.head_count (valeur unique seulement) + // Taille de l'état récurrent par couche (llama_hparams::n_embd_r / n_embd_s), + // pour estimer ce que pèse un état sauvegardé d'un modèle hybride. + SSMConv int // .ssm.conv_kernel + SSMInner int // .ssm.inner_size + SSMState int // .ssm.state_size + SSMGroups int // .ssm.group_count // HasNextNTensor : le tenseur blk.{BlockCount-1}.nextn.eh_proj.weight existe // dans l'une des tranches. C'est le test de llama.cpp lui-même pour reconnaître // une tête MTP : la clé nextn_predict_layers seule ne suffit pas, certains GGUF @@ -501,6 +510,12 @@ func ggufInfoFrom(first *ggufFile, all []*ggufFile) GGUFInfo { in.KeyLen = get("attention.key_length") in.ValLen = get("attention.value_length") in.FullAttnInterval = get("full_attention_interval") + in.EmbeddingLength = get("embedding_length") + in.HeadCount = get("attention.head_count") + in.SSMConv = get("ssm.conv_kernel") + in.SSMInner = get("ssm.inner_size") + in.SSMState = get("ssm.state_size") + in.SSMGroups = get("ssm.group_count") in.HeadCountKV = get("attention.head_count_kv") if arr, ok := first.arrays[p+"attention.head_count_kv"]; ok { in.HeadCountKVLayers = make([]int, len(arr)) diff --git a/internal/loki/backend_gguf_test.go b/internal/loki/backend_gguf_test.go index fcbcbe3..1c1d80a 100644 --- a/internal/loki/backend_gguf_test.go +++ b/internal/loki/backend_gguf_test.go @@ -184,7 +184,7 @@ func TestGGUFMetaQwenHybrideMTP(t *testing.T) { want := GGUFInfo{ Version: 3, Arch: "qwen35", BlockCount: 4, NextN: 1, ContextLength: 262144, KeyLen: 256, ValLen: 256, HeadCountKV: 4, HeadCountKVLayers: []int{0, 0, 0, 4}, - FullAttnInterval: 4, Hybrid: true, HasNextNTensor: true, TensorCount: 3, + FullAttnInterval: 4, Hybrid: true, SSMConv: 4, HasNextNTensor: true, TensorCount: 3, } if !reflect.DeepEqual(got, want) { t.Fatalf("ggufMeta =\n%+v\nattendu\n%+v", got, want) diff --git a/internal/loki/backend_serve.go b/internal/loki/backend_serve.go index a02f329..c7b3f3b 100644 --- a/internal/loki/backend_serve.go +++ b/internal/loki/backend_serve.go @@ -342,6 +342,10 @@ func cmdServe(args []string) error { probeServeGPUs(cfg, extra, &si) probeExpertEnv(&si) probeFidelityEnv(&si) + probeCacheRAM(cfg, extra, &si) + if strings.Contains(si.Help, "--slot-save-path") && !hasAnyFlag(extra, "--slot-save-path") { + si.SlotDir = prepareSlotDir(LokiHome()) + } llmArgs, env, notes := buildServeArgs(cfg, extra, bin, si) applyServeEnv(env) @@ -388,6 +392,15 @@ type serveSysInfo struct { GPUs int // GPU visibles (CUDA_VISIBLE_DEVICES, sinon nvidia-smi) ; 0 = inconnu ou non sondé ArgEnv map[string]string // LLAMA_ARG_* de l'environnement qui décident du pipeline, de --fit et du cache (serveArgEnv, fitArgEnv, fidelityArgEnv) UserEnv map[string]string // GGML_* des réglages d'expert déjà posés (voir expertEnvKeys) : intouchés + + // Cache de prompts (voir backend_serve_cache.go). Zéro = inconnu : le moteur + // garde son défaut. + GGUF *GGUFInfo // métadonnées du modèle ; nil = illisibles + ModelBytes int64 // taille du modèle, toutes tranches + RAMMiB int64 // RAM de la machine, limite du conteneur comprise + RAMAvailMiB int64 // RAM libre au lancement + VRAMMiB int64 // VRAM NVIDIA des cartes visibles + SlotDir string // dossier de --slot-save-path, créé et vérifié ; vide = pas de drapeau } // cudaDeviceEnv : sélection GPU (loki gpu), on filtre les devices visibles par @@ -569,6 +582,11 @@ func buildServeArgs(cfg map[string]string, extra []string, bin string, si serveS if si.APIKey != "" { args = append(args, "--api-key", si.APIKey) } + // Cache de prompts en RAM et dossier des slots (voir backend_serve_cache.go) : + // avant EXTRA_ARGS, qui garde le dernier mot. + cram, cramNotes := cacheRAMArgs(cfg, extra, si) + args = append(append(args, cram...), slotSaveArgs(cfg, extra, si)...) + notes = append(notes, cramNotes...) // EXTRA_ARGS (déjà découpé comme le ferait le shell — les guillemets gardent // ensemble un chemin qui contient des espaces) ferme la marche. args = append(args, extra...) diff --git a/internal/loki/backend_serve_cache.go b/internal/loki/backend_serve_cache.go new file mode 100644 index 0000000..6fa3b82 --- /dev/null +++ b/internal/loki/backend_serve_cache.go @@ -0,0 +1,535 @@ +package loki + +import ( + "context" + "fmt" + "io/fs" + "math" + "os" + "os/exec" + "path/filepath" + "runtime" + "strconv" + "strings" + "time" +) + +// Cache de prompts en RAM (--cache-ram) et isolation des travaux annexes. +// +// llama-server garde en RAM hôte une copie EXACTE des états de slot qu'il quitte +// (cache KV, état récurrent, points de reprise, contexte du brouillon MTP) et la +// recharge octet pour octet quand la conversation revient. Ce n'est pas une +// approximation : la seule différence avec un prefill complet est le bruit de +// virgule flottante dû à un découpage des lots différent, exactement comme le +// cache_prompt par défaut. Le type du cache KV n'est pas touché, il est copié +// tel quel. +// +// Le défaut du moteur (8 Gio) est trop petit pour une conversation longue d'un +// modèle dense : l'état principal (30 à 65 k jetons) ne tient plus à côté de +// celui d'un vérificateur ou d'un sous-agent, et la conversation est alors +// recalculée en entier — 30 à 80 s sur un 27B. Loki agrandit donc ce cache, +// mais seulement quand il sait mesurer le besoin et que la RAM est libre : il ne +// descend JAMAIS sous le défaut du moteur, et se tait dès que l'utilisateur a +// tranché (CACHE_RAM, -cram dans EXTRA_ARGS, LLAMA_ARG_CACHE_RAM). +// +// Tout ce qui décide est pur ; probeCacheRAM fait les lectures (GGUF, RAM, +// nvidia-smi) une fois, avant buildServeArgs. + +// engineCacheRAMDefault : cache_ram_mib par défaut de llama.cpp (common.h). +const engineCacheRAMDefault = 8192 + +// cacheArgEnv : la variable qui règle le cache sans drapeau. llama.cpp donne la +// priorité à la ligne de commande : un --cache-ram de Loki l'écraserait. +var cacheArgEnv = []string{"LLAMA_ARG_CACHE_RAM"} + +// probeCacheRAM complète si avec ce dont cacheRAMAuto a besoin. Chaque lecture +// ratée laisse un zéro, et un zéro veut dire « on ne sait pas » : le moteur +// garde alors son défaut. +func probeCacheRAM(cfg map[string]string, extra []string, si *serveSysInfo) { + if si.ArgEnv == nil { + si.ArgEnv = map[string]string{} + } + for _, k := range cacheArgEnv { + if v, ok := os.LookupEnv(k); ok && strings.TrimSpace(v) != "" { + si.ArgEnv[k] = v + } + } + if si.Model != "" { + si.ModelBytes = shardFamilySize(filepath.Dir(si.Model), filepath.Base(si.Model)) + if g, err := ggufMeta(si.Model); err == nil { + si.GGUF = &g + } + } + used, total := ramUsageMB() + if total > 0 { + si.RAMMiB, si.RAMAvailMiB = int64(total), int64(total-used) + } + if runtime.GOOS == "linux" { + if limit, cur := cgroupMemMiB(os.DirFS("/")); limit > 0 { + if si.RAMMiB == 0 || limit < si.RAMMiB { + si.RAMMiB = limit + } + if free := limit - cur; free >= 0 && (si.RAMAvailMiB == 0 || free < si.RAMAvailMiB) { + si.RAMAvailMiB = free + } + } + } + // Mémoire unifiée (macOS), AMD, Vulkan, --device choisi à la main : pas de + // total VRAM fiable, donc pas d'agrandissement. + if runtime.GOOS == "darwin" || hasAnyFlag(extra, "-dev", "--device") || si.ArgEnv["LLAMA_ARG_DEVICE"] != "" { + return + } + // La sélection du preset d'abord (cudaDeviceEnv l'accompagne toujours de + // PCI_BUS_ID), celle de l'environnement ensuite : l'aperçu de l'interface + // tourne dans un process où le preset n'a rien posé. + cvd, pci := cfg["CUDA_VISIBLE_DEVICES"], true + if cvd == "" { + cvd, pci = os.Getenv("CUDA_VISIBLE_DEVICES"), os.Getenv("CUDA_DEVICE_ORDER") == "PCI_BUS_ID" + } + si.VRAMMiB = nvidiaVRAMMiB(cvd, pci, 3*time.Second) +} + +// cgroupMemMiB : limite mémoire du conteneur et sa consommation, en Mio (v2, +// sinon v1). 0 = pas de limite. Le MemAvailable de /proc/meminfo voit l'hôte +// entier : sans ça, un conteneur limité à 16 Go sur une machine de 64 Go se +// croirait riche. +func cgroupMemMiB(fsys fs.FS) (limit, current int64) { + read := func(p string) int64 { + b, err := fs.ReadFile(fsys, p) + if err != nil { + return 0 + } + n, err := strconv.ParseInt(strings.TrimSpace(string(b)), 10, 64) + if err != nil || n <= 0 || n >= 1<<60 { // « max » ou la valeur géante du v1 sans limite + return 0 + } + return n >> 20 + } + if limit = read("sys/fs/cgroup/memory.max"); limit > 0 { + return limit, read("sys/fs/cgroup/memory.current") + } + if limit = read("sys/fs/cgroup/memory/memory.limit_in_bytes"); limit > 0 { + return limit, read("sys/fs/cgroup/memory/memory.usage_in_bytes") + } + return 0, 0 +} + +// nvidiaVRAMMiB : VRAM totale des cartes que llama-server verra. Borné dans le +// temps comme nvidiaGPUCount ; au moindre doute, 0. +func nvidiaVRAMMiB(cvd string, pciOrder bool, timeout time.Duration) int64 { + if !hasTool("nvidia-smi") { + return 0 + } + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + out, err := hideCmd(exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.total", + "--format=csv,noheader,nounits")).Output() + if err != nil { + return 0 + } + return vramFromSMI(string(out), cvd, pciOrder) +} + +// vramFromSMI additionne la VRAM des cartes sélectionnées par +// CUDA_VISIBLE_DEVICES. Les index de nvidia-smi sont dans l'ordre du bus PCI ; +// ceux de CUDA aussi, mais seulement avec CUDA_DEVICE_ORDER=PCI_BUS_ID. Sans +// lui, on ne sait pas quelle carte est laquelle : on prend la plus petite autant +// de fois qu'il y a de cartes — un minorant, qui ne fait que rendre Loki plus +// prudent. Un UUID ou une tranche MIG : on ne sait pas, 0. +func vramFromSMI(out, cvd string, pciOrder bool) int64 { + byIdx := map[int]int64{} + var minMiB int64 + for _, line := range strings.Split(strings.TrimSpace(out), "\n") { + f := strings.Split(line, ",") + if len(f) < 2 { + continue + } + i, err1 := strconv.Atoi(strings.TrimSpace(f[0])) + m, err2 := strconv.ParseInt(strings.TrimSpace(f[1]), 10, 64) + if err1 != nil || err2 != nil || m <= 0 { + return 0 + } + byIdx[i] = m + if minMiB == 0 || m < minMiB { + minMiB = m + } + } + if len(byIdx) == 0 { + return 0 + } + cvd = strings.TrimSpace(cvd) + if cvd == "" { + var sum int64 + for _, m := range byIdx { + sum += m + } + return sum + } + var sum int64 + n := 0 + for _, d := range strings.Split(cvd, ",") { + if d = strings.TrimSpace(d); d == "" { + continue + } + i, err := strconv.Atoi(d) + if err != nil { + return 0 + } + m, ok := byIdx[i] + if !ok { + return 0 + } + sum += m + n++ + } + if !pciOrder { + return minMiB * int64(n) + } + return sum +} + +// kvTypeBytes : octets par élément d'un type de cache (blocs de 32 pour les +// types quantifiés de ggml). Inconnu : compté comme f16, l'estimation reste +// bornée par le plafond de RAM. +func kvTypeBytes(t string) float64 { + switch strings.ToLower(strings.TrimSpace(t)) { + case "f32": + return 4 + case "q8_0": + return 34.0 / 32 + case "q5_1": + return 24.0 / 32 + case "q5_0": + return 22.0 / 32 + case "q4_1": + return 20.0 / 32 + case "q4_0", "iq4_nl": + return 18.0 / 32 + } + return 2 +} + +// ggufStateBytes estime ce que pèse l'état d'une séquence : octets de cache KV +// par jeton (couches d'attention), et octets d'état récurrent (couches SSM d'un +// hybride, en f32 comme dans llama.cpp). 0 par jeton = on ne sait pas. +func ggufStateBytes(g GGUFInfo, kt, vt string) (perToken float64, recurrent int64) { + kl, vl := g.KeyLen, g.ValLen + if g.EmbeddingLength > 0 && g.HeadCount > 0 { + if kl == 0 { + kl = g.EmbeddingLength / g.HeadCount + } + if vl == 0 { + vl = g.EmbeddingLength / g.HeadCount + } + } + if kl == 0 || vl == 0 || g.BlockCount == 0 { + return 0, 0 + } + heads, recLayers := 0, 0 + switch { + case len(g.HeadCountKVLayers) > 0: + for _, h := range g.HeadCountKVLayers { + heads += h + if h == 0 && g.Hybrid { + recLayers++ + } + } + case g.Hybrid && g.FullAttnInterval > 0: + attn := g.BlockCount / g.FullAttnInterval + heads, recLayers = attn*g.HeadCountKV, g.BlockCount-attn + default: + heads = g.BlockCount * g.HeadCountKV + if heads == 0 && g.Hybrid { + recLayers = g.BlockCount // récurrent pur + } + } + perToken = float64(heads) * (float64(kl)*kvTypeBytes(kt) + float64(vl)*kvTypeBytes(vt)) + if recLayers > 0 { + conv := int64(0) + if g.SSMConv > 1 { + conv = int64(g.SSMConv-1) * int64(g.SSMInner+2*g.SSMGroups*g.SSMState) + } + recurrent = int64(recLayers) * (conv + int64(g.SSMState)*int64(g.SSMInner)) * 4 + } + return perToken, recurrent +} + +// serveCtx : la taille de contexte que le moteur ouvrira (-c d'EXTRA_ARGS, sinon +// CTX, sinon le défaut de Loki ; 0 = celle d'entraînement du modèle). +func serveCtx(cfg map[string]string, extra []string, g *GGUFInfo) int { + s := flagValue(extra, "-c", "--ctx-size") + if s == "" { + s = strings.TrimSpace(cfg["CTX"]) + } + if s == "" { + s = "32768" + } + n, err := strconv.Atoi(s) + if err != nil || n < 0 { + return 0 + } + if n == 0 && g != nil { + n = g.ContextLength + } + return n +} + +// cpuWeights dit pourquoi des poids (ou le cache KV) vivent en RAM hôte. Le +// cache de prompts y ferait alors concurrence au cache de pages des experts +// mappés : on n'y touche pas. Vide = tout est sur GPU. +func cpuWeights(cfg map[string]string, extra []string, argEnv map[string]string, blocks int) string { + ot := flagValues(extra, "-ot", "--override-tensor") + if v := argEnv["LLAMA_ARG_OVERRIDE_TENSOR"]; v != "" { + ot = append(ot, v) + } + for _, v := range ot { + for _, el := range strings.Split(v, ",") { + if _, buft, ok := strings.Cut(el, "="); ok && strings.HasPrefix(strings.ToUpper(strings.TrimSpace(buft)), "CPU") { + return "tenseurs sur CPU (-ot)" + } + } + } + if hasAnyFlag(extra, "--cpu-moe", "-cmoe") || envTruthy(argEnv["LLAMA_ARG_CPU_MOE"]) { + return "experts sur CPU" + } + for _, n := range []struct{ env, short, long string }{ + {"LLAMA_ARG_N_CPU_MOE", "-ncmoe", "--n-cpu-moe"}, + {"LLAMA_ARG_N_CPU_FFN", "-ncffn", "--n-cpu-ffn"}, + } { + for _, v := range append(flagValues(extra, n.short, n.long), argEnv[n.env]) { + if v = strings.TrimSpace(v); v != "" && v != "0" { + return "couches " + n.long + " sur CPU" + } + } + } + if !kvOffloaded(extra, argEnv) { + return "cache KV sur CPU" + } + ngl := flagValue(extra, "-ngl", "--n-gpu-layers", "--gpu-layers") + if ngl == "" { + ngl = strings.TrimSpace(cfg["NGL"]) + if strings.EqualFold(ngl, "auto") { + ngl = argEnv["LLAMA_ARG_N_GPU_LAYERS"] + } + } + // 999 est la sentinelle « toutes » ; un nombre plus petit que le modèle (sa + // couche de sortie comprise) laisse des couches au CPU. + if n, err := strconv.Atoi(strings.TrimSpace(ngl)); err == nil && n != 999 && (blocks == 0 || n < blocks+1) { + return "couches GPU limitées à " + strconv.Itoa(n) + } + return "" +} + +// cacheRAMAuto calcule la taille du cache de prompts quand CACHE_RAM est vide +// ou « auto ». 0 = garder le défaut du moteur, why dit pourquoi. +// +// Besoin : 2,5 fois l'état d'une conversation pleine (cache KV à CTX jetons, +// plus l'état récurrent et ses points de reprise sur un hybride) — de quoi +// garder la conversation principale ET l'état d'un travail annexe. Plafond : 30 +// % de la RAM (limite du conteneur comprise) et la moitié de ce qui est libre au +// lancement. Ce plafond est un réglage, pas un filet : sous Linux, un +// dépassement finit en général par l'OOM killer, pas par le repli du moteur. +func cacheRAMAuto(cfg map[string]string, extra []string, si serveSysInfo) (mib int64, why string) { + switch { + case si.VRAMMiB <= 0: + return 0, "VRAM NVIDIA inconnue" + case si.GGUF == nil: + return 0, "métadonnées GGUF illisibles" + case si.RAMMiB <= 0 || si.RAMAvailMiB <= 0: + return 0, "RAM inconnue" + case si.ModelBytes <= 0: + return 0, "taille du modèle inconnue" + } + if r := cpuWeights(cfg, extra, si.ArgEnv, si.GGUF.BlockCount); r != "" { + return 0, r + } + kt, vt, _ := effectiveKVTypes(cfg, extra, si.ArgEnv) + perTok, rec := ggufStateBytes(*si.GGUF, kt, vt) + ctx := serveCtx(cfg, extra, si.GGUF) + if perTok <= 0 || ctx <= 0 { + return 0, "taille d'état inconnue" + } + kvBytes := perTok * float64(ctx) + if float64(si.ModelBytes)+kvBytes > 0.9*float64(si.VRAMMiB)*(1<<20) { + return 0, "modèle trop gros pour la VRAM : des poids restent en RAM" + } + state := kvBytes + if rec > 0 { + // Points de reprise : llama.cpp en garde jusqu'à 32 par séquence ; on en + // compte un par tranche de 8 k jetons, estimation grossière mais bornée. + state += float64(rec) * float64(min(32, ctx/8192+1)+1) + } + need := int64(math.Ceil(2.5 * state / (1 << 20))) + ceiling := min(si.RAMMiB*30/100, si.RAMAvailMiB/2) + if ceiling <= engineCacheRAMDefault { + return 0, fmt.Sprintf("RAM libre trop juste (plafond %d Mio)", ceiling) + } + if need <= engineCacheRAMDefault { + return 0, fmt.Sprintf("le défaut suffit (besoin estimé %d Mio)", need) + } + return min(need, ceiling), fmt.Sprintf("besoin estimé %d Mio, plafond %d Mio", need, ceiling) +} + +// cacheRAMSetting lit CACHE_RAM : auto (vide ou « auto »), ou un nombre de Mio +// passé tel quel (-1 = sans limite, 0 = cache coupé). ok=false : valeur +// illisible. +func cacheRAMSetting(cfg map[string]string) (manual string, auto, ok bool) { + v := strings.TrimSpace(cfg["CACHE_RAM"]) + if v == "" || strings.EqualFold(v, "auto") { + return "", true, true + } + if n, err := strconv.Atoi(v); err == nil && n >= -1 { + return strconv.Itoa(n), false, true + } + return "", false, false +} + +// cacheRAMArgs : le --cache-ram que Loki ajoute, et ce qu'il en dit. +func cacheRAMArgs(cfg map[string]string, extra []string, si serveSysInfo) (args, notes []string) { + manual, auto, ok := cacheRAMSetting(cfg) + if !ok { + return nil, []string{"CACHE_RAM=" + cfg["CACHE_RAM"] + " illisible (un nombre de Mio, -1, 0 ou auto) : ignoré"} + } + if !strings.Contains(si.Help, "--cache-ram") { + if !auto { + notes = append(notes, "ce moteur ne connaît pas --cache-ram : CACHE_RAM ignoré") + } + return nil, notes + } + if hasAnyFlag(extra, "-cram", "--cache-ram") { + return nil, nil + } + if si.ArgEnv["LLAMA_ARG_CACHE_RAM"] != "" { + if !auto { + notes = append(notes, "LLAMA_ARG_CACHE_RAM est posée : CACHE_RAM ignoré") + } + return nil, notes + } + if !auto { + return []string{"--cache-ram", manual}, nil + } + n, why := cacheRAMAuto(cfg, extra, si) + if n <= 0 { + return nil, nil + } + return []string{"--cache-ram", strconv.FormatInt(n, 10)}, + []string{fmt.Sprintf("cache de prompts : --cache-ram %d (RAM hôte, copie exacte de l'état, jamais la VRAM ; %s). "+ + "CACHE_RAM= pour fixer, 0 pour couper.", n, why)} +} + +// cacheRAMOff : le cache de prompts est-il coupé (0) par l'une des trois +// sources ? Sans lui, effacer un slot ne protège rien. +func cacheRAMOff(cfg map[string]string, extra []string, argEnv map[string]string) bool { + v := flagValue(extra, "-cram", "--cache-ram") + if v == "" { + v = argEnv["LLAMA_ARG_CACHE_RAM"] + } + if v == "" { + v, _, _ = cacheRAMSetting(cfg) + } + return strings.TrimSpace(v) == "0" +} + +// cacheIsolationOn : CACHE_ISOLATE=off coupe l'effacement du slot après les +// travaux annexes (et donc --slot-save-path), sans toucher au code. +func cacheIsolationOn(cfg map[string]string) bool { + return !strings.EqualFold(strings.TrimSpace(cfg["CACHE_ISOLATE"]), "off") +} + +// loopbackHost : le moteur n'écoute que sur la machine. +func loopbackHost(h string) bool { + switch strings.Trim(strings.TrimSpace(h), "[]") { + case "127.0.0.1", "::1", "localhost": + return true + } + return false +} + +// slotSaveArgs : --slot-save-path, sans lequel llama-server refuse toute action +// sur /slots — l'effacement compris. Mais le drapeau ouvre aussi save et +// restore, qui écrivent des Gio sur le disque, à quiconque joint le moteur. On +// ne le pose donc que si le moteur est fermé aux autres (boucle locale, ou clé +// d'API exigée), que le dossier existe (le moteur refuse de démarrer sinon), et +// que l'isolation servira : cache actif, CACHE_ISOLATE pas à off. Les proxys de +// Loki refusent de leur côté tout /slots qui n'est pas une lecture. +func slotSaveArgs(cfg map[string]string, extra []string, si serveSysInfo) []string { + if si.SlotDir == "" || !strings.Contains(si.Help, "--slot-save-path") || + hasAnyFlag(extra, "--slot-save-path") || !cacheIsolationOn(cfg) || cacheRAMOff(cfg, extra, si.ArgEnv) { + return nil + } + host := flagValue(extra, "--host") + if host == "" { + host = cfg["HOST"] + } + if host == "" { + host = "0.0.0.0" + } + if !loopbackHost(host) && si.APIKey == "" { + return nil + } + return []string{"--slot-save-path", si.SlotDir} +} + +// prepareSlotDir crée LOKI_HOME/slots (0700) et le vide : Loki n'y sauve jamais +// rien, un fichier présent ne peut venir que d'un appel qui n'aurait pas dû +// passer. Vide = dossier indisponible, le drapeau ne sera pas posé. +func prepareSlotDir(home string) string { + dir := filepath.Join(home, "slots") + if err := os.MkdirAll(dir, 0o700); err != nil { + return "" + } + if st, err := os.Stat(dir); err != nil || !st.IsDir() { + return "" + } + if entries, err := os.ReadDir(dir); err == nil { + for _, e := range entries { + if e.Type().IsRegular() { + _ = os.Remove(filepath.Join(dir, e.Name())) + } + } + } + return dir +} + +// cacheRAMPreview : la taille que le moteur prendra pour ce preset, pour +// l'éditeur. 0 = défaut du moteur (8192 Mio). Même calcul qu'au lancement, mais +// sans l'aide du moteur : l'aperçu dit ce que Loki demanderait. +func cacheRAMPreview(content string) (mib int64, why string) { + cfg := parseEnv(content) + if isExternalConfig(cfg) { + return 0, "preset externe" + } + extra := splitArgs(cfg["EXTRA_ARGS"]) + if v := flagValue(extra, "-cram", "--cache-ram"); v != "" { + n, _ := strconv.ParseInt(v, 10, 64) + return n, "fixé par EXTRA_ARGS" + } + if v := strings.TrimSpace(os.Getenv("LLAMA_ARG_CACHE_RAM")); v != "" { + n, _ := strconv.ParseInt(v, 10, 64) + return n, "fixé par LLAMA_ARG_CACHE_RAM" + } + manual, auto, ok := cacheRAMSetting(cfg) + if !ok { + return 0, "CACHE_RAM illisible" + } + if !auto { + n, _ := strconv.ParseInt(manual, 10, 64) + return n, "fixé par CACHE_RAM" + } + si := serveSysInfo{ArgEnv: map[string]string{}} + for _, k := range append(append(append([]string{}, serveArgEnv...), fidelityArgEnv...), cacheArgEnv...) { + if v, ok := os.LookupEnv(k); ok && v != "" { + si.ArgEnv[k] = v + } + } + if m := strings.TrimSpace(cfg["MODEL"]); m != "" { + if p, err := resolveServeModelPath(m); err == nil { + si.Model = p + } + } + if si.Model == "" { + return 0, "modèle introuvable" + } + probeCacheRAM(cfg, extra, &si) + return cacheRAMAuto(cfg, extra, si) +} diff --git a/internal/loki/backend_serve_cache_test.go b/internal/loki/backend_serve_cache_test.go new file mode 100644 index 0000000..4938eb4 --- /dev/null +++ b/internal/loki/backend_serve_cache_test.go @@ -0,0 +1,264 @@ +package loki + +import ( + "os" + "path/filepath" + "reflect" + "strings" + "testing" + "testing/fstest" +) + +const helpCacheRAM = helpRecent + `-cram, --cache-ram N set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable) +--slot-save-path PATH path to save slot kv cache (default: disabled) +` + +// denseGPU : un 27B dense de 17 Go sur deux cartes de 24 Go, machine de 64 Go. +// État : 64 couches × 4 têtes KV × (256+256) × 2 octets = 0,25 Mio par jeton. +func denseGPU() serveSysInfo { + return serveSysInfo{ + Help: helpCacheRAM, + Model: "/models/dense.gguf", + GGUF: &GGUFInfo{Arch: "qwen3", BlockCount: 64, HeadCountKV: 4, KeyLen: 256, ValLen: 256, + ContextLength: 262144}, + ModelBytes: 17 << 30, + RAMMiB: 65536, RAMAvailMiB: 60000, + VRAMMiB: 2 * 24576, + ArgEnv: map[string]string{}, + } +} + +func TestCacheRAMArgs(t *testing.T) { + cases := []struct { + name string + cfg map[string]string + extra []string + si func(*serveSysInfo) + want []string + notes int + }{ + {name: "dense sur GPU, 64 Go, 65k : besoin 40 Gio, plafonné à 30 % de la RAM", + cfg: map[string]string{"CTX": "65536"}, want: []string{"--cache-ram", "19660"}, notes: 1}, + {name: "contexte court : le défaut du moteur suffit", + cfg: map[string]string{"CTX": "8192"}}, + {name: "cache q8_0 : l'état pèse deux fois moins, besoin sous le plafond", + cfg: map[string]string{"CTX": "32768", "KV_TYPE": "q8_0"}, want: []string{"--cache-ram", "10880"}, notes: 1}, + {name: "même contexte en f16 : plafonné", + cfg: map[string]string{"CTX": "32768"}, want: []string{"--cache-ram", "19660"}, notes: 1}, + {name: "-c d'EXTRA_ARGS l'emporte sur CTX", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"-c", "8192"}}, + {name: "MoE, experts sur CPU par -ot : jamais touché (pas de 4096)", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"-ot", "exps=CPU"}}, + {name: "-ot mixte avec un élément sur CPU", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"-ot", "a=CUDA0,b=CPU"}}, + {name: "-ot vers une autre carte seulement : tout reste en VRAM", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"-ot", "blk\\.6.*=CUDA1"}, + want: []string{"--cache-ram", "19660"}, notes: 1}, + {name: "--n-cpu-moe", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"--n-cpu-moe", "30"}}, + {name: "LLAMA_ARG_N_CPU_MOE", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.ArgEnv["LLAMA_ARG_N_CPU_MOE"] = "20" }}, + {name: "cache KV sur CPU", + cfg: map[string]string{"CTX": "65536"}, extra: []string{"-nkvo"}}, + {name: "NGL plus petit que le modèle", + cfg: map[string]string{"CTX": "65536", "NGL": "40"}}, + {name: "NGL=999 : toutes les couches", + cfg: map[string]string{"CTX": "65536", "NGL": "999"}, want: []string{"--cache-ram", "19660"}, notes: 1}, + {name: "modèle trop gros pour la VRAM", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.ModelBytes = 40 << 30 }}, + {name: "pas de VRAM connue (macOS, AMD, sans nvidia-smi)", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.VRAMMiB = 0 }}, + {name: "GGUF illisible", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.GGUF = nil }}, + {name: "RAM libre trop juste", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.RAMAvailMiB = 12000 }}, + {name: "CACHE_RAM fixé : passé tel quel", + cfg: map[string]string{"CTX": "65536", "CACHE_RAM": "12000"}, want: []string{"--cache-ram", "12000"}}, + {name: "CACHE_RAM=-1 : sans limite, passé tel quel", + cfg: map[string]string{"CACHE_RAM": "-1"}, want: []string{"--cache-ram", "-1"}}, + {name: "CACHE_RAM=0 : cache coupé", + cfg: map[string]string{"CACHE_RAM": "0"}, want: []string{"--cache-ram", "0"}}, + {name: "CACHE_RAM=auto", + cfg: map[string]string{"CTX": "65536", "CACHE_RAM": "auto"}, want: []string{"--cache-ram", "19660"}, notes: 1}, + {name: "CACHE_RAM illisible : ignoré, et dit", + cfg: map[string]string{"CACHE_RAM": "beaucoup"}, notes: 1}, + {name: "-cram d'EXTRA_ARGS : Loki se tait", + cfg: map[string]string{"CTX": "65536", "CACHE_RAM": "12000"}, extra: []string{"-cram", "2048"}}, + {name: "LLAMA_ARG_CACHE_RAM : la ligne de commande l'écraserait, Loki se tait", + cfg: map[string]string{"CTX": "65536"}, si: func(s *serveSysInfo) { s.ArgEnv["LLAMA_ARG_CACHE_RAM"] = "4096" }}, + {name: "LLAMA_ARG_CACHE_RAM et CACHE_RAM : la variable gagne, et on le dit", + cfg: map[string]string{"CACHE_RAM": "12000"}, si: func(s *serveSysInfo) { s.ArgEnv["LLAMA_ARG_CACHE_RAM"] = "4096" }, notes: 1}, + {name: "moteur sans --cache-ram : rien, même fixé", + cfg: map[string]string{"CACHE_RAM": "12000"}, si: func(s *serveSysInfo) { s.Help = helpRecent }, notes: 1}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + si := denseGPU() + if c.si != nil { + c.si(&si) + } + got, notes := cacheRAMArgs(c.cfg, c.extra, si) + if !reflect.DeepEqual(got, c.want) { + t.Fatalf("args = %v, attendu %v (notes %v)", got, c.want, notes) + } + if len(notes) != c.notes { + t.Fatalf("notes = %v, attendu %d", notes, c.notes) + } + }) + } +} + +// Hybride (Qwen3.5) : 16 couches d'attention sur 64, état récurrent des 48 +// autres, compté avec ses points de reprise. +func TestCacheRAMAutoHybride(t *testing.T) { + si := denseGPU() + layers := make([]int, 64) + for i := range layers { + if (i+1)%4 == 0 { + layers[i] = 4 + } + } + si.GGUF = &GGUFInfo{Arch: "qwen35", BlockCount: 64, HeadCountKV: 4, HeadCountKVLayers: layers, + KeyLen: 256, ValLen: 256, FullAttnInterval: 4, Hybrid: true, + SSMConv: 4, SSMInner: 6144, SSMState: 128, SSMGroups: 16} + perTok, rec := ggufStateBytes(*si.GGUF, "", "") + if perTok != 65536 || rec != 156893184 { + t.Fatalf("état = %v o/jeton, %d o récurrents", perTok, rec) + } + // 4 Gio de KV + 10 états récurrents (9 points de reprise + l'état courant), + // × 2,5. + n, why := cacheRAMAuto(map[string]string{"CTX": "65536"}, nil, si) + if n != 13981 { + t.Fatalf("cacheRAMAuto = %d (%s), attendu 13981", n, why) + } +} + +// Sans key_length, llama.cpp prend n_embd / n_head : nous aussi. +func TestGGUFStateBytesReplis(t *testing.T) { + g := GGUFInfo{BlockCount: 2, HeadCountKV: 2, EmbeddingLength: 1024, HeadCount: 8} + if p, _ := ggufStateBytes(g, "", ""); p != 2*2*(128+128)*2 { + t.Fatalf("repli n_embd/n_head : %v", p) + } + if p, _ := ggufStateBytes(GGUFInfo{BlockCount: 2, HeadCountKV: 2}, "", ""); p != 0 { + t.Fatalf("dimensions inconnues : %v, attendu 0", p) + } +} + +func TestVRAMFromSMI(t *testing.T) { + const out = "0, 24576\n1, 16384\n" + cases := []struct { + cvd string + pci bool + want int64 + }{ + {"", false, 40960}, + {"1", true, 16384}, + {"0,1", true, 40960}, + {"1", false, 16384}, // ordre inconnu : la plus petite carte + {"0", false, 16384}, + {"GPU-1234", true, 0}, + {"2", true, 0}, + } + for _, c := range cases { + if got := vramFromSMI(out, c.cvd, c.pci); got != c.want { + t.Errorf("vramFromSMI(%q, %v) = %d, attendu %d", c.cvd, c.pci, got, c.want) + } + } + if got := vramFromSMI("[N/A]\n", "", false); got != 0 { + t.Errorf("sortie illisible : %d", got) + } +} + +func TestCgroupMemMiB(t *testing.T) { + v2 := fstest.MapFS{ + "sys/fs/cgroup/memory.max": {Data: []byte("17179869184\n")}, + "sys/fs/cgroup/memory.current": {Data: []byte("4294967296\n")}, + } + if l, c := cgroupMemMiB(v2); l != 16384 || c != 4096 { + t.Fatalf("v2 = %d/%d", l, c) + } + if l, _ := cgroupMemMiB(fstest.MapFS{"sys/fs/cgroup/memory.max": {Data: []byte("max\n")}}); l != 0 { + t.Fatalf("v2 sans limite = %d", l) + } + v1 := fstest.MapFS{"sys/fs/cgroup/memory/memory.limit_in_bytes": {Data: []byte("9223372036854771712\n")}} + if l, _ := cgroupMemMiB(v1); l != 0 { + t.Fatalf("v1 sans limite = %d", l) + } +} + +func TestSlotSaveArgs(t *testing.T) { + want := []string{"--slot-save-path", "/data/slots"} + cases := []struct { + name string + cfg map[string]string + extra []string + si func(*serveSysInfo) + want []string + }{ + {name: "boucle locale", cfg: map[string]string{"HOST": "127.0.0.1"}, want: want}, + {name: "--host ::1 dans EXTRA_ARGS", extra: []string{"--host", "::1"}, want: want}, + {name: "0.0.0.0 sans clé : ouvert à tous, pas de drapeau", cfg: map[string]string{}}, + {name: "0.0.0.0 avec clé d'API", cfg: map[string]string{"HOST": "0.0.0.0"}, + si: func(s *serveSysInfo) { s.APIKey = "k" }, want: want}, + {name: "EXTRA_ARGS le pose déjà", cfg: map[string]string{"HOST": "127.0.0.1"}, + extra: []string{"--slot-save-path", "/ailleurs"}}, + {name: "CACHE_ISOLATE=off", cfg: map[string]string{"HOST": "127.0.0.1", "CACHE_ISOLATE": "off"}}, + {name: "cache coupé : rien à isoler", cfg: map[string]string{"HOST": "127.0.0.1", "CACHE_RAM": "0"}}, + {name: "cache coupé par EXTRA_ARGS", cfg: map[string]string{"HOST": "127.0.0.1"}, extra: []string{"-cram", "0"}}, + {name: "moteur sans le drapeau", cfg: map[string]string{"HOST": "127.0.0.1"}, + si: func(s *serveSysInfo) { s.Help = helpRecent }}, + {name: "dossier indisponible", cfg: map[string]string{"HOST": "127.0.0.1"}, + si: func(s *serveSysInfo) { s.SlotDir = "" }}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + si := serveSysInfo{Help: helpCacheRAM, SlotDir: "/data/slots", ArgEnv: map[string]string{}} + if c.si != nil { + c.si(&si) + } + if got := slotSaveArgs(c.cfg, c.extra, si); !reflect.DeepEqual(got, c.want) { + t.Fatalf("slotSaveArgs = %v, attendu %v", got, c.want) + } + }) + } +} + +// Place dans la ligne complète : après --api-key, avant EXTRA_ARGS qui garde le +// dernier mot. Sans aide qui les connaisse, la ligne historique est intacte +// (TestBuildServeArgs). +func TestBuildServeArgsCacheRAM(t *testing.T) { + si := denseGPU() + si.APIKey, si.SlotDir = "k", "/data/slots" + args, _, _ := buildServeArgs(map[string]string{"CTX": "65536", "NGL": "999"}, []string{"--jinja"}, "/bin/llama-server", si) + got := strings.Join(args, " ") + if !strings.HasSuffix(got, "--api-key k --cache-ram 19660 --slot-save-path /data/slots --jinja") { + t.Fatalf("ligne = %s", got) + } +} + +func TestPrepareSlotDir(t *testing.T) { + home := t.TempDir() + dir := prepareSlotDir(home) + if dir != filepath.Join(home, "slots") { + t.Fatalf("dossier = %q", dir) + } + stray := filepath.Join(dir, "slot0.bin") + if err := os.WriteFile(stray, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + if prepareSlotDir(home) == "" { + t.Fatal("second appel") + } + if _, err := os.Stat(stray); !os.IsNotExist(err) { + t.Fatal("un fichier laissé dans slots/ doit être purgé au lancement") + } + // Un fichier à la place du dossier : pas de drapeau, plutôt qu'un moteur + // qui refuse de démarrer. + bad := t.TempDir() + if err := os.WriteFile(filepath.Join(bad, "slots"), nil, 0o600); err != nil { + t.Fatal(err) + } + if got := prepareSlotDir(bad); got != "" { + t.Fatalf("slots est un fichier : %q", got) + } +} diff --git a/internal/loki/chat_compact.go b/internal/loki/chat_compact.go index 7dfb444..9e56f83 100644 --- a/internal/loki/chat_compact.go +++ b/internal/loki/chat_compact.go @@ -660,6 +660,9 @@ Write the summary in the SAME language as the conversation.` } req.Header.Set("Content-Type", "application/json") ep.auth(req.Header.Set) + if !ep.External { + defer engineRequestStart()() // voir llm_slots.go + } resp, err := http.DefaultClient.Do(req) if err != nil { return "", friendlyLLMError(err) diff --git a/internal/loki/code_subagent.go b/internal/loki/code_subagent.go index b6cc77a..3da21d8 100644 --- a/internal/loki/code_subagent.go +++ b/internal/loki/code_subagent.go @@ -103,6 +103,10 @@ func toolSubagent(ctx context.Context, args map[string]any, parent Caps) string {Role: "user", Content: task + framing}, } var out strings.Builder + // Le sous-agent occupe le slot de la conversation du builder : une fois son + // rapport rendu, on l'efface pour que le builder reparte du cache RAM + // (llm_slots.go). Pas entre ses lectures : seulement à la fin. + defer engineSideJob()() // Trace jetable : aucun forwardStream, rien de persisté. L'utilisateur voit // la bulle de l'outil `subagent` et son rapport, pas les dix lectures. // Température 0.6 et non 0 : en glouton, Qwen avec réflexion tourne vite en diff --git a/internal/loki/code_verify.go b/internal/loki/code_verify.go index ef04aed..facb878 100644 --- a/internal/loki/code_verify.go +++ b/internal/loki/code_verify.go @@ -112,7 +112,9 @@ func (c *Conversation) runVerifyPass(ctx context.Context, caps Caps, temperature {Role: "user", Content: user.String()}, } // Trace ISOLÉE : ni persistée dans c.Messages, ni compactée — elle vit et - // meurt dans cette passe. Seuls les deltas d'affichage sortent. + // meurt dans cette passe. Seuls les deltas d'affichage sortent. Son état + // dans le moteur aussi : effacé à la fin de la passe (llm_slots.go). + defer engineSideJob()() _, _ = runChat(ctx, msgs, temperature, vcaps, func(ev StreamEvent) bool { c.forwardStream(ev, epoch) return true diff --git a/internal/loki/llm_bench.go b/internal/loki/llm_bench.go index 733072e..0d187a2 100644 --- a/internal/loki/llm_bench.go +++ b/internal/loki/llm_bench.go @@ -54,6 +54,9 @@ func runBench(nPrompt, nPredict int) (*benchResult, error) { if !healthCheck() { return nil, fmt.Errorf("serveur injoignable sur :%d", port) } + // Le bench prend le slot de la conversation : une fois fini, on l'efface + // pour qu'elle soit rechargée depuis le cache RAM (llm_slots.go). + defer engineSideJob()() if nPrompt <= 0 { nPrompt = 2000 } @@ -90,6 +93,7 @@ func runBench(nPrompt, nPredict int) (*benchResult, error) { req.Header.Set("Content-Type", "application/json") authHeader(req) client := &http.Client{Timeout: 5 * time.Minute} + defer engineRequestStart()() // avant l'effacement, différé plus haut resp, err := client.Do(req) if err != nil { return nil, err diff --git a/internal/loki/llm_client.go b/internal/loki/llm_client.go index 44addd9..f51af5b 100644 --- a/internal/loki/llm_client.go +++ b/internal/loki/llm_client.go @@ -728,6 +728,9 @@ type StatsEvent struct { // Taille TOTALE du prompt traité ce tour (préfixe caché compris), issue de // `usage.prompt_tokens`. 0 si le backend ne renvoie pas d'usage. PromptTokensTotal int `json:"prompt_tokens_total,omitempty"` + // Conseil ponctuel quand le cache de prompts n'a pas tenu après un travail + // annexe (voir noteEnginePrompt). Vide la plupart du temps. + CacheHint string `json:"cache_hint,omitempty"` } // ChatCallback receives stream events. Return false to abort the stream. @@ -756,9 +759,12 @@ type streamChunk struct { } `json:"timings"` // Chunk final (include_usage) : taille totale du prompt, hors choices. Usage *struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + PromptTokensDetails *struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` } `json:"usage"` } @@ -1046,8 +1052,16 @@ func runChat(ctx context.Context, messages []Message, temperature float64, caps // temps qu'un 401 garanti. Chaque endpoint porte la sienne. ep.auth(req.Header.Set) tReq := time.Now() // secours des stats quand le serveur n'envoie pas de timings + // Requête en vol vers le moteur local, jusqu'à la lecture complète du corps : + // l'isolation des travaux annexes n'efface jamais le slot pendant ce temps + // (llm_slots.go). Rien à compter pour une API externe. + endReq := func() {} + if !ep.External { + endReq = engineRequestStart() + } resp, err := http.DefaultClient.Do(req) if err != nil { + endReq() // Rien n'est encore parti à l'écran : le tour est rejouable tel quel. // C'est le cas du moteur qui redémarre (bascule de preset, rechargement // de modèle) — quelques secondes de connexion refusée qui faisaient @@ -1070,6 +1084,7 @@ func runChat(ctx context.Context, messages []Message, temperature float64, caps if resp.StatusCode != http.StatusOK { b, _ := io.ReadAll(io.LimitReader(resp.Body, 2000)) resp.Body.Close() + endReq() msg := strings.TrimSpace(string(b)) if msg == "" { msg = resp.Status @@ -1249,6 +1264,13 @@ func runChat(ctx context.Context, messages []Message, temperature float64, caps if chunk.Usage != nil && chunk.Usage.PromptTokens > 0 { stats.PromptTokensTotal = chunk.Usage.PromptTokens s := stats + if !ep.External { + cached, has := 0, chunk.Usage.PromptTokensDetails != nil + if has { + cached = chunk.Usage.PromptTokensDetails.CachedTokens + } + s.CacheHint = noteEnginePrompt(chunk.Usage.PromptTokens, cached, has) + } scb(StreamEvent{Stats: &s}) } if len(chunk.Choices) == 0 { @@ -1464,6 +1486,7 @@ func runChat(ctx context.Context, messages []Message, temperature float64, caps // arguments tronqués. On refuse donc le tour, en le disant. scanErr := sc.Err() resp.Body.Close() + endReq() if !sawTimings && !tFirst.IsZero() { // Décodage mesuré ici ; le 1er token tombe dans la lecture du prompt. // Pas de débit de lecture : le serveur a pu réutiliser une partie du diff --git a/internal/loki/llm_oai.go b/internal/loki/llm_oai.go index b18352d..30df2d1 100644 --- a/internal/loki/llm_oai.go +++ b/internal/loki/llm_oai.go @@ -104,7 +104,18 @@ func oaiHandler() http.Handler { } return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { p := r.URL.Path + // /slots en lecture seulement : save, restore et erase restent à Loki + // (llm_slots.go). + if slotsWrite(r) { + sendOAIError(w, http.StatusMethodNotAllowed, + "/slots est en lecture seule à travers Loki", "invalid_request_error", "method_not_allowed") + return + } if strings.HasPrefix(p, "/v1") || p == "/health" || p == "/props" || p == "/metrics" || strings.HasPrefix(p, "/slots") { + // Requête en vol vers le moteur : l'isolation des travaux annexes + // n'efface pas le slot pendant ce temps. Différé : ReverseProxy panique + // (ErrAbortHandler) quand le client coupe en plein flux. + defer engineRequestStart()() lp.ServeHTTP(w, r) return } diff --git a/internal/loki/llm_slots.go b/internal/loki/llm_slots.go new file mode 100644 index 0000000..ea48da8 --- /dev/null +++ b/internal/loki/llm_slots.go @@ -0,0 +1,236 @@ +package loki + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" +) + +// Isolation des travaux annexes dans le cache de prompts. +// +// llama-server tourne en --parallel 1 : un vérificateur, un sous-agent, une +// tâche planifiée ou un bench prennent LE slot de la conversation. Le moteur +// range alors l'état de la conversation dans son cache RAM (exact, octet pour +// octet) et le recharge quand elle revient — à condition qu'il y soit encore. +// Or, à la requête suivante, il sauve d'abord l'état du travail annexe, et pour +// lui faire de la place il évince les entrées les plus anciennes : souvent +// justement la conversation, qui est alors recalculée en entier. +// +// Effacer le slot une fois le travail annexe TERMINÉ (pas entre ses étapes) +// règle ça : un slot vide n'est pas sauvé, et la conversation est rechargée +// depuis la RAM. On ne jette qu'un état qui ne resservirait pas — rien de ce +// que voit le modèle ne change. +// +// Mais un effacement est différé par le moteur tant que le slot travaille, et +// s'exécute ensuite sur ce qui s'y trouve alors : il faut donc être SÛR que le +// slot est à nous. D'où les garde-fous, tous nécessaires : +// - moteur local (jamais une API externe) ; +// - build vanilla ≥ 8660 : avant (PR #20993), un slot vide ne rechargeait pas +// depuis le cache, et effacer aurait forcé un recalcul complet. Build +// inconnu ou fork : on s'abstient ; +// - un seul slot (/props total_slots, pas la clé PARALLEL qu'EXTRA_ARGS peut +// contredire) ; +// - aucune requête de Loki en vol vers le moteur (compteur ci-dessous, tenu +// par le chat, les résumés, le bench et les proxys /v1), et le slot 0 au +// repos d'après /slots ; +// - appel synchrone, borné à 2 s, sous le verrou du compteur : aucune +// requête de Loki ne peut partir entre la vérification et l'effacement. + +// slotEraseMinBuild : premier build de llama.cpp qui recharge un slot vide +// depuis le cache de prompts (PR #20993, 50e0ad08fb). +const slotEraseMinBuild = 8660 + +// engineGate compte les requêtes de Loki en vol vers le moteur local. +// L'effacement prend le verrou pendant sa vérification et son appel : une +// nouvelle requête attend au plus ces deux secondes. +var engineGate struct { + mu sync.Mutex + inflight int +} + +// engineRequestStart marque une requête vers le moteur local ; la fonction +// rendue la clôt (une seule fois, quel que soit le nombre d'appels), une fois +// le corps de la réponse lu. +func engineRequestStart() func() { + engineGate.mu.Lock() + engineGate.inflight++ + engineGate.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + engineGate.mu.Lock() + engineGate.inflight-- + engineGate.mu.Unlock() + }) + } +} + +// slotsWrite : une action sur /slots (save, restore, erase) plutôt qu'une +// lecture. Les proxys de Loki la refusent : seul Loki, en local, efface — et +// save ou restore écriraient des Gio sur le disque pour un client distant. +func slotsWrite(r *http.Request) bool { + return strings.HasPrefix(r.URL.Path, "/slots") && r.Method != http.MethodGet && r.Method != http.MethodHead +} + +// slotIsolator parle au moteur local ; séparé pour les tests. +type slotIsolator struct { + base string // http://localhost: + auth func(set func(k, v string)) + client *http.Client +} + +var reBuildInfo = regexp.MustCompile(`^b(\d+)`) + +// get lit un JSON du moteur, borné à 2 s. +func (s slotIsolator) get(ctx context.Context, path string, out any) error { + ctx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.base+path, nil) + if err != nil { + return err + } + s.auth(req.Header.Set) + resp, err := s.client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("GET %s : %s", path, resp.Status) + } + return json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(out) +} + +// eraseIfIdle efface le slot 0 si, et seulement si, tous les garde-fous +// passent. erased=false sans erreur : on s'est abstenu, sans dommage. +func (s slotIsolator) eraseIfIdle(ctx context.Context) (erased bool, err error) { + var props struct { + BuildInfo string `json:"build_info"` + TotalSlots int `json:"total_slots"` + } + if err := s.get(ctx, "/props", &props); err != nil { + return false, err + } + build := 0 + if m := reBuildInfo.FindStringSubmatch(props.BuildInfo); m != nil { + build, _ = strconv.Atoi(m[1]) + } + if build < slotEraseMinBuild || props.TotalSlots != 1 { + return false, nil + } + engineGate.mu.Lock() + defer engineGate.mu.Unlock() + if engineGate.inflight > 0 { + return false, nil + } + var slots []struct { + ID int `json:"id"` + IsProcessing bool `json:"is_processing"` + } + if err := s.get(ctx, "/slots", &slots); err != nil { + return false, err + } + for _, sl := range slots { + if sl.ID == 0 && sl.IsProcessing { + return false, nil + } + } + ctx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.base+"/slots/0?action=erase", nil) + if err != nil { + return false, err + } + s.auth(req.Header.Set) + resp, err := s.client.Do(req) + if err != nil { + return false, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(io.LimitReader(resp.Body, 300)) + return false, fmt.Errorf("effacement du slot : %s %s", resp.Status, strings.TrimSpace(string(b))) + } + return true, nil +} + +var slotEraseLogOnce sync.Once + +// engineSideJob encadre un travail annexe (sous-agent, vérification, tâche, +// bench) : à appeler AVANT, et la fonction rendue APRÈS — d'habitude +// « defer engineSideJob()() ». Elle efface le slot si c'est sans risque, et +// arme la vérification du cache sur la requête suivante (noteEnginePrompt). +func engineSideJob() func() { + before := lastEnginePrompt.Load() + return func() { + ep := resolveChatEndpoint() + if ep.External { + return + } + cfg := ReadConfig() + argEnv := map[string]string{} + for _, k := range cacheArgEnv { + argEnv[k] = os.Getenv(k) + } + if !cacheIsolationOn(cfg) || cacheRAMOff(cfg, splitArgs(cfg["EXTRA_ARGS"]), argEnv) { + return + } + iso := slotIsolator{base: fmt.Sprintf("http://localhost:%d", LLMPort()), auth: ep.auth, client: http.DefaultClient} + erased, err := iso.eraseIfIdle(context.Background()) + if err != nil { + // Moteur sans --slot-save-path (501), ou /slots coupé : l'erreur est + // sans conséquence et se répéterait à chaque travail annexe. + slotEraseLogOnce.Do(func() { + fmt.Fprintf(os.Stderr, "[cache] isolation des travaux annexes indisponible : %v\n", err) + }) + } + if erased && before > 0 { + slotHintPending.Store(before) + } + } +} + +// Vérification du cache après un effacement. Le moteur dit, dans l'usage de +// chaque réponse, combien de jetons du prompt il a repris du cache +// (prompt_tokens_details.cached_tokens). Si la conversation revient après un +// travail annexe avec bien moins que ce qu'elle avait, son état n'a pas tenu +// dans le cache RAM : c'est CACHE_RAM qu'il faut relever. +var ( + lastEnginePrompt atomic.Int64 // prompt_tokens de la dernière réponse locale + slotHintPending atomic.Int64 // prompt d'avant le travail annexe effacé ; 0 = rien à vérifier + slotHintShown atomic.Bool // le conseil n'apparaît qu'une fois par lancement +) + +// slotHintMinTokens : en dessous, un recalcul coûte trop peu pour en parler. +const slotHintMinTokens = 4096 + +// noteEnginePrompt enregistre l'usage d'une réponse du moteur local et rend le +// conseil à afficher, s'il y a lieu (une seule fois par lancement ; le journal, +// lui, le note à chaque fois). +func noteEnginePrompt(total, cached int, hasCached bool) string { + if total > 0 { + lastEnginePrompt.Store(int64(total)) + } + before := slotHintPending.Swap(0) + // Un prompt plus court que celui d'avant n'est pas la suite de la même + // conversation (nouvelle discussion, autre tâche) : rien à comparer. + if !hasCached || before < slotHintMinTokens || int64(total) < before || int64(cached) >= before/2 { + return "" + } + fmt.Fprintf(os.Stderr, "[cache] conversation recalculée après un travail annexe : %d jetons repris du cache sur %d\n", + cached, before) + if slotHintShown.Swap(true) { + return "" + } + return fmt.Sprintf("Conversation recalculée après un travail annexe (%d jetons repris du cache sur %d) : "+ + "le cache de prompts en RAM est trop petit — augmente CACHE_RAM dans le preset.", cached, before) +} diff --git a/internal/loki/llm_slots_test.go b/internal/loki/llm_slots_test.go new file mode 100644 index 0000000..e642d73 --- /dev/null +++ b/internal/loki/llm_slots_test.go @@ -0,0 +1,176 @@ +package loki + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" +) + +// fakeEngine : un llama-server réduit à /props, /slots et l'effacement. +type fakeEngine struct { + build string + totalSlots int + processing bool + eraseCode int + + mu sync.Mutex + erases []*http.Request +} + +func (f *fakeEngine) server(t *testing.T) *httptest.Server { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodGet && r.URL.Path == "/props": + _ = json.NewEncoder(w).Encode(map[string]any{"build_info": f.build, "total_slots": f.totalSlots}) + case r.Method == http.MethodGet && r.URL.Path == "/slots": + _ = json.NewEncoder(w).Encode([]map[string]any{{"id": 0, "is_processing": f.processing}}) + case strings.HasPrefix(r.URL.Path, "/slots/"): + f.mu.Lock() + f.erases = append(f.erases, r) + f.mu.Unlock() + w.WriteHeader(f.eraseCode) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func (f *fakeEngine) isolator(srv *httptest.Server) slotIsolator { + ep := chatEndpoint{Key: "sk-local"} + return slotIsolator{base: srv.URL, auth: ep.auth, client: srv.Client()} +} + +func TestSlotEraseAppel(t *testing.T) { + f := &fakeEngine{build: "b8700-50e0ad08", totalSlots: 1, eraseCode: 200} + srv := f.server(t) + erased, err := f.isolator(srv).eraseIfIdle(context.Background()) + if err != nil || !erased { + t.Fatalf("eraseIfIdle = %v, %v", erased, err) + } + if len(f.erases) != 1 { + t.Fatalf("%d effacements", len(f.erases)) + } + r := f.erases[0] + if r.Method != http.MethodPost || r.URL.Path != "/slots/0" || r.URL.Query().Get("action") != "erase" { + t.Fatalf("requête %s %s", r.Method, r.URL) + } + if got := r.Header.Get("Authorization"); got != "Bearer sk-local" { + t.Fatalf("Authorization = %q", got) + } +} + +// Moteur sans --slot-save-path : 501, toléré. +func TestSlotErase501Tolere(t *testing.T) { + f := &fakeEngine{build: "b10678-abc", totalSlots: 1, eraseCode: http.StatusNotImplemented} + erased, err := f.isolator(f.server(t)).eraseIfIdle(context.Background()) + if erased || err == nil { + t.Fatalf("501 : eraseIfIdle = %v, %v", erased, err) + } +} + +// Chaque garde-fou suffit à empêcher l'effacement — sans erreur : s'abstenir +// est le comportement normal. +func TestSlotEraseGardeFous(t *testing.T) { + cases := []struct { + name string + build string + slots int + processing bool + busy bool + }{ + {name: "build avant le rechargement d'un slot vide", build: "b8640-abc", slots: 1}, + {name: "build inconnu (fork)", build: "", slots: 1}, + {name: "plusieurs slots", build: "b9000-abc", slots: 2}, + {name: "slot 0 au travail", build: "b9000-abc", slots: 1, processing: true}, + {name: "requête de Loki en vol (proxy /v1, tour de chat)", build: "b9000-abc", slots: 1, busy: true}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + f := &fakeEngine{build: c.build, totalSlots: c.slots, processing: c.processing, eraseCode: 200} + srv := f.server(t) + if c.busy { + done := engineRequestStart() + defer done() + } + erased, err := f.isolator(srv).eraseIfIdle(context.Background()) + if erased || err != nil || len(f.erases) != 0 { + t.Fatalf("eraseIfIdle = %v, %v ; %d effacements", erased, err, len(f.erases)) + } + }) + } +} + +// La fonction de fin d'une requête ne décompte qu'une fois, même rappelée +// (runChat la rappelle sur plusieurs chemins de sortie). +func TestEngineRequestStartIdempotent(t *testing.T) { + done := engineRequestStart() + done() + done() + engineGate.mu.Lock() + n := engineGate.inflight + engineGate.mu.Unlock() + if n != 0 { + t.Fatalf("inflight = %d", n) + } +} + +func TestNoteEnginePrompt(t *testing.T) { + slotHintShown.Store(false) + slotHintPending.Store(0) + // Rien d'armé : jamais de conseil. + if h := noteEnginePrompt(30000, 100, true); h != "" { + t.Fatalf("sans travail annexe : %q", h) + } + if lastEnginePrompt.Load() != 30000 { + t.Fatal("dernier prompt non retenu") + } + // Conversation rechargée du cache : rien à dire. + slotHintPending.Store(30000) + if h := noteEnginePrompt(30500, 30000, true); h != "" { + t.Fatalf("cache tenu : %q", h) + } + // Moteur sans cached_tokens : on ne sait pas, on se tait. + slotHintPending.Store(30000) + if h := noteEnginePrompt(30500, 0, false); h != "" { + t.Fatalf("sans cached_tokens : %q", h) + } + // Prompt plus court : une autre conversation, pas un recalcul. + slotHintPending.Store(30000) + if h := noteEnginePrompt(5000, 1000, true); h != "" { + t.Fatalf("autre conversation : %q", h) + } + // Recalcul complet : le conseil, une seule fois. + slotHintPending.Store(30000) + if h := noteEnginePrompt(30500, 2000, true); !strings.Contains(h, "CACHE_RAM") { + t.Fatalf("recalcul : %q", h) + } + slotHintPending.Store(30000) + if h := noteEnginePrompt(30500, 2000, true); h != "" { + t.Fatalf("second conseil : %q", h) + } +} + +// Les proxys de Loki ne laissent passer que les lectures de /slots. +func TestProxysRefusentLesActionsSlots(t *testing.T) { + t.Setenv("LOKI_HOME", t.TempDir()) + t.Setenv("LOKI_LINK_ALLOW_OAI", "1") + for name, h := range map[string]http.Handler{ + "oaiHandler": oaiHandler(), + "newLinkHandler": newLinkHandler(http.NewServeMux()), + } { + for _, action := range []string{"save", "restore", "erase"} { + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/slots/0?action="+action, strings.NewReader(`{"filename":"x.bin"}`)) + h.ServeHTTP(rec, req) + if rec.Code != http.StatusMethodNotAllowed { + t.Errorf("%s : POST /slots/0?action=%s → %d, attendu 405", name, action, rec.Code) + } + } + } +} diff --git a/internal/loki/relay_link.go b/internal/loki/relay_link.go index 2d9cc00..c967529 100644 --- a/internal/loki/relay_link.go +++ b/internal/loki/relay_link.go @@ -519,6 +519,13 @@ func newLinkHandler(mux *http.ServeMux) http.Handler { http.Error(w, "endpoint OpenAI désactivé via le relais (transiterait en clair) ; LOKI_LINK_ALLOW_OAI=1 pour l'autoriser", http.StatusForbidden) return } + // /slots en lecture seule : un pair distant n'efface, ne sauve ni ne + // restaure rien sur ce moteur (llm_slots.go). + if slotsWrite(r) { + http.Error(w, "/slots est en lecture seule via le relais", http.StatusMethodNotAllowed) + return + } + defer engineRequestStart()() // voir oaiHandler lp.ServeHTTP(w, r) return } diff --git a/internal/loki/tasks_run.go b/internal/loki/tasks_run.go index 8aa21f0..bda4494 100644 --- a/internal/loki/tasks_run.go +++ b/internal/loki/tasks_run.go @@ -90,6 +90,9 @@ func (c *Conversation) RunAutonomous(ctx context.Context, taskID, taskName, prom } var content strings.Builder + // La tâche a pris le slot de la conversation : effacé une fois la tâche + // finie, verrou de génération encore tenu (llm_slots.go). + defer engineSideJob()() _, err := runChat(ctx, InjectSkills(final, caps), temperature, caps, func(ev StreamEvent) bool { if ev.Content != "" { content.WriteString(ev.Content) diff --git a/internal/loki/ui/index.html b/internal/loki/ui/index.html index 4ebb4bb..d94dcab 100644 --- a/internal/loki/ui/index.html +++ b/internal/loki/ui/index.html @@ -3456,6 +3456,13 @@ html[data-files="1"] #files-btn{color:var(--accent)} UBATCH +
+ + Cache de prompts (Mio)RAM hôte, copie exacte du contexte · jamais la VRAM + +
+ Cache de prompts (Mio)RAM hôte, copie exacte du contexte · jamais la VRAM + +