Replace lower-cased sidecar gpu resource with nvidia one (flyteorg#65)

eapolinario · Mar 9, 2020 · 5056ac3 · 5056ac3
1 parent 08c370c
commit 5056ac3
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go
@@ -19,7 +19,7 @@ import (
 
 var isAcceptableK8sName, _ = regexp.Compile("[a-z0-9]([-a-z0-9]*[a-z0-9])?")
 
-const resourceGPU = "GPU"
+const resourceGPU = "gpu"
 
 // ResourceNvidiaGPU is the name of the Nvidia GPU resource.
 // Copied from: k8s.io/autoscaler/cluster-autoscaler/utils/gpu/gpu.go

diff --git a/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar_test.go b/flyteplugins/go/tasks/plugins/k8s/sidecar/sidecar_test.go
@@ -156,6 +156,17 @@ func TestBuildSidecarResource(t *testing.T) {
 		primaryContainerKey: "a container",
 	}, res.GetAnnotations())
 	assert.Contains(t, res.(*v1.Pod).Spec.Tolerations, tolGPU)
+
+	// Test GPU overrides
+	expectedGpuRequest := resource.NewQuantity(2, resource.DecimalSI)
+	actualGpuRequest, ok := res.(*v1.Pod).Spec.Containers[0].Resources.Requests[ResourceNvidiaGPU]
+	assert.True(t, ok)
+	assert.True(t, expectedGpuRequest.Equal(actualGpuRequest))
+
+	expectedGpuLimit := resource.NewQuantity(3, resource.DecimalSI)
+	actualGpuLimit, ok := res.(*v1.Pod).Spec.Containers[0].Resources.Limits[ResourceNvidiaGPU]
+	assert.True(t, ok)
+	assert.True(t, expectedGpuLimit.Equal(actualGpuLimit))
 }
 
 func TestBuildSidecarResourceMissingPrimary(t *testing.T) {

diff --git a/flyteplugins/go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom b/flyteplugins/go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom
@@ -27,13 +27,16 @@
 			}],
 			"resources": {
 				"requests": {
+				    "gpu": {
+				    	"string": "2"
+				    },
 					"cpu": {
 						"string": "10"
 					}
 				},
 				"limits": {
 					"nvidia.com/gpu": {
-						"string": "2"
+						"string": "3"
 					},
 					"cpu": {
 						"string": "10"