diff --git a/sdk/python/kfp/cli/component.py b/sdk/python/kfp/cli/component.py index 079c200fe3c3..459265fe3931 100644 --- a/sdk/python/kfp/cli/component.py +++ b/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, diff --git a/sdk/python/kfp/cli/component_test.py b/sdk/python/kfp/cli/component_test.py index c7279559194b..7a40ac73822c 100644 --- a/sdk/python/kfp/cli/component_test.py +++ b/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() diff --git a/sdk/python/kfp/dsl/component_decorator.py b/sdk/python/kfp/dsl/component_decorator.py index 7c6589589dc2..6e0c70679d9d 100644 --- a/sdk/python/kfp/dsl/component_decorator.py +++ b/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) diff --git a/sdk/python/kfp/dsl/component_factory.py b/sdk/python/kfp/dsl/component_factory.py index 1af26d80bfc4..2f14aded2c47 100644 --- a/sdk/python/kfp/dsl/component_factory.py +++ b/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,24 +70,29 @@ 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 '' @@ -94,10 +100,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: + 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 diff --git a/sdk/python/kfp/dsl/component_factory_test.py b/sdk/python/kfp/dsl/component_factory_test.py index b602be241fd0..28ebbece2928 100644 --- a/sdk/python/kfp/dsl/component_factory_test.py +++ b/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):