mirror of
https://github.com/R0m1k3/Loki.git
synced 2026-10-11 17:26:57 +02:00
Interface : une seule lecture nvidia-smi à l'ouverture de l'éditeur de preset
À 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>
This commit is contained in:
1 parent
4b659ac186
commit
5f7203a935
2 files changed
+97
-4
No files matched your search
@@ -11,6 +11,7 @@ import (
|
|||||||
"runtime"
|
"runtime"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"time"
|
"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
|
// nvidiaVRAMMiB : VRAM totale des cartes que llama-server verra. Borné dans le
|
||||||
// temps comme nvidiaGPUCount ; au moindre doute, 0.
|
// temps comme nvidiaGPUCount ; au moindre doute, 0.
|
||||||
func nvidiaVRAMMiB(cvd string, pciOrder bool, timeout time.Duration) int64 {
|
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 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)
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
out, err := hideCmd(exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.total",
|
out, err := hideCmd(exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.total",
|
||||||
"--format=csv,noheader,nounits")).Output()
|
"--format=csv,noheader,nounits")).Output()
|
||||||
if err != nil {
|
return string(out), err
|
||||||
return 0
|
}
|
||||||
|
|
||||||
|
// 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
|
// vramFromSMI additionne la VRAM des cartes sélectionnées par
|
||||||
|
|||||||
@@ -1,12 +1,16 @@
|
|||||||
package loki
|
package loki
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"errors"
|
||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"reflect"
|
"reflect"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
"testing"
|
"testing"
|
||||||
"testing/fstest"
|
"testing/fstest"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
const helpCacheRAM = helpRecent + `-cram, --cache-ram N set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)
|
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)
|
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())
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user