benchmark.py

#!/usr/bin/env python3
"""
Orpheus TTS Benchmark Script

Comprehensive benchmarking for:
- TTFB (Time to First Byte) latency
- Concurrent stream handling (CCU)
- Throughput under load
- Tuning parameter exploration

Targets for RTX 4090:
- TTFB: <150ms (conversational), <200ms (long-form)
- CCU: 16-32 concurrent streams
- Throughput: 50+ requests/second
"""

import os
import sys
import time
import asyncio
import argparse
import statistics
import json
from pathlib import Path
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass, asdict
from concurrent.futures import ProcessPoolExecutor

import aiohttp

# Add parent to path
sys.path.insert(0, str(Path(__file__).parent))
from client import OrpheusTTSClient, TTSConfig


@dataclass
class BenchmarkResult:
    """Result from a single benchmark run."""
    ttfb_ms: float
    total_ms: float
    audio_size_bytes: int
    tokens_estimated: int
    success: bool
    error: str = None


@dataclass  
class BenchmarkSummary:
    """Summary statistics from benchmark."""
    num_requests: int
    successful: int
    failed: int
    
    # TTFB stats
    ttfb_avg_ms: float
    ttfb_min_ms: float
    ttfb_max_ms: float
    ttfb_p50_ms: float
    ttfb_p95_ms: float
    ttfb_p99_ms: float
    
    # Total time stats
    total_avg_ms: float
    total_min_ms: float
    total_max_ms: float
    
    # Throughput
    requests_per_second: float
    tokens_per_second: float
    concurrent_streams: int
    
    # Quality metrics
    target_ttfb_met: bool = False
    target_concurrent_met: bool = False


# Test text corpus
TEST_TEXTS = {
    "short": [
        "Hello, how are you?",
        "Yes, I can help with that.",
        "Thank you for calling.",
        "Please hold for a moment.",
        "Is there anything else?",
    ],
    "medium": [
        "Hello, thank you for calling our support line. My name is Sarah and I'll be helping you today. How may I assist you?",
        "I understand your concern about the recent billing issue. Let me pull up your account and take a look at what happened.",
        "The weather today will be sunny with a high of 75 degrees. Perfect weather for outdoor activities. Don't forget sunscreen!",
    ],
    "long": [
        """Welcome to our premium customer service line. My name is Michael and I'm a senior support specialist.
        I see you've been a valued customer with us for over three years, and I want to make sure we resolve 
        your issue completely today. Could you please describe the problem you're experiencing in detail?
        I'm here to help and will do everything I can to ensure your satisfaction.""",
        """Today we're going to discuss the quarterly financial results for our company. Revenue increased by 
        fifteen percent compared to the same period last year, driven primarily by strong performance in our 
        cloud services division. Operating expenses remained flat, resulting in improved profit margins.
        Looking ahead to the next quarter, we expect continued growth in all major business segments.""",
    ],
}


async def run_single_request(
    client: OrpheusTTSClient,
    text: str,
    voice: str,
) -> BenchmarkResult:
    """Run a single TTS request and measure timing."""
    start_time = time.perf_counter()
    first_chunk_time = None
    audio_size = 0
    
    try:
        async for chunk in client.synthesize_stream(text, voice):
            if first_chunk_time is None:
                first_chunk_time = time.perf_counter()
            audio_size += len(chunk)
        
        end_time = time.perf_counter()
        
        ttfb = (first_chunk_time - start_time) * 1000 if first_chunk_time else 0
        total = (end_time - start_time) * 1000
        
        # Estimate tokens: ~83 tokens per second of audio at 24kHz, 16-bit mono
        audio_duration_s = audio_size / (24000 * 2)  # 24kHz, 2 bytes per sample
        tokens_estimated = int(audio_duration_s * 83)
        
        return BenchmarkResult(
            ttfb_ms=ttfb,
            total_ms=total,
            audio_size_bytes=audio_size,
            tokens_estimated=tokens_estimated,
            success=True,
        )
        
    except Exception as e:
        return BenchmarkResult(
            ttfb_ms=0,
            total_ms=0,
            audio_size_bytes=0,
            tokens_estimated=0,
            success=False,
            error=str(e),
        )


def calculate_summary(
    results: List[BenchmarkResult],
    total_time: float,
    concurrent: int,
    target_ttfb_ms: float = 150,
    target_concurrent: int = 16,
) -> BenchmarkSummary:
    """Calculate benchmark summary from results."""
    successful = [r for r in results if r.success]
    ttfb_times = [r.ttfb_ms for r in successful]
    total_times = [r.total_ms for r in successful]
    tokens = [r.tokens_estimated for r in successful]
    
    if not ttfb_times:
        raise Exception("All requests failed")
    
    ttfb_sorted = sorted(ttfb_times)
    n = len(ttfb_sorted)
    
    return BenchmarkSummary(
        num_requests=len(results),
        successful=len(successful),
        failed=len(results) - len(successful),
        
        ttfb_avg_ms=statistics.mean(ttfb_times),
        ttfb_min_ms=min(ttfb_times),
        ttfb_max_ms=max(ttfb_times),
        ttfb_p50_ms=ttfb_sorted[n // 2],
        ttfb_p95_ms=ttfb_sorted[int(n * 0.95)] if n > 1 else ttfb_sorted[0],
        ttfb_p99_ms=ttfb_sorted[min(int(n * 0.99), n - 1)],
        
        total_avg_ms=statistics.mean(total_times),
        total_min_ms=min(total_times),
        total_max_ms=max(total_times),
        
        requests_per_second=len(successful) / total_time,
        tokens_per_second=sum(tokens) / total_time,
        concurrent_streams=concurrent,
        
        target_ttfb_met=statistics.mean(ttfb_times) <= target_ttfb_ms,
        target_concurrent_met=concurrent >= target_concurrent,
    )


async def benchmark_ttfb(
    url: str,
    num_requests: int = 50,
    voice: str = "tara",
    text_type: str = "medium",
    warmup: int = 3,
    target_ttfb_ms: float = 150,
) -> BenchmarkSummary:
    """Benchmark time-to-first-byte latency."""
    texts = TEST_TEXTS.get(text_type, TEST_TEXTS["medium"])
    results: List[BenchmarkResult] = []
    
    async with OrpheusTTSClient(url) as client:
        # Warmup
        print(f"  Warming up with {warmup} requests...")
        for i in range(warmup):
            await client.synthesize(texts[0], voice)
        
        # Benchmark
        print(f"  Running {num_requests} sequential requests...")
        start_time = time.perf_counter()
        
        for i in range(num_requests):
            text = texts[i % len(texts)]
            result = await run_single_request(client, text, voice)
            results.append(result)
            
            if (i + 1) % 10 == 0:
                avg_ttfb = statistics.mean([r.ttfb_ms for r in results if r.success])
                print(f"    {i + 1}/{num_requests} - Avg TTFB: {avg_ttfb:.1f}ms")
        
        total_time = time.perf_counter() - start_time
    
    return calculate_summary(results, total_time, 1, target_ttfb_ms)


async def benchmark_concurrent(
    url: str,
    concurrent: int = 16,
    requests_per_stream: int = 5,
    voice: str = "tara",
    text_type: str = "medium",
    target_ttfb_ms: float = 200,
) -> BenchmarkSummary:
    """Benchmark concurrent stream handling."""
    texts = TEST_TEXTS.get(text_type, TEST_TEXTS["medium"])
    results: List[BenchmarkResult] = []
    
    async def run_stream(stream_id: int) -> List[BenchmarkResult]:
        stream_results = []
        async with OrpheusTTSClient(url) as client:
            for i in range(requests_per_stream):
                text = texts[i % len(texts)]
                result = await run_single_request(client, text, voice)
                stream_results.append(result)
        return stream_results
    
    print(f"  Running {concurrent} concurrent streams, {requests_per_stream} requests each...")
    start_time = time.perf_counter()
    
    # Run concurrent streams
    tasks = [run_stream(i) for i in range(concurrent)]
    stream_results = await asyncio.gather(*tasks)
    
    total_time = time.perf_counter() - start_time
    
    # Flatten results
    for sr in stream_results:
        results.extend(sr)
    
    print(f"    Completed {len(results)} requests in {total_time:.1f}s")
    
    return calculate_summary(results, total_time, concurrent, target_ttfb_ms, concurrent)


async def benchmark_concurrent_scaling(
    url: str,
    max_concurrent: int = 32,
    requests_per_level: int = 20,
    voice: str = "tara",
) -> Dict[int, BenchmarkSummary]:
    """Benchmark how latency scales with concurrency."""
    print("\n>>> Concurrent Scaling Benchmark")
    print("    Testing how TTFB changes with increasing concurrency...")
    
    results = {}
    levels = [1, 2, 4, 8, 16, 24, 32]
    levels = [l for l in levels if l <= max_concurrent]
    
    for concurrent in levels:
        print(f"\n  Testing {concurrent} concurrent streams...")
        summary = await benchmark_concurrent(
            url,
            concurrent=concurrent,
            requests_per_stream=requests_per_level // concurrent,
            voice=voice,
        )
        results[concurrent] = summary
    
    return results


async def benchmark_throughput(
    url: str,
    duration_seconds: int = 30,
    concurrent: int = 16,
    voice: str = "tara",
) -> Dict[str, Any]:
    """Benchmark maximum throughput under sustained load."""
    text = TEST_TEXTS["short"][0]
    
    completed = 0
    total_audio_bytes = 0
    total_tokens = 0
    errors = 0
    ttfb_times = []
    
    start_time = time.perf_counter()
    end_time = start_time + duration_seconds
    
    semaphore = asyncio.Semaphore(concurrent)
    
    async def run_request():
        nonlocal completed, total_audio_bytes, total_tokens, errors, ttfb_times
        async with semaphore:
            try:
                async with OrpheusTTSClient(url) as client:
                    result = await run_single_request(client, text, voice)
                    if result.success:
                        completed += 1
                        total_audio_bytes += result.audio_size_bytes
                        total_tokens += result.tokens_estimated
                        ttfb_times.append(result.ttfb_ms)
                    else:
                        errors += 1
            except Exception:
                errors += 1
    
    print(f"  Running throughput test for {duration_seconds}s with {concurrent} concurrent...")
    
    tasks = []
    last_report = start_time
    
    while time.perf_counter() < end_time:
        task = asyncio.create_task(run_request())
        tasks.append(task)
        
        # Report progress every 5 seconds
        now = time.perf_counter()
        if now - last_report >= 5:
            elapsed = now - start_time
            print(f"    {elapsed:.0f}s: {completed} completed, {errors} errors, "
                  f"{completed/elapsed:.1f} req/s")
            last_report = now
        
        await asyncio.sleep(0.01)
    
    # Wait for remaining tasks
    await asyncio.gather(*tasks, return_exceptions=True)
    
    actual_duration = time.perf_counter() - start_time
    
    return {
        "duration_seconds": actual_duration,
        "completed_requests": completed,
        "errors": errors,
        "concurrent": concurrent,
        "requests_per_second": completed / actual_duration,
        "tokens_per_second": total_tokens / actual_duration,
        "total_audio_mb": total_audio_bytes / (1024 * 1024),
        "audio_mb_per_second": total_audio_bytes / (1024 * 1024) / actual_duration,
        "avg_ttfb_ms": statistics.mean(ttfb_times) if ttfb_times else 0,
        "p95_ttfb_ms": sorted(ttfb_times)[int(len(ttfb_times) * 0.95)] if ttfb_times else 0,
    }


async def benchmark_tuning(
    url: str,
    voice: str = "tara",
) -> Dict[str, Any]:
    """
    Explore tuning parameters to find optimal configuration.
    
    Tests different combinations of:
    - Batch sizes
    - Concurrent sessions
    - Text lengths
    """
    print("\n>>> Tuning Parameter Exploration")
    
    results = {
        "text_length_impact": {},
        "concurrent_scaling": {},
        "optimal_config": {},
    }
    
    # Test text length impact
    print("\n  Testing text length impact on TTFB...")
    for text_type in ["short", "medium", "long"]:
        print(f"    {text_type} texts...")
        summary = await benchmark_ttfb(
            url,
            num_requests=20,
            voice=voice,
            text_type=text_type,
            warmup=2,
        )
        results["text_length_impact"][text_type] = {
            "avg_ttfb_ms": summary.ttfb_avg_ms,
            "p95_ttfb_ms": summary.ttfb_p95_ms,
            "tokens_per_sec": summary.tokens_per_second,
        }
    
    # Test concurrent scaling
    print("\n  Testing concurrent scaling...")
    for concurrent in [4, 8, 16, 24, 32]:
        print(f"    {concurrent} concurrent streams...")
        try:
            summary = await benchmark_concurrent(
                url,
                concurrent=concurrent,
                requests_per_stream=3,
                voice=voice,
            )
            results["concurrent_scaling"][concurrent] = {
                "avg_ttfb_ms": summary.ttfb_avg_ms,
                "p95_ttfb_ms": summary.ttfb_p95_ms,
                "requests_per_sec": summary.requests_per_second,
                "success_rate": summary.successful / summary.num_requests,
            }
        except Exception as e:
            results["concurrent_scaling"][concurrent] = {"error": str(e)}
    
    # Find optimal config
    scaling = results["concurrent_scaling"]
    best_concurrent = max(
        [c for c in scaling.keys() if isinstance(c, int) and "error" not in scaling[c]],
        key=lambda c: scaling[c]["requests_per_sec"] if scaling[c]["avg_ttfb_ms"] < 200 else 0,
        default=8
    )
    
    results["optimal_config"] = {
        "recommended_concurrent": best_concurrent,
        "expected_ttfb_ms": scaling.get(best_concurrent, {}).get("avg_ttfb_ms", 0),
        "expected_throughput": scaling.get(best_concurrent, {}).get("requests_per_sec", 0),
    }
    
    return results


def print_summary(summary: BenchmarkSummary, title: str = "BENCHMARK RESULTS"):
    """Print benchmark summary."""
    print("\n" + "=" * 60)
    print(title)
    print("=" * 60)
    
    print(f"\nRequests: {summary.num_requests}")
    print(f"Successful: {summary.successful} ({100*summary.successful/summary.num_requests:.1f}%)")
    print(f"Failed: {summary.failed}")
    print(f"Concurrent Streams: {summary.concurrent_streams}")
    
    print(f"\n--- Time to First Byte (TTFB) ---")
    print(f"Average: {summary.ttfb_avg_ms:.1f} ms")
    print(f"Min: {summary.ttfb_min_ms:.1f} ms")
    print(f"Max: {summary.ttfb_max_ms:.1f} ms")
    print(f"P50: {summary.ttfb_p50_ms:.1f} ms")
    print(f"P95: {summary.ttfb_p95_ms:.1f} ms")
    print(f"P99: {summary.ttfb_p99_ms:.1f} ms")
    
    print(f"\n--- Total Request Time ---")
    print(f"Average: {summary.total_avg_ms:.1f} ms")
    print(f"Min: {summary.total_min_ms:.1f} ms")
    print(f"Max: {summary.total_max_ms:.1f} ms")
    
    print(f"\n--- Throughput ---")
    print(f"Requests/second: {summary.requests_per_second:.2f}")
    print(f"Tokens/second: {summary.tokens_per_second:.0f}")
    
    # Target comparison
    print(f"\n--- Target Comparison (RTX 4090) ---")
    
    if summary.target_ttfb_met:
        print(f"✓ TTFB target (<150ms avg): PASS ({summary.ttfb_avg_ms:.1f}ms)")
    else:
        print(f"✗ TTFB target (<150ms avg): FAIL ({summary.ttfb_avg_ms:.1f}ms)")
    
    if summary.concurrent_streams >= 16:
        print(f"✓ CCU target (16+): PASS ({summary.concurrent_streams} streams)")
    elif summary.concurrent_streams >= 8:
        print(f"~ CCU target (16+): PARTIAL ({summary.concurrent_streams} streams)")
    else:
        print(f"✗ CCU target (16+): FAIL ({summary.concurrent_streams} streams)")
    
    if summary.tokens_per_second >= 1000:
        print(f"✓ Token throughput (1000+/s): PASS ({summary.tokens_per_second:.0f}/s)")
    else:
        print(f"✗ Token throughput (1000+/s): FAIL ({summary.tokens_per_second:.0f}/s)")


def print_tuning_results(results: Dict[str, Any]):
    """Print tuning exploration results."""
    print("\n" + "=" * 60)
    print("TUNING RESULTS")
    print("=" * 60)
    
    print("\n--- Text Length Impact ---")
    for text_type, data in results["text_length_impact"].items():
        print(f"  {text_type:8s}: TTFB={data['avg_ttfb_ms']:.1f}ms, "
              f"P95={data['p95_ttfb_ms']:.1f}ms, "
              f"Tok/s={data['tokens_per_sec']:.0f}")
    
    print("\n--- Concurrent Scaling ---")
    print("  CCU   | Avg TTFB | P95 TTFB | Req/s  | Success")
    print("  ------|----------|----------|--------|--------")
    for ccu, data in sorted(results["concurrent_scaling"].items()):
        if "error" in data:
            print(f"  {ccu:5d} | ERROR: {data['error'][:30]}")
        else:
            print(f"  {ccu:5d} | {data['avg_ttfb_ms']:7.1f}ms | {data['p95_ttfb_ms']:7.1f}ms | "
                  f"{data['requests_per_sec']:6.1f} | {data['success_rate']*100:5.1f}%")
    
    print("\n--- Recommended Configuration ---")
    opt = results["optimal_config"]
    print(f"  Concurrent Sessions: {opt['recommended_concurrent']}")
    print(f"  Expected TTFB: {opt['expected_ttfb_ms']:.1f}ms")
    print(f"  Expected Throughput: {opt['expected_throughput']:.1f} req/s")


async def main():
    parser = argparse.ArgumentParser(description="Orpheus TTS Benchmark")
    parser.add_argument("--url", default="http://localhost:8000", help="Server URL")
    parser.add_argument("--voice", default="tara", help="Voice to use")
    parser.add_argument("--mode", choices=["ttfb", "concurrent", "throughput", "scaling", "tuning", "all"],
                        default="all", help="Benchmark mode")
    parser.add_argument("--requests", type=int, default=50, help="Number of requests for TTFB test")
    parser.add_argument("--concurrent", type=int, default=16, help="Concurrent streams")
    parser.add_argument("--duration", type=int, default=30, help="Throughput test duration (seconds)")
    parser.add_argument("--text", choices=["short", "medium", "long"], default="medium",
                        help="Text length for tests")
    parser.add_argument("--output", type=str, help="Save results to JSON file")
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("ORPHEUS TTS BENCHMARK")
    print("=" * 60)
    print(f"Server: {args.url}")
    print(f"Voice: {args.voice}")
    print(f"Mode: {args.mode}")
    print("-" * 60)
    
    all_results = {}
    
    if args.mode in ["ttfb", "all"]:
        print("\n>>> TTFB Benchmark (Sequential Requests)")
        summary = await benchmark_ttfb(
            args.url,
            num_requests=args.requests,
            voice=args.voice,
            text_type=args.text,
        )
        print_summary(summary, "TTFB BENCHMARK RESULTS")
        all_results["ttfb"] = asdict(summary)
    
    if args.mode in ["concurrent", "all"]:
        print("\n>>> Concurrent Streams Benchmark")
        summary = await benchmark_concurrent(
            args.url,
            concurrent=args.concurrent,
            requests_per_stream=max(1, args.requests // args.concurrent),
            voice=args.voice,
            text_type=args.text,
        )
        print_summary(summary, f"CONCURRENT ({args.concurrent} streams) BENCHMARK RESULTS")
        all_results["concurrent"] = asdict(summary)
    
    if args.mode in ["scaling", "all"]:
        scaling_results = await benchmark_concurrent_scaling(
            args.url,
            max_concurrent=args.concurrent,
            requests_per_level=args.requests,
            voice=args.voice,
        )
        
        print("\n" + "=" * 60)
        print("CONCURRENT SCALING RESULTS")
        print("=" * 60)
        print("\n  CCU  | Avg TTFB | P95 TTFB | Req/s ")
        print("  -----|----------|----------|-------")
        for ccu, summary in sorted(scaling_results.items()):
            print(f"  {ccu:4d} | {summary.ttfb_avg_ms:7.1f}ms | {summary.ttfb_p95_ms:7.1f}ms | "
                  f"{summary.requests_per_second:5.1f}")
        
        all_results["scaling"] = {k: asdict(v) for k, v in scaling_results.items()}
    
    if args.mode in ["throughput", "all"]:
        print("\n>>> Throughput Benchmark (Sustained Load)")
        results = await benchmark_throughput(
            args.url,
            duration_seconds=args.duration,
            concurrent=args.concurrent,
            voice=args.voice,
        )
        
        print("\n" + "=" * 60)
        print("THROUGHPUT BENCHMARK RESULTS")
        print("=" * 60)
        print(f"\nDuration: {results['duration_seconds']:.1f}s")
        print(f"Concurrent: {results['concurrent']}")
        print(f"Completed: {results['completed_requests']}")
        print(f"Errors: {results['errors']}")
        print(f"\n--- Performance ---")
        print(f"Requests/sec: {results['requests_per_second']:.2f}")
        print(f"Tokens/sec: {results['tokens_per_second']:.0f}")
        print(f"Audio MB/sec: {results['audio_mb_per_second']:.2f}")
        print(f"\n--- Latency Under Load ---")
        print(f"Avg TTFB: {results['avg_ttfb_ms']:.1f}ms")
        print(f"P95 TTFB: {results['p95_ttfb_ms']:.1f}ms")
        
        all_results["throughput"] = results
    
    if args.mode == "tuning":
        tuning_results = await benchmark_tuning(args.url, args.voice)
        print_tuning_results(tuning_results)
        all_results["tuning"] = tuning_results
    
    # Save results if requested
    if args.output:
        with open(args.output, 'w') as f:
            json.dump(all_results, f, indent=2)
        print(f"\nResults saved to: {args.output}")
    
    print("\n" + "=" * 60)
    print("BENCHMARK COMPLETE")
    print("=" * 60)