diff --git a/api/kubernetes/snapshot.go b/api/kubernetes/snapshot.go index c75d510846..c8f6c2e598 100644 --- a/api/kubernetes/snapshot.go +++ b/api/kubernetes/snapshot.go @@ -24,5 +24,5 @@ func (snapshotter *Snapshotter) CreateSnapshot(endpoint *portainer.Endpoint) (*p return nil, err } - return snapshot.CreateKubernetesSnapshot(client, endpoint.Kubernetes.Flags.GPUOperator) + return snapshot.CreateKubernetesSnapshot(client) } diff --git a/pkg/snapshot/kubernetes.go b/pkg/snapshot/kubernetes.go index c7b5e125f9..477475ec7b 100644 --- a/pkg/snapshot/kubernetes.go +++ b/pkg/snapshot/kubernetes.go @@ -22,14 +22,14 @@ import ( "k8s.io/client-go/kubernetes" ) -func CreateKubernetesSnapshot(cli *kubernetes.Clientset, gpuOperator bool) (*portainer.KubernetesSnapshot, error) { +func CreateKubernetesSnapshot(cli *kubernetes.Clientset) (*portainer.KubernetesSnapshot, error) { kubernetesSnapshot := &portainer.KubernetesSnapshot{} err := kubernetesSnapshotVersion(kubernetesSnapshot, cli) if err != nil { log.Warn().Err(err).Msg("unable to snapshot cluster version") } - err = kubernetesSnapshotNodes(kubernetesSnapshot, cli, gpuOperator) + err = kubernetesSnapshotNodes(kubernetesSnapshot, cli) if err != nil { log.Warn().Err(err).Msg("unable to snapshot cluster nodes") } @@ -48,7 +48,7 @@ func kubernetesSnapshotVersion(snapshot *portainer.KubernetesSnapshot, cli kuber return nil } -func kubernetesSnapshotNodes(snapshot *portainer.KubernetesSnapshot, cli kubernetes.Interface, gpuOperator bool) error { +func kubernetesSnapshotNodes(snapshot *portainer.KubernetesSnapshot, cli kubernetes.Interface) error { nodeList, err := cli.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{}) if err != nil { return err @@ -66,30 +66,25 @@ func kubernetesSnapshotNodes(snapshot *portainer.KubernetesSnapshot, cli kuberne totalCPUs += node.Status.Capacity.Cpu().Value() totalMemory += node.Status.Capacity.Memory().Value() - if gpuOperator { - nodeHasGPU := false - for resourceName, quantity := range node.Status.Capacity { - if strings.HasPrefix(string(resourceName), "nvidia.com/") { - totalGPU[string(resourceName)] += quantity.Value() - nodeHasGPU = true - } - } - if nodeHasGPU { - gpuNodeCount++ + nodeHasGPU := false + for resourceName, quantity := range node.Status.Capacity { + if strings.HasPrefix(string(resourceName), "nvidia.com/") { + totalGPU[string(resourceName)] += quantity.Value() + nodeHasGPU = true } } + if nodeHasGPU { + gpuNodeCount++ + } } snapshot.TotalCPU = totalCPUs snapshot.TotalMemory = totalMemory snapshot.NodeCount = len(nodeList.Items) snapshot.ClusterType = clusterTypeFromProviderID(nodeList.Items[0].Spec.ProviderID) - - if gpuOperator { - snapshot.GPUNodeCount = gpuNodeCount - if len(totalGPU) > 0 { - snapshot.TotalGPU = totalGPU - } + snapshot.GPUNodeCount = gpuNodeCount + if len(totalGPU) > 0 { + snapshot.TotalGPU = totalGPU } return nil diff --git a/pkg/snapshot/kubernetes_test.go b/pkg/snapshot/kubernetes_test.go index 79b9df87b1..265b36a3d5 100644 --- a/pkg/snapshot/kubernetes_test.go +++ b/pkg/snapshot/kubernetes_test.go @@ -95,7 +95,7 @@ func TestKubernetesSnapshotNodes(t *testing.T) { snapshot := &portainer.KubernetesSnapshot{} // Use the actual function now that it accepts kubernetes.Interface - err = kubernetesSnapshotNodes(snapshot, fakeClient, false) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) // Verify the results - these should match what kubernetesSnapshotNodes would produce @@ -117,7 +117,7 @@ func TestKubernetesSnapshotNodesEmptyCluster(t *testing.T) { fakeClient := kfake.NewClientset() snapshot := &portainer.KubernetesSnapshot{} - err := kubernetesSnapshotNodes(snapshot, fakeClient, false) + err := kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) // Values should remain at their zero state when no nodes exist @@ -173,7 +173,7 @@ func TestCreateKubernetesSnapshotIntegration(t *testing.T) { // Test that kubernetesSnapshotNodes logic works snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, false) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) // Verify the integration results @@ -201,7 +201,7 @@ func TestKubernetesSnapshotNodesWithAPIError(t *testing.T) { }) snapshot := &portainer.KubernetesSnapshot{} - err := kubernetesSnapshotNodes(snapshot, fakeClient, false) + err := kubernetesSnapshotNodes(snapshot, fakeClient) // Should return the API error require.Error(t, err) @@ -237,7 +237,7 @@ func TestKubernetesSnapshotNodesSingleNode(t *testing.T) { require.NoError(t, err) snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, false) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 1, snapshot.NodeCount) @@ -279,7 +279,7 @@ func TestKubernetesSnapshotNodesWithGPU(t *testing.T) { require.NoError(t, err) snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, true) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 2, snapshot.NodeCount) @@ -307,7 +307,7 @@ func TestKubernetesSnapshotNodesMultipleGPUTypes(t *testing.T) { require.NoError(t, err) snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, true) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 1, snapshot.GPUNodeCount) @@ -335,7 +335,7 @@ func TestKubernetesSnapshotNodesGPUAggregatedAcrossNodes(t *testing.T) { } snapshot := &portainer.KubernetesSnapshot{} - err := kubernetesSnapshotNodes(snapshot, fakeClient, true) + err := kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 2, snapshot.GPUNodeCount) @@ -359,7 +359,7 @@ func TestKubernetesSnapshotNodesNoGPULeavesTotalGPUNil(t *testing.T) { require.NoError(t, err) snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, false) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 0, snapshot.GPUNodeCount) @@ -387,7 +387,7 @@ func TestKubernetesSnapshotNodesZeroResources(t *testing.T) { require.NoError(t, err) snapshot := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snapshot, fakeClient, false) + err = kubernetesSnapshotNodes(snapshot, fakeClient) require.NoError(t, err) require.Equal(t, 1, snapshot.NodeCount) @@ -398,76 +398,50 @@ func TestKubernetesSnapshotNodesZeroResources(t *testing.T) { t.Log("Zero resources test passed - handles edge case correctly") } -func TestKubernetesSnapshotNodesGPUOperator(t *testing.T) { +func TestKubernetesSnapshotNodesGPUDetection(t *testing.T) { t.Parallel() const gpuMemoryBytes = int64(25769803776) // 16GiB + 8GiB - tests := []struct { - name string - gpuOperator bool - wantGPUCount int - wantTotalGPU map[string]int64 - }{ + fakeClient := kfake.NewClientset() + nodes := []*corev1.Node{ { - name: "disabled does not populate GPU fields even when GPU nodes exist", - gpuOperator: false, - wantGPUCount: 0, - wantTotalGPU: nil, + ObjectMeta: metav1.ObjectMeta{Name: "gpu-node"}, + Status: corev1.NodeStatus{ + Capacity: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("8"), + corev1.ResourceMemory: resource.MustParse("16Gi"), + "nvidia.com/gpu": resource.MustParse("4"), + }, + }, }, { - name: "enabled populates GPU fields from GPU nodes", - gpuOperator: true, - wantGPUCount: 1, - wantTotalGPU: map[string]int64{"nvidia.com/gpu": 4}, + ObjectMeta: metav1.ObjectMeta{Name: "cpu-node"}, + Status: corev1.NodeStatus{ + Capacity: corev1.ResourceList{ + corev1.ResourceCPU: resource.MustParse("4"), + corev1.ResourceMemory: resource.MustParse("8Gi"), + }, + }, }, } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - fakeClient := kfake.NewClientset() - nodes := []*corev1.Node{ - { - ObjectMeta: metav1.ObjectMeta{Name: "gpu-node"}, - Status: corev1.NodeStatus{ - Capacity: corev1.ResourceList{ - corev1.ResourceCPU: resource.MustParse("8"), - corev1.ResourceMemory: resource.MustParse("16Gi"), - "nvidia.com/gpu": resource.MustParse("4"), - }, - }, - }, - { - ObjectMeta: metav1.ObjectMeta{Name: "cpu-node"}, - Status: corev1.NodeStatus{ - Capacity: corev1.ResourceList{ - corev1.ResourceCPU: resource.MustParse("4"), - corev1.ResourceMemory: resource.MustParse("8Gi"), - }, - }, - }, - } - for _, n := range nodes { - _, err := fakeClient.CoreV1().Nodes().Create(t.Context(), n, metav1.CreateOptions{}) - require.NoError(t, err) - } - - snap := &portainer.KubernetesSnapshot{} - err := kubernetesSnapshotNodes(snap, fakeClient, tt.gpuOperator) - require.NoError(t, err) - - require.Equal(t, 2, snap.NodeCount) - require.Equal(t, int64(12), snap.TotalCPU) - require.Equal(t, gpuMemoryBytes, snap.TotalMemory) - require.Equal(t, tt.wantGPUCount, snap.GPUNodeCount) - require.Equal(t, tt.wantTotalGPU, snap.TotalGPU) - }) + for _, n := range nodes { + _, err := fakeClient.CoreV1().Nodes().Create(t.Context(), n, metav1.CreateOptions{}) + require.NoError(t, err) } + + snap := &portainer.KubernetesSnapshot{} + err := kubernetesSnapshotNodes(snap, fakeClient) + require.NoError(t, err) + + require.Equal(t, 2, snap.NodeCount) + require.Equal(t, int64(12), snap.TotalCPU) + require.Equal(t, gpuMemoryBytes, snap.TotalMemory) + require.Equal(t, 1, snap.GPUNodeCount) + require.Equal(t, map[string]int64{"nvidia.com/gpu": 4}, snap.TotalGPU) } -func TestKubernetesSnapshotNodesGPUOperatorEnabledNoGPUNodes(t *testing.T) { +func TestKubernetesSnapshotNodesNoGPUNodes(t *testing.T) { t.Parallel() fakeClient := kfake.NewClientset() @@ -484,7 +458,7 @@ func TestKubernetesSnapshotNodesGPUOperatorEnabledNoGPUNodes(t *testing.T) { require.NoError(t, err) snap := &portainer.KubernetesSnapshot{} - err = kubernetesSnapshotNodes(snap, fakeClient, true) + err = kubernetesSnapshotNodes(snap, fakeClient) require.NoError(t, err) require.Equal(t, 1, snap.NodeCount)