diff --git a/docs/notes/2.34.x.md b/docs/notes/2.34.x.md index d8354d98eec..97f9218c9b1 100644 --- a/docs/notes/2.34.x.md +++ b/docs/notes/2.34.x.md @@ -32,6 +32,8 @@ Options that are set in the docker image target such as `pull` or `build_no_cach #### Python +PEX builds using the `uv` resolver now selectively sync only the needed packages from the lockfile instead of downloading the entire resolve. This significantly reduces download size and build time for targets that depend on a small subset of a large lockfile (e.g., a script needing 2 packages from a resolve containing `torch`/CUDA). Existing lockfiles are fully backward-compatible; regenerate lockfiles to enable selective sync. + #### Shell #### Javascript diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py index 2f73739843b..fc49cdcf671 100644 --- a/src/python/pants/backend/python/util_rules/pex.py +++ b/src/python/pants/backend/python/util_rules/pex.py @@ -819,6 +819,7 @@ async def build_pex( VenvFromUvLockfileRequest( lockfile=requirements_setup.uv_lockfile, python=pex_python_setup.python, + subset_req_strings=req_strings if req_strings else None, ), **implicitly(), ) diff --git a/src/python/pants/backend/python/util_rules/pex_test.py b/src/python/pants/backend/python/util_rules/pex_test.py index d50dddb3151..77c90f1cebb 100644 --- a/src/python/pants/backend/python/util_rules/pex_test.py +++ b/src/python/pants/backend/python/util_rules/pex_test.py @@ -57,6 +57,7 @@ create_pex_and_get_pex_info, parse_requirements, ) +from pants.backend.python.util_rules.uv import generate_pyproject_toml from pants.backend.python.util_rules.uv import rules as uv_rules from pants.core.goals.generate_lockfiles import GenerateLockfileResult from pants.core.register import wrap_as_resources @@ -536,6 +537,103 @@ def test_build_pex_from_uv_lockfile(rule_runner: RuleRunner) -> None: assert b"ok" in res.stdout +def test_generate_pyproject_toml_includes_dependency_groups() -> None: + """Verify that generate_pyproject_toml produces per-requirement dependency groups.""" + ic = InterpreterConstraints(["CPython==3.14.*"]) + content = generate_pyproject_toml("test", ic, ["requests>=2.0", "torch==2.12.0"]) + + # Should contain the standard project section. + assert "[project]" in content + assert 'name = "pants-lockfile-for-test"' in content + assert '"requests>=2.0",' in content + assert '"torch==2.12.0",' in content + + # Should contain per-requirement dependency groups. + assert "[dependency-groups]" in content + assert 'requests = ["requests>=2.0"]' in content + assert 'torch = ["torch==2.12.0"]' in content + + +def test_generate_pyproject_toml_canonicalizes_group_names() -> None: + """Verify that group names are canonicalized (PEP 503).""" + ic = InterpreterConstraints(["CPython==3.14.*"]) + content = generate_pyproject_toml("test", ic, ["Foo-Bar[extra]>=1.0"]) + + assert "[dependency-groups]" in content + # PEP 503 canonicalization: Foo-Bar -> foo-bar + assert 'foo-bar = ["Foo-Bar[extra]>=1.0"]' in content + + +def test_generate_pyproject_toml_accumulates_duplicate_groups() -> None: + """Verify that multiple specifiers for the same package are accumulated in one group.""" + ic = InterpreterConstraints(["CPython==3.14.*"]) + content = generate_pyproject_toml("test", ic, ["torch==2.12.0", "torch[cuda]>=2.0"]) + + assert "[dependency-groups]" in content + # Both specifiers should appear in the same group, not overwrite each other. + assert 'torch = ["torch==2.12.0", "torch[cuda]>=2.0"]' in content + + +def test_build_pex_subset_from_uv_lockfile(rule_runner: RuleRunner) -> None: + """Build a PEX from a uv lockfile requesting only a subset of the locked packages. + + This tests the selective sync path: the lockfile contains multiple packages but + the PEX build only needs one of them. With selective sync, only the needed package + (and its transitive deps) should be installed. + """ + ic = InterpreterConstraints(["CPython==3.14.*"]) + rule_runner.set_options( + [ + "--python-resolves={'test': 'test.lock'}", + "--python-resolver=uv", + ], + env_inherit=PYTHON_BOOTSTRAP_ENV, + ) + gen_result = rule_runner.request( + GenerateLockfileResult, + [ + GenerateUvLockfile( + requirements=FrozenOrderedSet( + [ + "ansicolors==1.1.8", + "cowsay @ git+https://github.com/VaasuDevanS/cowsay-python@dcf7236f0b5ece9ed56e91271486e560526049cf", + ] + ), + find_links=FrozenOrderedSet([]), + interpreter_constraints=ic, + resolve_name="test", + lockfile_dest="test.lock", + diff=False, + ) + ], + ) + digest_contents = rule_runner.request(DigestContents, [gen_result.digest]) + rule_runner.write_files({fc.path: fc.content for fc in digest_contents}) + + # Build a PEX requesting only ansicolors (a subset of the lockfile). + pex_data = create_pex_and_get_all_data( + rule_runner, + requirements=PexRequirements( + ["ansicolors==1.1.8"], + from_superset=Resolve("test", use_entire_lockfile=False), + ), + interpreter_constraints=ic, + layout=PexLayout.ZIPAPP, + additional_pants_args=( + "--python-resolves={'test': 'test.lock'}", + "--python-resolver=uv", + ), + ) + + # The subset PEX should have ansicolors but not cowsay. + res = subprocess.run( + [pex_data.local_path, "-c", "import colors; print('subset-ok')"], + capture_output=True, + ) + assert res.returncode == 0 + assert b"subset-ok" in res.stdout + + def test_entry_point(rule_runner: RuleRunner) -> None: entry_point = "pydoc" pex_info = create_pex_and_get_pex_info(rule_runner, main=EntryPoint(entry_point)) diff --git a/src/python/pants/backend/python/util_rules/uv.py b/src/python/pants/backend/python/util_rules/uv.py index 945106679f2..5b88b4a99b3 100644 --- a/src/python/pants/backend/python/util_rules/uv.py +++ b/src/python/pants/backend/python/util_rules/uv.py @@ -13,6 +13,7 @@ from typing import ClassVar, cast from packaging.requirements import Requirement +from packaging.utils import canonicalize_name from pants.backend.python.subsystems import uv as uv_subsystem from pants.backend.python.subsystems.python_native_code import PythonNativeCodeSubsystem @@ -57,10 +58,16 @@ @dataclass(frozen=True) class VenvFromUvLockfileRequest: - """Request to install all packages from a uv lockfile into a virtualenv.""" + """Request to install packages from a uv lockfile into a virtualenv. + + If subset_req_strings is set, only those requirements (and their transitive deps) + will be installed via per-requirement dependency groups. If None, the entire lockfile + is synced. + """ lockfile: LoadedLockfile python: PythonExecutable + subset_req_strings: tuple[str, ...] | None = None @dataclass(frozen=True) @@ -118,6 +125,9 @@ async def get_uv_environment( # The synthetic project name (pants-lockfile-for-*) must not collide with any real requirement. # uv will include this project as a virtual package in the lockfile, and we set package = false, # so it won't try to install it. +# +# Per-requirement dependency groups are generated so that `uv sync --only-group ` can +# install a subset of requirements during PEX builds, avoiding downloading the entire resolve. def generate_pyproject_toml( resolve: str, ics: InterpreterConstraints, @@ -129,7 +139,8 @@ def escape_double_quotes(s: str) -> str: return s.replace('"', '\\"') requires_python = ",".join(str(constraint.specifier) for constraint in ics) - deps_lines = "\n".join(f' "{escape_double_quotes(r)}",' for r in sorted(reqs)) + sorted_reqs = sorted(reqs) + deps_lines = "\n".join(f' "{escape_double_quotes(r)}",' for r in sorted_reqs) content = dedent( """ @@ -165,6 +176,26 @@ def escape_double_quotes(s: str) -> str: content += "\n".join(extra_lines) + # Build per-requirement dependency groups. Each top-level requirement gets its own + # group named after the canonicalized package name (PEP 503). This allows selective + # install via `uv sync --only-group ` during PEX builds. + # Multiple specifiers for the same package (e.g. different extras) are accumulated + # into a single group. + groups: dict[str, list[str]] = {} + for r in sorted_reqs: + try: + parsed = Requirement(r) + group_name = canonicalize_name(parsed.name) + groups.setdefault(group_name, []).append(f'"{escape_double_quotes(r)}"') + except Exception: + logger.debug(f"Could not parse requirement {r!r} for dependency group generation") + + if groups: + group_lines = "\n".join( + f"{name} = [{', '.join(deps)}]" for name, deps in sorted(groups.items()) + ) + content += f"\n[dependency-groups]\n{group_lines}\n" + return content @@ -176,7 +207,13 @@ async def create_venv_repository_from_uv_lockfile( realpath_binary: RealpathBinary, buildroot: BuildRoot, ) -> VenvRepository: - """Install all packages from a uv lockfile into a virtualenv.""" + """Install packages from a uv lockfile into a virtualenv. + + When subset_req_strings is provided, only those packages (and their transitive deps) + are installed using per-requirement dependency groups and --only-group. The --inexact + flag prevents removal of previously-installed packages, allowing the shared venv to + accumulate packages across builds. + """ if request.lockfile.lockfile_format != LockfileFormat.UV: raise ValueError(f"Expected a uv lockfile, got {request.lockfile.lockfile_format}") if request.lockfile.metadata is None: @@ -209,9 +246,8 @@ async def create_venv_repository_from_uv_lockfile( ), get_digest_contents(request.lockfile.lockfile_digest), ) - uv_lock_digest = await create_digest( - CreateDigest([FileContent("uv.lock", uv_lock_contents[0].content)]) - ) + uv_lock_content = uv_lock_contents[0].content + uv_lock_digest = await create_digest(CreateDigest([FileContent("uv.lock", uv_lock_content)])) input_digest = await merge_digests( MergeDigests( @@ -223,11 +259,76 @@ async def create_venv_repository_from_uv_lockfile( ) ) - # We maintain one cached venv per buildroot+interpreter+resolve. uv will efficiently - # incrementally update the venv as the lockfile changes, and will handle concurrency of - # `uv sync` with appropriate locking. + # Build the uv sync arguments depending on whether we're doing a full or subset sync. + # Selective sync requires the lockfile to have been generated WITH dependency groups + # baked in (i.e., generated by the new generate_pyproject_toml that emits groups). + # Old lockfiles won't have [package.dev-dependencies] entries, so --only-group would + # fail. Detect this by checking for the dev-dependencies section in the lockfile. + lockfile_has_groups = b"[package.dev-dependencies]" in uv_lock_content + + # Also derive available group names from metadata to catch cases where a target + # requirement isn't a top-level resolve entry. + available_groups: set[str] = set() + for req_str in (str(req) for req in metadata.requirements): + try: + available_groups.add(canonicalize_name(Requirement(req_str).name)) + except Exception: + pass + + subset_group_args: list[str] = [] + if request.subset_req_strings: + if not lockfile_has_groups: + logger.warning( + "Lockfile does not contain dependency group metadata; " + "falling back to full sync. Regenerate the lockfile with " + f"`{bin_name()} generate-lockfiles` to enable selective sync." + ) + else: + for req_str in request.subset_req_strings: + try: + parsed = Requirement(req_str) + group_name = canonicalize_name(parsed.name) + if group_name not in available_groups: + logger.warning( + f"Requirement {req_str!r} (group {group_name!r}) not found " + "in lockfile dependency groups; falling back to full sync. " + "Consider regenerating the lockfile with " + f"`{bin_name()} generate-lockfiles`." + ) + subset_group_args = [] + break + subset_group_args.extend(["--only-group", group_name]) + except Exception: + logger.warning( + f"Could not parse requirement {req_str!r} for selective sync; " + "falling back to full sync." + ) + subset_group_args = [] + break + + if subset_group_args: + # Selective sync: install only the requested packages and their transitive deps. + # --inexact prevents removal of previously-installed packages, allowing the shared + # venv to accumulate packages across different PEX builds. + sync_mode_args = ("--inexact", *subset_group_args) + else: + # Full sync: install everything in the lockfile. + # TODO: extras can conflict, so we might need to be more selective. + sync_mode_args = ("--all-extras",) + + # When using --inexact, we include a lockfile content hash in the venv path so that + # stale packages from old lockfile versions are not visible to Pex. buildroot_entropy = hashlib.sha256(buildroot.path.encode()).hexdigest() - venv_path_suffix = os.path.join(buildroot_entropy, metadata.resolve, request.python.fingerprint) + if subset_group_args: + lock_hash = hashlib.sha256(uv_lock_content).hexdigest()[:12] + venv_path_suffix = os.path.join( + buildroot_entropy, metadata.resolve, request.python.fingerprint, lock_hash + ) + else: + # Full sync: uv manages the venv contents exactly, so no hash needed. + venv_path_suffix = os.path.join( + buildroot_entropy, metadata.resolve, request.python.fingerprint + ) uv_cmd = shlex.join( ( @@ -235,8 +336,7 @@ async def create_venv_repository_from_uv_lockfile( "sync", "--frozen", "--no-install-project", - # TODO: extras can conflict, so we might need to be more selective. - "--all-extras", + *sync_mode_args, "--no-progress", "--python", request.python.path,