fn_registry/functions/infra/get_gpu_info.go

package infra

import (
	"encoding/csv"
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

// GetGpuInfo queries NVIDIA GPUs via nvidia-smi and returns a slice of GpuInfo.
// If nvidia-smi is not installed or no NVIDIA GPU is present, returns an empty
// slice and a nil error (absence of NVIDIA hardware is not an error).
func GetGpuInfo() ([]GpuInfo, error) {
	out, err := exec.Command(
		"nvidia-smi",
		"--query-gpu=index,name,memory.total,memory.free,driver_version,cuda_version",
		"--format=csv,noheader,nounits",
	).Output()

	if err != nil {
		// nvidia-smi not installed or no NVIDIA device — not an error.
		var exitErr *exec.ExitError
		if errors.Is(err, exec.ErrNotFound) || errors.As(err, &exitErr) {
			return []GpuInfo{}, nil
		}
		return nil, fmt.Errorf("gpu_info: nvidia-smi: %w", err)
	}

	r := csv.NewReader(strings.NewReader(strings.TrimSpace(string(out))))
	r.TrimLeadingSpace = true

	records, err := r.ReadAll()
	if err != nil {
		return nil, fmt.Errorf("gpu_info: parse csv: %w", err)
	}

	gpus := make([]GpuInfo, 0, len(records))
	for _, rec := range records {
		if len(rec) < 6 {
			continue
		}

		idx, _ := strconv.Atoi(strings.TrimSpace(rec[0]))
		totalMb, _ := strconv.Atoi(strings.TrimSpace(rec[2]))
		freeMb, _ := strconv.Atoi(strings.TrimSpace(rec[3]))

		gpus = append(gpus, GpuInfo{
			Index:         idx,
			Name:          strings.TrimSpace(rec[1]),
			VramTotalMb:   totalMb,
			VramFreeMb:    freeMb,
			DriverVersion: strings.TrimSpace(rec[4]),
			CudaVersion:   strings.TrimSpace(rec[5]),
		})
	}

	return gpus, nil
}