clan_cli,clan_vm_manager: Moved profiler to clan_cli. Add cmd timer to see how long run commands take

pkgs/clan-cli/clan_cli/__init__.py

@@ -10,6 +10,7 @@ from . import backups, config, facts, flakes, flash, history, machines, secrets,
 from .custom_logger import setup_logging
 from .dirs import get_clan_flake_toplevel
 from .errors import ClanCmdError, ClanError
+from .profiler import profile
 from .ssh import cli as ssh_cli
 
 log = logging.getLogger(__name__)
@@ -117,6 +118,7 @@ def create_parser(prog: str | None = None) -> argparse.ArgumentParser:
 
 
 # this will be the entrypoint under /bin/clan (see pyproject.toml config)
+@profile
 def main() -> None:
     parser = create_parser()
     args = parser.parse_args()

pkgs/clan-cli/clan_cli/cmd.py

@@ -4,6 +4,7 @@ import select
 import shlex
 import subprocess
 import sys
+from datetime import datetime
 from enum import Enum
 from pathlib import Path
 from typing import IO, Any
@@ -67,6 +68,8 @@ def run(
     error_msg: str | None = None,
 ) -> CmdOut:
     glog.debug(f"$: {shlex.join(cmd)} \nCaller: {get_caller()}")
+    tstart = datetime.now()
+
     # Start the subprocess
     process = subprocess.Popen(
         cmd,
@@ -78,6 +81,9 @@ def run(
     stdout_buf, stderr_buf = handle_output(process, log)
 
     rc = process.wait()
+    tend = datetime.now()
+
+    glog.debug(f"Command took {tend - tstart}s to run")
 
     # Wait for the subprocess to finish
     cmd_out = CmdOut(

pkgs/clan-cli/clan_cli/profiler.py

@@ -0,0 +1,120 @@
+import cProfile
+import io
+import logging
+import os
+import pstats
+import re
+import weakref
+from collections.abc import Callable
+from typing import Any
+
+# Ensure you have a logger set up for logging exceptions
+log = logging.getLogger(__name__)
+explanation = """
+cProfile Output Columns Explanation:
+
+- ncalls: The number of calls to the function. This includes both direct and indirect (recursive) calls.
+
+- tottime: The total time spent in the given function alone, excluding time spent in calls to sub-functions. 
+  This measures the function's own overhead and execution time.
+
+- percall (first instance): Represents the average time spent in the function per call, calculated as tottime divided by ncalls.
+  This value excludes time spent in sub-function calls, focusing on the function's own processing time.
+
+- cumtime: The cumulative time spent in this function and all the sub-functions it calls (both directly and indirectly). 
+  This includes all execution time within the function, from the start of its invocation to its return, 
+  including all calls to other functions and the time those calls take.
+
+- percall (second instance): Represents the average time per call, including time spent in this function and in all sub-function calls.
+  It is calculated as cumtime divided by ncalls, providing an average over all calls that includes all nested function calls.
+"""
+
+
+def print_profile(profiler: cProfile.Profile, sortkey: pstats.SortKey) -> None:
+    s = io.StringIO()
+    ps = pstats.Stats(profiler, stream=s)
+    ps.sort_stats(sortkey)
+    ps.print_stats(12)
+
+    # Process the output to trim file paths
+    output_lines = s.getvalue().split("\n")
+    for line in output_lines:
+        try:
+            parts = re.split(r"\s+", line)[
+                1:
+            ]  # Split on the first space to separate the time from the path
+            fqpath = parts[-1]
+            fpath, line_num = fqpath.split(":")
+            if os.path.sep in fpath:  # Check if this looks like it includes a path
+                fpath = trim_path_to_three_levels(fpath)
+                prefix = f"{parts[0]:>7}"
+                prefix += f"{parts[1]:>9}"
+                prefix += f"{parts[2]:>9}"
+                prefix += f"{parts[3]:>9}"
+                prefix += f"{parts[4]:>9}"
+                new_line = f"{prefix:}   {fpath}:{line_num}"
+            else:
+                new_line = line
+        except (ValueError, IndexError):
+            new_line = line  # If there's no path, leave the line as is
+        print(new_line)
+
+
+# TODO: Add an RLock for every profiler, currently not thread safe
+class ProfilerStore:
+    profilers: dict[str, cProfile.Profile]
+
+    def __init__(self) -> None:
+        self.profilers = {}
+        self._exit_callback = weakref.finalize(self, self.on_exit)
+
+    def __getitem__(self, func: Callable) -> cProfile.Profile:
+        key = f"{func.__module__}:{func.__name__}"
+        if key not in self.profilers:
+            self.profilers[key] = cProfile.Profile()
+        return self.profilers[key]
+
+    def on_exit(self) -> None:
+        for key, profiler in self.profilers.items():
+            print("=" * 7 + key + "=" * 7)
+            print_profile(profiler, pstats.SortKey.TIME)
+            print_profile(profiler, pstats.SortKey.CUMULATIVE)
+
+        if len(self.profilers) > 0:
+            print(explanation)
+
+
+def trim_path_to_three_levels(path: str) -> str:
+    parts = path.split(os.path.sep)
+    if len(parts) > 4:
+        return os.path.sep.join(parts[-4:])
+    else:
+        return path
+
+
+PROFS = ProfilerStore()
+
+
+def profile(func: Callable) -> Callable:
+    """
+    A decorator that profiles the decorated function, printing out the profiling
+    results with paths trimmed to three directories deep.
+    """
+
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        global PROFS
+        profiler = PROFS[func]
+
+        try:
+            profiler.enable()
+            res = func(*args, **kwargs)
+            profiler.disable()
+        except Exception as ex:
+            profiler.disable()
+            raise ex
+        return res
+
+    if os.getenv("PERF", "0") == "1":
+        return wrapper
+    else:
+        return func