-
-
Notifications
You must be signed in to change notification settings - Fork 636
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor venv export #17098
Merged
Merged
Refactor venv export #17098
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit
Hold shift + click to select a range
ac302a4
Use a VenvPex when exporting python venvs.
benjyw 86da7cb
Use a regular packed Pex instead.
benjyw a91b4b3
Merge branch 'main' into refactor_venv_export
benjyw 8cb64f4
linr
benjyw c021f39
comment
benjyw bb286f3
use already-computed interpreter
benjyw File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -12,9 +12,8 @@ | |
from pants.backend.python.subsystems.setup import PythonSetup | ||
from pants.backend.python.target_types import PythonResolveField | ||
from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints | ||
from pants.backend.python.util_rules.pex import Pex, PexRequest | ||
from pants.backend.python.util_rules.pex import Pex, PexProcess, PexRequest | ||
from pants.backend.python.util_rules.pex_cli import PexPEX | ||
from pants.backend.python.util_rules.pex_environment import PythonExecutable | ||
from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest | ||
from pants.core.goals.export import ( | ||
ExportError, | ||
|
@@ -28,8 +27,8 @@ | |
from pants.engine.environment import EnvironmentName | ||
from pants.engine.internals.native_engine import AddPrefix, Digest, MergeDigests | ||
from pants.engine.internals.selectors import Get, MultiGet | ||
from pants.engine.process import Process, ProcessResult | ||
from pants.engine.rules import collect_rules, rule | ||
from pants.engine.process import ProcessResult | ||
from pants.engine.rules import collect_rules, rule, rule_helper | ||
from pants.engine.target import Target | ||
from pants.engine.unions import UnionMembership, UnionRule, union | ||
from pants.util.docutil import bin_name | ||
|
@@ -78,9 +77,69 @@ def debug_hint(self) -> str | None: | |
return self.resolve_name | ||
|
||
|
||
@rule_helper | ||
async def _do_export( | ||
requirements_pex: Pex, | ||
pex_pex: PexPEX, | ||
dest: str, | ||
resolve_name: str, | ||
qualify_path_with_python_version: bool, | ||
) -> ExportResult: | ||
# Get the path to the interpreter, and the full python version (including patch #). | ||
res = await Get( | ||
ProcessResult, | ||
PexProcess( | ||
pex=requirements_pex, | ||
description="Get interpreter path and version", | ||
argv=[ | ||
"-c", | ||
"import sys; print(sys.executable); print('.'.join(str(x) for x in sys.version_info[0:3]))", | ||
], | ||
extra_env={"PEX_INTERPRETER": "1"}, | ||
), | ||
) | ||
interpreter_path, py_version = res.stdout.strip().decode().split("\n") | ||
|
||
# NOTE: We add a unique prefix to the pex_pex path to avoid conflicts when multiple | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should probably materialize the Pex PEX once, to some location under .pants.d say, rather than once, in a temporary location, per export? But that can wait. |
||
# venvs are concurrently exporting. Without this prefix all the invocations write | ||
# the pex_pex to `python/virtualenvs/tools/pex`, and the `rm -f` of the pex_pex | ||
# path in one export will delete the binary out from under the others. | ||
pex_pex_dir = f".{resolve_name}.tmp" | ||
pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir)) | ||
pex_pex_dest = os.path.join("{digest_root}", pex_pex_dir) | ||
|
||
merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, requirements_pex.digest])) | ||
|
||
description = f"for {resolve_name} " if resolve_name else "" | ||
return ExportResult( | ||
f"virtualenv {description}(using Python {py_version})", | ||
dest, | ||
digest=merged_digest, | ||
post_processing_cmds=[ | ||
PostProcessingCommand( | ||
[ | ||
interpreter_path, | ||
os.path.join(pex_pex_dest, pex_pex.exe), | ||
os.path.join("{digest_root}", requirements_pex.name), | ||
"venv", | ||
"--pip", | ||
"--collisions-ok", | ||
"--remove=all", | ||
f"{{digest_root}}/{py_version if qualify_path_with_python_version else ''}", | ||
], | ||
{"PEX_MODULE": "pex.tools"}, | ||
), | ||
# Remove the PEX pex, to avoid confusion. | ||
PostProcessingCommand(["rm", "-rf", pex_pex_dest]), | ||
], | ||
) | ||
|
||
|
||
@rule | ||
async def export_virtualenv( | ||
request: _ExportVenvRequest, python_setup: PythonSetup, pex_pex: PexPEX | ||
async def export_virtualenv_for_targets( | ||
request: _ExportVenvRequest, | ||
python_setup: PythonSetup, | ||
pex_pex: PexPEX, | ||
) -> ExportResult: | ||
if request.resolve: | ||
interpreter_constraints = InterpreterConstraints( | ||
|
@@ -101,99 +160,42 @@ async def export_virtualenv( | |
), | ||
) | ||
|
||
# Note that an internal-only pex will always have the `python` field set. | ||
# See the build_pex() rule in pex.py. | ||
interpreter = cast(PythonExecutable, requirements_pex.python) | ||
|
||
# Get the full python version (including patch #), so we can use it as the venv name. | ||
res = await Get( | ||
ProcessResult, | ||
Process( | ||
description="Get interpreter version", | ||
argv=[ | ||
interpreter.path, | ||
"-c", | ||
"import sys; print('.'.join(str(x) for x in sys.version_info[0:3]))", | ||
], | ||
), | ||
) | ||
py_version = res.stdout.strip().decode() | ||
|
||
dest = ( | ||
os.path.join("python", "virtualenvs", path_safe(request.resolve)) | ||
if request.resolve | ||
else os.path.join("python", "virtualenv") | ||
) | ||
|
||
merged_digest = await Get(Digest, MergeDigests([pex_pex.digest, requirements_pex.digest])) | ||
pex_pex_path = os.path.join("{digest_root}", pex_pex.exe) | ||
maybe_resolve_str = f"for the resolve '{request.resolve}' " if request.resolve else "" | ||
return ExportResult( | ||
f"virtualenv {maybe_resolve_str}(using Python {py_version})", | ||
export_result = await _do_export( | ||
requirements_pex, | ||
pex_pex, | ||
dest, | ||
digest=merged_digest, | ||
post_processing_cmds=[ | ||
PostProcessingCommand( | ||
[ | ||
interpreter.path, | ||
pex_pex_path, | ||
os.path.join("{digest_root}", requirements_pex.name), | ||
"venv", | ||
"--pip", | ||
"--collisions-ok", | ||
"--remove=all", | ||
f"{{digest_root}}/{py_version}", | ||
], | ||
{"PEX_MODULE": "pex.tools"}, | ||
), | ||
PostProcessingCommand(["rm", "-f", pex_pex_path]), | ||
], | ||
request.resolve or "", | ||
qualify_path_with_python_version=True, | ||
) | ||
return export_result | ||
|
||
|
||
@rule | ||
async def export_tool(request: ExportPythonTool, pex_pex: PexPEX) -> ExportResult: | ||
assert request.pex_request is not None | ||
|
||
# TODO: Unify export_virtualenv() and export_tool(), since their implementations mostly overlap. | ||
dest = os.path.join("python", "virtualenvs", "tools") | ||
pex = await Get(Pex, PexRequest, request.pex_request) | ||
if not request.pex_request.internal_only: | ||
raise ExportError(f"The PexRequest for {request.resolve_name} must be internal_only.") | ||
|
||
# Note that an internal-only pex will always have the `python` field set. | ||
# See the build_pex() rule in pex.py. | ||
interpreter = cast(PythonExecutable, pex.python) | ||
|
||
# NOTE: We add a unique-per-tool prefix to the pex_pex path to avoid conflicts when | ||
# multiple tools are concurrently exporting. Without this prefix all the `export_tool` | ||
# invocations write the pex_pex to `python/virtualenvs/tools/pex`, and the `rm -f` of | ||
# the pex_pex path in one export will delete the binary out from under the others. | ||
pex_pex_dir = f".{request.resolve_name}.tmp" | ||
pex_pex_dest = os.path.join("{digest_root}", pex_pex_dir) | ||
pex_pex_digest = await Get(Digest, AddPrefix(pex_pex.digest, pex_pex_dir)) | ||
|
||
merged_digest = await Get(Digest, MergeDigests([pex_pex_digest, pex.digest])) | ||
return ExportResult( | ||
f"virtualenv for the tool '{request.resolve_name}'", | ||
# TODO: It seems unnecessary to qualify with "tools", since the tool resolve names don't collide | ||
# with user resolve names. We should get rid of this via a deprecation cycle. | ||
dest = os.path.join("python", "virtualenvs", "tools", request.resolve_name) | ||
pex = await Get(Pex, PexRequest, request.pex_request) | ||
export_result = await _do_export( | ||
pex, | ||
pex_pex, | ||
dest, | ||
digest=merged_digest, | ||
post_processing_cmds=[ | ||
PostProcessingCommand( | ||
[ | ||
interpreter.path, | ||
os.path.join(pex_pex_dest, pex_pex.exe), | ||
os.path.join("{digest_root}", pex.name), | ||
"venv", | ||
"--collisions-ok", | ||
"--remove=all", | ||
f"{{digest_root}}/{request.resolve_name}", | ||
], | ||
{"PEX_MODULE": "pex.tools"}, | ||
), | ||
PostProcessingCommand(["rm", "-rf", pex_pex_dest]), | ||
], | ||
request.resolve_name, | ||
# TODO: It is pretty ad-hoc that we do add the interpreter version for resolves but not for tools. | ||
# We should pick one and deprecate the other. | ||
qualify_path_with_python_version=False, | ||
) | ||
return export_result | ||
|
||
|
||
@rule | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This seemed like a more robust way to get a kosher interpreter then the previous line
interpreter = cast(PythonExecutable, requirements_pex.python)
and its clarifying comment below.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not so sure about that. The internal-onlyness is a critical thing here that is just as obscure still.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The internal-onlyness was critical in the previous method of using
requirements_pex.python
. Why is it critical here? This must give us some compatible interpreter in any case, surely?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm trying to point out you replaced a comment and re-use of an already computed and correct interpreter with another comment elsewhere about internal onlyness and a re-compute.
As far as criticality goes, it's critical always and until creating and storing a fully zipped PEX is fast and efficient. It never has been either. Zipping is always slow, and packed is needed for remote caching to work ~at all / not be storing both a wheel zip and a copy of that zipped up in 18 different non-packed PEX zips in LMDB.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah, fair enough. In this case I interpreted "critical" to mean "for correctness" - i.e.,
requirements_pex.python
is only guaranteed to be populated because of internal-onlyness. This code will blow up otherwise. Getting the path to an interpreter by running the pex instead, as this now does, sidesteps the need for this assumption and won't blow up if the Pex is not internal-only.You're referring to internal-onlyness being critical for performance, which it definitely is.
But since performance is important, it may be better to blow up, as it indicates a code bug. I'll change this.