Cleanup: improve comments, reduce test duplication, fix copyright year

Signed-off-by: Krystian Bednarczuk <krystian@cast.ai>

Author: krystiancastai

internal/pkg/nvmlprovider/provider.go

@@ -217,9 +217,8 @@ func (n nvmlProvider) GetDeviceProcessUtilization(gpuUUID string) (map[uint32]ui
 	return result, nil
 }
 
-// GetAllMIGDevicesProcessMemory returns memory usage for all MIG devices on a parent GPU.
-// Returns a map from GPU Instance ID to (PID -> memory used in bytes).
-// Note: Only memory info is available for MIG devices, not SM utilization.
+// GetAllMIGDevicesProcessMemory returns per-process memory usage for all MIG instances on a GPU.
+// Returns map[gpuInstanceID (MIG instance)]map[PID]memoryBytes.
 func (n nvmlProvider) GetAllMIGDevicesProcessMemory(parentGPUUUID string) (map[uint]map[uint32]uint64, error) {
 	if err := n.preCheck(); err != nil {
 		return nil, fmt.Errorf("failed to get MIG device process memory: %w", err)

internal/pkg/nvmlprovider/provider_test.go

@@ -30,6 +30,30 @@ func TestGetMIGDeviceInfoByID_When_NVML_Not_Initialized(t *testing.T) {
 	assert.Error(t, err, "uuid: %v, Device Info: %+v", validMIGUUID, deviceInfo)
 }
 
+func TestGetDeviceProcessMemory_When_NVML_Not_Initialized(t *testing.T) {
+	provider := nvmlProvider{}
+	result, err := provider.GetDeviceProcessMemory("GPU-test-uuid")
+	assert.Error(t, err)
+	assert.Nil(t, result)
+	assert.Contains(t, err.Error(), "failed to get device process memory")
+}
+
+func TestGetDeviceProcessUtilization_When_NVML_Not_Initialized(t *testing.T) {
+	provider := nvmlProvider{}
+	result, err := provider.GetDeviceProcessUtilization("GPU-test-uuid")
+	assert.Error(t, err)
+	assert.Nil(t, result)
+	assert.Contains(t, err.Error(), "failed to get device process utilization")
+}
+
+func TestGetAllMIGDevicesProcessMemory_When_NVML_Not_Initialized(t *testing.T) {
+	provider := nvmlProvider{}
+	result, err := provider.GetAllMIGDevicesProcessMemory("GPU-test-uuid")
+	assert.Error(t, err)
+	assert.Nil(t, result)
+	assert.Contains(t, err.Error(), "failed to get MIG device process memory")
+}
+
 func TestGetMIGDeviceInfoByID_When_DriverVersion_Below_R470(t *testing.T) {
 	Initialize()
 	assert.NotNil(t, Client(), "expected NVML Client to be not nil")

internal/pkg/nvmlprovider/types.go

@@ -18,7 +18,6 @@
 
 package nvmlprovider
 
-// NVML interface provides access to NVIDIA Management Library functionality
 type NVML interface {
 	GetMIGDeviceInfoByID(string) (*MIGDeviceInfo, error)
 	// GetDeviceProcessMemory returns memory usage for processes running on the GPU.
@@ -27,9 +26,8 @@ type NVML interface {
 	// GetDeviceProcessUtilization returns SM utilization for processes running on the GPU.
 	// Returns a map from PID to SM utilization percentage.
 	GetDeviceProcessUtilization(gpuUUID string) (map[uint32]uint32, error)
-	// GetAllMIGDevicesProcessMemory returns memory usage for all MIG devices on a parent GPU.
-	// Returns a map from GPU Instance ID to (PID -> memory used in bytes).
-	// Note: Only memory info is available for MIG devices, not SM utilization.
+	// GetAllMIGDevicesProcessMemory returns per-process memory usage for all MIG instances on a GPU.
+	// Returns map[gpuInstanceID (MIG instance)]map[PID]memoryBytes.
 	GetAllMIGDevicesProcessMemory(parentGPUUUID string) (map[uint]map[uint32]uint64, error)
 	Cleanup()
 }

internal/pkg/transformation/const.go

@@ -17,6 +17,7 @@
 package transformation
 
 const (
+	// Note standard resource attributes
 	podAttribute       = "pod"
 	namespaceAttribute = "namespace"
 	containerAttribute = "container"
@@ -39,7 +40,6 @@ const (
 
 	DRAGPUDriverName = "gpu.nvidia.com"
 
-	metricGPUUtil        = "DCGM_FI_DEV_GPU_UTIL"
-	metricFBUsed         = "DCGM_FI_DEV_FB_USED"
-	metricGREngineActive = "DCGM_FI_PROF_GR_ENGINE_ACTIVE"
+	metricGPUUtil = "DCGM_FI_DEV_GPU_UTIL"
+	metricFBUsed  = "DCGM_FI_DEV_FB_USED"
 )

internal/pkg/transformation/kubernetes.go

@@ -356,6 +356,9 @@ func (p *PodMapper) Process(metrics collector.MetricsByCounter, deviceInfo devic
 				// with the container info and the shared GPU label, if it exists.
 				// Notably, this will increase the number of unique metrics (i.e. labelsets)
 				// to by the number of containers sharing the GPU.
+				if len(podInfos) > 0 {
+					newmetrics = append(newmetrics, metrics[counter][j]) // original device-level metric
+				}
 				for _, pi := range podInfos {
 					metric, err := utils.DeepCopy(metrics[counter][j])
 					if err != nil {

internal/pkg/transformation/kubernetes_test.go

@@ -946,6 +946,7 @@ func TestProcessPodMapper_WithLabelsAndUID(t *testing.T) {
 }
 
 func TestBuildPodValueMap(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name      string
 		pidToPod  map[uint32]*PodInfo
@@ -960,6 +961,15 @@ func TestBuildPodValueMap(t *testing.T) {
 			fieldName: metricGPUUtil,
 			expected:  map[string]string{},
 		},
+		{
+			name:     "empty pidToPod returns empty map",
+			pidToPod: map[uint32]*PodInfo{},
+			data: &perProcessMetrics{
+				pidToSMUtil: map[uint32]uint32{1001: 50},
+			},
+			fieldName: metricGPUUtil,
+			expected:  map[string]string{},
+		},
 		{
 			name:     "maps PID values to pod UIDs for GPU util",
 			pidToPod: map[uint32]*PodInfo{1001: {UID: "uid1"}, 1002: {UID: "uid2"}},
@@ -991,13 +1001,15 @@ func TestBuildPodValueMap(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			result := buildPodValueMap(tc.pidToPod, tc.data, tc.fieldName)
 			assert.Equal(t, tc.expected, result)
 		})
 	}
 }
 
 func TestBuildIdlePodValues(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name           string
 		existingValues map[string]string
@@ -1026,13 +1038,15 @@ func TestBuildIdlePodValues(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			result := buildIdlePodValues(tc.existingValues, tc.devicePods)
 			assert.Equal(t, tc.expected, result)
 		})
 	}
 }
 
 func TestPodMapper_CreatePerProcessMetrics(t *testing.T) {
+	t.Parallel()
 	gpuUUID := "GPU-00000000-0000-0000-0000-000000000000"
 	podUID := "a9c80282-3f6b-4d5b-84d5-a137a6668011"
 
@@ -1226,6 +1240,7 @@ func TestPodMapper_CreatePerProcessMetrics(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			podMapper := &PodMapper{
 				Config: &appconfig.Config{
 					UseOldNamespace: tc.useOldNS,
@@ -1245,21 +1260,17 @@ func TestPodMapper_CreatePerProcessMetrics(t *testing.T) {
 }
 
 func TestStripVGPUSuffix(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name     string
 		deviceID string
 		expected string
 	}{
 		{
-			name:     "AWS MIG device ID with vgpu suffix",
+			name:     "MIG device ID with vgpu suffix",
 			deviceID: "MIG-2ce7a541-c516-5dbc-a76e-26cc100d9b55::7",
 			expected: "MIG-2ce7a541-c516-5dbc-a76e-26cc100d9b55",
 		},
-		{
-			name:     "AWS MIG device ID with different vgpu index",
-			deviceID: "MIG-a8d7e63b-588b-5fd8-826d-d1eab19c6f18::9",
-			expected: "MIG-a8d7e63b-588b-5fd8-826d-d1eab19c6f18",
-		},
 		{
 			name:     "Plain MIG UUID without suffix",
 			deviceID: "MIG-2ce7a541-c516-5dbc-a76e-26cc100d9b55",
@@ -1289,6 +1300,7 @@ func TestStripVGPUSuffix(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			result := stripVGPUSuffix(tc.deviceID)
 			assert.Equal(t, tc.expected, result)
 		})

internal/pkg/transformation/pidmapper.go

@@ -1,7 +1,5 @@
-//go:build linux
-
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

internal/pkg/transformation/pidmapper_stub.go

@@ -1,7 +1,5 @@
-//go:build linux
-
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,6 +23,7 @@ import (
 )
 
 func TestExtractPodUID(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name     string
 		path     string
@@ -64,13 +63,15 @@ func TestExtractPodUID(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			result := extractPodUID(tc.path)
 			assert.Equal(t, tc.expected, result)
 		})
 	}
 }
 
 func TestExtractPodUIDFromPaths(t *testing.T) {
+	t.Parallel()
 	tests := []struct {
 		name       string
 		subsystems map[string]string
@@ -104,13 +105,15 @@ func TestExtractPodUIDFromPaths(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
 			result := extractPodUIDFromPaths(tc.subsystems, tc.unified)
 			assert.Equal(t, tc.expected, result)
 		})
 	}
 }
 
 func TestBuildPIDToPodMap(t *testing.T) {
+	t.Parallel()
 	mapper := newPIDToPodMapper()
 
 	pods := []PodInfo{

internal/pkg/transformation/pidmapper_test.go

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,7 @@ import (
 )
 
 func isPerProcessMetric(fieldName string) bool {
-	return fieldName == metricGPUUtil || fieldName == metricFBUsed || fieldName == metricGREngineActive
+	return fieldName == metricGPUUtil || fieldName == metricFBUsed
 }
 
 // getGPUUUIDToDeviceID builds a mapping from GPU UUID to device ID based on the specified ID type.