From 5056ac398acfde3f015634c0ab70bee786d9ae63 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)

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

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go
index d1e204b61e..07074f8eb4 100755
--- 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
index 26fadc4092..f6ed1d4543 100755
--- 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
index c963af720c..499fc765ef 100755
--- 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"