SDK/Components/PythionContainerOp - Make the local output path config…

…urableThis is part of a bigger effort to make all output locations manageable in preparation for storage system.

sdk/python/kfp/compiler/_component_builder.py

@@ -306,9 +306,7 @@ def _generate_entrypoint(self, component_func):
     func_signature = 'def ' + new_func_name + '('
     for input_arg in input_args:
       func_signature += input_arg + ','
-    if len(input_args) > 0:
-      func_signature = func_signature[:-1]
-    func_signature += '):'
+    func_signature += '_output_file):'
     codegen.writeline(func_signature)
 
     # Call user function
@@ -321,9 +319,9 @@ def _generate_entrypoint(self, component_func):
     codegen.writeline(call_component_func)
 
     # Serialize output
-    codegen.writeline('with open("/output.txt", "w") as f:')
-    codegen.indent()
-    codegen.writeline('f.write(str(output))')
+    codegen.writeline('from pathlib import Path')
+    codegen.writeline('Path(_output_file).parent.mkdir(parents=True, exist_ok=True)')
+    codegen.writeline('Path(_output_file).write_text(str(output))')
     wrapper_code = codegen.end()
 
     # CLI codes
@@ -332,6 +330,7 @@ def _generate_entrypoint(self, component_func):
     codegen.writeline('parser = argparse.ArgumentParser(description="Parsing arguments")')
     for input_arg in input_args:
       codegen.writeline('parser.add_argument("' + input_arg + '", type=' + inputs[input_arg].__name__ + ')')
+    codegen.writeline('parser.add_argument("_output_file", type=str)')
     codegen.writeline('args = vars(parser.parse_args())')
     codegen.writeline('')
     codegen.writeline('if __name__ == "__main__":')
@@ -430,7 +429,6 @@ def _generate_pythonop(component_func, target_image, target_component_file=None)
   input_names = inspect.getfullargspec(component_func)[0]
 
   output_name = 'output'
-  output_file = '/output.txt' #TODO: change the output path to /outputs/output/file here and in code generator
   component_spec = ComponentSpec(
       name=component_name,
       description=component_description,
@@ -440,10 +438,7 @@ def _generate_pythonop(component_func, target_image, target_component_file=None)
           container=ContainerSpec(
               image=target_image,
               #command=['python3', program_file], #TODO: Include the command line
-              arguments=[{'value': input_name} for input_name in input_names],
-              file_outputs={ #TODO: Use proper output arguments (e.g. "{output: output_name}" ) instead of this workaround. Our 1st-party components should not be using the file_outputs workaround.
-                output_name: output_file,
-              }
+              arguments=[{'value': input_name} for input_name in input_names] + [{'output': output_name}],
           )
       )
   )

sdk/python/tests/compiler/component_builder_test.py

@@ -314,15 +314,17 @@ def sample_component_func(a: str, b: int) -> float:
     result = float(b + 5)
   return result
 
-def wrapper_sample_component_func(a,b):
+def wrapper_sample_component_func(a,b,_output_file):
   output = sample_component_func(str(a),int(b))
-  with open("/output.txt", "w") as f:
-    f.write(str(output))
+  from pathlib import Path
+  Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
+  Path(_output_file).write_text(str(output))
 
 import argparse
 parser = argparse.ArgumentParser(description="Parsing arguments")
 parser.add_argument("a", type=str)
 parser.add_argument("b", type=int)
+parser.add_argument("_output_file", type=str)
 args = vars(parser.parse_args())
 
 if __name__ == "__main__":
@@ -338,15 +340,17 @@ def sample_component_func_two(a: str, b: int) -> float:
     result = float(b + 5)
   return result
 
-def wrapper_sample_component_func_two(a,b):
+def wrapper_sample_component_func_two(a,b,_output_file):
   output = sample_component_func_two(str(a),int(b))
-  with open("/output.txt", "w") as f:
-    f.write(str(output))
+  from pathlib import Path
+  Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
+  Path(_output_file).write_text(str(output))
 
 import argparse
 parser = argparse.ArgumentParser(description="Parsing arguments")
 parser.add_argument("a", type=str)
 parser.add_argument("b", type=int)
+parser.add_argument("_output_file", type=str)
 args = vars(parser.parse_args())
 
 if __name__ == "__main__":
@@ -359,13 +363,15 @@ def wrapper_sample_component_func_two(a,b):
 def sample_component_func_three() -> float:
   return 1.0
 
-def wrapper_sample_component_func_three():
+def wrapper_sample_component_func_three(_output_file):
   output = sample_component_func_three()
-  with open("/output.txt", "w") as f:
-    f.write(str(output))
+  from pathlib import Path
+  Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
+  Path(_output_file).write_text(str(output))
 
 import argparse
 parser = argparse.ArgumentParser(description="Parsing arguments")
+parser.add_argument("_output_file", type=str)
 args = vars(parser.parse_args())
 
 if __name__ == "__main__":