Repository: NVIDIA/TransformerEngine

Issue:
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.


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- 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!")


Test patch: