chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,374 @@
|
||||
#!/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()
|
||||
|
||||
Reference in New Issue
Block a user