Repository: pex-tool/pex

Issue:
Large numbers of extras lead to exponential growth in PEX internal dependency resolution time.
This can be observed on main with the following PEX:
```
$ time pex "boto3-stubs[batch,codeartifact,codebuild,codepipeline,cognito-idp,dynamodb,kms,logs,ram,rds,route53,s3,secretsmanager,ssm,stepfunctions,sts]==1.24.78" -oslow.pex

real    0m6.487s
user    0m5.558s
sys     0m0.060s
$ ls -lh slow.pex
-rwxr-xr-x 1 jsirois jsirois 1.5M Oct  3 23:09 slow.pex
```

Although creating it is fast and the PEX is small, it takes an inordinate amount of time to resolve its internal distributions on boot:
```
$ time PEX_VERBOSE=1 ./slow.pex -c ''
pex: Laying out PEX zipfile /home/jsirois/dev/pantsbuild/jsirois-pex/slow.pex :: Installing /home/jsirois/dev/pantsbuild/jsirois-pex/slow.pex to /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2cpex: Laying out PEX zipfile /home/jsirois/dev/pantsbuild/jsirois-pex/slow.pex: 18.1ms

pex:   Installing /home/jsirois/dev/pantsbuild/jsirois-pex/slow.pex to /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f: 17.9ms
pex: Executing installed PEX for /home/jsirois/dev/pantsbuild/jsirois-pex/./slow.pex at /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f
pex:   Testing /usr/bin/python3.10 can resolve PEX at /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f: 70913.3ms
pex: Using the current interpreter /usr/bin/python3.10 since it matches constraints and PYTHONPATH is not set.
pex: Found site-library: /usr/lib/python3.10/dist-packages
pex: Found site-library: /usr/local/lib/python3.10/dist-packages
pex: Found site-library: /usr/lib/python3/dist-packages
pex: Tainted path element: /usr/local/lib/python3.10/dist-packages
pex: Tainted path element: /usr/lib/python3/dist-packages
pex: Scrubbing from user site: /home/jsirois/.local/lib/python3.10/site-packages
pex: Scrubbing from site-packages: /usr/local/lib/python3.10/dist-packages
pex: Scrubbing from site-packages: /usr/lib/python3/dist-packages
pex: New sys.path: ['/home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f/.bootstrap', '/home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f', '/usr/lib/python310.zip', '/usr/lib/python3.10', '/usr/lib/python3.10/lib-dynload']
pex: Activating PEX virtual environment from /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f: 1.7ms
pex: No entry point specified, dropping into interpreter
pex: Bootstrap complete, performing final sys.path modifications...
pex: PYTHONPATH contains:
pex:   *
pex:     /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f
pex:   * /usr/lib/python310.zip
pex:     /usr/lib/python3.10
pex:     /usr/lib/python3.10/lib-dynload
pex:     /home/jsirois/.pex/installed_wheels/b720424eba6739c9641cdf7c4674aecfa8aef259e7cf48802caec9f914a75bfa/boto3_stubs-1.24.78-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/c9c0304317455ab07685ec9d8e62a30f923b7552add346081846fb9efd6b71b9/botocore_stubs-1.27.85-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/f8b1bf874c2e38deb390fcadc1177df760eebe85116398a359c49423b70b1b64/types_awscrt-0.14.6-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/78f9ea029fedc97e1a5d323edd149a827d4ab361d11cd92bb1b5f235ff84ba72/types_s3transfer-0.6.0.post4-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/25642c956049920a5aa49edcdd6ab1e06d7e5d467fc00e0506c44ac86fbfca02/typing_extensions-4.3.0-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/6a94ffa7420d8f33fb7c10394f36e3508fc6f60670234f129020f6b09bbc2bb4/mypy_boto3_batch-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/9b1cfcbb716acbacddd7d65decf55da78dcd1da3befd4431dbe104b1dead2893/mypy_boto3_codeartifact-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/e8095b2a07c8287487bb6eb4fd64d21365102dcb825e63dd110a3639055bc443/mypy_boto3_codebuild-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/a27ec674962f42da6d2a233eef52bd4252b5e3a83e7667e97198977f90bddb75/mypy_boto3_codepipeline-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/6edda2d0194d72133da603cd05f2404603e0d86661e29beea839fc9b04daaa24/mypy_boto3_cognito_idp-1.24.66-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/1111c2eb8129bf8c09b5423d95ad483943f5fc08d745801c356f6a1b22b04b37/mypy_boto3_dynamodb-1.24.74-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/4f2638acf85aa26f618696680387635b378ad3c646746215cf6057a586d9c0b7/mypy_boto3_kms-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/f96257ec06099bfda1ce5f35b410e7fb93fb601bc312e8d7a09b13adaefd23f0/mypy_boto3_logs-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/c590661016684109c8837968b09dfffd25747f9fd69a713d97ea869b69a072ad/mypy_boto3_ram-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/d88bd90965f64e3344d1faf16273da6fb62f00bb539fb9ba5273feb750bb1f04/mypy_boto3_rds-1.24.76-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/252ac6fdb0e7e61c21b933a21c0b954ec0f60ef23002bfde088471eec778fd7a/mypy_boto3_route53-1.24.69-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/fd71cd3894a4da0dfcf1372984e667bfea8df930a1d27906206b6bb4d82ca418/mypy_boto3_s3-1.24.76-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/9ed3ec38a6c05961cb39a2d9fb891441d4cf22c63e34a6998fbd3d28ba290d9a/mypy_boto3_secretsmanager-1.24.83-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/a50fe448f3c18f76255e15878e21020001ec04a85b42996db721d9b89770ff11/mypy_boto3_ssm-1.24.81-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/7baf9b20fdbd797cd2bc09b56b61777cf0eb139b51d2fea938d238d806c85b69/mypy_boto3_stepfunctions-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/installed_wheels/cd47af04b29f9e9d4e400ab66e1723e72b3d4054d006772906dfc07b3c89280e/mypy_boto3_sts-1.24.36.post1-py3-none-any.whl
pex:     /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f/.bootstrap
pex:   * - paths that do not exist or will be imported via zipimport

real    1m11.209s
user    1m11.154s
sys     0m0.056s
```

Note the line:
```
pex:   Testing /usr/bin/python3.10 can resolve PEX at /home/jsirois/.pex/unzipped_pexes/b77f65d72adf1a7c31813f921974d2c60d71547f: 70913.3ms
```
That accounts for >99% of the wall time of 71.209s


Discussion:
The key property here is that the `boto3-stubs` requirement has 16 extras listed.

That results in 2^16 keyings of the `boto3-stubs` distribution during the PEX internal resolve process:
https://github.com/pantsbuild/pex/blob/ecf9b95544c53970628e982b7132a345bbc931cd/pex/environment.py#L195-L208
https://github.com/pantsbuild/pex/blob/ecf9b95544c53970628e982b7132a345bbc931cd/pex/environment.py#L377-L379

Where this becomes slow is collapsing the resulting resolved set of distributions into a set. The repeated hashing takes up the lion's share of the ~70s observed above on this one line of code:
https://github.com/pantsbuild/pex/blob/ecf9b95544c53970628e982b7132a345bbc931cd/pex/environment.py#L597

A quick experiment that caches the hash code of the distribution drops that ~70s to ~700ms.

Resolving patch:
diff --git a/pex/dist_metadata.py b/pex/dist_metadata.py
--- a/pex/dist_metadata.py
+++ b/pex/dist_metadata.py
@@ -508,7 +508,9 @@ def __str__(self):
         return self._str
 
 
-@attr.s(frozen=True)
+# N.B.: DistributionMetadata can have an expensive hash when a distribution has many requirements;
+# so we cache the hash. See: https://github.com/pantsbuild/pex/issues/1928
+@attr.s(frozen=True, cache_hash=True)
 class DistMetadata(object):
     @classmethod
     def load(cls, location):
diff --git a/pex/interpreter.py b/pex/interpreter.py
--- a/pex/interpreter.py
+++ b/pex/interpreter.py
@@ -123,9 +123,13 @@ def get(cls, binary=None):
         )
 
         # Pex identifies interpreters using a bit of Pex code injected via an extraction of that
-        # code under the `PEX_ROOT` adjoined to `sys.path` via `PYTHONPATH`. We ignore such adjoined
-        # `sys.path` entries to discover the true base interpreter `sys.path`.
-        pythonpath = frozenset(os.environ.get("PYTHONPATH", "").split(os.pathsep))
+        # code under the `PEX_ROOT` adjoined to `sys.path` via `PYTHONPATH`. Pex also exposes the
+        # vendored attrs distribution so that its `cache_hash=True` feature can work (see the
+        # bottom of pex/third_party/__init__.py where the vendor importer is installed). We ignore
+        # such adjoined `sys.path` entries to discover the true base interpreter `sys.path`.
+        pythonpath = frozenset(
+            os.environ.get("PYTHONPATH", "").split(os.pathsep) + list(third_party.exposed())
+        )
         sys_path = [item for item in sys.path if item and item not in pythonpath]
 
         return cls(
diff --git a/pex/third_party/__init__.py b/pex/third_party/__init__.py
--- a/pex/third_party/__init__.py
+++ b/pex/third_party/__init__.py
@@ -18,7 +18,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Container, Iterator, Optional, Sequence, Tuple
+    from typing import Container, Iterable, Iterator, List, Optional, Sequence, Tuple
 
 
 def _tracer():
@@ -57,11 +57,17 @@ class _Importable(namedtuple("_Importable", ["module", "is_pkg", "path", "prefix
     _exposed = False  # noqa: We want instance variable access defaulting to cls here.
 
     def expose(self):
+        # type: () -> None
         self._exposed = True
-        importlib.import_module(self.module)
         _tracer().log("Exposed {}".format(self), V=3)
 
+    @property
+    def exposed(self):
+        # type: () -> bool
+        return self._exposed
+
     def loader_for(self, fullname):
+        # type: (str) -> Optional[_Loader]
         if fullname.startswith(self.prefix + "."):
             target = fullname[len(self.prefix + ".") :]
         else:
@@ -78,6 +84,8 @@ def loader_for(self, fullname):
             vendor_module_name = vendor_path.replace(os.sep, ".")
             return _Loader(fullname, vendor_module_name)
 
+        return None
+
 
 class _DirIterator(namedtuple("_DirIterator", ["rootdir"])):
     def iter_root_modules(self, relpath):
@@ -141,6 +149,13 @@ def _filter_names(self, relpath, pattern, group):
                     yield match.group(group)
 
 
+def _vendored_path_items():
+    # type: () -> Iterable[str]
+    from pex import vendor
+
+    return tuple(spec.relpath for spec in vendor.iter_vendor_specs())
+
+
 class VendorImporter(object):
     """A `PEP-302 <https://www.python.org/dev/peps/pep-0302/>`_ meta_path importer for vendored
     code.
@@ -159,6 +174,7 @@ class VendorImporter(object):
 
     @staticmethod
     def _abs_root(root=None):
+        # type: (Optional[str]) -> str
         from pex import vendor
 
         return os.path.abspath(root or vendor.VendorSpec.ROOT)
@@ -181,35 +197,43 @@ def _iter_all_installed_vendor_importers(cls):
                 yield importer
 
     @classmethod
-    def _iter_installed_vendor_importers(cls, prefix, root, path_items):
+    def iter_installed_vendor_importers(
+        cls,
+        prefix,  # type: str
+        root=None,  # type: Optional[str]
+    ):
+        # type: (...) -> Iterator[VendorImporter]
+        root = cls._abs_root(root)
+        vendored_paths = set(_vendored_path_items())
         for importer in cls._iter_all_installed_vendor_importers():
             # All Importables for a given VendorImporter will have the same prefix.
             if importer._importables and importer._importables[0].prefix == prefix:
                 if importer._root == root:
-                    if {importable.path for importable in importer._importables} == set(path_items):
+                    if {importable.path for importable in importer._importables} == vendored_paths:
                         yield importer
 
     @classmethod
-    def install_vendored(cls, prefix, root=None, expose=None):
+    def install_vendored(
+        cls,
+        prefix,  # type: str
+        root=None,  # type: Optional[str]
+        expose=None,  # type: Optional[Iterable[str]]
+    ):
+        # type: (...) -> None
         """Install an importer for all vendored code with the given import prefix.
 
         All distributions listed in ``expose`` will also be made available for import in direct,
         un-prefixed form.
 
-        :param str prefix: The import prefix the installed importer will be responsible for.
-        :param str root: The root path of the distribution containing the vendored code. NB: This is the
-                         the path to the pex code, which serves as the root under which code is vendored
-                         at ``pex/vendor/_vendored``.
+        :param prefix: The import prefix the installed importer will be responsible for.
+        :param root: The root path of the distribution containing the vendored code. NB: This is the
+                     path to the pex code, which serves as the root under which code is vendored at
+                     ``pex/vendor/_vendored``.
         :param expose: Optional names of distributions to expose for direct, un-prefixed import.
-        :type expose: list of str
         :raise: :class:`ValueError` if any distributions to expose cannot be found.
         """
-        from pex import vendor
-
         root = cls._abs_root(root)
-        vendored_path_items = [spec.relpath for spec in vendor.iter_vendor_specs()]
-
-        installed = list(cls._iter_installed_vendor_importers(prefix, root, vendored_path_items))
+        installed = tuple(cls.iter_installed_vendor_importers(prefix, root=root))
         assert (
             len(installed) <= 1
         ), "Unexpected extra importers installed for vendored code:\n\t{}".format(
@@ -218,9 +242,10 @@ def install_vendored(cls, prefix, root=None, expose=None):
         if installed:
             vendor_importer = installed[0]
         else:
-            # Install all vendored code for pex internal access to it through the vendor import `prefix`.
+            # Install all vendored code for pex internal access to it through the vendor import
+            # `prefix`.
             vendor_importer = cls.install(
-                uninstallable=True, prefix=prefix, path_items=vendored_path_items, root=root
+                uninstallable=True, prefix=prefix, path_items=_vendored_path_items(), root=root
             )
 
         if expose:
@@ -233,7 +258,12 @@ def install_vendored(cls, prefix, root=None, expose=None):
             vendor_importer._expose(exposed_paths)
 
     @classmethod
-    def expose(cls, dists, root=None):
+    def expose(
+        cls,
+        dists,  # type: Iterable[str]
+        root=None,  # type: Optional[str]
+    ):
+        # type: (...) -> Iterator[str]
         from pex import vendor
 
         root = cls._abs_root(root)
@@ -288,13 +318,28 @@ def uninstall_all(cls):
         for vendor_importer in cls._iter_all_installed_vendor_importers():
             vendor_importer.uninstall()
 
-    def __init__(self, root, importables, uninstallable=True):
+    def __init__(
+        self,
+        root,  # type: str
+        importables,  # type: Tuple[_Importable, ...]
+        uninstallable=True,  # type: bool
+    ):
+        # type: (...) -> None
         self._root = root
         self._importables = importables
-
         self._uninstallable = uninstallable
 
-        self._loaders = []
+        self._loaders = []  # type: List[_Loader]
+
+    @property
+    def root(self):
+        # type: () -> str
+        return self._root
+
+    @property
+    def importables(self):
+        # type: () -> Iterable[_Importable]
+        return self._importables
 
     def uninstall(self):
         """Uninstall this importer if possible and un-import any modules imported by it."""
@@ -520,8 +565,19 @@ def install(root=None, expose=None):
     VendorImporter.install_vendored(prefix=import_prefix(), root=root, expose=expose)
 
 
+def exposed(root=None):
+    # type: (Optional[str]) -> Iterator[str]
+    """Returns the ``sys.path`` entries of distributions that have been exposed."""
+    for importer in VendorImporter.iter_installed_vendor_importers(
+        prefix=import_prefix(), root=root
+    ):
+        for importable in importer.importables:
+            if importable.exposed:
+                yield os.path.join(importer.root, importable.path)
+
+
 def expose(dists):
-    # type: (Sequence[str]) -> Iterator[str]
+    # type: (Iterable[str]) -> Iterator[str]
     """Exposes vendored code in isolated chroots.
 
     Any vendored distributions listed in ``dists`` will be unpacked to individual chroots for
@@ -537,4 +593,6 @@ def expose(dists):
 
 
 # Implicitly install an importer for vendored code on the first import of pex.third_party.
-install()
+# N.B.: attrs must be exposed to make use of `cache_hash=True` since that generates and compiles
+# code on the fly that generated code does a bare `import attr`.
+install(expose=["attrs"])


Test patch: