Selective uv sync: Install only needed packages on PEX builds

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

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(),
         )

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))

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 <name>` 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 <name>` 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,20 +259,84 @@ 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(
         (
             *downloaded_uv.args(),
             "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,