Lower noop wheel install overhead.

Previously, installing a wheel that was already installed incurred the
cost of hashing the installed wheel chroot every time. The overhead of
this wasted work for a warm cache was egregious for large distributions
like PyTorch, with gigabytes of files to hash taking seconds.

Work towards pex-tool#2312.

Author: jsirois

pex/jobs.py

@@ -732,7 +732,7 @@ def iter_map_parallel(
     # least two slots to ensure we process input items in parallel.
     pool_size = max(2, min(len(input_items) // min_average_load, _sanitize_max_jobs(max_jobs)))
 
-    perform_install = functools.partial(_apply_function, function)
+    apply_function = functools.partial(_apply_function, function)
 
     slots = defaultdict(list)  # type: DefaultDict[int, List[float]]
     with TRACER.timed(
@@ -741,7 +741,7 @@ def iter_map_parallel(
         )
     ):
         with _mp_pool(size=pool_size) as pool:
-            for pid, result, elapsed_secs in pool.imap_unordered(perform_install, input_items):
+            for pid, result, elapsed_secs in pool.imap_unordered(apply_function, input_items):
                 TRACER.log(
                     "[{pid}] {verbed} {result} in {elapsed_secs:.2f}s".format(
                         pid=pid,

pex/pep_376.py

@@ -17,6 +17,7 @@
 from pex.common import is_pyc_dir, is_pyc_file, safe_mkdir, safe_open
 from pex.interpreter import PythonInterpreter
 from pex.typing import TYPE_CHECKING, cast
+from pex.util import CacheHelper
 from pex.venv.virtualenv import Virtualenv
 
 if TYPE_CHECKING:
@@ -169,10 +170,24 @@ def save(
         record_relpath,  # type: Text
     ):
         # type: (...) -> InstalledWheel
-        layout = {"stash_dir": stash_dir, "record_relpath": record_relpath}
+
+        # We currently need the installed wheel chroot hash for PEX-INFO / boot purposes. It is
+        # expensive to calculate; so we do it here 1 time when saving the installed wheel.
+        fingerprint = CacheHelper.dir_hash(prefix_dir, hasher=hashlib.sha256)
+
+        layout = {
+            "stash_dir": stash_dir,
+            "record_relpath": record_relpath,
+            "fingerprint": fingerprint,
+        }
         with open(cls.layout_file(prefix_dir), "w") as fp:
             json.dump(layout, fp, sort_keys=True)
-        return cls(prefix_dir=prefix_dir, stash_dir=stash_dir, record_relpath=record_relpath)
+        return cls(
+            prefix_dir=prefix_dir,
+            stash_dir=stash_dir,
+            record_relpath=record_relpath,
+            fingerprint=fingerprint,
+        )
 
     @classmethod
     def load(cls, prefix_dir):
@@ -201,15 +216,20 @@ def load(cls, prefix_dir):
                     layout_file=layout_file, value=layout
                 )
             )
+
+        fingerprint = layout.get("fingerprint")
+
         return cls(
             prefix_dir=prefix_dir,
             stash_dir=cast(str, stash_dir),
             record_relpath=cast(str, record_relpath),
+            fingerprint=cast("Optional[str]", fingerprint),
         )
 
     prefix_dir = attr.ib()  # type: str
     stash_dir = attr.ib()  # type: str
     record_relpath = attr.ib()  # type: Text
+    fingerprint = attr.ib()  # type: Optional[str]
 
     def stashed_path(self, *components):
         # type: (*str) -> str

pex/resolver.py

@@ -23,6 +23,7 @@
 from pex.jobs import Raise, SpawnedJob, execute_parallel, iter_map_parallel
 from pex.network_configuration import NetworkConfiguration
 from pex.orderedset import OrderedSet
+from pex.pep_376 import InstalledWheel, LoadError
 from pex.pep_425 import CompatibilityTags
 from pex.pep_427 import InstallableType, WheelError, install_wheel_chroot
 from pex.pep_503 import ProjectName
@@ -485,7 +486,16 @@ def finalize_install(self, install_requests):
         # pex:   * - paths that do not exist or will be imported via zipimport
         # pex.pex 2.0.2
         #
-        wheel_dir_hash = fingerprint_path(self.install_chroot)
+        cached_fingerprint = None  # type: Optional[str]
+        try:
+            installed_wheel = InstalledWheel.load(self.install_chroot)
+        except LoadError:
+            # We support legacy chroots below by calculating the chroot fingerprint just in time.
+            pass
+        else:
+            cached_fingerprint = installed_wheel.fingerprint
+
+        wheel_dir_hash = cached_fingerprint or fingerprint_path(self.install_chroot)
         runtime_key_dir = os.path.join(self._installation_root, wheel_dir_hash)
         with atomic_directory(runtime_key_dir) as atomic_dir:
             if not atomic_dir.is_finalized():