truthound 1.0.8__py3-none-any.whl

truthound/benchmark/base.py

@@ -0,0 +1,757 @@
+"""Base classes and protocols for the benchmarking system.
+
+This module provides the foundational abstractions for:
+- Benchmark definition and execution
+- Result collection and metrics
+- Registry for benchmark discovery
+"""
+
+from __future__ import annotations
+
+import gc
+import hashlib
+import platform
+import statistics
+import threading
+import time
+import traceback
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Protocol,
+    TypeVar,
+    runtime_checkable,
+)
+
+import polars as pl
+
+
+# =============================================================================
+# Types and Enums
+# =============================================================================
+
+
+class BenchmarkCategory(str, Enum):
+    """Categories for benchmark classification."""
+
+    PROFILING = "profiling"
+    VALIDATION = "validation"
+    DRIFT = "drift"
+    LEARNING = "learning"
+    SCANNING = "scanning"
+    E2E = "e2e"
+    MEMORY = "memory"
+    THROUGHPUT = "throughput"
+
+
+class BenchmarkSize(str, Enum):
+    """Standard benchmark sizes for consistent comparison."""
+
+    TINY = "tiny"           # 1K rows
+    SMALL = "small"         # 10K rows
+    MEDIUM = "medium"       # 100K rows
+    LARGE = "large"         # 1M rows
+    XLARGE = "xlarge"       # 10M rows
+    STRESS = "stress"       # 100M rows
+
+    @property
+    def row_count(self) -> int:
+        """Get the row count for this size."""
+        sizes = {
+            "tiny": 1_000,
+            "small": 10_000,
+            "medium": 100_000,
+            "large": 1_000_000,
+            "xlarge": 10_000_000,
+            "stress": 100_000_000,
+        }
+        return sizes[self.value]
+
+
+class MetricUnit(str, Enum):
+    """Units for benchmark metrics."""
+
+    SECONDS = "seconds"
+    MILLISECONDS = "milliseconds"
+    MICROSECONDS = "microseconds"
+    BYTES = "bytes"
+    MEGABYTES = "megabytes"
+    ROWS_PER_SECOND = "rows/s"
+    MB_PER_SECOND = "MB/s"
+    OPERATIONS_PER_SECOND = "ops/s"
+    COUNT = "count"
+    RATIO = "ratio"
+    PERCENT = "percent"
+
+
+# =============================================================================
+# Metrics and Results
+# =============================================================================
+
+
+@dataclass
+class MetricValue:
+    """A single metric measurement."""
+
+    name: str
+    value: float
+    unit: MetricUnit
+    is_lower_better: bool = True
+    description: str = ""
+    tags: dict[str, str] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "name": self.name,
+            "value": self.value,
+            "unit": self.unit.value,
+            "is_lower_better": self.is_lower_better,
+            "description": self.description,
+            "tags": self.tags,
+        }
+
+
+@dataclass
+class BenchmarkMetrics:
+    """Collection of metrics from a benchmark run.
+
+    Provides statistical analysis of timing measurements.
+    """
+
+    # Timing metrics (in seconds)
+    durations: list[float] = field(default_factory=list)
+
+    # Memory metrics (in bytes)
+    peak_memory_bytes: int = 0
+    memory_delta_bytes: int = 0
+
+    # Throughput metrics
+    rows_processed: int = 0
+    bytes_processed: int = 0
+
+    # Custom metrics
+    custom_metrics: list[MetricValue] = field(default_factory=list)
+
+    @property
+    def mean_duration(self) -> float:
+        """Mean execution time in seconds."""
+        if not self.durations:
+            return 0.0
+        return statistics.mean(self.durations)
+
+    @property
+    def median_duration(self) -> float:
+        """Median execution time in seconds."""
+        if not self.durations:
+            return 0.0
+        return statistics.median(self.durations)
+
+    @property
+    def std_duration(self) -> float:
+        """Standard deviation of execution time."""
+        if len(self.durations) < 2:
+            return 0.0
+        return statistics.stdev(self.durations)
+
+    @property
+    def min_duration(self) -> float:
+        """Minimum execution time."""
+        if not self.durations:
+            return 0.0
+        return min(self.durations)
+
+    @property
+    def max_duration(self) -> float:
+        """Maximum execution time."""
+        if not self.durations:
+            return 0.0
+        return max(self.durations)
+
+    @property
+    def p95_duration(self) -> float:
+        """95th percentile execution time."""
+        if not self.durations:
+            return 0.0
+        sorted_durations = sorted(self.durations)
+        idx = int(len(sorted_durations) * 0.95)
+        return sorted_durations[min(idx, len(sorted_durations) - 1)]
+
+    @property
+    def p99_duration(self) -> float:
+        """99th percentile execution time."""
+        if not self.durations:
+            return 0.0
+        sorted_durations = sorted(self.durations)
+        idx = int(len(sorted_durations) * 0.99)
+        return sorted_durations[min(idx, len(sorted_durations) - 1)]
+
+    @property
+    def rows_per_second(self) -> float:
+        """Throughput in rows per second."""
+        if not self.mean_duration or not self.rows_processed:
+            return 0.0
+        return self.rows_processed / self.mean_duration
+
+    @property
+    def mb_per_second(self) -> float:
+        """Throughput in megabytes per second."""
+        if not self.mean_duration or not self.bytes_processed:
+            return 0.0
+        return (self.bytes_processed / (1024 * 1024)) / self.mean_duration
+
+    def add_duration(self, duration: float) -> None:
+        """Add a duration measurement."""
+        self.durations.append(duration)
+
+    def add_custom_metric(self, metric: MetricValue) -> None:
+        """Add a custom metric."""
+        self.custom_metrics.append(metric)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "timing": {
+                "mean_seconds": self.mean_duration,
+                "median_seconds": self.median_duration,
+                "std_seconds": self.std_duration,
+                "min_seconds": self.min_duration,
+                "max_seconds": self.max_duration,
+                "p95_seconds": self.p95_duration,
+                "p99_seconds": self.p99_duration,
+                "iterations": len(self.durations),
+            },
+            "memory": {
+                "peak_bytes": self.peak_memory_bytes,
+                "peak_mb": self.peak_memory_bytes / (1024 * 1024),
+                "delta_bytes": self.memory_delta_bytes,
+            },
+            "throughput": {
+                "rows_processed": self.rows_processed,
+                "bytes_processed": self.bytes_processed,
+                "rows_per_second": self.rows_per_second,
+                "mb_per_second": self.mb_per_second,
+            },
+            "custom": [m.to_dict() for m in self.custom_metrics],
+        }
+
+
+@dataclass
+class EnvironmentInfo:
+    """Information about the execution environment."""
+
+    python_version: str = ""
+    platform_system: str = ""
+    platform_release: str = ""
+    platform_machine: str = ""
+    cpu_count: int = 0
+    polars_version: str = ""
+    truthound_version: str = ""
+    timestamp: datetime = field(default_factory=datetime.now)
+
+    @classmethod
+    def capture(cls) -> "EnvironmentInfo":
+        """Capture current environment information."""
+        import truthound
+
+        return cls(
+            python_version=platform.python_version(),
+            platform_system=platform.system(),
+            platform_release=platform.release(),
+            platform_machine=platform.machine(),
+            cpu_count=platform.os.cpu_count() or 0,
+            polars_version=pl.__version__,
+            truthound_version=getattr(truthound, "__version__", "unknown"),
+            timestamp=datetime.now(),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "python_version": self.python_version,
+            "platform": {
+                "system": self.platform_system,
+                "release": self.platform_release,
+                "machine": self.platform_machine,
+            },
+            "cpu_count": self.cpu_count,
+            "polars_version": self.polars_version,
+            "truthound_version": self.truthound_version,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+
+@dataclass
+class BenchmarkResult:
+    """Result of a single benchmark execution.
+
+    Contains all metrics, environment info, and execution details.
+    """
+
+    benchmark_name: str
+    category: BenchmarkCategory
+    success: bool
+    metrics: BenchmarkMetrics
+    environment: EnvironmentInfo
+    parameters: dict[str, Any] = field(default_factory=dict)
+    error: str | None = None
+    error_traceback: str | None = None
+    started_at: datetime = field(default_factory=datetime.now)
+    completed_at: datetime | None = None
+    tags: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def duration_seconds(self) -> float:
+        """Total benchmark duration including setup/teardown."""
+        if self.completed_at is None:
+            return 0.0
+        return (self.completed_at - self.started_at).total_seconds()
+
+    @property
+    def result_id(self) -> str:
+        """Unique identifier for this result."""
+        content = f"{self.benchmark_name}:{self.started_at.isoformat()}"
+        return hashlib.sha256(content.encode()).hexdigest()[:12]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "id": self.result_id,
+            "benchmark_name": self.benchmark_name,
+            "category": self.category.value,
+            "success": self.success,
+            "metrics": self.metrics.to_dict(),
+            "environment": self.environment.to_dict(),
+            "parameters": self.parameters,
+            "error": self.error,
+            "started_at": self.started_at.isoformat(),
+            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
+            "duration_seconds": self.duration_seconds,
+            "tags": self.tags,
+        }
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+
+@dataclass
+class BenchmarkConfig:
+    """Configuration for benchmark execution.
+
+    Controls warmup, iterations, timeouts, and other execution parameters.
+    """
+
+    # Iteration control
+    warmup_iterations: int = 2
+    measure_iterations: int = 5
+    cooldown_seconds: float = 0.1
+
+    # Timeout and limits
+    timeout_seconds: float = 300.0
+    max_memory_mb: float = 0  # 0 = unlimited
+
+    # Data parameters
+    default_size: BenchmarkSize = BenchmarkSize.MEDIUM
+    gc_between_iterations: bool = True
+
+    # Output control
+    verbose: bool = False
+    capture_traceback: bool = True
+
+    # Tags for result filtering
+    tags: dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def quick(cls) -> "BenchmarkConfig":
+        """Quick benchmark configuration for fast feedback."""
+        return cls(
+            warmup_iterations=1,
+            measure_iterations=3,
+            cooldown_seconds=0.05,
+            default_size=BenchmarkSize.SMALL,
+        )
+
+    @classmethod
+    def standard(cls) -> "BenchmarkConfig":
+        """Standard benchmark configuration."""
+        return cls(
+            warmup_iterations=2,
+            measure_iterations=5,
+            cooldown_seconds=0.1,
+            default_size=BenchmarkSize.MEDIUM,
+        )
+
+    @classmethod
+    def thorough(cls) -> "BenchmarkConfig":
+        """Thorough benchmark configuration for CI/CD."""
+        return cls(
+            warmup_iterations=3,
+            measure_iterations=10,
+            cooldown_seconds=0.2,
+            default_size=BenchmarkSize.LARGE,
+            gc_between_iterations=True,
+        )
+
+    @classmethod
+    def stress(cls) -> "BenchmarkConfig":
+        """Stress test configuration."""
+        return cls(
+            warmup_iterations=1,
+            measure_iterations=3,
+            timeout_seconds=600.0,
+            default_size=BenchmarkSize.XLARGE,
+            gc_between_iterations=True,
+        )
+
+
+# =============================================================================
+# Benchmark Protocol and Base Class
+# =============================================================================
+
+
+@runtime_checkable
+class BenchmarkProtocol(Protocol):
+    """Protocol defining the benchmark interface."""
+
+    name: str
+    category: BenchmarkCategory
+    description: str
+
+    def setup(self, config: BenchmarkConfig, **kwargs: Any) -> None:
+        """Prepare benchmark resources."""
+        ...
+
+    def run_iteration(self) -> Any:
+        """Run a single benchmark iteration."""
+        ...
+
+    def teardown(self) -> None:
+        """Clean up benchmark resources."""
+        ...
+
+    def get_metrics(self) -> BenchmarkMetrics:
+        """Get collected metrics."""
+        ...
+
+
+class Benchmark(ABC):
+    """Abstract base class for benchmarks.
+
+    Provides the framework for defining benchmarks with:
+    - Setup and teardown phases
+    - Configurable warmup and measurement iterations
+    - Automatic metric collection
+    - Memory tracking
+
+    Example:
+        class MyBenchmark(Benchmark):
+            name = "my_benchmark"
+            category = BenchmarkCategory.PROFILING
+            description = "Benchmark for my operation"
+
+            def setup(self, config, **kwargs):
+                self.data = generate_data(kwargs.get("rows", 1000))
+
+            def run_iteration(self):
+                return process(self.data)
+
+            def teardown(self):
+                del self.data
+    """
+
+    name: str = "base"
+    category: BenchmarkCategory = BenchmarkCategory.E2E
+    description: str = ""
+
+    def __init__(self):
+        self._metrics = BenchmarkMetrics()
+        self._config: BenchmarkConfig | None = None
+        self._parameters: dict[str, Any] = {}
+        self._lock = threading.Lock()
+
+    def setup(self, config: BenchmarkConfig, **kwargs: Any) -> None:
+        """Prepare benchmark resources.
+
+        Override this method to set up any data or resources needed.
+
+        Args:
+            config: Benchmark configuration
+            **kwargs: Benchmark-specific parameters
+        """
+        self._config = config
+        self._parameters = kwargs
+
+    @abstractmethod
+    def run_iteration(self) -> Any:
+        """Run a single benchmark iteration.
+
+        This is the core operation being measured. It should be
+        deterministic and isolated from setup/teardown.
+
+        Returns:
+            Result of the iteration (for validation)
+        """
+        pass
+
+    def teardown(self) -> None:
+        """Clean up benchmark resources.
+
+        Override this method to release any resources.
+        """
+        pass
+
+    def get_metrics(self) -> BenchmarkMetrics:
+        """Get the collected metrics."""
+        return self._metrics
+
+    def validate_result(self, result: Any) -> bool:
+        """Validate the result of an iteration.
+
+        Override to add custom validation logic.
+
+        Args:
+            result: Result from run_iteration()
+
+        Returns:
+            True if result is valid
+        """
+        return result is not None
+
+    def execute(
+        self,
+        config: BenchmarkConfig | None = None,
+        **kwargs: Any,
+    ) -> BenchmarkResult:
+        """Execute the benchmark with full lifecycle.
+
+        Args:
+            config: Benchmark configuration (uses default if None)
+            **kwargs: Parameters passed to setup()
+
+        Returns:
+            BenchmarkResult with all metrics and metadata
+        """
+        config = config or BenchmarkConfig.standard()
+        environment = EnvironmentInfo.capture()
+        started_at = datetime.now()
+
+        try:
+            # Setup phase
+            self.setup(config, **kwargs)
+
+            # Track memory if psutil available
+            memory_before = self._get_memory_usage()
+
+            # Warmup iterations
+            for _ in range(config.warmup_iterations):
+                self.run_iteration()
+                if config.gc_between_iterations:
+                    gc.collect()
+                time.sleep(config.cooldown_seconds)
+
+            # Measurement iterations
+            self._metrics = BenchmarkMetrics()
+            peak_memory = memory_before
+
+            for _ in range(config.measure_iterations):
+                start_time = time.perf_counter()
+                result = self.run_iteration()
+                duration = time.perf_counter() - start_time
+
+                self._metrics.add_duration(duration)
+
+                # Validate result
+                if not self.validate_result(result):
+                    raise ValueError("Benchmark result validation failed")
+
+                # Track peak memory
+                current_memory = self._get_memory_usage()
+                peak_memory = max(peak_memory, current_memory)
+
+                if config.gc_between_iterations:
+                    gc.collect()
+                time.sleep(config.cooldown_seconds)
+
+            # Record memory metrics
+            memory_after = self._get_memory_usage()
+            self._metrics.peak_memory_bytes = peak_memory
+            self._metrics.memory_delta_bytes = memory_after - memory_before
+
+            return BenchmarkResult(
+                benchmark_name=self.name,
+                category=self.category,
+                success=True,
+                metrics=self._metrics,
+                environment=environment,
+                parameters=kwargs,
+                started_at=started_at,
+                completed_at=datetime.now(),
+                tags=config.tags,
+            )
+
+        except Exception as e:
+            return BenchmarkResult(
+                benchmark_name=self.name,
+                category=self.category,
+                success=False,
+                metrics=self._metrics,
+                environment=environment,
+                parameters=kwargs,
+                error=str(e),
+                error_traceback=traceback.format_exc() if config.capture_traceback else None,
+                started_at=started_at,
+                completed_at=datetime.now(),
+                tags=config.tags,
+            )
+
+        finally:
+            self.teardown()
+
+    def _get_memory_usage(self) -> int:
+        """Get current memory usage in bytes."""
+        try:
+            import psutil
+            return psutil.Process().memory_info().rss
+        except ImportError:
+            return 0
+
+
+# =============================================================================
+# Registry
+# =============================================================================
+
+
+class BenchmarkRegistry:
+    """Registry for benchmark discovery and management.
+
+    Allows benchmarks to be registered and discovered by name or category.
+
+    Example:
+        registry = BenchmarkRegistry()
+
+        @registry.register
+        class MyBenchmark(Benchmark):
+            name = "my_benchmark"
+            ...
+
+        # Get benchmark by name
+        benchmark = registry.get("my_benchmark")
+
+        # Get all benchmarks in category
+        benchmarks = registry.get_by_category(BenchmarkCategory.PROFILING)
+    """
+
+    def __init__(self) -> None:
+        self._benchmarks: dict[str, type[Benchmark]] = {}
+        self._lock = threading.Lock()
+
+    def register(
+        self,
+        benchmark_class: type[Benchmark],
+    ) -> type[Benchmark]:
+        """Register a benchmark class.
+
+        Can be used as a decorator.
+
+        Args:
+            benchmark_class: Benchmark class to register
+
+        Returns:
+            The registered class (for decorator use)
+        """
+        with self._lock:
+            name = benchmark_class.name
+            if name in self._benchmarks:
+                raise ValueError(f"Benchmark '{name}' is already registered")
+            self._benchmarks[name] = benchmark_class
+        return benchmark_class
+
+    def unregister(self, name: str) -> None:
+        """Unregister a benchmark by name."""
+        with self._lock:
+            if name in self._benchmarks:
+                del self._benchmarks[name]
+
+    def get(self, name: str) -> type[Benchmark]:
+        """Get a benchmark class by name.
+
+        Args:
+            name: Benchmark name
+
+        Returns:
+            Benchmark class
+
+        Raises:
+            KeyError: If benchmark not found
+        """
+        with self._lock:
+            if name not in self._benchmarks:
+                available = list(self._benchmarks.keys())
+                raise KeyError(
+                    f"Benchmark '{name}' not found. Available: {available}"
+                )
+            return self._benchmarks[name]
+
+    def get_by_category(
+        self,
+        category: BenchmarkCategory,
+    ) -> list[type[Benchmark]]:
+        """Get all benchmarks in a category.
+
+        Args:
+            category: Benchmark category
+
+        Returns:
+            List of benchmark classes
+        """
+        with self._lock:
+            return [
+                cls for cls in self._benchmarks.values()
+                if cls.category == category
+            ]
+
+    def list_names(self) -> list[str]:
+        """List all registered benchmark names."""
+        with self._lock:
+            return list(self._benchmarks.keys())
+
+    def list_all(self) -> list[type[Benchmark]]:
+        """List all registered benchmark classes."""
+        with self._lock:
+            return list(self._benchmarks.values())
+
+    def create(self, name: str) -> Benchmark:
+        """Create a benchmark instance by name.
+
+        Args:
+            name: Benchmark name
+
+        Returns:
+            Benchmark instance
+        """
+        benchmark_class = self.get(name)
+        return benchmark_class()
+
+
+# Global registry instance
+benchmark_registry = BenchmarkRegistry()
+
+
+def register_benchmark(cls: type[Benchmark]) -> type[Benchmark]:
+    """Decorator to register a benchmark with the global registry.
+
+    Example:
+        @register_benchmark
+        class MyBenchmark(Benchmark):
+            name = "my_benchmark"
+            ...
+    """
+    return benchmark_registry.register(cls)