diff --git a/src/python/pants/backend/codegen/protobuf/python/rules.py b/src/python/pants/backend/codegen/protobuf/python/rules.py index c2ac2cd5e8cc..f4326c46875d 100644 --- a/src/python/pants/backend/codegen/protobuf/python/rules.py +++ b/src/python/pants/backend/codegen/protobuf/python/rules.py @@ -21,7 +21,6 @@ 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 @@ -61,7 +60,6 @@ async def generate_python_from_protobuf( grpc_python_plugin: GrpcPythonPlugin, python_protobuf_subsystem: PythonProtobufSubsystem, python_protobuf_mypy_plugin: PythonProtobufMypyPlugin, - pex_environment: PexEnvironment, ) -> GeneratedSources: download_protoc_request = Get( DownloadedExternalTool, ExternalToolRequest, protoc.get_request(Platform.current) @@ -195,7 +193,6 @@ 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(), ), ) diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py index a8060dd8f859..58d0e97272b0 100644 --- a/src/python/pants/backend/python/util_rules/pex.py +++ b/src/python/pants/backend/python/util_rules/pex.py @@ -8,8 +8,6 @@ import itertools import json import logging -import os -import re import shlex from collections import defaultdict from dataclasses import dataclass @@ -844,33 +842,12 @@ async def create_venv_pex( seeded_venv_request = dataclasses.replace( pex_request, additional_args=pex_request.additional_args + ("--venv", "--seed") ) - venv_pex_result = await Get(BuildPexResult, PexRequest, seeded_venv_request) + 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. - 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"/(?Pvenvs/.*)$", 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_dir = PurePath(result.result.stdout.decode().strip()).parent - venv_script_writer = VenvScriptWriter(pex=venv_pex_result.create_pex(), venv_dir=venv_dir) + venv_script_writer = VenvScriptWriter(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 = { @@ -891,7 +868,7 @@ async def create_venv_pex( return VenvPex( digest=input_digest, - pex_filename=venv_pex_result.pex_filename, + pex_filename=result.pex_filename, pex=pex.script, python=python.script, bin=FrozenDict((bin_name, venv_script.script) for bin_name, venv_script in scripts.items()), @@ -1005,7 +982,6 @@ 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, @@ -1056,9 +1032,7 @@ def __init__( @rule -async def setup_venv_pex_process( - request: VenvPexProcess, pex_environment: PexEnvironment -) -> Process: +async def setup_venv_pex_process(request: VenvPexProcess) -> Process: venv_pex = request.venv_pex argv = (venv_pex.pex.argv0, *request.argv) input_digest = ( @@ -1074,7 +1048,6 @@ async def setup_venv_pex_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, diff --git a/src/python/pants/backend/python/util_rules/pex_cli.py b/src/python/pants/backend/python/util_rules/pex_cli.py index 2290d66908c2..e5111c208724 100644 --- a/src/python/pants/backend/python/util_rules/pex_cli.py +++ b/src/python/pants/backend/python/util_rules/pex_cli.py @@ -142,11 +142,14 @@ 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 = [ pex_binary.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 @@ -179,7 +182,7 @@ async def setup_pex_cli_process( env=env, output_files=request.output_files, output_directories=request.output_directories, - append_only_caches=pex_env.append_only_caches(), + append_only_caches={"pex_root": pex_root_path}, level=request.level, cache_scope=request.cache_scope, ) diff --git a/src/python/pants/backend/python/util_rules/pex_environment.py b/src/python/pants/backend/python/util_rules/pex_environment.py index a96d72f4cb55..9b514655e824 100644 --- a/src/python/pants/backend/python/util_rules/pex_environment.py +++ b/src/python/pants/backend/python/util_rules/pex_environment.py @@ -3,7 +3,7 @@ import os from dataclasses import dataclass -from pathlib import PurePath +from pathlib import Path from textwrap import dedent from typing import Mapping, Optional, Tuple, cast @@ -106,6 +106,7 @@ 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 +129,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=str(self.pex_root), + PEX_ROOT=os.path.join(self.named_caches_dir, "pex_root"), **self.subprocess_environment_dict, ) # NB: We only set `PEX_PYTHON_PATH` if the Python interpreter has not already been @@ -138,13 +139,6 @@ 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 @@ -233,6 +227,9 @@ def first_python_binary() -> Optional[PythonExecutable]: python_setup.interpreter_search_paths(pex_relevant_environment) ), 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(), )