Support complete_platforms for python_awslambda. (#14532)

Instead of (or in addition to) specifying a `runtime` for a
`python_awslambda`, you can now specify `complete_platforms`. This
allows creating a lambdex when an abbreviated platform string does not
provide enough information to evaluate environment markers during the
requirement resolution phase of building the lambdex.

src/python/pants/backend/awslambda/python/rules.py

@@ -1,6 +1,8 @@
 # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from __future__ import annotations
+
 import logging
 from dataclasses import dataclass
 
@@ -12,8 +14,10 @@
     ResolvePythonAwsHandlerRequest,
 )
 from pants.backend.python.subsystems.lambdex import Lambdex
+from pants.backend.python.target_types import PexCompletePlatformsField
 from pants.backend.python.util_rules import pex_from_targets
 from pants.backend.python.util_rules.pex import (
+    CompletePlatforms,
     Pex,
     PexPlatforms,
     PexRequest,
@@ -45,11 +49,12 @@
 
 @dataclass(frozen=True)
 class PythonAwsLambdaFieldSet(PackageFieldSet):
-    required_fields = (PythonAwsLambdaHandlerField, PythonAwsLambdaRuntime)
+    required_fields = (PythonAwsLambdaHandlerField,)
 
     handler: PythonAwsLambdaHandlerField
     include_requirements: PythonAwsLambdaIncludeRequirements
     runtime: PythonAwsLambdaRuntime
+    complete_platforms: PexCompletePlatformsField
     output_path: OutputPathField
 
 
@@ -78,13 +83,17 @@ async def package_python_awslambda(
     # We hardcode the platform value to the appropriate one for each AWS Lambda runtime.
     # (Running the "hello world" lambda in the example code will report the platform, and can be
     # used to verify correctness of these platform strings.)
-    py_major, py_minor = field_set.runtime.to_interpreter_version()
-    platform_str = f"linux_x86_64-cp-{py_major}{py_minor}-cp{py_major}{py_minor}"
-    # set pymalloc ABI flag - this was removed in python 3.8 https://bugs.python.org/issue36707
-    if py_major <= 3 and py_minor < 8:
-        platform_str += "m"
-    if (py_major, py_minor) == (2, 7):
-        platform_str += "u"
+    pex_platforms = []
+    interpreter_version = field_set.runtime.to_interpreter_version()
+    if interpreter_version:
+        py_major, py_minor = interpreter_version
+        platform_str = f"linux_x86_64-cp-{py_major}{py_minor}-cp{py_major}{py_minor}"
+        # set pymalloc ABI flag - this was removed in python 3.8 https://bugs.python.org/issue36707
+        if py_major <= 3 and py_minor < 8:
+            platform_str += "m"
+        if (py_major, py_minor) == (2, 7):
+            platform_str += "u"
+        pex_platforms.append(platform_str)
 
     additional_pex_args = (
         # Ensure we can resolve manylinux wheels in addition to any AMI-specific wheels.
@@ -93,12 +102,18 @@ async def package_python_awslambda(
         # available and matching the AMI platform.
         "--resolve-local-platforms",
     )
+
+    complete_platforms = await Get(
+        CompletePlatforms, PexCompletePlatformsField, field_set.complete_platforms
+    )
+
     pex_request = PexFromTargetsRequest(
         addresses=[field_set.address],
         internal_only=False,
         include_requirements=field_set.include_requirements.value,
         output_filename=output_filename,
-        platforms=PexPlatforms([platform_str]),
+        platforms=PexPlatforms(pex_platforms),
+        complete_platforms=complete_platforms,
         additional_args=additional_pex_args,
         additional_lockfile_args=additional_pex_args,
     )
@@ -145,13 +160,20 @@ async def package_python_awslambda(
             description=f"Setting up handler in {output_filename}",
         ),
     )
+
+    extra_log_data: list[tuple[str, str]] = []
+    if field_set.runtime.value:
+        extra_log_data.append(("Runtime", field_set.runtime.value))
+    extra_log_data.extend(("Complete platform", path) for path in complete_platforms)
+    # The AWS-facing handler function is always lambdex_handler.handler, which is the
+    # wrapper injected by lambdex that manages invocation of the actual handler.
+    extra_log_data.append(("Handler", "lambdex_handler.handler"))
+    first_column_width = 4 + max(len(header) for header, _ in extra_log_data)
+
     artifact = BuiltPackageArtifact(
         output_filename,
-        extra_log_lines=(
-            f"    Runtime: {field_set.runtime.value}",
-            # The AWS-facing handler function is always lambdex_handler.handler, which is the
-            # wrapper injected by lambdex that manages invocation of the actual handler.
-            "    Handler: lambdex_handler.handler",
+        extra_log_lines=tuple(
+            f"{header.rjust(first_column_width, ' ')}: {data}" for header, data in extra_log_data
         ),
     )
     return BuiltPackage(digest=result.output_digest, artifacts=(artifact,))

src/python/pants/backend/awslambda/python/rules_test.py

@@ -3,6 +3,8 @@
 
 from __future__ import annotations
 
+import os
+import subprocess
 import sys
 from io import BytesIO
 from textwrap import dedent
@@ -14,12 +16,23 @@
 from pants.backend.awslambda.python.rules import rules as awslambda_python_rules
 from pants.backend.awslambda.python.target_types import PythonAWSLambda
 from pants.backend.awslambda.python.target_types import rules as target_rules
+from pants.backend.python.goals import package_pex_binary
+from pants.backend.python.goals.package_pex_binary import PexBinaryFieldSet
 from pants.backend.python.subsystems.lambdex import Lambdex
 from pants.backend.python.subsystems.lambdex import rules as awslambda_python_subsystem_rules
-from pants.backend.python.target_types import PythonRequirementTarget, PythonSourcesGeneratorTarget
+from pants.backend.python.target_types import (
+    PexBinary,
+    PythonRequirementTarget,
+    PythonSourcesGeneratorTarget,
+)
 from pants.backend.python.target_types_rules import rules as python_target_types_rules
 from pants.core.goals.package import BuiltPackage
-from pants.core.target_types import FilesGeneratorTarget, RelocatedFiles, ResourcesGeneratorTarget
+from pants.core.target_types import (
+    FilesGeneratorTarget,
+    FileTarget,
+    RelocatedFiles,
+    ResourcesGeneratorTarget,
+)
 from pants.core.target_types import rules as core_target_types_rules
 from pants.engine.addresses import Address
 from pants.engine.fs import DigestContents
@@ -29,53 +42,85 @@
 
 @pytest.fixture
 def rule_runner() -> RuleRunner:
-    return RuleRunner(
+    rule_runner = RuleRunner(
         rules=[
             *awslambda_python_rules(),
             *awslambda_python_subsystem_rules(),
-            *target_rules(),
-            *python_target_types_rules(),
             *core_target_types_rules(),
+            *package_pex_binary.rules(),
+            *python_target_types_rules(),
+            *target_rules(),
             QueryRule(BuiltPackage, (PythonAwsLambdaFieldSet,)),
         ],
         target_types=[
+            FileTarget,
+            FilesGeneratorTarget,
+            PexBinary,
             PythonAWSLambda,
             PythonRequirementTarget,
+            PythonRequirementTarget,
             PythonSourcesGeneratorTarget,
-            FilesGeneratorTarget,
             RelocatedFiles,
             ResourcesGeneratorTarget,
         ],
     )
+    rule_runner.set_options([], env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+    return rule_runner
 
 
 def create_python_awslambda(
-    rule_runner: RuleRunner, addr: Address, *, extra_args: list[str] | None = None
+    rule_runner: RuleRunner,
+    addr: Address,
+    *,
+    expected_extra_log_lines: tuple[str, ...],
+    extra_args: list[str] | None = None,
 ) -> tuple[str, bytes]:
     rule_runner.set_options(
         ["--source-root-patterns=src/python", *(extra_args or ())],
         env_inherit={"PATH", "PYENV_ROOT", "HOME"},
     )
     target = rule_runner.get_target(addr)
     built_asset = rule_runner.request(BuiltPackage, [PythonAwsLambdaFieldSet.create(target)])
-    assert (
-        "    Runtime: python3.7",
-        "    Handler: lambdex_handler.handler",
-    ) == built_asset.artifacts[0].extra_log_lines
+    assert expected_extra_log_lines == built_asset.artifacts[0].extra_log_lines
     digest_contents = rule_runner.request(DigestContents, [built_asset.digest])
     assert len(digest_contents) == 1
     relpath = built_asset.artifacts[0].relpath
     assert relpath is not None
     return relpath, digest_contents[0].content
 
 
+@pytest.fixture
+def complete_platform(rule_runner: RuleRunner) -> bytes:
+    rule_runner.write_files(
+        {
+            "pex_exe/BUILD": dedent(
+                """\
+                python_requirement(name="req", requirements=["pex==2.1.66"])
+                pex_binary(dependencies=[":req"], script="pex")
+                """
+            ),
+        }
+    )
+    result = rule_runner.request(
+        BuiltPackage, [PexBinaryFieldSet.create(rule_runner.get_target(Address("pex_exe")))]
+    )
+    rule_runner.write_digest(result.digest)
+    pex_executable = os.path.join(rule_runner.build_root, "pex_exe/pex_exe.pex")
+    return subprocess.run(
+        args=[pex_executable, "interpreter", "inspect", "-mt"],
+        env=dict(PEX_MODULE="pex.cli", **os.environ),
+        check=True,
+        stdout=subprocess.PIPE,
+    ).stdout
+
+
 @pytest.mark.platform_specific_behavior
 @pytest.mark.parametrize(
     "major_minor_interpreter",
     all_major_minor_python_versions(Lambdex.default_interpreter_constraints),
 )
 def test_create_hello_world_lambda(
-    rule_runner: RuleRunner, major_minor_interpreter: str, caplog
+    rule_runner: RuleRunner, major_minor_interpreter: str, complete_platform: str, caplog
 ) -> None:
     rule_runner.write_files(
         {
@@ -87,16 +132,19 @@ def handler(event, context):
                     print('Hello, World!')
                 """
             ),
+            "src/python/foo/bar/platform.json": complete_platform,
             "src/python/foo/bar/BUILD": dedent(
                 """
                 python_requirement(name="mureq", requirements=["mureq==0.2"])
                 python_sources(name='lib')
 
+                file(name="platform", source="platform.json")
                 python_awslambda(
                     name='lambda',
                     dependencies=[':lib'],
                     handler='foo.bar.hello_world:handler',
                     runtime='python3.7',
+                    complete_platforms=[':platform'],
                 )
                 python_awslambda(
                     name='slimlambda',
@@ -112,6 +160,11 @@ def handler(event, context):
     zip_file_relpath, content = create_python_awslambda(
         rule_runner,
         Address("src/python/foo/bar", target_name="lambda"),
+        expected_extra_log_lines=(
+            "              Runtime: python3.7",
+            "    Complete platform: src/python/foo/bar/platform.json",
+            "              Handler: lambdex_handler.handler",
+        ),
         extra_args=[f"--lambdex-interpreter-constraints=['=={major_minor_interpreter}.*']"],
     )
     assert "src.python.foo.bar/lambda.zip" == zip_file_relpath
@@ -127,6 +180,10 @@ def handler(event, context):
     zip_file_relpath, content = create_python_awslambda(
         rule_runner,
         Address("src/python/foo/bar", target_name="slimlambda"),
+        expected_extra_log_lines=(
+            "    Runtime: python3.7",
+            "    Handler: lambdex_handler.handler",
+        ),
         extra_args=[f"--lambdex-interpreter-constraints=['=={major_minor_interpreter}.*']"],
     )
     assert "src.python.foo.bar/slimlambda.zip" == zip_file_relpath
@@ -182,7 +239,12 @@ def handler(event, context):
 
     assert not caplog.records
     zip_file_relpath, _ = create_python_awslambda(
-        rule_runner, Address("src/py/project", target_name="lambda")
+        rule_runner,
+        Address("src/py/project", target_name="lambda"),
+        expected_extra_log_lines=(
+            "    Runtime: python3.7",
+            "    Handler: lambdex_handler.handler",
+        ),
     )
     assert caplog.records
     assert "src.py.project/lambda.zip" == zip_file_relpath

src/python/pants/backend/awslambda/python/target_types.py

@@ -12,7 +12,7 @@
 )
 from pants.backend.python.dependency_inference.rules import PythonInferSubsystem, import_rules
 from pants.backend.python.subsystems.setup import PythonSetup
-from pants.backend.python.target_types import PythonResolveField
+from pants.backend.python.target_types import PexCompletePlatformsField, PythonResolveField
 from pants.core.goals.package import OutputPathField
 from pants.engine.addresses import Address
 from pants.engine.fs import GlobMatchErrorBehavior, PathGlobs, Paths
@@ -27,6 +27,7 @@
     InjectDependenciesRequest,
     InjectedDependencies,
     InvalidFieldException,
+    InvalidTargetException,
     SecondaryOwnerMixin,
     StringField,
     Target,
@@ -191,25 +192,28 @@ class PythonAwsLambdaRuntime(StringField):
     PYTHON_RUNTIME_REGEX = r"python(?P<major>\d)\.(?P<minor>\d+)"
 
     alias = "runtime"
-    required = True
-    value: str
+    default = None
     help = (
         "The identifier of the AWS Lambda runtime to target (pythonX.Y). See "
         "https://docs.aws.amazon.com/lambda/latest/dg/lambda-python.html."
     )
 
     @classmethod
-    def compute_value(cls, raw_value: Optional[str], address: Address) -> str:
-        value = cast(str, super().compute_value(raw_value, address))
+    def compute_value(cls, raw_value: Optional[str], address: Address) -> Optional[str]:
+        value = super().compute_value(raw_value, address)
+        if value is None:
+            return None
         if not re.match(cls.PYTHON_RUNTIME_REGEX, value):
             raise InvalidFieldException(
                 f"The `{cls.alias}` field in target at {address} must be of the form pythonX.Y, "
                 f"but was {value}."
             )
         return value
 
-    def to_interpreter_version(self) -> Tuple[int, int]:
+    def to_interpreter_version(self) -> Optional[Tuple[int, int]]:
         """Returns the Python version implied by the runtime, as (major, minor)."""
+        if self.value is None:
+            return None
         mo = cast(Match, re.match(self.PYTHON_RUNTIME_REGEX, self.value))
         return int(mo.group("major")), int(mo.group("minor"))
 
@@ -223,13 +227,22 @@ class PythonAWSLambda(Target):
         PythonAwsLambdaHandlerField,
         PythonAwsLambdaIncludeRequirements,
         PythonAwsLambdaRuntime,
+        PexCompletePlatformsField,
         PythonResolveField,
     )
     help = (
         "A self-contained Python function suitable for uploading to AWS Lambda.\n\n"
         f"See {doc_url('awslambda-python')}."
     )
 
+    def validate(self) -> None:
+        if self[PythonAwsLambdaRuntime].value is None and not self[PexCompletePlatformsField].value:
+            raise InvalidTargetException(
+                f"The `{self.alias}` target {self.address} must specify either a "
+                f"`{self[PythonAwsLambdaRuntime].alias}` or "
+                f"`{self[PexCompletePlatformsField].alias}` or both."
+            )
+
 
 def rules():
     return (