feat: add gpuCapacity and gpuAllocatable filters

replicatedhq · Dec 2, 2024 · 0faa536 · 0faa536
1 parent de15d02
commit 0faa536
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 0 deletions.
diff --git a/pkg/analyze/node_resources.go b/pkg/analyze/node_resources.go
@@ -18,6 +18,8 @@ import (
 	"github.com/replicatedhq/troubleshoot/pkg/constants"
 )
 
+const gpuResourceName = "nvidia.com/gpu"
+
 type AnalyzeNodeResources struct {
 	analyzer *troubleshootv1beta2.NodeResources
 }
@@ -329,6 +331,10 @@ func getQuantity(node corev1.Node, property string) *resource.Quantity {
 		return node.Status.Capacity.StorageEphemeral()
 	case "ephemeralStorageAllocatable":
 		return node.Status.Allocatable.StorageEphemeral()
+	case "gpuCapacity":
+		return node.Status.Capacity.Name(gpuResourceName, resource.DecimalSI)
+	case "gpuAllocatable":
+		return node.Status.Allocatable.Name(gpuResourceName, resource.DecimalSI)
 	}
 	return nil
 }
@@ -492,5 +498,26 @@ func nodeMatchesFilters(node corev1.Node, filters *troubleshootv1beta2.NodeResou
 		}
 	}
 
+	if filters.GPUCapacity != "" {
+		parsed, err := resource.ParseQuantity(filters.GPUCapacity)
+		if err != nil {
+			return false, errors.Wrap(err, "failed to parse gpu capacity")
+		}
+
+		if node.Status.Capacity.Name(gpuResourceName, resource.DecimalSI).Cmp(parsed) == -1 {
+			return false, nil
+		}
+	}
+	if filters.GPUAllocatable != "" {
+		parsed, err := resource.ParseQuantity(filters.GPUAllocatable)
+		if err != nil {
+			return false, errors.Wrap(err, "failed to parse gpu allocatable")
+		}
+
+		if node.Status.Allocatable.Name(gpuResourceName, resource.DecimalSI).Cmp(parsed) == -1 {
+			return false, nil
+		}
+	}
+
 	return true, nil
 }
diff --git a/pkg/analyze/node_resources_test.go b/pkg/analyze/node_resources_test.go
@@ -404,6 +404,7 @@ func Test_nodeMatchesFilters(t *testing.T) {
 				"hugepages-2Mi":              resource.MustParse("0"),
 				"memory":                     resource.MustParse("7951376Ki"),
 				"pods":                       resource.MustParse("29"),
+				gpuResourceName:              resource.MustParse("1"),
 			},
 			Allocatable: corev1.ResourceList{
 				"attachable-volumes-aws-ebs": resource.MustParse("25"),
@@ -413,6 +414,7 @@ func Test_nodeMatchesFilters(t *testing.T) {
 				"hugepages-2Mi":              resource.MustParse("0"),
 				"memory":                     resource.MustParse("7848976Ki"),
 				"pods":                       resource.MustParse("29"),
+				gpuResourceName:              resource.MustParse("1"),
 			},
 		},
 	}
@@ -626,6 +628,38 @@ func Test_nodeMatchesFilters(t *testing.T) {
 			},
 			expectResult: false,
 		},
+		{
+			name: "true when gpu capacity is available",
+			node: node,
+			filters: &troubleshootv1beta2.NodeResourceFilters{
+				GPUCapacity: "1",
+			},
+			expectResult: true,
+		},
+		{
+			name: "true when allocatable gpu is available",
+			node: node,
+			filters: &troubleshootv1beta2.NodeResourceFilters{
+				GPUAllocatable: "1",
+			},
+			expectResult: true,
+		},
+		{
+			name: "false when gpu capacity is not available",
+			node: node,
+			filters: &troubleshootv1beta2.NodeResourceFilters{
+				GPUCapacity: "2",
+			},
+			expectResult: false,
+		},
+		{
+			name: "false when allocatable gpu is not available",
+			node: node,
+			filters: &troubleshootv1beta2.NodeResourceFilters{
+				GPUAllocatable: "2",
+			},
+			expectResult: false,
+		},
 	}
 
 	for _, test := range tests {

diff --git a/pkg/apis/troubleshoot/v1beta2/analyzer_shared.go b/pkg/apis/troubleshoot/v1beta2/analyzer_shared.go
@@ -131,6 +131,8 @@ type NodeResourceFilters struct {
 	PodAllocatable              string                 `json:"podAllocatable,omitempty" yaml:"podAllocatable,omitempty"`
 	EphemeralStorageCapacity    string                 `json:"ephemeralStorageCapacity,omitempty" yaml:"ephemeralStorageCapacity,omitempty"`
 	EphemeralStorageAllocatable string                 `json:"ephemeralStorageAllocatable,omitempty" yaml:"ephemeralStorageAllocatable,omitempty"`
+	GPUCapacity                 string                 `json:"gpuCapacity,omitempty" yaml:"gpuCapacity,omitempty"`
+	GPUAllocatable              string                 `json:"gpuAllocatable,omitempty" yaml:"gpuAllocatable,omitempty"`
 	Selector                    *NodeResourceSelectors `json:"selector,omitempty" yaml:"selector,omitempty"`
 }