Make nvidia resource names configurable (#359)

* Make nvidia resource names configurable Signed-off-by: lx1036 <lx1036@126.com> * added unit test Signed-off-by: Vadym Fedorov <vfedorov@nvidia.com> --------- Signed-off-by: lx1036 <lx1036@126.com> Signed-off-by: Vadym Fedorov <vfedorov@nvidia.com> Co-authored-by: Vadym Fedorov <vfedorov@nvidia.com>
NVIDIA · Jul 15, 2024 · b4552f0 · b4552f0
1 parent 219eb08
commit b4552f0
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 1 deletion.
diff --git a/pkg/cmd/app.go b/pkg/cmd/app.go
@@ -74,6 +74,7 @@ const (
 	CLIDCGMLogLevel               = "dcgm-log-level"
 	CLIPodResourcesKubeletSocket  = "pod-resources-kubelet-socket"
 	CLIHPCJobMappingDir           = "hpc-job-mapping-dir"
+	CLINvidiaResourceNames        = "nvidia-resource-names"
 )
 
 func NewApp(buildVersion ...string) *cli.App {
@@ -237,6 +238,12 @@ func NewApp(buildVersion ...string) *cli.App {
 			Usage:   "Path to HPC job mapping file directory used for mapping GPUs to jobs.",
 			EnvVars: []string{"DCGM_HPC_JOB_MAPPING_DIR"},
 		},
+		&cli.StringSliceFlag{
+			Name:    CLINvidiaResourceNames,
+			Value:   cli.NewStringSlice(),
+			Usage:   "Nvidia resource names for specified GPU type like nvidia.com/a100, nvidia.com/a10.",
+			EnvVars: []string{"NVIDIA_RESOURCE_NAMES"},
+		},
 	}
 
 	if runtime.GOOS == "linux" {
@@ -631,5 +638,6 @@ func contextToConfig(c *cli.Context) (*dcgmexporter.Config, error) {
 		DCGMLogLevel:               dcgmLogLevel,
 		PodResourcesKubeletSocket:  c.String(CLIPodResourcesKubeletSocket),
 		HPCJobMappingDir:           c.String(CLIHPCJobMappingDir),
+		NvidiaResourceNames:        c.StringSlice(CLINvidiaResourceNames),
 	}, nil
 }
diff --git a/pkg/dcgmexporter/config.go b/pkg/dcgmexporter/config.go
@@ -58,4 +58,5 @@ type Config struct {
 	DCGMLogLevel               string
 	PodResourcesKubeletSocket  string
 	HPCJobMappingDir           string
+	NvidiaResourceNames        []string
 }
diff --git a/pkg/dcgmexporter/kubernetes.go b/pkg/dcgmexporter/kubernetes.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"net"
 	"regexp"
+	"slices"
 	"strings"
 	"time"
 
@@ -147,7 +148,7 @@ func (p *PodMapper) toDeviceToPod(
 			for _, device := range container.GetDevices() {
 
 				resourceName := device.GetResourceName()
-				if resourceName != nvidiaResourceName {
+				if resourceName != nvidiaResourceName && !slices.Contains(p.Config.NvidiaResourceNames, resourceName) {
 					// Mig resources appear differently than GPU resources
 					if !strings.HasPrefix(resourceName, nvidiaMigResourcePrefix) {
 						continue

diff --git a/pkg/dcgmexporter/kubernetes_test.go b/pkg/dcgmexporter/kubernetes_test.go
@@ -174,6 +174,7 @@ func TestProcessPodMapper_WithD_Different_Format_Of_DeviceID(t *testing.T) {
 		MetricGPUDevice     string
 		MetricMigProfile    string
 		PODGPUID            string
+		NvidiaResourceNames []string
 	}
 
 	testCases := []TestCase{
@@ -232,6 +233,13 @@ func TestProcessPodMapper_WithD_Different_Format_Of_DeviceID(t *testing.T) {
 			MetricGPUDevice:     "0",
 			GPUInstanceID:       3,
 		},
+		{
+			KubernetesGPUIDType: GPUUID,
+			ResourceName:        "nvidia.com/a100",
+			MetricGPUID:         "b8ea3855-276c-c9cb-b366-c6fa655957c5",
+			PODGPUID:            "b8ea3855-276c-c9cb-b366-c6fa655957c5",
+			NvidiaResourceNames: []string{"nvidia.com/a100"},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -272,6 +280,7 @@ func TestProcessPodMapper_WithD_Different_Format_Of_DeviceID(t *testing.T) {
 				podMapper, err := NewPodMapper(&Config{
 					KubernetesGPUIdType:       tc.KubernetesGPUIDType,
 					PodResourcesKubeletSocket: socketPath,
+					NvidiaResourceNames:       tc.NvidiaResourceNames,
 				})
 				require.NoError(t, err)
 				require.NotNil(t, podMapper)