From 97cac84faee9873a2367ac3b1f6b6d013216fb78 Mon Sep 17 00:00:00 2001 From: "ax.vivien" Date: Thu, 9 Nov 2023 11:15:21 +0100 Subject: [PATCH] always include-tools when using large_pex. This allow pex to be used as venv --- cluster_pack/packaging.py | 24 +++++++++++++++++++----- cluster_pack/uploader.py | 9 ++++++++- tests/test_packaging.py | 23 +++++++++++++++++++++++ tests/test_uploader.py | 6 +++--- 4 files changed, 53 insertions(+), 9 deletions(-) diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py index 92bdcf3..62aa514 100644 --- a/cluster_pack/packaging.py +++ b/cluster_pack/packaging.py @@ -106,6 +106,7 @@ def pack_spec_in_pex(spec_file: str, output: str, pex_inherit_path: str = "fallback", allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[List[str], str]] = None, additional_indexes: Optional[List[str]] = None) -> str: with open(spec_file, "r") as f: @@ -114,6 +115,7 @@ def pack_spec_in_pex(spec_file: str, _logger.debug(f"used requirements: {lines}") return pack_in_pex(lines, output, pex_inherit_path=pex_inherit_path, allow_large_pex=allow_large_pex, + include_pex_tools=include_pex_tools, additional_repo=additional_repo, additional_indexes=additional_indexes) @@ -124,6 +126,7 @@ def pack_in_pex(requirements: List[str], pex_inherit_path: str = "fallback", editable_requirements: Dict[str, str] = {}, allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None ) -> str: @@ -149,6 +152,8 @@ def pack_in_pex(requirements: List[str], tmp_ext = ".tmp" else: tmp_ext = "" + if include_pex_tools: + cmd.extend(["--include-tools"]) if editable_requirements and len(editable_requirements) > 0: for current_package in editable_requirements.values(): @@ -232,6 +237,7 @@ def pack(self, ignored_packages: Collection[str], editable_requirements: Dict[str, str], allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None) -> str: raise NotImplementedError @@ -239,7 +245,8 @@ def pack(self, def pack_from_spec(self, spec_file: str, output: str, - allow_large_pex: bool = False) -> str: + allow_large_pex: bool = False, + include_pex_tools: bool = False) -> str: raise NotImplementedError @@ -268,6 +275,7 @@ def pack(self, ignored_packages: Collection[str], editable_requirements: Dict[str, str], allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None) -> str: return conda.pack_venv_in_conda( @@ -281,7 +289,8 @@ def pack(self, def pack_from_spec(self, spec_file: str, output: str, - allow_large_pex: bool = False) -> str: + allow_large_pex: bool = False, + include_pex_tools: bool = False) -> str: return conda.create_and_pack_conda_env( spec_file=spec_file, reqs=None, @@ -302,6 +311,7 @@ def pack(self, ignored_packages: Collection[str], editable_requirements: Dict[str, str], allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None) -> str: return pack_in_pex(reqs, @@ -309,14 +319,18 @@ def pack(self, ignored_packages, editable_requirements=editable_requirements, allow_large_pex=allow_large_pex, + include_pex_tools=include_pex_tools, additional_repo=additional_repo, additional_indexes=additional_indexes) def pack_from_spec(self, spec_file: str, output: str, - allow_large_pex: bool = False) -> str: - return pack_spec_in_pex(spec_file=spec_file, output=output, allow_large_pex=allow_large_pex) + allow_large_pex: bool = False, + include_pex_tools: bool = False) -> str: + return pack_spec_in_pex(spec_file=spec_file, output=output, + allow_large_pex=allow_large_pex, + include_pex_tools=include_pex_tools) CONDA_PACKER = CondaPacker() @@ -329,7 +343,7 @@ def _get_editable_requirements(executable: str = sys.executable) -> List[str]: location = pkg.get("editable_project_location", pkg.get("location", "")) packages_found = set( setuptools.find_packages(location) + setuptools.find_packages(f"{location}/src") - ) + ) for _pkg in packages_found: if "." in _pkg: continue diff --git a/cluster_pack/uploader.py b/cluster_pack/uploader.py index 8151d38..53e7573 100644 --- a/cluster_pack/uploader.py +++ b/cluster_pack/uploader.py @@ -91,6 +91,7 @@ def upload_env( include_editable: bool = False, fs_args: Dict[str, Any] = {}, allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None ) -> Tuple[str, str]: @@ -105,6 +106,7 @@ def upload_env( :param include_editable: whether to include or not packages installed in editable mode :param fs_args: filesystem args :param allow_large_pex: whether to allow or not building a large pex + :param include_pex_tools: whether to include tools or not when building pex :param additional_repo: additional repositories compliant with PEP 503 or local directories laid out in the same format. ex: https://download.pytorch.org/whl/cu113 @@ -128,6 +130,7 @@ def upload_env( force_upload, include_editable, allow_large_pex=allow_large_pex, + include_pex_tools=include_pex_tools, additional_repo=additional_repo, additional_indexes=additional_indexes ) @@ -269,6 +272,7 @@ def _upload_env_from_venv( force_upload: bool = False, include_editable: bool = False, allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None ) -> None: @@ -287,7 +291,8 @@ def _upload_env_from_venv( with tempfile.TemporaryDirectory() as tempdir: local_package_path = _pack_from_venv(executable, reqs, tempdir, packer, additional_packages, ignored_packages, force_upload, include_editable, - allow_large_pex, additional_repo, additional_indexes) + allow_large_pex, include_pex_tools, additional_repo, + additional_indexes) dir = os.path.dirname(package_path) if not resolved_fs.exists(dir): @@ -319,6 +324,7 @@ def _pack_from_venv(executable: str, force_upload: bool = False, include_editable: bool = False, allow_large_pex: bool = False, + include_pex_tools: bool = False, additional_repo: Optional[Union[str, List[str]]] = None, additional_indexes: Optional[List[str]] = None) -> str: env_copied_from_fallback_location = False @@ -372,6 +378,7 @@ def _pack_from_venv(executable: str, ignored_packages=ignored_packages, editable_requirements=editable_requirements, allow_large_pex=allow_large_pex, + include_pex_tools=include_pex_tools, additional_repo=additional_repo, additional_indexes=additional_indexes ) diff --git a/tests/test_packaging.py b/tests/test_packaging.py index a145e7f..a245a1a 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -239,6 +239,29 @@ def test_pack_in_pex_with_allow_large(): )) +def test_pack_in_pex_with_include_tools(): + with tempfile.TemporaryDirectory() as tempdir: + requirements = ["pyarrow==6.0.1"] + packaging.pack_in_pex( + requirements, + f"{tempdir}/out.pex", + # make isolated pex from current pytest virtual env + pex_inherit_path="false", + include_pex_tools=True) + assert os.path.exists(f"{tempdir}/out.pex") + cmd = ('print("Start importing pyarrow.."); ' + 'import pyarrow; ' + 'print("Successfully imported pyarrow!")') + + with does_not_raise(): + print(subprocess.check_output( + (f"PEX_TOOLS=1 {tempdir}/out.pex venv {tempdir}/pex_venv " + f"&& . {tempdir}/pex_venv/bin/activate " + f"&& python -c '{cmd}'"), + shell=True + )) + + def test_pack_in_pex_with_large_correctly_retrieves_zip_archive(): with tempfile.TemporaryDirectory() as tempdir: current_packages = packaging.get_non_editable_requirements(sys.executable) diff --git a/tests/test_uploader.py b/tests/test_uploader.py index 8873706..d1e785e 100644 --- a/tests/test_uploader.py +++ b/tests/test_uploader.py @@ -119,7 +119,7 @@ def test_upload_env(): cluster_pack.upload_env(MYARCHIVE_FILENAME, cluster_pack.PEX_PACKER) mock_packer.assert_called_once_with( ["a==1.0", "b==2.0"], Any(str), [], additional_indexes=None, additional_repo=None, - allow_large_pex=False, editable_requirements={} + allow_large_pex=False, editable_requirements={}, include_pex_tools=False ) mock_fs.put.assert_called_once_with(MYARCHIVE_FILENAME, MYARCHIVE_FILENAME) @@ -131,7 +131,7 @@ def test_upload_env(): ) mock_packer.assert_called_once_with( ["b==2.0", "c==3.0"], Any(str), ["a"], additional_indexes=None, additional_repo=None, - allow_large_pex=False, editable_requirements={} + allow_large_pex=False, editable_requirements={}, include_pex_tools=False ) @@ -356,7 +356,7 @@ def test_format_pex_requirements(): pex_inherit_path="false") pex_info = PexInfo.from_pex(f"{tempdir}/out.pex") cleaned_requirements = uploader._format_pex_requirements(pex_info) - pip_version = 'pip==21.3.1' if sys.version_info.minor == 6 else 'pip==23.3.1' + pip_version = 'pip==21.3.1' if sys.version_info.minor == 6 else 'pip==23.3.2' assert [pip_version, 'pipdeptree==2.0.0', 'six==1.15.0'] == cleaned_requirements