wehub-resource-sync
375 lines
13 KiB
Python
375 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Load testing script for the Ray Serve video API.
|
|
|
|
Usage:
|
|
python -m client.load_test --video s3://bucket/path/to/video.mp4 --concurrency 4
|
|
python -m client.load_test --video s3://bucket/video.mp4 --concurrency 8 --url http://localhost:8000
|
|
python -m client.load_test --video s3://bucket/video.mp4 --concurrency 4 --token YOUR_TOKEN
|
|
|
|
Press Ctrl+C to stop and save results to CSV.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import csv
|
|
import signal
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field, asdict
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
|
|
@dataclass
|
|
class RequestResult:
|
|
"""Result of a single request."""
|
|
request_id: str
|
|
stream_id: str
|
|
start_time: float
|
|
end_time: float
|
|
latency_ms: float
|
|
success: bool
|
|
status_code: int
|
|
error: Optional[str] = None
|
|
# Server-side timing (if successful)
|
|
s3_download_ms: Optional[float] = None
|
|
decode_video_ms: Optional[float] = None
|
|
encode_ms: Optional[float] = None
|
|
decode_ms: Optional[float] = None
|
|
server_total_ms: Optional[float] = None
|
|
# Response data
|
|
num_chunks: Optional[int] = None
|
|
video_duration: Optional[float] = None
|
|
num_scene_changes: Optional[int] = None
|
|
|
|
|
|
@dataclass
|
|
class LoadTestStats:
|
|
"""Aggregated statistics for the load test."""
|
|
total_requests: int = 0
|
|
successful_requests: int = 0
|
|
failed_requests: int = 0
|
|
total_latency_ms: float = 0.0
|
|
min_latency_ms: float = float('inf')
|
|
max_latency_ms: float = 0.0
|
|
latencies: list = field(default_factory=list)
|
|
|
|
def add_result(self, result: RequestResult):
|
|
self.total_requests += 1
|
|
if result.success:
|
|
self.successful_requests += 1
|
|
self.total_latency_ms += result.latency_ms
|
|
self.min_latency_ms = min(self.min_latency_ms, result.latency_ms)
|
|
self.max_latency_ms = max(self.max_latency_ms, result.latency_ms)
|
|
self.latencies.append(result.latency_ms)
|
|
else:
|
|
self.failed_requests += 1
|
|
|
|
@property
|
|
def avg_latency_ms(self) -> float:
|
|
if not self.latencies:
|
|
return 0.0
|
|
return self.total_latency_ms / len(self.latencies)
|
|
|
|
@property
|
|
def p50_latency_ms(self) -> float:
|
|
return self._percentile(50)
|
|
|
|
@property
|
|
def p95_latency_ms(self) -> float:
|
|
return self._percentile(95)
|
|
|
|
@property
|
|
def p99_latency_ms(self) -> float:
|
|
return self._percentile(99)
|
|
|
|
def _percentile(self, p: float) -> float:
|
|
if not self.latencies:
|
|
return 0.0
|
|
sorted_latencies = sorted(self.latencies)
|
|
idx = int(len(sorted_latencies) * p / 100)
|
|
idx = min(idx, len(sorted_latencies) - 1)
|
|
return sorted_latencies[idx]
|
|
|
|
|
|
class LoadTester:
|
|
def __init__(
|
|
self,
|
|
video_path: str,
|
|
url: str,
|
|
concurrency: int,
|
|
num_frames: int = 16,
|
|
chunk_duration: float = 10.0,
|
|
timeout: float = 300.0,
|
|
token: Optional[str] = None,
|
|
):
|
|
self.video_path = video_path
|
|
self.url = url
|
|
self.concurrency = concurrency
|
|
self.num_frames = num_frames
|
|
self.chunk_duration = chunk_duration
|
|
self.timeout = timeout
|
|
self.token = token
|
|
|
|
self.results: list[RequestResult] = []
|
|
self.stats = LoadTestStats()
|
|
self.running = True
|
|
self.start_time: Optional[float] = None
|
|
self.request_counter = 0
|
|
self._lock = asyncio.Lock()
|
|
self._semaphore: asyncio.Semaphore = None
|
|
|
|
def _build_payload(self) -> dict:
|
|
"""Build the request payload."""
|
|
stream_id = uuid.uuid4().hex[:8]
|
|
return {
|
|
"stream_id": stream_id,
|
|
"video_path": self.video_path,
|
|
"num_frames": self.num_frames,
|
|
"chunk_duration": self.chunk_duration,
|
|
"use_batching": False,
|
|
}
|
|
|
|
async def _make_request(self, client: httpx.AsyncClient, request_id: str) -> RequestResult:
|
|
"""Make a single request and return the result."""
|
|
payload = self._build_payload()
|
|
stream_id = payload["stream_id"]
|
|
|
|
start = time.perf_counter()
|
|
start_timestamp = time.time()
|
|
|
|
headers = {}
|
|
if self.token:
|
|
headers["Authorization"] = f"Bearer {self.token}"
|
|
|
|
try:
|
|
response = await client.post(f"{self.url}/analyze", json=payload, headers=headers)
|
|
end = time.perf_counter()
|
|
latency_ms = (end - start) * 1000
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
timing = data.get("timing_ms", {})
|
|
return RequestResult(
|
|
request_id=request_id,
|
|
stream_id=stream_id,
|
|
start_time=start_timestamp,
|
|
end_time=time.time(),
|
|
latency_ms=latency_ms,
|
|
success=True,
|
|
status_code=response.status_code,
|
|
s3_download_ms=timing.get("s3_download_ms"),
|
|
decode_video_ms=timing.get("decode_video_ms"),
|
|
encode_ms=timing.get("encode_ms"),
|
|
decode_ms=timing.get("decode_ms"),
|
|
server_total_ms=timing.get("total_ms"),
|
|
num_chunks=data.get("num_chunks"),
|
|
video_duration=data.get("video_duration"),
|
|
num_scene_changes=data.get("num_scene_changes"),
|
|
)
|
|
else:
|
|
return RequestResult(
|
|
request_id=request_id,
|
|
stream_id=stream_id,
|
|
start_time=start_timestamp,
|
|
end_time=time.time(),
|
|
latency_ms=latency_ms,
|
|
success=False,
|
|
status_code=response.status_code,
|
|
error=response.text[:200],
|
|
)
|
|
except Exception as e:
|
|
end = time.perf_counter()
|
|
latency_ms = (end - start) * 1000
|
|
return RequestResult(
|
|
request_id=request_id,
|
|
stream_id=stream_id,
|
|
start_time=start_timestamp,
|
|
end_time=time.time(),
|
|
latency_ms=latency_ms,
|
|
success=False,
|
|
status_code=0,
|
|
error=str(e)[:200],
|
|
)
|
|
|
|
async def _request_task(self, client: httpx.AsyncClient, request_id: str):
|
|
"""Execute a single request and record the result. Releases semaphore when done."""
|
|
try:
|
|
result = await self._make_request(client, request_id)
|
|
|
|
async with self._lock:
|
|
self.results.append(result)
|
|
self.stats.add_result(result)
|
|
|
|
# Print progress
|
|
status = "✓" if result.success else "✗"
|
|
print(
|
|
f"{status} {result.request_id} | "
|
|
f"{result.latency_ms:7.1f}ms | "
|
|
f"Total: {self.stats.total_requests} "
|
|
f"(OK: {self.stats.successful_requests}, Fail: {self.stats.failed_requests})"
|
|
)
|
|
finally:
|
|
self._semaphore.release()
|
|
|
|
async def run(self):
|
|
"""Run the load test until interrupted."""
|
|
self.start_time = time.time()
|
|
self._semaphore = asyncio.Semaphore(self.concurrency)
|
|
pending_tasks: set[asyncio.Task] = set()
|
|
|
|
print("=" * 70)
|
|
print("🚀 Starting Load Test")
|
|
print("=" * 70)
|
|
print(f" Video: {self.video_path}")
|
|
print(f" URL: {self.url}")
|
|
print(f" Concurrency: {self.concurrency}")
|
|
print(f" Chunk dur: {self.chunk_duration}s")
|
|
print(f" Frames/chunk:{self.num_frames}")
|
|
print()
|
|
print("Press Ctrl+C to stop and save results...")
|
|
print("-" * 70)
|
|
|
|
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
|
while self.running:
|
|
# Block until a slot is available
|
|
await self._semaphore.acquire()
|
|
|
|
if not self.running:
|
|
self._semaphore.release()
|
|
break
|
|
# Spawn task (it will release semaphore when done)
|
|
self.request_counter += 1
|
|
request_id = f"req_{self.request_counter:06d}"
|
|
task = asyncio.create_task(self._request_task(client, request_id))
|
|
pending_tasks.add(task)
|
|
task.add_done_callback(pending_tasks.discard)
|
|
|
|
# Wait for in-flight requests
|
|
if pending_tasks:
|
|
print(f"\n⏳ Waiting for {len(pending_tasks)} in-flight requests...")
|
|
await asyncio.gather(*pending_tasks, return_exceptions=True)
|
|
|
|
def stop(self):
|
|
"""Stop the load test."""
|
|
self.running = False
|
|
|
|
def save_results(self, output_path: str):
|
|
"""Save results to CSV."""
|
|
if not self.results:
|
|
print("No results to save.")
|
|
return
|
|
|
|
fieldnames = [
|
|
"request_id", "stream_id", "start_time", "end_time", "latency_ms",
|
|
"success", "status_code", "error",
|
|
"s3_download_ms", "decode_video_ms", "encode_ms", "decode_ms", "server_total_ms",
|
|
"num_chunks", "video_duration", "num_scene_changes"
|
|
]
|
|
|
|
with open(output_path, "w", newline="") as f:
|
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
for result in self.results:
|
|
writer.writerow(asdict(result))
|
|
|
|
print(f"📁 Results saved to: {output_path}")
|
|
|
|
def print_summary(self):
|
|
"""Print summary statistics."""
|
|
if not self.results:
|
|
print("No results to summarize.")
|
|
return
|
|
|
|
duration = time.time() - self.start_time if self.start_time else 0
|
|
throughput = self.stats.total_requests / duration if duration > 0 else 0
|
|
|
|
print()
|
|
print("=" * 70)
|
|
print("📊 Load Test Summary")
|
|
print("=" * 70)
|
|
print(f" Duration: {duration:.1f}s")
|
|
print(f" Total requests: {self.stats.total_requests}")
|
|
print(f" Successful: {self.stats.successful_requests}")
|
|
print(f" Failed: {self.stats.failed_requests}")
|
|
print(f" Success rate: {self.stats.successful_requests / self.stats.total_requests * 100:.1f}%")
|
|
print(f" Throughput: {throughput:.2f} req/s")
|
|
print()
|
|
print("⏱️ Latency (successful requests):")
|
|
if self.stats.latencies:
|
|
print(f" Min: {self.stats.min_latency_ms:8.1f} ms")
|
|
print(f" Max: {self.stats.max_latency_ms:8.1f} ms")
|
|
print(f" Avg: {self.stats.avg_latency_ms:8.1f} ms")
|
|
print(f" P50: {self.stats.p50_latency_ms:8.1f} ms")
|
|
print(f" P95: {self.stats.p95_latency_ms:8.1f} ms")
|
|
print(f" P99: {self.stats.p99_latency_ms:8.1f} ms")
|
|
else:
|
|
print(" (no successful requests)")
|
|
print("=" * 70)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Load test the Ray Serve video API")
|
|
parser.add_argument("--video", type=str, required=True, help="S3 URI: s3://bucket/key")
|
|
parser.add_argument("--url", type=str, default="http://127.0.0.1:8000", help="Server URL")
|
|
parser.add_argument("--concurrency", type=int, default=4, help="Number of concurrent workers")
|
|
parser.add_argument("--num-frames", type=int, default=16, help="Frames per chunk")
|
|
parser.add_argument("--chunk-duration", type=float, default=10.0, help="Chunk duration in seconds")
|
|
parser.add_argument("--timeout", type=float, default=300.0, help="Request timeout in seconds")
|
|
parser.add_argument("--output", type=str, default=None, help="Output CSV path (default: load_test_<timestamp>.csv)")
|
|
parser.add_argument("--token", type=str, default=None, help="Bearer token for Authorization header")
|
|
args = parser.parse_args()
|
|
|
|
# Generate output path if not provided
|
|
if args.output is None:
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
args.output = f"load_test_{timestamp}_{args.concurrency}.csv"
|
|
|
|
tester = LoadTester(
|
|
video_path=args.video,
|
|
url=args.url,
|
|
concurrency=args.concurrency,
|
|
num_frames=args.num_frames,
|
|
chunk_duration=args.chunk_duration,
|
|
timeout=args.timeout,
|
|
token=args.token,
|
|
)
|
|
|
|
# Track interrupt count for force exit
|
|
interrupt_count = 0
|
|
|
|
# Handle Ctrl+C gracefully
|
|
def signal_handler(sig, frame):
|
|
nonlocal interrupt_count
|
|
interrupt_count += 1
|
|
|
|
if interrupt_count == 1:
|
|
print("\n\n🛑 Stopping load test (press Ctrl+C again to force exit)...")
|
|
tester.stop()
|
|
else:
|
|
print("\n\n⚠️ Force exiting...")
|
|
tester.print_summary()
|
|
tester.save_results(args.output)
|
|
sys.exit(1)
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
|
|
try:
|
|
asyncio.run(tester.run())
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
tester.print_summary()
|
|
tester.save_results(args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|