127 lines
3.9 KiB
Go
127 lines
3.9 KiB
Go
package platform
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
// satReadFile is a seam for tests to fake sysfs reads (numa_node files).
|
|
var satReadFile = os.ReadFile
|
|
|
|
// gpuBandwidthSocketGroups splits gpuIndices into NUMA-locality groups,
|
|
// ordered by ascending Linux NUMA node ID, for RunNvidiaBandwidthPack.
|
|
//
|
|
// Falls back to a single group containing all of gpuIndices — i.e. no split
|
|
// — whenever the NUMA node can't be resolved for every GPU, or all resolve
|
|
// to the same node: there's nothing meaningful to split in that case.
|
|
func gpuBandwidthSocketGroups(gpuIndices []int, logFunc func(string)) [][]int {
|
|
nodes, err := gpuNUMANodes(gpuIndices)
|
|
if err != nil {
|
|
if logFunc != nil {
|
|
logFunc(fmt.Sprintf("nvbandwidth: could not resolve GPU NUMA nodes (%v); running all GPUs as one group", err))
|
|
}
|
|
return [][]int{gpuIndices}
|
|
}
|
|
|
|
byNode := map[int][]int{}
|
|
for _, idx := range gpuIndices {
|
|
node, ok := nodes[idx]
|
|
if !ok {
|
|
if logFunc != nil {
|
|
logFunc(fmt.Sprintf("nvbandwidth: no NUMA node resolved for GPU %d; running all GPUs as one group", idx))
|
|
}
|
|
return [][]int{gpuIndices}
|
|
}
|
|
byNode[node] = append(byNode[node], idx)
|
|
}
|
|
// Fewer than two NUMA nodes means there is nothing meaningful to split.
|
|
if len(byNode) < 2 {
|
|
return [][]int{gpuIndices}
|
|
}
|
|
|
|
sortedNodes := make([]int, 0, len(byNode))
|
|
for node := range byNode {
|
|
sortedNodes = append(sortedNodes, node)
|
|
}
|
|
sort.Ints(sortedNodes)
|
|
|
|
groups := make([][]int, 0, len(sortedNodes))
|
|
for _, node := range sortedNodes {
|
|
groups = append(groups, dedupeSortedIndices(byNode[node]))
|
|
}
|
|
return groups
|
|
}
|
|
|
|
// gpuNUMANodes resolves the NUMA node each of gpuIndices' GPU is attached to,
|
|
// via nvidia-smi's PCI bus ID and the device's sysfs numa_node attribute.
|
|
// A GPU missing from the returned map means its node couldn't be resolved.
|
|
func gpuNUMANodes(gpuIndices []int) (map[int]int, error) {
|
|
out, err := satExecCommand("nvidia-smi", "--query-gpu=index,pci.bus_id", "--format=csv,noheader,nounits").Output()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("nvidia-smi: %w", err)
|
|
}
|
|
|
|
want := make(map[int]struct{}, len(gpuIndices))
|
|
for _, idx := range gpuIndices {
|
|
want[idx] = struct{}{}
|
|
}
|
|
|
|
nodes := make(map[int]int, len(gpuIndices))
|
|
for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
|
|
fields := strings.SplitN(line, ",", 2)
|
|
if len(fields) != 2 {
|
|
continue
|
|
}
|
|
idx, err := strconv.Atoi(strings.TrimSpace(fields[0]))
|
|
if err != nil {
|
|
continue
|
|
}
|
|
if _, ok := want[idx]; !ok {
|
|
continue
|
|
}
|
|
bdf := normalizeNvidiaBDF(strings.TrimSpace(fields[1]))
|
|
if node, ok := readPCINumaNode(bdf); ok {
|
|
nodes[idx] = node
|
|
}
|
|
}
|
|
return nodes, nil
|
|
}
|
|
|
|
// normalizeNvidiaBDF converts nvidia-smi's PCI bus ID to the exact form the
|
|
// sysfs paths under /sys/bus/pci/devices use: an 8-hex-digit domain is
|
|
// narrowed to 4 digits ("00000000:05:00.0" -> "0000:05:00.0"), and the hex
|
|
// is lower-cased ("0000:CB:00.0" -> "0000:cb:00.0"). nvidia-smi upper-cases
|
|
// the bus/device hex; sysfs directory names are always lower-case, so
|
|
// without this a BDF containing a hex letter (e.g. GPUs on bus 4b/cb/cf)
|
|
// would never match a sysfs entry and every numa_node / link-speed read
|
|
// would silently fail.
|
|
func normalizeNvidiaBDF(busID string) string {
|
|
busID = strings.ToLower(strings.TrimSpace(busID))
|
|
domain, rest, ok := strings.Cut(busID, ":")
|
|
if !ok {
|
|
return busID
|
|
}
|
|
if len(domain) > 4 {
|
|
domain = domain[len(domain)-4:]
|
|
}
|
|
return domain + ":" + rest
|
|
}
|
|
|
|
// readPCINumaNode reads a PCI device's NUMA affinity from sysfs. Returns
|
|
// false if the attribute is missing/unreadable or reports -1 (no affinity —
|
|
// common on single-socket or non-NUMA systems).
|
|
func readPCINumaNode(bdf string) (int, bool) {
|
|
data, err := satReadFile("/sys/bus/pci/devices/" + bdf + "/numa_node")
|
|
if err != nil {
|
|
return 0, false
|
|
}
|
|
node, err := strconv.Atoi(strings.TrimSpace(string(data)))
|
|
if err != nil || node < 0 {
|
|
return 0, false
|
|
}
|
|
return node, true
|
|
}
|