Repository: pantsbuild/pants

Issue:
Benchmark various defaults for PEX compression for internal PEXes
Pants uses lots of PEXes "internally", which are never exported for users to consume. Since these are implementation details, we have free-rein to choose an internal-default [PEX compression level](https://github.com/pantsbuild/pex/pull/1705)... and some values can make a considerable difference in performance.

To experiment/adjust for "internal only", logic similar to the logic for PEX's `layout` setting could be added, which set a default compression value for internal use: https://github.com/pantsbuild/pants/blob/902d3acb84c1b64c53c05e882200e4d578884c03/src/python/pants/backend/python/util_rules/pex.py#L196-L198


Discussion:
Are there internal PEX that aren't `PACKED` (grepping for `layout=`, `PexLayout.` and `internal_only=` suggests no)? If all of them are packed, will lowering the compression level do anything?
Yes. A packed layout consists of N + 1 zips. 1 PEX `.bootstrap` code zip and 1 per installed wheel contained in the PEX under `.deps/`. By default these N+1 are compressed, but turning off compression turns off their compression just as it does for a single zip zipapp layout.
Has anyone tried seeing whether PEX files might start up a little faster too if they're uncompressed?
All PEX files are ~uncompressed (unzipped). Whether `--venv` or normal (zipapp), and regardless of `--layout`, A PEX file is extracted in one way or another on 1st run and forever more just re-directs to the extracted location. The re-direct overhead is ~O(50ms) except when using `--sh-boot` which takes it down to ~1ms.
Oh yes! Thanks!
To underscore - the benchmark here is about PEX creation time then, not run time.
With https://github.com/pantsbuild/pex/releases/tag/v2.1.154, potentially we could be considering using `--no-pre-install-wheels` for internal PEXes, and maybe `--max-install-jobs`. Quoting https://github.com/pantsbuild/pex/issues/2292#issuecomment-1854582647, for perf results with PyTorch (i.e. large):

Status quo | With --no-pre-install-wheels | --no-pre-install-wheels | savings
-- | -- | -- | --
Cold build and run 1st local machine | 188.51s | 29.80s | 84% faster
Cold run 1st remote machine | 26.74s | 29.60s | 11% slower
Cold run 1st remote machine parallel | 12.79s | 14.06s | 10% slower
Size (bytes) | 2679282217 | 2678537958 | 0.02% smaller


(#20347 updates to pex 2.1.155 by default, but doesn't touch `version_constraints`, so if we actually start using that flag internally we'll need to either bump the constraints or make it conditional on version.)

Resolving patch:
diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py
--- a/src/python/pants/backend/python/util_rules/pex.py
+++ b/src/python/pants/backend/python/util_rules/pex.py
@@ -705,6 +705,12 @@ async def build_pex(
     if request.pex_path:
         argv.extend(["--pex-path", ":".join(pex.name for pex in request.pex_path)])
 
+    if request.internal_only:
+        # An internal-only runs on a single machine, and pre-installing wheels is wasted work in
+        # that case (see https://github.com/pex-tool/pex/issues/2292#issuecomment-1854582647 for
+        # analysis).
+        argv.append("--no-pre-install-wheels")
+
     argv.append(f"--sources-directory={source_dir_name}")
     sources_digest_as_subdir = await Get(
         Digest, AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name)
@@ -1157,6 +1163,9 @@ async def setup_pex_process(request: PexProcess, pex_environment: PexEnvironment
     complete_pex_env = pex_environment.in_sandbox(working_directory=request.working_directory)
     argv = complete_pex_env.create_argv(pex.name, *request.argv)
     env = {
+        # Set this in case this PEX was built with --no-pre-install-wheels, and thus parallelising
+        # the install on cold boot is handy.
+        "PEX_MAX_INSTALL_JOBS": str(request.concurrency_available),
         **complete_pex_env.environment_dict(python=pex.python),
         **request.extra_env,
     }
@@ -1196,7 +1205,7 @@ class VenvPexProcess:
     level: LogLevel
     input_digest: Digest | None
     working_directory: str | None
-    extra_env: FrozenDict[str, str] | None
+    extra_env: FrozenDict[str, str]
     output_files: tuple[str, ...] | None
     output_directories: tuple[str, ...] | None
     timeout_seconds: int | None
@@ -1229,7 +1238,7 @@ def __init__(
         object.__setattr__(self, "level", level)
         object.__setattr__(self, "input_digest", input_digest)
         object.__setattr__(self, "working_directory", working_directory)
-        object.__setattr__(self, "extra_env", FrozenDict(extra_env) if extra_env else None)
+        object.__setattr__(self, "extra_env", FrozenDict(extra_env or {}))
         object.__setattr__(self, "output_files", tuple(output_files) if output_files else None)
         object.__setattr__(
             self, "output_directories", tuple(output_directories) if output_directories else None
@@ -1252,6 +1261,12 @@ async def setup_venv_pex_process(
         else venv_pex.pex.argv0
     )
     argv = (pex_bin, *request.argv)
+    env = {
+        # Set this in case this PEX was built with --no-pre-install-wheels, and thus parallelising
+        # the install on cold boot is handy.
+        "PEX_MAX_INSTALL_JOBS": str(request.concurrency_available),
+        **request.extra_env,
+    }
     input_digest = (
         await Get(Digest, MergeDigests((venv_pex.digest, request.input_digest)))
         if request.input_digest
@@ -1270,7 +1285,7 @@ async def setup_venv_pex_process(
         level=request.level,
         input_digest=input_digest,
         working_directory=request.working_directory,
-        env=request.extra_env,
+        env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
         append_only_caches=append_only_caches,
diff --git a/src/python/pants/backend/python/util_rules/pex_cli.py b/src/python/pants/backend/python/util_rules/pex_cli.py
--- a/src/python/pants/backend/python/util_rules/pex_cli.py
+++ b/src/python/pants/backend/python/util_rules/pex_cli.py
@@ -40,7 +40,7 @@ class PexCli(TemplatedExternalTool):
 
     default_version = "v2.3.0"
     default_url_template = "https://github.com/pex-tool/pex/releases/download/{version}/pex"
-    version_constraints = ">=2.1.161,<3.0"
+    version_constraints = ">=2.3.0,<3.0"
 
     @classproperty
     def default_known_versions(cls):


Test patch:
diff --git a/src/python/pants/backend/python/goals/publish_test.py b/src/python/pants/backend/python/goals/publish_test.py
--- a/src/python/pants/backend/python/goals/publish_test.py
+++ b/src/python/pants/backend/python/goals/publish_test.py
@@ -136,7 +136,13 @@ def test_twine_upload(rule_runner, packages) -> None:
                 "my-package-0.1.0.tar.gz",
                 "my_package-0.1.0-py3-none-any.whl",
             ),
-            env=FrozenDict({"TWINE_USERNAME": "whoareyou", "TWINE_PASSWORD": "secret"}),
+            env=FrozenDict(
+                {
+                    "TWINE_USERNAME": "whoareyou",
+                    "TWINE_PASSWORD": "secret",
+                    "PEX_MAX_INSTALL_JOBS": "0",
+                }
+            ),
         ),
     )
     assert_package(
@@ -156,7 +162,7 @@ def test_twine_upload(rule_runner, packages) -> None:
                 "my-package-0.1.0.tar.gz",
                 "my_package-0.1.0-py3-none-any.whl",
             ),
-            env=FrozenDict({"TWINE_USERNAME": "whoami"}),
+            env=FrozenDict({"TWINE_USERNAME": "whoami", "PEX_MAX_INSTALL_JOBS": "0"}),
         ),
     )