add support for batchScheduler in Spark Tasks (flyteorg#216)

* add support for batchScheduler Signed-off-by: Miguel <mtoledo@lyft.com> * upd test Signed-off-by: Miguel <mtoledo@lyft.com>
eapolinario · Oct 26, 2021 · 886c838 · 886c838
1 parent b5074c7
commit 886c838
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go
@@ -210,6 +210,10 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
 		},
 	}
 
+	if val, ok := sparkConfig["spark.batchScheduler"]; ok {
+		j.Spec.BatchScheduler = &val
+	}
+
 	if sparkJob.MainApplicationFile != "" {
 		j.Spec.MainApplicationFile = &sparkJob.MainApplicationFile
 	}

diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go
@@ -47,6 +47,7 @@ var (
 		"spark.flyte.feature1.enabled":    "true",
 		"spark.flyteorg.feature2.enabled": "true",
 		"spark.flyteorg.feature3.enabled": "true",
+		"spark.batchScheduler":            "volcano",
 	}
 
 	dummyEnvVars = []*core.KeyValuePair{
@@ -384,6 +385,7 @@ func TestBuildResourceSpark(t *testing.T) {
 	assert.Equal(t, int32(execInstances), *sparkApp.Spec.Executor.Instances)
 	assert.Equal(t, dummySparkConf["spark.driver.memory"], *sparkApp.Spec.Driver.Memory)
 	assert.Equal(t, dummySparkConf["spark.executor.memory"], *sparkApp.Spec.Executor.Memory)
+	assert.Equal(t, dummySparkConf["spark.batchScheduler"], *sparkApp.Spec.BatchScheduler)
 
 	// Validate Interruptible Toleration and NodeSelector set for Executor but not Driver.
 	assert.Equal(t, 0, len(sparkApp.Spec.Driver.Tolerations))