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:
@@ -201,6 +197,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

@@ -10,6 +10,7 @@
     clickext,
     context,
     dependency_graph,
+    metrics,
     progress,
     requirements_file,
     resolver,
@@ -158,6 +159,8 @@ def bootstrap(
                 f"Could not produce a pip compatible constraints file. Please review {constraints_filename} for more details"
             )
 
+    metrics.summarize(wkctx, "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,
@@ -194,6 +195,7 @@ def build_sequence(
                 )
             )
             print(wheel_filename)
+    metrics.summarize(wkctx, "Building")
 
     _summary(wkctx, entries)
 

src/fromager/context.py

@@ -1,3 +1,4 @@
+import collections
 import logging
 import os
 import pathlib
@@ -72,6 +73,12 @@ def __init__(
 
         self.dependency_graph = dependency_graph.DependencyGraph()
 
+        # storing metrics
+        self.time_store: dict[str, dict[str, float]] = collections.defaultdict(
+            dict[str, float]
+        )
+        self.time_description_store: dict[str, str] = {}
+
     @property
     def pip_wheel_server_args(self) -> list[str]:
         args = ["--index-url", self.wheel_server_url]

src/fromager/metrics.py

@@ -0,0 +1,80 @@
+import functools
+import logging
+import time
+import typing
+from datetime import timedelta
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+from . import context
+
+
+def timeit(description: str) -> typing.Callable:
+    def timeit_decorator(func: typing.Callable) -> typing.Callable:
+        @functools.wraps(func)
+        def wrapper_timeit(
+            *,
+            ctx: context.WorkContext,
+            req: Requirement | None = None,
+            **kwargs: typing.Any,
+        ) -> typing.Any:
+            ctx.time_description_store[func.__name__] = description
+
+            start = time.perf_counter()
+            ret = func(ctx=ctx, 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}: {func.__name__} took {timedelta(seconds=runtime)} to {description}"
+                )
+            else:
+                logger.debug(
+                    f"{func.__name__} took {timedelta(seconds=runtime)} to {description}"
+                )
+
+            if req and version:
+                # store total time spent calling that function for a particular version of that req
+                ctx.time_store[f"{req.name}=={version}"][func.__name__] = (
+                    ctx.time_store[f"{req.name}=={version}"].get(func.__name__, 0)
+                    + runtime
+                )
+
+            return ret
+
+        return wrapper_timeit
+
+    return timeit_decorator
+
+
+def summarize(ctx: context.WorkContext, prefix: str) -> None:
+    logger = logging.getLogger(__name__)
+    for req in sorted(ctx.time_store.keys()):
+        total_time = sum(ctx.time_store[req].values())
+        log = f"{prefix} {req} took {timedelta(seconds=total_time)} total"
+        for fn_name, time_taken in ctx.time_store[req].items():
+            log += f", {timedelta(seconds=time_taken)} to {ctx.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,
     requirements_file,
@@ -47,6 +48,7 @@ def get_source_type(ctx: context.WorkContext, req: Requirement) -> str:
     return str(source_type)
 
 
+@metrics.timeit(description="download source")
 def download_source(
     *,
     ctx: context.WorkContext,
@@ -72,6 +74,7 @@ def download_source(
     return source_path
 
 
+@metrics.timeit(description="resolve source")
 def resolve_source(
     *,
     ctx: context.WorkContext,
@@ -212,17 +215,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
@@ -343,6 +335,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,
@@ -423,6 +416,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,7 +18,14 @@
 from packaging.utils import canonicalize_name, parse_wheel_filename
 from packaging.version import Version
 
-from . import external_commands, overrides, requirements_file, resolver, sources
+from . import (
+    external_commands,
+    metrics,
+    overrides,
+    requirements_file,
+    resolver,
+    sources,
+)
 
 if typing.TYPE_CHECKING:
     from . import build_environment, context
@@ -110,6 +116,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,
@@ -230,6 +237,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,
@@ -265,8 +273,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",
@@ -280,7 +286,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")
@@ -293,7 +298,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
 
 
@@ -375,6 +379,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,