Allow specifying cmake setup directory (#1186)

Allow specifying cmake directory Signed-off-by: Ryan Li <[email protected]> Co-authored-by: Ryan Li <[email protected]>
NVIDIA · Sep 17, 2024 · 28f95bd · 28f95bd
1 parent 528d44b
commit 28f95bd
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py
@@ -106,8 +106,12 @@ def run(self) -> None:
                 if isinstance(ext, CMakeExtension):
                     print(f"Building CMake extension {ext.name}")
                     # Set up incremental builds for CMake extensions
-                    setup_dir = Path(__file__).resolve().parent.parent
-                    build_dir = setup_dir / "build" / "cmake"
+                    build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR")
+                    if build_dir:
+                        build_dir = Path(build_dir).resolve()
+                    else:
+                        root_dir = Path(__file__).resolve().parent.parent
+                        build_dir = root_dir / "build" / "cmake"
 
                     # Ensure the directory exists
                     build_dir.mkdir(parents=True, exist_ok=True)