shotgun-sh 0.4.0.dev1__py3-none-any.whl → 0.6.2__py3-none-any.whl

shotgun/codebase/benchmarks/benchmark_runner.py

@@ -0,0 +1,309 @@
+"""Benchmark runner for codebase indexing performance analysis.
+
+This module provides the BenchmarkRunner class for running benchmark iterations
+and collecting performance statistics.
+"""
+
+from __future__ import annotations
+
+import gc
+import hashlib
+import json
+import shutil
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from shotgun.codebase.benchmarks.models import (
+    BenchmarkConfig,
+    BenchmarkMode,
+    BenchmarkResults,
+    BenchmarkRun,
+)
+from shotgun.codebase.core import Ingestor, SimpleGraphBuilder
+from shotgun.codebase.core.kuzu_compat import get_kuzu
+from shotgun.codebase.core.metrics_collector import MetricsCollector
+from shotgun.codebase.core.parser_loader import load_parsers
+from shotgun.logging_config import get_logger
+from shotgun.sdk.services import get_codebase_service
+from shotgun.utils.file_system_utils import get_shotgun_home
+
+logger = get_logger(__name__)
+
+
+def _compute_graph_id(codebase_path: Path) -> str:
+    """Compute a unique graph ID from the codebase path.
+
+    Args:
+        codebase_path: Path to the codebase
+
+    Returns:
+        A 12-character hex string identifying this codebase
+    """
+    return hashlib.sha256(str(codebase_path).encode()).hexdigest()[:12]
+
+
+class BenchmarkRunner:
+    """Runs benchmark iterations and collects statistics."""
+
+    def __init__(
+        self,
+        codebase_path: Path,
+        codebase_name: str,
+        iterations: int = 1,
+        warmup_iterations: int = 0,
+        parallel: bool = True,
+        worker_count: int | None = None,
+        collect_file_metrics: bool = True,
+        collect_worker_metrics: bool = True,
+        progress_callback: Callable[..., Any] | None = None,
+    ) -> None:
+        """Initialize benchmark runner.
+
+        Args:
+            codebase_path: Path to repository to benchmark
+            codebase_name: Human-readable name for the codebase
+            iterations: Number of measured benchmark runs
+            warmup_iterations: Number of warmup runs (not measured)
+            parallel: Whether to use parallel execution
+            worker_count: Number of workers (None = auto)
+            collect_file_metrics: Whether to collect per-file metrics
+            collect_worker_metrics: Whether to collect per-worker metrics
+            progress_callback: Optional callback for progress updates
+        """
+        self.codebase_path = codebase_path.resolve()
+        self.codebase_name = codebase_name
+        self.iterations = iterations
+        self.warmup_iterations = warmup_iterations
+        self.parallel = parallel
+        self.worker_count = worker_count
+        self.collect_file_metrics = collect_file_metrics
+        self.collect_worker_metrics = collect_worker_metrics
+        self.progress_callback = progress_callback
+
+        # Configuration object
+        self.config = BenchmarkConfig(
+            mode=BenchmarkMode.PARALLEL if parallel else BenchmarkMode.SEQUENTIAL,
+            worker_count=worker_count,
+            iterations=iterations,
+            warmup_iterations=warmup_iterations,
+            collect_file_metrics=collect_file_metrics,
+            collect_worker_metrics=collect_worker_metrics,
+        )
+
+        # Storage for database operations
+        self._storage_dir = get_shotgun_home() / "codebases"
+        self._service = get_codebase_service(self._storage_dir)
+
+    async def run(self) -> BenchmarkResults:
+        """Run all benchmark iterations and return aggregated results.
+
+        Returns:
+            BenchmarkResults with all run data and statistics
+        """
+        results = BenchmarkResults(
+            codebase_name=self.codebase_name,
+            codebase_path=str(self.codebase_path),
+            config=self.config,
+        )
+
+        # Run warmup iterations
+        for i in range(self.warmup_iterations):
+            logger.info(f"Running warmup iteration {i + 1}/{self.warmup_iterations}...")
+            if self.progress_callback:
+                self.progress_callback(
+                    f"Warmup {i + 1}/{self.warmup_iterations}", None, None
+                )
+
+            run = await self._run_single_iteration(
+                run_id=i,
+                is_warmup=True,
+            )
+            results.add_run(run)
+            await self._cleanup_database()
+
+        # Run measured iterations
+        for i in range(self.iterations):
+            logger.info(f"Running benchmark iteration {i + 1}/{self.iterations}...")
+            if self.progress_callback:
+                self.progress_callback(
+                    f"Benchmark {i + 1}/{self.iterations}", None, None
+                )
+
+            run = await self._run_single_iteration(
+                run_id=i,
+                is_warmup=False,
+            )
+            results.add_run(run)
+
+            # Clean up between iterations (but not after the last one)
+            if i < self.iterations - 1:
+                await self._cleanup_database()
+
+        # Register the codebase so it persists after benchmark
+        await self._register_codebase()
+
+        # Calculate statistics
+        results.calculate_statistics()
+
+        logger.info(
+            f"Benchmark complete: {self.iterations} iterations, "
+            f"avg {results.avg_duration_seconds:.2f}s"
+        )
+
+        return results
+
+    async def _run_single_iteration(
+        self,
+        run_id: int,
+        is_warmup: bool,
+    ) -> BenchmarkRun:
+        """Run a single benchmark iteration.
+
+        Args:
+            run_id: Run number
+            is_warmup: Whether this is a warmup run
+
+        Returns:
+            BenchmarkRun with collected metrics
+        """
+        # Create metrics collector
+        metrics_collector = MetricsCollector(
+            codebase_name=self.codebase_name,
+            collect_file_metrics=self.collect_file_metrics,
+            collect_worker_metrics=self.collect_worker_metrics,
+        )
+
+        # Generate unique graph ID for this run
+        graph_id = _compute_graph_id(self.codebase_path)
+
+        # Create database
+        kuzu = get_kuzu()
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+
+        # Ensure clean state
+        if graph_path.exists():
+            if graph_path.is_dir():
+                shutil.rmtree(graph_path)
+            else:
+                graph_path.unlink()
+        wal_path = self._storage_dir / f"{graph_id}.kuzu.wal"
+        if wal_path.exists():
+            wal_path.unlink()
+
+        # Create database and connection
+        db = kuzu.Database(str(graph_path))
+        conn = kuzu.Connection(db)
+
+        # Load parsers
+        parsers, queries = load_parsers()
+
+        # Create ingestor and builder
+        ingestor = Ingestor(conn)
+        ingestor.create_schema()
+
+        builder = SimpleGraphBuilder(
+            ingestor=ingestor,
+            repo_path=self.codebase_path,
+            parsers=parsers,
+            queries=queries,
+            metrics_collector=metrics_collector,
+            enable_parallel=self.parallel,
+            progress_callback=None,  # Disable TUI progress in benchmark mode
+        )
+
+        # Run indexing
+        await builder.run()
+
+        # Get metrics
+        metrics = metrics_collector.get_metrics()
+
+        # Close connection
+        del conn
+        del db
+
+        return BenchmarkRun(
+            run_id=run_id,
+            is_warmup=is_warmup,
+            metrics=metrics,
+        )
+
+    async def _cleanup_database(self) -> None:
+        """Delete database files and clear caches between runs."""
+        graph_id = _compute_graph_id(self.codebase_path)
+
+        # Delete database file
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+        if graph_path.exists():
+            if graph_path.is_dir():
+                shutil.rmtree(graph_path)
+            else:
+                graph_path.unlink()
+            logger.debug(f"Deleted database: {graph_path}")
+
+        # Delete WAL file
+        wal_path = self._storage_dir / f"{graph_id}.kuzu.wal"
+        if wal_path.exists():
+            wal_path.unlink()
+            logger.debug(f"Deleted WAL: {wal_path}")
+
+        # Force garbage collection
+        gc.collect()
+
+    async def _register_codebase(self) -> None:
+        """Register the codebase so it appears in `shotgun codebase list`.
+
+        This creates a Project node in the database with metadata that
+        identifies the indexed codebase.
+        """
+        graph_id = _compute_graph_id(self.codebase_path)
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+
+        if not graph_path.exists():
+            logger.warning("Cannot register codebase: database not found")
+            return
+
+        kuzu = get_kuzu()
+        db = kuzu.Database(str(graph_path))
+        conn = kuzu.Connection(db)
+
+        try:
+            # Check if Project node already exists
+            result = conn.execute("MATCH (p:Project) RETURN p.graph_id LIMIT 1")
+            if result.has_next():
+                logger.debug("Project node already exists, skipping registration")
+                return
+
+            # Create Project node with metadata
+            current_time = int(time.time())
+            conn.execute(
+                """
+                CREATE (p:Project {
+                    name: $name,
+                    repo_path: $repo_path,
+                    graph_id: $graph_id,
+                    created_at: $created_at,
+                    updated_at: $updated_at,
+                    schema_version: $schema_version,
+                    build_options: $build_options,
+                    indexed_from_cwds: $indexed_from_cwds
+                })
+                """,
+                {
+                    "name": self.codebase_name,
+                    "repo_path": str(self.codebase_path),
+                    "graph_id": graph_id,
+                    "created_at": current_time,
+                    "updated_at": current_time,
+                    "schema_version": "1.0.0",
+                    "build_options": json.dumps({}),
+                    "indexed_from_cwds": json.dumps([str(Path.cwd())]),
+                },
+            )
+            logger.info(
+                f"Registered codebase '{self.codebase_name}' with graph_id: {graph_id}"
+            )
+        finally:
+            del conn
+            del db

shotgun/codebase/benchmarks/exporters.py

@@ -0,0 +1,119 @@
+"""Metrics exporters for saving benchmark results to files.
+
+This module provides the MetricsExporter class for exporting benchmark
+results to various file formats.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from shotgun.codebase.benchmarks.formatters import (
+    JsonFormatter,
+    MarkdownFormatter,
+    MetricsDisplayOptions,
+)
+from shotgun.logging_config import get_logger
+
+if TYPE_CHECKING:
+    from shotgun.codebase.benchmarks.models import BenchmarkResults
+
+logger = get_logger(__name__)
+
+
+class MetricsExporter:
+    """Export benchmark metrics to files."""
+
+    def __init__(self) -> None:
+        """Initialize metrics exporter."""
+        self._format_map = {
+            ".json": self._export_json,
+            ".md": self._export_markdown,
+            ".markdown": self._export_markdown,
+        }
+
+    def export(
+        self,
+        results: BenchmarkResults,
+        filepath: Path | str,
+        format: str | None = None,
+        options: MetricsDisplayOptions | None = None,
+    ) -> None:
+        """Export benchmark results to file.
+
+        The format is auto-detected from the file extension if not specified.
+
+        Args:
+            results: Benchmark results to export
+            filepath: Path to export file
+            format: Optional format override ("json", "markdown")
+            options: Display options for controlling what to include
+
+        Raises:
+            ValueError: If format cannot be determined or is unsupported
+            OSError: If file cannot be written
+        """
+        filepath = Path(filepath)
+        options = options or MetricsDisplayOptions()
+
+        # Determine format
+        if format:
+            format_lower = format.lower()
+            if format_lower == "json":
+                export_func = self._export_json
+            elif format_lower in ("markdown", "md"):
+                export_func = self._export_markdown
+            else:
+                raise ValueError(f"Unknown export format: {format}")
+        else:
+            # Auto-detect from extension
+            suffix = filepath.suffix.lower()
+            if suffix not in self._format_map:
+                raise ValueError(
+                    f"Cannot determine format from extension '{suffix}'. "
+                    f"Supported extensions: {', '.join(self._format_map.keys())}. "
+                    f"Or specify format explicitly."
+                )
+            export_func = self._format_map[suffix]
+
+        # Ensure parent directory exists
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        # Export
+        export_func(results, filepath, options)
+        logger.info(f"Exported benchmark results to {filepath}")
+
+    def _export_json(
+        self,
+        results: BenchmarkResults,
+        filepath: Path,
+        options: MetricsDisplayOptions,
+    ) -> None:
+        """Export results to JSON file.
+
+        Args:
+            results: Benchmark results
+            filepath: Output path
+            options: Display options
+        """
+        formatter = JsonFormatter()
+        content = formatter.format_results(results, options)
+        filepath.write_text(content)
+
+    def _export_markdown(
+        self,
+        results: BenchmarkResults,
+        filepath: Path,
+        options: MetricsDisplayOptions,
+    ) -> None:
+        """Export results to Markdown file.
+
+        Args:
+            results: Benchmark results
+            filepath: Output path
+            options: Display options
+        """
+        formatter = MarkdownFormatter()
+        content = formatter.format_results(results, options)
+        filepath.write_text(content)

shotgun/codebase/benchmarks/formatters/__init__.py

@@ -0,0 +1,49 @@
+"""Result formatters for benchmark output.
+
+This package provides formatters for displaying benchmark results in various
+formats: JSON and Markdown.
+"""
+
+from shotgun.codebase.benchmarks.formatters.base import ResultFormatter
+from shotgun.codebase.benchmarks.formatters.json_formatter import JsonFormatter
+from shotgun.codebase.benchmarks.formatters.markdown import MarkdownFormatter
+
+# Re-export MetricsDisplayOptions from models for convenience
+from shotgun.codebase.benchmarks.models import MetricsDisplayOptions
+
+__all__ = [
+    "JsonFormatter",
+    "MarkdownFormatter",
+    "MetricsDisplayOptions",
+    "ResultFormatter",
+    "get_formatter",
+]
+
+
+def get_formatter(
+    output_format: str,
+) -> JsonFormatter | MarkdownFormatter:
+    """Get appropriate formatter for output format.
+
+    Args:
+        output_format: Format name - "json" or "markdown"
+
+    Returns:
+        Formatter instance
+
+    Raises:
+        ValueError: If output format is unknown
+    """
+    formatters: dict[str, type[JsonFormatter | MarkdownFormatter]] = {
+        "json": JsonFormatter,
+        "markdown": MarkdownFormatter,
+    }
+
+    format_lower = output_format.lower()
+    if format_lower not in formatters:
+        raise ValueError(
+            f"Unknown output format: {output_format}. "
+            f"Valid formats: {', '.join(formatters.keys())}"
+        )
+
+    return formatters[format_lower]()

shotgun/codebase/benchmarks/formatters/base.py

@@ -0,0 +1,34 @@
+"""Base classes and protocols for formatters.
+
+This module provides the Protocol interface that all formatters implement.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from shotgun.codebase.benchmarks.models import (
+        BenchmarkResults,
+        MetricsDisplayOptions,
+    )
+
+
+class ResultFormatter(Protocol):
+    """Protocol for formatting benchmark results."""
+
+    def format_results(
+        self,
+        results: BenchmarkResults,
+        options: MetricsDisplayOptions,
+    ) -> str:
+        """Format benchmark results for display.
+
+        Args:
+            results: Benchmark results to format
+            options: Display options
+
+        Returns:
+            Formatted string ready for output
+        """
+        ...

shotgun/codebase/benchmarks/formatters/json_formatter.py

@@ -0,0 +1,106 @@
+"""JSON formatter for benchmark results.
+
+This module provides the JsonFormatter class for displaying benchmark results
+as JSON.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from shotgun.codebase.benchmarks.models import (
+        BenchmarkResults,
+        MetricsDisplayOptions,
+    )
+
+
+class JsonFormatter:
+    """Format benchmark results as JSON."""
+
+    def format_results(
+        self,
+        results: BenchmarkResults,
+        options: MetricsDisplayOptions,
+    ) -> str:
+        """Format benchmark results as JSON.
+
+        Args:
+            results: Benchmark results to format
+            options: Display options
+
+        Returns:
+            JSON string
+        """
+        data = {
+            "codebase_name": results.codebase_name,
+            "codebase_path": results.codebase_path,
+            "config": {
+                "mode": results.config.mode,
+                "worker_count": results.config.worker_count,
+                "iterations": results.config.iterations,
+                "warmup_iterations": results.config.warmup_iterations,
+            },
+            "statistics": {
+                "avg_duration_seconds": results.avg_duration_seconds,
+                "min_duration_seconds": results.min_duration_seconds,
+                "max_duration_seconds": results.max_duration_seconds,
+                "std_dev_seconds": results.std_dev_seconds,
+                "avg_throughput": results.avg_throughput,
+                "avg_memory_mb": results.avg_memory_mb,
+                "speedup_factor": results.speedup_factor,
+                "efficiency": results.efficiency,
+            },
+            "runs": [],
+        }
+
+        # Add run data
+        for run in results.measured_runs:
+            run_data: dict[str, object] = {
+                "run_id": run.run_id,
+                "duration_seconds": run.metrics.total_duration_seconds,
+                "total_files": run.metrics.total_files,
+                "total_nodes": run.metrics.total_nodes,
+                "total_relationships": run.metrics.total_relationships,
+                "throughput": run.metrics.avg_throughput,
+                "peak_memory_mb": run.metrics.peak_memory_mb,
+            }
+
+            # Add phase metrics
+            if options.show_phase_metrics:
+                phase_data: dict[str, dict[str, object]] = {}
+                for name, phase in run.metrics.phase_metrics.items():
+                    phase_data[name] = {
+                        "duration_seconds": phase.duration_seconds,
+                        "items_processed": phase.items_processed,
+                        "throughput": phase.throughput,
+                        "memory_mb": phase.memory_mb,
+                    }
+                run_data["phase_metrics"] = phase_data
+
+            # Add file metrics
+            if options.show_file_metrics and run.metrics.file_metrics:
+                file_metrics_list = run.metrics.file_metrics
+                if options.top_n_files:
+                    file_metrics_list = sorted(
+                        file_metrics_list,
+                        key=lambda f: f.parse_time_ms,
+                        reverse=True,
+                    )[: options.top_n_files]
+
+                run_data["file_metrics"] = [
+                    {
+                        "file_path": f.file_path,
+                        "language": f.language,
+                        "file_size_bytes": f.file_size_bytes,
+                        "parse_time_ms": f.parse_time_ms,
+                        "definitions_extracted": f.definitions_extracted,
+                    }
+                    for f in file_metrics_list
+                ]
+
+            runs_list: list[dict[str, object]] = data["runs"]  # type: ignore[assignment]
+            runs_list.append(run_data)
+
+        return json.dumps(data, indent=2)