mirrors/uv
1
0
Fork 0
mirror of https://github.com/astral-sh/uv.git synced 2025-08-04 19:08:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Automatically infer the PyTorch index via --torch-backend=auto (#12070)

Browse source 
## Summary

This is a prototype that I'm considering shipping under `--preview`,
based on [`light-the-torch`](https://github.com/pmeier/light-the-torch).

`light-the-torch` patches pip to pull PyTorch packages from the PyTorch
indexes automatically. And, in particular, `light-the-torch` will query
the installed CUDA drivers to determine which indexes are compatible
with your system.

This PR implements equivalent behavior under `--torch-backend auto`,
though you can also set `--torch-backend cpu`, etc. for convenience.
When enabled, the registry client will fetch from the appropriate
PyTorch index when it sees a package from the PyTorch ecosystem (and
ignore any other configured indexes, _unless_ the package is explicitly
pinned to a different index).

Right now, this is only implemented in the `uv pip` CLI, since it
doesn't quite fit into the lockfile APIs given that it relies on feature
detection on the currently-running machine.

## Test Plan

On macOS, you can test this with (e.g.):

```shell
UV_TORCH_BACKEND=auto UV_CUDA_DRIVER_VERSION=450.80.2 cargo run \
  pip install torch --python-platform linux --python-version 3.12
```

On a GPU-enabled EC2 machine:

```shell
ubuntu@ip-172-31-47-149:~/uv$ UV_TORCH_BACKEND=auto cargo run pip install torch -v
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.31s
     Running `target/debug/uv pip install torch -v`
DEBUG uv 0.6.6 (e95ca063b 2025-03-14)
DEBUG Searching for default Python interpreter in virtual environments
DEBUG Found `cpython-3.13.0-linux-x86_64-gnu` at `/home/ubuntu/uv/.venv/bin/python3` (virtual environment)
DEBUG Using Python 3.13.0 environment at: .venv
DEBUG Acquired lock for `.venv`
DEBUG At least one requirement is not satisfied: torch
warning: The `--torch-backend` setting is experimental and may change without warning. Pass `--preview` to disable this warning.
DEBUG Detected CUDA driver version from `/sys/module/nvidia/version`: 550.144.3
...
```
This commit is contained in:
Charlie Marsh 2025-03-19 07:37:08 -07:00 • committed by GitHub
parent e40c551b80
commit 5173b59b50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 1289 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

2
crates/uv-resolver/src/resolver/mod.rs
View file

@ -1332,7 +1332,7 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
/// apply it in two cases:
///
/// 1. Local versions, where the non-local version has greater platform coverage. The intent is
/// such that, if we're resolving `PyTorch`, and we choose `torch==2.5.2+cpu`, we want to
/// such that, if we're resolving PyTorch, and we choose `torch==2.5.2+cpu`, we want to
/// fork so that we can select `torch==2.5.2` on macOS (since the `+cpu` variant doesn't
/// include any macOS wheels).
/// 2. Platforms that the user explicitly marks as "required" (opt-in). For example, the user