fix(bug): fix underscore param by checking pipeline param instead of task param (kubeflow#1077)

* fix underscore param by checking pipeline param instead of task param

* update test cases

Author: Tomcli

sdk/python/kfp_tekton/compiler/_tekton_handler.py

@@ -315,10 +315,10 @@ def _handle_tekton_custom_task(custom_task: dict, workflow: dict, recursive_task
                                     '$(params.%s)' % param_result[1]))
                             else:
                                 new_param_ref = '%s-%s' % param_result
-                                # Nested task reference name must be the same as the one in the params
+                                # Nested task reference name must be the same as the one in the pipeline params
                                 # Otherwise, use the default param ref since it will be generated
                                 # by the custom task functions such as condition and loop
-                                for i in task.get('params', []):
+                                for i in custom_task_cr['spec']['pipelineSpec'].get('params', []):
                                     if sanitize_k8s_name(i['name']) == new_param_ref:
                                         new_param_ref = i['name']
                                         break

sdk/python/tests/compiler/testdata/nested_loop_with_underscore.py

@@ -12,9 +12,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from kfp import dsl
+import itertools
+import yaml
+from kfp import dsl, components
 from kfp.components import load_component_from_text
 from kfp_tekton.compiler import TektonCompiler
+from kfp_tekton.tekton import TEKTON_CUSTOM_TASK_IMAGES
+
+
+ARTIFACT_FETCHER_IMAGE_NAME = "fetcher/image:latest"
+TEKTON_CUSTOM_TASK_IMAGES = TEKTON_CUSTOM_TASK_IMAGES.append(ARTIFACT_FETCHER_IMAGE_NAME)
+
+_artifact_fetcher_no = 0
+
+
+def artifact_fetcher(**artifact_paths: str):
+  '''A containerOp template resolving some artifacts, given their paths.'''
+  global _artifact_fetcher_no
+  template_yaml = {
+    'name': f'artifact-fetcher-{_artifact_fetcher_no}',
+    'description': 'Artifact Fetch',
+    'inputs': [
+      {'name': name, 'type': 'String'}
+      for name in artifact_paths.keys()
+    ],
+    'outputs': [
+      {'name': name, 'type': 'Artifact'}
+      for name in artifact_paths.keys()
+    ],
+    'implementation': {
+      'container': {
+        'image': ARTIFACT_FETCHER_IMAGE_NAME,
+        'command': ['sh', '-c'],  # irrelevant
+        'args': [
+          '--apiVersion', 'fetcher.tekton.dev/v1alpha1',
+          '--kind', 'FETCHER',
+          '--name', 'artifact_fetcher',
+          *itertools.chain(*[
+            (f'--{name}', {'inputValue': name})
+            for name in artifact_paths.keys()
+          ])
+        ]
+      }
+    }
+  }
+  _artifact_fetcher_no += 1
+  template_str = yaml.dump(template_yaml, Dumper=yaml.SafeDumper)
+  template = components.load_component_from_text(template_str)
+  op = template(**artifact_paths)
+  op.add_pod_annotation("valid_container", "false")
+  return op
 
 
 class Coder:
@@ -57,6 +104,7 @@ def double_loop_with_underscore(param_a: list = [1, 2, 3], param_b: list = ["a",
   with dsl.ParallelFor(param_a):
     with dsl.ParallelFor(param_b):
       op1 = PrintOp('print-1', f"print {op0.output}")
+      op2 = artifact_fetcher(path=op0.output)
 
 
 if __name__ == '__main__':

sdk/python/tests/compiler/testdata/nested_loop_with_underscore.yaml

@@ -150,6 +150,14 @@ spec:
                               "print-1", "outputs": [{"description": "Represents an
                               output paramter.", "name": "output_value", "type": "String"}],
                               "version": "print-1@sha256=3b81342bc143f625b58ebdb01e7c83b145880dee807be35c1e16fdb835d46580"}'
+                    - name: artifact-fetcher-0
+                      params:
+                      - name: path
+                        value: $(params.print-0-output_value)
+                      taskRef:
+                        name: artifact_fetcher
+                        apiVersion: fetcher.tekton.dev/v1alpha1
+                        kind: FETCHER
                   iterateParam: param_b-loop-item
                 metadata:
                   labels:

sdk/python/tests/compiler/testdata/nested_loop_with_underscore_noninlined.yaml

@@ -57,7 +57,9 @@ metadata:
       "print-0-output_value", "type": "string"}], "results": [{"description": "/tmp/outputs/output_value/data",
       "name": "output-value", "type": "string"}], "steps": [{"command": ["sh", "-c",
       "set -e\necho $0 > $1\n", "print $(inputs.params.print-0-output_value)", "$(results.output-value.path)"],
-      "image": "alpine:3.6", "name": "main"}]}}]}}}]'
+      "image": "alpine:3.6", "name": "main"}]}}, {"name": "artifact-fetcher-1", "params":
+      [{"name": "path", "value": "$(params.print-0-output_value)"}], "taskRef": {"apiVersion":
+      "fetcher.tekton.dev/v1alpha1", "kind": "FETCHER", "name": "artifact_fetcher"}}]}}}]'
   labels:
     pipelines.kubeflow.org/pipelinename: ''
     pipelines.kubeflow.org/generation: ''