Refactor to expose command

Author: ArjunS07

asv/commands/initialize_diffcheck.py

@@ -1,51 +1,11 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import argparse
-import itertools
-import os
 
-from .. import util
-from ..benchmarks import Benchmarks
 from ..console import log
-from ..environment import get_environments
-from ..repo import NoRepository
-from ..runner import run_benchmarks
-from ..contrib.lightspeed.deps_db import BenchmarkId, LightspeedDB
-from ..contrib.lightspeed.survey import run_survey
+from ..contrib.lightspeed.session import LightspeedSession
 from . import Command, common_args
 
-_DEPS_DB_FILENAME = ".lightspeed_deps.db"
-
-
-def _get_all_bids(benchmarks):
-    bids = []
-    for name, benchmark in benchmarks.items():
-        params = benchmark.get('params')
-        if params:
-            for i in range(len(list(itertools.product(*params)))):
-                bids.append(BenchmarkId(name, i))
-        else:
-            bids.append(BenchmarkId(name))
-    return bids
-
-
-def _store_baseline_from_results(asv_results, benchmarks, db):
-    """Write timing stats from a Results object into the baseline table."""
-    for name, benchmark in benchmarks.items():
-        result_vals = asv_results._results.get(name)
-        stats_list = asv_results._stats.get(name)
-        if result_vals is None or stats_list is None:
-            continue
-        if benchmark.get('params'):
-            for param_idx, (val, stat) in enumerate(zip(result_vals, stats_list)):
-                if val is not None and stat is not None:
-                    db.store_baseline(BenchmarkId(name, param_idx), val, stat)
-        else:
-            val = result_vals[0] if result_vals else None
-            stat = stats_list[0] if stats_list else None
-            if val is not None and stat is not None:
-                db.store_baseline(BenchmarkId(name), val, stat)
-
 
 class InitializeDiffcheck(Command):
     @classmethod
@@ -85,57 +45,27 @@ def run_from_conf_args(cls, conf, args):
         return cls.run(
             conf=conf,
             source_root=args.source_root,
-            env_spec=args.env_spec,
             force=args.force,
-            bench=args.bench,
-            launch_method=getattr(args, 'launch_method', None),
+            launch_method=getattr(args, "launch_method", None),
         )
 
     @classmethod
-    def run(cls, conf, source_root, env_spec=None, force=False, bench=None, launch_method=None):
-        source_root = os.path.abspath(source_root)
-        if not os.path.isdir(source_root):
-            raise util.UserError(f"--source-root is not a directory: {source_root}")
-
-        env_spec = env_spec or ["existing:same"]
-        environments = list(get_environments(conf, env_spec))
-        if not environments:
-            raise util.UserError("No environments available")
-        conf.dvcs = "none"
-
-        try:
-            benchmarks = Benchmarks.load(conf, regex=bench)
-            log.info(f"Loaded {len(benchmarks)} benchmark(s) from benchmarks.json")
-        except util.UserError:
-            log.info("benchmarks.json not found — discovering benchmarks...")
-            benchmarks = Benchmarks.discover(conf, NoRepository(), environments, [None], regex=bench)
-            benchmarks.save()
-            log.info(f"Discovered and saved {len(benchmarks)} benchmark(s)")
+    def run(cls, conf, source_root, force=False, launch_method=None):
+        if launch_method:
+            conf.launch_method = launch_method
 
-        if not benchmarks:
-            log.error("No benchmarks found")
-            return 1
+        session = LightspeedSession._from_conf(conf)
+        result = session.initialize_diffcheck(source_root, force=force)
 
-        db_path = os.path.join(conf.results_dir, _DEPS_DB_FILENAME)
-        db = LightspeedDB(db_path)
-        all_bids = _get_all_bids(benchmarks)
-
-        if not force and all(db.has_baseline(bid) for bid in all_bids):
-            log.info("Baseline already exists for all benchmarks. Use --force to re-run.")
-            return 0
-
-        log.info("Pass 1: coverage survey")
-        with log.indent():
-            run_survey(conf.benchmark_dir, source_root, db, all_bids, verbose=True)
-
-        env = environments[0]
-        env.create()
-        lm = launch_method or getattr(conf, 'launch_method', None) or 'auto'
-        log.info("Pass 2: baseline timing")
-        results = run_benchmarks(benchmarks, env, launch_method=lm)
-
-        _store_baseline_from_results(results, benchmarks, db)
-
-        stored = sum(1 for bid in all_bids if db.has_baseline(bid))
-        log.info(f"Done. Baseline stored for {stored}/{len(all_bids)} benchmark(s) in {db_path}")
+        log.info(
+            f"Done. {len(result.benchmarks_discovered)} benchmark(s) discovered, "
+            f"{len(result.benchmarks_impactable)} impactable, "
+            f"{result.source_files_covered} source file(s) covered."
+        )
+        log.info(f"Dependency DB: {result.deps_db_path}")
+        if result.timing.phases:
+            log.info(
+                f"Timing: survey={result.timing.phases['coverage']:.1f}s  "
+                f"benchmarks={result.timing.phases['benchmarking']:.1f}s"
+            )
         return 0

asv/commands/measure_impacted.py

@@ -1,77 +1,27 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import argparse
-import os
-import subprocess
 
-from .. import util
-from ..benchmarks import Benchmarks
 from ..console import log
-from ..environment import get_environments
-from ..runner import run_benchmarks
-from ..contrib.lightspeed.deps_db import BenchmarkId, LightspeedDB
-from ..contrib.lightspeed.fingerprint import changed_files_with_fingerprints
+from ..contrib.lightspeed.session import LightspeedSession, MeasureResult
 from . import Command, common_args
 
-_DEPS_DB_FILENAME = ".lightspeed_deps.db"
 
-
-def _get_git_changed_files():
-    try:
-        out = subprocess.check_output(
-            ["git", "diff", "HEAD", "--name-only"], stderr=subprocess.DEVNULL,
-        ).decode().strip()
-        return [os.path.abspath(p) for p in out.splitlines()] if out else []
-    except subprocess.CalledProcessError as exc:
-        raise util.UserError(f"'git diff HEAD' failed: {exc}")
-
-
-def _extract_results(asv_results, benchmarks):
-    """Pull timing stats from a Results object into dict[BenchmarkId, dict]."""
-    out = {}
-    for name, benchmark in benchmarks.items():
-        result_vals = asv_results._results.get(name)
-        stats_list = asv_results._stats.get(name)
-        if result_vals is None or stats_list is None:
-            continue
-        if benchmark.get('params'):
-            for param_idx, (val, stat) in enumerate(zip(result_vals, stats_list)):
-                if val is not None and stat is not None:
-                    out[BenchmarkId(name, param_idx)] = {'median': val, **stat}
-        else:
-            val = result_vals[0] if result_vals else None
-            stat = stats_list[0] if stats_list else None
-            if val is not None and stat is not None:
-                out[BenchmarkId(name)] = {'median': val, **stat}
-    return out
-
-
-def _format_time(s):
-    if s is None:
-        return "n/a"
-    if s >= 1.0:
-        return f"{s:.3f}s"
-    if s >= 1e-3:
-        return f"{s * 1e3:.3f}ms"
-    if s >= 1e-6:
-        return f"{s * 1e6:.3f}us"
-    return f"{s * 1e9:.3f}ns"
-
-
-def _print_delta_table(deltas):
-    if not deltas:
+def _print_delta_table(result: MeasureResult):
+    if not result.benchmarks:
         return
-    col_w = max(len(str(bid)) for bid in deltas) + 2
+    col_w = max(len(name) for name in result.benchmarks) + 2
     header = f"{'benchmark':<{col_w}} {'baseline':>12} {'current':>12} {'delta':>10}"
     print(header)
     print("-" * len(header))
-    for bid, info in sorted(deltas.items(), key=lambda x: str(x[0])):
-        pct = info['delta'] * 100
+    for name, d in sorted(result.benchmarks.items()):
+        pct = d.delta_pct
+        pct_str = f"{'+' if pct >= 0 else ''}{pct:.1f}%" if pct is not None else "n/a"
         print(
-            f"{str(bid):<{col_w}} "
-            f"{_format_time(info['baseline_median']):>12} "
-            f"{_format_time(info['current_median']):>12} "
-            f"{'+' if pct >= 0 else ''}{pct:>8.1f}%"
+            f"{name:<{col_w}} "
+            f"{d.baseline_str:>12} "
+            f"{d.current_str:>12} "
+            f"{pct_str:>10}"
         )
 
 
@@ -88,8 +38,7 @@ def setup_arguments(cls, subparsers):
                 "run only those benchmarks, and compare against the stored baseline.\n\n"
                 "examples:\n"
                 "  asv measure_impacted --changed-files src/foo.py src/bar.py\n"
-                "  asv measure_impacted --from-git-diff\n"
-                "  asv measure_impacted --from-git-diff --factor 1.05"
+                "  asv measure_impacted --from-git-diff"
             ),
         )
         common_args.add_environment(parser, default_same=True)
@@ -107,10 +56,6 @@ def setup_arguments(cls, subparsers):
             "--step-id", default=None, metavar="ID",
             help="Optional step label passed to _on_step_results() hook.",
         )
-        parser.add_argument(
-            "--factor", type=float, default=1.1,
-            help="Exit non-zero if any benchmark regresses by more than this factor (default: 1.1).",
-        )
         common_args.add_bench(parser)
         common_args.add_launch_method(parser)
         parser.set_defaults(func=cls.run_from_args)
@@ -123,87 +68,39 @@ def run_from_conf_args(cls, conf, args):
             changed_files=args.changed_files,
             from_git_diff=args.from_git_diff,
             step_id=args.step_id,
-            factor=args.factor,
-            env_spec=args.env_spec,
-            bench=args.bench,
-            launch_method=getattr(args, 'launch_method', None),
+            launch_method=getattr(args, "launch_method", None),
         )
 
     @classmethod
     def run(
         cls, conf, changed_files=None, from_git_diff=False,
-        step_id=None, factor=1.1, env_spec=None, bench=None, launch_method=None,
+        step_id=None, launch_method=None,
     ):
-        env_spec = env_spec or ["existing:same"]
-        environments = list(get_environments(conf, env_spec))
-        if not environments:
-            raise util.UserError("No environments available")
-        conf.dvcs = "none"
+        if launch_method:
+            conf.launch_method = launch_method
 
-        benchmarks = Benchmarks.load(conf, regex=bench)
+        session = LightspeedSession._from_conf(conf)
+        result = session.measure_impacted(
+            from_git_diff=from_git_diff,
+            changed_files=changed_files,
+        )
 
-        db_path = os.path.join(conf.results_dir, _DEPS_DB_FILENAME)
-        if not os.path.exists(db_path):
-            raise util.UserError(
-                f"Dependency database not found: {db_path}\n"
-                "Run 'asv initialize_diffcheck --source-root <path>' first."
+        if not result.benchmarks:
+            log.info(
+                f"No benchmarks affected "
+                f"({result.total_count} total, 0 selected)."
             )
-        db = LightspeedDB(db_path)
-
-        baseline = {str(bid): b for bid in db.get_all_benchmark_ids() if (b := db.get_baseline(bid))}
-
-        if from_git_diff:
-            changed_files = _get_git_changed_files()
-        else:
-            changed_files = [os.path.abspath(p) for p in (changed_files or [])]
-
-        if not changed_files:
-            log.info("No changed files — nothing to do.")
-            return 0
-
-        changes = changed_files_with_fingerprints(changed_files, db.get_stored_fshas())
-        affected_bids = db.get_affected_benchmark_ids(changes)
-        if not affected_bids:
-            log.info("No benchmarks are affected by the changed files.")
             return 0
 
-        affected_names = {bid.name for bid in affected_bids}
-        filtered = benchmarks.filter_out(set(benchmarks.keys()) - affected_names)
-
-        log.info(f"Running {len(filtered)} affected benchmark(s)...")
-        env = environments[0]
-        env.create()
-        lm = launch_method or getattr(conf, 'launch_method', None) or 'auto'
-        results = run_benchmarks(filtered, env, launch_method=lm)
-
-        step_results = _extract_results(results, filtered)
-        deltas = {}
-        for bid, timing in step_results.items():
-            base = baseline.get(str(bid))
-            if base is None:
-                log.warning(f"No baseline for {bid} — skipping delta")
-                continue
-            delta = (timing['median'] - base['median']) / base['median']
-            deltas[bid] = {
-                'baseline_median': base['median'],
-                'current_median': timing['median'],
-                'delta': delta,
-                **timing,
-            }
-
-        cls._on_step_results(results, deltas, step_id)
-        _print_delta_table(deltas)
-
-        regressions = [bid for bid, d in deltas.items() if d['delta'] > (factor - 1)]
-        if regressions:
-            log.warning(f"{len(regressions)} benchmark(s) regressed by more than factor {factor}:")
-            with log.indent():
-                for bid in regressions:
-                    log.warning(str(bid))
-            return 1
+        log.info(
+            f"Ran {result.selected_count}/{result.total_count} benchmark(s) "
+            f"in {result.timing.total_s:.1f}s"
+        )
+        cls._on_step_results(result, step_id)
+        _print_delta_table(result)
         return 0
 
     @classmethod
-    def _on_step_results(cls, results, deltas, step_id):
+    def _on_step_results(cls, result: MeasureResult, step_id):
         """No-op hook. Override in a fork subclass to persist step results."""
-        pass
+        pass

asv/contrib/lightspeed/__init__.py

@@ -0,0 +1,11 @@
+from .session import (  # noqa: F401
+    LightspeedSession,
+    InitResult,
+    MeasureResult,
+    BenchmarkDelta,
+    TimingInfo,
+    ASVError,
+    ConfigError,
+    BenchmarkError,
+    NoBenchmarksError,
+)