keep EnvFrom from pod template (#5423)

* keep EnvFrom from pod template not complete nor tested, just as hint for potential fix for regression in #4969 * Add test and fix build error Signed-off-by: Eduardo Apolinario <[email protected]> * Fix test Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Jun 21, 2024 · c10346d · c10346d
1 parent f87a049
commit c10346d
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 6 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go
@@ -297,7 +297,7 @@ func AddFlyteCustomizationsToContainer(ctx context.Context, parameters template.
 	if parameters.IncludeConsoleURL {
 		consoleURL = parameters.TaskExecMetadata.GetConsoleURL()
 	}
-	container.Env, container.EnvFrom = DecorateEnvVars(ctx, container.Env, parameters.TaskExecMetadata.GetEnvironmentVariables(), parameters.TaskExecMetadata.GetTaskExecutionID(), consoleURL)
+	container.Env, container.EnvFrom = DecorateEnvVars(ctx, container.Env, container.EnvFrom, parameters.TaskExecMetadata.GetEnvironmentVariables(), parameters.TaskExecMetadata.GetTaskExecutionID(), consoleURL)
 
 	// retrieve platformResources and overrideResources to use when aggregating container resources
 	platformResources := parameters.TaskExecMetadata.GetPlatformResources()

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper_test.go
@@ -695,3 +695,40 @@ func TestAddFlyteCustomizationsToContainer_ValidateExistingResources(t *testing.
 	assert.True(t, container.Resources.Requests.Cpu().Equal(resource.MustParse("10")))
 	assert.True(t, container.Resources.Limits.Cpu().Equal(resource.MustParse("10")))
 }
+
+func TestAddFlyteCustomizationsToContainer_ValidateEnvFrom(t *testing.T) {
+	configMapSource := v1.EnvFromSource{
+		ConfigMapRef: &v1.ConfigMapEnvSource{
+			LocalObjectReference: v1.LocalObjectReference{
+				Name: "my-configmap",
+			},
+		},
+	}
+	secretSource := v1.EnvFromSource{
+		SecretRef: &v1.SecretEnvSource{
+			LocalObjectReference: v1.LocalObjectReference{
+				Name: "my-secret",
+			},
+		},
+	}
+
+	container := &v1.Container{
+		Command: []string{
+			"{{ .Input }}",
+		},
+		Args: []string{
+			"{{ .OutputPrefix }}",
+		},
+		EnvFrom: []v1.EnvFromSource{
+			configMapSource,
+			secretSource,
+		},
+	}
+
+	err := AddFlyteCustomizationsToContainer(context.TODO(), getTemplateParametersForTest(nil, nil, false, ""), ResourceCustomizationModeEnsureExistingResourcesInRange, container)
+	assert.NoError(t, err)
+
+	assert.Len(t, container.EnvFrom, 2)
+	assert.Equal(t, container.EnvFrom[0], configMapSource)
+	assert.Equal(t, container.EnvFrom[1], secretSource)
+}
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds.go
@@ -127,7 +127,7 @@ func GetExecutionEnvVars(id pluginsCore.TaskExecutionID, consoleURL string) []v1
 	return envVars
 }
 
-func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, taskEnvironmentVariables map[string]string, id pluginsCore.TaskExecutionID, consoleURL string) ([]v1.EnvVar, []v1.EnvFromSource) {
+func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, envFroms []v1.EnvFromSource, taskEnvironmentVariables map[string]string, id pluginsCore.TaskExecutionID, consoleURL string) ([]v1.EnvVar, []v1.EnvFromSource) {
 	envVars = append(envVars, GetContextEnvVars(ctx)...)
 	envVars = append(envVars, GetExecutionEnvVars(id, consoleURL)...)
 
@@ -142,8 +142,6 @@ func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, taskEnvironmentVa
 		envVars = append(envVars, v1.EnvVar{Name: k, Value: value})
 	}
 
-	envFroms := []v1.EnvFromSource{}
-
 	for _, secretName := range config.GetK8sPluginConfig().DefaultEnvFromSecrets {
 		optional := true
 		secretRef := v1.SecretEnvSource{LocalObjectReference: v1.LocalObjectReference{Name: secretName}, Optional: &optional}

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go
@@ -330,7 +330,7 @@ func TestDecorateEnvVars(t *testing.T) {
 				DefaultEnvVars:        tt.additionEnvVar,
 				DefaultEnvVarsFromEnv: tt.additionEnvVarFromEnv,
 			}))
-			if got, _ := DecorateEnvVars(ctx, tt.args.envVars, tt.executionEnvVar, tt.args.id, tt.consoleURL); !reflect.DeepEqual(got, tt.want) {
+			if got, _ := DecorateEnvVars(ctx, tt.args.envVars, nil, tt.executionEnvVar, tt.args.id, tt.consoleURL); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("DecorateEnvVars() = %v, want %v", got, tt.want)
 			}
 		})

diff --git a/flyteplugins/go/tasks/plugins/array/awsbatch/transformer.go b/flyteplugins/go/tasks/plugins/array/awsbatch/transformer.go
@@ -138,7 +138,7 @@ func UpdateBatchInputForArray(_ context.Context, batchInput *batch.SubmitJobInpu
 
 func getEnvVarsForTask(ctx context.Context, execID pluginCore.TaskExecutionID, containerEnvVars []*core.KeyValuePair,
 	defaultEnvVars map[string]string) []v1.EnvVar {
-	envVars, _ := flytek8s.DecorateEnvVars(ctx, flytek8s.ToK8sEnvVar(containerEnvVars), nil, execID, "")
+	envVars, _ := flytek8s.DecorateEnvVars(ctx, flytek8s.ToK8sEnvVar(containerEnvVars), nil, nil, execID, "")
 	m := make(map[string]string, len(envVars))
 	for _, envVar := range envVars {
 		m[envVar.Name] = envVar.Value