Repository: pantsbuild/pants

Issue:
Improve hit rates by resolving once per lockfile, and then subsetting per target
Currently the `package` and `run` goals  use a "precise" resolve which is a subset of any lockfile/`constraints.txt` file: this means that each binary is likely to need to re-run the resolve (`package` and `run` also use different resolves from one another, as seen in #10513). This is a different strategy than the one employed for `test`, which obeys the [--python-setup-resolve-all-constraints](https://www.pantsbuild.org/docs/reference-python-setup) flag, and will thus do one repo-wide resolve that is used for all tests.

Each approach has downsides: a per-target resolve is precise, and results in the smallest cache key for the binary, but there are way more resolves than there would be otherwise. The "resolve all constraints" strategy drastically reduces the number of resolves that need to happen (...to "one"), but means that cache keys are larger than they should be, and dependencies lists aren't strictly accurate for testing.

But there is a way to unify these approaches and get both performance and correctness for this case: we should implement a single resolve per lockfile (similar to the `--resolve-all-constraints` approach), but then we should filter down the relevant subset per dependee in a graph aware fashion. This would get us optimal cache hits and exactly one resolve, while also being more correct, and thus a good default.

_(getting the maximum benefit of this idea requires #11165)_


Discussion:
@benjyw , @jsirois : We discussed this one a while back, but the additional usecase mentioned in the description came up today in a conversation with @asherf , so I think we have even more reason to do it.
Resolves are the most expensive thing we do for most repos, especially repos with data science libraries like TensorFlow. I'm not totally following the proposal (I haven't looked closely), but the best of both worlds sounds great.
Spelling this out a bit more, given:
1. a resolve - a set of distributions
2. a target

We should gather the target source dependency closure and from that set of sources determine the 3rdparty set of dependencies - a set of PyPI "project names".  For each project name, find it in the resolved set of distributions (1) or fail. Look at the metadata of the found distribution and collect all distributions it depends on from set 1, recursively or fail. If there are no failures use the collected set of distributions which will be a subset of 1 as the final resolved set of distributions for 2.
> Spelling this out a bit more, given:

Yea, I think that that sounds right. And until we support more than one lockfile, I think that the "identity" of the resolve 
(your 1) is basically just the `constraints.txt` file that is locked in by the [--python-setup-resolve-all-constraints flag](https://www.pantsbuild.org/docs/reference-python-setup) flag.
Hm. I think that I was wrong about this being orthogonal to #11165. Unless #11165 is implemented (to effectively make the `--resolve-all-constraints` feature the default per-resolve), this change will only improve cache hit rates when `--resolve-all-constraints` is enabled (which by definition only has one resolve), but will _not_ reduce the number of resolves without that flag.

Resolving patch:
diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py
--- a/src/python/pants/backend/python/util_rules/pex.py
+++ b/src/python/pants/backend/python/util_rules/pex.py
@@ -325,6 +325,7 @@ class PexRequest(EngineAwareParameter):
     sources: Digest | None
     additional_inputs: Digest | None
     main: MainSpecification | None
+    repository_pex: Pex | None
     additional_args: Tuple[str, ...]
     pex_path: Tuple[Pex, ...]
     apply_requirement_constraints: bool
@@ -341,6 +342,7 @@ def __init__(
         sources: Digest | None = None,
         additional_inputs: Digest | None = None,
         main: MainSpecification | None = None,
+        repository_pex: Pex | None = None,
         additional_args: Iterable[str] = (),
         pex_path: Iterable[Pex] = (),
         apply_requirement_constraints: bool = True,
@@ -364,6 +366,8 @@ def __init__(
             directly in the Pex, but should be present in the environment when building the Pex.
         :param main: The main for the built Pex, equivalent to Pex's `-e` or '-c' flag. If
             left off, the Pex will open up as a REPL.
+        :param repository_pex: An optional PEX to resolve requirements from via the Pex CLI
+            `--pex-repository` option.
         :param additional_args: Any additional Pex flags.
         :param pex_path: Pex files to add to the PEX_PATH.
         :param apply_requirement_constraints: Whether to apply any configured
@@ -379,6 +383,7 @@ def __init__(
         self.sources = sources
         self.additional_inputs = additional_inputs
         self.main = main
+        self.repository_pex = repository_pex
         self.additional_args = tuple(additional_args)
         self.pex_path = tuple(pex_path)
         self.apply_requirement_constraints = apply_requirement_constraints
@@ -561,17 +566,25 @@ async def build_pex(
     argv = [
         "--output-file",
         request.output_filename,
+        *request.additional_args,
+    ]
+
+    if request.repository_pex:
+        argv.extend(["--pex-repository", request.repository_pex.name])
+    else:
         # NB: In setting `--no-pypi`, we rely on the default value of `--python-repos-indexes`
         # including PyPI, which will override `--no-pypi` and result in using PyPI in the default
         # case. Why set `--no-pypi`, then? We need to do this so that
         # `--python-repos-repos=['custom_url']` will only point to that index and not include PyPI.
-        "--no-pypi",
-        *(f"--index={index}" for index in python_repos.indexes),
-        *(f"--repo={repo}" for repo in python_repos.repos),
-        "--resolver-version",
-        "pip-2020-resolver",
-        *request.additional_args,
-    ]
+        argv.extend(
+            [
+                "--no-pypi",
+                *(f"--index={index}" for index in python_repos.indexes),
+                *(f"--repo={repo}" for repo in python_repos.repos),
+                "--resolver-version",
+                "pip-2020-resolver",
+            ]
+        )
 
     python: PythonExecutable | None = None
 
@@ -619,6 +632,9 @@ async def build_pex(
         Digest, AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name)
     )
     additional_inputs_digest = request.additional_inputs or EMPTY_DIGEST
+    repository_pex_digest = (
+        request.repository_pex.digest if request.repository_pex else EMPTY_DIGEST
+    )
 
     merged_digest = await Get(
         Digest,
@@ -627,6 +643,7 @@ async def build_pex(
                 sources_digest_as_subdir,
                 additional_inputs_digest,
                 constraint_file_digest,
+                repository_pex_digest,
                 *(pex.digest for pex in request.pex_path),
             )
         ),
@@ -886,39 +903,31 @@ async def two_step_create_pex(two_step_pex_request: TwoStepPexRequest) -> TwoSte
     request = two_step_pex_request.pex_request
     req_pex_name = "__requirements.pex"
 
-    additional_inputs: Digest | None
+    repository_pex: Pex | None = None
 
     # Create a pex containing just the requirements.
     if request.requirements:
-        requirements_pex_request = 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)}"
+        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)}"
+                ),
             ),
         )
-        requirements_pex = await Get(Pex, PexRequest, requirements_pex_request)
-        additional_inputs = requirements_pex.digest
-        additional_args = (*request.additional_args, f"--requirements-pex={req_pex_name}")
-    else:
-        additional_inputs = None
-        additional_args = request.additional_args
 
     # Now create a full PEX on top of the requirements PEX.
-    full_pex_request = dataclasses.replace(
-        request,
-        requirements=PexRequirements(),
-        additional_inputs=additional_inputs,
-        additional_args=additional_args,
-    )
+    full_pex_request = dataclasses.replace(request, repository_pex=repository_pex)
     full_pex = await Get(Pex, PexRequest, full_pex_request)
     return TwoStepPex(pex=full_pex)
 
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
--- a/src/python/pants/backend/python/util_rules/pex_from_targets.py
+++ b/src/python/pants/backend/python/util_rules/pex_from_targets.py
@@ -19,6 +19,7 @@
 )
 from pants.backend.python.util_rules.pex import (
     MaybeConstraintsFile,
+    Pex,
     PexInterpreterConstraints,
     PexPlatforms,
     PexRequest,
@@ -228,6 +229,7 @@ async def pex_from_targets(
     )
 
     requirements = exact_reqs
+    repository_pex: Pex | None = None
     description = request.description
 
     if constraints_file.path:
@@ -270,8 +272,17 @@ async def pex_from_targets(
                     "file does not cover all requirements."
                 )
             else:
-                requirements = PexRequirements(str(req) for req in constraints_file_reqs)
-                description = description or f"Resolving {python_setup.requirement_constraints}"
+                repository_pex = await Get(
+                    Pex,
+                    PexRequest(
+                        description=f"Resolving {python_setup.requirement_constraints}",
+                        output_filename="repository.pex",
+                        internal_only=request.internal_only,
+                        requirements=PexRequirements(str(req) for req in constraints_file_reqs),
+                        interpreter_constraints=interpreter_constraints,
+                        platforms=request.platforms,
+                    ),
+                )
     elif (
         python_setup.resolve_all_constraints != ResolveAllConstraintsOption.NEVER
         and python_setup.resolve_all_constraints_was_set_explicitly()
@@ -292,6 +303,7 @@ async def pex_from_targets(
         main=request.main,
         sources=merged_input_digest,
         additional_inputs=request.additional_inputs,
+        repository_pex=repository_pex,
         additional_args=request.additional_args,
         description=description,
     )


Test patch:
diff --git a/src/python/pants/backend/python/util_rules/pex_from_targets_test.py b/src/python/pants/backend/python/util_rules/pex_from_targets_test.py
--- a/src/python/pants/backend/python/util_rules/pex_from_targets_test.py
+++ b/src/python/pants/backend/python/util_rules/pex_from_targets_test.py
@@ -1,14 +1,20 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import json
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path, PurePath
 from textwrap import dedent
-from typing import Optional
+from typing import List, Optional, cast
 
 import pytest
+from _pytest.tmpdir import TempPathFactory
 
 from pants.backend.python.target_types import PythonLibrary, PythonRequirementLibrary
 from pants.backend.python.util_rules import pex_from_targets
-from pants.backend.python.util_rules.pex import PexRequest, PexRequirements
+from pants.backend.python.util_rules.pex import Pex, PexRequest, PexRequirements
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
 from pants.build_graph.address import Address
 from pants.engine.internals.scheduler import ExecutionError
@@ -27,7 +33,86 @@ def rule_runner() -> RuleRunner:
     )
 
 
-def test_constraints_validation(rule_runner: RuleRunner) -> None:
+@dataclass(frozen=True)
+class Project:
+    name: str
+    version: str
+
+
+def create_dists(workdir: Path, project: Project, *projects: Project) -> PurePath:
+    project_dirs = []
+    for proj in (project, *projects):
+        project_dir = workdir / "projects" / proj.name
+        project_dir.mkdir(parents=True)
+        project_dirs.append(project_dir)
+
+        (project_dir / "pyproject.toml").write_text(
+            dedent(
+                """\
+                [build-system]
+                requires = ["setuptools==54.1.2", "wheel==0.36.2"]
+                build-backend = "setuptools.build_meta"
+                """
+            )
+        )
+        (project_dir / "setup.cfg").write_text(
+            dedent(
+                f"""\
+                [metadata]
+                name = {proj.name}
+                version = {proj.version}
+                """
+            )
+        )
+
+    pex = workdir / "pex"
+    subprocess.run(
+        args=[sys.executable, "-m", "pex", *project_dirs, "--include-tools", "-o", pex], check=True
+    )
+
+    find_links = workdir / "find-links"
+    subprocess.run(
+        args=[
+            sys.executable,
+            "-m",
+            "pex.tools",
+            pex,
+            "repository",
+            "extract",
+            "--find-links",
+            find_links,
+        ],
+        check=True,
+    )
+    return find_links
+
+
+def requirements(rule_runner: RuleRunner, pex: Pex) -> List[str]:
+    rule_runner.scheduler.write_digest(pex.digest)
+    completed_process = subprocess.run(
+        args=[
+            sys.executable,
+            "-m",
+            "pex.tools",
+            pex.name,
+            "info",
+        ],
+        cwd=rule_runner.build_root,
+        stdout=subprocess.PIPE,
+        check=True,
+    )
+    return cast(List[str], json.loads(completed_process.stdout.decode())["requirements"])
+
+
+def test_constraints_validation(tmp_path_factory: TempPathFactory, rule_runner: RuleRunner) -> None:
+    find_links = create_dists(
+        tmp_path_factory.mktemp("sdists"),
+        Project("Foo-Bar", "1.0.0"),
+        Project("Bar", "5.5.5"),
+        Project("baz", "2.2.2"),
+        Project("QUX", "3.4.5"),
+    )
+
     rule_runner.add_to_build_file(
         "",
         dedent(
@@ -53,16 +138,6 @@ def test_constraints_validation(rule_runner: RuleRunner) -> None:
         """
         ),
     )
-    rule_runner.create_file(
-        "constraints2.txt",
-        dedent(
-            """
-            foo==1.0.0
-            bar==5.5.5
-            qux==3.4.5
-        """
-        ),
-    )
 
     def get_pex_request(
         constraints_file: Optional[str],
@@ -81,28 +156,34 @@ def get_pex_request(
             args.append(f"--python-setup-resolve-all-constraints={resolve_all.value}")
         if constraints_file:
             args.append(f"--python-setup-requirement-constraints={constraints_file}")
-        rule_runner.set_options(args)
+        args.append("--python-repos-indexes=[]")
+        args.append(f"--python-repos-repos={find_links}")
+        rule_runner.set_options(args, env_inherit={"PATH"})
         return rule_runner.request(PexRequest, [request])
 
     pex_req1 = get_pex_request("constraints1.txt", ResolveAllConstraintsOption.NEVER)
     assert pex_req1.requirements == PexRequirements(["foo-bar>=0.1.2", "bar==5.5.5", "baz"])
+    assert pex_req1.repository_pex is None
 
     pex_req1_direct = get_pex_request(
         "constraints1.txt", ResolveAllConstraintsOption.NEVER, direct_deps_only=True
     )
     assert pex_req1_direct.requirements == PexRequirements(["baz"])
+    assert pex_req1_direct.repository_pex is None
 
     pex_req2 = get_pex_request("constraints1.txt", ResolveAllConstraintsOption.ALWAYS)
-    assert pex_req2.requirements == PexRequirements(
-        ["Foo._-BAR==1.0.0", "bar==5.5.5", "baz==2.2.2", "qux==3.4.5"]
+    assert pex_req2.requirements == PexRequirements(["foo-bar>=0.1.2", "bar==5.5.5", "baz"])
+    assert pex_req2.repository_pex is not None
+    repository_pex = pex_req2.repository_pex
+    assert ["Foo._-BAR==1.0.0", "bar==5.5.5", "baz==2.2.2", "qux==3.4.5"] == requirements(
+        rule_runner, repository_pex
     )
 
     pex_req2_direct = get_pex_request(
         "constraints1.txt", ResolveAllConstraintsOption.ALWAYS, direct_deps_only=True
     )
-    assert pex_req2_direct.requirements == PexRequirements(
-        ["Foo._-BAR==1.0.0", "bar==5.5.5", "baz==2.2.2", "qux==3.4.5"]
-    )
+    assert pex_req2_direct.requirements == PexRequirements(["baz"])
+    assert pex_req2_direct.repository_pex == repository_pex
 
     with pytest.raises(ExecutionError) as err:
         get_pex_request(None, ResolveAllConstraintsOption.ALWAYS)