Set PEX_MAX_INSTALL_JOBS on pants invocations of pexes

pantsbuild · Apr 3, 2024 · e9c5ec1 · e9c5ec1
1 parent 01916cc
commit e9c5ec1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git 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"}),
         ),
     )
 

diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py
@@ -1163,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,
     }
@@ -1202,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
@@ -1235,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
@@ -1258,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
@@ -1276,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,