fix traversal func

kubeflow · May 4, 2022 · a7e839d · a7e839d
1 parent 9f8e541
commit a7e839d
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 52 deletions.
diff --git a/sdk/python/kfp/compiler_cli_tests/test_data/pipeline_with_importer_dynamic_variables.yaml b/sdk/python/kfp/compiler_cli_tests/test_data/pipeline_with_importer_dynamic_variables.yaml
@@ -19,7 +19,7 @@ components:
     executorLabel: exec-importer-2
     inputDefinitions:
       parameters:
-        metadata_0:
+        metadata_2:
           parameterType: STRING
         uri:
           parameterType: STRING
@@ -48,8 +48,8 @@ deploymentSpec:
           runtimeParameter: uri
         metadata:
           containerSpec:
-            imageUri: "{{$.inputs.parameters['metadata_1']}}"
-          name: "{{$.inputs.parameters['metadata_0']}}"
+            imageUri: '{{$.inputs.parameters[''metadata_1'']}}'
+          name: '{{$.inputs.parameters[''metadata_0'']}}'
         typeSchema:
           schemaTitle: system.Artifact
           schemaVersion: 0.0.1
@@ -58,36 +58,37 @@ deploymentSpec:
         artifactUri:
           constant: gs://ml-pipeline-playground/shakespeare1.txt
         metadata:
-          containerSpec: {}
-          name: "{{$.inputs.parameters['metadata_0']}}"
+          containerSpec:
+            imageUri: us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-5:latest
+          name: '{{$.inputs.parameters[''metadata_2'']}}'
         typeSchema:
           schemaTitle: system.Artifact
           schemaVersion: 0.0.1
     exec-make-name:
       container:
         args:
-          - --executor_input
-          - "{{$}}"
-          - --function_to_execute
-          - make_name
+        - --executor_input
+        - '{{$}}'
+        - --function_to_execute
+        - make_name
         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\
-            \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-alpha.2'\
-            \ && \"$0\" \"$@\"\n"
-          - sh
-          - -ec
-          - 'program_path=$(mktemp -d)
+        - 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\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-alpha.2'\
+          \ && \"$0\" \"$@\"\n"
+        - sh
+        - -ec
+        - 'program_path=$(mktemp -d)
 
-            printf "%s" "$0" > "$program_path/ephemeral_component.py"
+          printf "%s" "$0" > "$program_path/ephemeral_component.py"
 
-            python3 -m kfp.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
+          python3 -m kfp.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
 
-            '
-          - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
-            \ *\n\ndef make_name(name: str) -> str:\n    return name\n\n"
+          '
+        - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
+          \ *\n\ndef make_name(name: str) -> str:\n    return name\n\n"
         image: python:3.7
 pipelineInfo:
   name: pipeline-with-importer
@@ -115,10 +116,10 @@ root:
         componentRef:
           name: comp-importer-2
         dependentTasks:
-          - make-name
+        - make-name
         inputs:
           parameters:
-            metadata_0:
+            metadata_2:
               taskOutputParameter:
                 outputParameterKey: Output
                 producerTask: make-name

diff --git a/sdk/python/kfp/components/importer_node.py b/sdk/python/kfp/components/importer_node.py
@@ -25,6 +25,7 @@
 INPUT_KEY = 'uri'
 OUTPUT_KEY = 'artifact'
 METADATA_KEY_PREFIX = 'metadata_'
+_metadata_counter = 0
 
 
 def importer(
@@ -50,39 +51,33 @@ def importer(
         constant string value.
     """
     metadata_inputs = {}
-    _metadata_counter = 0
     call_inputs = {}
 
     def traverse_dict_and_create_metadata_inputs(d: Any) -> Dict:
         nonlocal metadata_inputs
-        nonlocal _metadata_counter
         nonlocal call_inputs
-        res = {}
-        for k, v in d.items():
-            if isinstance(k, pipeline_channel.PipelineParameterChannel):
-                res[k] = placeholders.input_parameter_placeholder(
-                    METADATA_KEY_PREFIX + str(_metadata_counter))
-                metadata_inputs[METADATA_KEY_PREFIX +
-                                str(_metadata_counter)] = structures.InputSpec(
-                                    type='String')
-                call_inputs[METADATA_KEY_PREFIX + str(_metadata_counter)] = v
-                _metadata_counter += 1
-            if isinstance(v, dict):
-                res[k] = traverse_dict_and_create_metadata_inputs(v)
-            elif isinstance(v, list):
-                res[k] = [
-                    traverse_dict_and_create_metadata_inputs(el) for el in v
-                ]
-            elif isinstance(v, pipeline_channel.PipelineParameterChannel):
-                res[k] = placeholders.input_parameter_placeholder(
-                    METADATA_KEY_PREFIX + str(_metadata_counter))
-                metadata_inputs[METADATA_KEY_PREFIX +
-                                str(_metadata_counter)] = structures.InputSpec(
-                                    type='String')
-                call_inputs[METADATA_KEY_PREFIX + str(_metadata_counter)] = v
-                _metadata_counter += 1
+        global _metadata_counter
 
-        return res
+        if isinstance(d, pipeline_channel.PipelineParameterChannel):
+            metadata_inputs[METADATA_KEY_PREFIX +
+                            str(_metadata_counter)] = structures.InputSpec(
+                                type='String')
+            call_inputs[METADATA_KEY_PREFIX + str(_metadata_counter)] = d
+            result = placeholders.input_parameter_placeholder(
+                METADATA_KEY_PREFIX + str(_metadata_counter))
+            _metadata_counter += 1
+            return result
+        elif isinstance(d, dict):
+            return {
+                traverse_dict_and_create_metadata_inputs(k):
+                traverse_dict_and_create_metadata_inputs(v)
+                for k, v in d.items()
+            }
+
+        elif isinstance(d, list):
+            return [traverse_dict_and_create_metadata_inputs(el) for el in d]
+        else:
+            return d
 
     metadata_with_placeholders = traverse_dict_and_create_metadata_inputs(
         metadata)