Timing for build (#1048)

* add timing for build

* using perf_counter

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

build_tools/build_ext.py

@@ -10,6 +10,7 @@ import subprocess
 import sys
 import sysconfig
 import copy
+import time
 
 from pathlib import Path
 from subprocess import CalledProcessError
@@ -81,6 +82,7 @@ class CMakeExtension(setuptools.Extension):
             build_command.append(str(max_jobs))
 
         # Run CMake commands
+        start_time = time.perf_counter()
         for command in [configure_command, build_command, install_command]:
             print(f"Running command {' '.join(command)}")
             try:
@@ -88,6 +90,9 @@ class CMakeExtension(setuptools.Extension):
             except (CalledProcessError, OSError) as e:
                 raise RuntimeError(f"Error when running CMake: {e}")
 
+        total_time = time.perf_counter() - start_time
+        print(f"Time for build_ext: {total_time:.2f} seconds")
+
 
 def get_build_ext(extension_cls: Type[setuptools.Extension]):
     class _CMakeBuildExtension(extension_cls):

setup.py

@@ -5,10 +5,12 @@
 """Installation script."""
 
 import os
+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.utils import (
@@ -39,10 +41,23 @@ elif "jax" in frameworks:
     install_and_import("pybind11[global]")
     from pybind11.setup_helpers import build_ext as BuildExtension
 
+# Start timing
+start_time = time.perf_counter()
+
 
 CMakeBuildExtension = get_build_ext(BuildExtension)
 
 
+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"Time for bdist_wheel: {total_time:.2f} seconds")
+
+
 def setup_common_extension() -> CMakeExtension:
     """Setup CMake extension for common library"""
     # Project directory root
@@ -141,7 +156,7 @@ if __name__ == "__main__":
         },
         description="Transformer acceleration library",
         ext_modules=ext_modules,
-        cmdclass={"build_ext": CMakeBuildExtension},
+        cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": TimedBdist},
         python_requires=">=3.8, <3.13",
         classifiers=[
             "Programming Language :: Python :: 3.8",
@@ -156,3 +171,8 @@ if __name__ == "__main__":
         include_package_data=True,
         package_data={"": ["VERSION.txt"]},
     )
+
+    # End timing
+    end_time = time.perf_counter()
+    total_time = end_time - start_time
+    print(f"Total build time: {total_time:.2f} seconds")