Spaces:
Configuration error
Configuration error
| """ARIA Comprehensive Load Testing & Capacity Benchmark Engine. | |
| Orchestrates all load test phases, collects telemetry, calculates statistical | |
| percentiles (p50, p95, p99, max), evaluates safety thresholds, and produces | |
| comprehensive capacity measurements. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import random | |
| import time | |
| from dataclasses import dataclass, field | |
| from typing import List, Optional, Tuple | |
| import httpx | |
| import numpy as np | |
| from tests.load.sse_client import consume_chat_stream, consume_analyze_stream | |
| from tests.load.system_monitor import SystemMonitor | |
| from tests.load.scenarios import ( | |
| BENCHMARK_REPO, | |
| REALISTIC_CHAT_QUERIES, | |
| BROWSING_ENDPOINTS, | |
| LIGHTWEIGHT_ENDPOINTS, | |
| ) | |
| class ConcurrencyLevelResult: | |
| """Metrics captured for a specific concurrency level in a scenario.""" | |
| concurrency: int | |
| duration_s: float | |
| total_requests: int | |
| successful_requests: int | |
| failed_requests: int | |
| timeout_requests: int | |
| error_rate_pct: float | |
| throughput_rps: float | |
| p50_latency_ms: float | |
| p95_latency_ms: float | |
| p99_latency_ms: float | |
| max_latency_ms: float | |
| avg_ttft_ms: Optional[float] = None | |
| p95_ttft_ms: Optional[float] = None | |
| avg_stream_duration_ms: Optional[float] = None | |
| total_tokens_delivered: int = 0 | |
| tokens_per_sec: float = 0.0 | |
| fallback_count: int = 0 | |
| cpu_avg_pct: float = 0.0 | |
| cpu_peak_pct: float = 0.0 | |
| mem_rss_mb_peak: float = 0.0 | |
| mem_growth_mb: float = 0.0 | |
| peak_open_connections: int = 0 | |
| errors: List[str] = field(default_factory=list) | |
| class ScenarioBenchmarkResult: | |
| """Complete benchmark results across progressive concurrency levels for a scenario.""" | |
| scenario_name: str | |
| levels: List[ConcurrencyLevelResult] = field(default_factory=list) | |
| safe_capacity: int = 0 | |
| degradation_point: int = 0 | |
| hard_limit: int = 0 | |
| primary_bottleneck: str = "" | |
| class BenchmarkEngine: | |
| """Automated benchmark executor for ARIA load and capacity testing.""" | |
| def __init__( | |
| self, | |
| base_url: str = "http://127.0.0.1:8001", | |
| target_pid: Optional[int] = None, | |
| timeout: float = 60.0, | |
| api_key: str = "aria-benchmark-key", | |
| ): | |
| self.base_url = base_url.rstrip("/") | |
| self.target_pid = target_pid | |
| self.timeout = timeout | |
| self.api_key = api_key | |
| def _calc_percentiles(values: List[float]) -> Tuple[float, float, float, float]: | |
| if not values: | |
| return 0.0, 0.0, 0.0, 0.0 | |
| arr = np.array(values) | |
| return ( | |
| float(round(np.percentile(arr, 50), 1)), | |
| float(round(np.percentile(arr, 95), 1)), | |
| float(round(np.percentile(arr, 99), 1)), | |
| float(round(np.max(arr), 1)), | |
| ) | |
| async def run_chat_batch( | |
| self, | |
| concurrency: int, | |
| duration_s: float = 10.0, | |
| repo: str = BENCHMARK_REPO, | |
| ) -> ConcurrencyLevelResult: | |
| """Run concurrent Chat SSE requests.""" | |
| monitor = SystemMonitor(target_pid=self.target_pid) | |
| await monitor.start() | |
| latencies: List[float] = [] | |
| ttfts: List[float] = [] | |
| stream_durations: List[float] = [] | |
| errors: List[str] = [] | |
| successful = 0 | |
| failed = 0 | |
| timeouts = 0 | |
| total_tokens = 0 | |
| fallback_count = 0 | |
| stop_time = time.time() + duration_s | |
| limits = httpx.Limits( | |
| max_connections=concurrency + 50, | |
| max_keepalive_connections=concurrency + 50, | |
| ) | |
| async with httpx.AsyncClient(limits=limits) as client: | |
| async def worker(): | |
| nonlocal successful, failed, timeouts, total_tokens, fallback_count | |
| while time.time() < stop_time: | |
| question = random.choice(REALISTIC_CHAT_QUERIES) | |
| res = await consume_chat_stream( | |
| client=client, | |
| base_url=self.base_url, | |
| repo=repo, | |
| message=question, | |
| timeout=self.timeout, | |
| ) | |
| latencies.append(res.duration_ms) | |
| if res.ttft_ms is not None: | |
| ttfts.append(res.ttft_ms) | |
| if res.duration_ms > 0: | |
| stream_durations.append(res.duration_ms) | |
| total_tokens += res.tokens_count | |
| if res.fallback_mode: | |
| fallback_count += 1 | |
| if res.success: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if "timed out" in str(res.error).lower(): | |
| timeouts += 1 | |
| if res.error and len(errors) < 10: | |
| errors.append(str(res.error)) | |
| # Brief realistic user pause | |
| await asyncio.sleep(0.05) | |
| tasks = [asyncio.create_task(worker()) for _ in range(concurrency)] | |
| await asyncio.gather(*tasks, return_exceptions=True) | |
| res_summary = await monitor.stop() | |
| total_reqs = successful + failed | |
| actual_duration = max(0.1, res_summary.duration_s) | |
| p50, p95, p99, max_lat = self._calc_percentiles(latencies) | |
| _, p95_ttft, _, _ = self._calc_percentiles(ttfts) | |
| return ConcurrencyLevelResult( | |
| concurrency=concurrency, | |
| duration_s=actual_duration, | |
| total_requests=total_reqs, | |
| successful_requests=successful, | |
| failed_requests=failed, | |
| timeout_requests=timeouts, | |
| error_rate_pct=round((failed / total_reqs * 100) if total_reqs else 0.0, 2), | |
| throughput_rps=round(total_reqs / actual_duration, 2), | |
| p50_latency_ms=p50, | |
| p95_latency_ms=p95, | |
| p99_latency_ms=p99, | |
| max_latency_ms=max_lat, | |
| avg_ttft_ms=round(sum(ttfts) / len(ttfts), 1) if ttfts else None, | |
| p95_ttft_ms=p95_ttft if ttfts else None, | |
| avg_stream_duration_ms=round( | |
| sum(stream_durations) / len(stream_durations), 1 | |
| ) | |
| if stream_durations | |
| else None, | |
| total_tokens_delivered=total_tokens, | |
| tokens_per_sec=round(total_tokens / actual_duration, 1), | |
| fallback_count=fallback_count, | |
| cpu_avg_pct=res_summary.cpu_avg, | |
| cpu_peak_pct=res_summary.cpu_peak, | |
| mem_rss_mb_peak=res_summary.memory_rss_mb_peak, | |
| mem_growth_mb=res_summary.memory_growth_mb, | |
| peak_open_connections=res_summary.peak_open_connections, | |
| errors=errors, | |
| ) | |
| async def run_sse_concurrency_test( | |
| self, | |
| concurrency: int, | |
| duration_s: float = 10.0, | |
| ) -> ConcurrencyLevelResult: | |
| """Measure pure concurrent SSE streams.""" | |
| return await self.run_chat_batch( | |
| concurrency=concurrency, | |
| duration_s=duration_s, | |
| ) | |
| async def run_repo_analysis_batch( | |
| self, | |
| concurrency: int, | |
| duration_s: float = 15.0, | |
| ) -> ConcurrencyLevelResult: | |
| """Measure concurrent repository analysis / indexing capacity.""" | |
| monitor = SystemMonitor(target_pid=self.target_pid) | |
| await monitor.start() | |
| latencies: List[float] = [] | |
| errors: List[str] = [] | |
| successful = 0 | |
| failed = 0 | |
| timeouts = 0 | |
| stop_time = time.time() + duration_s | |
| limits = httpx.Limits( | |
| max_connections=concurrency + 20, | |
| max_keepalive_connections=concurrency + 20, | |
| ) | |
| async with httpx.AsyncClient(limits=limits) as client: | |
| async def worker(): | |
| nonlocal successful, failed, timeouts | |
| while time.time() < stop_time: | |
| res = await consume_analyze_stream( | |
| client=client, | |
| base_url=self.base_url, | |
| repo_url=f"https://github.com/{BENCHMARK_REPO}", | |
| force_rebuild=False, | |
| timeout=90.0, | |
| ) | |
| latencies.append(res.duration_ms) | |
| if res.success: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if "timed out" in str(res.error).lower(): | |
| timeouts += 1 | |
| if res.error and len(errors) < 10: | |
| errors.append(str(res.error)) | |
| await asyncio.sleep(0.2) | |
| tasks = [asyncio.create_task(worker()) for _ in range(concurrency)] | |
| await asyncio.gather(*tasks, return_exceptions=True) | |
| res_summary = await monitor.stop() | |
| total_reqs = successful + failed | |
| actual_duration = max(0.1, res_summary.duration_s) | |
| p50, p95, p99, max_lat = self._calc_percentiles(latencies) | |
| return ConcurrencyLevelResult( | |
| concurrency=concurrency, | |
| duration_s=actual_duration, | |
| total_requests=total_reqs, | |
| successful_requests=successful, | |
| failed_requests=failed, | |
| timeout_requests=timeouts, | |
| error_rate_pct=round((failed / total_reqs * 100) if total_reqs else 0.0, 2), | |
| throughput_rps=round(total_reqs / actual_duration, 2), | |
| p50_latency_ms=p50, | |
| p95_latency_ms=p95, | |
| p99_latency_ms=p99, | |
| max_latency_ms=max_lat, | |
| cpu_avg_pct=res_summary.cpu_avg, | |
| cpu_peak_pct=res_summary.cpu_peak, | |
| mem_rss_mb_peak=res_summary.memory_rss_mb_peak, | |
| mem_growth_mb=res_summary.memory_growth_mb, | |
| peak_open_connections=res_summary.peak_open_connections, | |
| errors=errors, | |
| ) | |
| async def run_mixed_workload_batch( | |
| self, | |
| concurrency: int, | |
| duration_s: float = 12.0, | |
| ) -> ConcurrencyLevelResult: | |
| """Measure Mixed Production Workload: 60% Chat SSE, 20% Repo Analysis, 10% Browsing, 10% Ops.""" | |
| monitor = SystemMonitor(target_pid=self.target_pid) | |
| await monitor.start() | |
| latencies: List[float] = [] | |
| ttfts: List[float] = [] | |
| errors: List[str] = [] | |
| successful = 0 | |
| failed = 0 | |
| timeouts = 0 | |
| total_tokens = 0 | |
| stop_time = time.time() + duration_s | |
| limits = httpx.Limits( | |
| max_connections=concurrency + 50, | |
| max_keepalive_connections=concurrency + 50, | |
| ) | |
| async with httpx.AsyncClient(limits=limits) as client: | |
| async def worker(): | |
| nonlocal successful, failed, timeouts, total_tokens | |
| while time.time() < stop_time: | |
| r = random.random() | |
| t_req_start = time.perf_counter() | |
| if r < 0.60: | |
| # 60% Chat SSE | |
| question = random.choice(REALISTIC_CHAT_QUERIES) | |
| res = await consume_chat_stream( | |
| client=client, | |
| base_url=self.base_url, | |
| repo=BENCHMARK_REPO, | |
| message=question, | |
| timeout=self.timeout, | |
| ) | |
| latencies.append(res.duration_ms) | |
| if res.ttft_ms: | |
| ttfts.append(res.ttft_ms) | |
| total_tokens += res.tokens_count | |
| if res.success: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if "timed out" in str(res.error).lower(): | |
| timeouts += 1 | |
| if res.error and len(errors) < 10: | |
| errors.append(str(res.error)) | |
| elif r < 0.80: | |
| # 20% Repo Analysis | |
| res = await consume_analyze_stream( | |
| client=client, | |
| base_url=self.base_url, | |
| repo_url=f"https://github.com/{BENCHMARK_REPO}", | |
| force_rebuild=False, | |
| timeout=60.0, | |
| ) | |
| latencies.append(res.duration_ms) | |
| if res.success: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if "timed out" in str(res.error).lower(): | |
| timeouts += 1 | |
| if res.error and len(errors) < 10: | |
| errors.append(str(res.error)) | |
| elif r < 0.90: | |
| # 10% Repo Browsing / Intelligence APIs | |
| ep = random.choice(BROWSING_ENDPOINTS) | |
| try: | |
| resp = await client.request( | |
| method=ep["method"], | |
| url=f"{self.base_url}{ep['path']}", | |
| headers={"X-API-Key": "aria-benchmark-key"}, | |
| timeout=15.0, | |
| ) | |
| lat = (time.perf_counter() - t_req_start) * 1000.0 | |
| latencies.append(lat) | |
| if resp.status_code == 200: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if len(errors) < 10: | |
| errors.append( | |
| f"HTTP {resp.status_code}: {ep['path']}" | |
| ) | |
| except Exception as exc: | |
| lat = (time.perf_counter() - t_req_start) * 1000.0 | |
| latencies.append(lat) | |
| failed += 1 | |
| if len(errors) < 10: | |
| errors.append(f"Browsing err: {exc}") | |
| else: | |
| # 10% Lightweight Ops APIs | |
| ep = random.choice(LIGHTWEIGHT_ENDPOINTS) | |
| try: | |
| resp = await client.request( | |
| method=ep["method"], | |
| url=f"{self.base_url}{ep['path']}", | |
| headers={"X-API-Key": "aria-benchmark-key"}, | |
| timeout=10.0, | |
| ) | |
| lat = (time.perf_counter() - t_req_start) * 1000.0 | |
| latencies.append(lat) | |
| if resp.status_code == 200: | |
| successful += 1 | |
| else: | |
| failed += 1 | |
| if len(errors) < 10: | |
| errors.append( | |
| f"HTTP {resp.status_code}: {ep['path']}" | |
| ) | |
| except Exception as exc: | |
| lat = (time.perf_counter() - t_req_start) * 1000.0 | |
| latencies.append(lat) | |
| failed += 1 | |
| if len(errors) < 10: | |
| errors.append(f"Ops err: {exc}") | |
| await asyncio.sleep(0.05) | |
| tasks = [asyncio.create_task(worker()) for _ in range(concurrency)] | |
| await asyncio.gather(*tasks, return_exceptions=True) | |
| res_summary = await monitor.stop() | |
| total_reqs = successful + failed | |
| actual_duration = max(0.1, res_summary.duration_s) | |
| p50, p95, p99, max_lat = self._calc_percentiles(latencies) | |
| return ConcurrencyLevelResult( | |
| concurrency=concurrency, | |
| duration_s=actual_duration, | |
| total_requests=total_reqs, | |
| successful_requests=successful, | |
| failed_requests=failed, | |
| timeout_requests=timeouts, | |
| error_rate_pct=round((failed / total_reqs * 100) if total_reqs else 0.0, 2), | |
| throughput_rps=round(total_reqs / actual_duration, 2), | |
| p50_latency_ms=p50, | |
| p95_latency_ms=p95, | |
| p99_latency_ms=p99, | |
| max_latency_ms=max_lat, | |
| avg_ttft_ms=round(sum(ttfts) / len(ttfts), 1) if ttfts else None, | |
| total_tokens_delivered=total_tokens, | |
| tokens_per_sec=round(total_tokens / actual_duration, 1), | |
| cpu_avg_pct=res_summary.cpu_avg, | |
| cpu_peak_pct=res_summary.cpu_peak, | |
| mem_rss_mb_peak=res_summary.memory_rss_mb_peak, | |
| mem_growth_mb=res_summary.memory_growth_mb, | |
| peak_open_connections=res_summary.peak_open_connections, | |
| errors=errors, | |
| ) | |
| def analyze_scenario_capacity( | |
| self, | |
| results: List[ConcurrencyLevelResult], | |
| scenario_name: str, | |
| ) -> ScenarioBenchmarkResult: | |
| """Determine Safe Capacity, Degradation Point, Hard Limit, and Bottleneck.""" | |
| safe_capacity = 0 | |
| degradation_point = 0 | |
| hard_limit = 0 | |
| bottleneck = "Single ASGI Worker / Python GIL CPU Event-Loop Contention" | |
| for res in results: | |
| # Safe condition: error_rate < 1.0% and p95 latency under threshold | |
| latency_thresh = 15000.0 if "analysis" in scenario_name.lower() else 5000.0 | |
| if res.error_rate_pct < 1.0 and res.p95_latency_ms < latency_thresh: | |
| safe_capacity = res.concurrency | |
| # Degradation condition: latency jumps > 2.5x baseline or error > 1% or CPU > 85% | |
| if degradation_point == 0: | |
| if ( | |
| res.error_rate_pct >= 1.0 | |
| or res.p95_latency_ms > latency_thresh | |
| or res.cpu_peak_pct > 85.0 | |
| ): | |
| degradation_point = res.concurrency | |
| # Hard limit: error rate > 10% or timeouts > 5% or sustained failure | |
| if hard_limit == 0: | |
| if res.error_rate_pct >= 10.0 or res.timeout_requests > 0: | |
| hard_limit = res.concurrency | |
| if degradation_point == 0 and results: | |
| degradation_point = results[-1].concurrency | |
| if hard_limit == 0 and results: | |
| hard_limit = int(results[-1].concurrency * 1.5) | |
| # Identify bottleneck characteristics | |
| peak_cpu = max((r.cpu_peak_pct for r in results), default=0) | |
| peak_mem = max((r.mem_rss_mb_peak for r in results), default=0) | |
| timeouts = sum(r.timeout_requests for r in results) | |
| if peak_cpu > 80.0: | |
| bottleneck = ( | |
| "CPU Saturation (Single Worker Event Loop + Local Embedding Generation)" | |
| ) | |
| elif timeouts > 0: | |
| bottleneck = "HTTP Connection Queue Backlog / Worker Starvation" | |
| elif peak_mem > 4000: | |
| bottleneck = "Memory Saturation (In-memory Store / Model Buffers)" | |
| else: | |
| bottleneck = "Single ASGI Event-Loop Concurrency Bound (Uvicorn 1 Worker)" | |
| return ScenarioBenchmarkResult( | |
| scenario_name=scenario_name, | |
| levels=results, | |
| safe_capacity=safe_capacity, | |
| degradation_point=degradation_point, | |
| hard_limit=hard_limit, | |
| primary_bottleneck=bottleneck, | |
| ) | |