Repository: kubeflow/pipelines

Issue:
Improve stack traces of lightweight components
Currently the code for the Lightweight python components is executed using `python3 -c <code>`. This has good portability (only needs `python3`), but causes an issue with exception stack traces. Python does not seem to show any stack traces for the inline code.

One solution  is to first write the code to a temporary file and then execute that file.


Discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

This issue has been automatically closed because it has not had recent activity. Please comment "/reopen" to reopen it.


Resolving patch:
diff --git a/sdk/python/kfp/components/_python_op.py b/sdk/python/kfp/components/_python_op.py
--- a/sdk/python/kfp/components/_python_op.py
+++ b/sdk/python/kfp/components/_python_op.py
@@ -33,6 +33,7 @@
 
 import inspect
 from pathlib import Path
+import textwrap
 from typing import Callable, List, Optional, TypeVar
 import warnings
 
@@ -234,8 +235,6 @@ def get_fixers(self):
 
 
 def _capture_function_code_using_source_copy(func) -> str:	
-    import textwrap
-
     func_code = inspect.getsource(func)
 
     #Function might be defined in some indented scope (e.g. in another function).
@@ -640,7 +639,18 @@ def get_serializer_and_register_definitions(type_name) -> str:
     component_spec.implementation=ContainerImplementation(
         container=ContainerSpec(
             image=base_image,
-            command=package_preinstallation_command + ['python3', '-u', '-c', full_source],
+            command=package_preinstallation_command + [
+                'sh',
+                '-ec',
+                # Writing the program code to a file.
+                # This is needed for Python to show stack traces and for `inspect.getsource` to work (used by PyTorch JIT and this module for example).
+                textwrap.dedent('''\
+                    program_path=$(mktemp)
+                    echo -n "$0" > "$program_path"
+                    python3 -u "$program_path" "$@"
+                '''),
+                full_source,
+            ],
             args=arguments,
         )
     )
diff --git a/sdk/python/kfp/containers/_component_builder.py b/sdk/python/kfp/containers/_component_builder.py
--- a/sdk/python/kfp/containers/_component_builder.py
+++ b/sdk/python/kfp/containers/_component_builder.py
@@ -217,24 +217,25 @@ def build_python_component(component_func, target_image, base_image=None, depend
   component_spec = _func_to_component_spec(component_func, base_image=base_image)
   command_line_args = component_spec.implementation.container.command
 
-  dash_c_index = command_line_args.index('-c')
-  program_code_index = dash_c_index + 1
+  program_launcher_index = command_line_args.index('program_path=$(mktemp)\necho -n "$0" > "$program_path"\npython3 -u "$program_path" "$@"\n')
+  assert program_launcher_index in [2, 3]
+
+  program_code_index = program_launcher_index + 1
   program_code = command_line_args[program_code_index]
   program_rel_path = 'ml/main.py'
   program_container_path = '/' + program_rel_path
 
   # Replacing the inline code with calling a local program
-  # Before: python3 -u -c 'import sys ...' --param1 ...
+  # Before: sh -ec '... && python3 -u ...' 'import sys ...' --param1 ...
   # After:  python3 -u main.py --param1 ...
   command_line_args[program_code_index] = program_container_path
-  command_line_args.pop(dash_c_index)
+  command_line_args.pop(program_launcher_index)
+  command_line_args[program_launcher_index - 1] = '-u'  # -ec => -u
+  command_line_args[program_launcher_index - 2] = python_version  # sh => python3
 
   if python_version == 'python2':
     import warnings
     warnings.warn('Python2 is not longer supported')
-    # Replacing the python interpreter
-    python_interpreter_index = command_line_args.index('python3')
-    command_line_args[python_interpreter_index] = python_version
 
   arc_docker_filename = 'Dockerfile'
   arc_requirement_filename = 'requirements.txt'


Test patch:
diff --git a/sdk/python/tests/compiler/testdata/parallelfor_item_argument_resolving.yaml b/sdk/python/tests/compiler/testdata/parallelfor_item_argument_resolving.yaml
--- a/sdk/python/tests/compiler/testdata/parallelfor_item_argument_resolving.yaml
+++ b/sdk/python/tests/compiler/testdata/parallelfor_item_argument_resolving.yaml
@@ -16,9 +16,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-strings-Output}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -44,9 +47,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-strings-Output-loop-item}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -72,9 +78,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-str-Output}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -100,9 +109,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-ints-Output}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -128,9 +140,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-ints-Output-loop-item}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -156,9 +171,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-dicts-Output}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -184,9 +202,12 @@ spec:
           - "--param1"
           - "{{inputs.parameters.produce-list-of-dicts-Output-loop-item-subvar-aaa}}"
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def consume(param1):
                   print(param1)
@@ -362,9 +383,12 @@ spec:
           - "----output-paths"
           - /tmp/outputs/Output/data
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def produce_list_of_dicts():
                   return ([{"aaa": "aaa1", "bbb": "bbb1"}, {"aaa": "aaa2", "bbb": "bbb2"}],)
@@ -424,9 +448,12 @@ spec:
           - "----output-paths"
           - /tmp/outputs/Output/data
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def produce_list_of_ints():
                   return ([1234567890, 987654321],)
@@ -486,9 +513,12 @@ spec:
           - "----output-paths"
           - /tmp/outputs/Output/data
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def produce_list_of_strings():
                   return (["a", "z"],)
@@ -548,9 +578,12 @@ spec:
           - "----output-paths"
           - /tmp/outputs/Output/data
         command: 
-          - python3
-          - "-u"
-          - "-c"
+          - sh
+          - "-ec"
+          - |
+              program_path=$(mktemp)
+              echo -n "$0" > "$program_path"
+              python3 -u "$program_path" "$@"
           - |
               def produce_str():
                   return "Hello"