mirror of https://github.com/astral-sh/uv
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:
parent
965d5b338f
commit
3e34aee63e
|
|
@ -1839,6 +1839,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
|
||||||
.torch_backend
|
.torch_backend
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.filter(|torch_backend| matches!(torch_backend, TorchStrategy::Cuda { .. }))
|
.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))
|
.and_then(|_| pins.get(name, version).and_then(ResolvedDist::index))
|
||||||
.map(IndexUrl::url)
|
.map(IndexUrl::url)
|
||||||
.and_then(SystemDependency::from_index)
|
.and_then(SystemDependency::from_index)
|
||||||
|
|
|
||||||
|
|
@ -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`].
|
/// Return the appropriate index URLs for the given [`TorchStrategy`].
|
||||||
pub fn index_urls(&self) -> impl Iterator<Item = &IndexUrl> {
|
pub fn index_urls(&self) -> impl Iterator<Item = &IndexUrl> {
|
||||||
match self {
|
match self {
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue