diff --git a/internal/loki/backend_serve_cache.go b/internal/loki/backend_serve_cache.go index ed536af..a095048 100644 --- a/internal/loki/backend_serve_cache.go +++ b/internal/loki/backend_serve_cache.go @@ -11,6 +11,7 @@ import ( "runtime" "strconv" "strings" + "sync" "time" ) @@ -121,17 +122,54 @@ func cgroupMemMiB(fsys fs.FS) (limit, current int64) { // 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") { + out, ok := smiTotals.get(time.Now, timeout) + if !ok { return 0 } + return vramFromSMI(out, cvd, pciOrder) +} + +// nvidiaTotalsQuery : la lecture brute des mémoires totales. Variable pour +// les tests. +var nvidiaTotalsQuery = func(timeout time.Duration) (string, error) { + if !hasTool("nvidia-smi") { + return "", exec.ErrNotFound + } 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 string(out), err +} + +// smiTotalsCache : la lecture des mémoires totales, partagée quelques +// secondes. L'éditeur de preset demande à l'ouverture la taille du cache de +// prompts (et celle du second slot) et les avis MoE, en parallèle : chacun +// lançait SON nvidia-smi pour une réponse identique — la mémoire totale d'une +// carte ne bouge pas. Le verrou tenu pendant la lecture fait attendre les +// appels simultanés sur la même ; la sélection des cartes (CUDA_VISIBLE_DEVICES) +// s'applique après, sur la sortie brute. Un échec est gardé aussi, le temps du +// délai : la même salve ne relance pas trois lectures qui échoueront. +type smiTotalsCache struct { + mu sync.Mutex + exp time.Time + out string + ok bool +} + +const smiTotalsTTL = 10 * time.Second + +var smiTotals smiTotalsCache + +func (c *smiTotalsCache) get(now func() time.Time, timeout time.Duration) (string, bool) { + c.mu.Lock() + defer c.mu.Unlock() + if now().Before(c.exp) { + return c.out, c.ok } - return vramFromSMI(string(out), cvd, pciOrder) + out, err := nvidiaTotalsQuery(timeout) + c.out, c.ok, c.exp = out, err == nil, now().Add(smiTotalsTTL) + return c.out, c.ok } // vramFromSMI additionne la VRAM des cartes sélectionnées par diff --git a/internal/loki/backend_serve_cache_test.go b/internal/loki/backend_serve_cache_test.go index d31610a..7553071 100644 --- a/internal/loki/backend_serve_cache_test.go +++ b/internal/loki/backend_serve_cache_test.go @@ -1,12 +1,16 @@ 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) @@ -262,3 +266,54 @@ func TestPrepareSlotDir(t *testing.T) { 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()) + } +}