Spaces:
Sleeping
Sleeping
| """ | |
| 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) | |