fix(sdk): Kfp support for pip trusted host (#11151)

Signed-off-by: Diego Lovison <diegolovison@gmail.com>
Signed-off-by: KevinGrantLee <kglee@google.com>

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,35 +70,47 @@ 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:
-    """Generates index url options for pip install command based on provided
-    pip_index_urls.
+def make_index_url_options(pip_index_urls: Optional[List[str]],
+                           pip_trusted_hosts: Optional[List[str]]) -> str:
+    """Generates index URL options for the pip install command based on the
+    provided pip_index_urls and pip_trusted_hosts.
 
     Args:
-        pip_index_urls: Optional list of pip index urls
+        pip_index_urls (Optional[List[str]]): Optional list of pip index URLs.
+        pip_trusted_hosts (Optional[List[str]]): Optional list of pip trusted hosts.
 
     Returns:
-        - Empty string if pip_index_urls is empty/None.
-        - '--index-url url --trusted-host url ' if pip_index_urls contains 1
-        url
-        - the above followed by '--extra-index-url url --trusted-host url '
-        for
-        each next url in pip_index_urls if pip_index_urls contains more than 1
-        url
-
-        Note: In case pip_index_urls is not empty, the returned string will
-        contain space at the end.
+        str:
+            - An empty string if pip_index_urls is empty or None.
+            - '--index-url url ' if pip_index_urls contains 1 URL.
+            - The above followed by '--extra-index-url url ' for each additional URL in pip_index_urls
+            if pip_index_urls contains more than 1 URL.
+            - If pip_trusted_hosts is None:
+                - The above followed by '--trusted-host url ' for each URL in pip_index_urls.
+            - If pip_trusted_hosts is an empty List.
+                - No --trusted-host information will be added
+            - If pip_trusted_hosts contains any URLs:
+                - 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 a space at the end.
     """
     if not pip_index_urls:
         return ''
 
     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:
+        options.extend([f'--trusted-host {index_url}'])
+        options.extend(f'--trusted-host {extra_index_url}'
+                       for extra_index_url in extra_index_urls)
+    elif len(pip_trusted_hosts) > 0:
+        options.extend(f'--trusted-host {trusted_host}'
+                       for trusted_host in pip_trusted_hosts)
 
     return ' '.join(options) + ' '
 
@@ -126,6 +139,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 +150,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 +532,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 +546,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 +592,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

sdk/python/kfp/dsl/component_factory_test.py

@@ -152,6 +152,42 @@ def test_with_packages_to_install_with_pip_index_url(self):
                 '\nif ! [ -x "$(command -v pip)" ]; then\n    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host https://myurl.org/simple \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host https://myurl.org/simple \'package1\' \'package2\' && "$0" "$@"\n'
             ]))
 
+    def test_with_packages_to_install_with_pip_index_url_and_trusted_host(self):
+        packages_to_install = ['package1', 'package2']
+        pip_index_urls = ['https://myurl.org/simple']
+        pip_trusted_hosts = ['myurl.org']
+
+        command = component_factory._get_packages_to_install_command(
+            packages_to_install=packages_to_install,
+            pip_index_urls=pip_index_urls,
+            pip_trusted_hosts=pip_trusted_hosts,
+        )
+
+        self.assertEqual(
+            strip_kfp_version(command),
+            strip_kfp_version([
+                'sh', '-c',
+                '\nif ! [ -x "$(command -v pip)" ]; then\n    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host myurl.org \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple --trusted-host myurl.org \'package1\' \'package2\' && "$0" "$@"\n'
+            ]))
+
+    def test_with_packages_to_install_with_pip_index_url_and_empty_trusted_host(
+            self):
+        packages_to_install = ['package1', 'package2']
+        pip_index_urls = ['https://myurl.org/simple']
+
+        command = component_factory._get_packages_to_install_command(
+            packages_to_install=packages_to_install,
+            pip_index_urls=pip_index_urls,
+            pip_trusted_hosts=[],
+        )
+
+        self.assertEqual(
+            strip_kfp_version(command),
+            strip_kfp_version([
+                'sh', '-c',
+                '\nif ! [ -x "$(command -v pip)" ]; then\n    python3 -m ensurepip || python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple \'kfp==2.1.3\' \'--no-deps\' \'typing-extensions>=3.7.4,<5; python_version<"3.9"\'  &&  python3 -m pip install --quiet --no-warn-script-location --index-url https://myurl.org/simple \'package1\' \'package2\' && "$0" "$@"\n'
+            ]))
+
 
 class TestInvalidParameterName(unittest.TestCase):