Add env FLYTE_FAIL_ON_ERROR to aws batch job (flyteorg#306)

* Correct aws batch job state Signed-off-by: Kevin Su <[email protected]> * Add env FAST_ON_ERROR to aws batch job Signed-off-by: Kevin Su <[email protected]> * Add tests Signed-off-by: Kevin Su <[email protected]> * update tests Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * update test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Kevin Su <[email protected]> Signed-off-by: Kevin Su <[email protected]>
eapolinario · Jan 23, 2023 · 2518b61 · 2518b61
1 parent be58134
commit 2518b61
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/flyteplugins/go/tasks/plugins/array/awsbatch/transformer.go b/flyteplugins/go/tasks/plugins/array/awsbatch/transformer.go
@@ -25,6 +25,7 @@ import (
 const (
 	ArrayJobIndex       = "BATCH_JOB_ARRAY_INDEX_VAR_NAME"
 	arrayJobIDFormatter = "%v:%v"
+	failOnError         = "FLYTE_FAIL_ON_ERROR"
 )
 
 const assignResources = true
@@ -118,7 +119,6 @@ func UpdateBatchInputForArray(_ context.Context, batchInput *batch.SubmitJobInpu
 		envVars = append(envVars, &batch.KeyValuePair{Name: refStr(ArrayJobIndex), Value: refStr("FAKE_JOB_ARRAY_INDEX")},
 			&batch.KeyValuePair{Name: refStr("FAKE_JOB_ARRAY_INDEX"), Value: refStr("0")})
 	}
-
 	batchInput.ArrayProperties = arrayProps
 	batchInput.ContainerOverrides.Environment = envVars
 
@@ -136,15 +136,14 @@ func getEnvVarsForTask(ctx context.Context, execID pluginCore.TaskExecutionID, c
 	for key, value := range defaultEnvVars {
 		m[key] = value
 	}
-
+	m[failOnError] = "true"
 	finalEnvVars := make([]v1.EnvVar, 0, len(m))
 	for key, val := range m {
 		finalEnvVars = append(finalEnvVars, v1.EnvVar{
 			Name:  key,
 			Value: val,
 		})
 	}
-
 	return finalEnvVars
 }
 

diff --git a/flyteplugins/go/tasks/plugins/array/awsbatch/transformer_test.go b/flyteplugins/go/tasks/plugins/array/awsbatch/transformer_test.go
@@ -130,6 +130,7 @@ func TestArrayJobToBatchInput(t *testing.T) {
 		ContainerOverrides: &batch.ContainerOverrides{
 			Command: []*string{ref("cmd"), ref("/inputs/prefix")},
 			Environment: []*batch.KeyValuePair{
+				{Name: ref(failOnError), Value: refStr("true")},
 				{Name: refStr("BATCH_JOB_ARRAY_INDEX_VAR_NAME"), Value: refStr("AWS_BATCH_JOB_ARRAY_INDEX")},
 			},
 			Memory: refInt(1074),
@@ -237,5 +238,9 @@ func Test_getEnvVarsForTask(t *testing.T) {
 			Name:  "MyKey",
 			Value: "MyVal",
 		},
+		{
+			Name:  "FLYTE_FAIL_ON_ERROR",
+			Value: "true",
+		},
 	}, envVars)
 }