Use append_only_caches in Pex processes. (#11760) (#11765)

Previously we cheated and exported PEX_ROOT as the absolute path of the
named caches dir + "pex_root". Since we obtained the absolute path from
the local machine, this cheat was unmasked by remote execution where
the absolute path of the named caches dir was different.

Now we use a relative PEX_ROOT that ties in with append_only_caches.
The VenvPex scripts are updated to use relative paths as well.

Fixes #11753
Fixes #11742

(cherry picked from commit 4210c6c)
___
And in support of the above:

Simplify Black file specification. (#11762)
    
Black supports passing an explicit list of files and we have that. This
clears the way for having the pex_root named cache symlinked into the
Black Process chroot without having Black trying to recursively format
everything in there.

(cherry picked from commit c927cee)

src/python/pants/backend/codegen/protobuf/python/rules.py

@@ -17,6 +17,7 @@
     VenvPex,
     VenvPexRequest,
 )
+from pants.backend.python.util_rules.pex_environment import PexEnvironment
 from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
 from pants.core.util_rules.source_files import SourceFilesRequest
 from pants.core.util_rules.stripped_source_files import StrippedSourceFiles
@@ -55,6 +56,7 @@ async def generate_python_from_protobuf(
     protoc: Protoc,
     grpc_python_plugin: GrpcPythonPlugin,
     python_protobuf_subsystem: PythonProtobufSubsystem,
+    pex_environment: PexEnvironment,
 ) -> GeneratedSources:
     download_protoc_request = Get(
         DownloadedExternalTool, ExternalToolRequest, protoc.get_request(Platform.current)
@@ -161,6 +163,7 @@ async def generate_python_from_protobuf(
             description=f"Generating Python sources from {request.protocol_target.address}.",
             level=LogLevel.DEBUG,
             output_directories=(output_dir,),
+            append_only_caches=pex_environment.append_only_caches(),
         ),
     )
 

src/python/pants/backend/python/lint/black/rules.py

@@ -1,9 +1,7 @@
 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-import re
 from dataclasses import dataclass
-from pathlib import PurePath
 from typing import Tuple
 
 from pants.backend.python.lint.black.subsystem import Black
@@ -62,13 +60,7 @@ def generate_args(*, source_files: SourceFiles, black: Black, check_only: bool)
     if black.config:
         args.extend(["--config", black.config])
     args.extend(black.args)
-    # NB: For some reason, Black's --exclude option only works on recursive invocations, meaning
-    # calling Black on a directory(s) and letting it auto-discover files. However, we don't want
-    # Black to run over everything recursively under the directory of our target, as Black should
-    # only touch files directly specified. We can use `--include` to ensure that Black only
-    # operates on the files we actually care about.
-    args.extend(["--include", "|".join(re.escape(f) for f in source_files.files)])
-    args.extend(PurePath(f).parent.as_posix() for f in source_files.files)
+    args.extend(source_files.files)
     return tuple(args)
 
 

src/python/pants/backend/python/util_rules/pex.py

@@ -7,6 +7,8 @@
 import functools
 import itertools
 import logging
+import os
+import re
 import shlex
 from collections import defaultdict
 from dataclasses import dataclass
@@ -793,12 +795,33 @@ async def create_venv_pex(
     seeded_venv_request = dataclasses.replace(
         pex_request, additional_args=pex_request.additional_args + ("--venv", "--seed")
     )
-    result = await Get(BuildPexResult, PexRequest, seeded_venv_request)
+    venv_pex_result = await Get(BuildPexResult, PexRequest, seeded_venv_request)
     # Pex --seed mode outputs the path of the PEX executable. In the --venv case this is the `pex`
     # script in the venv root directory.
-    venv_dir = PurePath(result.result.stdout.decode().strip()).parent
+    abs_venv_dir = PurePath(venv_pex_result.result.stdout.decode().strip()).parent
+
+    # TODO(John Sirois): Instead of performing this calculation, which uses knowledge of the
+    #  PEX_ROOT layout (It assumes a `venvs` dir), rely on verbose seeding or a pex tool to get
+    #  the relevant PEX_ROOT / venv pex directory combination in a single Process execution so
+    #  we can relpath them.
+
+    pex_root = pex_environment.pex_root
+    if pex_root.is_absolute():
+        if not os.path.commonprefix([pex_root, abs_venv_dir]) == str(pex_root):
+            raise AssertionError(
+                f"The PEX_ROOT is absolute {pex_root} but does not match the venv pex "
+                f"{abs_venv_dir}."
+            )
+        venv_dir = abs_venv_dir.relative_to(pex_root)
+    else:
+        match = re.search(r"/(?P<venv_dir>venvs/.*)$", str(abs_venv_dir))
+        if match is None:
+            raise AssertionError(
+                f"Failed to find PEX_ROOT {pex_root} in venv_pex path {abs_venv_dir}."
+            )
+        venv_dir = pex_root / match["venv_dir"]
 
-    venv_script_writer = VenvScriptWriter(pex=result.create_pex(), venv_dir=venv_dir)
+    venv_script_writer = VenvScriptWriter(pex=venv_pex_result.create_pex(), venv_dir=venv_dir)
     pex = venv_script_writer.exe(bash, pex_environment)
     python = venv_script_writer.python(bash, pex_environment)
     scripts = {
@@ -932,6 +955,7 @@ async def setup_pex_process(request: PexProcess, pex_environment: PexEnvironment
         env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
+        append_only_caches=pex_environment.append_only_caches(),
         timeout_seconds=request.timeout_seconds,
         execution_slot_variable=request.execution_slot_variable,
         cache_scope=request.cache_scope,
@@ -982,7 +1006,9 @@ def __init__(
 
 
 @rule
-async def setup_venv_pex_process(request: VenvPexProcess) -> Process:
+async def setup_venv_pex_process(
+    request: VenvPexProcess, pex_environment: PexEnvironment
+) -> Process:
     venv_pex = request.venv_pex
     argv = (venv_pex.pex.argv0, *request.argv)
     input_digest = (
@@ -998,6 +1024,7 @@ async def setup_venv_pex_process(request: VenvPexProcess) -> Process:
         env=request.extra_env,
         output_files=request.output_files,
         output_directories=request.output_directories,
+        append_only_caches=pex_environment.append_only_caches(),
         timeout_seconds=request.timeout_seconds,
         execution_slot_variable=request.execution_slot_variable,
         cache_scope=request.cache_scope,

src/python/pants/backend/python/util_rules/pex_cli.py

@@ -133,14 +133,11 @@ async def setup_pex_cli_process(
         digests_to_merge.append(request.additional_input_digest)
     input_digest = await Get(Digest, MergeDigests(digests_to_merge))
 
-    pex_root_path = ".cache/pex_root"
     argv = [
         downloaded_pex_bin.exe,
         *cert_args,
         "--python-path",
         create_path_env_var(pex_env.interpreter_search_paths),
-        "--pex-root",
-        pex_root_path,
         # Ensure Pex and its subprocesses create temporary files in the the process execution
         # sandbox. It may make sense to do this generally for Processes, but in the short term we
         # have known use cases where /tmp is too small to hold large wheel downloads Pex is asked to
@@ -173,7 +170,7 @@ async def setup_pex_cli_process(
         env=env,
         output_files=request.output_files,
         output_directories=request.output_directories,
-        append_only_caches={"pex_root": pex_root_path},
+        append_only_caches=pex_env.append_only_caches(),
         level=request.level,
         cache_scope=request.cache_scope,
     )

src/python/pants/backend/python/util_rules/pex_environment.py

@@ -3,7 +3,7 @@
 
 import os
 from dataclasses import dataclass
-from pathlib import Path
+from pathlib import PurePath
 from textwrap import dedent
 from typing import Mapping, Optional, Tuple, cast
 
@@ -105,7 +105,6 @@ class PexEnvironment(EngineAwareReturnType):
     path: Tuple[str, ...]
     interpreter_search_paths: Tuple[str, ...]
     subprocess_environment_dict: FrozenDict[str, str]
-    named_caches_dir: str
     bootstrap_python: Optional[PythonExecutable] = None
 
     def create_argv(
@@ -128,7 +127,7 @@ def environment_dict(self, *, python_configured: bool) -> Mapping[str, str]:
             PATH=create_path_env_var(self.path),
             PEX_INHERIT_PATH="false",
             PEX_IGNORE_RCFILES="true",
-            PEX_ROOT=os.path.join(self.named_caches_dir, "pex_root"),
+            PEX_ROOT=str(self.pex_root),
             **self.subprocess_environment_dict,
         )
         # NB: We only set `PEX_PYTHON_PATH` if the Python interpreter has not already been
@@ -138,6 +137,13 @@ def environment_dict(self, *, python_configured: bool) -> Mapping[str, str]:
             d["PEX_PYTHON_PATH"] = create_path_env_var(self.interpreter_search_paths)
         return d
 
+    @property
+    def pex_root(self) -> PurePath:
+        return PurePath(".cache/pex_root")
+
+    def append_only_caches(self) -> Mapping[str, str]:
+        return {"pex_root": str(self.pex_root)}
+
     def level(self) -> LogLevel:
         return LogLevel.DEBUG if self.bootstrap_python else LogLevel.WARN
 
@@ -221,9 +227,6 @@ def first_python_binary() -> Optional[PythonExecutable]:
         path=pex_runtime_env.path,
         interpreter_search_paths=tuple(python_setup.interpreter_search_paths),
         subprocess_environment_dict=subprocess_env_vars.vars,
-        # TODO: This path normalization is duplicated with `engine_initializer.py`. How can we do
-        #  the normalization only once, via the options system?
-        named_caches_dir=Path(global_options.options.named_caches_dir).resolve().as_posix(),
         bootstrap_python=first_python_binary(),
     )