Run more rules concurrently when building a pex (#20988)

This refactors the setup rules invoked to build a pex to run
concurrently where they can.

This probably doesn't make too much difference in the common case at the
moment, since these seem to currently run very fast (for the single
sample I did, of running `pants --no-local-cache test
src/python/pants/backend/python/util_rules/pex_test.py`), but:

- being concurrent seems better than being unnecessarily sequential
- I'm planning to make `_setup_pex_requirements` start invoking external
processes (`pex3 lock export-subset`) for #15694, which'll make it much
more expensive.

There's a few moving parts across separate individually-sensible
commits. This includes switching to use call-by-name syntax as a bit of
an experiment.

I've labelled this as
https://github.com/pantsbuild/pants/labels/category%3Ainternal, not
https://github.com/pantsbuild/pants/labels/category%3Aperformance,
because it doesn't seem like it has much impact at the moment.

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

@@ -63,8 +63,9 @@
 from pants.engine.fs import EMPTY_DIGEST, AddPrefix, CreateDigest, Digest, FileContent, MergeDigests
 from pants.engine.internals.native_engine import Snapshot
 from pants.engine.internals.selectors import MultiGet
+from pants.engine.intrinsics import add_prefix_request_to_digest
 from pants.engine.process import Process, ProcessCacheScope, ProcessResult
-from pants.engine.rules import Get, collect_rules, rule
+from pants.engine.rules import Get, collect_rules, concurrently, implicitly, rule
 from pants.engine.target import (
     HydratedSources,
     HydrateSourcesRequest,
@@ -418,6 +419,7 @@ class _BuildPexPythonSetup:
     argv: list[str]
 
 
+@rule
 async def _determine_pex_python_and_platforms(request: PexRequest) -> _BuildPexPythonSetup:
     # NB: If `--platform` is specified, this signals that the PEX should not be built locally.
     # `--interpreter-constraint` only makes sense in the context of building locally. These two
@@ -509,6 +511,7 @@ async def get_req_strings(pex_reqs: PexRequirements) -> PexRequirementsInfo:
     return PexRequirementsInfo(tuple(sorted(req_strings)), tuple(sorted(find_links)))
 
 
+@rule
 async def _setup_pex_requirements(
     request: PexRequest, python_setup: PythonSetup
 ) -> _BuildPexRequirementsSetup:
@@ -665,17 +668,42 @@ async def build_pex(
                 ),
             )
 
+    source_dir_name = "source_files"
+
+    pex_python_setup_req = _determine_pex_python_and_platforms(request)
+    requirements_setup_req = _setup_pex_requirements(**implicitly({request: PexRequest}))
+    sources_digest_as_subdir_req = add_prefix_request_to_digest(
+        AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name)
+    )
+    if isinstance(request.requirements, PexRequirements):
+        (
+            pex_python_setup,
+            requirements_setup,
+            sources_digest_as_subdir,
+            req_info,
+        ) = await concurrently(
+            pex_python_setup_req,
+            requirements_setup_req,
+            sources_digest_as_subdir_req,
+            get_req_strings(request.requirements),
+        )
+        req_strings = req_info.req_strings
+    else:
+        pex_python_setup, requirements_setup, sources_digest_as_subdir = await concurrently(
+            pex_python_setup_req,
+            requirements_setup_req,
+            sources_digest_as_subdir_req,
+        )
+        req_strings = ()
+
     argv = [
         "--output-file",
         request.output_filename,
         *request.additional_args,
     ]
 
-    pex_python_setup = await _determine_pex_python_and_platforms(request)
     argv.extend(pex_python_setup.argv)
 
-    source_dir_name = "source_files"
-
     if request.main is not None:
         argv.extend(request.main.iter_pex_args())
         if isinstance(request.main, Executable):
@@ -705,12 +733,8 @@ async def build_pex(
         argv.append("--no-pre-install-wheels")
 
     argv.append(f"--sources-directory={source_dir_name}")
-    sources_digest_as_subdir = await Get(
-        Digest, AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name)
-    )
 
     # Include any additional arguments and input digests required by the requirements.
-    requirements_setup = await _setup_pex_requirements(request, python_setup)
     argv.extend(requirements_setup.argv)
 
     merged_digest = await Get(
@@ -734,18 +758,13 @@ async def build_pex(
     else:
         output_directories = [request.output_filename]
 
-    req_strings = (
-        (await Get(PexRequirementsInfo, PexRequirements, request.requirements)).req_strings
-        if isinstance(request.requirements, PexRequirements)
-        else []
-    )
     result = await Get(
         ProcessResult,
         PexCliProcess(
             subcommand=(),
             extra_args=argv,
             additional_input_digest=merged_digest,
-            description=await _build_pex_description(request, req_strings, python_setup.resolves),
+            description=_build_pex_description(request, req_strings, python_setup.resolves),
             output_files=output_files,
             output_directories=output_directories,
             concurrency_available=requirements_setup.concurrency_available,
@@ -771,7 +790,7 @@ async def build_pex(
     )
 
 
-async def _build_pex_description(
+def _build_pex_description(
     request: PexRequest, req_strings: Sequence[str], resolve_to_lockfile: Mapping[str, str]
 ) -> str:
     if request.description:

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

@@ -826,16 +826,12 @@ def assert_description(
             requirements=requirements,
             description=description,
         )
-        req_strings = (
-            requirements.req_strings_or_addrs if isinstance(requirements, PexRequirements) else []
-        )
-        assert (
-            run_rule_with_mocks(
-                _build_pex_description,
-                rule_args=[request, req_strings, {}],
-            )
-            == expected
-        )
+        req_strings = []
+        if isinstance(requirements, PexRequirements):
+            for s in requirements.req_strings_or_addrs:
+                assert isinstance(s, str)
+                req_strings.append(s)
+        assert _build_pex_description(request, req_strings, {}) == expected
 
     repo_pex = Pex(EMPTY_DIGEST, "repo.pex", None)