pool_node.py

#!/usr/bin/env python3
"""QuickSilver Pro compute-pool node — reference client.

Contributes one machine to the QuickSilver Pro compute pool. It does NOT run a
model; it is a thin adapter that connects your machine to the pool and forwards
each inbound request to a local OpenAI-compatible server you already run (vLLM,
llama.cpp's server, Ollama, LM Studio, MLX, TGI — anything that speaks
/v1/chat/completions). Point it at that server with --upstream.

There is no Apple/Mac/MLX assumption anywhere in the protocol: a Linux box with
an NVIDIA GPU serving GLM through vLLM contributes exactly the same way a Mac
does.

How it works (the whole protocol, in one paragraph): you register once with your
account's provider key and get back a short-lived share key plus the relay's
coordinates. You open ONE outbound WebSocket to the relay — no inbound ports,
no port forwarding. The relay pushes each customer request down that socket as a
JSON frame; you replay it against your local server and stream the response back
up the same socket. A heartbeat POST every few seconds makes your node's spare
capacity visible. That's it: no polling, no job queue, no result upload.

Dependencies: `pip install websockets httpx`. Python 3.9+.

Run:
    export QSP_PROVIDER_KEY=qsppk-...            # from the dashboard, shown once
    python pool_node.py --model glm-5.3-flash --upstream http://localhost:8000

If your local server serves the model under a different name than the pool's
catalog id (e.g. your vLLM reports "GLM-5.3-Flash-EXL3" but the pool calls it
"glm-5.3-flash"), add --local-model GLM-5.3-Flash-EXL3. The client rewrites the
model name outbound to your server and back to the catalog id in the response,
so the pool's readiness probe — which requires the response to echo the catalog
id exactly — passes. Without this, a name mismatch means your node connects but
never receives traffic.
"""
from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import platform
import random
import socket
import sys
from typing import Any, Optional

try:
    import httpx
    import websockets
except ImportError:
    sys.exit(
        "missing dependencies. Run:  pip install websockets httpx\n"
        "(both are pure-Python and install in seconds.)"
    )

# Cloudflare fronts the QuickSilver hosts and rejects requests carrying a default
# HTTP-library User-Agent (Python's included) with a 403 whose body is HTML
# 'error code: 1010' — before the request ever reaches the API, so there is no
# JSON error to read. A product User-Agent on EVERY call (register, heartbeat,
# and the WebSocket connect) avoids it. Do not remove this.
VERSION = "0.1.0"
USER_AGENT = f"quicksilver-pool-node/{VERSION} (+https://quicksilverpro.io)"

DEFAULT_REGISTER_URL = "https://pay.quicksilverpro.io/v1/pool/nodes/register"


def log(msg: str) -> None:
    print(msg, flush=True)


# --------------------------------------------------------------------------- #
# Registration                                                                #
# --------------------------------------------------------------------------- #

class Registration:
    """The coordinates the relay hands back. Everything the node needs to run is
    here — never hard-code a host; the relay may move behind these fields."""

    def __init__(self, data: dict):
        self.node_id: str = data["node_id"]
        self.share_key: str = data["share_key"]
        self.model: str = data["model"]
        self.relay_url: str = data["relay_url"]          # wss://<relay>/up
        self.heartbeat_url: str = data["heartbeat_url"]
        self.heartbeat_interval_s: float = float(data.get("heartbeat_interval_s", 10))


async def register(client: httpx.AsyncClient, args: argparse.Namespace) -> Registration:
    """POST /v1/pool/nodes/register. Idempotent for a given (account, model,
    worker): re-calling it returns the same node_id and ROTATES the share key,
    so we always use the value from the most recent response."""
    body = {
        "model": args.model,
        "worker": args.worker,
        "max_concurrency": args.max_concurrency,
        "hardware": {
            "host": platform.node(),
            "platform": platform.platform(),
            "python": platform.python_version(),
            "client": f"quicksilver-pool-node/{VERSION}",
        },
    }
    resp = await client.post(
        args.register_url,
        headers={"Authorization": f"Bearer {args.provider_key}", "User-Agent": USER_AGENT},
        json=body,
        timeout=30,
    )
    if resp.status_code != 200:
        # Registration errors use {"error": {"code", "message"}}; a body that
        # failed schema binding comes back as FastAPI's {"detail": ...}. Read
        # both shapes so the message is useful.
        detail = ""
        try:
            j = resp.json()
            detail = (j.get("error") or {}).get("message") or j.get("detail") or ""
        except Exception:
            detail = (resp.text or "")[:200]
        raise SystemExit(f"registration failed: HTTP {resp.status_code} {detail}")
    reg = Registration(resp.json())
    log(f"registered node {reg.node_id} for model {reg.model}")
    return reg


# --------------------------------------------------------------------------- #
# Request forwarding                                                           #
# --------------------------------------------------------------------------- #

def _rewrite_model(payload: bytes, frm: str, to: str) -> bytes:
    """Swap a JSON body's top-level `model` field frm→to. Used to bridge the
    catalog id the pool speaks and the name your local server serves under.
    Returns the body unchanged if it is not JSON or carries no matching model."""
    try:
        obj = json.loads(payload)
    except Exception:
        return payload
    if isinstance(obj, dict) and obj.get("model") == frm:
        obj["model"] = to
        return json.dumps(obj).encode("utf-8")
    return payload


def _rewrite_sse_model(chunk: bytes, frm: str, to: str) -> bytes:
    """Same swap, applied to the `model` field inside SSE `data:` lines of a
    streamed completion. Cheap string replace scoped to the exact JSON token so
    it cannot touch message content."""
    if frm == to or b'"model"' not in chunk:
        return chunk
    return chunk.replace(
        f'"model":"{frm}"'.encode(), f'"model":"{to}"'.encode()
    ).replace(
        f'"model": "{frm}"'.encode(), f'"model": "{to}"'.encode()
    )


async def handle_request(
    ws, upstream: httpx.AsyncClient, upstream_base: str,
    catalog_model: str, local_model: Optional[str], frame: dict,
) -> None:
    """Replay one `req` frame against the local server and stream the response
    back as head / chunk* / (end | err). One of these runs per inbound request;
    the pool bounds how many at once via max_concurrency."""
    req_id = frame["id"]
    try:
        path = frame.get("path", "/v1/chat/completions")
        method = frame.get("method", "POST")
        headers = dict(frame.get("headers") or {})
        headers.pop("Content-Length", None)
        headers.pop("content-length", None)
        body = base64.b64decode(frame.get("body") or b"")

        # Outbound rewrite: catalog id -> the name the local server serves under.
        if local_model:
            body = _rewrite_model(body, catalog_model, local_model)

        url = upstream_base.rstrip("/") + path
        async with upstream.stream(method, url, headers=headers, content=body, timeout=None) as resp:
            await _send(ws, {"t": "head", "id": req_id, "status": resp.status_code,
                             "headers": dict(resp.headers)})
            async for chunk in resp.aiter_bytes():
                if not chunk:
                    continue
                if local_model:
                    chunk = _rewrite_sse_model(chunk, local_model, catalog_model)
                await _send(ws, {"t": "chunk", "id": req_id,
                                 "data": base64.b64encode(chunk).decode("ascii")})
        await _send(ws, {"t": "end", "id": req_id})
    except asyncio.CancelledError:
        # The relay sent `abort` (the customer gave up). Stop; the upstream
        # stream is closed by the context manager. Do not send `end`.
        raise
    except Exception as exc:  # noqa: BLE001 — any upstream failure is one request's problem
        await _send(ws, {"t": "err", "id": req_id, "msg": f"{type(exc).__name__}: {exc}"})


async def _send(ws, obj: dict) -> None:
    await ws.send(json.dumps(obj))


# --------------------------------------------------------------------------- #
# Heartbeat                                                                    #
# --------------------------------------------------------------------------- #

async def heartbeat_loop(client: httpx.AsyncClient, reg: Registration, inflight) -> None:
    """POST the heartbeat at the interval the relay asked for. It is pure
    observability — it never gates traffic — but reporting `inflight` is what
    makes the node's spare capacity visible. A 401/404 means the share key was
    rotated or the node removed: raise so the outer loop re-registers."""
    while True:
        await asyncio.sleep(reg.heartbeat_interval_s)
        try:
            resp = await client.post(
                reg.heartbeat_url,
                headers={"Authorization": f"Bearer {reg.share_key}", "User-Agent": USER_AGENT},
                json={"inflight": len(inflight), "client": f"quicksilver-pool-node/{VERSION}"},
                timeout=15,
            )
            if resp.status_code in (401, 404):
                raise ReRegister(f"heartbeat {resp.status_code}: share key no longer bound")
        except (httpx.HTTPError,) as exc:
            log(f"heartbeat transient error ({exc}); will retry")


class ReRegister(Exception):
    """Signals the outer loop to register again (rotates the share key)."""


# --------------------------------------------------------------------------- #
# Tunnel                                                                       #
# --------------------------------------------------------------------------- #

async def run_tunnel(
    client: httpx.AsyncClient, upstream: httpx.AsyncClient,
    reg: Registration, args: argparse.Namespace,
) -> None:
    """Open the tunnel and serve until it drops. The share key goes in the URL
    PATH — treat the assembled URL as a secret (never log it)."""
    url = f"{reg.relay_url}/{reg.share_key}?id={reg.node_id}"
    inflight: dict[str, asyncio.Task] = {}

    async with websockets.connect(
        url, additional_headers={"User-Agent": USER_AGENT},
        max_size=None, ping_interval=20, ping_timeout=20,
    ) as ws:
        await _send(ws, {"t": "ready", "v": 1})
        log(f"tunnel up; serving {reg.model} from {args.upstream}")
        hb = asyncio.create_task(heartbeat_loop(client, reg, inflight))
        try:
            async for raw in ws:
                if isinstance(raw, bytes):
                    raw = raw.decode("utf-8", "replace")
                try:
                    frame = json.loads(raw)
                except Exception:
                    continue
                t = frame.get("t")
                if t == "req":
                    req_id = frame.get("id")
                    if not req_id:
                        continue
                    task = asyncio.create_task(
                        handle_request(ws, upstream, args.upstream, reg.model,
                                       args.local_model, frame)
                    )
                    inflight[req_id] = task
                    task.add_done_callback(lambda _t, i=req_id: inflight.pop(i, None))
                elif t == "abort":
                    task = inflight.get(frame.get("id"))
                    if task:
                        task.cancel()
        finally:
            hb.cancel()
            for task in list(inflight.values()):
                task.cancel()


# --------------------------------------------------------------------------- #
# Main loop                                                                    #
# --------------------------------------------------------------------------- #

async def main_async(args: argparse.Namespace) -> None:
    backoff = 1.0
    async with httpx.AsyncClient() as client, httpx.AsyncClient() as upstream:
        reg = await register(client, args)
        while True:
            try:
                await run_tunnel(client, upstream, reg, args)
                # A clean tunnel close is a transient drop: reconnect with the
                # key we already hold before re-registering.
                log("tunnel closed; reconnecting")
                backoff = 1.0
            except ReRegister as exc:
                log(f"re-registering: {exc}")
                reg = await register(client, args)
                backoff = 1.0
                continue
            except (websockets.InvalidStatus, websockets.InvalidStatusCode) if hasattr(
                websockets, "InvalidStatusCode"
            ) else websockets.InvalidStatus as exc:  # type: ignore[misc]
                # 401 on connect => the share key is no longer bound; re-register.
                # Anything else (e.g. 409 claim conflict) is retryable with backoff.
                status = getattr(getattr(exc, "response", None), "status_code", None)
                if status == 401:
                    log("tunnel 401; re-registering")
                    reg = await register(client, args)
                    backoff = 1.0
                    continue
                log(f"tunnel connect failed ({exc}); retrying")
            except Exception as exc:  # noqa: BLE001
                log(f"tunnel error ({exc}); retrying")
            # Reconnect with jittered backoff, capped.
            await asyncio.sleep(backoff + random.uniform(0, backoff))
            backoff = min(backoff * 2, 30.0)


def parse_args(argv: Optional[list] = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Contribute a machine to the QuickSilver Pro compute pool.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--model", required=True,
                   help="Pool catalog id to serve, e.g. glm-5.3-flash.")
    p.add_argument("--upstream", default="http://localhost:8000",
                   help="Base URL of your local OpenAI-compatible server. The "
                        "relay's request path (e.g. /v1/chat/completions) is "
                        "appended verbatim, so give the host WITHOUT /v1.")
    p.add_argument("--local-model", default=None,
                   help="If your server serves the model under a different name "
                        "than --model, its local name. The client rewrites the "
                        "model field both ways so the readiness probe passes.")
    p.add_argument("--provider-key", default=os.environ.get("QSP_PROVIDER_KEY", ""),
                   help="Account provider key (qsppk-...). Prefer the "
                        "QSP_PROVIDER_KEY env var over passing it on the CLI.")
    p.add_argument("--worker", default=socket.gethostname(),
                   help="Per-machine label; node identity is (account, model, "
                        "worker). Two machines on one account need distinct values.")
    p.add_argument("--max-concurrency", type=int, default=2,
                   help="Concurrent request slots to advertise (clamped 1–64).")
    p.add_argument("--register-url", default=DEFAULT_REGISTER_URL,
                   help="Registration endpoint (the pay.* host, not api.*).")
    args = p.parse_args(argv)
    if not args.provider_key.startswith("qsppk-"):
        p.error("a provider key (qsppk-...) is required via --provider-key or "
                "QSP_PROVIDER_KEY. Create one in the dashboard; it is shown once.")
    return args


def main() -> None:
    args = parse_args()
    try:
        asyncio.run(main_async(args))
    except KeyboardInterrupt:
        log("\nstopped.")


if __name__ == "__main__":
    main()