Use nv-gpu-info on Windows to get VRAM size

Signed-off-by: Piotr Stankiewicz <piotr.stankiewicz@docker.com>

Author: Piotr Stankiewicz

main.go

@@ -90,7 +90,7 @@ func main() {
 		log.Fatalf("unable to initialize %s backend: %v", llamacpp.Name, err)
 	}
 
-	gpuInfo := gpuinfo.New()
+	gpuInfo := gpuinfo.New(llamaServerPath)
 
 	scheduler := scheduling.NewScheduler(
 		log,

pkg/gpuinfo/gpuinfo.go

@@ -1,11 +1,17 @@
 package gpuinfo
 
-type GPUInfo struct{}
+type GPUInfo struct {
+	// modelRuntimeInstallPath is the location where DMR installed it's llama-server
+	// and accompanying tools
+	modelRuntimeInstallPath string
+}
 
-func New() *GPUInfo {
-	return &GPUInfo{}
+func New(modelRuntimeInstallPath string) *GPUInfo {
+	return &GPUInfo{
+		modelRuntimeInstallPath: modelRuntimeInstallPath,
+	}
 }
 
 func (g *GPUInfo) GetVRAMSize() (uint64, error) {
-	return getVRAMSize()
+	return getVRAMSize(g.modelRuntimeInstallPath)
 }

pkg/gpuinfo/memory_darwin.go

@@ -8,7 +8,7 @@ import "C"
 import "errors"
 
 // getVRAMSize returns total system GPU memory in bytes
-func getVRAMSize() (uint64, error) {
+func getVRAMSize(_ string) (uint64, error) {
 	vramSize := C.getVRAMSize()
 	if vramSize == 0 {
 		return 0, errors.New("could not get metal VRAM size")

pkg/gpuinfo/memory_linux.go

@@ -8,7 +8,7 @@ import "C"
 import "errors"
 
 // getVRAMSize returns total system GPU memory in bytes
-func getVRAMSize() (uint64, error) {
+func getVRAMSize(_ string) (uint64, error) {
 	vramSize := C.getVRAMSize()
 	if vramSize == 0 {
 		return 0, errors.New("could not get nvidia VRAM size")

pkg/gpuinfo/memory_windows.go

@@ -1,16 +1,31 @@
 package gpuinfo
 
-/*
-#include "nvapi.h"
-*/
-import "C"
-import "errors"
+import (
+	"bufio"
+	"context"
+	"errors"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+)
 
 // getVRAMSize returns total system GPU memory in bytes
-func getVRAMSize() (uint64, error) {
-	vramSize := C.getVRAMSize()
-	if vramSize == 0 {
-		return 0, errors.New("could not get nvapi VRAM size")
+func getVRAMSize(ctx context.Context, modelRuntimeInstallPath string) (uint64, error) {
+	nvGPUInfoBin := filepath.Join(modelRuntimeInstallPath, "com.docker.nv-gpu-info.exe")
+
+	cmd := exec.CommandContext(ctx, nvGPUInfoBin)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return 0, err
+	}
+	sc := bufio.NewScanner(strings.NewReader(string(out)))
+	for sc.Scan() {
+		vram, found := strings.CutPrefix(sc.Text(), "GPU[0]: dedicated memory:")
+		if found {
+			vram = strings.TrimSpace(vram)
+			return strconv.ParseUint(vram, 10, 64)
+		}
 	}
-	return uint64(vramSize), nil
+	return 0, errors.New("unexpected nv-gpu-info output format")
 }