From 66971172f87c74ad16d256fae171731a4e7e0eb1 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Wed, 6 Mar 2019 19:14:01 -0800 Subject: [PATCH] Added support for loading zip-packed components The zip-packed components are supported in all load_component APIs: kfp.components.load_component kfp.components.load_component_from_file kfp.components.load_component_from_url kfp.components.ComponentStore.load_component --- sdk/python/kfp/components/_component_store.py | 2 +- sdk/python/kfp/components/_components.py | 31 ++++++++++++++-- .../tests/components/test_components.py | 34 +++++++++++------- .../test_data/python_add.component.zip | Bin 0 -> 547 bytes 4 files changed, 50 insertions(+), 17 deletions(-) create mode 100644 sdk/python/tests/components/test_data/python_add.component.zip diff --git a/sdk/python/kfp/components/_component_store.py b/sdk/python/kfp/components/_component_store.py index cff2187287c..208943c5b69 100644 --- a/sdk/python/kfp/components/_component_store.py +++ b/sdk/python/kfp/components/_component_store.py @@ -85,6 +85,6 @@ def load_component(self, name, digest=None, tag=None): except: continue if response.content: - return comp._create_task_factory_from_component_text(response.content, url) + return comp._load_component_from_yaml_or_zip_bytes(response.content, url) raise RuntimeError('Component {} was not found. Tried the following locations:\n{}'.format(name, '\n'.join(tried_locations))) diff --git a/sdk/python/kfp/components/_components.py b/sdk/python/kfp/components/_components.py index 9d1df43de87..6f531dde4ef 100644 --- a/sdk/python/kfp/components/_components.py +++ b/sdk/python/kfp/components/_components.py @@ -77,7 +77,7 @@ def load_component_from_url(url): import requests resp = requests.get(url) resp.raise_for_status() - return _create_task_factory_from_component_text(resp.content, url) + return _load_component_from_yaml_or_zip_bytes(resp.content, url) def load_component_from_file(filename): @@ -93,8 +93,8 @@ def load_component_from_file(filename): ''' if filename is None: raise TypeError - with open(filename, 'r') as yaml_file: - return _create_task_factory_from_component_text(yaml_file, filename) + with open(filename, 'rb') as component_stream: + return _load_component_from_yaml_or_zip_stream(component_stream, filename) def load_component_from_text(text): @@ -113,6 +113,31 @@ def load_component_from_text(text): return _create_task_factory_from_component_text(text, None) +_COMPONENT_FILE_NAME_IN_ARCHIVE = 'component.yaml' + + +def _load_component_from_yaml_or_zip_bytes(bytes, component_filename=None): + import io + component_stream = io.BytesIO(bytes) + return _load_component_from_yaml_or_zip_stream(component_stream, component_filename) + + +def _load_component_from_yaml_or_zip_stream(stream, component_filename=None): + '''Loads component from a stream and creates a task factory function. + The stream can be YAML or a zip file with a component.yaml file inside. + ''' + import zipfile + stream.seek(0) + if zipfile.is_zipfile(stream): + stream.seek(0) + with zipfile.ZipFile(stream) as zip_obj: + with zip_obj.open(_COMPONENT_FILE_NAME_IN_ARCHIVE) as component_stream: + return _create_task_factory_from_component_text(component_stream, component_filename) + else: + stream.seek(0) + return _create_task_factory_from_component_text(stream, component_filename) + + def _create_task_factory_from_component_text(text_or_file, component_filename=None): component_dict = load_yaml(text_or_file) return _create_task_factory_from_component_dict(component_dict, component_filename) diff --git a/sdk/python/tests/components/test_components.py b/sdk/python/tests/components/test_components.py index c5336671de8..10c16659eed 100644 --- a/sdk/python/tests/components/test_components.py +++ b/sdk/python/tests/components/test_components.py @@ -23,24 +23,32 @@ from kfp.components._yaml_utils import load_yaml class LoadComponentTestCase(unittest.TestCase): - def test_load_component_from_file(self): - _this_file = Path(__file__).resolve() - _this_dir = _this_file.parent - _test_data_dir = _this_dir.joinpath('test_data') - component_path_obj = _test_data_dir.joinpath('python_add.component.yaml') - component_text = component_path_obj.read_text() - component_dict = load_yaml(component_text) - task_factory1 = comp.load_component_from_file(str(component_path_obj)) - assert task_factory1.__doc__ == component_dict['description'] + def _test_load_component_from_file(self, component_path: str): + task_factory1 = comp.load_component_from_file(component_path) arg1 = 3 arg2 = 5 task1 = task_factory1(arg1, arg2) - assert task1.human_name == component_dict['name'] - assert task1.image == component_dict['implementation']['container']['image'] - assert task1.arguments[0] == str(arg1) - assert task1.arguments[1] == str(arg2) + self.assertEqual(task1.human_name, 'Add') + self.assertEqual(task_factory1.__doc__.strip(), 'Returns sum of two arguments') + self.assertEqual(task1.image, 'python:3.5') + self.assertEqual(task1.arguments[0], str(arg1)) + self.assertEqual(task1.arguments[1], str(arg2)) + + def test_load_component_from_yaml_file(self): + _this_file = Path(__file__).resolve() + _this_dir = _this_file.parent + _test_data_dir = _this_dir.joinpath('test_data') + component_path = _test_data_dir.joinpath('python_add.component.yaml') + self._test_load_component_from_file(str(component_path)) + + def test_load_component_from_zipped_yaml_file(self): + _this_file = Path(__file__).resolve() + _this_dir = _this_file.parent + _test_data_dir = _this_dir.joinpath('test_data') + component_path = _test_data_dir.joinpath('python_add.component.zip') + self._test_load_component_from_file(str(component_path)) @unittest.skip @unittest.expectedFailure #The repo is non-public and will change soon. TODO: Update the URL and enable the test once we move to a public repo diff --git a/sdk/python/tests/components/test_data/python_add.component.zip b/sdk/python/tests/components/test_data/python_add.component.zip new file mode 100644 index 0000000000000000000000000000000000000000..951d7589740d181472118e9491f1e6a890d7bb43 GIT binary patch literal 547 zcmWIWW@Zs#W?Wn}wBKt3 zfxFMO)2_R0fBWoJo>rl%S&k#eM#mLJuQ(%)@b@0i68v(1<`KV6r55pvHhs_M$n*z` zl^3inU1}(Ig2n9K+XtJOcDVaJQ96=xu>62Ke^=Zj*W(+F^i=wGXEDoHwnQ=SXXo5? zPus5E{Fdus^|+Y5vfA7;HZDBg=-u1pa{i!qlFTHPqMlHFCjR$-wEphDT^}Uta7}Bu zaZ2zm)tfS23iamm`u4Og__5Vomt&!No3-caOV@8SUfLDrdR6L{km-h>^B13)az1P8 zud^nZ=MR`DIUbbH5O4Xdyi#C7OU0wPO>ge)GqyD=e)UKzCiP&^;+dts?B|SUWW=OA z<%E6%R_(fx)2vjo zXKVhgfVq<^zB+JTZ%`Mme1HGgmY=m}Ch|O<_wh(+%B&Fk&Qs5qY5w!ezax1lM)L66 z-kMi4-_P{Ar|Vb!#qizl#jh@Jc=de#s<%sL?PRaLoHu{9et