Ensure Pip cache operations are atomic. (#1778)

By forcing the temporary directories Pip creates to be on the same
filesystem as the Pip cache, we indirectly ensure the caching Pip does
is atomic and safe in the presence of parallel Pex runs.

Fixes #1776

pex/pip/tool.py

@@ -15,7 +15,7 @@
 
 from pex import dist_metadata, targets, third_party
 from pex.auth import PasswordEntry
-from pex.common import atomic_directory, safe_mkdtemp
+from pex.common import atomic_directory, safe_mkdir, safe_mkdtemp
 from pex.compatibility import unquote, urlparse
 from pex.dist_metadata import ProjectNameAndVersion, Requirement
 from pex.interpreter import PythonInterpreter
@@ -690,6 +690,15 @@ def _spawn_pip_isolated(
         if package_index_configuration:
             extra_env.update(package_index_configuration.env)
 
+        # Ensure the pip cache (`http/` and `wheels/` dirs) is housed in the same partition as the
+        # temporary directories it creates. This is needed to ensure atomic filesystem operations
+        # since Pip relies upon `shutil.move` which is only atomic when `os.rename` can be used.
+        # See https://github.com/pantsbuild/pex/issues/1776 for an example of the issues non-atomic
+        # moves lead to in the `pip wheel` case.
+        pip_tmpdir = os.path.join(cache or ENV.PEX_ROOT, "tmp")
+        safe_mkdir(pip_tmpdir)
+        extra_env.update(TMPDIR=pip_tmpdir)
+
         with ENV.strip().patch(
             PEX_ROOT=cache or ENV.PEX_ROOT,
             PEX_VERBOSE=str(ENV.PEX_VERBOSE),