Use the default multiprocessing pool configuration. (#2596)

Previously we were forcing use of fork for all OSes to try to juice
performance at the cost of running into threading issues at some point
when things got more complex. They may have reached that point with CI
lockups happening on macOS. Revert to the default pool setup to attempt
to fix macOS hangs.

To support this, also propagate any CLI global `Variables` adjustments
to subprocesses. Previously these were only propagated for the current
process; now the adjustments propagate through the environment so that
subprocesses see the same adjustments. In particular, this ensures
subprocesses use the same `PEX_ROOT` when a fallback has to be used to
work around lack of write perms.

Author: jsirois

pex/commands/command.py

@@ -17,7 +17,7 @@
 from pex import pex_warnings
 from pex.argparse import HandleBoolAction
 from pex.cache import access as cache_access
-from pex.common import safe_mkdtemp, safe_open
+from pex.common import environment_as, safe_mkdtemp, safe_open
 from pex.compatibility import shlex_quote
 from pex.result import Error, Ok, Result
 from pex.typing import TYPE_CHECKING, Generic, cast
@@ -379,7 +379,7 @@ def warn_ignore_pex_root(set_via):
             else:
                 pex_root = options.cache_dir or options.pex_root or ENV.PEX_ROOT
 
-            with ENV.patch(PEX_ROOT=pex_root, TMPDIR=tmpdir) as env:
+            with ENV.patch(PEX_ROOT=pex_root, TMPDIR=tmpdir) as env, environment_as(**env):
                 cache_access.read_write()
                 yield env
 

pex/common.py

@@ -17,6 +17,7 @@
 import time
 import zipfile
 from collections import defaultdict, namedtuple
+from contextlib import contextmanager
 from datetime import datetime
 from uuid import uuid4
 from zipfile import ZipFile, ZipInfo
@@ -917,3 +918,28 @@ def iter_copytree(
         if copy_mode is CopyMode.SYMLINK:
             # Once we've symlinked the top-level directories and files, we've "copied" everything.
             return
+
+
+@contextmanager
+def environment_as(**kwargs):
+    # type: (**Any) -> Iterator[None]
+    """Mutates the `os.environ` for the duration of the context.
+
+    Keyword arguments with None values are removed from os.environ (if present) and all other
+    keyword arguments are added or updated in `os.environ` with the values taken from the
+    stringification (`str(...)`) of each value.
+    """
+    existing = {key: os.environ.get(key) for key in kwargs}
+
+    def adjust_environment(mapping):
+        for key, value in mapping.items():
+            if value is not None:
+                os.environ[key] = str(value)
+            else:
+                os.environ.pop(key, None)
+
+    adjust_environment(kwargs)
+    try:
+        yield
+    finally:
+        adjust_environment(existing)

pex/jobs.py

@@ -17,7 +17,7 @@
 from pex.common import pluralize
 from pex.compatibility import Queue, cpu_count
 from pex.tracer import TRACER
-from pex.typing import TYPE_CHECKING, Generic, cast
+from pex.typing import TYPE_CHECKING, Generic
 
 if TYPE_CHECKING:
     from typing import (
@@ -679,12 +679,8 @@ def join(self):
 @contextmanager
 def _mp_pool(size):
     # type: (int) -> Iterator[Pool]
-    try:
-        context = multiprocessing.get_context("fork")  # type: ignore[attr-defined]
-        pool = cast("Pool", context.Pool(processes=size))
-    except (AttributeError, ValueError):
-        pool = multiprocessing.Pool(processes=size)
 
+    pool = multiprocessing.Pool(processes=size)
     try:
         yield pool
     finally:

testing/__init__.py

@@ -674,25 +674,6 @@ def all_python_venvs(system_site_packages=False):
     )
 
 
-@contextmanager
-def environment_as(**kwargs):
-    # type: (**Any) -> Iterator[None]
-    existing = {key: os.environ.get(key) for key in kwargs}
-
-    def adjust_environment(mapping):
-        for key, value in mapping.items():
-            if value is not None:
-                os.environ[key] = str(value)
-            else:
-                os.environ.pop(key, None)
-
-    adjust_environment(kwargs)
-    try:
-        yield
-    finally:
-        adjust_environment(existing)
-
-
 @contextmanager
 def pushd(directory):
     # type: (Text) -> Iterator[None]

tests/integration/cli/commands/test_cache_prune.py

@@ -27,13 +27,13 @@
     VenvDirs,
 )
 from pex.cli.commands.cache.du import DiskUsage
-from pex.common import safe_open
+from pex.common import environment_as, safe_open
 from pex.pep_503 import ProjectName
 from pex.pex_info import PexInfo
 from pex.pip.version import PipVersion, PipVersionValue
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV
-from testing import environment_as, run_pex_command
+from testing import run_pex_command
 from testing.cli import run_pex3
 from testing.pytest.tmp import Tempdir
 

tests/integration/test_integration.py

@@ -19,7 +19,15 @@
 
 from pex import targets
 from pex.cache.dirs import CacheDir, InterpreterDir
-from pex.common import is_exe, safe_mkdir, safe_open, safe_rmtree, temporary_dir, touch
+from pex.common import (
+    environment_as,
+    is_exe,
+    safe_mkdir,
+    safe_open,
+    safe_rmtree,
+    temporary_dir,
+    touch,
+)
 from pex.compatibility import WINDOWS, commonpath
 from pex.dist_metadata import Distribution, Requirement, is_wheel
 from pex.fetcher import URLFetcher
@@ -46,7 +54,6 @@
     IntegResults,
     built_wheel,
     ensure_python_interpreter,
-    environment_as,
     get_dep_dist_names_from_pex,
     make_env,
     run_pex_command,

tests/integration/test_issue_157.py

@@ -14,10 +14,11 @@
 import pytest
 from colors import color  # vendor:skip
 
+from pex.common import environment_as
 from pex.pex_info import PexInfo
 from pex.typing import TYPE_CHECKING
 from pex.version import __version__
-from testing import IS_PYPY, environment_as, make_env, run_pex_command, scie
+from testing import IS_PYPY, make_env, run_pex_command, scie
 
 if TYPE_CHECKING:
     from typing import Any, Iterable, Iterator, List, Tuple

tests/resolve/test_target_options.py

@@ -10,6 +10,7 @@
 import pytest
 
 import pex.resolve.target_configuration
+from pex.common import environment_as
 from pex.interpreter import PythonInterpreter
 from pex.pep_425 import CompatibilityTags
 from pex.pep_508 import MarkerEnvironment
@@ -20,7 +21,7 @@
 from pex.targets import CompletePlatform, Targets
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV
-from testing import IS_MAC, environment_as
+from testing import IS_MAC
 
 if TYPE_CHECKING:
     from typing import Any, Dict, Iterable, List, Optional, Tuple, Type

tests/test_atomic_directory.py

@@ -9,9 +9,8 @@
 import pytest
 
 from pex.atomic_directory import AtomicDirectory, FileLockStyle, _is_bsd_lock, atomic_directory
-from pex.common import temporary_dir, touch
+from pex.common import environment_as, temporary_dir, touch
 from pex.typing import TYPE_CHECKING
-from testing import environment_as
 
 try:
     from unittest import mock

tests/test_interpreter.py

@@ -14,7 +14,7 @@
 import pytest
 
 from pex.cache.dirs import InterpreterDir
-from pex.common import chmod_plus_x, safe_mkdir, safe_mkdtemp, temporary_dir, touch
+from pex.common import chmod_plus_x, environment_as, safe_mkdir, safe_mkdtemp, temporary_dir, touch
 from pex.executor import Executor
 from pex.interpreter import PythonInterpreter, create_shebang
 from pex.jobs import Job
@@ -30,7 +30,6 @@
     ensure_python_distribution,
     ensure_python_interpreter,
     ensure_python_venv,
-    environment_as,
     pushd,
 )
 from testing.pytest.tmp import TempdirFactory

tests/test_pex.py

@@ -16,7 +16,7 @@
 import pytest
 
 from pex import resolver
-from pex.common import safe_mkdir, safe_open, temporary_dir
+from pex.common import environment_as, safe_mkdir, safe_open, temporary_dir
 from pex.compatibility import PY2, WINDOWS, to_bytes
 from pex.dist_metadata import Distribution
 from pex.interpreter import PythonIdentity, PythonInterpreter
@@ -33,7 +33,6 @@
     PY_VER,
     WheelBuilder,
     ensure_python_interpreter,
-    environment_as,
     install_wheel,
     make_bdist,
     run_simple_pex,

tests/test_pip.py

@@ -13,7 +13,7 @@
 
 import pytest
 
-from pex.common import safe_rmtree
+from pex.common import environment_as, safe_rmtree
 from pex.dist_metadata import Distribution, Requirement
 from pex.interpreter import PythonInterpreter
 from pex.jobs import Job
@@ -29,7 +29,7 @@
 from pex.typing import TYPE_CHECKING
 from pex.variables import ENV
 from pex.venv.virtualenv import Virtualenv
-from testing import IS_LINUX, PY310, ensure_python_interpreter, environment_as
+from testing import IS_LINUX, PY310, ensure_python_interpreter
 from testing.pytest.tmp import Tempdir
 
 if TYPE_CHECKING:

tests/test_requirements.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from pex.common import safe_open, temporary_dir, touch
+from pex.common import environment_as, safe_open, temporary_dir, touch
 from pex.dist_metadata import Requirement
 from pex.fetcher import URLFetcher
 from pex.requirements import (
@@ -30,7 +30,6 @@
 )
 from pex.third_party.packaging.markers import Marker
 from pex.typing import TYPE_CHECKING
-from testing import environment_as
 
 if TYPE_CHECKING:
     from typing import Any, Iterable, Iterator, List, Optional, Union

tests/test_variables.py

@@ -7,12 +7,12 @@
 import pytest
 
 from pex import pex_warnings
+from pex.common import environment_as
 from pex.compatibility import PY2
 from pex.pex_warnings import PEXWarning
 from pex.typing import TYPE_CHECKING
 from pex.util import named_temporary_file
 from pex.variables import NoValueError, Variables
-from testing import environment_as
 from testing.pytest.tmp import Tempdir
 
 if TYPE_CHECKING: