diff --git a/backend/src/v2/driver/driver.go b/backend/src/v2/driver/driver.go index ebb194f646e..de0fcbc3525 100644 --- a/backend/src/v2/driver/driver.go +++ b/backend/src/v2/driver/driver.go @@ -17,10 +17,12 @@ import ( "context" "encoding/json" "fmt" - "github.com/kubeflow/pipelines/backend/src/v2/objectstore" "strconv" + "strings" "time" + "github.com/kubeflow/pipelines/backend/src/v2/objectstore" + "github.com/golang/glog" "github.com/golang/protobuf/ptypes/timestamp" "github.com/google/uuid" @@ -553,6 +555,31 @@ func extendPodSpecPatch( // Get secret env information for _, secretAsEnv := range kubernetesExecutorConfig.GetSecretAsEnv() { for _, keyToEnv := range secretAsEnv.GetKeyToEnv() { + secretName := secretAsEnv.GetSecretName() + + // Check if the secret name is dynamic (wrapped in "{{}}") + if strings.HasPrefix(secretName, "{{") && strings.HasSuffix(secretName, "}}") { + // Strip the braces + key := secretName[2 : len(secretName)-2] + + // Check if the key exists in the parameter inputs map + inputParams, _, err := dag.Execution.GetParameters() + if err != nil { + return fmt.Errorf("failed to get input parameters: %v", err) + } + + val, ok := inputParams[key] + if !ok { + return fmt.Errorf("dynamic secret name key '%s' not found in input parameters", key) + } + + secretName = val.GetStringValue() + if secretName == "" { + return fmt.Errorf("secret name for key '%s' is empty", key) + } + } else if strings.TrimSpace(secretName) == "" { + return fmt.Errorf("secret name is empty or invalid") + } secretEnvVar := k8score.EnvVar{ Name: keyToEnv.GetEnvVar(), ValueFrom: &k8score.EnvVarSource{ @@ -561,7 +588,7 @@ func extendPodSpecPatch( }, }, } - secretEnvVar.ValueFrom.SecretKeyRef.LocalObjectReference.Name = secretAsEnv.GetSecretName() + secretEnvVar.ValueFrom.SecretKeyRef.LocalObjectReference.Name = secretName podSpec.Containers[0].Env = append(podSpec.Containers[0].Env, secretEnvVar) } } diff --git a/kubernetes_platform/python/kfp/kubernetes/secret.py b/kubernetes_platform/python/kfp/kubernetes/secret.py index d4a21257954..d8429e072d1 100644 --- a/kubernetes_platform/python/kfp/kubernetes/secret.py +++ b/kubernetes_platform/python/kfp/kubernetes/secret.py @@ -16,13 +16,15 @@ from google.protobuf import json_format from kfp.dsl import PipelineTask +from typing import Union +from kfp.dsl.pipeline_channel import PipelineParameterChannel from kfp.kubernetes import common from kfp.kubernetes import kubernetes_executor_config_pb2 as pb def use_secret_as_env( task: PipelineTask, - secret_name: str, + secret_name: Union[str, PipelineParameterChannel], secret_key_to_env: Dict[str, str], ) -> PipelineTask: """Use a Kubernetes Secret as an environment variable as described by the `Kubernetes documentation @@ -39,6 +41,12 @@ def use_secret_as_env( msg = common.get_existing_kubernetes_config_as_message(task) + val = secret_name + # if secret_name is a PipelineParameterChannel, then we don't know what secret to mount until RUNTIME + # so, treat is as a map KEY instead of a secret name + if isinstance(secret_name, PipelineParameterChannel): + val = "{{" + secret_name.name + "}}" + key_to_env = [ pb.SecretAsEnv.SecretKeyToEnvMap( secret_key=secret_key, @@ -46,7 +54,7 @@ def use_secret_as_env( ) for secret_key, env_var in secret_key_to_env.items() ] secret_as_env = pb.SecretAsEnv( - secret_name=secret_name, + secret_name=val, key_to_env=key_to_env, ) diff --git a/kubernetes_platform/python/test/unit/test_secret.py b/kubernetes_platform/python/test/unit/test_secret.py index a22e6be4520..3890e9c17bc 100644 --- a/kubernetes_platform/python/test/unit/test_secret.py +++ b/kubernetes_platform/python/test/unit/test_secret.py @@ -319,6 +319,36 @@ def my_pipeline(): } } + def test_with_secret_name_param_env(self): + @dsl.pipeline + def my_pipeline(secret_name: str = 'my-secret'): + task = comp() + kubernetes.use_secret_as_env( + task, + secret_name=secret_name, + secret_key_to_env={'password': 'PASSWORD'} + ) + + assert json_format.MessageToDict(my_pipeline.platform_spec) == { + 'platforms': { + 'kubernetes': { + 'deploymentSpec': { + 'executors': { + 'exec-comp': { + 'secretAsEnv': [{ + 'secretName': '{{secret_name}}', + 'keyToEnv': [{ + 'secretKey': 'password', + 'envVar': 'PASSWORD' + }] + }] + } + } + } + } + } + } + def test_preserves_secret_as_volume(self): # checks that use_secret_as_env respects previously set secrets as vol