Skip to content

wheelnext/nvidia-variant-provider

Repository files navigation

NVIDIA Variant Provider Plugin

Python Version License Wheel Variant

A variant provider plugin for the Wheel Variant upcoming proposed standard that enables automatic detection and selection of NVIDIA GPU-optimized Python packages.

Table of Contents

Overview

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.

Key Benefits

  • 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

The Problem

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

The Solution

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.

Features

Detected Hardware Properties

  1. User-Mode Driver (UMD) Version

    • System NVIDIA driver version (e.g., "12.9")
    • Follows CUDA version compatibility rules
  2. 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]

Installation

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:

Install from PyPI (when available)

[uv] pip install nvidia-variant-provider

Install from Source

[uv] pip install "nvidia-variant-provider @ git+https://github.com/wheelnext/nvidia-variant-provider.git"

How It Works

Variant Property Format

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 flag 90_real

Detection Process

  1. Initialization: The plugin uses NVIDIA Management Library (NVML) to query system information
  2. Hardware Detection: Identifies all NVIDIA GPUs and their capabilities, read the NVIDIA User-Mode Driver version.
  3. Property Generation: Creates variant properties in the format nvidia :: feature :: value
  4. Priority Ordering: Returns features in order of importance for package selection

Usage

For End Users

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

DEBUG - Overwrite Detection

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

Configuring Your Project

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"

API Reference

NvidiaVariantPlugin

The main plugin class that implements hardware detection.

from nvidia_variant_provider.plugin import NvidiaVariantPlugin

plugin = NvidiaVariantPlugin()
assert plugin.namespace == "nvidia"

Methods

get_supported_configs() -> list[VariantFeatureConfig]

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']),
# ]
get_all_configs() -> list[VariantFeatureConfig]

Returns all possible variant configurations (for validation/testing).

Example 3: Project Configuration

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"]

Troubleshooting

Common Issues

1. NVML Initialization Failed

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

2. No Variants Found

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())

3. Wrong Variant Selected

Problem: Incorrect variant being chosen

Solution: Check variant priorities in package's pyproject.toml and use explicit selection:

[uv] pip install package-name#cu126

Reporting Issues

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

License

This project is licensed under the Apache 2 License - see the LICENSE file for details.

Related Projects

  • 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.

Packages

No packages published

Contributors 2

  •  
  •  

Languages