fix(api): add GPU resource allocation support and fix pod request/limit calculations (#10368)

* feat(pod): add GPU resource allocation support and refactor request/limit functions

* test(pod): add GPU allocation resource validation in unit tests

Author: floreks

modules/api/pkg/resource/node/detail.go

@@ -170,7 +170,7 @@ func getNodeAllocatedResources(node v1.Node, podList *v1.PodList) (NodeAllocated
 	reqs, limits := map[v1.ResourceName]resource.Quantity{}, map[v1.ResourceName]resource.Quantity{}
 
 	for _, p := range podList.Items {
-		podReqs, podLimits, err := pod.PodRequestsAndLimits(&p)
+		podReqs, podLimits, err := pod.RequestsAndLimits(&p)
 		if err != nil {
 			return NodeAllocatedResources{}, err
 		}

modules/api/pkg/resource/pod/common.go

@@ -17,6 +17,7 @@ package pod
 import (
 	"fmt"
 
+	"github.com/samber/lo"
 	v1 "k8s.io/api/core/v1"
 
 	metricapi "k8s.io/dashboard/api/pkg/integration/metric/api"
@@ -289,11 +290,11 @@ func maxResourceList(list, new v1.ResourceList) {
 	}
 }
 
-// PodRequestsAndLimits returns a dictionary of all defined resources summed up for all
+// RequestsAndLimits returns a dictionary of all defined resources summed up for all
 // containers of the pod. If pod overhead is non-nil, the pod overhead is added to the
 // total container resource requests and to the total container limits which have a
 // non-zero quantity.
-func PodRequestsAndLimits(pod *v1.Pod) (reqs, limits v1.ResourceList, err error) {
+func RequestsAndLimits(pod *v1.Pod) (reqs, limits v1.ResourceList, err error) {
 	reqs, limits = v1.ResourceList{}, v1.ResourceList{}
 	for _, container := range pod.Spec.Containers {
 		addResourceList(reqs, container.Resources.Requests)
@@ -320,39 +321,49 @@ func PodRequestsAndLimits(pod *v1.Pod) (reqs, limits v1.ResourceList, err error)
 }
 
 func getPodAllocatedResources(pod *v1.Pod) (PodAllocatedResources, error) {
-	reqs, limits, err := PodRequestsAndLimits(pod)
+	reqs, limits, err := RequestsAndLimits(pod)
 	if err != nil {
 		return PodAllocatedResources{}, err
 	}
 
-	for podReqName, podReqValue := range reqs {
-		if value, ok := reqs[podReqName]; !ok {
-			reqs[podReqName] = podReqValue.DeepCopy()
-		} else {
-			value.Add(podReqValue)
-			reqs[podReqName] = value
-		}
-	}
-
-	for podLimitName, podLimitValue := range limits {
-		if value, ok := limits[podLimitName]; !ok {
-			limits[podLimitName] = podLimitValue.DeepCopy()
-		} else {
-			value.Add(podLimitValue)
-			limits[podLimitName] = value
-		}
-	}
-
-	cpuRequests, cpuLimits, memoryRequests, memoryLimits := reqs[v1.ResourceCPU], limits[v1.ResourceCPU], reqs[v1.ResourceMemory], limits[v1.ResourceMemory]
+	cpuRequests := reqs[v1.ResourceCPU]
+	cpuLimits := limits[v1.ResourceCPU]
+	memoryRequests := reqs[v1.ResourceMemory]
+	memoryLimits := limits[v1.ResourceMemory]
 
 	return PodAllocatedResources{
 		CPURequests:    cpuRequests.MilliValue(),
 		CPULimits:      cpuLimits.MilliValue(),
 		MemoryRequests: memoryRequests.Value(),
 		MemoryLimits:   memoryLimits.Value(),
+		GPURequests:    toGPUAllocations(reqs),
+		GPULimits:      toGPUAllocations(limits),
 	}, nil
 }
 
+func toGPUAllocations(resources v1.ResourceList) []GPUAllocation {
+	result := make([]GPUAllocation, 0)
+	for resource, quantity := range resources {
+		if gpuType := ToGPU(string(resource)); gpuType != NoGPU {
+			result = append(result, GPUAllocation{
+				Quantity: quantity.Value(),
+				Type:     gpuType,
+			})
+		}
+	}
+
+	return lo.Reduce(result, func(acc []GPUAllocation, item GPUAllocation, _ int) []GPUAllocation {
+		for i, existing := range acc {
+			if existing.Type == item.Type {
+				acc[i].Quantity += item.Quantity
+				return acc
+			}
+		}
+
+		return append(acc, item)
+	}, []GPUAllocation{})
+}
+
 func isPodInitializedConditionTrue(status *v1.PodStatus) bool {
 	for _, condition := range status.Conditions {
 		if condition.Type != v1.PodInitialized {

modules/api/pkg/resource/pod/common_test.go

@@ -44,6 +44,10 @@ func TestToPodPodStatusFailed(t *testing.T) {
 		Status:            string(v1.PodFailed),
 		Warnings:          []common.Event{},
 		ContainerStatuses: make([]ContainerStatus, 0),
+		AllocatedResources: PodAllocatedResources{
+			GPURequests: []GPUAllocation{},
+			GPULimits:   []GPUAllocation{},
+		},
 	}
 
 	actual := toPod(pod, &MetricsByPod{}, []common.Event{})
@@ -72,6 +76,10 @@ func TestToPodPodStatusSucceeded(t *testing.T) {
 		Status:            string(v1.PodSucceeded),
 		Warnings:          []common.Event{},
 		ContainerStatuses: make([]ContainerStatus, 0),
+		AllocatedResources: PodAllocatedResources{
+			GPURequests: []GPUAllocation{},
+			GPULimits:   []GPUAllocation{},
+		},
 	}
 
 	actual := toPod(pod, &MetricsByPod{}, []common.Event{})
@@ -104,6 +112,10 @@ func TestToPodPodStatusRunning(t *testing.T) {
 		Status:            string(v1.PodRunning),
 		Warnings:          []common.Event{},
 		ContainerStatuses: make([]ContainerStatus, 0),
+		AllocatedResources: PodAllocatedResources{
+			GPURequests: []GPUAllocation{},
+			GPULimits:   []GPUAllocation{},
+		},
 	}
 
 	actual := toPod(pod, &MetricsByPod{}, []common.Event{})
@@ -132,6 +144,10 @@ func TestToPodPodStatusPending(t *testing.T) {
 		Status:            string(v1.PodPending),
 		Warnings:          []common.Event{},
 		ContainerStatuses: make([]ContainerStatus, 0),
+		AllocatedResources: PodAllocatedResources{
+			GPURequests: []GPUAllocation{},
+			GPULimits:   []GPUAllocation{},
+		},
 	}
 
 	actual := toPod(pod, &MetricsByPod{}, []common.Event{})
@@ -177,6 +193,10 @@ func TestToPodContainerStates(t *testing.T) {
 				State: Waiting,
 			},
 		},
+		AllocatedResources: PodAllocatedResources{
+			GPURequests: []GPUAllocation{},
+			GPULimits:   []GPUAllocation{},
+		},
 	}
 
 	actual := toPod(pod, &MetricsByPod{}, []common.Event{})
@@ -199,6 +219,10 @@ func TestToPod(t *testing.T) {
 				TypeMeta:          types.TypeMeta{Kind: types.ResourceKindPod},
 				Warnings:          []common.Event{},
 				ContainerStatuses: make([]ContainerStatus, 0),
+				AllocatedResources: PodAllocatedResources{
+					GPURequests: []GPUAllocation{},
+					GPULimits:   []GPUAllocation{},
+				},
 			},
 		}, {
 			pod: &v1.Pod{
@@ -214,6 +238,10 @@ func TestToPod(t *testing.T) {
 				},
 				Warnings:          []common.Event{},
 				ContainerStatuses: make([]ContainerStatus, 0),
+				AllocatedResources: PodAllocatedResources{
+					GPURequests: []GPUAllocation{},
+					GPULimits:   []GPUAllocation{},
+				},
 			},
 		},
 	}

modules/api/pkg/resource/pod/list.go

@@ -15,6 +15,8 @@
 package pod
 
 import (
+	"strings"
+
 	v1 "k8s.io/api/core/v1"
 	metaV1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	k8sClient "k8s.io/client-go/kubernetes"
@@ -91,6 +93,55 @@ type PodAllocatedResources struct {
 
 	// MemoryLimits is defined memory limit.
 	MemoryLimits int64 `json:"memoryLimits"`
+
+	// GPURequests is a number and type of requested GPUs.
+	GPURequests []GPUAllocation `json:"gpuRequests"`
+
+	// GPULimits is a limit number and type of requested GPUs.
+	GPULimits []GPUAllocation `json:"gpuLimits"`
+}
+
+type GPU string
+
+const (
+	NoGPU      GPU = ""
+	UnknownGPU GPU = "unknown"
+	NvidiaGPU  GPU = "nvidia"
+	AMDGPU     GPU = "amd"
+	IntelGPU   GPU = "intel"
+
+	NvidiaLabel = "nvidia.com/gpu"
+	AMDLabel    = "amd.com/gpu"
+	// IntelLabel is for a partial match only, and it should be checked if it starts with this prefix.
+	IntelLabel = "gpu.intel.com"
+)
+
+func ToGPU(gpuType string) GPU {
+	switch gpuType {
+	case NvidiaLabel:
+		return NvidiaGPU
+	case AMDLabel:
+		return AMDGPU
+	}
+
+	if strings.HasPrefix(gpuType, IntelLabel) {
+		return IntelGPU
+	}
+
+	if strings.Contains(gpuType, "gpu") {
+		return UnknownGPU
+	}
+
+	return NoGPU
+}
+
+// GPUAllocation describes GPU allocation.
+type GPUAllocation struct {
+	// Quantity is a number of requested GPUs.
+	Quantity int64 `json:"quantity"`
+
+	// Type of GPU.
+	Type GPU `json:"type"`
 }
 
 var EmptyPodList = &PodList{

modules/api/pkg/resource/pod/list_test.go

@@ -105,6 +105,10 @@ func TestGetPodListFromChannels(t *testing.T) {
 					TypeMeta:          types.TypeMeta{Kind: types.ResourceKindPod},
 					Warnings:          []common.Event{},
 					ContainerStatuses: make([]pod.ContainerStatus, 0),
+					AllocatedResources: pod.PodAllocatedResources{
+						GPURequests: []pod.GPUAllocation{},
+						GPULimits:   []pod.GPUAllocation{},
+					},
 				}},
 				Errors: []error{},
 			},
@@ -139,3 +143,23 @@ func TestGetPodListFromChannels(t *testing.T) {
 		}
 	}
 }
+
+func TestToGPU(t *testing.T) {
+	cases := []struct {
+		in  string
+		out pod.GPU
+	}{
+		{"nvidia.com/gpu", pod.NvidiaGPU},
+		{"amd.com/gpu", pod.AMDGPU},
+		{"gpu.intel.com/xe", pod.IntelGPU},
+		{"gpu.intel.com/iris", pod.IntelGPU},
+		{"unknown.gpu.type", pod.UnknownGPU},
+		{"", pod.NoGPU},
+	}
+
+	for _, c := range cases {
+		if gpuType := pod.ToGPU(c.in); gpuType != c.out {
+			t.Errorf("ToGPU(%q) == %q, expected %q", c.in, gpuType, c.out)
+		}
+	}
+}

modules/api/pkg/resource/service/pods_test.go

@@ -61,14 +61,18 @@ func TestGetServicePods(t *testing.T) {
 						TypeMeta:          types.TypeMeta{Kind: types.ResourceKindPod},
 						Warnings:          []common.Event{},
 						ContainerStatuses: make([]pod.ContainerStatus, 0),
+						AllocatedResources: pod.PodAllocatedResources{
+							GPURequests: []pod.GPUAllocation{},
+							GPULimits:   []pod.GPUAllocation{},
+						},
 					},
 				},
 				Errors: []error{},
 			},
 		},
 	}
 	for _, c := range cases {
-		fakeClient := fake.NewSimpleClientset(c.service, c.podList)
+		fakeClient := fake.NewClientset(c.service, c.podList)
 
 		actual, _ := GetServicePods(fakeClient, nil, c.namespace, c.name, dataselect.NoDataSelect)
 		if !reflect.DeepEqual(actual, c.expected) {