[civo] Add Gpu count to node template

Signed-off-by: Vishal Anarase <vishalanarse11@gmail.com>
kubernetes · Dec 20, 2023 · 8703ff9 · 8703ff9
1 parent c420cea
commit 8703ff9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/cluster-autoscaler/cloudprovider/civo/civo_manager.go b/cluster-autoscaler/cloudprovider/civo/civo_manager.go
@@ -234,6 +234,7 @@ func getCivoNodeTemplate(pool civocloud.KubernetesPool, client nodeGroupClient)
 	template.Labels = pool.Labels
 	template.Region = pool.Region
 	template.Taints = pool.Taints
+	template.GpuCount = size.GPUCount
 
 	return template
 }
diff --git a/cluster-autoscaler/cloudprovider/civo/civo_node_group.go b/cluster-autoscaler/cloudprovider/civo/civo_node_group.go
@@ -28,6 +28,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
 	autoscaler "k8s.io/autoscaler/cluster-autoscaler/config"
+	"k8s.io/autoscaler/cluster-autoscaler/utils/gpu"
 	"k8s.io/klog/v2"
 	schedulerframework "k8s.io/kubernetes/pkg/scheduler/framework"
 )
@@ -55,6 +56,7 @@ type CivoNodeTemplate struct {
 	DiskGigabytes int               `json:"disk_gb,omitempty"`
 	Labels        map[string]string `json:"labels,omitempty"`
 	Taints        []apiv1.Taint     `json:"taint,omitempty"`
+	GpuCount      int               `json:"gpu_count,omitempty"`
 	Region        string            `json:"region,omitempty"`
 }
 
@@ -301,7 +303,8 @@ func (n *NodeGroup) buildNodeFromTemplate(name string, template *CivoNodeTemplat
 	node.Status.Capacity[apiv1.ResourcePods] = *resource.NewQuantity(110, resource.DecimalSI)
 	node.Status.Capacity[apiv1.ResourceCPU] = *resource.NewQuantity(int64(template.CPUCores*1000), resource.DecimalSI)
 	node.Status.Capacity[apiv1.ResourceMemory] = *resource.NewQuantity(int64(template.RAMMegabytes*1024*1024), resource.DecimalSI)
-	node.Status.Capacity[apiv1.ResourceEphemeralStorage] = *resource.NewQuantity(int64(template.DiskGigabytes*1024*1024), resource.DecimalSI)
+	node.Status.Capacity[apiv1.ResourceEphemeralStorage] = *resource.NewQuantity(int64(template.DiskGigabytes*1024*1024*1024), resource.DecimalSI)
+	node.Status.Capacity[gpu.ResourceNvidiaGPU] = *resource.NewQuantity(int64(template.GpuCount), resource.DecimalSI)
 
 	node.Status.Allocatable = node.Status.Capacity
 

diff --git a/cluster-autoscaler/cloudprovider/civo/civo_node_group_test.go b/cluster-autoscaler/cloudprovider/civo/civo_node_group_test.go
@@ -543,7 +543,7 @@ func TestNodeGroup_TemplateNodeInfo(t *testing.T) {
 		assert.Equal(t, len(nodeInfo.Pods), 1, "should have one template pod")
 		assert.Equal(t, nodeInfo.Node().Status.Capacity.Cpu().ToDec().Value(), int64(1000), "should match cpu capacity ")
 		assert.Equal(t, nodeInfo.Node().Status.Capacity.Memory().ToDec().Value(), int64(1073741824), "should match memory capacity")
-		assert.Equal(t, nodeInfo.Node().Status.Capacity.StorageEphemeral().ToDec().Value(), int64(20971520), "should match epheral storage capacity")
+		assert.Equal(t, nodeInfo.Node().Status.Capacity.StorageEphemeral().ToDec().Value(), int64(21474836480), "should match epheral storage capacity")
 		assert.Equal(t, nodeInfo.Node().Labels["id"], "id", "should match labels")
 		assert.Equal(t, nodeInfo.Node().Labels["kubernetes.civo.com/civo-node-pool"], "1", "should match labels")
 		assert.Equal(t, nodeInfo.Node().Labels["kubernetes.io/os"], "linux", "should match labels")