Skip to content

Commit

Permalink
Refactor build_pex rule with @rule_helper (#16465)
Browse files Browse the repository at this point in the history 
This allows us to leverage early returns & inlining variables. It also allows creating comprehensive unit tests.

[ci skip-rust]
  • Loading branch information
@Eric-Arellano
Eric-Arellano authored Aug 10, 2022
1 parent 7765408 commit d5df1b1
Show file tree
Hide file tree
Showing 2 changed files with 334 additions and 99 deletions.
244 changes: 146 additions & 98 deletions src/python/pants/backend/python/util_rules/pex.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@
from pants.engine.internals.selectors import MultiGet
from pants.engine.platform import Platform
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, rule, rule_helper
from pants.engine.target import HydratedSources, HydrateSourcesRequest, SourcesField, Targets
from pants.util.frozendict import FrozenDict
from pants.util.logging import LogLevel
Expand Down Expand Up @@ -354,35 +354,30 @@ def create_pex(self) -> Pex:
return Pex(digest=self.digest, name=self.pex_filename, python=self.python)


@rule(level=LogLevel.DEBUG)
async def build_pex(
request: PexRequest,
python_setup: PythonSetup,
python_repos: PythonRepos,
platform: Platform,
pex_runtime_env: PexRuntimeEnvironment,
) -> BuildPexResult:
"""Returns a PEX with the given settings."""
argv = [
"--output-file",
request.output_filename,
"--no-emit-warnings",
*python_setup.manylinux_pex_args,
*request.additional_args,
]
@dataclass
class _BuildPexPythonSetup:
python: PythonExecutable | None
argv: list[str]

python: PythonExecutable | None = None

@rule_helper
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
# flags are mutually exclusive. See https://github.com/pantsbuild/pex/issues/957.
if request.platforms or request.complete_platforms:
# Note that this means that this is not an internal-only pex.
# TODO(#9560): consider validating that these platforms are valid with the interpreter
# constraints.
argv.extend(request.platforms.generate_pex_arg_list())
argv.extend(request.complete_platforms.generate_pex_arg_list())
elif request.python:
return _BuildPexPythonSetup(
None,
[
*request.platforms.generate_pex_arg_list(),
*request.complete_platforms.generate_pex_arg_list(),
],
)

if request.python:
python = request.python
elif request.internal_only:
# NB: If it's an internal_only PEX, we do our own lookup of the interpreter based on the
Expand All @@ -394,30 +389,25 @@ async def build_pex(
else:
# `--interpreter-constraint` options are mutually exclusive with the `--python` option,
# so we only specify them if we have not already located a concrete Python.
argv.extend(request.interpreter_constraints.generate_pex_arg_list())
return _BuildPexPythonSetup(None, request.interpreter_constraints.generate_pex_arg_list())

if python:
argv.extend(["--python", python.path])
return _BuildPexPythonSetup(python, ["--python", python.path])

if request.main is not None:
argv.extend(request.main.iter_pex_args())

# TODO(John Sirois): Right now any request requirements will shadow corresponding pex path
# requirements, which could lead to problems. Support shading python binaries.
# See: https://github.com/pantsbuild/pants/issues/9206
if request.pex_path:
argv.extend(["--pex-path", ":".join(pex.name for pex in request.pex_path)])
@dataclass
class _BuildPexRequirementsSetup:
digests: list[Digest]
argv: list[str]
concurrency_available: int

source_dir_name = "source_files"
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_digests = []
pex_lock_resolver_args = [*python_repos.pex_args]
@rule_helper
async def _setup_pex_requirements(
request: PexRequest, python_repos: PythonRepos, python_setup: PythonSetup
) -> _BuildPexRequirementsSetup:
pex_lock_resolver_args = list(python_repos.pex_args)
pip_resolver_args = [*python_repos.pex_args, "--resolver-version", "pip-2020-resolver"]

if isinstance(request.requirements, EntireLockfile):
lockfile, resolve_config = await MultiGet(
Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile)),
Expand All @@ -426,15 +416,13 @@ async def build_pex(
ResolvePexConfigRequest(request.requirements.lockfile.resolve_name),
),
)
concurrency_available = lockfile.requirement_estimate
requirements_digests.append(lockfile.lockfile_digest)
if lockfile.is_pex_native:
argv.extend(["--lock", lockfile.lockfile_path])
argv.extend(pex_lock_resolver_args)
else:
argv = (
["--lock", lockfile.lockfile_path, *pex_lock_resolver_args]
if lockfile.is_pex_native
else
# We use pip to resolve a requirements.txt pseudo-lockfile, possibly with hashes.
argv.extend(["--requirement", lockfile.lockfile_path, "--no-transitive"])
argv.extend(pip_resolver_args)
["--requirement", lockfile.lockfile_path, "--no-transitive", *pip_resolver_args]
)
if lockfile.metadata and request.requirements.complete_req_strings:
validate_metadata(
lockfile.metadata,
Expand All @@ -444,54 +432,111 @@ async def build_pex(
python_setup,
constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash,
)
else:
# TODO: This is not the best heuristic for available concurrency, since the
# requirements almost certainly have transitive deps which also need building, but it
# is better than using something hardcoded.
concurrency_available = len(request.requirements.req_strings)
argv.extend(request.requirements.req_strings)

if isinstance(request.requirements.from_superset, Pex):
repository_pex = request.requirements.from_superset
argv.extend(["--pex-repository", repository_pex.name])
requirements_digests.append(repository_pex.digest)
elif isinstance(request.requirements.from_superset, LoadedLockfile):
loaded_lockfile = request.requirements.from_superset
resolve_config = await Get(
ResolvePexConfig,
ResolvePexConfigRequest(loaded_lockfile.original_lockfile.resolve_name),

return _BuildPexRequirementsSetup(
[lockfile.lockfile_digest], argv, lockfile.requirement_estimate
)

# TODO: This is not the best heuristic for available concurrency, since the
# requirements almost certainly have transitive deps which also need building, but it
# is better than using something hardcoded.
concurrency_available = len(request.requirements.req_strings)

if isinstance(request.requirements.from_superset, Pex):
repository_pex = request.requirements.from_superset
return _BuildPexRequirementsSetup(
[repository_pex.digest],
[*request.requirements.req_strings, "--pex-repository", repository_pex.name],
concurrency_available,
)

if isinstance(request.requirements.from_superset, LoadedLockfile):
loaded_lockfile = request.requirements.from_superset
# NB: This is also validated in the constructor.
assert loaded_lockfile.is_pex_native
resolve_config = await Get(
ResolvePexConfig,
ResolvePexConfigRequest(loaded_lockfile.original_lockfile.resolve_name),
)
if not request.requirements.req_strings:
return _BuildPexRequirementsSetup([], [], concurrency_available)

if loaded_lockfile.metadata:
validate_metadata(
loaded_lockfile.metadata,
request.interpreter_constraints,
loaded_lockfile.original_lockfile,
request.requirements.req_strings,
python_setup,
constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash,
)
# NB: This is also validated in the constructor.
assert loaded_lockfile.is_pex_native
if request.requirements.req_strings:
requirements_digests.append(loaded_lockfile.lockfile_digest)
argv.extend(["--lock", loaded_lockfile.lockfile_path])
argv.extend(pex_lock_resolver_args)

if loaded_lockfile.metadata:
validate_metadata(
loaded_lockfile.metadata,
request.interpreter_constraints,
loaded_lockfile.original_lockfile,
request.requirements.req_strings,
python_setup,
constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash,
)
else:
assert request.requirements.from_superset is None

# We use pip to perform a normal resolve.
argv.extend(pip_resolver_args)
if request.requirements.constraints_strings:
constraints_file = "__constraints.txt"
constraints_content = "\n".join(request.requirements.constraints_strings)
requirements_digests.append(
await Get(
Digest,
CreateDigest([FileContent(constraints_file, constraints_content.encode())]),
)
)
argv.extend(["--constraints", constraints_file])

return _BuildPexRequirementsSetup(
[loaded_lockfile.lockfile_digest],
[
*request.requirements.req_strings,
"--lock",
loaded_lockfile.lockfile_path,
*pex_lock_resolver_args,
],
concurrency_available,
)

# We use pip to perform a normal resolve.
assert request.requirements.from_superset is None
digests = []
argv = [*request.requirements.req_strings, *pip_resolver_args]
if request.requirements.constraints_strings:
constraints_file = "__constraints.txt"
constraints_content = "\n".join(request.requirements.constraints_strings)
digests.append(
await Get(
Digest,
CreateDigest([FileContent(constraints_file, constraints_content.encode())]),
)
)
argv.extend(["--constraints", constraints_file])
return _BuildPexRequirementsSetup(digests, argv, concurrency_available=concurrency_available)


@rule(level=LogLevel.DEBUG)
async def build_pex(
request: PexRequest,
python_setup: PythonSetup,
python_repos: PythonRepos,
platform: Platform,
pex_runtime_env: PexRuntimeEnvironment,
) -> BuildPexResult:
"""Returns a PEX with the given settings."""
argv = [
"--output-file",
request.output_filename,
"--no-emit-warnings",
*python_setup.manylinux_pex_args,
*request.additional_args,
]

pex_python_setup = await _determine_pex_python_and_platforms(request)
argv.extend(pex_python_setup.argv)

if request.main is not None:
argv.extend(request.main.iter_pex_args())

# TODO(John Sirois): Right now any request requirements will shadow corresponding pex path
# requirements, which could lead to problems. Support shading python binaries.
# See: https://github.com/pantsbuild/pants/issues/9206
if request.pex_path:
argv.extend(["--pex-path", ":".join(pex.name for pex in request.pex_path)])

source_dir_name = "source_files"
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_repos, python_setup)
argv.extend(requirements_setup.argv)

merged_digest = await Get(
Digest,
Expand All @@ -500,7 +545,7 @@ async def build_pex(
request.complete_platforms.digest,
sources_digest_as_subdir,
request.additional_inputs,
*requirements_digests,
*requirements_setup.digests,
*(pex.digest for pex in request.pex_path),
)
),
Expand All @@ -517,14 +562,14 @@ async def build_pex(
process = await Get(
Process,
PexCliProcess(
python=python,
python=pex_python_setup.python,
subcommand=(),
extra_args=argv,
additional_input_digest=merged_digest,
description=_build_pex_description(request),
output_files=output_files,
output_directories=output_directories,
concurrency_available=concurrency_available,
concurrency_available=requirements_setup.concurrency_available,
),
)

Expand All @@ -549,7 +594,10 @@ async def build_pex(
)

return BuildPexResult(
result=result, pex_filename=request.output_filename, digest=digest, python=python
result=result,
pex_filename=request.output_filename,
digest=digest,
python=pex_python_setup.python,
)


Expand Down
Loading

0 comments on commit d5df1b1

Please sign in to comment.