always include-tools when using large_pex. This allow pex to be used …

…as venv

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,14 +237,16 @@ 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
 
     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,21 +311,26 @@ 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,
                            output,
                            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

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
         )

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"&& source {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)

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