bee/audit/internal/platform/nvidia_settings.go

package platform

import (
	"fmt"
	"strconv"
	"strings"
)

// NvidiaGPUSetting reports the persistent GPU settings that nvidia-smi
// exposes: ECC mode, MIG mode, and Confidential Computing mode all persist
// via the GPU's inforom/firmware and survive a reboot, but only take effect
// after a GPU reset (nvidia-smi -r) or host reboot (hence separate
// current/pending fields for ECC and MIG). Power limit does NOT persist
// across reboot/driver reload — it is re-applied here on demand, same as
// nvidia-smi's own behavior.
type NvidiaGPUSetting struct {
	Index              int     `json:"index"`
	Name               string  `json:"name"`
	ECCCurrent         string  `json:"ecc_current"` // "Enabled" / "Disabled" / "N/A"
	ECCPending         string  `json:"ecc_pending"` // differs from current until a GPU reset/reboot
	MIGCurrent         string  `json:"mig_current"` // "Enabled" / "Disabled" / "N/A" (N/A = not MIG-capable)
	MIGPending         string  `json:"mig_pending"`
	CCState            string  `json:"cc_state"` // "ON" / "OFF" / "" (unsupported or query failed)
	PowerLimitW        float64 `json:"power_limit_w"`
	PowerMinLimitW     float64 `json:"power_min_limit_w"`
	PowerMaxLimitW     float64 `json:"power_max_limit_w"`
	PowerDefaultLimitW float64 `json:"power_default_limit_w"`
}

// ListNvidiaGPUSettings reports current ECC, MIG, Confidential Computing and
// power-limit configuration per GPU.
func (s *System) ListNvidiaGPUSettings() ([]NvidiaGPUSetting, error) {
	out, err := satExecCommand("nvidia-smi",
		"--query-gpu=index,name,ecc.mode.current,ecc.mode.pending,mig.mode.current,mig.mode.pending,power.limit,power.min_limit,power.max_limit,power.default_limit",
		"--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, fmt.Errorf("nvidia-smi: %w", err)
	}
	var settings []NvidiaGPUSetting
	for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		parts := strings.Split(line, ",")
		if len(parts) < 10 {
			continue
		}
		idx, err := strconv.Atoi(strings.TrimSpace(parts[0]))
		if err != nil {
			continue
		}
		g := NvidiaGPUSetting{
			Index:              idx,
			Name:               strings.TrimSpace(parts[1]),
			ECCCurrent:         strings.TrimSpace(parts[2]),
			ECCPending:         strings.TrimSpace(parts[3]),
			MIGCurrent:         strings.TrimSpace(parts[4]),
			MIGPending:         strings.TrimSpace(parts[5]),
			PowerLimitW:        parseNvidiaFloat(parts[6]),
			PowerMinLimitW:     parseNvidiaFloat(parts[7]),
			PowerMaxLimitW:     parseNvidiaFloat(parts[8]),
			PowerDefaultLimitW: parseNvidiaFloat(parts[9]),
		}
		g.CCState = s.queryNvidiaGPUCCState(idx)
		settings = append(settings, g)
	}
	return settings, nil
}

// queryNvidiaGPUCCState reads the per-GPU Confidential Computing state via
// `nvidia-smi conf-compute -i <index> -q`, reusing the same field-name
// convention as RunConfidentialComputingCheckPack's parseConfComputeFields.
// Returns "" if the driver/GPU doesn't support conf-compute at all.
func (s *System) queryNvidiaGPUCCState(index int) string {
	out, err := satExecCommand("nvidia-smi", "conf-compute", "-i", strconv.Itoa(index), "-q").Output()
	if err != nil {
		return ""
	}
	return parseConfComputeFields(out)["CC State"]
}

func parseNvidiaFloat(v string) float64 {
	f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
	if err != nil {
		return 0
	}
	return f
}

// SetNvidiaGPUECC sets pending ECC mode for one GPU. Takes effect after the
// GPU is reset (nvidia-smi -r) or the host is rebooted.
func (s *System) SetNvidiaGPUECC(index int, enabled bool) (string, error) {
	val := "0"
	if enabled {
		val = "1"
	}
	out, err := satExecCommand("nvidia-smi", "-i", strconv.Itoa(index), "-e", val).CombinedOutput()
	return string(out), err
}

// SetNvidiaGPUMIG sets pending MIG (Multi-Instance GPU) mode for one GPU.
// Persists via inforom but takes effect only after a GPU reset (nvidia-smi
// -r) or host reboot. Only supported on MIG-capable GPUs (Ampere/Hopper
// data-center parts); fails harmlessly on others.
func (s *System) SetNvidiaGPUMIG(index int, enabled bool) (string, error) {
	val := "0"
	if enabled {
		val = "1"
	}
	out, err := satExecCommand("nvidia-smi", "-i", strconv.Itoa(index), "-mig", val).CombinedOutput()
	return string(out), err
}

// SetNvidiaGPUCCMode sets Confidential Computing mode for one GPU
// (nvidia-smi conf-compute -scc). Persists via inforom, reversible, but
// takes effect only after a GPU reset (nvidia-smi -r) or host reboot.
func (s *System) SetNvidiaGPUCCMode(index int, enabled bool) (string, error) {
	val := "0"
	if enabled {
		val = "1"
	}
	out, err := satExecCommand("nvidia-smi", "conf-compute", "-i", strconv.Itoa(index), "-scc", val).CombinedOutput()
	return string(out), err
}

// SetNvidiaGPUPowerLimit sets the power limit (watts) for one GPU. Does not
// persist across reboot/driver reload — callers that want a durable change
// must re-apply it (e.g. at boot), same as nvidia-smi's own semantics.
func (s *System) SetNvidiaGPUPowerLimit(index int, watts float64) (string, error) {
	out, err := satExecCommand("nvidia-smi", "-i", strconv.Itoa(index), "-pl", strconv.FormatFloat(watts, 'f', 0, 64)).CombinedOutput()
	return string(out), err
}

// ResetNvidiaGPUDefaults restores factory-default settings on every visible
// NVIDIA GPU that has drifted from them: power limit back to
// power.default_limit (as reported by nvidia-smi), ECC re-enabled (NVIDIA
// data-center GPUs ship with ECC on; nvidia-smi has no queryable "ECC
// default" field, so this is the documented factory default, not a value
// read back from hardware), and MIG/Confidential-Computing mode disabled
// (both ship off from the factory). GPUs already at defaults are left
// alone.
func (s *System) ResetNvidiaGPUDefaults() (string, error) {
	settings, err := s.ListNvidiaGPUSettings()
	if err != nil {
		return "", err
	}
	var out strings.Builder
	for _, g := range settings {
		if g.PowerDefaultLimitW > 0 && g.PowerLimitW != g.PowerDefaultLimitW {
			res, err := s.SetNvidiaGPUPowerLimit(g.Index, g.PowerDefaultLimitW)
			out.WriteString(res)
			if err != nil {
				fmt.Fprintf(&out, "GPU %d: power limit reset failed: %v\n", g.Index, err)
			}
		}
		if strings.EqualFold(g.ECCCurrent, "disabled") {
			res, err := s.SetNvidiaGPUECC(g.Index, true)
			out.WriteString(res)
			if err != nil {
				fmt.Fprintf(&out, "GPU %d: ECC re-enable failed: %v\n", g.Index, err)
			}
		}
		if strings.EqualFold(g.MIGCurrent, "enabled") {
			res, err := s.SetNvidiaGPUMIG(g.Index, false)
			out.WriteString(res)
			if err != nil {
				fmt.Fprintf(&out, "GPU %d: MIG disable failed: %v\n", g.Index, err)
			}
		}
		if strings.EqualFold(g.CCState, "on") {
			res, err := s.SetNvidiaGPUCCMode(g.Index, false)
			out.WriteString(res)
			if err != nil {
				fmt.Fprintf(&out, "GPU %d: CC mode disable failed: %v\n", g.Index, err)
			}
		}
	}
	if out.Len() == 0 {
		out.WriteString("All GPUs already at factory-default settings.\n")
	}
	return out.String(), nil
}