diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/container_helper.go index cffc57a77d..268819473a 100644 --- 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 index fa9f4d66ba..55f81256b0 100644 --- 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 index d6398fb35d..2093639d97 100644 --- 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 index 8c252a7c06..c75ca61c42 100644 --- 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