measure time for different steps in bootstrapping and build-sequence

Signed-off-by: Shubh Bapna <[email protected]>

src/fromager/build_environment.py

@@ -13,11 +13,7 @@
 
 from packaging.requirements import Requirement
 
-from . import (
-    dependencies,
-    external_commands,
-    resolver,
-)
+from . import dependencies, external_commands, metrics, resolver
 from .requirements_file import RequirementType
 
 if typing.TYPE_CHECKING:
@@ -198,6 +194,7 @@ def _createenv(self) -> None:
         logger.info("installed dependencies into build environment in %s", self.path)
 
 
+@metrics.timeit(description="prepare build environment")
 def prepare_build_environment(
     *,
     ctx: context.WorkContext,

src/fromager/commands/bootstrap.py

@@ -9,6 +9,7 @@
     clickext,
     context,
     dependency_graph,
+    metrics,
     progress,
     requirements_file,
     resolver,
@@ -162,6 +163,8 @@ def bootstrap(
                 f"Could not produce a pip compatible constraints file. Please review {constraints_filename} for more details"
             )
 
+    metrics.summarize("Bootstrapping")
+
 
 def write_constraints_file(
     graph: dependency_graph.DependencyGraph,

src/fromager/commands/build.py

@@ -22,6 +22,7 @@
     clickext,
     context,
     hooks,
+    metrics,
     overrides,
     progress,
     server,
@@ -191,6 +192,7 @@ def build_sequence(
                 )
             )
             print(wheel_filename)
+    metrics.summarize("Building")
 
     _summary(wkctx, entries)
 

src/fromager/metrics.py

@@ -0,0 +1,77 @@
+import functools
+import logging
+import time
+import typing
+from collections import defaultdict
+from datetime import timedelta
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+_time_store: dict[str, dict[str, float]] = defaultdict(dict[str, float])
+_time_description_store: dict[str, str] = {}
+
+
+def timeit(description: str):
+    def timeit_decorator(func: typing.Callable) -> typing.Callable:
+        _time_description_store[func.__name__] = description
+
+        @functools.wraps(func)
+        def wrapper_timeit(
+            *, req: Requirement | None = None, **kwargs: typing.Any
+        ) -> typing.Any:
+            start = time.perf_counter()
+            ret = func(req=req, **kwargs)
+            end = time.perf_counter()
+            # get the logger for the module from which this function was called
+            logger = logging.getLogger(func.__module__)
+            version = (
+                kwargs.get("version")
+                or kwargs.get("dist_version")
+                or kwargs.get("resolved_version")
+            )
+            if not version:
+                version = _extract_version_from_return(ret)
+
+            runtime = end - start
+
+            if req:
+                logger.debug(
+                    f"{req.name}: took {timedelta(seconds=runtime)} to {description}"
+                )
+            else:
+                logger.debug(f"took {timedelta(seconds=runtime)} to {description}")
+
+            if req and version:
+                # store total time spent calling that function for a particular version of that req
+                _time_store[f"{req.name}=={version}"][func.__name__] = (
+                    _time_store[f"{req.name}=={version}"].get(func.__name__, 0)
+                    + runtime
+                )
+
+            return ret
+
+        return wrapper_timeit
+
+    return timeit_decorator
+
+
+def summarize(prefix: str) -> None:
+    logger = logging.getLogger(__name__)
+    for req in sorted(_time_store.keys()):
+        total_time = sum(_time_store[req].values())
+        log = f"{prefix} {req} took {timedelta(seconds=total_time)} total"
+        for fn_name, time_taken in _time_store[req].items():
+            log += f", {timedelta(seconds=time_taken)} to {_time_description_store[fn_name]}"
+        logger.info(log)
+
+
+def _extract_version_from_return(ret: typing.Any) -> Version | None:
+    try:
+        for r in ret:
+            if isinstance(r, Version):
+                return r
+    except Exception:
+        if isinstance(ret, Version):
+            return ret
+    return None

src/fromager/sources.py

@@ -18,6 +18,7 @@
 from . import (
     dependencies,
     external_commands,
+    metrics,
     overrides,
     pyproject,
     resolver,
@@ -45,6 +46,7 @@ def get_source_type(ctx: context.WorkContext, req: Requirement) -> str:
     return source_type
 
 
+@metrics.timeit(description="download source")
 def download_source(
     *,
     ctx: context.WorkContext,
@@ -70,6 +72,7 @@ def download_source(
     return source_path
 
 
+@metrics.timeit(description="resolve source")
 def resolve_source(
     *,
     ctx: context.WorkContext,
@@ -204,17 +207,6 @@ def download_url(
     return outfile
 
 
-def _sdist_root_name(source_filename: pathlib.Path) -> str:
-    base_name = pathlib.Path(source_filename).name
-    if base_name.endswith(".tar.gz"):
-        ext_to_strip = ".tar.gz"
-    elif base_name.endswith(".zip"):
-        ext_to_strip = ".zip"
-    else:
-        raise ValueError(f"Do not know how to work with {source_filename}")
-    return base_name[: -len(ext_to_strip)]
-
-
 def _takes_arg(f: typing.Callable, arg_name: str) -> bool:
     sig = inspect.signature(f)
     return arg_name in sig.parameters
@@ -354,6 +346,7 @@ def read_build_meta(unpack_dir: pathlib.Path) -> dict:
         return json.load(f)
 
 
+@metrics.timeit(description="prepare source")
 def prepare_source(
     *,
     ctx: context.WorkContext,
@@ -434,6 +427,7 @@ def prepare_new_source(
     vendor_rust.vendor_rust(req, source_root_dir)
 
 
+@metrics.timeit(description="build sdist")
 def build_sdist(
     *,
     ctx: context.WorkContext,

src/fromager/wheels.py

@@ -10,7 +10,6 @@
 import textwrap
 import typing
 import zipfile
-from datetime import datetime
 from urllib.parse import urlparse
 
 import elfdeps
@@ -19,12 +18,7 @@
 from packaging.utils import canonicalize_name, parse_wheel_filename
 from packaging.version import Version
 
-from . import (
-    external_commands,
-    overrides,
-    resolver,
-    sources,
-)
+from . import external_commands, metrics, overrides, resolver, sources
 
 if typing.TYPE_CHECKING:
     from . import build_environment, context
@@ -115,6 +109,7 @@ def default_add_extra_metadata_to_wheels(
     raise NotImplementedError
 
 
+@metrics.timeit(description="add extra metadata to wheels")
 def add_extra_metadata_to_wheels(
     *,
     ctx: context.WorkContext,
@@ -235,6 +230,7 @@ def add_extra_metadata_to_wheels(
     raise FileNotFoundError("Could not locate new wheels file")
 
 
+@metrics.timeit(description="build wheels")
 def build_wheel(
     *,
     ctx: context.WorkContext,
@@ -270,8 +266,6 @@ def build_wheel(
         )
         extra_environ["DIST_EXTRA_CONFIG"] = str(dist_extra_cfg)
 
-    # Start the timer
-    start = datetime.now().replace(microsecond=0)
     overrides.find_and_invoke(
         req.name,
         "build_wheel",
@@ -285,7 +279,6 @@ def build_wheel(
         version=version,
     )
     # End the timer
-    end = datetime.now().replace(microsecond=0)
     wheels = list(ctx.wheels_build.glob("*.whl"))
     if len(wheels) != 1:
         raise FileNotFoundError("Could not locate built wheels")
@@ -298,7 +291,6 @@ def build_wheel(
         sdist_root_dir=sdist_root_dir,
         wheel_file=wheels[0],
     )
-    logger.info(f"{req.name}: built wheel '{wheel}' in {end - start}")
     return wheel
 
 
@@ -380,6 +372,7 @@ def get_wheel_server_urls(ctx: context.WorkContext, req: Requirement) -> list[st
     return servers
 
 
+@metrics.timeit(description="resolve wheel")
 def resolve_prebuilt_wheel(
     *,
     ctx: context.WorkContext,