Avoid introducing unnecessary system dependency on CUDA (#15449)

## Summary

Packages like `triton` should come from the PyTorch index, but they
don't actually vary across (e.g.) the `cu128` or `cu129` indexes.

Closes https://github.com/astral-sh/uv/issues/15446.

## Test Plan

Validate that the following pins to `cu128`, rather than `cpu`:

```
echo "vllm\ntorch==2.7.1+cu128" | cargo run pip compile --torch-backend=auto --extra-index-url https://wheels.vllm.ai/b2f6c247a9b84556a8ea0e75bb4a2db765ff3315 - --python-platform linux --python-version 3.13 -v
```
This commit is contained in:
Charlie Marsh 2025-08-22 12:06:23 +01:00 committed by GitHub
parent 965d5b338f
commit 3e34aee63e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 0 deletions

View File

@ -1839,6 +1839,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
.torch_backend
.as_ref()
.filter(|torch_backend| matches!(torch_backend, TorchStrategy::Cuda { .. }))
.filter(|torch_backend| torch_backend.has_system_dependency(name))
.and_then(|_| pins.get(name, version).and_then(ResolvedDist::index))
.map(IndexUrl::url)
.and_then(SystemDependency::from_index)

View File

@ -277,6 +277,25 @@ impl TorchStrategy {
)
}
/// Returns `true` if the given [`PackageName`] has a system dependency (e.g., CUDA or ROCm).
///
/// For example, `triton` is hosted on the PyTorch indexes, but does not have a system
/// dependency on the associated CUDA version (i.e., the `triton` on the `cu128` index doesn't
/// depend on CUDA 12.8).
pub fn has_system_dependency(&self, package_name: &PackageName) -> bool {
matches!(
package_name.as_str(),
"torch"
| "torcharrow"
| "torchaudio"
| "torchcsprng"
| "torchdata"
| "torchdistx"
| "torchtext"
| "torchvision"
)
}
/// Return the appropriate index URLs for the given [`TorchStrategy`].
pub fn index_urls(&self) -> impl Iterator<Item = &IndexUrl> {
match self {