Fixing issues with env vs envFrom

Signed-off-by: Neil Stout <[email protected]>
flyteorg · Mar 14, 2024 · 2be920c · 2be920c
1 parent 7709167
commit 2be920c
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 59 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go
@@ -291,7 +291,7 @@ func AddFlyteCustomizationsToContainer(ctx context.Context, parameters template.
 	}
 	container.Args = modifiedArgs
 
-	container.Env = DecorateEnvVars(ctx, container.Env, parameters.TaskExecMetadata.GetEnvironmentVariables(), parameters.TaskExecMetadata.GetTaskExecutionID())
+	container.Env, container.EnvFrom = DecorateEnvVars(ctx, container.Env, parameters.TaskExecMetadata.GetEnvironmentVariables(), parameters.TaskExecMetadata.GetTaskExecutionID())
 
 	// retrieve platformResources and overrideResources to use when aggregating container resources
 	platformResources := parameters.TaskExecMetadata.GetPlatformResources()

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds.go
@@ -113,7 +113,7 @@ func GetExecutionEnvVars(id pluginsCore.TaskExecutionID) []v1.EnvVar {
 	return envVars
 }
 
-func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, taskEnvironmentVariables map[string]string, id pluginsCore.TaskExecutionID) []v1.EnvVar {
+func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, taskEnvironmentVariables map[string]string, id pluginsCore.TaskExecutionID) ([]v1.EnvVar, []v1.EnvFromSource) {
 	envVars = append(envVars, GetContextEnvVars(ctx)...)
 	envVars = append(envVars, GetExecutionEnvVars(id)...)
 
@@ -127,26 +127,22 @@ func DecorateEnvVars(ctx context.Context, envVars []v1.EnvVar, taskEnvironmentVa
 		value := os.Getenv(envVarName)
 		envVars = append(envVars, v1.EnvVar{Name: k, Value: value})
 	}
-	for _, configMapName := range config.GetK8sPluginConfig().DefaultEnvVarsFromConfigMaps {
-		envVars = append(envVars, v1.EnvVar{
-			ValueFrom: &v1.EnvVarSource{
-				ConfigMapKeyRef: &v1.ConfigMapKeySelector{
-					Key: configMapName,
-				},
-			},
-		})
-	}
+
+	envFroms := []v1.EnvFromSource{}
+
 	for _, secretName := range config.GetK8sPluginConfig().DefaultEnvVarsFromSecrets {
-		envVars = append(envVars, v1.EnvVar{
-			ValueFrom: &v1.EnvVarSource{
-				SecretKeyRef: &v1.SecretKeySelector{
-					Key: secretName,
-				},
-			},
-		})
+		requireExist := true
+		secretRef := v1.SecretEnvSource{v1.LocalObjectReference{Name: secretName}, &requireExist}
+		envFroms = append(envFroms, v1.EnvFromSource{SecretRef: &secretRef})
 	}
 
-	return envVars
+	for _, cmName := range config.GetK8sPluginConfig().DefaultEnvVarsFromConfigMaps {
+		requireExist := true
+		cmRef := v1.ConfigMapEnvSource{v1.LocalObjectReference{Name: cmName}, &requireExist}
+		envFroms = append(envFroms, v1.EnvFromSource{ConfigMapRef: &cmRef})
+	}
+
+	return envVars, envFroms
 }
 
 func GetPodTolerations(interruptible bool, resourceRequirements ...v1.ResourceRequirements) []v1.Toleration {

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go
@@ -249,16 +249,6 @@ func TestDecorateEnvVars(t *testing.T) {
 		"k": "value",
 	}
 
-	var emptyConfigMaps []string
-	envVarConfigMaps := []string{
-		"test-cm-1",
-	}
-
-	var emptySecrets []string
-	envVarSecrets := []string{
-		"test-secret-1",
-	}
-
 	originalEnvVal := os.Getenv("value")
 	err := os.Setenv("value", "v")
 	if err != nil {
@@ -269,22 +259,6 @@ func TestDecorateEnvVars(t *testing.T) {
 	expected := append(defaultEnv, GetContextEnvVars(ctx)...)
 	expected = append(expected, GetExecutionEnvVars(mockTaskExecutionIdentifier{})...)
 
-	expectedPlusCMs := append(expected, v12.EnvVar{
-		ValueFrom: &v12.EnvVarSource{
-			ConfigMapKeyRef: &v12.ConfigMapKeySelector{
-				Key: envVarConfigMaps[0],
-			},
-		},
-	})
-
-	expectedPlusSecrets := append(expected, v12.EnvVar{
-		ValueFrom: &v12.EnvVarSource{
-			SecretKeyRef: &v12.SecretKeySelector{
-				Key: envVarSecrets[0],
-			},
-		},
-	})
-
 	aggregated := append(expected, v12.EnvVar{Name: "k", Value: "v"})
 	type args struct {
 		envVars []v12.EnvVar
@@ -296,26 +270,20 @@ func TestDecorateEnvVars(t *testing.T) {
 		additionEnvVar        map[string]string
 		additionEnvVarFromEnv map[string]string
 		executionEnvVar       map[string]string
-		configMapEnvNames     []string
-		secretEnvNames        []string
 		want                  []v12.EnvVar
 	}{
-		{"no-additional", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, emptyEnvVar, emptyConfigMaps, emptySecrets, expected},
-		{"with-additional", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, additionalEnv, emptyEnvVar, emptyEnvVar, emptyConfigMaps, emptySecrets, aggregated},
-		{"from-env", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, envVarsFromEnv, emptyEnvVar, emptyConfigMaps, emptySecrets, aggregated},
-		{"from-execution-metadata", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, additionalEnv, emptyConfigMaps, emptySecrets, aggregated},
-		{"with-from-configmap", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, emptyEnvVar, envVarConfigMaps, emptySecrets, expectedPlusCMs},
-		{"with-from-secret", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, emptyEnvVar, emptyConfigMaps, envVarSecrets, expectedPlusSecrets},
+		{"no-additional", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, emptyEnvVar, expected},
+		{"with-additional", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, additionalEnv, emptyEnvVar, emptyEnvVar, aggregated},
+		{"from-env", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, envVarsFromEnv, emptyEnvVar, aggregated},
+		{"from-execution-metadata", args{envVars: defaultEnv, id: mockTaskExecutionIdentifier{}}, emptyEnvVar, emptyEnvVar, additionalEnv, aggregated},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
-				DefaultEnvVars:               tt.additionEnvVar,
-				DefaultEnvVarsFromEnv:        tt.additionEnvVarFromEnv,
-				DefaultEnvVarsFromConfigMaps: tt.configMapEnvNames,
-				DefaultEnvVarsFromSecrets:    tt.secretEnvNames,
+				DefaultEnvVars:        tt.additionEnvVar,
+				DefaultEnvVarsFromEnv: tt.additionEnvVarFromEnv,
 			}))
-			if got := DecorateEnvVars(ctx, tt.args.envVars, tt.executionEnvVar, tt.args.id); !reflect.DeepEqual(got, tt.want) {
+			if got, _ := DecorateEnvVars(ctx, tt.args.envVars, tt.executionEnvVar, tt.args.id); !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
@@ -136,7 +136,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, execID)
 	m := make(map[string]string, len(envVars))
 	for _, envVar := range envVars {
 		m[envVar.Name] = envVar.Value