[internal] Remove TwoStepPex abstraction (pantsbuild#12343)

We already use repository PEXes when `[python-setup].resolve_all_constraints` or experimental lockfile are set, thanks to code in `pex_from_targets.py`. Those benefit from using the same repository PEX as the rest of the codebase, e.g. when using MyPy or Pytest. However, we were ignoring this repository pex with `python_awslambda` and `pex_binary` and resolving more granular repository PEXes after having already resolved the global one. Iiuc, those more granular ones are not coming from the global one. This is wasted work. While a `TwoStepPex` could be useful in those two callsites if you are not using constraints/lockfiles, we strongly encourage using lockfiles. Removing `TwoStepPex` allows us to remove complexity in a hairy part of the codebase, especially as we add new complexity for lockfile support. [ci skip-rust] [ci skip-build-wheels]
Eric-Arellano · Jul 14, 2021 · 37401bf · 37401bf
1 parent ee5fe18
commit 37401bf
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 130 deletions.
diff --git a/src/python/pants/backend/awslambda/python/rules.py b/src/python/pants/backend/awslambda/python/rules.py
@@ -14,17 +14,14 @@
 from pants.backend.python.util_rules import pex_from_targets
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import (
+    Pex,
     PexPlatforms,
     PexRequest,
     PexRequirements,
-    TwoStepPex,
     VenvPex,
     VenvPexProcess,
 )
-from pants.backend.python.util_rules.pex_from_targets import (
-    PexFromTargetsRequest,
-    TwoStepPexFromTargetsRequest,
-)
+from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.core.goals.package import (
     BuiltPackage,
     BuiltPackageArtifact,
@@ -75,21 +72,20 @@ async def package_python_awslambda(
         platform += "m"
     if (py_major, py_minor) == (2, 7):
         platform += "u"
-    pex_request = TwoStepPexFromTargetsRequest(
-        PexFromTargetsRequest(
-            addresses=[field_set.address],
-            internal_only=False,
-            main=None,
-            output_filename=output_filename,
-            platforms=PexPlatforms([platform]),
-            additional_args=[
-                # Ensure we can resolve manylinux wheels in addition to any AMI-specific wheels.
-                "--manylinux=manylinux2014",
-                # When we're executing Pex on Linux, allow a local interpreter to be resolved if
-                # available and matching the AMI platform.
-                "--resolve-local-platforms",
-            ],
-        )
+
+    pex_request = PexFromTargetsRequest(
+        addresses=[field_set.address],
+        internal_only=False,
+        main=None,
+        output_filename=output_filename,
+        platforms=PexPlatforms([platform]),
+        additional_args=[
+            # Ensure we can resolve manylinux wheels in addition to any AMI-specific wheels.
+            "--manylinux=manylinux2014",
+            # When we're executing Pex on Linux, allow a local interpreter to be resolved if
+            # available and matching the AMI platform.
+            "--resolve-local-platforms",
+        ],
     )
 
     lambdex_request = PexRequest(
@@ -102,7 +98,7 @@ async def package_python_awslambda(
 
     lambdex_pex, pex_result, handler, transitive_targets = await MultiGet(
         Get(VenvPex, PexRequest, lambdex_request),
-        Get(TwoStepPex, TwoStepPexFromTargetsRequest, pex_request),
+        Get(Pex, PexFromTargetsRequest, pex_request),
         Get(ResolvedPythonAwsHandler, ResolvePythonAwsHandlerRequest(field_set.handler)),
         Get(TransitiveTargets, TransitiveTargetsRequest([field_set.address])),
     )
@@ -129,7 +125,7 @@ async def package_python_awslambda(
         VenvPexProcess(
             lambdex_pex,
             argv=("build", "-e", handler.val, output_filename),
-            input_digest=pex_result.pex.digest,
+            input_digest=pex_result.digest,
             output_files=(output_filename,),
             description=f"Setting up handler in {output_filename}",
         ),

diff --git a/src/python/pants/backend/python/goals/package_pex_binary.py b/src/python/pants/backend/python/goals/package_pex_binary.py
@@ -24,11 +24,8 @@
     ResolvedPexEntryPoint,
     ResolvePexEntryPointRequest,
 )
-from pants.backend.python.util_rules.pex import PexPlatforms, TwoStepPex
-from pants.backend.python.util_rules.pex_from_targets import (
-    PexFromTargetsRequest,
-    TwoStepPexFromTargetsRequest,
-)
+from pants.backend.python.util_rules.pex import Pex, PexPlatforms
+from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.core.goals.package import (
     BuiltPackage,
     BuiltPackageArtifact,
@@ -126,22 +123,20 @@ async def package_pex_binary(
         )
 
     output_filename = field_set.output_path.value_or_default(field_set.address, file_ending="pex")
-    two_step_pex = await Get(
-        TwoStepPex,
-        TwoStepPexFromTargetsRequest(
-            PexFromTargetsRequest(
-                addresses=[field_set.address],
-                internal_only=False,
-                # TODO(John Sirois): Support ConsoleScript in PexBinary targets:
-                #  https://github.com/pantsbuild/pants/issues/11619
-                main=resolved_entry_point.val,
-                platforms=PexPlatforms.create_from_platforms_field(field_set.platforms),
-                output_filename=output_filename,
-                additional_args=field_set.generate_additional_args(pex_binary_defaults),
-            )
+    pex = await Get(
+        Pex,
+        PexFromTargetsRequest(
+            addresses=[field_set.address],
+            internal_only=False,
+            # TODO(John Sirois): Support ConsoleScript in PexBinary targets:
+            #  https://github.com/pantsbuild/pants/issues/11619
+            main=resolved_entry_point.val,
+            platforms=PexPlatforms.create_from_platforms_field(field_set.platforms),
+            output_filename=output_filename,
+            additional_args=field_set.generate_additional_args(pex_binary_defaults),
         ),
     )
-    return BuiltPackage(two_step_pex.pex.digest, (BuiltPackageArtifact(output_filename),))
+    return BuiltPackage(pex.digest, (BuiltPackageArtifact(output_filename),))
 
 
 def rules():

diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py
@@ -172,21 +172,6 @@ def debug_hint(self) -> str:
         return self.output_filename
 
 
-@dataclass(frozen=True)
-class TwoStepPexRequest:
-    """A request to create a PEX in two steps.
-
-    First we create a requirements-only pex. Then we create the full pex on top of that
-    requirements pex, instead of having the full pex directly resolve its requirements.
-
-    This allows us to re-use the requirements-only pex when no requirements have changed (which is
-    the overwhelmingly common case), thus avoiding spurious re-resolves of the same requirements
-    over and over again.
-    """
-
-    pex_request: PexRequest
-
-
 @dataclass(frozen=True)
 class Pex:
     """Wrapper for a digest containing a pex file created with some filename."""
@@ -196,17 +181,6 @@ class Pex:
     python: PythonExecutable | None
 
 
-@dataclass(frozen=True)
-class TwoStepPex:
-    """The result of creating a PEX in two steps.
-
-    TODO(9320): A workaround for https://github.com/pantsbuild/pants/issues/9320. Really we
-      just want the rules to directly return a Pex.
-    """
-
-    pex: Pex
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -630,41 +604,6 @@ async def create_venv_pex(
     )
 
 
-@rule(level=LogLevel.DEBUG)
-async def two_step_create_pex(two_step_pex_request: TwoStepPexRequest) -> TwoStepPex:
-    """Create a PEX in two steps: a requirements-only PEX and then a full PEX from it."""
-    request = two_step_pex_request.pex_request
-    req_pex_name = "__requirements.pex"
-
-    repository_pex: Pex | None = None
-
-    # Create a pex containing just the requirements.
-    if request.requirements:
-        repository_pex = await Get(
-            Pex,
-            PexRequest(
-                output_filename=req_pex_name,
-                internal_only=request.internal_only,
-                requirements=request.requirements,
-                interpreter_constraints=request.interpreter_constraints,
-                platforms=request.platforms,
-                # TODO: Do we need to pass all the additional args to the requirements pex creation?
-                #  Some of them may affect resolution behavior, but others may be irrelevant.
-                #  For now we err on the side of caution.
-                additional_args=request.additional_args,
-                description=(
-                    f"Resolving {pluralize(len(request.requirements), 'requirement')}: "
-                    f"{', '.join(request.requirements)}"
-                ),
-            ),
-        )
-
-    # Now create a full PEX on top of the requirements PEX.
-    full_pex_request = dataclasses.replace(request, repository_pex=repository_pex)
-    full_pex = await Get(Pex, PexRequest, full_pex_request)
-    return TwoStepPex(pex=full_pex)
-
-
 @frozen_after_init
 @dataclass(unsafe_hash=True)
 class PexProcess:

diff --git a/src/python/pants/backend/python/util_rules/pex_from_targets.py b/src/python/pants/backend/python/util_rules/pex_from_targets.py
@@ -18,13 +18,7 @@
     parse_requirements_file,
 )
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.backend.python.util_rules.pex import (
-    Pex,
-    PexPlatforms,
-    PexRequest,
-    PexRequirements,
-    TwoStepPexRequest,
-)
+from pants.backend.python.util_rules.pex import Pex, PexPlatforms, PexRequest, PexRequirements
 from pants.backend.python.util_rules.pex import rules as pex_rules
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFilesRequest,
@@ -163,21 +157,6 @@ def for_requirements(
         )
 
 
-@dataclass(frozen=True)
-class TwoStepPexFromTargetsRequest:
-    """Request to create a PEX from the closure of a set of targets, in two steps.
-
-    First we create a requirements-only pex. Then we create the full pex on top of that
-    requirements pex, instead of having the full pex directly resolve its requirements.
-
-    This allows us to re-use the requirements-only pex when no requirements have changed (which is
-    the overwhelmingly common case), thus avoiding spurious re-resolves of the same requirements
-    over and over again.
-    """
-
-    pex_from_targets_request: PexFromTargetsRequest
-
-
 @rule(level=LogLevel.DEBUG)
 async def pex_from_targets(request: PexFromTargetsRequest, python_setup: PythonSetup) -> PexRequest:
     if request.direct_deps_only:
@@ -327,11 +306,5 @@ async def pex_from_targets(request: PexFromTargetsRequest, python_setup: PythonS
     )
 
 
-@rule
-async def two_step_pex_from_targets(req: TwoStepPexFromTargetsRequest) -> TwoStepPexRequest:
-    pex_request = await Get(PexRequest, PexFromTargetsRequest, req.pex_from_targets_request)
-    return TwoStepPexRequest(pex_request=pex_request)
-
-
 def rules():
     return (*collect_rules(), *pex_rules(), *python_sources_rules())