llmgw/internal/agent/discovery/gpu_fabric.go

package discovery

import (
	"context"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// GPUFabric is a GPU's place in a multi-node NVLink domain, as the driver
// reports it: the NVL72 rack, or the NVLink Switch system, it belongs to.
type GPUFabric struct {
	// State is Completed once the fabric manager has joined the GPU to its
	// domain; In Progress, Not Started or Not Supported otherwise.
	State string `json:"state"`
	// Status is Success, or the reason the join failed.
	Status string `json:"status,omitempty"`
	// ClusterUUID and CliqueID name the domain: GPUs sharing both can reach
	// each other over NVLink whatever node they are in.
	ClusterUUID string `json:"cluster_uuid,omitempty"`
	CliqueID    string `json:"clique_id,omitempty"`
	// Bandwidth is Full, or how the domain's NVLink is degraded.
	Bandwidth string `json:"bandwidth,omitempty"`
}

// Domain is the NVLink domain the GPU is joined to, empty when it is in none
// that crosses nodes: a card with no fabric, one not joined, or one whose
// cluster is all zeros, which is what a single machine's NVSwitch reports.
func (f *GPUFabric) Domain() string {
	if f == nil || !strings.EqualFold(f.State, "Completed") {
		return ""
	}
	u := strings.ToLower(strings.TrimSpace(f.ClusterUUID))
	if u == "" || strings.Trim(u, "0-") == "" {
		return ""
	}
	return u + "/" + f.CliqueID
}

// ParseGPUFabric reads the Fabric sections of `nvidia-smi -q`, by the PCI
// address each GPU's section starts with.
func ParseGPUFabric(out string) map[string]GPUFabric {
	res := map[string]GPUFabric{}
	var bus string
	fabricIndent := -1
	var cur *GPUFabric
	flush := func() {
		if bus != "" && cur != nil && cur.State != "" {
			res[bus] = *cur
		}
		cur, fabricIndent = nil, -1
	}
	for _, line := range strings.Split(out, "\n") {
		trim := strings.TrimSpace(line)
		indent := len(line) - len(strings.TrimLeft(line, " \t"))
		if strings.HasPrefix(line, "GPU ") && indent == 0 {
			flush()
			bus = strings.ToLower(strings.TrimSpace(strings.TrimPrefix(line, "GPU ")))
			continue
		}
		if trim == "Fabric" {
			flush()
			cur, fabricIndent = &GPUFabric{}, indent
			continue
		}
		if cur == nil {
			continue
		}
		if trim != "" && indent <= fabricIndent {
			flush()
			continue
		}
		k, v, ok := strings.Cut(trim, ":")
		if !ok {
			continue
		}
		k, v = strings.TrimSpace(k), strings.TrimSpace(v)
		if strings.EqualFold(v, "N/A") {
			v = ""
		}
		switch k {
		case "State":
			cur.State = v
		case "Status":
			cur.Status = v
		case "ClusterUUID", "Cluster UUID":
			cur.ClusterUUID = v
		case "CliqueId", "Clique ID":
			cur.CliqueID = v
		case "Bandwidth":
			cur.Bandwidth = v
		}
	}
	flush()
	return res
}

// DiscoverGPUFabric runs `nvidia-smi -q` for the Fabric sections. Nothing
// is reported on a driver or a card without them.
func DiscoverGPUFabric(ctx context.Context, binary string) (map[string]GPUFabric, error) {
	if binary == "" {
		if _, err := exec.LookPath("nvidia-smi"); err != nil {
			return nil, nil
		}
		binary = "nvidia-smi"
	}
	ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
	defer cancel()
	out, err := exec.CommandContext(ctx, binary, "-q").Output()
	if err != nil {
		return nil, err
	}
	return ParseGPUFabric(string(out)), nil
}

// ApplyFabric puts each GPU's fabric on it, matched by PCI address.
func ApplyFabric(gpus []GPU, fab map[string]GPUFabric) {
	for i := range gpus {
		if f, ok := fab[strings.ToLower(gpus[i].PCIBusID)]; ok {
			f := f
			gpus[i].Fabric = &f
		}
	}
}

// IMEX is the service that lets GPUs in different nodes of one NVLink
// domain share memory: without it, multi-node NVLink carries nothing.
type IMEX struct {
	// State is systemd's word for nvidia-imex: active, inactive, or empty
	// where it is not installed.
	State string `json:"state,omitempty"`
	// Channels is how many IMEX channels exist; a process needs one.
	Channels int `json:"channels"`
}

// DiscoverIMEX reads the IMEX service and its channels.
func DiscoverIMEX(ctx context.Context) IMEX {
	var m IMEX
	if entries, err := os.ReadDir("/dev/nvidia-caps-imex-channels"); err == nil {
		m.Channels = len(entries)
	}
	if _, err := exec.LookPath("systemctl"); err == nil {
		ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
		defer cancel()
		out, _ := exec.CommandContext(ctx, "systemctl", "is-active", "nvidia-imex").Output()
		m.State = strings.TrimSpace(string(out))
		if m.State == "inactive" {
			// "inactive" is also what systemd says of a unit that does not
			// exist; the binary says which.
			if _, err := exec.LookPath("nvidia-imex"); err != nil {
				if _, err := os.Stat(filepath.Join("/usr/bin", "nvidia-imex")); err != nil {
					m.State = ""
				}
			}
		}
	}
	return m
}