Fixes and extra logging.

Working through the perf analysis in pex-tool#2292 brought these to light.
jsirois · Dec 13, 2023 · af4749d · af4749d
1 parent 645d49f
commit af4749d
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 32 deletions.
diff --git a/pex/layout.py b/pex/layout.py
@@ -237,11 +237,11 @@ def install_distribution(item):
 
                         from pex.atomic_directory import atomic_directory
                         from pex.common import safe_mkdir
-                        from pex.layout import _identify_layout
+                        from pex.layout import identify_layout
                         from pex.tracer import TRACER
 
 
-                        with _identify_layout({pex!r}) as layout, atomic_directory(
+                        with identify_layout({pex!r}) as layout, atomic_directory(
                             {spread_dest!r}, source={source!r}
                         ) as spread_chroot:
                             if not spread_chroot.is_finalized():

diff --git a/pex/pex_builder.py b/pex/pex_builder.py
@@ -819,16 +819,17 @@ def zip_cache_dir(path):
         cached_bootstrap_zip_dir = zip_cache_dir(
             os.path.join(pex_info.pex_root, "bootstrap_zips", bootstrap_hash)
         )
-        with atomic_directory(cached_bootstrap_zip_dir) as atomic_bootstrap_zip_dir:
-            if not atomic_bootstrap_zip_dir.is_finalized():
-                self._chroot.zip(
-                    os.path.join(atomic_bootstrap_zip_dir.work_dir, pex_info.bootstrap),
-                    deterministic_timestamp=deterministic_timestamp,
-                    exclude_file=is_pyc_temporary_file,
-                    strip_prefix=pex_info.bootstrap,
-                    labels=("bootstrap",),
-                    compress=compress,
-                )
+        with TRACER.timed("Zipping PEX .bootstrap/ code."):
+            with atomic_directory(cached_bootstrap_zip_dir) as atomic_bootstrap_zip_dir:
+                if not atomic_bootstrap_zip_dir.is_finalized():
+                    self._chroot.zip(
+                        os.path.join(atomic_bootstrap_zip_dir.work_dir, pex_info.bootstrap),
+                        deterministic_timestamp=deterministic_timestamp,
+                        exclude_file=is_pyc_temporary_file,
+                        strip_prefix=pex_info.bootstrap,
+                        labels=("bootstrap",),
+                        compress=compress,
+                    )
         safe_copy(
             os.path.join(cached_bootstrap_zip_dir, pex_info.bootstrap),
             os.path.join(dirname, pex_info.bootstrap),
@@ -839,26 +840,32 @@ def zip_cache_dir(path):
         if pex_info.distributions:
             internal_cache = os.path.join(dirname, pex_info.internal_cache)
             os.mkdir(internal_cache)
-            for location, fingerprint in pex_info.distributions.items():
-                dest = os.path.join(internal_cache, location)
-                if pex_info.deps_are_wheel_files:
-                    for path in self._chroot.filesets[location]:
-                        safe_copy(os.path.join(self._chroot.chroot, path), dest)
-                else:
-                    cached_installed_wheel_zip_dir = zip_cache_dir(
-                        os.path.join(pex_info.pex_root, "packed_wheels", fingerprint)
-                    )
-                    with atomic_directory(cached_installed_wheel_zip_dir) as atomic_zip_dir:
-                        if not atomic_zip_dir.is_finalized():
-                            self._chroot.zip(
-                                os.path.join(atomic_zip_dir.work_dir, location),
-                                deterministic_timestamp=deterministic_timestamp,
-                                exclude_file=is_pyc_temporary_file,
-                                strip_prefix=os.path.join(pex_info.internal_cache, location),
-                                labels=(location,),
-                                compress=compress,
-                            )
-                    safe_copy(os.path.join(cached_installed_wheel_zip_dir, location), dest)
+            with TRACER.timed(
+                "{action} {count} distributions.".format(
+                    action="Copying" if pex_info.deps_are_wheel_files else "Zipping",
+                    count=len(pex_info.distributions)
+                )
+            ):
+                for location, fingerprint in pex_info.distributions.items():
+                    dest = os.path.join(internal_cache, location)
+                    if pex_info.deps_are_wheel_files:
+                        for path in self._chroot.filesets[location]:
+                            safe_copy(os.path.join(self._chroot.chroot, path), dest)
+                    else:
+                        cached_installed_wheel_zip_dir = zip_cache_dir(
+                            os.path.join(pex_info.pex_root, "packed_wheels", fingerprint)
+                        )
+                        with atomic_directory(cached_installed_wheel_zip_dir) as atomic_zip_dir:
+                            if not atomic_zip_dir.is_finalized():
+                                self._chroot.zip(
+                                    os.path.join(atomic_zip_dir.work_dir, location),
+                                    deterministic_timestamp=deterministic_timestamp,
+                                    exclude_file=is_pyc_temporary_file,
+                                    strip_prefix=os.path.join(pex_info.internal_cache, location),
+                                    labels=(location,),
+                                    compress=compress,
+                                )
+                        safe_copy(os.path.join(cached_installed_wheel_zip_dir, location), dest)
 
     def _build_zipapp(
         self,

diff --git a/pex/venv/installer.py b/pex/venv/installer.py
@@ -672,6 +672,7 @@ def sys_executable_paths():
                     "PEX_PYTHON_PATH",
                     "PEX_VERBOSE",
                     "PEX_EMIT_WARNINGS",
+                    "PEX_MAX_INSTALL_JOBS",
                     # This is used by the vendoring system.
                     "__PEX_UNVENDORED__",
                     # These are _not_ used at runtime, but are present under testing / CI and