diff --git a/sdk/python/kfp/compiler/_component_builder.py b/sdk/python/kfp/compiler/_component_builder.py index 1f47791ac52..1022ff3365d 100644 --- a/sdk/python/kfp/compiler/_component_builder.py +++ b/sdk/python/kfp/compiler/_component_builder.py @@ -157,7 +157,6 @@ def _generate_dockerfile_with_py(self, target_file, base_image, python_filepath, with open(target_file, 'w') as f: f.write('FROM ' + base_image + '\n') f.write('RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools\n') - f.write('RUN pip3 install fire\n') if has_requirement_file: f.write('ADD ' + self._ARC_REQUIREMENT_FILE + ' /ml/\n') f.write('RUN pip3 install -r /ml/' + self._ARC_REQUIREMENT_FILE + '\n') @@ -329,10 +328,15 @@ def _generate_entrypoint(self, component_func): # CLI codes codegen.begin() - codegen.writeline('import fire') + codegen.writeline('import argparse') + 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('args = vars(parser.parse_args())') + codegen.writeline('') codegen.writeline('if __name__ == "__main__":') codegen.indent() - codegen.writeline('fire.Fire(' + new_func_name + ')') + codegen.writeline(new_func_name + '(**args)') # Remove the decorator from the component source src_lines = component_src.split('\n') diff --git a/sdk/python/tests/compiler/component_builder_test.py b/sdk/python/tests/compiler/component_builder_test.py index df1ad2a1e74..faa694c1567 100644 --- a/sdk/python/tests/compiler/component_builder_test.py +++ b/sdk/python/tests/compiler/component_builder_test.py @@ -151,13 +151,11 @@ def test_generate_dockerfile(self): golden_dockerfile_payload_one = '''\ FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0 RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools -RUN pip3 install fire ADD main.py /ml/ ENTRYPOINT ["python3", "/ml/main.py"]''' golden_dockerfile_payload_two = '''\ FROM gcr.io/ngao-mlpipeline-testing/tensorflow:1.10.0 RUN apt-get update -y && apt-get install --no-install-recommends -y -q python3 python3-pip python3-setuptools -RUN pip3 install fire ADD requirements.txt /ml/ RUN pip3 install -r /ml/requirements.txt ADD main.py /ml/ @@ -321,9 +319,14 @@ def wrapper_sample_component_func(a,b): with open("/output.txt", "w") as f: f.write(str(output)) -import fire +import argparse +parser = argparse.ArgumentParser(description="Parsing arguments") +parser.add_argument("a", type=str) +parser.add_argument("b", type=int) +args = vars(parser.parse_args()) + if __name__ == "__main__": - fire.Fire(wrapper_sample_component_func) + wrapper_sample_component_func(**args) ''' self.assertEqual(golden, generated_codes) @@ -340,9 +343,14 @@ def wrapper_sample_component_func_two(a,b): with open("/output.txt", "w") as f: f.write(str(output)) -import fire +import argparse +parser = argparse.ArgumentParser(description="Parsing arguments") +parser.add_argument("a", type=str) +parser.add_argument("b", type=int) +args = vars(parser.parse_args()) + if __name__ == "__main__": - fire.Fire(wrapper_sample_component_func_two) + wrapper_sample_component_func_two(**args) ''' self.assertEqual(golden, generated_codes) @@ -356,8 +364,11 @@ def wrapper_sample_component_func_three(): with open("/output.txt", "w") as f: f.write(str(output)) -import fire +import argparse +parser = argparse.ArgumentParser(description="Parsing arguments") +args = vars(parser.parse_args()) + if __name__ == "__main__": - fire.Fire(wrapper_sample_component_func_three) + wrapper_sample_component_func_three(**args) ''' self.assertEqual(golden, generated_codes) \ No newline at end of file