support pipelinechannel on set env vars

Signed-off-by: JerT33 <trestjeremiah@gmail.com>

compiler updates

Signed-off-by: JerT33 <trestjeremiah@gmail.com>

address copilot feedback

Signed-off-by: JerT33 <trestjeremiah@gmail.com>

fix tests

Signed-off-by: JerT33 <trestjeremiah@gmail.com>

fix formatting

Signed-off-by: JerT33 <trestjeremiah@gmail.com>

Author: JerT33

backend/src/v2/driver/driver.go

@@ -267,7 +267,11 @@ func initPodSpecPatch(
 	// Convert environment variables
 	userEnvVar := make([]k8score.EnvVar, 0)
 	for _, envVar := range container.GetEnv() {
-		userEnvVar = append(userEnvVar, k8score.EnvVar{Name: envVar.GetName(), Value: envVar.GetValue()})
+		resolvedValue, err := resolveInputParameterPlaceholders(envVar.GetValue(), executorInput)
+		if err != nil {
+			return nil, fmt.Errorf("failed to resolve environment variable %q: %w", envVar.GetName(), err)
+		}
+		userEnvVar = append(userEnvVar, k8score.EnvVar{Name: envVar.GetName(), Value: resolvedValue})
 	}
 
 	userEnvVar = append(userEnvVar, proxy.GetConfig().GetEnvVars()...)

sdk/python/kfp/compiler/compiler_test.py

@@ -1081,6 +1081,80 @@ def simple_pipeline():
                     'empty-component']
                 self.assertTrue('inputs' not in dag_task)
 
+    def test_pipeline_with_pipeline_channel_env_variable(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+
+            @dsl.component(base_image='python:3.11')
+            def empty_component():
+                pass
+
+            @dsl.pipeline()
+            def simple_pipeline(my_var: str):
+                task = empty_component()
+                task.set_env_variable('MY_VAR', my_var)
+
+            output_yaml = os.path.join(tmpdir, 'result.yaml')
+            compiler.Compiler().compile(
+                pipeline_func=simple_pipeline,
+                package_path=output_yaml,
+                pipeline_parameters={'my_var': 'default'})
+            self.assertTrue(os.path.exists(output_yaml))
+
+            with open(output_yaml, 'r') as f:
+                pipeline_spec = yaml.safe_load(f)
+                container = pipeline_spec['deploymentSpec']['executors'][
+                    'exec-empty-component']['container']
+                env_vars = {
+                    e['name']: e['value'] for e in container.get('env', [])
+                }
+                self.assertEqual(
+                    env_vars['MY_VAR'],
+                    "{{$.inputs.parameters['pipelinechannel--my_var']}}")
+                input_parameters = pipeline_spec['root']['dag']['tasks'][
+                    'empty-component']['inputs']['parameters']
+                self.assertIn('pipelinechannel--my_var', input_parameters)
+
+    def test_pipeline_with_task_output_env_variable(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+
+            @dsl.component(base_image='python:3.11')
+            def producer() -> str:
+                return 'hello'
+
+            @dsl.component(base_image='python:3.11')
+            def consumer():
+                pass
+
+            @dsl.pipeline()
+            def task_output_pipeline():
+                prod_task = producer()
+                cons_task = consumer()
+                cons_task.set_env_variable('MY_VAR', prod_task.output)
+
+            output_yaml = os.path.join(tmpdir, 'result.yaml')
+            compiler.Compiler().compile(
+                pipeline_func=task_output_pipeline, package_path=output_yaml)
+            self.assertTrue(os.path.exists(output_yaml))
+
+            with open(output_yaml, 'r') as f:
+                pipeline_spec = yaml.safe_load(f)
+                container = pipeline_spec['deploymentSpec']['executors'][
+                    'exec-consumer']['container']
+                env_vars = {
+                    e['name']: e['value'] for e in container.get('env', [])
+                }
+                self.assertEqual(
+                    env_vars['MY_VAR'],
+                    "{{$.inputs.parameters['pipelinechannel--producer-Output']}}"
+                )
+                input_parameters = pipeline_spec['root']['dag']['tasks'][
+                    'consumer']['inputs']['parameters']
+                self.assertIn('pipelinechannel--producer-Output',
+                              input_parameters)
+                self.assertEqual(
+                    input_parameters['pipelinechannel--producer-Output']
+                    ['taskOutputParameter']['producerTask'], 'producer')
+
     def test_compile_with_kubernetes_manifest_format(self):
         with tempfile.TemporaryDirectory() as tmpdir:
 

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -145,6 +145,17 @@ def build_task_spec_for_task(
         if val and pipeline_channel.extract_pipeline_channels_from_any(val):
             task.inputs['base_image'] = val
 
+    if task.container_spec and task.container_spec.env:
+        for env_name, env_val in task.container_spec.env.items():
+            if env_val and pipeline_channel.extract_pipeline_channels_from_any(
+                    env_val):
+                if env_name in task.inputs:
+                    raise ValueError(
+                        f'Environment variable name "{env_name}" collides '
+                        f'with an existing task input. Please rename either '
+                        f'the input or the environment variable.')
+                task.inputs[env_name] = env_val
+
     for input_name, input_value in task.inputs.items():
         # Since LoopParameterArgument and LoopArtifactArgument and LoopArgumentVariable are narrower
         # types than PipelineParameterChannel, start with them.
@@ -729,9 +740,7 @@ def convert_to_placeholder(input_value: str) -> str:
         compiler injected input name."""
         pipeline_channels = (
             pipeline_channel.extract_pipeline_channels_from_any(input_value))
-        if pipeline_channels:
-            assert len(pipeline_channels) == 1
-            channel = pipeline_channels[0]
+        for channel in pipeline_channels:
             additional_input_name = (
                 compiler_utils.additional_input_name_for_pipeline_channel(
                     channel))
@@ -748,7 +757,7 @@ def convert_to_placeholder(input_value: str) -> str:
             args=task.container_spec.args,
             env=[
                 pipeline_spec_pb2.PipelineDeploymentConfig.PipelineContainerSpec
-                .EnvVar(name=name, value=value)
+                .EnvVar(name=name, value=convert_to_placeholder(value))
                 for name, value in (task.container_spec.env or {}).items()
             ]))
 

sdk/python/kfp/dsl/pipeline_task.py

@@ -623,22 +623,43 @@ def set_display_name(self, name: str) -> 'PipelineTask':
         return self
 
     @block_if_final()
-    def set_env_variable(self, name: str, value: str) -> 'PipelineTask':
+    def set_env_variable(
+        self,
+        name: str,
+        value: Union[str, pipeline_channel.PipelineChannel],
+    ) -> 'PipelineTask':
         """Sets environment variable for the task.
 
         Args:
             name: Environment variable name.
-            value: Environment variable value.
+            value: Environment variable value. Supports dynamic values such as
+                Pipeline Parameters or outputs from previous tasks, which are
+                resolved at runtime.
 
         Returns:
             Self return to allow chained setting calls.
         """
         self._ensure_container_spec_exists()
 
+        pipeline_channels = (
+            pipeline_channel.extract_pipeline_channels_from_any(value))
+
+        if isinstance(value, pipeline_channel.PipelineChannel):
+            value = str(value)
+
         if self.container_spec.env is not None:
             self.container_spec.env[name] = value
         else:
             self.container_spec.env = {name: value}
+
+        if pipeline_channels:
+            existing_channel_patterns = {
+                channel.pattern for channel in self._channel_inputs
+            }
+            for channel in pipeline_channels:
+                if channel.pattern not in existing_channel_patterns:
+                    self._channel_inputs.append(channel)
+                    existing_channel_patterns.add(channel.pattern)
         return self
 
     @block_if_final()

sdk/python/kfp/dsl/pipeline_task_test.py

@@ -18,6 +18,7 @@
 
 from absl.testing import parameterized
 from kfp import dsl
+from kfp.dsl import pipeline_channel
 from kfp.dsl import pipeline_task
 from kfp.dsl import placeholders
 from kfp.dsl import structures
@@ -366,6 +367,20 @@ def test_set_env_variable(self):
         task.set_env_variable('env_name', 'env_value')
         self.assertEqual({'env_name': 'env_value'}, task.container_spec.env)
 
+    def test_set_env_variable_with_pipeline_channel(self):
+        task = pipeline_task.PipelineTask(
+            component_spec=structures.ComponentSpec.from_yaml_documents(
+                V2_YAML),
+            args={'input1': 'value'},
+        )
+        channel = pipeline_channel.PipelineParameterChannel(
+            name='param', channel_type='String', task_name='upstream')
+        task.set_env_variable('MY_VAR', channel)
+        self.assertIn('MY_VAR', task.container_spec.env)
+        self.assertEqual(str(channel), task.container_spec.env['MY_VAR'])
+        self.assertEqual(1, len(task._channel_inputs))
+        self.assertEqual(channel.pattern, task._channel_inputs[0].pattern)
+
     def test_set_display_name(self):
         task = pipeline_task.PipelineTask(
             component_spec=structures.ComponentSpec.from_yaml_documents(