Kfp support for pip trusted host

Signed-off-by: Diego Lovison <[email protected]>

sdk/python/kfp/cli/component.py

@@ -155,6 +155,7 @@ def __init__(
         self._base_image = None
         self._target_image = None
         self._pip_index_urls = None
+        self._pip_trusted_hosts = None
         self._load_components()
 
     def _load_components(self):
@@ -214,11 +215,16 @@ def _load_components(self):
         logging.info(f'Using target image: {self._target_image}')
 
         pip_index_urls = []
+        pip_trusted_hosts = []
         for comp in self._components:
             if comp.pip_index_urls is not None:
                 pip_index_urls.extend(comp.pip_index_urls)
+            if comp.pip_trusted_hosts is not None:
+                pip_trusted_hosts.extend(comp.pip_trusted_hosts)
         if pip_index_urls:
             self._pip_index_urls = list(dict.fromkeys(pip_index_urls))
+        if pip_trusted_hosts:
+            self._pip_trusted_hosts = list(dict.fromkeys(pip_trusted_hosts))
 
     def _maybe_write_file(self,
                           filename: str,
@@ -277,7 +283,7 @@ def generate_kfp_config(self):
 
     def maybe_generate_dockerfile(self, overwrite_dockerfile: bool = False):
         index_urls_options = component_factory.make_index_url_options(
-            self._pip_index_urls)
+            self._pip_index_urls, self._pip_trusted_hosts)
         dockerfile_contents = _DOCKERFILE_TEMPLATE.format(
             base_image=self._base_image,
             maybe_copy_kfp_package=self._maybe_copy_kfp_package,

sdk/python/kfp/cli/component_test.py

@@ -37,6 +37,7 @@ def _make_component(
     packages_to_install: Optional[List[str]] = None,
     output_component_file: Optional[str] = None,
     pip_index_urls: Optional[List[str]] = None,
+    pip_trusted_hosts: Optional[List[str]] = None,
 ) -> str:
     return textwrap.dedent('''
     from kfp.dsl import *
@@ -46,7 +47,8 @@ def _make_component(
         target_image={target_image},
         packages_to_install={packages_to_install},
         output_component_file={output_component_file},
-        pip_index_urls={pip_index_urls})
+        pip_index_urls={pip_index_urls},
+        pip_trusted_hosts={pip_trusted_hosts})
     def {func_name}():
         pass
     ''').format(
@@ -55,7 +57,8 @@ def {func_name}():
         packages_to_install=repr(packages_to_install),
         output_component_file=repr(output_component_file),
         pip_index_urls=repr(pip_index_urls),
-        func_name=func_name)
+        func_name=func_name,
+        pip_trusted_hosts=repr(pip_trusted_hosts))
 
 
 def _write_file(filename: str, file_contents: str):
@@ -527,9 +530,9 @@ def test_docker_file_is_created_correctly_with_two_urls(self):
 
                 WORKDIR /usr/local/src/kfp/components
                 COPY runtime-requirements.txt runtime-requirements.txt
-                RUN pip install --index-url https://pypi.org/simple --trusted-host https://pypi.org/simple --extra-index-url https://example.com/pypi/simple --trusted-host https://example.com/pypi/simple --no-cache-dir -r runtime-requirements.txt
+                RUN pip install --index-url https://pypi.org/simple --extra-index-url https://example.com/pypi/simple --trusted-host https://pypi.org/simple --trusted-host https://example.com/pypi/simple --no-cache-dir -r runtime-requirements.txt
 
-                RUN pip install --index-url https://pypi.org/simple --trusted-host https://pypi.org/simple --extra-index-url https://example.com/pypi/simple --trusted-host https://example.com/pypi/simple --no-cache-dir kfp==1.2.3
+                RUN pip install --index-url https://pypi.org/simple --extra-index-url https://example.com/pypi/simple --trusted-host https://pypi.org/simple --trusted-host https://example.com/pypi/simple --no-cache-dir kfp==1.2.3
                 COPY . .
                 '''))
 
@@ -614,6 +617,66 @@ def test_dockerfile_can_contain_custom_kfp_package(self):
         self.assertTrue(contents.startswith(file_start))
         self.assertRegex(contents, 'RUN pip install --no-cache-dir kfp-*')
 
+    @mock.patch('kfp.__version__', '1.2.3')
+    def test_docker_file_is_created_one_trusted_host(self):
+        component = _make_component(
+            func_name='train',
+            target_image='custom-image',
+            pip_index_urls=['https://pypi.org/simple'],
+            pip_trusted_hosts=['pypi.org'])
+        _write_components('components.py', component)
+
+        result = self.runner.invoke(
+            self.cli,
+            ['build', str(self._working_dir)],
+        )
+        self.assertEqual(result.exit_code, 0)
+        self._docker_client.api.build.assert_called_once()
+        self.assert_file_exists_and_contains(
+            'Dockerfile',
+            textwrap.dedent('''\
+                    # Generated by KFP.
+    
+                    FROM python:3.8
+    
+                    WORKDIR /usr/local/src/kfp/components
+                    COPY runtime-requirements.txt runtime-requirements.txt
+                    RUN pip install --index-url https://pypi.org/simple --trusted-host pypi.org --no-cache-dir -r runtime-requirements.txt
+    
+                    RUN pip install --index-url https://pypi.org/simple --trusted-host pypi.org --no-cache-dir kfp==1.2.3
+                    COPY . .
+                    '''))
+
+    @mock.patch('kfp.__version__', '1.2.3')
+    def test_docker_file_is_created_two_trusted_host(self):
+        component = _make_component(
+            func_name='train',
+            target_image='custom-image',
+            pip_index_urls=['https://pypi.org/simple'],
+            pip_trusted_hosts=['pypi.org', 'pypi.org:8888'])
+        _write_components('components.py', component)
+
+        result = self.runner.invoke(
+            self.cli,
+            ['build', str(self._working_dir)],
+        )
+        self.assertEqual(result.exit_code, 0)
+        self._docker_client.api.build.assert_called_once()
+        self.assert_file_exists_and_contains(
+            'Dockerfile',
+            textwrap.dedent('''\
+                    # Generated by KFP.
+    
+                    FROM python:3.8
+    
+                    WORKDIR /usr/local/src/kfp/components
+                    COPY runtime-requirements.txt runtime-requirements.txt
+                    RUN pip install --index-url https://pypi.org/simple --trusted-host pypi.org --trusted-host pypi.org:8888 --no-cache-dir -r runtime-requirements.txt
+    
+                    RUN pip install --index-url https://pypi.org/simple --trusted-host pypi.org --trusted-host pypi.org:8888 --no-cache-dir kfp==1.2.3
+                    COPY . .
+                    '''))
+
 
 if __name__ == '__main__':
     unittest.main()

sdk/python/kfp/dsl/component_decorator.py

@@ -27,7 +27,8 @@ def component(func: Optional[Callable] = None,
               pip_index_urls: Optional[List[str]] = None,
               output_component_file: Optional[str] = None,
               install_kfp_package: bool = True,
-              kfp_package_path: Optional[str] = None):
+              kfp_package_path: Optional[str] = None,
+              pip_trusted_hosts: Optional[List[str]] = None):
     """Decorator for Python-function based components.
 
     A KFP component can either be a lightweight component or a containerized
@@ -114,7 +115,8 @@ def pipeline():
             pip_index_urls=pip_index_urls,
             output_component_file=output_component_file,
             install_kfp_package=install_kfp_package,
-            kfp_package_path=kfp_package_path)
+            kfp_package_path=kfp_package_path,
+            pip_trusted_hosts=pip_trusted_hosts)
 
     return component_factory.create_component_from_func(
         func,
@@ -124,4 +126,5 @@ def pipeline():
         pip_index_urls=pip_index_urls,
         output_component_file=output_component_file,
         install_kfp_package=install_kfp_package,
-        kfp_package_path=kfp_package_path)
+        kfp_package_path=kfp_package_path,
+        pip_trusted_hosts=pip_trusted_hosts)

sdk/python/kfp/dsl/component_factory.py

@@ -56,6 +56,7 @@ class ComponentInfo():
     base_image: str = _DEFAULT_BASE_IMAGE
     packages_to_install: Optional[List[str]] = None
     pip_index_urls: Optional[List[str]] = None
+    pip_trusted_hosts: Optional[List[str]] = None
 
 
 # A map from function_name to components.  This is always populated when a
@@ -69,22 +70,34 @@ def _python_function_name_to_component_name(name):
     return name_with_spaces[0].upper() + name_with_spaces[1:]
 
 
-def make_index_url_options(pip_index_urls: Optional[List[str]]) -> str:
+def make_index_url_options(pip_index_urls: Optional[List[str]],
+                           pip_trusted_hosts: Optional[List[str]]) -> str:
     """Generates index url options for pip install command based on provided
     pip_index_urls.
 
     Args:
         pip_index_urls: Optional list of pip index urls
+        pip_trusted_hosts: Optional list of pip trust hosts
 
     Returns:
         - Empty string if pip_index_urls is empty/None.
-        - '--index-url url --trusted-host url ' if pip_index_urls contains 1
+        - '--index-url url ' if pip_index_urls contains 1
         url
-        - the above followed by '--extra-index-url url --trusted-host url '
+        - the above followed by '--extra-index-url url '
         for
         each next url in pip_index_urls if pip_index_urls contains more than 1
         url
 
+        - pip_trusted_hosts was added later. In this case, if pip_trusted_hosts is None or empty
+        - the above followed by '--trusted-host url '
+        for
+        each url in pip_index_urls
+
+        - if pip_trusted_hosts is greater than 0
+        - the above followed by '--trusted-host url '
+        for
+        each url in pip_trusted_hosts
+
         Note: In case pip_index_urls is not empty, the returned string will
         contain space at the end.
     """
@@ -94,10 +107,17 @@ def make_index_url_options(pip_index_urls: Optional[List[str]]) -> str:
     index_url = pip_index_urls[0]
     extra_index_urls = pip_index_urls[1:]
 
-    options = [f'--index-url {index_url} --trusted-host {index_url}']
-    options.extend(
-        f'--extra-index-url {extra_index_url} --trusted-host {extra_index_url}'
-        for extra_index_url in extra_index_urls)
+    options = [f'--index-url {index_url}']
+    options.extend(f'--extra-index-url {extra_index_url}'
+                   for extra_index_url in extra_index_urls)
+
+    if pip_trusted_hosts is None or len(pip_trusted_hosts) == 0:
+        options.extend([f'--trusted-host {index_url}'])
+        options.extend(f'--trusted-host {extra_index_url}'
+                       for extra_index_url in extra_index_urls)
+    else:
+        options.extend(f'--trusted-host {trusted_host}'
+                       for trusted_host in pip_trusted_hosts)
 
     return ' '.join(options) + ' '
 
@@ -126,6 +146,7 @@ def _get_packages_to_install_command(
     packages_to_install: Optional[List[str]] = None,
     install_kfp_package: bool = True,
     target_image: Optional[str] = None,
+    pip_trusted_hosts: Optional[List[str]] = None,
 ) -> List[str]:
     packages_to_install = packages_to_install or []
     kfp_in_user_pkgs = any(pkg.startswith('kfp') for pkg in packages_to_install)
@@ -136,7 +157,8 @@ def _get_packages_to_install_command(
     if not inject_kfp_install and not packages_to_install:
         return []
     pip_install_strings = []
-    index_url_options = make_index_url_options(pip_index_urls)
+    index_url_options = make_index_url_options(pip_index_urls,
+                                               pip_trusted_hosts)
 
     if inject_kfp_install:
         if kfp_package_path:
@@ -517,6 +539,7 @@ def create_component_from_func(
     output_component_file: Optional[str] = None,
     install_kfp_package: bool = True,
     kfp_package_path: Optional[str] = None,
+    pip_trusted_hosts: Optional[List[str]] = None,
 ) -> python_component.PythonComponent:
     """Implementation for the @component decorator.
 
@@ -530,6 +553,7 @@ def create_component_from_func(
         kfp_package_path=kfp_package_path,
         packages_to_install=packages_to_install,
         pip_index_urls=pip_index_urls,
+        pip_trusted_hosts=pip_trusted_hosts,
     )
 
     command = []
@@ -575,7 +599,8 @@ def create_component_from_func(
         output_component_file=output_component_file,
         base_image=base_image,
         packages_to_install=packages_to_install,
-        pip_index_urls=pip_index_urls)
+        pip_index_urls=pip_index_urls,
+        pip_trusted_hosts=pip_trusted_hosts)
 
     if REGISTERED_MODULES is not None:
         REGISTERED_MODULES[component_name] = component_info