refactor: Support user supplied SM count for NVIDIA GPUs

pkg/collector/gpu.go

@@ -43,6 +43,16 @@ var (
 	).Hidden().Default("").String()
 )
 
+// User facing opts.
+var (
+	nvidiaGPUSMCount = CEEMSExporterApp.Flag(
+		"collector.gpu.nvidia.sm-count",
+		`SM count for NVIDIA GPUs. 
+It should be of format <gpu_minor>: <sm_count> delimited by ",".
+Example: 0:160,1:160.`,
+	).Default("").String()
+)
+
 // Regexes.
 var (
 	pciBusIDRegex = regexp.MustCompile(`(?P<domain>[0-9a-fA-F]+):(?P<bus>[0-9a-fA-F]+):(?P<slot>[0-9a-fA-F]+)\.(?P<function>[0-9a-fA-F]+)`)
@@ -66,6 +76,7 @@ var (
 
 // Nvidia SM count for different architectures
 // Fetched from https://www.techpowerup.com/gpu-specs/
+// and https://cputronic.com/gpu
 var (
 	nvidiaSMCount = map[string]uint64{
 		"V100": 80,
@@ -80,7 +91,7 @@ var (
 		"H200": 132,
 		"H800": 114,
 		"B100": 132,
-		"B200": 132,
+		"B200": 160,
 	}
 )
 
@@ -910,6 +921,11 @@ func (g *GPUSMI) nvidiaGPUDevices(vendor vendor) ([]Device, error) {
 				return nil, err
 			}
 
+			// Emit a warning log if SM count is not supplied by the user
+			if *nvidiaGPUSMCount == "" {
+				g.logger.Warn("--collector.gpu.nvidia.sm-count is not provided by the user. Using a default value which might not be accurate.", "gpu_model", device.Name, "current_sm_value", device.NumSMs)
+			}
+
 			return parseNvidiaSmiListOutput(string(nvidiaSmiListOutput), devices), nil
 		}
 	}
@@ -1107,6 +1123,36 @@ func parseNvidiaSmiOutput(cmdOutput []byte) ([]Device, error) {
 		return nil, fmt.Errorf("failed to parse nvidia-smi xml log %w", err)
 	}
 
+	// If nvidiaGPUSMCount is provided, use it instead of default values
+	userSuppliedSMCount := make(map[int]uint64)
+
+	if *nvidiaGPUSMCount != "" {
+		for smCountMap := range strings.SplitSeq(*nvidiaGPUSMCount, ",") {
+			countMap := strings.Split(smCountMap, ":")
+			if len(countMap) < 2 {
+				return nil, fmt.Errorf(`--collector.gpu.nvidia.sm-count %s is invalid. It should be of format <gpu_minor>:<sm_count> delimited by ","`, *nvidiaGPUSMCount)
+			}
+
+			// Convert minor to int and count to unin64
+			minor, err := strconv.Atoi(countMap[0])
+			if err != nil {
+				return nil, fmt.Errorf(`--collector.gpu.nvidia.sm-count %s is invalid. Failed to convert minor to int: %w`, *nvidiaGPUSMCount, err)
+			}
+
+			count, err := strconv.ParseUint(countMap[1], 10, 64)
+			if err != nil {
+				return nil, fmt.Errorf(`--collector.gpu.nvidia.sm-count %s is invalid. Failed to convert sm count to uint64: %w`, *nvidiaGPUSMCount, err)
+			}
+
+			userSuppliedSMCount[minor] = count
+		}
+
+		// Check if SM count for all GPUs are provided
+		if len(userSuppliedSMCount) < len(nvidiaSMILog.GPUs) {
+			return nil, fmt.Errorf("--collector.gpu.nvidia.sm-count %s does not provide sm count for all gpus. Detected %d gpus whereas sm count is provided for %d gpus", *nvidiaGPUSMCount, len(nvidiaSMILog.GPUs), len(userSuppliedSMCount))
+		}
+	}
+
 	// NOTE: Ensure that we sort the devices using PCI address
 	// Seems like nvidia-smi most of the times returns them in correct order.
 	var globalIndex uint64
@@ -1121,20 +1167,24 @@ func parseNvidiaSmiOutput(cmdOutput []byte) ([]Device, error) {
 		}
 
 		// Attempt to get total number of SMs
-		for model, numSMs := range nvidiaSMCount {
-			// Model names can be as follows:
-			// Telsa V100-PCIe-40GB
-			// NVIDIA A100-SXM-80GB
-			// NVIDIA H100 80GB HB3
-			// NVIDIA B100 and so on..
-			// So we try to split by "space" and hypen and attempt to test
-			// against model names
-			for s := range strings.SplitSeq(gpu.ProductName, " ") {
-				for ss := range strings.SplitSeq(s, "-") {
-					if strings.TrimSpace(ss) == model {
-						dev.NumSMs = numSMs
-
-						break
+		if len(userSuppliedSMCount) > 0 {
+			dev.NumSMs = userSuppliedSMCount[igpu]
+		} else {
+			for model, numSMs := range nvidiaSMCount {
+				// Model names can be as follows:
+				// Telsa V100-PCIe-40GB
+				// NVIDIA A100-SXM-80GB
+				// NVIDIA H100 80GB HB3
+				// NVIDIA B100 and so on..
+				// So we try to split by "space" and hypen and attempt to test
+				// against model names
+				for s := range strings.SplitSeq(gpu.ProductName, " ") {
+					for ss := range strings.SplitSeq(s, "-") {
+						if strings.TrimSpace(ss) == model {
+							dev.NumSMs = numSMs
+
+							break
+						}
 					}
 				}
 			}
@@ -1708,6 +1758,11 @@ func parseTopologyProperties(path string, re *regexp.Regexp) (uint64, error) {
 		return strconv.ParseUint(match[1], 10, 64)
 	}
 
+	// If there is any error scanning the file, return it
+	if scanner.Err() != nil {
+		return 0, scanner.Err()
+	}
+
 	return 0, fmt.Errorf("no property found in the file match %s", re.String())
 }
 

pkg/collector/gpu_test.go

@@ -310,6 +310,74 @@ func TestParseNvidiaSmiOutput(t *testing.T) {
 	}
 }
 
+func TestParseNvidiaSmiOutputWithCustomSMCount(t *testing.T) {
+	_, err := CEEMSExporterApp.Parse(
+		[]string{
+			"--collector.gpu.nvidia-smi-path", "testdata/nvidia-smi",
+			"--collector.gpu.type", "nvidia",
+			"--collector.gpu.nvidia.sm-count", "0:108,1:108,2:160,3:160,4:108,5:108,6:108,7:108",
+		},
+	)
+	require.NoError(t, err)
+
+	g, err := NewGPUSMI(nil, noOpLogger)
+	require.NoError(t, err)
+
+	// Select nvidia vendor
+	for _, vendor := range g.vendors {
+		if vendor.id == nvidia {
+			gpuDevices, err := g.nvidiaGPUDevices(vendor)
+			require.NoError(t, err)
+			assert.Equal(t, uint64(160), gpuDevices[2].NumSMs)
+			assert.Equal(t, uint64(108), gpuDevices[4].NumSMs)
+		}
+	}
+}
+
+func TestParseNvidiaSmiOutputWithCustomSMCountFails(t *testing.T) {
+	_, err := CEEMSExporterApp.Parse(
+		[]string{
+			"--collector.gpu.nvidia-smi-path", "testdata/nvidia-smi",
+			"--collector.gpu.type", "nvidia",
+			"--collector.gpu.nvidia.sm-count", "0:108,1:108,2:160,4:108,5:108,6:108,7:108",
+		},
+	)
+	require.NoError(t, err)
+
+	g, err := NewGPUSMI(nil, noOpLogger)
+	require.NoError(t, err)
+
+	// Select nvidia vendor
+	for _, vendor := range g.vendors {
+		if vendor.id == nvidia {
+			// Fewer GPUs in SM count than actual number of GPUs
+			_, err := g.nvidiaGPUDevices(vendor)
+			require.Error(t, err)
+		}
+	}
+
+	_, err = CEEMSExporterApp.Parse(
+		[]string{
+			"--collector.gpu.nvidia-smi-path", "testdata/nvidia-smi",
+			"--collector.gpu.type", "nvidia",
+			"--collector.gpu.nvidia.sm-count", "0:108,1:108,2:160,3:1084:108,5:108,6:108,7:108",
+		},
+	)
+	require.NoError(t, err)
+
+	g, err = NewGPUSMI(nil, noOpLogger)
+	require.NoError(t, err)
+
+	// Select nvidia vendor
+	for _, vendor := range g.vendors {
+		if vendor.id == nvidia {
+			// Malformed string
+			_, err := g.nvidiaGPUDevices(vendor)
+			require.Error(t, err)
+		}
+	}
+}
+
 func TestNvidiaMIGAtLowerAddr(t *testing.T) {
 	nvidiaSmiLog := `<?xml version="1.0" ?>
 <!DOCTYPE nvidia_smi_log SYSTEM "nvsmi_device_v12.dtd">

website/docs/configuration/ceems-exporter.md

@@ -21,6 +21,30 @@ a consistent styling. They will be removed in `v1.0.0`.
 
 :::
 
+## GPUs
+
+When MIG is enabled on NVIDIA GPUs, current exporter will need the total number of SMs
+in each MIG partition and physical GPU to be able to estimate the fraction of GPU compute
+used by the MIG instance. Currently, `nvidia-smi` reports only the SM counts for each
+MIG instance and it does not report the SM count for entire physical GPU. If MIG is enabled
+on the NVIDIA GPUs, we **strongly recommend** the users to provide the SM count for each
+GPU using `--collector.gpu.nvidia.sm-count` CLI flag. It takes the format
+`<gpu_minor>:<sm_count>` delimited by ",".
+For example, if there are 2 GPUs on the node with 140 SMs on each GPU, then the CLI flag
+must be configured as follows:
+
+```bash
+ceems_exporter --collector.gpu.nvidia.sm-count="0:140,1:140" --collector.slurm
+```
+
+if CLI flag is not provided, exporter will use a default value based on the GPU model.
+However, same GPU model can have different SM count based on the form factor. For instance,
+[H100 PCIe 80GB](https://www.techpowerup.com/gpu-specs/h100-pcie-80-gb.c3899) model has
+114 SMs whereas [H100 SXM5 80GB](https://www.techpowerup.com/gpu-specs/h100-sxm5-80-gb.c3900) has
+132 SMs. Exporter will always use a default value of 132 SMs for H100 GPUs regardless of
+form factor which affects the power estimation for MIG instance. To know more about
+power estimation modelling when MIG is enabled, please consult [power estimation docs](../advanced/power-consumption.md).
+
 ## Collectors
 
 The following collectors are supported by the Prometheus exporter and they can be configured