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..863e3f9f9285 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): 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..4adac6a9876b 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,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