Repository: pex-tool/pex

Issue:
Regression in `_populate_deps`(?) for `--venv` PEXes
It seems f72330201421be3ab69efe7cb26758409106e562 was a significant performance regression for `--venv` PEXes,especially for PEXes with many small files. Potentially it's a correctness fix for which we just have to eat the performance impact. In addition to a regression, this time doesn't seem to be bracketed by any of the tracing spans, and so isn't obvious in logs.

The table compares the cold-cache start-up time for two packages that are both fairly large, but with drastically different numbers of files. The table compares 2.1.119 and 2.1.120, but I separately bisected the regression to f72330201421be3ab69efe7cb26758409106e562.

| package | wheel size | number of files | time 2.1.119 (s) | time 2.1.120 (s) |
|---|---|---|---|---|
| `semgrep==1.14.0` | 27MB | 123  | 1.4 | 1.8 |
| `plotly==5.13.1` | 15MB  | ~12k | 2.5 | 6.6 |

Reproducer, using `ts -i` to record the incremental timings in text (leading number on each line of the output):

```shell
# 2.1.119
curl -L https://github.com/pantsbuild/pex/releases/download/v2.1.119/pex -o pex-2.1.119

python pex-2.1.119 --venv -m platform semgrep==1.14.0 -o semgrep_via_119.pex
time PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python semgrep_via_119.pex 2>&1 | ts -i '%.s'

python pex-2.1.119 --venv -m platform plotly==5.13.1 -o plotly_via_119.pex
time PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python plotly_via_119.pex 2>&1 | ts -i '%.s'

# 2.1.120
curl -L https://github.com/pantsbuild/pex/releases/download/v2.1.120/pex -o pex-2.1.120

python pex-2.1.120 --venv -m platform semgrep==1.14.0 -o semgrep_via_120.pex
time PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python semgrep_via_120.pex 2>&1 | ts -i '%.s'

python pex-2.1.120 --venv -m platform plotly==5.13.1 -o plotly_via_120.pex
time PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python plotly_via_120.pex 2>&1 | ts -i '%.s'
```

Output for plotly, with 2.1.119:

```
...
0.158249 pex: Laying out PEX zipfile /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex
0.000613 pex: Laying out PEX zipfile /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex :: Installing /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex to /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/unzipped_pexes/9b545f5764cfad9084f31d8fa169a9d3b12acb9d
1.749525 pex: Laying out PEX zipfile /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex: 1750.2ms
0.000070 pex:   Installing /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex to /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/unzipped_pexes/9b545f5764cfad9084f31d8fa169a9d3b12acb9d: 1749.7ms
0.001437 pex: Executing installed PEX for /Users/huon/tmp/pex-perf-regression/plotly_via_119.pex at /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/unzipped_pexes/9b545f5764cfad9084f31d8fa169a9d3b12acb9d
0.242556 pex: Testing /Users/huon/.pyenv/versions/3.9.10/bin/python3.9 can resolve PEX at /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/unzipped_pexes/9b545f5764cfad9084f31d8fa169a9d3b12acb9d
0.008913 pex:   Testing /Users/huon/.pyenv/versions/3.9.10/bin/python3.9 can resolve PEX at /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/unzipped_pexes/9b545f5764cfad9084f31d8fa169a9d3b12acb9d: 8.9ms
0.052860 pex: Hashing pex
0.010083 pex: Hashing pex: 10.1ms
0.000098 pex: Isolating pex
0.000277 pex: Isolating pex :: Extracting pex to /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/isolated/f2df3e7415f34f911478882907d6c86e4750ccaa
0.064839 pex: Isolating pex: 65.1ms
0.000036 pex:   Extracting pex to /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/isolated/f2df3e7415f34f911478882907d6c86e4750ccaa: 64.7ms
0.149548 pex: Re-writing /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/venvs/9b545f5764cfad9084f31d8fa169a9d3b12acb9d/547bc0ed92f40a50fbeafded4a247435b07c2533.lck.work/bin/activate
0.000239 pex: Re-writing /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/venvs/9b545f5764cfad9084f31d8fa169a9d3b12acb9d/547bc0ed92f40a50fbeafded4a247435b07c2533.lck.work/bin/activate.csh
0.000226 pex: Re-writing /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.K8DBvXx7/venvs/9b545f5764cfad9084f31d8fa169a9d3b12acb9d/547bc0ed92f40a50fbeafded4a247435b07c2533.lck.work/bin/activate.fish
0.042920 macOS-12.5.1-arm64-arm-64bit
PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python plotly_via_119.pex 2>&1  1.07s user 1.25s system 92% cpu 2.503 total
...
```

Output for plotly with 2.1.120 is broadly the same, except for the last few lines:
```
...
0.000250 pex: Re-writing /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.dnHmgqeE/venvs/d3179052216c59d19f39c057259d9b6f94635d28/547bc0ed92f40a50fbeafded4a247435b07c2533.lck.work/bin/activate.csh
0.000249 pex: Re-writing /private/var/folders/sv/vd266m4d4lvctgs2wpnhjs9w0000gn/T/tmp.dnHmgqeE/venvs/d3179052216c59d19f39c057259d9b6f94635d28/547bc0ed92f40a50fbeafded4a247435b07c2533.lck.work/bin/activate.fish
4.125079 macOS-12.5.1-arm64-arm-64bit
PEX_VERBOSE=1 PEX_ROOT=$(mktemp -d) python plotly_via_120.pex 2>&1  1.52s user 4.61s system 93% cpu 6.568 total
...
```

In particular there's 4.1s of extra work happening _silently_ (even at `PEX_VERBOSE=9`) after the last "Re-writing" line.

A `py-spy` flame graph of a similar PEX (built with 2.1.125) hints at some possibilities for the location of the regression. Unfortunately, I haven't been able to get this recorded reliably, so it's not the PEX in the reproducer above, nor do I have a trace of 2.1.119 to compare.

![17565-2023-03-07T13-27-24+11-00](https://user-images.githubusercontent.com/1203825/223680555-21816d28-7842-42f3-a5f3-614a685adb35.svg)

 

(As always, thanks for PEX; it is a reliable workhorse in our deploys. 👍 )


Discussion:
You bisected the right change. As you may notice, that change purports to be a fix for loose layout PEXes - which you don't use. So likely the Pex misidentifes itself.
```diff
diff --git a/pex/pex.py b/pex/pex.py
index eebab84..feb50d4 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -172,6 +172,7 @@ class PEX(object):  # noqa: T000
         # type: () -> Layout.Value
         if self._layout is None:
             self._layout = Layout.identify(self._pex)
+        print(">>> I'm {} and my layout is: {}".format(self._pex, self._layout), file=sys.stderr)
         return self._layout

     def pex_info(self, include_env_overrides=True):
```

Shows:
```
$ python -mpex cowsay -c cowsay -ocowsay.pex --venv
>>> I'm /home/jsirois/.pex/pip-20.3.4-patched.pex/3a5d8c60b769730f991f150aaa45975637d8ae30 and my layout is: loose
$ rm -rf ~/.pex
$ ./cowsay.pex 'Moo!'
>>> I'm /home/jsirois/.pex/unzipped_pexes/33d2aea7a351e1778d02fa3b66ec67f811445092 and my layout is: loose
  ____
| Moo! |
  ====
    \
     \
       ^__^
       (oo)\_______
       (__)\       )\/\
           ||----w |
           ||     ||
```

I'm ~afk through March 14th and will fix this then if no-one else has. But that's the issue.

Resolving patch:
diff --git a/pex/layout.py b/pex/layout.py
--- a/pex/layout.py
+++ b/pex/layout.py
@@ -21,11 +21,24 @@
 BOOTSTRAP_DIR = ".bootstrap"
 DEPS_DIR = ".deps"
 PEX_INFO_PATH = "PEX-INFO"
+PEX_LAYOUT_PATH = "PEX-LAYOUT"
 
 
 class Layout(Enum["Layout.Value"]):
     class Value(Enum.Value):
-        pass
+        @classmethod
+        def try_load(cls, pex_directory):
+            # type: (str) -> Optional[Layout.Value]
+            layout = os.path.join(pex_directory, PEX_LAYOUT_PATH)
+            if not os.path.isfile(layout):
+                return None
+            with open(layout) as fp:
+                return Layout.for_value(fp.read().strip())
+
+        def record(self, pex_directory):
+            # type: (str) -> None
+            with open(os.path.join(pex_directory, PEX_LAYOUT_PATH), "w") as fp:
+                fp.write(self.value)
 
     ZIPAPP = Value("zipapp")
     PACKED = Value("packed")
@@ -47,10 +60,23 @@ def identify(cls, pex):
 
         return cls.LOOSE
 
+    @classmethod
+    def identify_original(cls, pex):
+        # type: (str) -> Layout.Value
+        layout = cls.identify(pex)
+        if layout != Layout.LOOSE:
+            return layout
+        return cls.Value.try_load(pex) or Layout.LOOSE
+
 
 class _Layout(object):
-    def __init__(self, path):
-        # type: (str) -> None
+    def __init__(
+        self,
+        layout,  # type: Layout.Value
+        path,  # type: str
+    ):
+        # type: (...) -> None
+        self._layout = layout
         self._path = os.path.normpath(path)
 
     @property
@@ -95,6 +121,10 @@ def extract_main(self, dest_dir):
         # type: (str) -> None
         raise NotImplementedError()
 
+    def record(self, dest_dir):
+        # type: (str) -> None
+        self._layout.record(dest_dir)
+
 
 def _install(
     layout,  # type: _Layout
@@ -165,7 +195,7 @@ def _install(
 
                     layout.extract_pex_info(chroot.work_dir)
                     layout.extract_main(chroot.work_dir)
-
+                    layout.record(chroot.work_dir)
         return install_to
 
 
@@ -176,7 +206,7 @@ def __init__(
         zfp,  # type: zipfile.ZipFile
     ):
         # type: (...) -> None
-        super(_ZipAppPEX, self).__init__(path)
+        super(_ZipAppPEX, self).__init__(Layout.ZIPAPP, path)
         self._zfp = zfp
         self._names = tuple(zfp.namelist())
 
@@ -224,6 +254,10 @@ def __str__(self):
 
 
 class _PackedPEX(_Layout):
+    def __init__(self, path):
+        # type: (str) -> None
+        super(_PackedPEX, self).__init__(Layout.PACKED, path)
+
     def extract_bootstrap(self, dest_dir):
         # type: (str) -> None
         with open_zip(os.path.join(self._path, BOOTSTRAP_DIR)) as zfp:
diff --git a/pex/pex.py b/pex/pex.py
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -171,7 +171,7 @@ def __init__(
     def layout(self):
         # type: () -> Layout.Value
         if self._layout is None:
-            self._layout = Layout.identify(self._pex)
+            self._layout = Layout.identify_original(self._pex)
         return self._layout
 
     def pex_info(self, include_env_overrides=True):
diff --git a/pex/venv/pex.py b/pex/venv/pex.py
--- a/pex/venv/pex.py
+++ b/pex/venv/pex.py
@@ -10,7 +10,7 @@
 from collections import Counter, defaultdict
 from textwrap import dedent
 
-from pex import pex_warnings
+from pex import layout, pex_warnings
 from pex.common import chmod_plus_x, pluralize, safe_mkdir
 from pex.compatibility import is_valid_python_identifier
 from pex.dist_metadata import Distribution
@@ -313,11 +313,12 @@ def _populate_sources(
         src=PEXEnvironment.mount(pex.path()).path,
         dst=venv.site_packages_dir,
         exclude=(
-            pex_info.internal_cache,
-            pex_info.bootstrap,
             "__main__.py",
             "__pycache__",
-            pex_info.PATH,
+            layout.BOOTSTRAP_DIR,
+            layout.DEPS_DIR,
+            layout.PEX_INFO_PATH,
+            layout.PEX_LAYOUT_PATH,
         ),
         symlink=False,
     ):


Test patch:
diff --git a/tests/integration/test_issue_2088.py b/tests/integration/test_issue_2088.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/test_issue_2088.py
@@ -0,0 +1,90 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+import os.path
+import shutil
+import subprocess
+import sys
+
+import pytest
+
+from pex.common import safe_rmtree
+from pex.compatibility import commonpath
+from pex.layout import Layout
+from pex.testing import run_pex_command
+from pex.typing import TYPE_CHECKING
+from pex.venv.virtualenv import Virtualenv
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@pytest.mark.parametrize(
+    "layout", [pytest.param(layout, id=layout.value) for layout in Layout.values()]
+)
+@pytest.mark.parametrize(
+    "symlink_site_packages",
+    [
+        pytest.param(True, id="--no-venv-site-packages-copies"),
+        pytest.param(False, id="--venv-site-packages-copies"),
+    ],
+)
+def test_venv_symlink_site_packages(
+    tmpdir,  # type: Any
+    layout,  # type: Layout.Value
+    symlink_site_packages,  # type: bool
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+    venv_site_packages_copies_arg = (
+        "--no-venv-site-packages-copies" if symlink_site_packages else "--venv-site-packages-copies"
+    )
+    result = run_pex_command(
+        args=[
+            "ansicolors==1.1.8",
+            "--pex-root",
+            pex_root,
+            "--runtime-pex-root",
+            pex_root,
+            "--venv",
+            venv_site_packages_copies_arg,
+            "--layout",
+            layout.value,
+            "--seed",
+            "-o",
+            pex,
+        ]
+    )
+    result.assert_success()
+    venv_pex_path = str(result.output.strip())
+
+    safe_rmtree(pex_root)
+    colors_module_realpath = os.path.realpath(
+        subprocess.check_output(
+            args=[sys.executable, pex, "-c", "import colors; print(colors.__file__)"]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+
+    venv_dir = os.path.dirname(venv_pex_path)
+    virtualenv = Virtualenv(venv_dir)
+    site_packages_dir = os.path.realpath(virtualenv.site_packages_dir)
+
+    ansicolors_venv_package_dir_realpath = os.path.join(site_packages_dir, "colors")
+    assert os.path.isdir(ansicolors_venv_package_dir_realpath)
+
+    symlinks_expected = symlink_site_packages and layout != Layout.LOOSE
+    assert os.path.islink(ansicolors_venv_package_dir_realpath) == symlinks_expected
+
+    if symlinks_expected:
+        installed_wheels_dir_realpath = os.path.realpath(os.path.join(pex_root, "installed_wheels"))
+        assert installed_wheels_dir_realpath == commonpath(
+            (installed_wheels_dir_realpath, colors_module_realpath)
+        )
+    else:
+        assert (
+            os.path.join(ansicolors_venv_package_dir_realpath, "__init__.py")
+            == colors_module_realpath
+        )