test_scaling.py

"""
test_scaling.py — RunPod load-balancing worker scaling test
============================================================

Reproduces the user's real-world topology:

    Client (Python) ──WS, Bearer + optional ?token=──> Traefik (or direct) ──> FastAPI/uvicorn

Tests three things:

1. HTTP health-check polling  — does repeated GET /ping trigger scaling?
2. WebSocket-only load        — do concurrent WS connections trigger scaling?
3. Mixed HTTP + WS            — does a POST /generate alongside a WS session
                                trigger scaling as expected?

Usage
-----
Local (no RunPod):
    python3 app.py &         # start the worker on localhost:80
    python3 test_scaling.py --url http://localhost:80

Against a RunPod load-balancing endpoint (Authorization: Bearer key on /ping, /generate, WS):
    python3 test_scaling.py \
        --url https://ENDPOINT_ID.api.runpod.ai \
        --token "$RUNPOD_API_KEY" \
        --connections 10 \
        --duration 30

Concurrent WebSocket test only (skip /ping and mixed HTTP+WS):

    python3 test_scaling.py --url ... --token ... --ws-only --connections 20

The script prints a summary table at the end showing whether each test
category received responses — use this to confirm whether RunPod is
counting WS connections toward its scaling metric.
"""

import argparse
import asyncio
import json
import statistics
import time
from dataclasses import dataclass, field
from typing import Optional

import httpx
import websockets
import websockets.exceptions


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--url", default="http://localhost:80",
                   help="Base URL of the worker, e.g. https://abc123.api.runpod.ai")
    p.add_argument("--token", default=None,
                   help="API key or bearer token: Authorization: Bearer on HTTP and WS; "
                        "also ?token= on WS URL for workers that read the query string")
    p.add_argument("--connections", type=int, default=5,
                   help="Number of concurrent WebSocket connections to open")
    p.add_argument("--duration", type=int, default=20,
                   help="How long (seconds) to hold each WS connection open")
    p.add_argument("--health-polls", type=int, default=30,
                   help="Number of GET /ping polls during the health-check test")
    p.add_argument("--health-interval", type=float, default=1.0,
                   help="Seconds between health-check polls")
    p.add_argument("--ws-only", action="store_true",
                   help="Run only the concurrent WebSocket test (skip /ping and mixed HTTP+WS)")
    p.add_argument("--ws-timeout", type=float, default=60.0,
                   help="WebSocket open_timeout in seconds (default 60 — allows for cold starts)")
    p.add_argument("--stagger", type=float, default=0.3,
                   help="Seconds between launching each WS connection (default 0.3)")
    p.add_argument("--warmup", action="store_true", default=True,
                   help="Poll /ping until a worker responds 200 before opening WS connections")
    p.add_argument("--no-warmup", dest="warmup", action="store_false",
                   help="Skip the /ping warmup wait")
    p.add_argument("--warmup-timeout", type=float, default=120.0,
                   help="Max seconds to wait for a worker to become healthy (default 120)")
    return p.parse_args()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def bearer_headers(token: Optional[str]) -> dict[str, str]:
    if not token:
        return {}
    return {"Authorization": f"Bearer {token}"}


def ws_url(base: str, path: str, token: Optional[str]) -> str:
    """Convert http(s) base URL to ws(s) and append path + optional token."""
    base = base.rstrip("/")
    if base.startswith("https://"):
        ws_base = "wss://" + base[len("https://"):]
    elif base.startswith("http://"):
        ws_base = "ws://" + base[len("http://"):]
    else:
        ws_base = base
    url = f"{ws_base}{path}"
    if token:
        url += f"?token={token}"
    return url


async def wait_for_worker(
    base_url: str,
    token: Optional[str],
    timeout: float = 120.0,
    poll_interval: float = 3.0,
) -> bool:
    """Poll GET /ping until the worker returns 200 or timeout is reached.

    Returns True when a healthy worker is found, False on timeout.
    Workers return 204 while still initialising, then 200 when ready.
    """
    deadline = time.monotonic() + timeout
    attempt = 0
    async with httpx.AsyncClient(timeout=10, headers=bearer_headers(token)) as client:
        while time.monotonic() < deadline:
            attempt += 1
            try:
                r = await client.get(f"{base_url}/ping")
                if r.status_code == 200:
                    print(f"  worker ready after {attempt} poll(s)")
                    return True
                print(f"  poll {attempt}: status {r.status_code} (still initialising) ...")
            except Exception as exc:
                print(f"  poll {attempt}: {exc} — retrying ...")
            await asyncio.sleep(poll_interval)
    return False


@dataclass
class TestResult:
    name: str
    success: bool
    notes: list[str] = field(default_factory=list)
    latencies_ms: list[float] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Test 1: HTTP health-check polling
# (reproduces the reverse-proxy behaviour that does NOT scale workers)
# ---------------------------------------------------------------------------

async def test_health_polling(
    base_url: str,
    polls: int,
    interval: float,
    token: Optional[str],
) -> TestResult:
    result = TestResult(name="HTTP /ping polling", success=False)
    print(f"\n[1] Sending {polls} GET {base_url}/ping polls (interval={interval}s) ...")

    async with httpx.AsyncClient(timeout=10, headers=bearer_headers(token)) as client:
        for i in range(polls):
            t0 = time.monotonic()
            try:
                r = await client.get(f"{base_url}/ping")
                elapsed = (time.monotonic() - t0) * 1000
                result.latencies_ms.append(elapsed)
                status = r.status_code
                if status == 200:
                    result.success = True
                    result.notes.append(f"poll {i+1}: {status} ({elapsed:.1f} ms)")
                else:
                    result.notes.append(f"poll {i+1}: unexpected status {status}")
            except Exception as exc:
                result.notes.append(f"poll {i+1}: ERROR — {exc}")
            await asyncio.sleep(interval)

    if result.latencies_ms:
        result.notes.append(
            f"latency p50={statistics.median(result.latencies_ms):.1f} ms  "
            f"p95={sorted(result.latencies_ms)[int(len(result.latencies_ms)*0.95)]:.1f} ms"
        )
    return result


# ---------------------------------------------------------------------------
# Test 2: Concurrent WebSocket connections (WS-only load)
# ---------------------------------------------------------------------------

async def _ws_session(
    url: str,
    session_id: int,
    duration: float,
    results: list[dict],
    token: Optional[str] = None,
    open_timeout: float = 60.0,
    max_retries: int = 3,
) -> None:
    """Open a WS, send prompts every 2 s, collect responses.

    Retries the connection up to max_retries times to handle cases where
    the worker is still warming up when this session starts.
    """
    session_results: list[dict] = []
    t_start = time.monotonic()

    connect_kwargs: dict = {"open_timeout": open_timeout}
    if token:
        connect_kwargs["additional_headers"] = [("Authorization", f"Bearer {token}")]

    last_exc: Optional[Exception] = None
    for attempt in range(1, max_retries + 1):
        try:
            async with websockets.connect(url, **connect_kwargs) as ws:
                session_results.append({"session": session_id, "event": "connected", "t": 0.0})

                async def keep_sending():
                    seq = 0
                    while time.monotonic() - t_start < duration:
                        payload = {"prompt": f"session {session_id} message {seq}", "max_tokens": 10}
                        await ws.send(json.dumps(payload))
                        seq += 1
                        await asyncio.sleep(2)

                async def keep_receiving():
                    while time.monotonic() - t_start < duration:
                        try:
                            raw = await asyncio.wait_for(ws.recv(), timeout=5)
                            data = json.loads(raw)
                            session_results.append({
                                "session": session_id,
                                "event": "recv",
                                "data": data,
                                "t": round(time.monotonic() - t_start, 2),
                            })
                        except asyncio.TimeoutError:
                            pass

                await asyncio.gather(keep_sending(), keep_receiving())
                break  # success — no more retries needed

        except websockets.exceptions.InvalidStatus as exc:
            # 4xx/5xx from the server — retrying won't help
            session_results.append({"session": session_id, "event": "error", "detail": str(exc)})
            break
        except Exception as exc:
            last_exc = exc
            if attempt < max_retries:
                backoff = 2 ** attempt
                session_results.append({
                    "session": session_id,
                    "event": "retry",
                    "attempt": attempt,
                    "detail": str(exc),
                })
                await asyncio.sleep(backoff)
            else:
                session_results.append({"session": session_id, "event": "error", "detail": str(exc)})

    results.extend(session_results)


async def test_concurrent_websockets(
    base_url: str,
    token: Optional[str],
    n_connections: int,
    duration: float,
    open_timeout: float = 60.0,
    stagger: float = 0.3,
    do_warmup: bool = True,
    warmup_timeout: float = 120.0,
) -> TestResult:
    result = TestResult(
        name=f"Concurrent WebSocket connections ({n_connections}x {duration}s)",
        success=False,
    )
    url = ws_url(base_url, "/ws", token)
    print(f"\n[2] Opening {n_connections} concurrent WS connections to {url} for {duration}s ...")

    if do_warmup:
        print(f"  Waiting for a healthy worker (up to {warmup_timeout:.0f}s) ...")
        ready = await wait_for_worker(base_url, token, timeout=warmup_timeout)
        if not ready:
            result.notes.append(
                f"WARMUP TIMEOUT: no healthy worker found after {warmup_timeout:.0f}s — "
                "all WS connections will likely time out too."
            )
        else:
            result.notes.append("warmup: worker healthy before WS connections opened")

    # Stagger connection opens to avoid a thundering-herd timeout on the same worker
    all_events: list[dict] = []

    async def staggered_session(i: int) -> None:
        await asyncio.sleep(i * stagger)
        await _ws_session(url, i, duration, all_events, token=token, open_timeout=open_timeout)

    tasks = [staggered_session(i) for i in range(n_connections)]
    await asyncio.gather(*tasks)

    connected = sum(1 for e in all_events if e["event"] == "connected")
    responses = sum(1 for e in all_events if e["event"] == "recv")
    errors = sum(1 for e in all_events if e["event"] == "error")
    retries = sum(1 for e in all_events if e["event"] == "retry")

    result.success = connected > 0
    result.notes.append(f"sessions connected:    {connected}/{n_connections}")
    result.notes.append(f"total frames received: {responses}")
    result.notes.append(f"connection errors:     {errors}")
    if retries:
        result.notes.append(f"connect retries:       {retries}")

    if errors > 0:
        error_details = [e["detail"] for e in all_events if e["event"] == "error"]
        for detail in set(error_details):
            result.notes.append(f"  error detail: {detail}")

    # Scaling diagnosis hint
    if connected == 0:
        result.notes.append(
            "DIAGNOSIS: No WS connections succeeded. "
            "This strongly suggests RunPod is not treating WS connections as active "
            "requests for scaling purposes — workers may have scaled to 0."
        )
    elif responses == 0:
        result.notes.append(
            "DIAGNOSIS: Connections opened but no data received. "
            "Workers are alive but may not be counted toward scaling metrics."
        )
    else:
        result.notes.append(
            "DIAGNOSIS: WS connections and data flow OK. "
            "If you still see scaling issues, the problem is likely at the proxy layer."
        )

    return result


# ---------------------------------------------------------------------------
# Test 3: Mixed HTTP + WS (does HTTP POST /generate trigger scaling?)
# ---------------------------------------------------------------------------

async def test_mixed_http_ws(
    base_url: str,
    token: Optional[str],
    duration: float,
    open_timeout: float = 60.0,
) -> TestResult:
    result = TestResult(name="Mixed HTTP POST + WebSocket", success=False)
    ws_uri = ws_url(base_url, "/ws", token)
    print(f"\n[3] Mixed test — HTTP POST /generate + 1 WS connection for {duration}s ...")

    http_results: list[tuple[int, float]] = []
    ws_events: list[dict] = []

    async def http_loop():
        async with httpx.AsyncClient(timeout=30, headers=bearer_headers(token)) as client:
            t0 = time.monotonic()
            seq = 0
            while time.monotonic() - t0 < duration:
                t_req = time.monotonic()
                try:
                    r = await client.post(
                        f"{base_url}/generate",
                        json={"prompt": f"http test {seq}", "max_tokens": 10},
                    )
                    elapsed = (time.monotonic() - t_req) * 1000
                    http_results.append((r.status_code, elapsed))
                    result.latencies_ms.append(elapsed)
                except Exception as exc:
                    http_results.append((-1, 0.0))
                    result.notes.append(f"HTTP error: {exc}")
                seq += 1
                await asyncio.sleep(2)

    await asyncio.gather(
        http_loop(),
        _ws_session(
            ws_uri,
            session_id=0,
            duration=duration,
            results=ws_events,
            token=token,
            open_timeout=open_timeout,
        ),
    )

    ok_http = sum(1 for s, _ in http_results if s == 200)
    ws_connected = any(e["event"] == "connected" for e in ws_events)
    ws_frames = sum(1 for e in ws_events if e["event"] == "recv")

    result.success = ok_http > 0
    result.notes.append(f"HTTP 200 responses:   {ok_http}/{len(http_results)}")
    result.notes.append(f"WS connected:         {ws_connected}")
    result.notes.append(f"WS frames received:   {ws_frames}")

    if result.latencies_ms:
        result.notes.append(
            f"HTTP latency p50={statistics.median(result.latencies_ms):.1f} ms  "
            f"p95={sorted(result.latencies_ms)[int(len(result.latencies_ms)*0.95)]:.1f} ms"
        )

    return result


# ---------------------------------------------------------------------------
# Summary printer
# ---------------------------------------------------------------------------

def print_summary(results: list[TestResult]) -> None:
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    for r in results:
        status = "PASS" if r.success else "FAIL"
        print(f"\n[{status}] {r.name}")
        for note in r.notes:
            print(f"       {note}")
    print("\n" + "=" * 60)
    print(
        "\nNOTE: If WS-only tests fail to scale but HTTP tests succeed,\n"
        "RunPod's load-balancer is likely not counting open WebSocket\n"
        "connections as active work. Workaround: send a lightweight\n"
        "HTTP keepalive (e.g. GET /ping every 30 s) from the client\n"
        "or proxy to keep workers alive, or switch to min_workers >= 1.\n"
    )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

async def main() -> None:
    args = parse_args()
    base = args.url.rstrip("/")

    results: list[TestResult] = []

    if not args.ws_only:
        results.append(
            await test_health_polling(base, args.health_polls, args.health_interval, args.token)
        )

    results.append(
        await test_concurrent_websockets(
            base,
            args.token,
            args.connections,
            args.duration,
            open_timeout=args.ws_timeout,
            stagger=args.stagger,
            do_warmup=args.warmup,
            warmup_timeout=args.warmup_timeout,
        )
    )

    if not args.ws_only:
        results.append(
            await test_mixed_http_ws(base, args.token, args.duration, open_timeout=args.ws_timeout)
        )

    print_summary(results)


if __name__ == "__main__":
    asyncio.run(main())