Implementation:NVIDIA DALI QA Setup Packages

Knowledge Sources	NVIDIA_DALI
Domains	Testing, Dependency_Management
Last Updated	2026-02-08 16:00 GMT

Overview

QA test dependency management script that generates pip install commands for DALI's test suite dependencies, with version resolution based on CUDA version, Python version, and free-threaded build status.

Description

This script (qa/setup_packages.py) implements a sophisticated package version management system for DALI's quality assurance test infrastructure. It defines a class hierarchy for describing test dependencies: PckgVer for version specifications with Python version bounds and free-threaded build constraints, BasePackage as the abstract base with version resolution logic, PlainPackage for CUDA-independent packages, CudaPackage for CUDA-version-specific packages, CudaPackageExtraIndex for packages requiring alternate pip indices, and CudaHttpPackage for packages distributed as direct HTTP downloads with platform tag compatibility checking.

The script maintains a comprehensive registry of all test dependencies including: jupyter, opencv-python-headless, cupy (with CUDA-specific variants), tensorflow-gpu (with protobuf and urllib3 constraints), torch and torchvision (via PyTorch's CUDA-specific pip index), paddlepaddle-gpu, jax with CUDA local backend, flax, clu, numba, and numba-cuda. Each package specifies exact versions, Python version compatibility bounds, optional installation aliases, dependency packages, pip constraints, and whether it supports free-threaded Python builds.

The command-line interface supports multiple operations: listing all available configurations (--list), computing the number of possible configurations (--num), generating install strings for a specific configuration index (--install N), returning all versions (--all), generating uninstall strings (--remove), retrieving extra pip index URLs (--extra_index), retrieving direct link indices (--links_index), and generating pip constraints files (--constraints). These operations are parameterized by CUDA version and a selectable subset of packages.

Usage

This script is used by DALI CI/CD pipelines and QA scripts to determine which versions of test framework dependencies (TensorFlow, PyTorch, JAX, etc.) to install for a given CUDA version and Python version combination. It is typically called from shell scripts that parse its output to construct pip install commands.

Code Reference

Source Location

• Repository: NVIDIA_DALI
• File: qa/setup_packages.py
• Lines: 1-904

Signature

class PckgVer:
    def __init__(self, ver, python_min_ver=None, python_max_ver=None,
                 alias=None, dependencies=None, constraints=[], python_free_threaded=None): ...
    def __bool__(self): ...

class BasePackage:
    def __init__(self, key, versions, name=None): ...
    def get_version(self, idx, cuda_version=None): ...
    def get_install_string(self, idx, cuda_version=None): ...
    def get_all_versions(self, cuda_version=None): ...

class PlainPackage(BasePackage): ...
class CudaPackage(BasePackage): ...
class CudaPackageExtraIndex(CudaPackage): ...
class CudaHttpPackage(CudaPackage): ...

# Top-level functions
def print_configs(cuda_version): ...
def cal_num_of_configs(packages, cuda_version): ...
def get_install_string(idx, packages, cuda_version): ...
def get_extra_indices(packages, cuda_version): ...
def gen_constraints(packages, cuda_version) -> str: ...
def main(): ...

Import

# Run as standalone script
python qa/setup_packages.py [options]

# Example: get install string for torch with CUDA 12.0
python qa/setup_packages.py --install 0 --cuda 120 --use torch

I/O Contract

Inputs

Name	Type	Required	Description
--list / -l	flag	No	List all available package configurations
--num / -n	flag	No	Print number of possible configurations
--install / -i	int	No	Get the Nth configuration install string
--all / -a	flag	No	Return all versions of requested packages
--remove / -r	flag	No	List package names for uninstall
--cuda	str	No	CUDA version string, e.g. "120" for CUDA 12.0 (default: "90")
--use / -u	list of str	No	Package keys to include (e.g., torch tensorflow-gpu)
--extra_index / -e	flag	No	Return extra pip index URLs
--links_index / -k	flag	No	Return direct link index URLs
--constraints / -c	flag	No	Generate and return path to constraints file

Outputs

Name	Type	Description
Install string	stdout	pip-compatible install string with package==version format
Package count	stdout	Number of available version configurations (with --num)
Remove string	stdout	Space-separated package names for pip uninstall (with --remove)
Extra index URL	stdout	Extra pip index URL for packages like PyTorch (with --extra_index)
Constraints file path	stdout	Path to temporary constraints file (with --constraints)

Usage Examples

List all available configurations for CUDA 12.0

python qa/setup_packages.py --list --cuda 120

Get PyTorch install string for CUDA 12.0

# Get install string
INSTALL_STR=$(python qa/setup_packages.py -i 0 --cuda 120 -u torch torchvision)
EXTRA_INDEX=$(python qa/setup_packages.py -e --cuda 120 -u torch torchvision)

# Use in pip install
pip install $INSTALL_STR --extra-index-url $EXTRA_INDEX

Get number of TensorFlow configs

NUM_CONFIGS=$(python qa/setup_packages.py -n --cuda 120 -u tensorflow-gpu)
echo "Available TF configurations: $NUM_CONFIGS"

Related Pages

• Environment:NVIDIA_DALI_CUDA_GPU_Environment