ilker/llmgw
2
0 Fork 1
Code Issues Pull Requests Actions Packages Projects Releases 95 Wiki Activity
Files
a8bb417618d269f2c2f35d901a82ae2cfefda01e
llmgw/internal/agent/discovery/mig_test.go
T
Ilker ManapandClaude Opus 5.5 0fd55a862b MIG: a card in MIG mode is placed as its instances, each a GPU of class <card>-mig-<profile> 
Discovery reports a card in MIG mode as its instances instead of the whole
card, on both paths: nvidia-smi (-L names them, -q -x gives their memory,
ids and processes) and NVML (GetMigMode, GetMaxMigDeviceCount,
GetMigDeviceHandleByIndex, behind an interface so it is tested with a fake).
Each instance has its MIG UUID, its own memory, and a class from the card's
and the profile: a100-80g-mig-3g.40gb. Placement numbers instances
1000 + 100 x card + instance and every view shows them as card/migN.

A replica on MIG runs on exactly one instance: a process opens one MIG
device and instances have no NVLink between them, so tp, pp, gpus or nodes
above one on a MIG class are refused in catalog validation and again in
placement, and the planner offers one instance or none. The agent gives
the engine CUDA_VISIBLE_DEVICES=MIG-<uuid> (no Vulkan device variable on
MIG). Memory and processes are per instance; heat, power and wear the
card's; utilisation is left out rather than reported as zero. A card in MIG
mode with nothing carved stays out of placement with the reason, and
llmgw doctor reports each card's mode and instances and warns of the states
in which nothing can run. ctl nodes, ctl status, the console and the plan
count devices by class. docs/install.md has a MIG section.

Written against fakes of nvidia-smi and NVML output; not yet run on MIG
hardware.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-05 00:01:14 +02:00

419 lines
17 KiB
Go
Raw Blame History

package discovery
import (
"encoding/json"
"os"
"path/filepath"
"reflect"
"slices"
"strings"
"testing"
)
// Two A100 80GB cards: card 0 in MIG mode, cut into two 3g.40gb instances,
// card 1 whole. What nvidia-smi prints for that machine, by command. None of
// it is from real MIG hardware, which this was written without; the shapes
// follow NVIDIA's documentation and the output of drivers R470 to R550.
const (
migQuery = `0, GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77, NVIDIA A100-SXM4-80GB, 81920, 0, 8.0, 00000000:07:00.0, 550.54.15, 31, 0, [N/A]
1, GPU-1f2e3d4c-5b6a-7980-a1b2-c3d4e5f60718, NVIDIA A100-SXM4-80GB, 81920, 0, 8.0, 00000000:0F:00.0, 550.54.15, 33, 0, 0
`
migList = `GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77)
MIG 3g.40gb Device 0: (UUID: MIG-c6d4f1ef-42e4-5de3-91c7-45d71c87eb3f)
MIG 3g.40gb Device 1: (UUID: MIG-0f7c3a2e-9d1b-5a8e-b3c4-1e2f3a4b5c6d)
GPU 1: NVIDIA A100-SXM4-80GB (UUID: GPU-1f2e3d4c-5b6a-7980-a1b2-c3d4e5f60718)
`
migReport = `<?xml version="1.0" ?>
<nvidia_smi_log>
<driver_version>550.54.15</driver_version>
<gpu id="00000000:07:00.0">
<product_name>NVIDIA A100-SXM4-80GB</product_name>
<mig_mode><current_mig>Enabled</current_mig><pending_mig>Enabled</pending_mig></mig_mode>
<mig_devices>
<mig_device>
<index>1</index>
<gpu_instance_id>1</gpu_instance_id>
<compute_instance_id>0</compute_instance_id>
<device_attributes><shared><multiprocessor_count>42</multiprocessor_count></shared></device_attributes>
<fb_memory_usage><total>40192 MiB</total><reserved>0 MiB</reserved><used>21000 MiB</used><free>19192 MiB</free></fb_memory_usage>
</mig_device>
<mig_device>
<index>0</index>
<gpu_instance_id>2</gpu_instance_id>
<compute_instance_id>0</compute_instance_id>
<fb_memory_usage><total>40192 MiB</total><reserved>0 MiB</reserved><used>37 MiB</used><free>40155 MiB</free></fb_memory_usage>
</mig_device>
</mig_devices>
<fb_memory_usage><total>81920 MiB</total><used>21037 MiB</used></fb_memory_usage>
<processes>
<process_info>
<gpu_instance_id>1</gpu_instance_id>
<compute_instance_id>0</compute_instance_id>
<pid>4100</pid><type>C</type><process_name>vllm</process_name>
<used_memory>20900 MiB</used_memory>
</process_info>
</processes>
</gpu>
<gpu id="00000000:0F:00.0">
<product_name>NVIDIA A100-SXM4-80GB</product_name>
<mig_mode><current_mig>Disabled</current_mig><pending_mig>Disabled</pending_mig></mig_mode>
<mig_devices>None</mig_devices>
<fb_memory_usage><total>81920 MiB</total><used>300 MiB</used></fb_memory_usage>
<processes>
<process_info>
<gpu_instance_id>N/A</gpu_instance_id>
<compute_instance_id>N/A</compute_instance_id>
<pid>5200</pid><type>C</type><used_memory>250 MiB</used_memory>
</process_info>
</processes>
</gpu>
</nvidia_smi_log>`
migLGI = `+-------------------------------------------------------+
| GPU instances: |
| GPU Name Profile Instance Placement |
| ID ID Start:Size |
|=======================================================|
| 0 MIG 3g.40gb 9 1 4:4 |
+-------------------------------------------------------+
| 0 MIG 3g.40gb 9 2 0:4 |
+-------------------------------------------------------+
`
migLCI = `+--------------------------------------------------------------------+
| Compute instances: |
| GPU GPU Name Profile Instance Placement |
| Instance ID ID Start:Size |
| ID |
|====================================================================|
| 0 1 MIG 3g.40gb 2 0 0:3 |
+--------------------------------------------------------------------+
| 0 2 MIG 3g.40gb 2 0 0:3 |
+--------------------------------------------------------------------+
`
)
func migMachine(t *testing.T) []GPU {
t.Helper()
cards, err := parseSMI(migQuery)
if err != nil {
t.Fatal(err)
}
report, err := parseMIGReport([]byte(migReport))
if err != nil {
t.Fatal(err)
}
return expandMIG(cards, parseSMIList(migList), report)
}
func TestACardInMIGModeIsReportedAsItsInstancesInsteadOfItself(t *testing.T) {
gpus := migMachine(t)
var got []string
for _, g := range gpus {
got = append(got, DeviceLabel(g.Index)+" "+g.Class())
}
want := []string{"0/mig0 a100-80g-mig-3g.40gb", "0/mig1 a100-80g-mig-3g.40gb", "1 a100-80g"}
if !slices.Equal(got, want) {
t.Fatalf("devices = %q, want %q", got, want)
}
mig := gpus[1]
if mig.Index != 1001 || mig.UUID != "MIG-0f7c3a2e-9d1b-5a8e-b3c4-1e2f3a4b5c6d" {
t.Errorf("instance 1 is %d %s; an engine is given it by that UUID", mig.Index, mig.UUID)
}
if mig.MemoryTotalMB != 40192 || mig.MemoryUsedMB != 21000 {
t.Errorf("instance 1 memory %d/%d, want its own 21000/40192 rather than the card's",
mig.MemoryUsedMB, mig.MemoryTotalMB)
}
if m := mig.MIG; m == nil || m.ParentIndex != 0 || m.Device != 1 || m.Profile != "3g.40gb" ||
m.GPUInstance != 1 || m.ComputeInstance != 0 || m.ParentUUID != "GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77" {
t.Errorf("instance 1's parentage: %+v", mig.MIG)
}
if mig.ComputeCapability != "8.0" || mig.DriverVersion != "550.54.15" || mig.Health.TemperatureC != 31 {
t.Errorf("what the card and its instances share was not carried: %+v", mig)
}
// Twenty gigabytes in use on an instance is a foreign tenant there, and
// only there: instance 0 next to it is free.
if mig.Health.OK() || !gpus[0].Health.OK() {
t.Errorf("foreign memory judged per instance: 0 %v, 1 %v", gpus[0].Health.Problems, mig.Health.Problems)
}
// No driver reports utilisation per instance, and none is made up.
if mig.UtilizationPercent != 0 {
t.Errorf("an instance reported utilisation %d", mig.UtilizationPercent)
}
if gpus[2].IsMIG() || gpus[2].Index != 1 || gpus[2].CardIndex() != 1 || gpus[0].CardIndex() != 0 {
t.Errorf("the whole card beside it changed: %+v", gpus[2])
}
}
func TestAMIGInstanceSurvivesTheTripToTheControlPlaneWithItsClass(t *testing.T) {
g := migMachine(t)[0]
b, err := json.Marshal(g)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(b), `"class":"a100-80g-mig-3g.40gb"`) {
t.Errorf("the JSON carries no class for the instance: %s", b)
}
var back GPU
if err := json.Unmarshal(b, &back); err != nil {
t.Fatal(err)
}
if back.Class() != g.Class() || !reflect.DeepEqual(back.MIG, g.MIG) {
t.Errorf("came back as %s %+v", back.Class(), back.MIG)
}
}
func TestMIGClassesAreNamedFromTheCardAndTheProfile(t *testing.T) {
for _, c := range []struct {
name string
cardMB int
profile string
want string
}{
{"NVIDIA A100-SXM4-80GB MIG 3g.40gb", 81920, "3g.40gb", "a100-80g-mig-3g.40gb"},
// The same profile name on a 40 GB card is another class: its
// instance is the same size, but the card is not, and neither is
// the set of profiles that card can be cut to.
{"NVIDIA A100-SXM4-40GB MIG 3g.20gb", 40960, "3g.20gb", "a100-40g-mig-3g.20gb"},
{"NVIDIA H100 80GB HBM3 MIG 1g.10gb+me", 81559, "1g.10gb+me", "h100-80g-mig-1g.10gb-me"},
{"NVIDIA H100 80GB HBM3 MIG 7g.80gb", 81559, "7g.80gb", "h100-80g-mig-7g.80gb"},
} {
g := GPU{Name: c.name, MemoryTotalMB: 10000, MIG: &MIGDevice{ParentMemoryTotalMB: c.cardMB, Profile: c.profile}}
if got := g.Class(); got != c.want {
t.Errorf("%s -> %s, want %s", c.name, got, c.want)
}
if !IsMIGClass(g.Class()) {
t.Errorf("%s is not recognised as a MIG class", g.Class())
}
}
if IsMIGClass("a100-80g") || IsMIGClass("6000-24g") {
t.Error("a whole card's class was taken for a MIG one")
}
if MIGProfileMemoryGB("a100-80g-mig-3g.40gb") != 40 || MIGProfileMemoryGB("1g.10gb+me") != 10 || MIGProfileMemoryGB("a100-80g") != 0 {
t.Error("the instance's memory is misread from its profile")
}
}
func TestDeviceNumbersAreShownAsCardAndInstance(t *testing.T) {
for i, want := range map[int]string{0: "0", 7: "7", MIGIndex(0, 2): "0/mig2", MIGIndex(3, 0): "3/mig0", 1302: "3/mig2"} {
if got := DeviceLabel(i); got != want {
t.Errorf("DeviceLabel(%d) = %s, want %s", i, got, want)
}
}
if IsMIGIndex(15) || !IsMIGIndex(MIGIndex(0, 0)) {
t.Error("whole cards and instances are told apart wrongly")
}
}
func TestACardInMIGModeWithNoInstancesIsKeptOutOfPlacement(t *testing.T) {
cards, _ := parseSMI(migQuery)
report, _ := parseMIGReport([]byte(strings.Replace(migReport,
migReport[strings.Index(migReport, "<mig_devices>"):strings.Index(migReport, "</mig_devices>")+len("</mig_devices>")],
"<mig_devices>None</mig_devices>", 1)))
gpus := expandMIG(cards, parseSMIList("GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-x)\nGPU 1: NVIDIA A100-SXM4-80GB (UUID: GPU-y)\n"), report)
if len(gpus) != 2 || gpus[0].IsMIG() {
t.Fatalf("got %+v", gpus)
}
// CUDA will not open the whole card while MIG is on, and there is no
// instance to open instead: placing anything here fails at start.
if gpus[0].Health.OK() || !strings.Contains(gpus[0].Health.Problems[0], "no instances") {
t.Errorf("card 0's problems: %v", gpus[0].Health.Problems)
}
if !gpus[1].Health.OK() {
t.Errorf("card 1, not in MIG mode, was flagged: %v", gpus[1].Health.Problems)
}
// Instances the report sees but -L could not name cannot be given to an
// engine either; the card says why.
report, _ = parseMIGReport([]byte(migReport))
gpus = expandMIG(cards, nil, report)
if len(gpus) != 2 || gpus[0].Health.OK() || !strings.Contains(gpus[0].Health.Problems[0], "did not name") {
t.Errorf("with -L failing: %+v", gpus)
}
}
func TestInstancesAreFoundFromTheListAloneWhenTheReportCannotBeRead(t *testing.T) {
cards, _ := parseSMI(migQuery)
gpus := expandMIG(cards, parseSMIList(migList), nil)
if len(gpus) != 3 || !gpus[0].IsMIG() || gpus[0].MemoryTotalMB != 40*1024 {
t.Errorf("without the report: %+v", gpus)
}
}
func TestAnOlderDriversInstanceUUIDIsTakenAsItIs(t *testing.T) {
// R450 and R460 name an instance by its card and ids, and CUDA of that
// era takes exactly that in CUDA_VISIBLE_DEVICES.
l := parseSMIList("GPU 0: A100-SXM4-40GB (UUID: GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77)\n" +
" MIG 1g.5gb Device 0: (UUID: MIG-GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77/7/0)\n")
if len(l[0]) != 1 || l[0][0].UUID != "MIG-GPU-5d5ba0d6-d33d-2b2c-524d-9e3d8d2b8a77/7/0" || l[0][0].Profile != "1g.5gb" {
t.Errorf("parsed %+v", l)
}
}
func TestProcessesOnAMIGCardAreChargedToTheirOwnInstance(t *testing.T) {
mem, err := parseGPUMemory([]byte(migReport))
if err != nil {
t.Fatal(err)
}
want := []GPUMemory{
{Index: 1000, UsedMB: 37},
{Index: 1001, UsedMB: 21000, Procs: []GPUProcess{{PID: 4100, UsedMB: 20900}}},
{Index: 1, UsedMB: 300, Procs: []GPUProcess{{PID: 5200, UsedMB: 250}}},
}
if !reflect.DeepEqual(mem, want) {
t.Errorf("got %+v\nwant %+v", mem, want)
}
// In the order discovery reports the devices, so a heartbeat's readings
// line up with the inventory one for one.
gpus := migMachine(t)
for i := range gpus {
if mem[i].Index != gpus[i].Index {
t.Errorf("reading %d is for device %d, inventory has %d", i, mem[i].Index, gpus[i].Index)
}
}
procs, _ := parseGPUProcesses([]byte(migReport))
if !reflect.DeepEqual(procs, map[int][]int{1001: {4100}, 1: {5200}}) {
t.Errorf("processes by device: %v", procs)
}
}
func TestAMIGInstanceTakesItsCardsPlaceInTheTopologyButNoNVLink(t *testing.T) {
gpus := migMachine(t)
topo := Topology{
NVLink: map[int]map[int]int{0: {1: 12}, 1: {0: 12}},
NUMA: map[int]int{0: 0, 1: 1},
NICPath: map[int]map[string]string{0: {"mlx5_0": "PIX"}},
}
topo.Apply(gpus)
if len(gpus[0].NVLinkPeers) != 0 || len(gpus[1].NVLinkPeers) != 0 {
t.Errorf("MIG instances were given NVLink peers: %v %v", gpus[0].NVLinkPeers, gpus[1].NVLinkPeers)
}
if gpus[1].NUMANode != 0 || gpus[1].NICPaths["mlx5_0"] != "PIX" {
t.Errorf("instance 1 did not take its card's NUMA node and adapters: %+v", gpus[1])
}
if gpus[2].NUMANode != 1 || !slices.Equal(gpus[2].NVLinkPeers, []int{0}) {
t.Errorf("the whole card's topology: %+v", gpus[2])
}
}
func TestMIGInstancesAreDiscoveredThroughNvidiaSMI(t *testing.T) {
// The real code path against a stub binary that answers each command as
// the machine above would.
dir := t.TempDir()
stub := filepath.Join(dir, "nvidia-smi")
script := "#!/bin/sh\ncase \"$1\" in\n" +
"-L) cat <<'EOF'\n" + migList + "EOF\n;;\n" +
"-q) cat <<'EOF'\n" + migReport + "\nEOF\n;;\n" +
"*) cat <<'EOF'\n" + migQuery + "EOF\n;;\nesac\n"
if err := os.WriteFile(stub, []byte(script), 0o755); err != nil {
t.Fatal(err)
}
gpus, err := DiscoverGPUsSMI(t.Context(), stub)
if err != nil {
t.Fatal(err)
}
inv := Inventory{GPUs: gpus}
if !slices.Equal(inv.Classes(), []string{"a100-80g", "a100-80g-mig-3g.40gb"}) || len(gpus) != 3 {
t.Errorf("classes %v from %d devices", inv.Classes(), len(gpus))
}
}
// fakeMIGCard is NVML's view of one card, as nvmlDeviceGetMigMode,
// nvmlDeviceGetMaxMigDeviceCount and nvmlDeviceGetMigDeviceHandleByIndex
// would give it.
type fakeMIGCard struct {
mode, pending, supported bool
max int
devs map[int]fakeMIGInstance
}
func (c fakeMIGCard) MIGMode() (bool, bool, bool) { return c.mode, c.pending, c.supported }
func (c fakeMIGCard) MaxMIGDevices() int { return c.max }
func (c fakeMIGCard) MIGDevice(i int) (nvmlMIGInstance, bool) {
d, ok := c.devs[i]
return d, ok
}
type fakeMIGInstance struct {
uuid, name string
total, used int
gi, ci int
slices, attrMem int
}
func (m fakeMIGInstance) UUID() string { return m.uuid }
func (m fakeMIGInstance) Name() string { return m.name }
func (m fakeMIGInstance) Memory() (int, int, bool) {
return m.total, m.used, m.total > 0
}
func (m fakeMIGInstance) Instances() (int, int) { return m.gi, m.ci }
func (m fakeMIGInstance) Attributes() (int, int, bool) {
return m.slices, m.attrMem, m.slices > 0
}
func TestNVMLReportsACardInMIGModeAsItsInstances(t *testing.T) {
card := GPU{Index: 2, UUID: "GPU-card", Name: "NVIDIA H100 80GB HBM3", MemoryTotalMB: 81559,
ComputeCapability: "9.0", NUMANode: 1, Health: GPUHealth{TemperatureC: 40}}
// Seven slots, three filled, one of them on a driver that does not put
// the profile in the device's name.
dev := fakeMIGCard{mode: true, pending: true, supported: true, max: 7, devs: map[int]fakeMIGInstance{
0: {uuid: "MIG-aaa", name: "NVIDIA H100 80GB HBM3 MIG 3g.40gb", total: 40448, used: 10, gi: 2, ci: 0},
1: {uuid: "MIG-bbb", name: "NVIDIA H100 80GB HBM3 MIG 3g.40gb", total: 40448, used: 10, gi: 1, ci: 0},
4: {uuid: "MIG-ccc", name: "NVIDIA H100 80GB HBM3", total: 9984, gi: 9, ci: 0, slices: 1, attrMem: 9984},
}}
gpus, isMIG := migFromNVML(card, dev)
if !isMIG || len(gpus) != 3 {
t.Fatalf("got %d devices, MIG %v", len(gpus), isMIG)
}
var classes []string
for _, g := range gpus {
classes = append(classes, DeviceLabel(g.Index)+" "+g.UUID+" "+g.Class())
}
want := []string{"2/mig0 MIG-aaa h100-80g-mig-3g.40gb", "2/mig1 MIG-bbb h100-80g-mig-3g.40gb",
"2/mig4 MIG-ccc h100-80g-mig-1g.10gb"}
if !slices.Equal(classes, want) {
t.Errorf("devices %q, want %q", classes, want)
}
if g := gpus[1]; g.MemoryTotalMB != 40448 || g.MIG.GPUInstance != 1 || g.NUMANode != 1 ||
g.Health.TemperatureC != 40 || len(g.NVLinkPeers) != 0 {
t.Errorf("instance 1: %+v", g)
}
// MIG on and nothing carved: the card, kept out of placement.
gpus, isMIG = migFromNVML(card, fakeMIGCard{mode: true, supported: true, max: 7})
if !isMIG || len(gpus) != 1 || gpus[0].IsMIG() || gpus[0].Health.OK() {
t.Errorf("an empty MIG card: %+v", gpus)
}
// MIG off, or a card without it: the card, as it always was.
for _, d := range []fakeMIGCard{{supported: true, max: 7}, {}} {
gpus, isMIG = migFromNVML(card, d)
if isMIG || len(gpus) != 1 || !reflect.DeepEqual(gpus[0], card) {
t.Errorf("a card not in MIG mode came back as %+v", gpus)
}
}
}
func TestMIGInstanceTablesAreRead(t *testing.T) {
gis := ParseMIGInstances(migLGI)
if len(gis) != 2 || gis[0] != (MIGInstance{GPU: 0, Profile: "3g.40gb", ProfileID: 9, GPUInstance: 1,
ComputeInstance: -1, Placement: "4:4"}) {
t.Errorf("-lgi: %+v", gis)
}
cis := ParseMIGInstances(migLCI)
if len(cis) != 2 || cis[1] != (MIGInstance{GPU: 0, GPUInstance: 2, Profile: "3g.40gb", ProfileID: 2,
ComputeInstance: 0, Placement: "0:3"}) {
t.Errorf("-lci: %+v", cis)
}
if got := ParseMIGInstances("No GPU instances found: Not Found\n"); len(got) != 0 {
t.Errorf("nothing to list read as %+v", got)
}
cards, err := buildMIGCards([]byte(migReport), migList, migLGI, migLCI)
if err != nil {
t.Fatal(err)
}
if len(cards) != 2 || !cards[0].Enabled || len(cards[0].Instances) != 2 || len(cards[0].GPUInstances) != 2 ||
len(cards[0].ComputeInstances) != 2 || cards[1].Enabled {
t.Errorf("cards: %+v", cards)
}
}
Reference in new issue View Git Blame Copy Permalink