diff --git a/sdk/python/kfp/compiler/_component_builder.py b/sdk/python/kfp/compiler/_component_builder.py index 373ddc5d0d41..d3b89ac5ee06 100644 --- a/sdk/python/kfp/compiler/_component_builder.py +++ b/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}], ) ) ) diff --git a/sdk/python/tests/compiler/component_builder_test.py b/sdk/python/tests/compiler/component_builder_test.py index faa694c15677..b02e1a4bba60 100644 --- a/sdk/python/tests/compiler/component_builder_test.py +++ b/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__":