Spaces:
Sleeping
Sleeping
File size: 2,731 Bytes
8b83582 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """
Send a burst of requests to the FastAPI endpoint to populate the Grafana dashboard.
Usage:
python scripts/load_test.py # default: 30 requests
python scripts/load_test.py --requests 60 # more requests
python scripts/load_test.py --url http://localhost:8000 # custom URL
Requires: pip install requests
"""
import os
import sys
import time
import random
import argparse
import requests
SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "..", "samples")
NOISE_TYPES = ["gaussian", "salt_pepper", "speckle"]
def get_images():
images = [
os.path.join(SAMPLES_DIR, f)
for f in os.listdir(SAMPLES_DIR)
if f.endswith(".png")
]
if not images:
print("No sample images found. Run: python scripts/generate_samples.py")
sys.exit(1)
return images
def send_request(url: str, image_path: str, noise_type: str) -> dict:
with open(image_path, "rb") as f:
files = {"file": (os.path.basename(image_path), f, "image/png")}
data = {"noise_type": noise_type}
t0 = time.perf_counter()
resp = requests.post(f"{url}/predict", files=files, data=data, timeout=60)
elapsed = time.perf_counter() - t0
return {
"status": resp.status_code,
"noise": noise_type,
"image": os.path.basename(image_path),
"time": elapsed,
"ok": resp.status_code == 200,
}
def run(url: str, n_requests: int):
images = get_images()
print(f"Sending {n_requests} requests to {url}")
print(f"Images: {[os.path.basename(i) for i in images]}")
print("-" * 55)
ok = err = 0
for i in range(1, n_requests + 1):
image = random.choice(images)
noise_type = random.choice(NOISE_TYPES)
try:
r = send_request(url, image, noise_type)
status = "OK " if r["ok"] else "ERR"
print(f"[{i:3d}/{n_requests}] {status} {noise_type:12s} "
f"{r['image']:20s} {r['time']:.2f}s")
if r["ok"]:
ok += 1
else:
err += 1
except Exception as e:
print(f"[{i:3d}/{n_requests}] FAIL {noise_type:12s} {e}")
err += 1
# Vary delay so Grafana graphs look natural, not like a flat line
delay = random.uniform(0.3, 1.2)
time.sleep(delay)
print("-" * 55)
print(f"Done — {ok} success / {err} errors out of {n_requests} requests")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="http://localhost:8000")
parser.add_argument("--requests", type=int, default=30)
args = parser.parse_args()
run(args.url, args.requests)
|