diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 554d8ffef..fd43a9a3b 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -17,6 +17,9 @@ import os import platform import sys +import sysconfig +import time +import warnings from argparse import ArgumentParser from contextlib import contextmanager from functools import partial @@ -67,21 +70,31 @@ def clear_folder(folder_path): os.remove(item_path) -def setup_venv(project_dir: Path, requirements_file: Path): - """Creates/updates a venv and installs requirements. +def sysconfig_scheme(override_vars=None): + # Backported 'venv' scheme from Python 3.11+ + if os.name == 'nt': + scheme = { + 'purelib': '{base}/Lib/site-packages', + 'scripts': '{base}/Scripts', + } + else: + scheme = { + 'purelib': '{base}/lib/python{py_version_short}/site-packages', + 'scripts': '{base}/bin', + } - Args: - project_dir: The root directory of the project. - requirements_file: Path to the requirements file. + vars_ = sysconfig.get_config_vars() + if override_vars: + vars_.update(override_vars) + return {key: value.format(**vars_) for key, value in scheme.items()} - Returns: - Tuple[Path, Path]: Paths to the python and conan executables in the venv. - """ + +def create_venv(project_dir: Path): py_major = sys.version_info.major py_minor = sys.version_info.minor - venv_dir = project_dir / f".venv-{py_major}.{py_minor}" + venv_prefix = project_dir / f".venv-{py_major}.{py_minor}" print( - f"-- Using virtual environment at: {venv_dir} (Python {py_major}.{py_minor})" + f"-- Using virtual environment at: {venv_prefix} (Python {py_major}.{py_minor})" ) # Ensure compatible virtualenv version is installed (>=20.29.1, <22.0) @@ -89,24 +102,60 @@ def setup_venv(project_dir: Path, requirements_file: Path): build_run(f'"{sys.executable}" -m pip install "virtualenv>=20.29.1,<22.0"') # Create venv if it doesn't exist - if not venv_dir.exists(): - print(f"-- Creating virtual environment in {venv_dir}...") + if not venv_prefix.exists(): + print(f"-- Creating virtual environment in {venv_prefix}...") build_run( - f'"{sys.executable}" -m virtualenv --system-site-packages "{venv_dir}"' + f'"{sys.executable}" -m virtualenv --system-site-packages "{venv_prefix}"' ) else: print("-- Virtual environment already exists.") + return venv_prefix + + +def setup_venv(project_dir: Path, requirements_file: Path, no_venv: bool): + """Creates/updates a venv and installs requirements. + + Args: + project_dir: The root directory of the project. + requirements_file: Path to the requirements file. + no_venv: Use current Python environment as is. + + Returns: + Tuple[Path, Path]: Paths to the python and conan executables in the venv. + """ + if no_venv or sys.prefix != sys.base_prefix: + reason = "Explicitly requested by user" if no_venv else "Already inside virtual environment" + print(f"-- {reason}, using environment {sys.prefix} as is.") + venv_prefix = Path(sys.prefix) + else: + venv_prefix = create_venv(project_dir) + + scheme = sysconfig_scheme({'base': venv_prefix}) # Determine venv executable paths - scripts_dir = venv_dir / "bin" - venv_python = scripts_dir / "python" + scripts_dir = Path(scheme["scripts"]) + venv_python = venv_prefix / sys.executable.removeprefix(sys.prefix)[1:] # Install/update requirements print( - f"-- Installing requirements from {requirements_file} into {venv_dir}..." + f"-- Installing requirements from {requirements_file} into {venv_prefix}..." ) build_run(f'"{venv_python}" -m pip install -r "{requirements_file}"') + purelib_dir = Path(scheme["purelib"]) + pytorch_package_dir = purelib_dir / "torch" + has_non_system_pytorch = venv_prefix != sys.base_prefix and pytorch_package_dir.exists( + ) + if os.environ.get( + "NVIDIA_PYTORCH_VERSION") is not None and has_non_system_pytorch: + warnings.warn( + f"Using the NVIDIA PyTorch container without leveraging the included PyTorch installation " + f"and other packages may lead to compatibility issues.\n" + f"If you encounter any problems, please delete the environment at `{venv_prefix}` so that " + f"`build_wheel.py` can recreate the virtual environment.") + print("^^^^^^^^^^ IMPORTANT WARNING ^^^^^^^^^^", file=sys.stderr) + time.sleep(15) + venv_conan = setup_conan(scripts_dir, venv_python) return venv_python, venv_conan @@ -218,7 +267,8 @@ def main(*, micro_benchmarks: bool = False, nvtx: bool = False, skip_stubs: bool = False, - generate_fmha: bool = False): + generate_fmha: bool = False, + no_venv: bool = False): if clean: clean_wheel = True @@ -241,7 +291,8 @@ def main(*, # Setup venv and install requirements venv_python, venv_conan = setup_venv(project_dir, - project_dir / requirements_filename) + project_dir / requirements_filename, + no_venv) # Ensure base TRT is installed (check inside the venv) reqs = check_output([str(venv_python), "-m", "pip", "freeze"]) @@ -694,6 +745,10 @@ def add_arguments(parser: ArgumentParser): parser.add_argument("--generate_fmha", action="store_true", help="Generate the FMHA cu files.") + parser.add_argument( + "--no-venv", + action="store_true", + help="Use the current Python interpreter without creating a virtual environment.") if __name__ == "__main__":