setup.py

# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""Installation script."""

from importlib import metadata
import os
import shutil
import subprocess
import time
from pathlib import Path
from typing import List, Tuple

import setuptools
from wheel.bdist_wheel import bdist_wheel

from build_tools.build_ext import CMakeExtension, get_build_ext
from build_tools.te_version import te_version
from build_tools.utils import (
    cuda_archs,
    cuda_version,
    cusolvermp_pypi_package_name,
    get_frameworks,
    remove_dups,
    min_python_version_str,
)

frameworks = get_frameworks()
current_file_path = Path(__file__).parent.resolve()


from setuptools.command.build_ext import build_ext as BuildExtension

os.environ["NVTE_PROJECT_BUILDING"] = "1"

if "pytorch" in frameworks:
    from torch.utils.cpp_extension import BuildExtension
elif "jax" in frameworks:
    from pybind11.setup_helpers import build_ext as BuildExtension


CMakeBuildExtension = get_build_ext(BuildExtension)
archs = cuda_archs()


class TimedBdist(bdist_wheel):
    """Helper class to measure build time"""

    def run(self):
        start_time = time.perf_counter()
        super().run()
        total_time = time.perf_counter() - start_time
        print(f"Total time for bdist_wheel: {total_time:.2f} seconds")


def setup_common_extension() -> CMakeExtension:
    """Setup CMake extension for common library"""
    cmake_flags = ["-DCMAKE_CUDA_ARCHITECTURES={}".format(archs)]
    if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))):
        assert (
            os.getenv("MPI_HOME") is not None
        ), "MPI_HOME must be set when compiling with NVTE_UB_WITH_MPI=1"
        cmake_flags.append("-DNVTE_UB_WITH_MPI=ON")

    if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", "0"))):
        assert (
            os.getenv("NVSHMEM_HOME") is not None
        ), "NVSHMEM_HOME must be set when compiling with NVTE_ENABLE_NVSHMEM=1"
        cmake_flags.append("-DNVTE_ENABLE_NVSHMEM=ON")

    if bool(int(os.getenv("NVTE_BUILD_ACTIVATION_WITH_FAST_MATH", "0"))):
        cmake_flags.append("-DNVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON")

    if bool(int(os.getenv("NVTE_WITH_CUBLASMP", "0"))):
        cmake_flags.append("-DNVTE_WITH_CUBLASMP=ON")
        cublasmp_dir = os.getenv("CUBLASMP_HOME") or metadata.distribution(
            f"nvidia-cublasmp-cu{cuda_version()[0]}"
        ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}")
        cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}")

    if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))):
        cmake_flags.append("-DNVTE_WITH_CUSOLVERMP=ON")
        cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr")
        cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}")

    # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is
    # targeted. Set NVTE_WITH_NCCL_EP=0 to force off.
    nccl_ep_env = os.getenv("NVTE_WITH_NCCL_EP")
    nccl_ep_explicit = nccl_ep_env is not None
    build_with_nccl_ep = bool(int(nccl_ep_env if nccl_ep_explicit else "1"))
    if build_with_nccl_ep:
        arch_tokens = [a.strip() for a in str(archs or "").split(";") if a.strip()]
        has_hopper_or_newer = any(
            t.lower() == "native" or (t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90)
            for t in arch_tokens
        )
        if not has_hopper_or_newer:
            if nccl_ep_explicit:
                raise RuntimeError(
                    f"NVTE_WITH_NCCL_EP=1 was set but NVTE_CUDA_ARCHS ('{archs}') "
                    "contains no arch >= 90. NCCL EP requires Hopper or newer."
                )
            print(f"[NCCL EP] No arch >= 90 in NVTE_CUDA_ARCHS ('{archs}'); skipping build.")
            build_with_nccl_ep = False
    if build_with_nccl_ep:
        nccl_home = build_nccl_ep_submodule()
        cmake_flags.append(f"-DNCCL_INCLUDE_DIR={nccl_home}/include")
    else:
        cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF")

    # Add custom CMake arguments from environment variable
    nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS")
    if nvte_cmake_extra_args:
        cmake_flags.extend(nvte_cmake_extra_args.split())

    # Project directory root
    root_path = Path(__file__).resolve().parent

    return CMakeExtension(
        name="transformer_engine",
        cmake_path=root_path / Path("transformer_engine/common"),
        cmake_flags=cmake_flags,
    )


def setup_requirements() -> Tuple[List[str], List[str]]:
    """Setup Python dependencies

    Returns dependencies for runtime and testing.
    """

    # Common requirements
    install_reqs: List[str] = [
        "pydantic",
        "importlib-metadata>=1.0",
        "packaging",
        cusolvermp_pypi_package_name(),
    ]
    test_reqs: List[str] = ["pytest>=8.2.1"]

    # Framework-specific requirements
    if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
        if "pytorch" in frameworks:
            from build_tools.pytorch import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())
        if "jax" in frameworks:
            from build_tools.jax import install_requirements, test_requirements

            install_reqs.extend(install_requirements())
            test_reqs.extend(test_requirements())

    return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]]


def _discover_nccl_home() -> str:
    """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig."""
    env_home = os.environ.get("NCCL_HOME")
    if env_home:
        if (Path(env_home) / "include" / "nccl.h").exists():
            return env_home
        print(
            f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but "
            f"'{env_home}/include/nccl.h' was not found; falling back to system probes."
        )

    lib_names = ("libnccl.so", "libnccl.so.2")
    # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu).
    lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu")

    # pip-installed NCCL (nvidia-nccl-cu* wheel) lives under nvidia/nccl in
    # site-packages and has no top-level include/lib layout.
    try:
        import importlib.util

        spec = importlib.util.find_spec("nvidia.nccl")
        if spec is not None and spec.submodule_search_locations:
            pip_root = Path(next(iter(spec.submodule_search_locations)))
            if (pip_root / "include" / "nccl.h").exists() and any(
                (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names
            ):
                return str(pip_root)
    except (ImportError, ValueError):
        pass

    for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"):
        p = Path(cand)
        if (p / "include" / "nccl.h").exists() and any(
            (p / sub / name).exists() for sub in lib_subdirs for name in lib_names
        ):
            return str(p)

    try:
        out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode()
        for line in out.splitlines():
            if "libnccl.so" in line and "=>" in line:
                lib_path = Path(line.split("=>")[-1].strip())
                # Walk upward so multiarch layouts (.../lib/<triplet>/libnccl.so)
                # resolve to the prefix that contains include/nccl.h.
                for root in (lib_path.parent.parent, lib_path.parent.parent.parent):
                    if (root / "include" / "nccl.h").exists():
                        return str(root)
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass

    raise RuntimeError(
        "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix."
    )


def build_nccl_ep_submodule() -> str:
    """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME."""
    nccl_root = current_file_path / "3rdparty" / "nccl"
    if not (nccl_root / "Makefile").exists():
        raise RuntimeError(
            f"NCCL submodule not found at {nccl_root}. "
            "Run `git submodule update --init --recursive`."
        )

    build_dir = nccl_root / "build"
    nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a"
    gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode"

    # Caller gates on arch >= 90 or "native"; expand "native" to the host's
    # actual sm_XX so the build stamp distinguishes machines.
    arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()]
    arch_list: list[str] = []
    for t in arch_tokens:
        if t.lower() == "native":
            try:
                out = subprocess.check_output(
                    ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
                    stderr=subprocess.DEVNULL,
                ).decode()
            except (subprocess.CalledProcessError, FileNotFoundError) as e:
                raise RuntimeError(
                    "NVTE_CUDA_ARCHS=native requires nvidia-smi to resolve the host arch."
                ) from e
            for line in out.splitlines():
                cap = line.strip().replace(".", "")
                if cap.isdigit() and int(cap) >= 90 and cap not in arch_list:
                    arch_list.append(cap)
        else:
            bare = t.rstrip("af")
            if bare.isdigit() and int(bare) >= 90 and bare not in arch_list:
                arch_list.append(bare)
    if not arch_list:
        raise RuntimeError(
            "NCCL EP requires Hopper or newer (SM >= 90); none found in"
            f" NVTE_CUDA_ARCHS={cuda_archs()!r}. Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL"
            " EP build (the rest of TE still builds)."
        )
    gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list)

    nproc = os.cpu_count() or 8
    env = os.environ.copy()
    env["NVCC_GENCODE"] = gencode
    # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build
    # outputs to the submodule's local build/ tree.
    nccl_home = _discover_nccl_home()
    env["NCCL_HOME"] = nccl_home
    env["NCCL_EP_BUILDDIR"] = str(build_dir)

    prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None
    if not nccl_ep_lib.exists() or prev_gencode != gencode:
        if nccl_ep_lib.exists() and prev_gencode != gencode:
            print(
                f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); "
                "rebuilding libnccl_ep.a"
            )
            subprocess.check_call(
                ["make", "-C", "contrib/nccl_ep", "clean"],
                cwd=str(nccl_root),
                env=env,
            )
        print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')")
        subprocess.check_call(
            ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"],
            cwd=str(nccl_root),
            env=env,
        )
        gencode_stamp.parent.mkdir(parents=True, exist_ok=True)
        gencode_stamp.write_text(gencode)

    return nccl_home


def git_check_submodules() -> None:
    """
    Attempt to checkout git submodules automatically during setup.

    This runs successfully only if the submodules are
    either in the correct or uninitialized state.

    Note to devs: With this, any updates to the submodules itself, e.g. moving to a newer
    commit, must be commited before build. This also ensures that stale submodules aren't
    being silently used by developers.
    """

    # Provide an option to skip these checks for development.
    if bool(int(os.getenv("NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD", "0"))):
        return

    # Require git executable.
    if shutil.which("git") is None:
        return

    # Require a .gitmodules file.
    if not (current_file_path / ".gitmodules").exists():
        return

    try:
        submodules = subprocess.check_output(
            ["git", "submodule", "status", "--recursive"],
            cwd=str(current_file_path),
            text=True,
        ).splitlines()

        for submodule in submodules:
            # '-' start is for an uninitialized submodule.
            # ' ' start is for a submodule on the correct commit.
            assert submodule[0] in (
                " ",
                "-",
            ), (
                "Submodules are initialized incorrectly. If this is intended, set the "
                "environment variable `NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD` to a "
                "non-zero value to skip these checks during development. Otherwise, "
                "run `git submodule update --init --recursive` to checkout the correct"
                " submodule commits."
            )

        subprocess.check_call(
            ["git", "submodule", "update", "--init", "--recursive"],
            cwd=str(current_file_path),
        )
    except subprocess.CalledProcessError:
        return


if __name__ == "__main__":
    __version__ = te_version()

    git_check_submodules()

    with open("README.rst", encoding="utf-8") as f:
        long_description = f.read()

    # Settings for building top level empty package for dependency management.
    if bool(int(os.getenv("NVTE_BUILD_METAPACKAGE", "0"))):
        assert bool(
            int(os.getenv("NVTE_RELEASE_BUILD", "0"))
        ), "NVTE_RELEASE_BUILD env must be set for metapackage build."
        ext_modules = []
        package_data = {}
        include_package_data = False
        install_requires = []
        extras_require = {
            "core": [f"transformer_engine_cu12=={__version__}"],
            "core_cu12": [f"transformer_engine_cu12=={__version__}"],
            "core_cu13": [f"transformer_engine_cu13=={__version__}"],
            "pytorch": [f"transformer_engine_torch=={__version__}"],
            "jax": [f"transformer_engine_jax=={__version__}"],
        }
    else:
        install_requires, test_requires = setup_requirements()
        ext_modules = [setup_common_extension()]
        package_data = {"": ["VERSION.txt"]}
        include_package_data = True
        extras_require = {"test": test_requires}

        if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):
            if "pytorch" in frameworks:
                from build_tools.pytorch import setup_pytorch_extension

                ext_modules.append(
                    setup_pytorch_extension(
                        "transformer_engine/pytorch/csrc",
                        current_file_path / "transformer_engine" / "pytorch" / "csrc",
                        current_file_path / "transformer_engine",
                    )
                )
            if "jax" in frameworks:
                from build_tools.jax import setup_jax_extension

                ext_modules.append(
                    setup_jax_extension(
                        "transformer_engine/jax/csrc",
                        current_file_path / "transformer_engine" / "jax" / "csrc",
                        current_file_path / "transformer_engine",
                    )
                )

    # Configure package
    setuptools.setup(
        name="transformer_engine",
        version=__version__,
        packages=setuptools.find_packages(
            include=[
                "transformer_engine",
                "transformer_engine.*",
                "transformer_engine/build_tools",
            ],
        ),
        extras_require=extras_require,
        description="Transformer acceleration library",
        long_description=long_description,
        long_description_content_type="text/x-rst",
        ext_modules=ext_modules,
        cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist},
        python_requires=f">={min_python_version_str()}",
        classifiers=["Programming Language :: Python :: 3"],
        install_requires=install_requires,
        license_files=("LICENSE",),
        include_package_data=include_package_data,
        package_data=package_data,
    )