vram.go

 1package services
 2
 3import (
 4	"fmt"
 5	"os"
 6	"path/filepath"
 7	"strconv"
 8	"strings"
 9)
10
11const vramTotalGlob = "/sys/class/drm/card*/device/mem_info_vram_total"
12
13// checkVRAM sums free VRAM across all AMD cards from sysfs and requires at
14// least minFreeVRAMBytes free. It fails closed: no readable AMD VRAM sysfs is
15// an error (only AMD GPUs are supported).
16func checkVRAM() error {
17	total, used, err := readVRAM()
18	if err != nil {
19		return fmt.Errorf("services: %w", err)
20	}
21	free := total - used
22	if free < minFreeVRAMBytes {
23		return fmt.Errorf("services: insufficient free VRAM to start models: need %d bytes, have %d bytes", minFreeVRAMBytes, free)
24	}
25	return nil
26}
27
28func readVRAM() (int64, int64, error) {
29	totals, err := filepath.Glob(vramTotalGlob)
30	if err != nil {
31		return 0, 0, fmt.Errorf("glob %s: %w", vramTotalGlob, err)
32	}
33	if len(totals) == 0 {
34		return 0, 0, fmt.Errorf("no AMD GPU VRAM sysfs found at %s (AMD only)", vramTotalGlob)
35	}
36	var total, used int64
37	for _, tf := range totals {
38		t, err := readMemInfoBytes(tf)
39		if err != nil {
40			return 0, 0, err
41		}
42		uf := strings.Replace(tf, "mem_info_vram_total", "mem_info_vram_used", 1)
43		u, err := readMemInfoBytes(uf)
44		if err != nil {
45			return 0, 0, err
46		}
47		total += t
48		used += u
49	}
50	return total, used, nil
51}
52
53func readMemInfoBytes(path string) (int64, error) {
54	raw, err := os.ReadFile(path)
55	if err != nil {
56		return 0, fmt.Errorf("read %s: %w", path, err)
57	}
58	fields := strings.Fields(string(raw))
59	if len(fields) == 0 {
60		return 0, fmt.Errorf("%s is empty", path)
61	}
62	n, err := strconv.ParseInt(fields[0], 10, 64)
63	if err != nil {
64		return 0, fmt.Errorf("parse %s: %w", path, err)
65	}
66	return n, nil
67}