Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SDK/Components/PythonContainerOp - Make the local output path configurable #424

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 6 additions & 11 deletions sdk/python/kfp/compiler/_component_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,9 +301,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
Expand All @@ -316,9 +314,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
Expand All @@ -327,6 +325,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__":')
Expand Down Expand Up @@ -425,7 +424,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,
Expand All @@ -435,10 +433,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
args=[{'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,
}
args=[{'value': input_name} for input_name in input_names] + [{'output': output_name}],
gaoning777 marked this conversation as resolved.
Show resolved Hide resolved
)
)
)
Expand Down
24 changes: 15 additions & 9 deletions sdk/python/tests/compiler/component_builder_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__":
Expand All @@ -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__":
Expand All @@ -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__":
Expand Down