Repository: pex-tool/pex

Issue:
Parallell pex creation fails due to possible race conditions in package installation
We're using pex with makefiles and parallell jobs in CI and are noticing spurious errors like the one below
```
pex --find-links out/wheels -o out/pexes/pex1.pex ./pex1 -e pex1:main
pex --find-links out/wheels -o out/pexes/pex2.pex -r pex2/requirements.txt --python-script pex2/entrypoint.py
Could not find the installation RECORD for zipp 3.10.0 under /root/.pex/installed_wheels/4fcb6f278987a6605757302a6e40e896257570d11c51628968ccb2a47e80c6c1/zipp-3.10.0-py3-none-any.whl.workdir
```
The other pex built fine and if I retry it will work. Problem is probably made worse by having all wheels cached so no time is spent downloading them which makes it more likely to run into races when installing.

This would probably need some locking around the installed_wheels folder or more fine grained locking/atomik operations for the real installation folder.

It's not urgent for us as I can avoid making those in parallell or ensure separate pexdirs but it was somewhat of a nasty issue to debug and would perhaps save someone the time and frustration in the future.


Discussion:
> This would probably need some locking around the installed_wheels folder or more fine grained locking/atomik operations for the real installation folder.

Yeah, Pex has an atomic_directory abstraction it uses for all such things:
https://github.com/pantsbuild/pex/blob/22e651589b3f2da26da6786dcbd1553fd7c8d4a6/pex/atomic_directory.py#L102-L122

Clearly something is broken there though. Pex is used by Pants quite heavily and it runs many Pex processes in parallel as well. That had smoked out these issues ~2 years ago, but there is clearly something wrong (see #1969). I explained a bit over on that issue, but if you're on Pex 2.1.112 can you try exporting `_PEX_FILE_LOCK_STYLE=bsd` and report back if that helps? You'd need to wipe the PEX_ROOT (default ~/.pex) before trying the parallel make to start from a clean slate with this maybe-fix.
@tapetersen, @christopherfrieler had success using `_PEX_FILE_LOCK_STYLE=bsd` over in #1969. If you get a chance to try this and report back, I'd love to have all the data points I can get to close this out in #1971.
I'm running the builds with pex installed in a .venv and get the following (along with the same errors when trying with the env-var)
```
Ignoring the following environment variables in Pex venv mode:
_PEX_FILE_LOCK_STYLE=bsd
```
Will try to install it outside venv and try again. and see if I get any difference
Hmm even Outside of a venv with `_PEX_FILE_LOCK_STYLE=bsd` I get failed builds and sometimes still the message about it being ignored but sometimes it also fails without that warning

Error without that warning
```
Could not find the installation RECORD for commonmark 0.9.1 under /root/.pex/installed_wheels/da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9/commonmark-0.9.1-py2.py3-none-any.whl.workdir
```

Error with that warning (both having the variable set from commandline
```
pex --find-links out/wheels -o out/pexes/pex1.pex ./pex1 -e pex1:main
pex --find-links out/wheels -o out/pexes/pex2.pex -r pex2/requirements.txt --python-script pex2entrypoint.py
pid 3227 -> /root/.pex/venvs/3ff4eec5a4c6144123a0484f8822cd11473e74e9/5fd7049af63e03f347278c89401424cd9731df9a/bin/python -sE /root/.pex/venvs/3ff4eec5a4c6144123a0484f8822cd11473e74e9/5fd7049af63e03f347278c89401424cd9731df9a/pex --disable-pip-version-check --no-python-version-warning --exists-action a --no-input --use-deprecated legacy-resolver --isolated -q --cache-dir /root/.pex/pip_cache install --no-deps --no-index --only-binary :all: --no-warn-script-location --force-reinstall --ignore-installed --ignore-requires-python --prefix /root/.pex/installed_wheels/da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9/commonmark-0.9.1-py2.py3-none-any.whl.workdir --no-compile /root/.pex/downloads/resolver_download.ispj28lb/usr.local.bin.python3.8/commonmark-0.9.1-py2.py3-none-any.whl exited with 1 and STDERR:
Ignoring the following environment variables in Pex venv mode:
_PEX_FILE_LOCK_STYLE=bsd

ERROR: Could not install packages due to an EnvironmentError: [Errno 2] No such file or directory: '/root/.pex/installed_wheels/da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9/commonmark-0.9.1-py2.py3-none-any.whl.workdir/lib/python3.8/site-packages/commonmark/tests/unit_tests.py'
```
`pex --version
2.1.112`

Also running with `PIP_NO_INDEX=1 PIP_FIND_LINKS=/root/wheels` so all wheels are pre-downloaded and are just installed which as mentioned I suspect makes it more likely to run into it.
Ok, thanks for that data. The warning is expected, the continued failures are not. I need to think more here.
@jsirois No problem, like I said no rush for us. We can work around it in the makefile without any problems but happy to provide data if you want to resolve it.
Actually - one question presents itself. Are your errors confined to missing files in paths that are prefixed by `/root/.pex/installed_wheels/`? If so, that might be a special case I can look at. I think that is the oldest use of atomic_directory and it may not use the locking at all, just an atomic rename. Your examples indicate this happens when Pex is calling `pip install` (Pex dogfoods itself and runs Pip as a venv PEX in there).
Hmm not sure about all errors but from a quick glance all I can see do indeed appear under that directory.
. I prepared a small example for reproduction as well that may help you here in a gist
https://gist.github.com/tapetersen/32d564029be43022502e814a8d3e3057

Excellent - thank you.
I repro - which is great. Thank you for setting that gist up. I should be able to get to the bottom of this pretty quickly.

Resolving patch:
diff --git a/pex/atomic_directory.py b/pex/atomic_directory.py
--- a/pex/atomic_directory.py
+++ b/pex/atomic_directory.py
@@ -8,6 +8,7 @@
 import os
 import threading
 from contextlib import contextmanager
+from uuid import uuid4
 
 from pex import pex_warnings
 from pex.common import safe_mkdir, safe_rmtree
@@ -19,10 +20,50 @@
 
 
 class AtomicDirectory(object):
-    def __init__(self, target_dir):
-        # type: (str) -> None
+    """A directory whose contents are populated atomically.
+
+    By default, an atomic directory allows racing processes to populate a directory atomically.
+    Each gets its own unique work directory to populate non-atomically, but the final target
+    directory is swapped to atomically from one of the racing work directories.
+
+    In order to lock the atomic directory so that only 1 process works to populate it, use the
+    `atomic_directory` context manager.
+
+    If the target directory will have immutable contents, either approach will do. If not, the
+    exclusively locked `atomic_directory` context manager should be used.
+
+    The common case for a non-obvious mutable directory in Python is any directory `.py` files are
+    populated to. Those files will later be bytecode compiled to adjacent `.pyc` files on the fly
+    by the Python interpreter and can go missing underneath a process looking at that directory if
+    AtomicDirectory is used directly. For the target directory `sys_path_entry` that failure mode
+    looks like:
+
+    Process A -> Starts work to create `sys_path_entry`.
+    Process B -> Starts work to create `sys_path_entry`.
+    Process A -> Atomically creates `sys_path_entry`.
+    Process C -> Sees `sys_path_entry` from Process A and starts running Python code in that dir.
+    Process D -> Sees `sys_path_entry` from Process A and starts running Python code in that dir.
+    Process D -> Succeeds importing `foo`.
+    Process C -> Starts to import `sys_path_entry/foo.py` and Python sees the corresponding .pyc
+                 file already exists (Process D created it).
+    Process B -> Atomically creates `sys_path_entry`, replacing the result from Process A and
+                 disappearing any `.pyc` files.
+    Process C -> Goes to import from the `.pyc` file it found and errors since it is gone.
+
+    The background facts in this case are that CPython reasonably does a check then act surrounding
+    .pyc files and uses no lock. It assumes Python source trees will not be disturbed during its
+    run. Without an exclusively locked `atomic_directory` Pex can allow the check-then-act window to
+    be observed by racing processes.
+    """
+
+    def __init__(
+        self,
+        target_dir,  # type: str
+        locked=False,  # type: bool
+    ):
+        # type: (...) -> None
         self._target_dir = target_dir
-        self._work_dir = "{}.workdir".format(target_dir)
+        self._work_dir = "{}.{}.work".format(target_dir, "lck" if locked else uuid4().hex)
 
     @property
     def work_dir(self):
@@ -124,7 +165,7 @@ def atomic_directory(
     # We use double-checked locking with the check being target_dir existence and the lock being an
     # exclusive blocking file lock.
 
-    atomic_dir = AtomicDirectory(target_dir=target_dir)
+    atomic_dir = AtomicDirectory(target_dir=target_dir, locked=True)
     if atomic_dir.is_finalized():
         # Our work is already done for us so exit early.
         yield atomic_dir
diff --git a/pex/venv/pex.py b/pex/venv/pex.py
--- a/pex/venv/pex.py
+++ b/pex/venv/pex.py
@@ -426,6 +426,8 @@ def sys_executable_paths():
                     # These are used by Pex's Pip venv to implement universal locks.
                     "_PEX_PYTHON_VERSIONS_FILE",
                     "_PEX_TARGET_SYSTEMS_FILE",
+                    # This is used as an experiment knob for atomic_directory locking.
+                    "_PEX_FILE_LOCK_STYLE",
                 )
             ]
             if ignored_pex_env_vars:


Test patch:
diff --git a/tests/test_atomic_directory.py b/tests/test_atomic_directory.py
--- a/tests/test_atomic_directory.py
+++ b/tests/test_atomic_directory.py
@@ -1,8 +1,24 @@
 # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-from pex.atomic_directory import FileLockStyle, _is_bsd_lock
+import errno
+import os
+from contextlib import contextmanager
+
+import pytest
+
+from pex.atomic_directory import AtomicDirectory, FileLockStyle, _is_bsd_lock, atomic_directory
+from pex.common import temporary_dir, touch
 from pex.testing import environment_as
+from pex.typing import TYPE_CHECKING
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock  # type: ignore[no-redef,import]
+
+if TYPE_CHECKING:
+    from typing import Iterator, Optional, Type
 
 
 def test_is_bsd_lock():
@@ -26,3 +42,100 @@ def test_is_bsd_lock():
         ), "Expected the default lock style to be taken from the environment when defined."
         assert not _is_bsd_lock(lock_style=FileLockStyle.POSIX)
         assert _is_bsd_lock(lock_style=FileLockStyle.BSD)
+
+
+@contextmanager
+def maybe_raises(exception=None):
+    # type: (Optional[Type[Exception]]) -> Iterator[None]
+    @contextmanager
+    def noop():
+        yield
+
+    context = noop() if exception is None else pytest.raises(exception)
+    with context:
+        yield
+
+
+def atomic_directory_finalize_test(errno, expect_raises=None):
+    # type: (int, Optional[Type[Exception]]) -> None
+    with mock.patch("os.rename", spec_set=True, autospec=True) as mock_rename:
+        mock_rename.side_effect = OSError(errno, os.strerror(errno))
+        with maybe_raises(expect_raises):
+            AtomicDirectory("to.dir").finalize()
+
+
+def test_atomic_directory_finalize_eexist():
+    # type: () -> None
+    atomic_directory_finalize_test(errno.EEXIST)
+
+
+def test_atomic_directory_finalize_enotempty():
+    # type: () -> None
+    atomic_directory_finalize_test(errno.ENOTEMPTY)
+
+
+def test_atomic_directory_finalize_eperm():
+    # type: () -> None
+    atomic_directory_finalize_test(errno.EPERM, expect_raises=OSError)
+
+
+def test_atomic_directory_empty_workdir_finalize():
+    # type: () -> None
+    with temporary_dir() as sandbox:
+        target_dir = os.path.join(sandbox, "target_dir")
+        assert not os.path.exists(target_dir)
+
+        with atomic_directory(target_dir) as atomic_dir:
+            assert not atomic_dir.is_finalized()
+            assert target_dir == atomic_dir.target_dir
+            assert os.path.exists(atomic_dir.work_dir)
+            assert os.path.isdir(atomic_dir.work_dir)
+            assert [] == os.listdir(atomic_dir.work_dir)
+
+            touch(os.path.join(atomic_dir.work_dir, "created"))
+
+            assert not os.path.exists(target_dir)
+
+        assert not os.path.exists(atomic_dir.work_dir), "The work_dir should always be cleaned up."
+        assert os.path.exists(os.path.join(target_dir, "created"))
+
+
+def test_atomic_directory_empty_workdir_failure():
+    # type: () -> None
+    class SimulatedRuntimeError(RuntimeError):
+        pass
+
+    with temporary_dir() as sandbox:
+        target_dir = os.path.join(sandbox, "target_dir")
+        with pytest.raises(SimulatedRuntimeError):
+            with atomic_directory(target_dir) as atomic_dir:
+                assert not atomic_dir.is_finalized()
+                touch(os.path.join(atomic_dir.work_dir, "created"))
+                raise SimulatedRuntimeError()
+
+        assert not os.path.exists(  # type: ignore[unreachable]
+            atomic_dir.work_dir
+        ), "The work_dir should always be cleaned up."
+        assert not os.path.exists(target_dir), (
+            "When the context raises the work_dir it was given should not be moved to the "
+            "target_dir."
+        )
+
+
+def test_atomic_directory_empty_workdir_finalized():
+    # type: () -> None
+    with temporary_dir() as target_dir:
+        with atomic_directory(target_dir) as work_dir:
+            assert (
+                work_dir.is_finalized()
+            ), "When the target_dir exists no work_dir should be created."
+
+
+def test_atomic_directory_locked_mode():
+    # type: () -> None
+
+    assert AtomicDirectory("unlocked").work_dir != AtomicDirectory("unlocked").work_dir
+    assert (
+        AtomicDirectory("locked", locked=True).work_dir
+        == AtomicDirectory("locked", locked=True).work_dir
+    )
diff --git a/tests/test_common.py b/tests/test_common.py
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -4,11 +4,9 @@
 import contextlib
 import errno
 import os
-from contextlib import contextmanager
 
 import pytest
 
-from pex.atomic_directory import AtomicDirectory, atomic_directory
 from pex.common import (
     Chroot,
     PermPreservingZipFile,
@@ -29,94 +27,7 @@
     import mock  # type: ignore[no-redef,import]
 
 if TYPE_CHECKING:
-    from typing import Any, Iterator, Optional, Tuple, Type
-
-
-@contextmanager
-def maybe_raises(exception=None):
-    # type: (Optional[Type[Exception]]) -> Iterator[None]
-    @contextmanager
-    def noop():
-        yield
-
-    context = noop() if exception is None else pytest.raises(exception)
-    with context:
-        yield
-
-
-def atomic_directory_finalize_test(errno, expect_raises=None):
-    # type: (int, Optional[Type[Exception]]) -> None
-    with mock.patch("os.rename", spec_set=True, autospec=True) as mock_rename:
-        mock_rename.side_effect = OSError(errno, os.strerror(errno))
-        with maybe_raises(expect_raises):
-            AtomicDirectory("to.dir").finalize()
-
-
-def test_atomic_directory_finalize_eexist():
-    # type: () -> None
-    atomic_directory_finalize_test(errno.EEXIST)
-
-
-def test_atomic_directory_finalize_enotempty():
-    # type: () -> None
-    atomic_directory_finalize_test(errno.ENOTEMPTY)
-
-
-def test_atomic_directory_finalize_eperm():
-    # type: () -> None
-    atomic_directory_finalize_test(errno.EPERM, expect_raises=OSError)
-
-
-def test_atomic_directory_empty_workdir_finalize():
-    # type: () -> None
-    with temporary_dir() as sandbox:
-        target_dir = os.path.join(sandbox, "target_dir")
-        assert not os.path.exists(target_dir)
-
-        with atomic_directory(target_dir) as atomic_dir:
-            assert not atomic_dir.is_finalized()
-            assert target_dir == atomic_dir.target_dir
-            assert os.path.exists(atomic_dir.work_dir)
-            assert os.path.isdir(atomic_dir.work_dir)
-            assert [] == os.listdir(atomic_dir.work_dir)
-
-            touch(os.path.join(atomic_dir.work_dir, "created"))
-
-            assert not os.path.exists(target_dir)
-
-        assert not os.path.exists(atomic_dir.work_dir), "The work_dir should always be cleaned up."
-        assert os.path.exists(os.path.join(target_dir, "created"))
-
-
-def test_atomic_directory_empty_workdir_failure():
-    # type: () -> None
-    class SimulatedRuntimeError(RuntimeError):
-        pass
-
-    with temporary_dir() as sandbox:
-        target_dir = os.path.join(sandbox, "target_dir")
-        with pytest.raises(SimulatedRuntimeError):
-            with atomic_directory(target_dir) as atomic_dir:
-                assert not atomic_dir.is_finalized()
-                touch(os.path.join(atomic_dir.work_dir, "created"))
-                raise SimulatedRuntimeError()
-
-        assert not os.path.exists(  # type: ignore[unreachable]
-            atomic_dir.work_dir
-        ), "The work_dir should always be cleaned up."
-        assert not os.path.exists(target_dir), (
-            "When the context raises the work_dir it was given should not be moved to the "
-            "target_dir."
-        )
-
-
-def test_atomic_directory_empty_workdir_finalized():
-    # type: () -> None
-    with temporary_dir() as target_dir:
-        with atomic_directory(target_dir) as work_dir:
-            assert (
-                work_dir.is_finalized()
-            ), "When the target_dir exists no work_dir should be created."
+    from typing import Any, Iterator, Tuple
 
 
 def extract_perms(path):