mirror of
https://github.com/R0m1k3/Loki.git
synced 2026-10-11 17:26:57 +02:00
MichaelandClaude Opus 5.5
5f7203a935
À l'ouverture, l'éditeur demande en parallèle la taille du cache de prompts (et celle du second slot) et les avis MoE. Chaque aperçu lançait son nvidia-smi pour lire les mémoires totales — trois processus pour une réponse identique, la mémoire d'une carte ne bougeant pas. - La lecture brute des mémoires totales est partagée dix secondes ; le verrou tenu pendant la lecture fait attendre les demandes simultanées sur la même. La sélection des cartes (CUDA_VISIBLE_DEVICES) s'applique ensuite, comme avant. - Un échec est gardé le même délai : une salve ne relance pas trois lectures vouées à échouer. « loki serve » ne lit qu'une fois : rien n'y change. - Test : trois demandes simultanées, une lecture, chaque sélection juste ; échec non relancé dans la salve. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
320 lines
12 KiB
Go
320 lines
12 KiB
Go
package loki
|
||
|
||
import (
|
||
"errors"
|
||
"os"
|
||
"path/filepath"
|
||
"reflect"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"testing/fstest"
|
||
"time"
|
||
)
|
||
|
||
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, false)
|
||
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, false) == "" {
|
||
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, false); got != "" {
|
||
t.Fatalf("slots est un fichier : %q", got)
|
||
}
|
||
}
|
||
|
||
// L'éditeur demande en parallèle le cache de prompts, le second slot et les
|
||
// avis MoE : une seule lecture nvidia-smi pour tous, la sélection des cartes
|
||
// appliquée ensuite. Après le délai, une lecture neuve ; un échec n'est pas
|
||
// relancé dans la même salve.
|
||
func TestNvidiaVRAMShared(t *testing.T) {
|
||
var calls atomic.Int32
|
||
fail := false
|
||
prevQ := nvidiaTotalsQuery
|
||
nvidiaTotalsQuery = func(time.Duration) (string, error) {
|
||
calls.Add(1)
|
||
time.Sleep(20 * time.Millisecond) // les appels simultanés attendent la même lecture
|
||
if fail {
|
||
return "", errors.New("exit status 9")
|
||
}
|
||
return "0, 16311\n1, 12288\n", nil
|
||
}
|
||
t.Cleanup(func() {
|
||
nvidiaTotalsQuery = prevQ
|
||
smiTotals.mu.Lock()
|
||
smiTotals.exp = time.Time{} // les tests suivants relisent nvidia-smi
|
||
smiTotals.mu.Unlock()
|
||
})
|
||
expire := func() {
|
||
smiTotals.mu.Lock()
|
||
smiTotals.exp = time.Time{}
|
||
smiTotals.mu.Unlock()
|
||
}
|
||
expire()
|
||
var wg sync.WaitGroup
|
||
got := make([]int64, 3)
|
||
for i, cvd := range []string{"", "0", "1"} {
|
||
wg.Add(1)
|
||
go func() { defer wg.Done(); got[i] = nvidiaVRAMMiB(cvd, true, time.Second) }()
|
||
}
|
||
wg.Wait()
|
||
if n := calls.Load(); n != 1 {
|
||
t.Errorf("%d lectures nvidia-smi pour trois demandes simultanées", n)
|
||
}
|
||
if !reflect.DeepEqual(got, []int64{16311 + 12288, 16311, 12288}) {
|
||
t.Errorf("sélection des cartes : %v", got)
|
||
}
|
||
expire()
|
||
fail = true
|
||
if v := nvidiaVRAMMiB("", true, time.Second); v != 0 {
|
||
t.Errorf("échec : %d, attendu 0", v)
|
||
}
|
||
if v := nvidiaVRAMMiB("", true, time.Second); v != 0 || calls.Load() != 2 {
|
||
t.Errorf("échec relancé dans la même salve : %d lectures", calls.Load())
|
||
}
|
||
}
|