SDK - Added support for loading zip-packed components

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)))

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)

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