feat: better build_wheel.py venv handling

Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>

Author: tongyuantongyu

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,46 +70,90 @@ 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():
+    # 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()
+    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)
     print("-- Ensuring virtualenv version >=20.29.1,<22.0 is installed...")
     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()
     # 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(
+            "You are using the NVIDIA PyTorch container, but not using the provided PyTorch installation. "
+            "This will very likely cause issues. "
+            f"If anything goes wrong, please remove the environment at {venv_prefix} "
+            "to let build_wheel.py rebuild 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 +265,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 +289,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 +743,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 current Python interpreter as is, don't create venv.")
 
 
 if __name__ == "__main__":