Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
897163042c | ||
|
|
e58ff33912 | ||
|
|
2040b30c94 | ||
|
|
a07b6726ea | ||
|
|
d2c49e52fa |
+58
-22
@@ -813,7 +813,7 @@ func run(ctx context.Context, cfg config.Config, log *slog.Logger, logOut io.Wri
|
||||
// AUTO_UPDATE.
|
||||
if isService {
|
||||
serveControl(ctx, log, u, exePath, applyStaged,
|
||||
statusProvider(cfg, lk, comfySup, health, started, gw),
|
||||
statusProvider(cfg, lk, comfySup, health, started, gw, ollamaClient),
|
||||
reloadHandler(cfg, configPath, restartWhenIdle))
|
||||
}
|
||||
|
||||
@@ -847,14 +847,17 @@ type gpuWatch struct {
|
||||
enabled bool
|
||||
usedMB int
|
||||
total int
|
||||
tempC int
|
||||
fanPct int
|
||||
known bool
|
||||
foreign string // last Check result: external holders, "" when none
|
||||
at time.Time
|
||||
}
|
||||
|
||||
func (g *gpuWatch) setVRAM(used, total int) {
|
||||
func (g *gpuWatch) setVRAM(st game.GPUStats) {
|
||||
g.mu.Lock()
|
||||
g.usedMB, g.total, g.known = used, total, true
|
||||
g.usedMB, g.total = st.UsedMB, st.TotalMB
|
||||
g.tempC, g.fanPct, g.known = st.TempC, st.FanPct, true
|
||||
g.mu.Unlock()
|
||||
}
|
||||
|
||||
@@ -864,14 +867,15 @@ func (g *gpuWatch) setCheck(foreign string) {
|
||||
g.mu.Unlock()
|
||||
}
|
||||
|
||||
func (g *gpuWatch) get() (used, total int, known bool, foreign string, ageS int64) {
|
||||
func (g *gpuWatch) get() (st game.GPUStats, known bool, foreign string, ageS int64) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
ageS = -1
|
||||
if !g.at.IsZero() {
|
||||
ageS = int64(time.Since(g.at).Seconds())
|
||||
}
|
||||
return g.usedMB, g.total, g.known, g.foreign, ageS
|
||||
return game.GPUStats{UsedMB: g.usedMB, TotalMB: g.total, TempC: g.tempC, FanPct: g.fanPct},
|
||||
g.known, g.foreign, ageS
|
||||
}
|
||||
|
||||
// gameLoop polls for foreign GPU holders (a game, another ML job). While one
|
||||
@@ -894,8 +898,8 @@ func gameLoop(ctx context.Context, cfg config.Config, log *slog.Logger, det *gam
|
||||
log.Warn("game detection failed", "err", err)
|
||||
}
|
||||
gw.setCheck(summarizeHolders(holders))
|
||||
if used, total, verr := game.QueryVRAMMB(ctx); verr == nil {
|
||||
gw.setVRAM(used, total)
|
||||
if st, verr := game.QueryGPUStats(ctx); verr == nil {
|
||||
gw.setVRAM(st)
|
||||
}
|
||||
switch {
|
||||
case len(holders) > 0 && !held:
|
||||
@@ -1106,6 +1110,16 @@ type statusDownstream struct {
|
||||
URL string `json:"url"`
|
||||
Up bool `json:"up"`
|
||||
Managed string `json:"managed,omitempty"`
|
||||
// Models lists Ollama's loaded models with their VRAM footprint. Null
|
||||
// when unknown (query failed / not applicable); [] means none loaded —
|
||||
// deliberately no omitempty so the two stay distinguishable.
|
||||
Models []statusModel `json:"models"`
|
||||
}
|
||||
|
||||
// statusModel is one loaded Ollama model.
|
||||
type statusModel struct {
|
||||
Name string `json:"name"`
|
||||
VRAMMB int64 `json:"vram_mb"` // 0 = resident in RAM, not VRAM
|
||||
}
|
||||
|
||||
type statusLock struct {
|
||||
@@ -1138,7 +1152,11 @@ type statusGPU struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
UsedMB int `json:"used_mb"`
|
||||
TotalMB int `json:"total_mb"`
|
||||
Known bool `json:"known"`
|
||||
// TempC/FanPct are -1 when unknown (never sampled or nvidia-smi
|
||||
// reported N/A).
|
||||
TempC int `json:"temp_c"`
|
||||
FanPct int `json:"fan_pct"`
|
||||
Known bool `json:"known"`
|
||||
// Foreign is the last detector finding (external GPU holders), empty
|
||||
// when the last check found none.
|
||||
Foreign string `json:"foreign,omitempty"`
|
||||
@@ -1183,23 +1201,41 @@ func diffConfig(a, b config.Config) []string {
|
||||
return out
|
||||
}
|
||||
|
||||
// statusProvider assembles the one-line JSON snapshot for CmdStatus.
|
||||
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch) func() string {
|
||||
// statusProvider assembles the one-line JSON snapshot for CmdStatus. The
|
||||
// loaded-model query to Ollama gets a short timeout so a wedged upstream
|
||||
// cannot stall the status channel for long.
|
||||
func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Process, health *healthTracker, started time.Time, gw *gpuWatch, ollamaClient *ollama.Client) func() string {
|
||||
return func() string {
|
||||
snap := statusSnapshot{
|
||||
Version: version,
|
||||
UptimeS: int64(time.Since(started).Seconds()),
|
||||
}
|
||||
used, total, known, foreign, ageS := gw.get()
|
||||
st, known, foreign, ageS := gw.get()
|
||||
snap.GPU = statusGPU{
|
||||
Enabled: gw.enabled,
|
||||
UsedMB: used, TotalMB: total, Known: known,
|
||||
UsedMB: st.UsedMB, TotalMB: st.TotalMB, Known: known,
|
||||
TempC: -1, FanPct: -1,
|
||||
Foreign: foreign, AgeS: ageS,
|
||||
}
|
||||
if known {
|
||||
snap.GPU.TempC, snap.GPU.FanPct = st.TempC, st.FanPct
|
||||
}
|
||||
if cfg.OllamaURL != "" {
|
||||
snap.Downstreams = append(snap.Downstreams, statusDownstream{
|
||||
d := statusDownstream{
|
||||
Name: "ollama", URL: cfg.OllamaURL, Up: health.get("ollama"),
|
||||
})
|
||||
}
|
||||
if ollamaClient != nil {
|
||||
mctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
models, err := ollamaClient.LoadedModelDetails(mctx)
|
||||
cancel()
|
||||
if err == nil {
|
||||
d.Models = make([]statusModel, 0, len(models))
|
||||
for _, m := range models {
|
||||
d.Models = append(d.Models, statusModel{Name: m.Name, VRAMMB: m.SizeVRAM / (1024 * 1024)})
|
||||
}
|
||||
}
|
||||
}
|
||||
snap.Downstreams = append(snap.Downstreams, d)
|
||||
}
|
||||
if cfg.ComfyURL != "" {
|
||||
d := statusDownstream{Name: "comfy", URL: cfg.ComfyURL, Up: health.get("comfy")}
|
||||
@@ -1208,15 +1244,15 @@ func statusProvider(cfg config.Config, lk *lock.Lock, comfySup *supervise.Proces
|
||||
}
|
||||
snap.Downstreams = append(snap.Downstreams, d)
|
||||
}
|
||||
st := lk.Status()
|
||||
lst := lk.Status()
|
||||
snap.Lock = statusLock{
|
||||
State: string(st.State),
|
||||
Detail: st.Detail,
|
||||
LLMInflight: st.LLMInflight,
|
||||
LLMWaiting: st.LLMWaiting,
|
||||
ImageQueue: st.ImageQueue,
|
||||
External: st.External,
|
||||
SinceS: int64(time.Since(st.Since).Seconds()),
|
||||
State: string(lst.State),
|
||||
Detail: lst.Detail,
|
||||
LLMInflight: lst.LLMInflight,
|
||||
LLMWaiting: lst.LLMWaiting,
|
||||
ImageQueue: lst.ImageQueue,
|
||||
External: lst.External,
|
||||
SinceS: int64(time.Since(lst.Since).Seconds()),
|
||||
}
|
||||
b, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
|
||||
@@ -22,13 +22,14 @@ const (
|
||||
)
|
||||
|
||||
// hotkeysLine is the monitor's footer.
|
||||
const hotkeysLine = " " + cDim + "q quit · u update now" + cReset + "\x1b[K\n"
|
||||
const hotkeysLine = " " + cDim + "q quit · u update now · r reload config" + cReset + "\x1b[K\n"
|
||||
|
||||
// monitorCommand renders a live status view of the running service,
|
||||
// refreshed every second from the control channel. When the service
|
||||
// reports a different version and the executable on disk changed (the
|
||||
// updater replaced it), the monitor restarts itself onto the new binary.
|
||||
// Hotkeys: q quits, u triggers an update check on the service.
|
||||
// Hotkeys: q quits, u triggers an update check on the service, r asks the
|
||||
// service to reload its config file.
|
||||
func monitorCommand() int {
|
||||
if !stdoutIsTerminal() {
|
||||
fmt.Fprintln(os.Stderr, "gpu-turnstile: --monitor needs an interactive terminal")
|
||||
@@ -57,7 +58,27 @@ func monitorCommand() int {
|
||||
var note string
|
||||
var noteAt time.Time
|
||||
noteCh := make(chan string, 1)
|
||||
updatePending := false
|
||||
askPending := false
|
||||
|
||||
// ask sends a one-shot command to the service and reports the reply in
|
||||
// the note line. Only one ask runs at a time.
|
||||
ask := func(cmd, busy, label string) {
|
||||
if askPending {
|
||||
return
|
||||
}
|
||||
askPending = true
|
||||
note, noteAt = busy, time.Now()
|
||||
go func() {
|
||||
reply, err := control.Ask(cmd)
|
||||
if err != nil {
|
||||
noteCh <- label + ": no answer from the service"
|
||||
return
|
||||
}
|
||||
msg := strings.TrimPrefix(reply, "OK ")
|
||||
msg = strings.TrimPrefix(msg, "ERR ")
|
||||
noteCh <- label + ": " + msg
|
||||
}()
|
||||
}
|
||||
|
||||
poll := func() string {
|
||||
frame := renderWaiting()
|
||||
@@ -104,23 +125,12 @@ func monitorCommand() int {
|
||||
case 'q', 'Q', 3: // q or Ctrl+C (raw mode delivers it as a byte)
|
||||
return 0
|
||||
case 'u', 'U':
|
||||
if !updatePending {
|
||||
updatePending = true
|
||||
note, noteAt = "checking for updates…", time.Now()
|
||||
go func() {
|
||||
reply, err := control.Ask(control.CmdUpdateNow)
|
||||
if err != nil {
|
||||
noteCh <- "update: no answer from the service"
|
||||
return
|
||||
}
|
||||
msg := strings.TrimPrefix(reply, "OK ")
|
||||
msg = strings.TrimPrefix(msg, "ERR ")
|
||||
noteCh <- "update: " + msg
|
||||
}()
|
||||
}
|
||||
ask(control.CmdUpdateNow, "checking for updates…", "update")
|
||||
case 'r', 'R':
|
||||
ask(control.CmdReloadEnv, "reloading config…", "reload")
|
||||
}
|
||||
case n := <-noteCh:
|
||||
updatePending = false
|
||||
askPending = false
|
||||
note, noteAt = n, time.Now()
|
||||
}
|
||||
}
|
||||
@@ -188,7 +198,9 @@ func renderMonitor(snap statusSnapshot, width int) string {
|
||||
b.WriteString(cDim + " " + strings.Repeat("─", width-2) + cReset + "\x1b[K\n")
|
||||
|
||||
for _, d := range snap.Downstreams {
|
||||
b.WriteString(renderDownstream(d) + "\x1b[K\n")
|
||||
busy := (d.Name == "ollama" && snap.Lock.State == "llm" && snap.Lock.LLMInflight > 0) ||
|
||||
(d.Name == "comfy" && snap.Lock.State == "image")
|
||||
b.WriteString(renderDownstream(d, busy) + "\x1b[K\n")
|
||||
}
|
||||
b.WriteString("\x1b[K\n")
|
||||
b.WriteString(renderLock(snap.Lock) + "\x1b[K\n")
|
||||
@@ -213,6 +225,12 @@ func renderGPU(g statusGPU) string {
|
||||
s := " GPU: "
|
||||
if g.Known {
|
||||
s += renderVRAM(g.UsedMB, g.TotalMB)
|
||||
if g.TempC >= 0 {
|
||||
s += fmt.Sprintf(" · %d°C", g.TempC)
|
||||
}
|
||||
if g.FanPct >= 0 {
|
||||
s += fmt.Sprintf(" · fan %d%%", g.FanPct)
|
||||
}
|
||||
} else {
|
||||
s += cDim + "VRAM unknown (nvidia-smi not answering)" + cReset
|
||||
}
|
||||
@@ -227,25 +245,46 @@ func renderGPU(g statusGPU) string {
|
||||
return s
|
||||
}
|
||||
|
||||
// formatMB renders MiB as "4.2 GiB" (or MiB below 1 GiB).
|
||||
func formatMB(mb int64) string {
|
||||
if mb >= 1024 {
|
||||
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
|
||||
}
|
||||
return fmt.Sprintf("%d MiB", mb)
|
||||
}
|
||||
|
||||
// renderVRAM renders "4.2 / 16.0 GiB used" (or MiB below 1 GiB).
|
||||
func renderVRAM(used, total int) string {
|
||||
format := func(mb int) string {
|
||||
if mb >= 1024 {
|
||||
return fmt.Sprintf("%.1f GiB", float64(mb)/1024)
|
||||
}
|
||||
return fmt.Sprintf("%d MiB", mb)
|
||||
}
|
||||
if total > 0 {
|
||||
return format(used) + " / " + format(total) + " used"
|
||||
return formatMB(int64(used)) + " / " + formatMB(int64(total)) + " used"
|
||||
}
|
||||
return format(used) + " used"
|
||||
return formatMB(int64(used)) + " used"
|
||||
}
|
||||
|
||||
// printableLen counts characters without ANSI escapes (ASCII-only content).
|
||||
func printableLen(s string) int { return len(s) }
|
||||
|
||||
func renderDownstream(d statusDownstream) string {
|
||||
func renderDownstream(d statusDownstream, busy bool) string {
|
||||
url := cDim + d.URL + cReset
|
||||
detail := ""
|
||||
if d.Models != nil {
|
||||
if len(d.Models) == 0 {
|
||||
detail = cDim + " · no models loaded" + cReset
|
||||
} else {
|
||||
parts := make([]string, 0, len(d.Models))
|
||||
for _, m := range d.Models {
|
||||
if m.VRAMMB > 0 {
|
||||
parts = append(parts, fmt.Sprintf("%s (%s VRAM)", m.Name, formatMB(m.VRAMMB)))
|
||||
} else {
|
||||
parts = append(parts, m.Name+" (in RAM)")
|
||||
}
|
||||
}
|
||||
detail = " · " + strings.Join(parts, ", ")
|
||||
}
|
||||
}
|
||||
if busy {
|
||||
detail += " · " + cCyan + "busy" + cReset
|
||||
}
|
||||
switch d.Managed {
|
||||
case "stopped":
|
||||
return fmt.Sprintf(" %s○%s %-8s %sstopped (managed — starts on demand)%s %s",
|
||||
@@ -259,7 +298,7 @@ func renderDownstream(d statusDownstream) string {
|
||||
suffix = " (external)"
|
||||
}
|
||||
if d.Up {
|
||||
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, url)
|
||||
return fmt.Sprintf(" %s●%s %-8s %sUP%s%s%s %s", cGreen, cReset, d.Name, cGreen, cReset, suffix, detail, url)
|
||||
}
|
||||
return fmt.Sprintf(" %s●%s %-8s %sDOWN%s %s", cRed, cReset, d.Name, cRed, cReset, url)
|
||||
}
|
||||
|
||||
@@ -32,9 +32,9 @@ func TestRenderMonitor(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
|
||||
snap.GPU = statusGPU{Enabled: true, Known: true, UsedMB: 4300, TotalMB: 16384, TempC: 55, FanPct: 42, Foreign: "cyberpunk2077.exe (pid 1234)", AgeS: 12}
|
||||
frame = renderMonitor(snap, 80)
|
||||
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "external:", "checked 12s ago"} {
|
||||
for _, want := range []string{"GPU:", "4.2 GiB / 16.0 GiB used", "55°C", "fan 42%", "external:", "checked 12s ago"} {
|
||||
if !strings.Contains(frame, want) {
|
||||
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||
}
|
||||
@@ -47,6 +47,25 @@ func TestRenderMonitor(t *testing.T) {
|
||||
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||
}
|
||||
}
|
||||
|
||||
snap.Downstreams[0].Models = []statusModel{{Name: "llama3.1:8b", VRAMMB: 4900}, {Name: "embed", VRAMMB: 0}}
|
||||
snap.Lock = statusLock{State: "llm", LLMInflight: 1}
|
||||
frame = renderMonitor(snap, 80)
|
||||
for _, want := range []string{"llama3.1:8b (4.8 GiB VRAM)", "embed (in RAM)", "busy"} {
|
||||
if !strings.Contains(frame, want) {
|
||||
t.Errorf("frame missing %q:\n%s", want, frame)
|
||||
}
|
||||
}
|
||||
|
||||
snap.Downstreams[0].Models = []statusModel{}
|
||||
snap.Lock = statusLock{State: "idle"}
|
||||
frame = renderMonitor(snap, 80)
|
||||
if !strings.Contains(frame, "no models loaded") {
|
||||
t.Errorf("frame missing %q:\n%s", "no models loaded", frame)
|
||||
}
|
||||
if strings.Contains(frame, "busy") {
|
||||
t.Errorf("idle lock still shows busy:\n%s", frame)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFmtDur(t *testing.T) {
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
# ComfyUI --listen 0.0.0.0 --port 8189).
|
||||
services:
|
||||
gpu-turnstile:
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.2.9
|
||||
image: git.rambossek.at/public/gpu-turnstile:v0.3.1
|
||||
restart: unless-stopped
|
||||
environment:
|
||||
# Each consumer is enabled by setting its URL; leave one unset to
|
||||
|
||||
+42
-15
@@ -186,27 +186,54 @@ func queryComputeApps(ctx context.Context) ([]computeApp, error) {
|
||||
return parseComputeApps(string(out))
|
||||
}
|
||||
|
||||
// QueryVRAMMB returns used and total GPU VRAM in MiB via nvidia-smi.
|
||||
// Unlike the per-process list this works under WDDM too.
|
||||
func QueryVRAMMB(ctx context.Context) (used, total int, err error) {
|
||||
// GPUStats is one nvidia-smi reading of the whole card.
|
||||
type GPUStats struct {
|
||||
UsedMB int
|
||||
TotalMB int
|
||||
TempC int // -1 when nvidia-smi reports N/A
|
||||
FanPct int // -1 when N/A (some cards don't expose the fan)
|
||||
}
|
||||
|
||||
// QueryGPUStats returns VRAM usage, temperature and fan speed via
|
||||
// nvidia-smi. Unlike the per-process list this works under WDDM too.
|
||||
func QueryGPUStats(ctx context.Context) (GPUStats, error) {
|
||||
out, err := exec.CommandContext(ctx, "nvidia-smi",
|
||||
"--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits").Output()
|
||||
"--query-gpu=memory.used,memory.total,temperature.gpu,fan.speed", "--format=csv,noheader,nounits").Output()
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
return GPUStats{}, err
|
||||
}
|
||||
usedStr, totalStr, ok := strings.Cut(strings.TrimSpace(string(out)), ",")
|
||||
if !ok {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(string(out)))
|
||||
return parseGPUStats(string(out))
|
||||
}
|
||||
|
||||
// parseGPUStats parses one "used, total, temp, fan" CSV line (MiB, °C,
|
||||
// percent). The memory fields must be numeric; temperature and fan fall
|
||||
// back to -1 on "N/A" and friends.
|
||||
func parseGPUStats(out string) (GPUStats, error) {
|
||||
fields := strings.Split(strings.TrimSpace(out), ",")
|
||||
if len(fields) != 4 {
|
||||
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected output %q", strings.TrimSpace(out))
|
||||
}
|
||||
used, err = strconv.Atoi(strings.TrimSpace(usedStr))
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(string(out)))
|
||||
num := func(s string) (int, error) {
|
||||
return strconv.Atoi(strings.TrimSpace(s))
|
||||
}
|
||||
total, err = strconv.Atoi(strings.TrimSpace(totalStr))
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(string(out)))
|
||||
optional := func(s string) int {
|
||||
n, err := num(s)
|
||||
if err != nil {
|
||||
return -1
|
||||
}
|
||||
return n
|
||||
}
|
||||
return used, total, nil
|
||||
var st GPUStats
|
||||
var err error
|
||||
if st.UsedMB, err = num(fields[0]); err != nil {
|
||||
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected used memory in %q", strings.TrimSpace(out))
|
||||
}
|
||||
if st.TotalMB, err = num(fields[1]); err != nil {
|
||||
return GPUStats{}, fmt.Errorf("nvidia-smi: unexpected total memory in %q", strings.TrimSpace(out))
|
||||
}
|
||||
st.TempC = optional(fields[2])
|
||||
st.FanPct = optional(fields[3])
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// parseComputeApps parses "pid, used_memory" CSV lines (no header, MiB
|
||||
|
||||
@@ -116,6 +116,31 @@ func TestParseGPUEngineInstance(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseGPUStats(t *testing.T) {
|
||||
st, err := parseGPUStats("4300, 16384, 55, 42\n")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.UsedMB != 4300 || st.TotalMB != 16384 || st.TempC != 55 || st.FanPct != 42 {
|
||||
t.Errorf("got %+v", st)
|
||||
}
|
||||
|
||||
// Cards that don't expose temperature/fan report N/A.
|
||||
st, err = parseGPUStats("1024, 16384, N/A, N/A")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.TempC != -1 || st.FanPct != -1 {
|
||||
t.Errorf("got %+v, want -1 for N/A fields", st)
|
||||
}
|
||||
|
||||
for _, bad := range []string{"", "1, 2", "x, 16384, 55, 42", "1024, x, 55, 42", "1, 2, 3, 4, 5"} {
|
||||
if _, err := parseGPUStats(bad); err == nil {
|
||||
t.Errorf("%q parsed, want failure", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessesLive(t *testing.T) {
|
||||
if runtime.GOOS != "windows" && runtime.GOOS != "linux" {
|
||||
t.Skip("no process listing on this platform")
|
||||
|
||||
@@ -60,13 +60,21 @@ func (c *Client) Probe(ctx context.Context) error {
|
||||
|
||||
type psResponse struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
Model string `json:"model"`
|
||||
Name string `json:"name"`
|
||||
Model string `json:"model"`
|
||||
SizeVRAM int64 `json:"size_vram"` // bytes resident in VRAM (0 = RAM-only)
|
||||
} `json:"models"`
|
||||
}
|
||||
|
||||
// LoadedModels returns the names of models currently held in memory.
|
||||
func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
|
||||
// LoadedModel is one model currently held in memory.
|
||||
type LoadedModel struct {
|
||||
Name string
|
||||
SizeVRAM int64 // bytes resident in VRAM; 0 when the model sits in RAM
|
||||
}
|
||||
|
||||
// LoadedModelDetails returns the models currently held in memory with
|
||||
// their VRAM footprint.
|
||||
func (c *Client) LoadedModelDetails(ctx context.Context) ([]LoadedModel, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+"/api/ps", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -84,13 +92,26 @@ func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
|
||||
if err := json.NewDecoder(resp.Body).Decode(&ps); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
models := make([]string, 0, len(ps.Models))
|
||||
models := make([]LoadedModel, 0, len(ps.Models))
|
||||
for _, m := range ps.Models {
|
||||
if m.Name != "" {
|
||||
models = append(models, m.Name)
|
||||
} else {
|
||||
models = append(models, m.Model)
|
||||
name := m.Name
|
||||
if name == "" {
|
||||
name = m.Model
|
||||
}
|
||||
models = append(models, LoadedModel{Name: name, SizeVRAM: m.SizeVRAM})
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
// LoadedModels returns the names of models currently held in memory.
|
||||
func (c *Client) LoadedModels(ctx context.Context) ([]string, error) {
|
||||
details, err := c.LoadedModelDetails(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
models := make([]string, 0, len(details))
|
||||
for _, m := range details {
|
||||
models = append(models, m.Name)
|
||||
}
|
||||
return models, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user