A variant provider plugin for the Wheel Variant upcoming proposed standard that enables automatic detection and selection of NVIDIA GPU-optimized Python packages.
- Overview
- Features
- Installation
- How It Works
- Usage
- Building Variant Wheels
- Configuration
- API Reference
- Examples
- Troubleshooting
- Contributing
- License
The NVIDIA Variant Provider Plugin is part of the work conducted under WheelNext initiative to "Re-invent the Wheel" and Python package distribution for scientific computing and hardware-accelerated computing.
This package provides the logic to automatically detects NVIDIA GPU capabilities and CUDA environments to select the most optimized package variants for your system.
- Automatic GPU Detection: Detects CUDA driver versions, NVIDIA GPU compute capabilities
- Seamless Integration: Works transparently with pip, uv, and other Python package installers (when they will support Wheel Variants)
- Backward Compatible: Maintains compatibility with existing Python packaging infrastructure
Current Python wheels for GPU-accelerated packages like PyTorch can be non trivial to install. Users must manually:
- Check their CUDA version
- Navigate to special package indexes
- Deal with compatibility matrices
With Wheel Variants and this plugin, users can simply run:
[uv] pip install torch
The plugin automatically:
- Detects your GPU (e.g., RTX 4090 with compute capability 8.9)
- Identifies your CUDA driver version
- Downloads the right version of
torch
(or other requested package) compatible with your machine.
-
User-Mode Driver (UMD) Version
- System NVIDIA driver version (e.g., "12.9")
- Follows CUDA version compatibility rules
-
GPU Architecture (Compute Capability)
- Determine the compute capability available on the system.
- Resolve with compute capability compatibility in mind.
- Follows the
CMAKE
flag standard for clarity:{major}{minor}_[real|virtual]
This package is automatically installed when necessary, it is not necessary to download and install this package. However, if you still wish to do so, here is how:
[uv] pip install nvidia-variant-provider
[uv] pip install "nvidia-variant-provider @ git+https://github.com/wheelnext/nvidia-variant-provider.git"
Variant Properties emitted follow a three-tuple structure
namespace :: feature :: value
Examples:
nvidia :: cuda_version_lower_bound :: 12.6
- Means:CUDA version >= 12.6
nvidia :: cuda_version_upper_bound :: 13
- Means:CUDA version < 13
nvidia :: sm_arch :: 90_real
- Means compatible with CMAKE flag90_real
- Initialization: The plugin uses NVIDIA Management Library (NVML) to query system information
- Hardware Detection: Identifies all NVIDIA GPUs and their capabilities, read the NVIDIA User-Mode Driver version.
- Property Generation: Creates variant properties in the format
nvidia :: feature :: value
- Priority Ordering: Returns features in order of importance for package selection
Once installed, the plugin works automatically with variant-aware package installers:
# Automatic variant selection
[uv] pip install torch # Automatically selects GPU-optimized variant
# Force specific variant (if needed)
[uv] pip install "torch#cu129"
[uv] pip install "torch==2.8.0#cu129"
# Disable variant selection
[uv] pip install --no-variant torch
This plugin includes 2 environments variable to overwrite the detection mechanism and force the resolution to anything you wish.
This can be used to either work around a known problem with this plugin, debug the installer toolchain without an actual NVIDIA GPU or any other purpose.
**Disclaimer:•• Using this feature may lead you to a non functional installation.
export NV_VARIANT_PROVIDER_FORCE_CUDA_DRIVER_VERSION = "12.9" # replace with the value you want to try
export NV_VARIANT_PROVIDER_FORCE_SM_ARCH = "9.0" # replace with the value you want to try
Example Usage:
# Test with specific CUDA version
export NV_VARIANT_PROVIDER_FORCE_CUDA_DRIVER_VERSION=12.8
[uv] pip install torch
# Test with specific architecture
export NV_VARIANT_PROVIDER_FORCE_SM_ARCH=10.0
[uv] pip install torch
Add variant configuration to your pyproject.toml
:
[variant.default-priorities]
namespace = ["nvidia"]
[variant.providers.nvidia]
requires = ["nvidia-variant-provider>=0.0.1,<1.0.0"]
enable-if = "platform_system == 'Linux' or platform_system == 'Windows'"
plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin"
The main plugin class that implements hardware detection.
from nvidia_variant_provider.plugin import NvidiaVariantPlugin
plugin = NvidiaVariantPlugin()
assert plugin.namespace == "nvidia"
Returns the list of supported variant configurations based on the current system.
configs = plugin.get_supported_configs()
# Returns:
# [
# VariantFeatureConfig(name='cuda_version_lower_bound', values=['12.8', '12.7', ...]),
# VariantFeatureConfig(name='cuda_version_upper_bound', values=[..., '13.0', '12.9']),
# VariantFeatureConfig(name='sm_arch', values=['95_real', ..., '90_real', '90_virtual']),
# ]
Returns all possible variant configurations (for validation/testing).
Complete pyproject.toml
example:
[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "my-cuda-package"
version = "1.0.0"
dependencies = [
"numpy",
"nvidia-cuda-runtime-cu12 ; 'nvidia::cuda::12' in variant_properties",
"nvidia-cuda-runtime-cu11 ; 'nvidia::cuda::11' in variant_properties",
]
[variant.providers.nvidia]
requires = ["nvidia-variant-provider>=0.0.1,<1.0.0"]
enable-if = "platform_system == 'Linux' or platform_system == 'Windows'"
plugin-api = "nvidia_variant_provider.plugin:NvidiaVariantPlugin"
[variant.default-priorities]
namespace = ["nvidia"]
Problem: Error message about NVML
NVMLError: Initialization error
Solution:
- Ensure NVIDIA drivers are properly installed
- Check if
nvidia-smi
command works - May need to install NVIDIA CUDA Driver
Problem: Falls back to generic wheel despite having GPU
Solution: Check detection output:
from nvidia_variant_provider.detect_cuda import CudaEnvironment
print(CudaEnvironment.from_system())
Problem: Incorrect variant being chosen
Solution: Check variant priorities in package's pyproject.toml
and use explicit selection:
[uv] pip install package-name#cu126
Please report issues on our GitHub Issues page with:
- System information (OS, Python version)
- GPU information (
nvidia-smi
output) - Complete error messages
- Steps to reproduce
This project is licensed under the Apache 2 License - see the LICENSE file for details.
- variantlib - Core library for variant support
- WheelNext - The broader initiative for next-generation Python packaging
- PEP XXX - The Wheel Variants proposal
For more information about the WheelNext initiative and Wheel Variants, visit wheelnext.dev.