fix(sdk): Update custom task arg parsing for dash arguments (kubeflow#1114)

Author: Tomcli

sdk/python/kfp_tekton/compiler/compiler.py

@@ -1025,12 +1025,18 @@ def map_cel_vars(a):
             container_args = i.get('args', [])
             custom_task_command = {}
             container_command = i.get('command', [])
-            for index, item in enumerate(container_args):
-              if item.startswith('--'):
-                custom_task_args[item[2:]] = container_args[index + 1]
-            for index, item in enumerate(container_command):
-              if item.startswith('--'):
-                custom_task_command[item[2:]] = container_command[index + 1]
+
+            def find_parameter(arg_list, arg_task_list):
+              skip_index = False
+              for index, item in enumerate(arg_list):
+                if skip_index:
+                  skip_index = False
+                  continue
+                if item.startswith('--'):
+                  arg_task_list[item[2:]] = arg_list[index + 1]
+                  skip_index = True
+            find_parameter(container_args, custom_task_args)
+            find_parameter(container_command, custom_task_command)
             non_param_keys = ['name', 'apiVersion', 'kind', 'taskSpec', 'taskRef']
             task_params = []
             command_params = []

sdk/python/tests/compiler/compiler_tests.py

@@ -134,6 +134,13 @@ def test_custom_task_output_workflow(self):
     from .testdata.custom_task_output import uppercase_vs_lowercase
     self._test_pipeline_workflow(uppercase_vs_lowercase, 'custom_task_output.yaml', skip_noninlined=True)
 
+  def test_custom_task_dash_args_workflow(self):
+    """
+    Test compiling a custom task with dash arguments workflow.
+    """
+    from .testdata.custom_task_dash_args import pipeline_dash_args
+    self._test_pipeline_workflow(pipeline_dash_args, 'custom_task_dash_args.yaml', skip_noninlined=True)
+
   def test_parallel_join_workflow(self):
     """
     Test compiling a parallel join workflow.

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

@@ -0,0 +1,45 @@
+# Copyright 2023 kubeflow.org
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from kfp import dsl
+from kfp.components import load_component_from_text
+
+
+@dsl.pipeline()
+def pipeline_dash_args():
+    cel_run_bash_script = load_component_from_text(r"""
+        name: cel-run-bash-script
+        inputs: []
+        outputs:
+        - {name: env-variables}
+        implementation:
+          container:
+            image: aipipeline/cel-eval:latest
+            command: [cel]
+            args: [--apiVersion, custom.tekton.dev/v1alpha1, --kind, VariableStore, --name,
+              1234567-var-store, --lit_0, '---', --taskSpec,
+              '{}']
+            fileOutputs: {}
+    """)()
+
+    cel_run_bash_script.add_pod_annotation("valid_container", "false")
+
+
+if __name__ == '__main__':
+  from kfp_tekton.compiler import TektonCompiler as Compiler
+  from kfp_tekton.compiler.pipeline_utils import TektonPipelineConf
+  tekton_pipeline_conf = TektonPipelineConf()
+  tekton_pipeline_conf.set_tekton_inline_spec(True)
+  Compiler().compile(pipeline_dash_args, __file__.replace('.py', '.yaml'), tekton_pipeline_conf=tekton_pipeline_conf, type_check=False)

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

@@ -0,0 +1,50 @@
+# Copyright 2021 kubeflow.org
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: tekton.dev/v1beta1
+kind: PipelineRun
+metadata:
+  name: pipeline-dash-args
+  annotations:
+    tekton.dev/output_artifacts: '{}'
+    tekton.dev/input_artifacts: '{}'
+    tekton.dev/artifact_bucket: mlpipeline
+    tekton.dev/artifact_endpoint: minio-service.kubeflow:9000
+    tekton.dev/artifact_endpoint_scheme: http://
+    tekton.dev/artifact_items: '{}'
+    sidecar.istio.io/inject: "false"
+    tekton.dev/template: ''
+    pipelines.kubeflow.org/big_data_passing_format: $(workspaces.$TASK_NAME.path)/artifacts/$ORIG_PR_NAME/$TASKRUN_NAME/$TASK_PARAM_NAME
+    pipelines.kubeflow.org/pipeline_spec: '{"name": "Pipeline dash args"}'
+  labels:
+    pipelines.kubeflow.org/pipelinename: ''
+    pipelines.kubeflow.org/generation: ''
+spec:
+  pipelineSpec:
+    tasks:
+    - name: cel-run-bash-script
+      params:
+      - name: lit_0
+        value: '---'
+      taskSpec:
+        apiVersion: custom.tekton.dev/v1alpha1
+        kind: VariableStore
+        spec: {}
+        metadata:
+          labels:
+            pipelines.kubeflow.org/cache_enabled: "true"
+          annotations:
+            valid_container: "false"
+            pipelines.kubeflow.org/component_spec_digest: '{"name": "cel-run-bash-script",
+              "outputs": [{"name": "env-variables"}], "version": "cel-run-bash-script@sha256=ea3b98e707ecd49f3cd88c4abde7b9dedc0154dd4c1b10a3a3feff3b85fd82ad"}'