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

import atexit
import os
import sys
import subprocess
import io
import re
import copy
import tempfile
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
from distutils.version import LooseVersion
from distutils.file_util import copy_file


path = os.path.dirname(os.path.realpath(__file__))
with open(path + "/VERSION", "r") as f:
    te_version = f.readline()
CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")

def get_cuda_bare_metal_version(cuda_dir):
    raw_output = subprocess.check_output(
        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
    )
    output = raw_output.split()
    release_idx = output.index("release") + 1
    release = output[release_idx].split(".")
    bare_metal_major = release[0]
    bare_metal_minor = release[1][0]
    return raw_output, bare_metal_major, bare_metal_minor


def append_nvcc_threads(nvcc_extra_args):
    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
        return nvcc_extra_args + ["--threads", "4"]
    return nvcc_extra_args


def extra_gencodes(cc_flag):
    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
    if int(bare_metal_major) >= 11:
        cc_flag.append("-gencode")
        cc_flag.append("arch=compute_80,code=sm_80")
        if int(bare_metal_minor) >= 8:
            cc_flag.append("-gencode")
            cc_flag.append("arch=compute_90,code=sm_90")


def extra_compiler_flags():
    return [
        "-O3",
        "-gencode",
        "arch=compute_70,code=sm_70",
        "-U__CUDA_NO_HALF_OPERATORS__",
        "-U__CUDA_NO_HALF_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT162_OPERATORS__",
        "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
        "-I./transformer_engine/common/layer_norm/",
        "--expt-relaxed-constexpr",
        "--expt-extended-lambda",
        "--use_fast_math",
    ]


cc_flag = []
extra_gencodes(cc_flag)


def make_abs_path(l):
    return [os.path.join(path, p) for p in l]


include_dirs = [
    "transformer_engine/common/include",
    "transformer_engine/pytorch/csrc",
]
include_dirs = make_abs_path(include_dirs)

pytorch_sources = [
    "transformer_engine/pytorch/csrc/extensions.cu",
    "transformer_engine/pytorch/csrc/common.cu",
    "transformer_engine/pytorch/csrc/ts_fp8_op.cpp",
]
pytorch_sources = make_abs_path(pytorch_sources)

all_sources = pytorch_sources

supported_frameworks = {
    "all": all_sources,
    "pytorch": pytorch_sources,
    "jax": None, # JAX use transformer_engine/CMakeLists.txt
    "tensorflow": None, # tensorflow use transformer_engine/CMakeLists.txt
}

framework = os.environ.get("NVTE_FRAMEWORK", "pytorch")

args = sys.argv.copy()
for s in args:
    if s.startswith("--framework="):
        framework = s.replace("--framework=", "")
        sys.argv.remove(s)
if framework not in supported_frameworks.keys():
    raise ValueError("Unsupported framework " + framework)


class CMakeExtension(Extension):
    def __init__(self, name, cmake_path, sources, **kwargs):
        super(CMakeExtension, self).__init__(name, sources=sources, **kwargs)
        self.cmake_path = cmake_path

class FrameworkBuilderBase:
    def __init__(self, *args, **kwargs) -> None:
        pass

    def cmake_flags(self):
        return []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self, extensions):
        pass

    @staticmethod
    def install_requires():
        return []

class PyTorchBuilder(FrameworkBuilderBase):
    def __init__(self, *args, **kwargs) -> None:
        pytorch_args = copy.deepcopy(args)
        pytorch_kwargs = copy.deepcopy(kwargs)
        from torch.utils.cpp_extension import BuildExtension
        self.pytorch_build_extensions = BuildExtension(*pytorch_args, **pytorch_kwargs)

    def initialize_options(self):
        self.pytorch_build_extensions.initialize_options()

    def finalize_options(self):
        self.pytorch_build_extensions.finalize_options()

    def run(self, extensions):
        other_ext = [
            ext for ext in extensions if not isinstance(ext, CMakeExtension)
        ]
        self.pytorch_build_extensions.extensions = other_ext
        print("Building pyTorch extensions!")
        self.pytorch_build_extensions.run()

    @staticmethod
    def install_requires():
        return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class TensorFlowBuilder(FrameworkBuilderBase):
    def cmake_flags(self):
        return ["-DENABLE_TENSORFLOW=ON"]

    def run(self, extensions):
        print("Building TensorFlow extensions!")

class JaxBuilder(FrameworkBuilderBase):
    def cmake_flags(self):
        p = [d for d in sys.path if 'dist-packages' in d][0]
        return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]

    def run(self, extensions):
        print("Building jax extensions!")

    def install_requires():
        # TODO: find a way to install pybind11 and ninja directly.
        return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []

ext_modules.append(
    CMakeExtension(
        name="transformer_engine",
        cmake_path=os.path.join(path, "transformer_engine"),
        sources=[],
        include_dirs=include_dirs,
    )
)

if framework in ("all", "pytorch"):
    from torch.utils.cpp_extension import CUDAExtension
    ext_modules.append(
        CUDAExtension(
            name="transformer_engine_extensions",
            sources=supported_frameworks[framework],
            extra_compile_args={
                "cxx": ["-O3"],
                "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
            },
            include_dirs=include_dirs,
        )
    )
    dlfw_builder_funcs.append(PyTorchBuilder)

if framework in ("all", "jax"):
    dlfw_builder_funcs.append(JaxBuilder)
    # Trigger a better error when pybind11 isn't present.
    # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
    # This doesn't install a python packages. So the line bellow is too strict.
    # When it fail, we need to detect if cmake will find pybind11.
    # import pybind11

if framework in ("all", "tensorflow"):
    dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
    dlfw_install_requires = dlfw_install_requires + builder.install_requires()


def get_cmake_bin():
    cmake_bin = "cmake"
    try:
        out = subprocess.check_output([cmake_bin, "--version"])
    except OSError:
        cmake_installed_version = LooseVersion("0.0")
    else:
        cmake_installed_version = LooseVersion(
            re.search(r"version\s*([\d.]+)", out.decode()).group(1)
        )

    if cmake_installed_version < LooseVersion("3.18.0"):
        print(
            "Could not find a recent CMake to build Transformer Engine. "
            "Attempting to install CMake 3.18 to a temporary location via pip.",
            flush=True,
        )
        cmake_temp_dir = tempfile.TemporaryDirectory(prefix="nvte-cmake-tmp")
        atexit.register(cmake_temp_dir.cleanup)
        try:
            _ = subprocess.check_output(
                ["pip", "install", "--target", cmake_temp_dir.name, "cmake~=3.18.0"]
            )
        except Exception:
            raise RuntimeError(
                "Failed to install temporary CMake. "
                "Please update your CMake to 3.18+."
            )
        cmake_bin = os.path.join(cmake_temp_dir.name, "bin", "run_cmake")
        with io.open(cmake_bin, "w") as f_run_cmake:
            f_run_cmake.write(
                f"#!/bin/sh\nPYTHONPATH={cmake_temp_dir.name} {os.path.join(cmake_temp_dir.name, 'bin', 'cmake')} \"$@\""
            )
        os.chmod(cmake_bin, 0o755)

    return cmake_bin


class CMakeBuildExtension(build_ext, object):
    def __init__(self, *args, **kwargs) -> None:
        self.dlfw_flags = kwargs["dlfw_flags"]
        super(CMakeBuildExtension, self).__init__(*args, **kwargs)

    def build_extensions(self) -> None:
        print("Building CMake extensions!")

        cmake_bin = get_cmake_bin()
        config = "Debug" if self.debug else "Release"

        ext_name = self.extensions[0].name
        build_dir = self.get_ext_fullpath(ext_name).replace(
            self.get_ext_filename(ext_name), ""
        )
        build_dir = os.path.abspath(build_dir)

        cmake_args = [
            "-DCMAKE_BUILD_TYPE=" + config,
            "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
        ]
        try:
            import ninja
        except ImportError:
            pass
        else:
            cmake_args.append("-GNinja")

        cmake_args = cmake_args + self.dlfw_flags

        cmake_build_args = ["--config", config]

        cmake_build_dir = os.path.join(self.build_temp, config)
        if not os.path.exists(cmake_build_dir):
            os.makedirs(cmake_build_dir)

        config_and_build_commands = [
            [cmake_bin, self.extensions[0].cmake_path] + cmake_args,
            [cmake_bin, "--build", "."] + cmake_build_args,
        ]

        if True:
            print(f"Running CMake in {cmake_build_dir}:")
            for command in config_and_build_commands:
                print(" ".join(command))
            sys.stdout.flush()

        # Config and build the extension
        try:
            for command in config_and_build_commands:
                subprocess.check_call(command, cwd=cmake_build_dir)
        except OSError as e:
            raise RuntimeError("CMake failed: {}".format(str(e)))

class TEBuildExtension(build_ext, object):
    def __init__(self, *args, **kwargs) -> None:

        self.dlfw_builder = []
        for functor in dlfw_builder_funcs:
            self.dlfw_builder.append(functor(*args, **kwargs))

        flags = []
        for builder in self.dlfw_builder:
            flags = flags + builder.cmake_flags()

        cmake_args = copy.deepcopy(args)
        cmake_kwargs = copy.deepcopy(kwargs)
        cmake_kwargs["dlfw_flags"] = flags
        self.cmake_build_extensions = CMakeBuildExtension(*cmake_args, **cmake_kwargs)

        self.all_outputs = None
        super(TEBuildExtension, self).__init__(*args, **kwargs)

    def initialize_options(self):
        self.cmake_build_extensions.initialize_options()
        for builder in self.dlfw_builder:
            builder.initialize_options()
        super(TEBuildExtension, self).initialize_options()

    def finalize_options(self):
        self.cmake_build_extensions.finalize_options()
        for builder in self.dlfw_builder:
            builder.finalize_options()
        super(TEBuildExtension, self).finalize_options()

    def run(self) -> None:
        old_inplace, self.inplace = self.inplace, 0
        cmake_ext = [ext for ext in self.extensions if isinstance(ext, CMakeExtension)]
        self.cmake_build_extensions.extensions = cmake_ext
        self.cmake_build_extensions.run()

        for builder in self.dlfw_builder:
            builder.run(self.extensions)

        self.all_outputs = []
        for f in os.scandir(self.build_lib):
            if f.is_file():
                self.all_outputs.append(f.path)

        self.inplace = old_inplace
        if old_inplace:
            self.copy_extensions_to_source()

    def copy_extensions_to_source(self):
        ext = self.extensions[0]
        build_py = self.get_finalized_command("build_py")
        fullname = self.get_ext_fullname(ext.name)
        modpath = fullname.split(".")
        package = ".".join(modpath[:-1])
        package_dir = build_py.get_package_dir(package)

        for f in os.scandir(self.build_lib):
            if f.is_file():
                src_filename = f.path
                dest_filename = os.path.join(
                    package_dir, os.path.basename(src_filename)
                )
                # Always copy, even if source is older than destination, to ensure
                # that the right extensions for the current Python/platform are
                # used.
                copy_file(
                    src_filename,
                    dest_filename,
                    verbose=self.verbose,
                    dry_run=self.dry_run,
                )

    def get_outputs(self):
        return self.all_outputs


setup(
    name="transformer_engine",
    version=te_version,
    packages=find_packages(
        exclude=(
            "build",
            "csrc",
            "include",
            "tests",
            "dist",
            "docs",
            "tests",
            "examples",
            "transformer_engine.egg-info",
        )
    ),
    description="Transformer acceleration library",
    ext_modules=ext_modules,
    cmdclass={"build_ext": TEBuildExtension},
    install_requires=dlfw_install_requires,
    extras_require={
        'test': ['pytest',
                 'tensorflow_datasets'],
        'test_pytest': ['onnxruntime',],
    },
    license_files=("LICENSE",),
)

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

import atexit
import os
import sys
import subprocess
import io
import re
import copy
import tempfile
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
from distutils.version import LooseVersion
from distutils.file_util import copy_file


path = os.path.dirname(os.path.realpath(__file__))
with open(path + "/VERSION", "r") as f:
    te_version = f.readline()
CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")

def get_cuda_bare_metal_version(cuda_dir):
    raw_output = subprocess.check_output(
        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
    )
    output = raw_output.split()
    release_idx = output.index("release") + 1
    release = output[release_idx].split(".")
    bare_metal_major = release[0]
    bare_metal_minor = release[1][0]
    return raw_output, bare_metal_major, bare_metal_minor


def append_nvcc_threads(nvcc_extra_args):
    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
        return nvcc_extra_args + ["--threads", "4"]
    return nvcc_extra_args


def extra_gencodes(cc_flag):
    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
    if int(bare_metal_major) >= 11:
        cc_flag.append("-gencode")
        cc_flag.append("arch=compute_80,code=sm_80")
        if int(bare_metal_minor) >= 8:
            cc_flag.append("-gencode")
            cc_flag.append("arch=compute_90,code=sm_90")


def extra_compiler_flags():
    return [
        "-O3",
        "-gencode",
        "arch=compute_70,code=sm_70",
        "-U__CUDA_NO_HALF_OPERATORS__",
        "-U__CUDA_NO_HALF_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
        "-U__CUDA_NO_BFLOAT162_OPERATORS__",
        "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
        "-I./transformer_engine/common/layer_norm/",
        "--expt-relaxed-constexpr",
        "--expt-extended-lambda",
        "--use_fast_math",
    ]


cc_flag = []
extra_gencodes(cc_flag)


def make_abs_path(l):
    return [os.path.join(path, p) for p in l]


include_dirs = [
    "transformer_engine/common/include",
    "transformer_engine/pytorch/csrc",
]
include_dirs = make_abs_path(include_dirs)

pytorch_sources = [
    "transformer_engine/pytorch/csrc/extensions.cu",
    "transformer_engine/pytorch/csrc/common.cu",
    "transformer_engine/pytorch/csrc/ts_fp8_op.cpp",
]
pytorch_sources = make_abs_path(pytorch_sources)

all_sources = pytorch_sources

supported_frameworks = {
    "all": all_sources,
    "pytorch": pytorch_sources,
    "jax": None, # JAX use transformer_engine/CMakeLists.txt
    "tensorflow": None, # tensorflow use transformer_engine/CMakeLists.txt
}

framework = os.environ.get("NVTE_FRAMEWORK", "pytorch")

args = sys.argv.copy()
for s in args:
    if s.startswith("--framework="):
        framework = s.replace("--framework=", "")
        sys.argv.remove(s)
if framework not in supported_frameworks.keys():
    raise ValueError("Unsupported framework " + framework)


class CMakeExtension(Extension):
    def __init__(self, name, cmake_path, sources, **kwargs):
        super(CMakeExtension, self).__init__(name, sources=sources, **kwargs)
        self.cmake_path = cmake_path

class FrameworkBuilderBase:
    def __init__(self, *args, **kwargs) -> None:
        pass

    def cmake_flags(self):
        return []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self, extensions):
        pass

    @staticmethod
    def install_requires():
        return []

class PyTorchBuilder(FrameworkBuilderBase):
    def __init__(self, *args, **kwargs) -> None:
        pytorch_args = copy.deepcopy(args)
        pytorch_kwargs = copy.deepcopy(kwargs)
        from torch.utils.cpp_extension import BuildExtension
        self.pytorch_build_extensions = BuildExtension(*pytorch_args, **pytorch_kwargs)

    def initialize_options(self):
        self.pytorch_build_extensions.initialize_options()

    def finalize_options(self):
        self.pytorch_build_extensions.finalize_options()

    def run(self, extensions):
        other_ext = [
            ext for ext in extensions if not isinstance(ext, CMakeExtension)
        ]
        self.pytorch_build_extensions.extensions = other_ext
        print("Building pyTorch extensions!")
        self.pytorch_build_extensions.run()

    @staticmethod
    def install_requires():
        return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class TensorFlowBuilder(FrameworkBuilderBase):
    def cmake_flags(self):
        p = [d for d in sys.path if 'dist-packages' in d][0]
        return ["-DENABLE_TENSORFLOW=ON", "-DCMAKE_PREFIX_PATH="+p]

    def run(self, extensions):
        print("Building TensorFlow extensions!")

class JaxBuilder(FrameworkBuilderBase):
    def cmake_flags(self):
        p = [d for d in sys.path if 'dist-packages' in d][0]
        return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]

    def run(self, extensions):
        print("Building jax extensions!")

    def install_requires():
        # TODO: find a way to install pybind11 and ninja directly.
        return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []

ext_modules.append(
    CMakeExtension(
        name="transformer_engine",
        cmake_path=os.path.join(path, "transformer_engine"),
        sources=[],
        include_dirs=include_dirs,
    )
)

if framework in ("all", "pytorch"):
    from torch.utils.cpp_extension import CUDAExtension
    ext_modules.append(
        CUDAExtension(
            name="transformer_engine_extensions",
            sources=supported_frameworks[framework],
            extra_compile_args={
                "cxx": ["-O3"],
                "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
            },
            include_dirs=include_dirs,
        )
    )
    dlfw_builder_funcs.append(PyTorchBuilder)

if framework in ("all", "jax"):
    dlfw_builder_funcs.append(JaxBuilder)
    # Trigger a better error when pybind11 isn't present.
    # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
    # This doesn't install a python packages. So the line bellow is too strict.
    # When it fail, we need to detect if cmake will find pybind11.
    # import pybind11

if framework in ("all", "tensorflow"):
    dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
    dlfw_install_requires = dlfw_install_requires + builder.install_requires()


def get_cmake_bin():
    cmake_bin = "cmake"
    try:
        out = subprocess.check_output([cmake_bin, "--version"])
    except OSError:
        cmake_installed_version = LooseVersion("0.0")
    else:
        cmake_installed_version = LooseVersion(
            re.search(r"version\s*([\d.]+)", out.decode()).group(1)
        )

    if cmake_installed_version < LooseVersion("3.18.0"):
        print(
            "Could not find a recent CMake to build Transformer Engine. "
            "Attempting to install CMake 3.18 to a temporary location via pip.",
            flush=True,
        )
        cmake_temp_dir = tempfile.TemporaryDirectory(prefix="nvte-cmake-tmp")
        atexit.register(cmake_temp_dir.cleanup)
        try:
            _ = subprocess.check_output(
                ["pip", "install", "--target", cmake_temp_dir.name, "cmake~=3.18.0"]
            )
        except Exception:
            raise RuntimeError(
                "Failed to install temporary CMake. "
                "Please update your CMake to 3.18+."
            )
        cmake_bin = os.path.join(cmake_temp_dir.name, "bin", "run_cmake")
        with io.open(cmake_bin, "w") as f_run_cmake:
            f_run_cmake.write(
                f"#!/bin/sh\nPYTHONPATH={cmake_temp_dir.name} {os.path.join(cmake_temp_dir.name, 'bin', 'cmake')} \"$@\""
            )
        os.chmod(cmake_bin, 0o755)

    return cmake_bin


class CMakeBuildExtension(build_ext, object):
    def __init__(self, *args, **kwargs) -> None:
        self.dlfw_flags = kwargs["dlfw_flags"]
        super(CMakeBuildExtension, self).__init__(*args, **kwargs)

    def build_extensions(self) -> None:
        print("Building CMake extensions!")

        cmake_bin = get_cmake_bin()
        config = "Debug" if self.debug else "Release"

        ext_name = self.extensions[0].name
        build_dir = self.get_ext_fullpath(ext_name).replace(
            self.get_ext_filename(ext_name), ""
        )
        build_dir = os.path.abspath(build_dir)

        cmake_args = [
            "-DCMAKE_BUILD_TYPE=" + config,
            "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
        ]
        try:
            import ninja
        except ImportError:
            pass
        else:
            cmake_args.append("-GNinja")

        cmake_args = cmake_args + self.dlfw_flags

        cmake_build_args = ["--config", config]

        cmake_build_dir = os.path.join(self.build_temp, config)
        if not os.path.exists(cmake_build_dir):
            os.makedirs(cmake_build_dir)

        config_and_build_commands = [
            [cmake_bin, self.extensions[0].cmake_path] + cmake_args,
            [cmake_bin, "--build", "."] + cmake_build_args,
        ]

        if True:
            print(f"Running CMake in {cmake_build_dir}:")
            for command in config_and_build_commands:
                print(" ".join(command))
            sys.stdout.flush()

        # Config and build the extension
        try:
            for command in config_and_build_commands:
                subprocess.check_call(command, cwd=cmake_build_dir)
        except OSError as e:
            raise RuntimeError("CMake failed: {}".format(str(e)))

class TEBuildExtension(build_ext, object):
    def __init__(self, *args, **kwargs) -> None:

        self.dlfw_builder = []
        for functor in dlfw_builder_funcs:
            self.dlfw_builder.append(functor(*args, **kwargs))

        flags = []
        for builder in self.dlfw_builder:
            flags = flags + builder.cmake_flags()

        cmake_args = copy.deepcopy(args)
        cmake_kwargs = copy.deepcopy(kwargs)
        cmake_kwargs["dlfw_flags"] = flags
        self.cmake_build_extensions = CMakeBuildExtension(*cmake_args, **cmake_kwargs)

        self.all_outputs = None
        super(TEBuildExtension, self).__init__(*args, **kwargs)

    def initialize_options(self):
        self.cmake_build_extensions.initialize_options()
        for builder in self.dlfw_builder:
            builder.initialize_options()
        super(TEBuildExtension, self).initialize_options()

    def finalize_options(self):
        self.cmake_build_extensions.finalize_options()
        for builder in self.dlfw_builder:
            builder.finalize_options()
        super(TEBuildExtension, self).finalize_options()

    def run(self) -> None:
        old_inplace, self.inplace = self.inplace, 0
        cmake_ext = [ext for ext in self.extensions if isinstance(ext, CMakeExtension)]
        self.cmake_build_extensions.extensions = cmake_ext
        self.cmake_build_extensions.run()

        for builder in self.dlfw_builder:
            builder.run(self.extensions)

        self.all_outputs = []
        for f in os.scandir(self.build_lib):
            if f.is_file():
                self.all_outputs.append(f.path)

        self.inplace = old_inplace
        if old_inplace:
            self.copy_extensions_to_source()

    def copy_extensions_to_source(self):
        ext = self.extensions[0]
        build_py = self.get_finalized_command("build_py")
        fullname = self.get_ext_fullname(ext.name)
        modpath = fullname.split(".")
        package = ".".join(modpath[:-1])
        package_dir = build_py.get_package_dir(package)

        for f in os.scandir(self.build_lib):
            if f.is_file():
                src_filename = f.path
                dest_filename = os.path.join(
                    package_dir, os.path.basename(src_filename)
                )
                # Always copy, even if source is older than destination, to ensure
                # that the right extensions for the current Python/platform are
                # used.
                copy_file(
                    src_filename,
                    dest_filename,
                    verbose=self.verbose,
                    dry_run=self.dry_run,
                )

    def get_outputs(self):
        return self.all_outputs


setup(
    name="transformer_engine",
    version=te_version,
    packages=find_packages(
        exclude=(
            "build",
            "csrc",
            "include",
            "tests",
            "dist",
            "docs",
            "tests",
            "examples",
            "transformer_engine.egg-info",
        )
    ),
    description="Transformer acceleration library",
    ext_modules=ext_modules,
    cmdclass={"build_ext": TEBuildExtension},
    install_requires=dlfw_install_requires,
    extras_require={
        'test': ['pytest',
                 'tensorflow_datasets'],
        'test_pytest': ['onnxruntime',],
    },
    license_files=("LICENSE",),
)

setup.py
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 7bf1efff04..ff64f1de72 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -43,8 +43,7 @@ jobs:
         uses: actions/checkout@v3
       - name: 'Build'
         run: |
-          apt-get -y update && \
-          apt -y install ninja-build pybind11-dev && \
+          pip install ninja pybind11 && \
           pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
           mkdir -p wheelhouse && \
           NVTE_FRAMEWORK=jax pip wheel -w wheelhouse . -v
@@ -69,8 +68,7 @@ jobs:
         uses: actions/checkout@v3
       - name: 'Build'
         run: |
-          apt-get -y update && \
-          apt -y install ninja-build pybind11-dev && \
+          pip install ninja pybind11 && \
           mkdir -p wheelhouse && \
           NVTE_FRAMEWORK=tensorflow pip wheel -w wheelhouse . -v
       - name: 'Upload wheel'
diff --git a/setup.py b/setup.py
index df83504226..cba7f8f610 100644
--- a/setup.py
+++ b/setup.py
@@ -161,7 +161,8 @@ def install_requires():
 
 class TensorFlowBuilder(FrameworkBuilderBase):
     def cmake_flags(self):
-        return ["-DENABLE_TENSORFLOW=ON"]
+        p = [d for d in sys.path if 'dist-packages' in d][0]
+        return ["-DENABLE_TENSORFLOW=ON", "-DCMAKE_PREFIX_PATH="+p]
 
     def run(self, extensions):
         print("Building TensorFlow extensions!")

TF installation failed when using pip install but not apt-get
It seems we updated the installation instructions from
```
apt-get install ninja-build pybind11-dev
```
to
```
pip install pybind11
pip install ninja
```
This would cause the TF build failure with:
```
Building wheels for collected packages: transformer-engine
  Building wheel for transformer-engine (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py bdist_wheel did not run successfully.
  │ exit code: 1
  ╰─> [75 lines of output]
      /usr/local/lib/python3.8/dist-packages/setuptools/dist.py:547: UserWarning: Normalizing '0.8.0dev
      ' to '0.8.0.dev0'
        warnings.warn(tmpl.format(**locals()))
      running bdist_wheel
      running build
      running build_py
      copying transformer_engine/tensorflow/module.py -> build/lib.linux-x86_64-3.8/transformer_engine/tensorflow
      copying transformer_engine/tensorflow/transformer.py -> build/lib.linux-x86_64-3.8/transformer_engine/tensorflow
      running build_ext
      Building CMake extensions!
      Could not find a recent CMake to build Transformer Engine. Attempting to install CMake 3.18 to a temporary location via pip.
      WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
      Running CMake in build/temp.linux-x86_64-3.8/Release:
      /tmp/nvte-cmake-tmp1wzpr34f/bin/run_cmake /home/workspace/repo_zoo/TransformerEngine/transformer_engine -DCMAKE_BUILD_TYPE=Release -DCMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE=/home/workspace/repo_zoo/TransformerEngine/build/lib.linux-x86_64-3.8 -GNinja -DENABLE_TENSORFLOW=ON
      /tmp/nvte-cmake-tmp1wzpr34f/bin/run_cmake --build . --config Release
      CMake Error at CMakeLists.txt:38 (find_package):
        Could not find a package configuration file provided by "pybind11" with any
        of the following names:
                                                                                                                                                                                                                                                                                                                                 pybind11Config.cmake
          pybind11-config.cmake

        Add the installation prefix of "pybind11" to CMAKE_PREFIX_PATH or set
        "pybind11_DIR" to a directory containing one of the above files.  If
        "pybind11" provides a separate development package or SDK, be sure it has
        been installed.
      -- Configuring incomplete, errors occurred!
      See also "/home/workspace/repo_zoo/TransformerEngine/build/temp.linux-x86_64-3.8/Release/CMakeFiles/CMakeOutput.log".
      See also "/home/workspace/repo_zoo/TransformerEngine/build/temp.linux-x86_64-3.8/Release/CMakeFiles/CMakeError.log".
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "/home/workspace/repo_zoo/TransformerEngine/setup.py", line 393, in <module>
          setup(
        File "/usr/local/lib/python3.8/dist-packages/setuptools/__init__.py", line 108, in setup
          return distutils.core.setup(**attrs)
        File "/usr/lib/python3.8/distutils/core.py", line 148, in setup
          dist.run_commands()
        File "/usr/lib/python3.8/distutils/dist.py", line 966, in run_commands
          self.run_command(cmd)
        File "/usr/local/lib/python3.8/dist-packages/setuptools/dist.py", line 1221, in run_command
          super().run_command(command)
        File "/usr/lib/python3.8/distutils/dist.py", line 985, in run_command
          cmd_obj.run()
        File "/usr/local/lib/python3.8/dist-packages/wheel/bdist_wheel.py", line 343, in run
          self.run_command("build")
        File "/usr/lib/python3.8/distutils/cmd.py", line 313, in run_command
          self.distribution.run_command(command)
        File "/usr/local/lib/python3.8/dist-packages/setuptools/dist.py", line 1221, in run_command
          super().run_command(command)
        File "/usr/lib/python3.8/distutils/dist.py", line 985, in run_command
          cmd_obj.run()
        File "/usr/lib/python3.8/distutils/command/build.py", line 135, in run
          self.run_command(cmd_name)
        File "/usr/lib/python3.8/distutils/cmd.py", line 313, in run_command
          self.distribution.run_command(command)
        File "/usr/local/lib/python3.8/dist-packages/setuptools/dist.py", line 1221, in run_command
          super().run_command(command)
        File "/usr/lib/python3.8/distutils/dist.py", line 985, in run_command
          cmd_obj.run()
        File "/home/workspace/repo_zoo/TransformerEngine/setup.py", line 351, in run
          self.cmake_build_extensions.run()
        File "/usr/local/lib/python3.8/dist-packages/setuptools/command/build_ext.py", line 84, in run
          _build_ext.run(self)
        File "/usr/local/lib/python3.8/dist-packages/Cython/Distutils/old_build_ext.py", line 186, in run
          _build_ext.build_ext.run(self)
        File "/usr/lib/python3.8/distutils/command/build_ext.py", line 340, in run
          self.build_extensions()
        File "/home/workspace/repo_zoo/TransformerEngine/setup.py", line 312, in build_extensions
          subprocess.check_call(command, cwd=cmake_build_dir)
        File "/usr/lib/python3.8/subprocess.py", line 364, in check_call
          raise CalledProcessError(retcode, cmd)
      subprocess.CalledProcessError: Command '['/tmp/nvte-cmake-tmp1wzpr34f/bin/run_cmake', '/home/workspace/repo_zoo/TransformerEngine/transformer_engine', '-DCMAKE_BUILD_TYPE=Release', '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE=/home/workspace/repo_zoo/TransformerEngine/build/lib.linux-x86_64-3.8', '-GNinja',
 '-DENABLE_TENSORFLOW=ON']' returned non-zero exit status 1.\      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for transformer-engine
  Running setup.py clean for transformer-engine

```

Also, cc @trevor-m I remembered we have met the similar things before about the version issues of ninja or pybind11.
