From 5027eb83d98e37b575f232e8df945ad0d7b57fe6 Mon Sep 17 00:00:00 2001
From: qzhu <qzhu@cloudera.com>
Date: Fri, 15 Dec 2023 22:00:07 +0800
Subject: [PATCH] Fix error

---
 pkg/common/resource.go      | 20 ++++++++++++++-----
 pkg/common/resource_test.go | 38 +++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/pkg/common/resource.go b/pkg/common/resource.go
index 96c8c7a96..826ddcffb 100644
--- a/pkg/common/resource.go
+++ b/pkg/common/resource.go
@@ -57,11 +57,6 @@ func GetPodResource(pod *v1.Pod) (resource *si.Resource) {
 		Resources: map[string]*si.Quantity{"pods": {Value: 1}},
 	}
 
-	// A QosBestEffort pod does not request any resources, just a single pod
-	if qos.GetPodQOS(pod) == v1.PodQOSBestEffort {
-		return podResource
-	}
-
 	for _, c := range pod.Spec.Containers {
 		resourceList := c.Resources.Requests
 		containerResource := getResource(resourceList)
@@ -74,6 +69,21 @@ func GetPodResource(pod *v1.Pod) (resource *si.Resource) {
 		checkInitContainerRequest(pod, podResource)
 	}
 
+	// iterate the pod resources when resource is zero, remove it from the pod resource
+	for k, v := range podResource.Resources {
+		if v.Value == 0 {
+			delete(podResource.Resources, k)
+		}
+	}
+
+	// A QosBestEffort pod does not request any cpu/mem resources, just a single pod
+	// But with other resources requested, it is not the best effort pod
+	if qos.GetPodQOS(pod) == v1.PodQOSBestEffort && (len(podResource.Resources) == 1) {
+		return &si.Resource{
+			Resources: map[string]*si.Quantity{"pods": {Value: 1}},
+		}
+	}
+
 	// K8s pod EnableOverHead from:
 	// alpha: v1.16
 	// beta: v1.18
diff --git a/pkg/common/resource_test.go b/pkg/common/resource_test.go
index f15554491..3e0065458 100644
--- a/pkg/common/resource_test.go
+++ b/pkg/common/resource_test.go
@@ -329,6 +329,44 @@ func TestBestEffortPod(t *testing.T) {
 	assert.Equal(t, res.Resources["pods"].GetValue(), int64(1))
 }
 
+func TestGPUOnlyResources(t *testing.T) {
+	containers := make([]v1.Container, 0)
+
+	// container 01
+	c1Resources := make(map[v1.ResourceName]resource.Quantity)
+	c1Resources[v1.ResourceName("nvidia.com/gpu")] = resource.MustParse("1")
+	containers = append(containers, v1.Container{
+		Name: "container-01",
+		Resources: v1.ResourceRequirements{
+			Requests: c1Resources,
+		},
+	})
+
+	pod := &v1.Pod{
+		TypeMeta: apis.TypeMeta{
+			Kind:       "Pod",
+			APIVersion: "v1",
+		},
+		ObjectMeta: apis.ObjectMeta{
+			Name: "pod-resource-test-00001",
+			UID:  "UID-00001",
+		},
+		Spec: v1.PodSpec{
+			Containers: containers,
+		},
+	}
+
+	res := GetPodResource(pod)
+	assert.Equal(t, len(res.Resources), 2)
+	assert.Equal(t, res.Resources["pods"].GetValue(), int64(1))
+	assert.Equal(t, res.Resources["nvidia.com/gpu"].GetValue(), int64(1))
+
+	c1Resources[v1.ResourceName("nvidia.com/gpu")] = resource.MustParse("0")
+	res = GetPodResource(pod)
+	assert.Equal(t, len(res.Resources), 1)
+	assert.Equal(t, res.Resources["pods"].GetValue(), int64(1))
+}
+
 func TestNodeResource(t *testing.T) {
 	nodeCapacity := make(map[v1.ResourceName]resource.Quantity)
 	nodeCapacity[v1.ResourceCPU] = resource.MustParse("14500m")