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
|
||||
.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)
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Reference in New Issue