Add prototype implementation of wheel variant specification

[charliermarsh]

(Branch taken over by @konstin)

This branch contains the prototype implementation of wheel variants proposal. It support uv pip install and uv sync for a basic test case and torch. Expect that things break, please report problems. It support static and dynamic providers as well as markers. Naturally, it is not ready for any serious usage yet.

The best way to test is to clone the repo, checkout the branch (charlie/wheel-variant) and run:

export RUST_LOG=uv_distribution_types=debug,uv_distribution::distribution_database=debug
cargo run pip install torch --index https://wheelnext.github.io/variants-index/v0.0.2/ --refresh

If you don’t have Rust installed, there's a bootstrapping process from current uv, though it's slow:

export RUST_LOG=uv_distribution_types=debug,uv_distribution::distribution_database=debug
uvx --from "uv @ git+https://github.com/astral-sh/uv@charlie/wheel-variant" uv pip install torch --index https://wheelnext.github.io/variants-index/v0.0.2/ --refresh

To use a static variant declaration, i.e. not query any providers, you pass a JSON file by settings its path in UV_VARIANT_LOCK. It assumes that you don't want to run any additional providers with it, i.e. no third party code execution, unless you set UV_VARIANT_LOCK_INCOMPLETE=1. The file format has to follow the schema below:

[metadata]
created-by = "tool-that-created-the-file" # Freeform string
version = "0.1" # PEP 440 version

[[provider]]
resolved = ["cpu_level_provider==0.1"]
plugin-api = "cpu_level_provider"
namespace = "cpu_level"

[provider.properties]
x86_64_level = ["v1", "v2"]

[[provider]]
resolved = ["nvidia-variant-provider==0.0.1"]
plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin"
namespace = "nvidia"

[provider.properties]
cuda_version_lower_bound = ["12.8"]
sm_arch = ["100_real", "120_real", "70_real", "75_real", "80_real", "86_real", "90_real"]

[warsawnv]

The build-wheels.sh step gives:

Successfully built wheels/provider_fictional_hw-1.0.0-py3-none-any.whl
+ uv build --out-dir ./wheels --project ./vendor/variantlib
warning: Failed to parse `pyproject.toml` during settings discovery:
  TOML parse error at line 101, column 11
      |
  101 | [[tool.uv.variant]]
      |           ^^^^^^^
  unknown field `variant`, expected one of `required-version`, `native-tls`, `offline`, `no-cache`, `cache-dir`, `preview`, `python-preference`, `python-downloads`, `concurrent-downloads`, `concurrent-builds`, `concurrent-installs`, `index`, `index-url`, `extra-index-url`, `no-index`, `find-links`, `index-strategy`, `keyring-provider`, `allow-insecure-host`, `resolution`, `prerelease`, `fork-strategy`, `dependency-metadata`, `config-settings`, `no-build-isolation`, `no-build-isolation-package`, `exclude-newer`, `link-mode`, `compile-bytecode`, `no-sources`, `upgrade`, `upgrade-package`, `reinstall`, `reinstall-package`, `no-build`, `no-build-package`, `no-binary`, `no-binary-package`, `python-install-mirror`, `pypy-install-mirror`, `publish-url`, `trusted-publishing`, `check-url`, `pip`, `cache-keys`, `override-dependencies`, `constraint-dependencies`, `build-constraint-dependencies`, `environments`, `required-environments`, `conflicts`, `workspace`, `sources`, `managed`, `package`, `default-groups`, `dev-dependencies`, `build-backend`

But then it seems to succeed.

[charliermarsh]

Ah sorry. I’ll update the instructions. You can either remove that section from the pyproject.toml, then run the build script, then re-add it. Or run cargo build, then point the build script to ./target/debug/uv instead of your global uv.

[warsawnv]

And the cargo run step fails with:

     Running `target/debug/uv pip install dummy_project --find-links ./wheels --extra-index-url 'https://mockhouse.wheelnext.dev/pep-xxx-variants/' --no-cache --reinstall`
error: No virtual environment found; run `uv venv` to create an environment, or pass `--system` to install into a non-virtual environment

[warsawnv]

Or run cargo build, then point the build script to ./target/debug/uv instead of your global uv

I had a feeling about that!

[charliermarsh]

You can run uv venv prior to the cargo run command to create a virtual environment in the current working directory. cargo run will then install into that virtual environment by default.

[warsawnv]

% git diff 
diff --git a/build-wheels.sh b/build-wheels.sh
index 96d89524d..193eed495 100755
--- a/build-wheels.sh
+++ b/build-wheels.sh
@@ -13,12 +13,14 @@
 
 set -euxo pipefail
 
+UV=./target/debug/uv
+
 # Create the destination directory if it doesn't exist.
 rm -rf wheels
 mkdir wheels
 
 # Build the wheels for the fictional hardware provider package.
-uv build --out-dir ./wheels --project ./vendor/provider_fictional_hw
+$UV build --out-dir ./wheels --project ./vendor/provider_fictional_hw
 
 # Build the wheels for the variant library.
-uv build --out-dir ./wheels --project ./vendor/variantlib
+$UV build --out-dir ./wheels --project ./vendor/variantlib

[seemethere]

[[tool.uv.variant]]
backend = "provider_fictional_hw.plugin:build"
requires = ["provider_fictional_hw"]

For this configuration I'm curious if this would be end-user facing or would this be package maintainer facing?

From my POV I feel like this would be best served at the package maintainer side where a potential PyTorch packaging config might look like:

[[tool.uv.variant]]
backend = "torch.build.variants:detect_variants"
requires = ["cuda_variants", "rocm_variants", "xpu_variants", "your_fav_hardware_variants"]

fwiw I don't know if requiring people to install variant plugins themselves is necessarily going to be a way to adopt this heavily since this should be something that reduces the friction towards someone installing the correct package for their machine instead of them having to remember which variant plugins they're installing vs. not.

[charliermarsh]

I was thinking it would be declared by the package maintainer, but possibly overridable by the end user (so, like, if they don't care about your_fav_hardware_variants, they can "skip" that variant provider).

[charliermarsh]

(Applied your changes, thanks @warsawnv.)