Use pyproject.toml to specify build requirements

.github/workflows/build.yml

@@ -43,7 +43,7 @@ jobs:
         with:
           submodules: recursive
       - name: 'Build'
-        run: pip install . -v --no-deps
+        run: pip install . -v --no-build-isolation
         env:
           NVTE_FRAMEWORK: pytorch
           MAX_JOBS: 1
@@ -61,7 +61,7 @@ jobs:
         with:
           submodules: recursive
       - name: 'Build'
-        run: pip install . -v
+        run: pip install . -v --no-build-isolation
         env:
           NVTE_FRAMEWORK: jax
       - name: 'Sanity check'
@@ -81,7 +81,7 @@ jobs:
         run: |
           apt-get update
           apt-get install -y libgoogle-glog-dev
-          pip install . -v
+          pip install . -v --no-build-isolation
         env:
           NVTE_FRAMEWORK: paddle
       - name: 'Sanity check'

README.rst

@@ -172,7 +172,7 @@ To install the latest stable version of Transformer Engine,
 
 .. code-block:: bash
 
-    pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
 This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable NVTE_FRAMEWORK to a comma-separated list (e.g. NVTE_FRAMEWORK=jax,pytorch).
 

build_tools/utils.py

@@ -262,15 +262,16 @@ def copy_common_headers(te_src, dst):
         shutil.copy(file_path, new_path)
 
 
-def install_and_import(package):
-    """Install a package via pip (if not already installed) and import into globals."""
-    main_package = package.split("[")[0]
-    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
-    globals()[main_package] = importlib.import_module(main_package)
+def install_packages(packages):
+    """Install a package via pip (if not already installed)."""
+    for package in packages:
+        main_package = package.split(">")[0]
+        subprocess.run([sys.executable, "-m", "pip", "install", package], check=True)
+        globals()[main_package] = importlib.import_module(main_package)
 
 
 def uninstall_te_fw_packages():
-    subprocess.check_call(
+    subprocess.run(
         [
             sys.executable,
             "-m",

docs/installation.rst

@@ -44,7 +44,7 @@ Execute the following command to install the latest stable version of Transforme
 
 .. code-block:: bash
 
-  pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+  pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
 This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,pytorch`).
 
@@ -61,15 +61,15 @@ Execute the following command to install the latest development build of Transfo
 
 .. code-block:: bash
 
-  pip install git+https://github.com/NVIDIA/TransformerEngine.git@main
+  pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@main
 
 This will automatically detect if any supported deep learning frameworks are installed and build Transformer Engine support for them. To explicitly specify frameworks, set the environment variable `NVTE_FRAMEWORK` to a comma-separated list (e.g. `NVTE_FRAMEWORK=jax,pytorch`). To only build the framework-agnostic C++ API, set `NVTE_FRAMEWORK=none`.
 
 In order to install a specific PR, execute after changing NNN to the PR number:
 
 .. code-block:: bash
 
-  pip install git+https://github.com/NVIDIA/TransformerEngine.git@refs/pull/NNN/merge
+  pip install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@refs/pull/NNN/merge
 
 
 Installation (from source)
@@ -83,8 +83,8 @@ Execute the following commands to install Transformer Engine from source:
   git clone --branch stable --recursive https://github.com/NVIDIA/TransformerEngine.git
 
   cd TransformerEngine
-  export NVTE_FRAMEWORK=pytorch   # Optionally set framework
-  pip install .                   # Build and install
+  export NVTE_FRAMEWORK=pytorch         # Optionally set framework
+  pip install . --no-build-isolation    # Build and install
 
 If the Git repository has already been cloned, make sure to also clone the submodules:
 
@@ -96,10 +96,10 @@ Extra dependencies for testing can be installed by setting the "test" option:
 
 .. code-block:: bash
 
-  pip install .[test]
+  pip install .[test] --no-build-isolation
 
 To build the C++ extensions with debug symbols, e.g. with the `-g` flag:
 
 .. code-block:: bash
 
-  pip install . --global-option=--debug
+  pip install . --global-option=--debug --no-build-isolation

pyproject.toml

 args: [--line-length=100, --preview, --enable-unstable-feature=string_processing] 
 args: [--line-length=100, --preview, --enable-unstable-feature=string_processing] 
@@ -0,0 +1,5 @@
+[build-system]
+requires = ["setuptools>=61.0", "cmake>=3.21", "pybind11", "ninja", "pip"]
+
+# Use legacy backend to import local packages in setup.py
+build-backend = "setuptools.build_meta:__legacy__"

requirements.txt

@@ -0,0 +1,9 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+setuptools>=61.0
+cmake>=3.21
+pybind11
+ninja
+pip

setup.py

@@ -17,8 +17,8 @@
     found_pybind11,
     remove_dups,
     get_frameworks,
-    install_and_import,
     uninstall_te_fw_packages,
+    install_packages,
 )
 from build_tools.te_version import te_version
 
@@ -29,14 +29,14 @@
 
 from setuptools.command.build_ext import build_ext as BuildExtension
 
+install_packages(["pybind11"])
 os.environ["NVTE_PROJECT_BUILDING"] = "1"
 
 if "pytorch" in frameworks:
     from torch.utils.cpp_extension import BuildExtension
 elif "paddle" in frameworks:
     from paddle.utils.cpp_extension import BuildExtension
 elif "jax" in frameworks:
-    install_and_import("pybind11[global]")
     from pybind11.setup_helpers import build_ext as BuildExtension
 
 
@@ -54,35 +54,27 @@ def setup_common_extension() -> CMakeExtension:
     )
 
 
-def setup_requirements() -> Tuple[List[str], List[str], List[str]]:
+def setup_requirements() -> Tuple[List[str], List[str]]:
     """Setup Python dependencies
 
     Returns dependencies for build, runtime, and testing.
     """
 
     # Common requirements
-    setup_reqs: List[str] = []
     install_reqs: List[str] = [
         "pydantic",
         "importlib-metadata>=1.0",
         "packaging",
     ]
     test_reqs: List[str] = ["pytest>=8.2.1"]
 
-    # Requirements that may be installed outside of Python
-    if not found_cmake():
-        setup_reqs.append("cmake>=3.21")
-    if not found_ninja():
-        setup_reqs.append("ninja")
-    if not found_pybind11():
-        setup_reqs.append("pybind11")
-
-    return [remove_dups(reqs) for reqs in [setup_reqs, install_reqs, test_reqs]]
+    return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]]
 
 
 if __name__ == "__main__":
     # Dependencies
-    setup_requires, install_requires, test_requires = setup_requirements()
+
+    install_requires, test_requires = setup_requirements()
 
     __version__ = te_version()
 
@@ -150,7 +142,6 @@ def setup_requirements() -> Tuple[List[str], List[str], List[str]]:
             "Programming Language :: Python :: 3.11",
             "Programming Language :: Python :: 3.12",
         ],
-        setup_requires=setup_requires,
         install_requires=install_requires,
         license_files=("LICENSE",),
         include_package_data=True,

transformer_engine/jax/setup.py

@@ -29,11 +29,11 @@
 
 
 from build_tools.build_ext import get_build_ext
-from build_tools.utils import copy_common_headers, install_and_import
+from build_tools.utils import copy_common_headers, install_packages
 from build_tools.te_version import te_version
 from build_tools.jax import setup_jax_extension
 
-install_and_import("pybind11")
+install_packages(["pybind11"])
 from pybind11.setup_helpers import build_ext as BuildExtension
 
 os.environ["NVTE_PROJECT_BUILDING"] = "1"