From 7f8e7d446f1cc5303ea1d5de808039a67bc209cb Mon Sep 17 00:00:00 2001 From: Mahendra Paipuri Date: Sat, 15 Aug 2026 12:34:55 +0200 Subject: [PATCH] refactor: Support user supplied SM count for NVIDIA GPUs * For NVIDIA GPUs, SM count depends on the form factor even for the same GPU model. It is not easy to keep list of all SMs and even if we managed to do it, the way NVIDIA exposes the model name is not standardised, so there is no easy way to map the models to SM count. So, we delegate this task to operators where they can supply SM count for each GPU using CLI argument. When MIG is enabled and this CLI argument is not provided, we emit a warning log so that operators will know that we are using default SM count for current GPUs Signed-off-by: Mahendra Paipuri --- pkg/collector/gpu.go | 85 ++++++++++++++++---- pkg/collector/gpu_test.go | 68 ++++++++++++++++ website/docs/configuration/ceems-exporter.md | 24 ++++++ 3 files changed, 162 insertions(+), 15 deletions(-) diff --git a/pkg/collector/gpu.go b/pkg/collector/gpu.go index 98ac4c53..e0af8fb0 100644 --- a/pkg/collector/gpu.go +++ b/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 : delimited by ",". +Example: 0:160,1:160.`, + ).Default("").String() +) + // Regexes. var ( pciBusIDRegex = regexp.MustCompile(`(?P[0-9a-fA-F]+):(?P[0-9a-fA-F]+):(?P[0-9a-fA-F]+)\.(?P[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 : 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()) } diff --git a/pkg/collector/gpu_test.go b/pkg/collector/gpu_test.go index 2767b45b..97758d11 100644 --- a/pkg/collector/gpu_test.go +++ b/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 := ` diff --git a/website/docs/configuration/ceems-exporter.md b/website/docs/configuration/ceems-exporter.md index 43ff10df..9aafb894 100644 --- a/website/docs/configuration/ceems-exporter.md +++ b/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 +`:` 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