Add support for setting pod tolerations (#181)

lyft · Mar 3, 2020 · e191024 · e191024
1 parent c3c2663
commit e191024
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 2 deletions.
diff --git a/pkg/apis/app/v1beta1/types.go b/pkg/apis/app/v1beta1/types.go
@@ -109,6 +109,7 @@ type JobManagerConfig struct {
 	Replicas              *int32                      `json:"replicas,omitempty"`
 	OffHeapMemoryFraction *float64                    `json:"offHeapMemoryFraction,omitempty"`
 	NodeSelector          map[string]string           `json:"nodeSelector,omitempty"`
+	Tolerations           []apiv1.Toleration          `json:"tolerations,omitempty"`
 }
 
 type TaskManagerConfig struct {
@@ -117,6 +118,7 @@ type TaskManagerConfig struct {
 	TaskSlots             *int32                      `json:"taskSlots,omitempty"`
 	OffHeapMemoryFraction *float64                    `json:"offHeapMemoryFraction,omitempty"`
 	NodeSelector          map[string]string           `json:"nodeSelector,omitempty"`
+	Tolerations           []apiv1.Toleration          `json:"tolerations,omitempty"`
 }
 
 type EnvironmentConfig struct {

diff --git a/pkg/apis/app/v1beta1/zz_generated.deepcopy.go b/pkg/apis/app/v1beta1/zz_generated.deepcopy.go
diff --git a/pkg/controller/flink/job_manager_controller.go b/pkg/controller/flink/job_manager_controller.go
@@ -333,6 +333,7 @@ func jobmanagerTemplate(app *v1beta1.FlinkApplication) *v1.Deployment {
 					Volumes:          app.Spec.Volumes,
 					ImagePullSecrets: app.Spec.ImagePullSecrets,
 					NodeSelector:     app.Spec.JobManagerConfig.NodeSelector,
+					Tolerations:      app.Spec.JobManagerConfig.Tolerations,
 				},
 			},
 		},

diff --git a/pkg/controller/flink/job_manager_controller_test.go b/pkg/controller/flink/job_manager_controller_test.go
@@ -50,12 +50,19 @@ func TestJobManagerCreateSuccess(t *testing.T) {
 	app.Spec.JarName = testJarName
 	app.Spec.EntryClass = testEntryClass
 	app.Spec.ProgramArgs = testProgramArgs
+	app.Spec.JobManagerConfig.Tolerations = []coreV1.Toleration{{
+		Key:      "key",
+		Operator: "Equal",
+		Value:    "Value",
+		Effect:   "NoSchedule",
+	}}
+
 	annotations := map[string]string{
 		"key":                  "annotation",
 		"flink-job-properties": "jarName: " + testJarName + "\nparallelism: 8\nentryClass:" + testEntryClass + "\nprogramArgs:\"" + testProgramArgs + "\"",
 	}
 	app.Annotations = annotations
-	hash := "c3c0af0b"
+	hash := "5e7c7283"
 	expectedLabels := map[string]string{
 		"flink-app":             "app-name",
 		"flink-app-hash":        hash,
@@ -75,6 +82,7 @@ func TestJobManagerCreateSuccess(t *testing.T) {
 			assert.Equal(t, annotations, deployment.Spec.Template.Annotations)
 			assert.Equal(t, app.Namespace, deployment.Spec.Template.Namespace)
 			assert.Equal(t, expectedLabels, deployment.Labels)
+			assert.Equal(t, app.Spec.JobManagerConfig.Tolerations, deployment.Spec.Template.Spec.Tolerations)
 			assert.Equal(t, int32(1), *deployment.Spec.Replicas)
 			assert.Equal(t, "app-name", deployment.OwnerReferences[0].Name)
 			assert.Equal(t, "flink.k8s.io/v1beta1", deployment.OwnerReferences[0].APIVersion)

diff --git a/pkg/controller/flink/task_manager_controller.go b/pkg/controller/flink/task_manager_controller.go
@@ -209,6 +209,7 @@ func taskmanagerTemplate(app *v1beta1.FlinkApplication) *v1.Deployment {
 					Volumes:          app.Spec.Volumes,
 					ImagePullSecrets: app.Spec.ImagePullSecrets,
 					NodeSelector:     app.Spec.TaskManagerConfig.NodeSelector,
+					Tolerations:      app.Spec.TaskManagerConfig.Tolerations,
 				},
 			},
 		},

diff --git a/pkg/controller/flink/task_manager_controller_test.go b/pkg/controller/flink/task_manager_controller_test.go
@@ -56,12 +56,19 @@ func TestTaskManagerCreateSuccess(t *testing.T) {
 	app.Spec.JarName = testJarName
 	app.Spec.EntryClass = testEntryClass
 	app.Spec.ProgramArgs = testProgramArgs
+	app.Spec.TaskManagerConfig.Tolerations = []coreV1.Toleration{{
+		Key:      "key",
+		Operator: "Equal",
+		Value:    "Value",
+		Effect:   "NoSchedule",
+	}}
+
 	annotations := map[string]string{
 		"key":                  "annotation",
 		"flink-job-properties": "jarName: test.jar\nparallelism: 8\nentryClass:com.test.MainClass\nprogramArgs:\"--test\"",
 	}
 
-	hash := "c3c0af0b"
+	hash := "6b5e9b61"
 
 	app.Annotations = annotations
 	expectedLabels := map[string]string{
@@ -79,6 +86,7 @@ func TestTaskManagerCreateSuccess(t *testing.T) {
 		assert.Equal(t, annotations, deployment.Spec.Template.Annotations)
 		assert.Equal(t, app.Namespace, deployment.Spec.Template.Namespace)
 		assert.Equal(t, expectedLabels, deployment.Labels)
+		assert.Equal(t, app.Spec.TaskManagerConfig.Tolerations, deployment.Spec.Template.Spec.Tolerations)
 
 		assert.Equal(t, "blob.server.port: 6125\njobmanager.heap.size: 1572864k\n"+
 			"jobmanager.rpc.port: 6123\n"+