[WheelVariants - Discussion] Dependencies on variants of other packages

[mgorny]

The first of the open issues on the pre-GLEP is:

How should dependencies be expressed when a package depends on a variant of another package?

I think it makes sense to start discussing that. I'm going to list some specific problems and ideas I can think of.

Dependencies in metadata

I think the first question that needs to be answered is what kind of dependencies do we need to express. I'd like to propose a few for initial discussion, with some pseudocode syntax. For now, let's start with the actual dependencies that will appear in the generated metadata for a distribution, rather than the specification used in pyproject.toml, etc.

• Requiring a specific variant:

  blas::variant::openblas
  

  • wheels without blas::variant or with any other variant are not permitted
• Requiring one of permitted variants:

  blas::variant::{mkl,openblas}
  

  • either mkl or openblas variant is permitted
  • wheels without blas::variant or with any other variant are not permitted
• Requiring any variant to be present:

  blas::variant::*
  

  • any blas::variant is permitted
  • wheels without blas::variant are not permitted
• Disallowing a specific variant:

  !blas::variant::mkl
  

  • any blas::variant other than mkl is permitted
  • wheels without blas::variant are permitted
• Requiring a variant other than the specified set:

  blas::variant::^{mkl,openblas}
  

  • mkl and openblas are not permitted, other variants (including new variants) are permitted
  • wheels without blas::variant are not permitted
  • kinda redundant, can be more verbosely expressed as:

    blas::variant::*
    !blas::variant::mkl
    !blas::variant::openblas
    

These are all the simple cases I can think of. Technically, we could consider more complex solutions — perhaps even being able to express stuff like "I'm fine with any version of torch but if a CUDA-enabled version is installed, I also need triton" (presuming torch wouldn't depend on triton itself). But given that the requirement syntax in wheels is already quite limited and we managed, I don't think we really need to overthink that — that is, unless we are actually trying to extend dependencies in general.

Dependency specifications in pyproject.toml

The above options are good for generated metadata, but pyproject.toml would probably need some more:

• Requiring the same variant as used to build this package:

  blas::variant::=
  

  • creates blas::variant::mkl when building a mkl variant, openblas when building openblas variant, etc.
  • should we just skip it when building with no blas::variant?
• Requiring any supported variant when building against any other variant:

  cuda::runtime::?
  

  • creates cuda::runtime::* when building a CUDA variant, no dependency otherwise
  • could be used to require a CUDA-enabled dependency without fixing on a specific version
  • abuses the fact that two incompatible versions cannot be "supported" on the same system, so a wheel with a compatible runtime will always be selected

Full syntax

I would lean towards reusing the "extras" syntax. For example:

torch[optree,cuda::runtime::*]

Since names of "extras" don't permit :, the variant syntax would not collide with existing extras. I think it also fits reasonably well mentally, as a similar enough concept: "I want optree support" and "I want a CUDA-enabled variant".

[rgommers]

I'd say my first thought here is: try really hard to avoid doing that.

The draft PEP doesn't explain why it would be needed. I think if this becomes necessary, there's a good chance the design failed somewhere else. I am thinking of build variants primarily as "a more optimized package, or one that runs on specific hardware not captured by the platform tag of wheels". In the former case, it should be fine to get the base variant. In the latter case, the base variant may not work at all - but then it's the user who should be selecting the desired hardware config.

Can we make this more concrete? Thinking about PyTorch as the familiar example, the problem today is that dependencies = ['torch'] gets you the cu124 variant, which may not be what you want. Once we have access to build variants on PyPI, we'll get for examplecpu and cu126 variants. The goal is for dependencies = ['torch'] to become reliable, with the user-facing knobs for variant selection in control. The goal is not to start doing dependencies = ['torch+cu126'] in another package, that kinda defeats composability.

I can think of hypothetical cases where depending on a build variant may be somewhat helpful, but not concrete cases. Who has those use cases? Once we have them, we can think of other options to resolve them.

[DEKHTIARJonathan]

@rgommers I think I have some examples for you.

In most cases - I would think - "depending on a specific Variant" would not be the case - but you could depend on "a family of variants"

For example

• Imagine building an extension for PyTorch:
  • it implements some custom CUDA Ops/Kernels.
  • Well it probably means that you should depend on a CUDA accelerated version of PyTorch and probably the same CUDA.
  • Is that problem though ? Not really - the variant resolution should "revolve" both the plugin and base PyT to the same Variant family (a CUDA enabled build of PyT).

In summary => well build your extension exclusively as a CUDA variant and it will only be installable by someone who installed PyT as a CUDA Variant

Now I could imagine a scenario where you have CUDA 11 installed - upgrade your driver to CUDA 12 - and then install that plugin. But honestly at that point the whole thing will collapse (just as it would today).

---

I don't think we really need "explicit dependency on variant" because they should be a "compiler specialized" version of the package. And you don't need a "fine match". An AVX512 enabled package will function with a non-AVX512 enabled package

---

On the other hand - one part that is true is that different variants will have different dependencies. And I'm not sure how we can specify that.

PyT CUDA will depend on cuBLAS, not PyT ROCM.

And I don't see an "easy syntax to state that" outside of "dynamically/programatically" defining that field in a setup.py file

[rgommers]

On the other hand - one part that is true is that different variants will have different dependencies. And I'm not sure how we can specify that.

In that case in pyproject.toml you'll have:

[project]
# To make this better, we need the ability to specify *both* static and dynamic dependencies,
# which needs the PEP that Henry talked about writing.
dynamic = ['dependencies']

# Control the dependencies that end up in metadata for a specific build via a
# backend-specific mechanism, either in code or in a `[tool.backend-name.xxx]`
# table in this file.

And then frontends know that they have to check the actual wheel metadata to get the correct dependencies.

EDIT: pre-PEP draft from April 2025 at https://discuss.python.org/t/partially-dynamic-project-metadata-proposal-pre-pep/88608

[DEKHTIARJonathan]

Yeap do you think we should start working ASAP on:

• The PEP you mentioned above from Henry
• Staged Releases that @warsawnv was proposing

We might need some "prep work" done before proposing the Wheel Variants. Or is it okay to publish 3 PEPs simultaneously that work together hand in hand ?

[rgommers]

Is that problem though ? Not really - the variant resolution should "revolve" both the plugin and base PyT to the same Variant family (a CUDA enabled build of PyT).

Yep, that. I'd hope we can make the installer smart enough to know the following:

• If nothing is installed yet, find the matching torch version
• If torch+cpu is already installed, reinstall torch+cu12x
• If there is no matching variant at all, error out.

What may be needed for the second bullet is to say "needs matching variant of dependencies: torch". That is more flexible - it'll work for the next PyTorch release that changes from cu124 to cu126, while depending explicitly on torch+cu124 doesn't - and it avoids coupling metadata of one package to detailed build info of another package.

The default recommendation I think should be that if torch builds variants A, B and C, then packages that depend on it should also build variants A, B, and C. And if they can't, that needs declaring. That way we get coherent collections to the extent possible. If instead we have each package pick A, B, or C at random, we're in for a huge mess.

[rgommers]

Or is it okay to publish 3 PEPs simultaneously that work together hand in hand ?

Not forbidden, but the other two PEPs are good ideas anyway, and simpler (if still complex), so I'd try to aim for one by one. 3 at once will lead to reviewer overload as well.

[DEKHTIARJonathan]

The default recommendation I think should be that if torch builds variants A, B and C, then packages that depend on it should also build variants A, B, and C. And if they can't, that needs declaring. That way we get coherent collections to the extent possible. If instead we have each package pick A, B, or C at random, we're in for a huge mess.

Yes and no :D
Saying yes is safe but if you understand what you do - this is not a hard requirement.

Example:
Let's say PyTorch builds CUDA + SM90, the extension can build CUDA + SM90/92/95/100.

Blanket statement like above are difficult to generalize.

What if you install PyTorch CPU + AVX512, an extension without AVX512 would absolutely function 👍

I agree with the spirit of what you said. But it's not that easy. Or at least you can make that recommendation and that rule can be broken carefully if you understand what it means ;)

[DEKHTIARJonathan]

Not forbidden, but the other two PEPs are good ideas anyway, and simpler (if still complex), so I'd try to aim for one by one. 3 at once will lead to reviewer overload as well.

Alright, after the summit let's talk to Henry & Barry and get that going if everybody agrees.

[rgommers]

What if you install PyTorch CPU + AVX512, an extension without AVX512 would absolutely function

Sure, I was focused on incompatible variants. A SIMD variant can always be ignored. The point was to make depending on torch to work. The recommendation I suggested is to not publish a set of variant builds that don't work with one of the variants of a dependency. There may be exceptions, but they should be rare.

Or at least you can make that recommendation and that rule can be broken carefully if you understand what it means ;)

Yes, I think we agree here. The remaining question is how you will be allowed to break it, since there's multiple ways to go about that:

• let the user manage it manually like today,
• let the package declare "I need variant X", or
• let the package declare "I don't work with variant Y"

[DEKHTIARJonathan]

The way I, myself and only myself, see it, as a user I don't care about "Variants" I just want that it works:

pip install pluginA
pip install torch

That's it I don't want it more complex than this. Eventually maybe I'll allow myself to do a "multi plugins environment" which implies to specify pluginA > pluginB.

I'm not sure the user perspective of "I want Variant X" is going to be the average. I've rarely rarely heard a user saying I want CUDA 12.4.1 + SM92. They say "I want Package X for CUDA".
And then obviously some power users who would totally say the more specific request.

From my perspective - we should look at it from the perspective "I want pluginA" and not "I want Variant A".
We are designing (currently) things so that the user will have the variant selection automatic and plugin selection manual. And I do think it's the right perspective - very little knowledge is required to install a plugin. Anything else is ambitious.

I do believe (and maybe I'm wrong) we need to **converge to a design that automagically works. It doesn't need to be optimal in all cases. But works

[DEKHTIARJonathan]

• let the user manage it manually like today,
• let the package declare "I need variant X", or
• let the package declare "I don't work with variant Y"

I think we are gonna need real world examples of these to make progress.

let the package declare "I need variant X"

I don't see how this is useful - to install a PyTorch extension that needs CUDA12 for example you would need first to have installed the NVIDIA plugin. And that plugin would have already resolved PyTorch to CUDA 12 variant.

let the package declare "I don't work with variant Y"

I don't see how you can have the ROCM version of PyTorch and a CUDA plugin installed. Plugin resolution would prevent that.
I mean I can craft a situation like this on purpose ... But that would imply having 2 GPU drivers, 2 different brands of GPUs and changing the plugin priority between PyTorch and the Extension ... I don't think it's realistic.
I will need an example here to imagine a realistic scenario of this specific case.

---

From my perspective - this will be WAY WAY better sorted by a runtime validation

from mytorch_custom_extension import AttentionBlock
>>> raise RuntimeError("We detected an inconsistent environment. This extension requires a CUDA enabled build of PyTorch. Please check documentation ...")

Something you would inject inside the __init__.py of mytorch_custom_extension. It's not ideal. But it gives you a very clear error message and tells you how to fix the issue.

1. Wipe the virtualenv
2. Recreate the venv
3. Install the NVIDIA plugin
4. Reinstall everything

And with uv that could be lightning fast.

1. rm -rf .venv && uv venv
2. uv sync

In the background uv reads a local project settings stating you need the variant plugin NVIDIA. Done

[rgommers]

let the package declare "I need variant X"

I don't see how this is useful

This is exactly the title of this issue, and the open issue in the draft PEP. Or let me rephrase again: "I need value X of variant cuda". The ROCm stuff above is a distraction, we were only talking about cpu/cu124/cu126. This is hard to discuss async without having an established vocabulary yet:)

My main point is that the open issue isn't necessarily a good idea to resolve at all; depending on exactly torch+cu124 or some such thing is fairly likely a footgun that will come with major extra complexity. So I'd start from not including such functionality at all, and iterating when multiple serious use cases show up later on after all the basics are working.

From my perspective - this will be WAY WAY better sorted by a runtime validation

Yep, I like that. 👍🏼

[DEKHTIARJonathan]

I would need Leo / Piotr to chime in - but I don't believe the minor version of the CUDA driver matters (too much). I'm low confidence here ... But essentially that's the reason why NV doesn't publish nvidia-cublas-cu12.1 but nvidia-cublas-cu12.

My main point is that the open issue isn't necessarily a good idea to resolve at all

Yeap I agree here. Unless someone can come up with a scenario to explain how it's essential. I think this will dramatically increase design complexity and I'm not sure of the benefit really.