fix(sdk): Fix environment variable set in component yaml lost during …

…compilation. Fixes #8884 (#8885) * Fix env set in component yaml not captured * release note * clean up unnecessary code in test
kubeflow · Feb 23, 2023 · 8d3e0dd · 8d3e0dd
1 parent 386ff88
commit 8d3e0dd
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 39 deletions.
diff --git a/sdk/RELEASE.md b/sdk/RELEASE.md
@@ -12,6 +12,7 @@
 ## Bug fixes and other changes
 * Enables output definitions when compiling components as pipelines. [\#8848](https://github.com/kubeflow/pipelines/pull/8848)
 * Fix bug when passing data between tasks using f-strings [\#8879](https://github.com/kubeflow/pipelines/pull/8879)
+* Fix environment variable set in component yaml lost during compilation [\#8885](https://github.com/kubeflow/pipelines/pull/8885)
 
 ## Documentation updates
 

diff --git a/sdk/python/kfp/components/structures.py b/sdk/python/kfp/components/structures.py
@@ -343,7 +343,7 @@ def from_container_dict(
             image=container_dict['image'],
             command=container_dict.get('command'),
             args=container_dict.get('args'),
-            env=None,  # can only be set on tasks
+            env=container_dict.get('env'),
             resources=None)  # can only be set on tasks
 
 

diff --git a/sdk/python/kfp/components/structures_test.py b/sdk/python/kfp/components/structures_test.py
@@ -436,49 +436,27 @@ def test_env(self):
         self.assertEqual(obj.env, None)
 
     def test_from_container_dict_no_placeholders(self):
-        component_spec = structures.ComponentSpec(
-            name='test',
-            implementation=structures.Implementation(
-                container=structures.ContainerSpecImplementation(
-                    image='python:3.7',
-                    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)\nprintf "%s" "$0" > "$program_path/ephemeral_component.py"\npython3 -m kfp.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"\n',
-                        '\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import *\n\ndef concat_message(first: str, second: str) -> str:\n    return first + second\n\n'
-                    ],
-                    args=[
-                        '--executor_input',
-                        placeholders.ExecutorInputPlaceholder(),
-                        '--function_to_execute', 'concat_message'
-                    ],
-                    env=None,
-                    resources=None),
-                graph=None,
-                importer=None),
-            description=None,
-            inputs={
-                'first': structures.InputSpec(type='String', default=None),
-                'second': structures.InputSpec(type='String', default=None)
-            },
-            outputs={'Output': structures.OutputSpec(type='String')})
+        expected_container_spec = structures.ContainerSpecImplementation(
+            image='python:3.7',
+            command=['sh', '-c', 'dummy'],
+            args=['--executor_input', '{{$}}', '--function_to_execute', 'func'],
+            env={'ENV1': 'val1'},
+            resources=None)
+
         container_dict = {
             'args': [
-                '--executor_input', '{{$}}', '--function_to_execute', 'fail_op'
+                '--executor_input', '{{$}}', '--function_to_execute', 'func'
             ],
-            '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)\nprintf "%s" "$0" > "$program_path/ephemeral_component.py"\npython3 -m kfp.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"\n',
-                '\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import *\n\ndef fail_op(message: str):\n    """Fails."""\n    import sys\n    print(message)\n    sys.exit(1)\n\n'
-            ],
-            'image': 'python:3.7'
+            'command': ['sh', '-c', 'dummy'],
+            'image': 'python:3.7',
+            'env': {
+                'ENV1': 'val1'
+            },
         }
 
         loaded_container_spec = structures.ContainerSpecImplementation.from_container_dict(
             container_dict)
+        self.assertEqual(expected_container_spec, loaded_container_spec)
 
     def test_raise_error_if_access_artifact_by_itself(self):
 

diff --git a/sdk/python/test_data/pipelines/pipeline_with_env.py b/sdk/python/test_data/pipelines/pipeline_with_env.py
@@ -35,9 +35,11 @@ def print_env_op():
     - -c
     - |
       set -e -x
+      echo "$ENV1"
       echo "$ENV2"
       echo "$ENV3"
     env:
+      ENV1: val0
       ENV2: val0
 """)
 

diff --git a/sdk/python/test_data/pipelines/pipeline_with_env.yaml b/sdk/python/test_data/pipelines/pipeline_with_env.yaml
@@ -15,12 +15,16 @@ deploymentSpec:
         - -c
         - 'set -e -x
 
+          echo "$ENV1"
+
           echo "$ENV2"
 
           echo "$ENV3"
 
           '
         env:
+        - name: ENV1
+          value: val0
         - name: ENV2
           value: val2
         - name: ENV3
@@ -38,7 +42,7 @@ deploymentSpec:
         - -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-beta.8'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.12'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -76,4 +80,4 @@ root:
         taskInfo:
           name: print-env-op
 schemaVersion: 2.1.0
-sdkVersion: kfp-2.0.0-beta.8
+sdkVersion: kfp-2.0.0-beta.12