fixing formatting issues and resolving pr comments

Signed-off-by: Nelesh Singla <117123879+nsingla@users.noreply.github.com>

Author: nsingla

.github/workflows/kfp-sdk-tests.yml

@@ -20,7 +20,7 @@ concurrency:
 jobs:
   sdk-tests:
     runs-on: ubuntu-latest
-    timeout-minutes: 45
+    timeout-minutes: 50
     strategy:
       matrix:
         python-version: ['3.9', '3.13']

backend/test/testutil/kubernetes_utils.go

@@ -77,7 +77,12 @@ func ReadPodLogs(client *kubernetes.Clientset, namespace string, podName string,
 			podLogsRequest := client.CoreV1().Pods(namespace).GetLogs(podFromPodName.Name, podLogOptions)
 			podLogs, err := podLogsRequest.Stream(context.Background()) // Pass a context for cancellation
 			if err != nil {
-				logger.Log("Failed to stream pod logs due to %v", err)
+				logger.Log("Failed to stream pod logs for container '%s' due to %v", container.Name, err)
+				continue
+			}
+			if podLogs == nil {
+				logger.Log("Pod log stream is nil for container '%s'", container.Name)
+				continue
 			}
 			defer func(podLogs io.ReadCloser) {
 				err = podLogs.Close()
@@ -87,7 +92,7 @@ func ReadPodLogs(client *kubernetes.Clientset, namespace string, podName string,
 			}(podLogs)
 			_, err = io.Copy(buf, podLogs)
 			if err != nil {
-				logger.Log("Failed to add pod logs to buffer due to: %v", err)
+				logger.Log("Failed to add pod logs to buffer for container '%s' due to: %v", container.Name, err)
 			}
 		}
 	} else {

kubernetes_platform/python/kfp/kubernetes/common.py

@@ -47,22 +47,6 @@ def get_existing_kubernetes_config_as_message(
     return json_format.ParseDict(cur_k8_config_dict, k8_config_msg)
 
 
-def ensure_channel_input(task: PipelineTask,
-                         channel: pipeline_channel.PipelineChannel) -> None:
-    """Adds a channel to the task's tracked inputs if not already present.
-
-    This ensures the compiler propagates the channel through sub-DAG
-    boundaries (e.g. ParallelFor) even when the channel is only
-    referenced from Kubernetes platform config and not from normal
-    task arguments.
-    """
-    existing_channel_patterns = {
-        existing.pattern for existing in task._channel_inputs
-    }
-    if channel.pattern not in existing_channel_patterns:
-        task._channel_inputs.append(channel)
-
-
 def parse_k8s_parameter_input(
     input_param: Union[pipeline_channel.PipelineParameterChannel, str, dict],
     task: PipelineTask,

kubernetes_platform/python/kfp/kubernetes/volume.py

@@ -88,7 +88,7 @@ def mount_pvc(
     pvc_mount.pvc_name_parameter.CopyFrom(pvc_name_parameter)
 
     # deprecated: for backwards compatibility
-    pvc_name_from_task = _assign_pvc_name_to_msg(task, pvc_mount, pvc_name)
+    pvc_name_from_task = _assign_pvc_name_to_msg(pvc_mount, pvc_name)
 
     if pvc_name_from_task:
         task.after(pvc_name.task)
@@ -109,11 +109,8 @@ def DeletePVC(pvc_name: str):
     return dsl.ContainerSpec(image='argostub/deletepvc')
 
 
-def _assign_pvc_name_to_msg(
-    task: PipelineTask,
-    msg: message.Message,
-    pvc_name: Union[str, 'PipelineChannel'],
-) -> bool:
+def _assign_pvc_name_to_msg(msg: message.Message,
+                            pvc_name: Union[str, 'PipelineChannel']) -> bool:
     """Assigns pvc_name to the msg's pvc_reference oneof.
 
     Returns True if pvc_name is an upstream task output; otherwise, False.
@@ -122,7 +119,6 @@ def _assign_pvc_name_to_msg(
         msg.constant = pvc_name
         return False
     elif hasattr(pvc_name, 'task_name'):
-        common.ensure_channel_input(task=task, channel=pvc_name)
         if pvc_name.task_name is None:
             msg.component_input_parameter = pvc_name.name
             return False

kubernetes_platform/python/test/unit/test_secret.py

@@ -973,29 +973,29 @@ def my_pipeline():
 
 class TestEnsureChannelInputPropagation:
     """Tests that PipelineParameterChannel values used in Kubernetes platform
-    config are correctly registered in task._channel_inputs for sub-DAG
+    config are correctly registered in task.channel_inputs for sub-DAG
     propagation."""
 
     def test_pipeline_param_added_to_channel_inputs(self):
         """When a PipelineParameterChannel is passed as secret_name,
-        it should be appended to the task's _channel_inputs."""
+        it should be appended to the task's channel_inputs."""
 
         @dsl.pipeline
         def my_pipeline(secret_name: str):
             task = comp()
-            initial_count = len(task._channel_inputs)
+            initial_count = len(task.channel_inputs)
             kubernetes.use_secret_as_env(
                 task,
                 secret_name=secret_name,
                 secret_key_to_env={"key": "VAR"},
             )
-            assert len(task._channel_inputs) == initial_count + 1
-            channel_patterns = {ch.pattern for ch in task._channel_inputs}
+            assert len(task.channel_inputs) == initial_count + 1
+            channel_patterns = {ch.pattern for ch in task.channel_inputs}
             assert secret_name.pattern in channel_patterns
 
     def test_duplicate_pipeline_param_not_added_twice(self):
         """When the same PipelineParameterChannel is used in both
-        use_secret_as_env and use_secret_as_volume, _channel_inputs
+        use_secret_as_env and use_secret_as_volume, channel_inputs
         should not contain duplicates."""
 
         @dsl.pipeline
@@ -1006,44 +1006,44 @@ def my_pipeline(secret_name: str):
                 secret_name=secret_name,
                 secret_key_to_env={"key1": "VAR1"},
             )
-            count_after_first = len(task._channel_inputs)
+            count_after_first = len(task.channel_inputs)
             kubernetes.use_secret_as_volume(
                 task,
                 secret_name=secret_name,
                 mount_path="/mnt/secret",
             )
-            assert len(task._channel_inputs) == count_after_first
+            assert len(task.channel_inputs) == count_after_first
 
     def test_string_input_does_not_add_channel_inputs(self):
-        """When a literal string is passed, _channel_inputs should not
+        """When a literal string is passed, channel_inputs should not
         be modified."""
 
         @dsl.pipeline
         def my_pipeline():
             task = comp()
-            initial_count = len(task._channel_inputs)
+            initial_count = len(task.channel_inputs)
             kubernetes.use_secret_as_env(
                 task,
                 secret_name="literal-secret",
                 secret_key_to_env={"key": "VAR"},
             )
-            assert len(task._channel_inputs) == initial_count
+            assert len(task.channel_inputs) == initial_count
 
     def test_task_output_added_to_channel_inputs(self):
         """When a task output PipelineChannel is passed as secret_name,
-        it should be appended to _channel_inputs."""
+        it should be appended to channel_inputs."""
 
         @dsl.pipeline
         def my_pipeline():
             name_task = comp_with_output()
             task = comp()
-            initial_count = len(task._channel_inputs)
+            initial_count = len(task.channel_inputs)
             kubernetes.use_secret_as_env(
                 task,
                 secret_name=name_task.output,
                 secret_key_to_env={"key": "VAR"},
             )
-            assert len(task._channel_inputs) == initial_count + 1
+            assert len(task.channel_inputs) == initial_count + 1
 
 
 @dsl.component

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -1698,6 +1698,12 @@ def _rewrite(data: Any) -> Any:
                 if surfaced_name in parent_component_inputs.parameters:
                     rewritten.pop('taskOutputParameter')
                     rewritten['componentInputParameter'] = surfaced_name
+                else:
+                    raise compiler_utils.InvalidTopologyException(
+                        'Failed to rewrite cross-DAG platform config reference '
+                        f'for task output {producer_task}.{output_key}. '
+                        f'Expected surfaced input {surfaced_name!r} in parent '
+                        'component inputs, but it was not found.')
 
         return rewritten
 

sdk/python/kfp/compiler/pipeline_spec_builder_test.py

@@ -24,6 +24,7 @@
 from kfp import dsl
 from kfp import kubernetes
 from kfp.compiler import compiler
+from kfp.compiler import compiler_utils
 from kfp.compiler import pipeline_spec_builder
 from kfp.dsl import TaskConfigField
 from kfp.pipeline_spec import pipeline_spec_pb2
@@ -724,6 +725,82 @@ def pipe():
         self.assertNotIn('componentInputParameter', secret_name_fields)
         self.assertNotIn('taskOutputParameter', secret_name_fields)
 
+    def test_parallelfor_pipeline_input_mount_pvc(self):
+        """Pipeline pvc_name param inside ParallelFor is correctly surfaced and
+        rewritten for mount_pvc platform config."""
+
+        @dsl.component
+        def my_comp(item: str):
+            print(item)
+
+        @dsl.pipeline
+        def pipe(pvc_name: str):
+            with dsl.ParallelFor(items=['a', 'b'], parallelism=1) as item:
+                t = my_comp(item=item)
+                kubernetes.mount_pvc(
+                    t,
+                    pvc_name=pvc_name,
+                    mount_path='/mnt/data',
+                )
+
+        pipeline_spec, platform_spec = self._compile_and_parse(pipe)
+
+        loop_component = pipeline_spec.components['comp-for-loop-2']
+        self.assertIn('pipelinechannel--pvc_name',
+                      loop_component.input_definitions.parameters)
+
+        root_task_params = pipeline_spec.root.dag.tasks[
+            'for-loop-2'].inputs.parameters
+        self.assertEqual(
+            root_task_params['pipelinechannel--pvc_name']
+            .component_input_parameter,
+            'pvc_name',
+        )
+
+        pvc_param = (
+            platform_spec.platforms['kubernetes'].deployment_spec
+            .executors['exec-my-comp'].fields['pvcMount'].list_value.values[0]
+            .struct_value.fields['pvcNameParameter'].struct_value
+            .fields['componentInputParameter'].string_value)
+        self.assertEqual(pvc_param, 'pipelinechannel--pvc_name')
+
+    def test_exit_handler_platform_config_rewrite_path(self):
+        """Exit handler task platform config uses rewrite path with parent
+        component context."""
+
+        @dsl.component
+        def cleanup():
+            print('cleanup')
+
+        @dsl.component
+        def main_task():
+            print('main')
+
+        @dsl.pipeline
+        def pipe(secret_name: str):
+            exit_task = cleanup()
+            kubernetes.use_secret_as_env(
+                exit_task,
+                secret_name=secret_name,
+                secret_key_to_env={'key': 'VAL'},
+            )
+            with dsl.ExitHandler(exit_task=exit_task):
+                main_task()
+
+        _, platform_spec = self._compile_and_parse(pipe)
+
+        cleanup_executors = [
+            executor for executor in platform_spec.platforms['kubernetes']
+            .deployment_spec.executors.values()
+            if 'secretAsEnv' in executor.fields
+        ]
+        self.assertEqual(len(cleanup_executors), 1)
+        secret_param = (
+            cleanup_executors[0].fields['secretAsEnv'].list_value.values[0]
+            .struct_value.fields['secretNameParameter'].struct_value
+            .fields['componentInputParameter'].string_value)
+        self.assertEqual(secret_param, 'secret_name')
+
 
 class TestRewritePlatformConfigInputReferences(unittest.TestCase):
     """Unit tests for the _rewrite_platform_config_input_references helper."""
@@ -878,6 +955,29 @@ def test_rewrites_multiple_params(self):
             'pipelinechannel--pvc_name',
         )
 
+    def test_raises_when_cross_dag_output_missing_surfaced_input(self):
+        platform_config = {
+            'kubernetes': {
+                'secretAsEnv': [{
+                    'secretNameParameter': {
+                        'taskOutputParameter': {
+                            'producerTask': 'emit-secret',
+                            'outputParameterKey': 'Output',
+                        }
+                    },
+                }]
+            }
+        }
+        parent_inputs = pipeline_spec_pb2.ComponentInputsSpec()
+
+        with self.assertRaisesRegex(compiler_utils.InvalidTopologyException,
+                                    'Expected surfaced input'):
+            pipeline_spec_builder._rewrite_platform_config_input_references(
+                platform_config,
+                parent_inputs,
+                tasks_in_current_dag=['worker-task'],
+            )
+
 
 def pipeline_spec_from_file(filepath: str) -> str:
     with open(filepath, 'r') as f:

sdk/python/kfp/dsl/pipeline_task.py

@@ -328,6 +328,7 @@ def _register_pipeline_channels(
             pipeline_channels: List[pipeline_channel.PipelineChannel]) -> None:
         """Backwards-compatible wrapper for ``register_pipeline_channels``."""
         self.register_pipeline_channels(pipeline_channels)
+
     @block_if_final()
     def set_caching_options(self,
                             enable_caching: bool,

test_data/compiled-workflows/nested_parallel_for_secret.yaml

@@ -21,9 +21,9 @@ spec:
         \"test-secret-1\"\n\n"],"image":"python:3.11"}'
     - name: kubernetes-comp-worker-component
       value: '{"secretAsEnv":[{"keyToEnv":[{"envVar":"MY_SECRET_KEY","secretKey":"username"}],"optional":false,"secretNameParameter":{"componentInputParameter":"pipelinechannel--emit-secret-name-Output"}}]}'
-    - name: components-490eb648ead5e4c6f02bebcc44125df7cf0178498ee770cc67d4bd310900a4fe
+    - name: components-3b4506509514761c798c91195d4caf336df6411cbc2327c68dd7902b1bcad98c
       value: '{"executorLabel":"exec-worker-component","inputDefinitions":{"parameters":{"item":{"parameterType":"STRING"}}},"outputDefinitions":{"parameters":{"Output":{"parameterType":"STRING"}}}}'
-    - name: implementations-490eb648ead5e4c6f02bebcc44125df7cf0178498ee770cc67d4bd310900a4fe
+    - name: implementations-3b4506509514761c798c91195d4caf336df6411cbc2327c68dd7902b1bcad98c
       value: '{"args":["--executor_input","{{$}}","--function_to_execute","worker_component"],"command":["sh","-c","\nif
         ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip || python3
         -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1
@@ -35,7 +35,7 @@ spec:
         kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import *\n\ndef
         worker_component(item: str) -\u003e str:\n    import os\n    secret_val =
         os.environ.get(\"MY_SECRET_KEY\", \"not-set\")\n    print(f\"Item: {item},
-        Secret value: {secret_val}\")\n    return secret_val\n\n"],"image":"python:3.11"}'
+        Secret set: {secret_val != ''not-set''}\")\n    return secret_val\n\n"],"image":"python:3.11"}'
     - name: components-comp-for-loop-2
       value: '{"dag":{"tasks":{"worker-component":{"cachingOptions":{"enableCache":true},"componentRef":{"name":"comp-worker-component"},"inputs":{"parameters":{"item":{"componentInputParameter":"pipelinechannel--loop-item-param-1"}}},"taskInfo":{"name":"worker-component"}}}},"inputDefinitions":{"parameters":{"pipelinechannel--emit-secret-name-Output":{"parameterType":"STRING"},"pipelinechannel--loop-item-param-1":{"parameterType":"STRING"}}}}'
     - name: components-root
@@ -272,11 +272,11 @@ spec:
       - arguments:
           parameters:
           - name: component
-            value: '{{workflow.parameters.components-490eb648ead5e4c6f02bebcc44125df7cf0178498ee770cc67d4bd310900a4fe}}'
+            value: '{{workflow.parameters.components-3b4506509514761c798c91195d4caf336df6411cbc2327c68dd7902b1bcad98c}}'
           - name: task
             value: '{"cachingOptions":{"enableCache":true},"componentRef":{"name":"comp-worker-component"},"inputs":{"parameters":{"item":{"componentInputParameter":"pipelinechannel--loop-item-param-1"}}},"taskInfo":{"name":"worker-component"}}'
           - name: container
-            value: '{{workflow.parameters.implementations-490eb648ead5e4c6f02bebcc44125df7cf0178498ee770cc67d4bd310900a4fe}}'
+            value: '{{workflow.parameters.implementations-3b4506509514761c798c91195d4caf336df6411cbc2327c68dd7902b1bcad98c}}'
           - name: task-name
             value: worker-component
           - name: parent-dag-id