|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Background monitor for multi-key rate limit test. |
| 4 | +Polls API key usage 5 times per second and logs to file. |
| 5 | +""" |
| 6 | +import asyncio |
| 7 | +import json |
| 8 | +import time |
| 9 | +import sys |
| 10 | +from dataclasses import dataclass, asdict |
| 11 | +from datetime import datetime |
| 12 | + |
| 13 | +import httpx |
| 14 | + |
| 15 | +BASE_URL = "http://localhost:3000/api" |
| 16 | +ADMIN_SECRET = "admin" |
| 17 | +OUTPUT_FILE = "/tmp/rate_limit_monitor.jsonl" |
| 18 | +POLL_INTERVAL = 0.2 # 5 times per second |
| 19 | + |
| 20 | + |
| 21 | +@dataclass |
| 22 | +class UsageSnapshot: |
| 23 | + timestamp: float |
| 24 | + datetime_str: str |
| 25 | + api_key_id: int |
| 26 | + api_key_short: str |
| 27 | + rpm_current: int |
| 28 | + rpm_limit: int |
| 29 | + rpm_remaining: int |
| 30 | + tpm_current: int |
| 31 | + tpm_limit: int |
| 32 | + tpm_remaining: int |
| 33 | + |
| 34 | + |
| 35 | +async def get_all_api_keys(client: httpx.AsyncClient) -> list[dict]: |
| 36 | + """Get all API keys""" |
| 37 | + try: |
| 38 | + resp = await client.get( |
| 39 | + f"{BASE_URL}/admin/apiKey", |
| 40 | + headers={"Authorization": f"Bearer {ADMIN_SECRET}"}, |
| 41 | + timeout=5.0 |
| 42 | + ) |
| 43 | + if resp.status_code == 200: |
| 44 | + return resp.json() |
| 45 | + except Exception as e: |
| 46 | + print(f"[Monitor] Error getting API keys: {e}", file=sys.stderr) |
| 47 | + return [] |
| 48 | + |
| 49 | + |
| 50 | +async def get_usage(client: httpx.AsyncClient, key_info: dict) -> UsageSnapshot | None: |
| 51 | + """Get usage for a single API key""" |
| 52 | + try: |
| 53 | + resp = await client.get( |
| 54 | + f"{BASE_URL}/admin/apiKey/{key_info['key']}/usage", |
| 55 | + headers={"Authorization": f"Bearer {ADMIN_SECRET}"}, |
| 56 | + timeout=5.0 |
| 57 | + ) |
| 58 | + if resp.status_code == 200: |
| 59 | + data = resp.json() |
| 60 | + now = time.time() |
| 61 | + return UsageSnapshot( |
| 62 | + timestamp=now, |
| 63 | + datetime_str=datetime.fromtimestamp(now).strftime("%H:%M:%S.%f")[:-3], |
| 64 | + api_key_id=key_info["id"], |
| 65 | + api_key_short=key_info["key"][:15] + "...", |
| 66 | + rpm_current=data.get("usage", {}).get("rpm", {}).get("current", 0), |
| 67 | + rpm_limit=data.get("limits", {}).get("rpm", 0), |
| 68 | + rpm_remaining=data.get("usage", {}).get("rpm", {}).get("remaining", 0), |
| 69 | + tpm_current=data.get("usage", {}).get("tpm", {}).get("current", 0), |
| 70 | + tpm_limit=data.get("limits", {}).get("tpm", 0), |
| 71 | + tpm_remaining=data.get("usage", {}).get("tpm", {}).get("remaining", 0), |
| 72 | + ) |
| 73 | + except Exception as e: |
| 74 | + pass # Silent fail for monitoring |
| 75 | + return None |
| 76 | + |
| 77 | + |
| 78 | +async def monitor_loop(stop_file: str = "/tmp/stop_monitor"): |
| 79 | + """Main monitoring loop""" |
| 80 | + print(f"[Monitor] Starting rate limit monitor") |
| 81 | + print(f"[Monitor] Output file: {OUTPUT_FILE}") |
| 82 | + print(f"[Monitor] Poll interval: {POLL_INTERVAL}s (5x per second)") |
| 83 | + print(f"[Monitor] Stop by creating file: {stop_file}") |
| 84 | + print() |
| 85 | + |
| 86 | + # Clear output file |
| 87 | + with open(OUTPUT_FILE, "w") as f: |
| 88 | + f.write("") |
| 89 | + |
| 90 | + snapshot_count = 0 |
| 91 | + last_keys_check = 0 |
| 92 | + api_keys: list[dict] = [] |
| 93 | + |
| 94 | + async with httpx.AsyncClient() as client: |
| 95 | + while True: |
| 96 | + # Check stop condition |
| 97 | + import os |
| 98 | + if os.path.exists(stop_file): |
| 99 | + print(f"\n[Monitor] Stop file detected, exiting...") |
| 100 | + try: |
| 101 | + os.remove(stop_file) |
| 102 | + except: |
| 103 | + pass |
| 104 | + break |
| 105 | + |
| 106 | + # Refresh API keys list every 2 seconds |
| 107 | + now = time.time() |
| 108 | + if now - last_keys_check > 2: |
| 109 | + api_keys = await get_all_api_keys(client) |
| 110 | + last_keys_check = now |
| 111 | + if api_keys: |
| 112 | + # Filter to recent test keys (created in last 10 minutes) |
| 113 | + # This is a simple heuristic based on comment pattern |
| 114 | + test_keys = [k for k in api_keys if k.get("comment", "").startswith("test-multikey-")] |
| 115 | + if test_keys: |
| 116 | + api_keys = test_keys |
| 117 | + print(f"[Monitor] Tracking {len(api_keys)} test API keys: {[k['id'] for k in api_keys]}") |
| 118 | + |
| 119 | + if not api_keys: |
| 120 | + await asyncio.sleep(POLL_INTERVAL) |
| 121 | + continue |
| 122 | + |
| 123 | + # Get usage for all keys concurrently |
| 124 | + tasks = [get_usage(client, k) for k in api_keys] |
| 125 | + snapshots = await asyncio.gather(*tasks, return_exceptions=True) |
| 126 | + |
| 127 | + # Write snapshots to file |
| 128 | + with open(OUTPUT_FILE, "a") as f: |
| 129 | + for snapshot in snapshots: |
| 130 | + if isinstance(snapshot, UsageSnapshot): |
| 131 | + f.write(json.dumps(asdict(snapshot)) + "\n") |
| 132 | + snapshot_count += 1 |
| 133 | + |
| 134 | + # Print summary every 10 snapshots |
| 135 | + if snapshot_count % 10 == 0: |
| 136 | + print(f"[Monitor] {snapshot.datetime_str} Key-{snapshot.api_key_id}: " |
| 137 | + f"RPM {snapshot.rpm_current}/{snapshot.rpm_limit} " |
| 138 | + f"(rem={snapshot.rpm_remaining}), " |
| 139 | + f"TPM {snapshot.tpm_current}/{snapshot.tpm_limit}") |
| 140 | + |
| 141 | + await asyncio.sleep(POLL_INTERVAL) |
| 142 | + |
| 143 | + print(f"\n[Monitor] Total snapshots collected: {snapshot_count}") |
| 144 | + print(f"[Monitor] Output saved to: {OUTPUT_FILE}") |
| 145 | + |
| 146 | + |
| 147 | +async def main(): |
| 148 | + # Remove stop file if exists |
| 149 | + import os |
| 150 | + stop_file = "/tmp/stop_monitor" |
| 151 | + if os.path.exists(stop_file): |
| 152 | + os.remove(stop_file) |
| 153 | + |
| 154 | + try: |
| 155 | + await monitor_loop(stop_file) |
| 156 | + except KeyboardInterrupt: |
| 157 | + print("\n[Monitor] Interrupted by user") |
| 158 | + except Exception as e: |
| 159 | + print(f"\n[Monitor] Error: {e}") |
| 160 | + |
| 161 | + |
| 162 | +if __name__ == "__main__": |
| 163 | + asyncio.run(main()) |
0 commit comments