From e79dce8ea6086b183b5f0c762a69a923b662f7b6 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <krogan@lyft.com>
Date: Mon, 9 Mar 2020 11:54:11 -0700
Subject: [PATCH] Replace lower-cased sidecar gpu resource with nvidia one
 (#65)

---
 go/tasks/pluginmachinery/flytek8s/container_helper.go |  2 +-
 go/tasks/plugins/k8s/sidecar/sidecar_test.go          | 11 +++++++++++
 go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom  |  5 ++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/go/tasks/pluginmachinery/flytek8s/container_helper.go b/go/tasks/pluginmachinery/flytek8s/container_helper.go
index d1e204b61e..07074f8eb4 100755
--- a/go/tasks/pluginmachinery/flytek8s/container_helper.go
+++ b/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/go/tasks/plugins/k8s/sidecar/sidecar_test.go b/go/tasks/plugins/k8s/sidecar/sidecar_test.go
index 26fadc4092..f6ed1d4543 100755
--- a/go/tasks/plugins/k8s/sidecar/sidecar_test.go
+++ b/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/go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom b/go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom
index c963af720c..499fc765ef 100755
--- a/go/tasks/plugins/k8s/sidecar/testdata/sidecar_custom
+++ b/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"