aiondemandcluster 0.1.0__py3-none-any.whl

aiod/__init__.py

@@ -0,0 +1,4 @@
+"""AI on Demand — spin up a HuggingFace model on vast.ai GPUs and drive it from
+Claude Code via Claude Code Router."""
+
+__version__ = "0.1.0"

aiod/bench.py

@@ -0,0 +1,173 @@
+"""A small benchmark for a running OpenAI-compatible endpoint.
+
+Measures the numbers you actually want to report:
+  * TTFT  — time to first token (streaming), p50/p95
+  * decode speed — output tokens/sec per request
+  * throughput   — aggregate output tokens/sec under concurrency
+  * $/1M tokens  — derived from the instance price and throughput
+
+Works against vLLM or llama.cpp (both stream OpenAI-style SSE).
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+
+import httpx
+
+DEFAULT_PROMPT = (
+    "Write a Python function `merge_intervals(intervals)` that merges overlapping "
+    "intervals and returns the merged list, with a short docstring."
+)
+
+
+@dataclass
+class ReqResult:
+    ok: bool
+    ttft: float | None = None  # seconds to first token
+    total: float | None = None  # total wall time
+    completion_tokens: int = 0
+    error: str | None = None
+
+    @property
+    def decode_tok_s(self) -> float | None:
+        if self.ok and self.ttft is not None and self.total and self.completion_tokens:
+            gen_time = max(self.total - self.ttft, 1e-6)
+            return self.completion_tokens / gen_time
+        return None
+
+
+@dataclass
+class BenchResult:
+    n: int
+    concurrency: int
+    wall_time: float
+    results: list[ReqResult] = field(default_factory=list)
+    price_per_hr: float | None = None
+
+    @property
+    def ok(self) -> list[ReqResult]:
+        return [r for r in self.results if r.ok]
+
+    @property
+    def ttft_p50(self) -> float | None:
+        return _pct([r.ttft for r in self.ok if r.ttft is not None], 50)
+
+    @property
+    def ttft_p95(self) -> float | None:
+        return _pct([r.ttft for r in self.ok if r.ttft is not None], 95)
+
+    @property
+    def avg_decode_tok_s(self) -> float | None:
+        vals = [r.decode_tok_s for r in self.ok if r.decode_tok_s]
+        return sum(vals) / len(vals) if vals else None
+
+    @property
+    def total_completion_tokens(self) -> int:
+        return sum(r.completion_tokens for r in self.ok)
+
+    @property
+    def throughput_tok_s(self) -> float | None:
+        if self.wall_time > 0 and self.total_completion_tokens:
+            return self.total_completion_tokens / self.wall_time
+        return None
+
+    @property
+    def cost_per_million(self) -> float | None:
+        """$/1M output tokens at this throughput."""
+        tps = self.throughput_tok_s
+        if tps and self.price_per_hr:
+            return (self.price_per_hr / 3600.0) / tps * 1_000_000
+        return None
+
+
+def _pct(values: list[float], p: float) -> float | None:
+    if not values:
+        return None
+    s = sorted(values)
+    k = max(0, min(len(s) - 1, round((p / 100.0) * (len(s) - 1))))
+    return s[k]
+
+
+def _one_request(
+    base_url: str, model: str, api_key: str | None, prompt: str, max_tokens: int, timeout: float
+) -> ReqResult:
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "temperature": 0.2,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+    }
+    url = f"{base_url.rstrip('/')}/chat/completions"
+    start = time.time()
+    ttft: float | None = None
+    chunk_tokens = 0
+    usage_tokens = 0
+    try:
+        with httpx.stream("POST", url, headers=headers, json=payload, timeout=timeout) as r:
+            if r.status_code != 200:
+                r.read()
+                return ReqResult(ok=False, error=f"HTTP {r.status_code}: {r.text[:160]}")
+            for line in r.iter_lines():
+                if not line.startswith("data: "):
+                    continue
+                data = line[len("data: "):].strip()
+                if data == "[DONE]":
+                    break
+                try:
+                    obj = json.loads(data)
+                except json.JSONDecodeError:
+                    continue
+                choices = obj.get("choices") or []
+                if choices:
+                    delta = choices[0].get("delta") or {}
+                    # reasoning models (GLM-5, DeepSeek-R1, ...) stream into
+                    # reasoning_content; count it as generated work too.
+                    if delta.get("content") or delta.get("reasoning_content"):
+                        if ttft is None:
+                            ttft = time.time() - start
+                        chunk_tokens += 1
+                if obj.get("usage"):
+                    usage_tokens = obj["usage"].get("completion_tokens", 0)
+    except httpx.HTTPError as e:
+        return ReqResult(ok=False, error=str(e))
+
+    total = time.time() - start
+    return ReqResult(
+        ok=True, ttft=ttft, total=total, completion_tokens=usage_tokens or chunk_tokens
+    )
+
+
+def run_benchmark(
+    base_url: str,
+    model: str,
+    api_key: str | None = None,
+    n: int = 5,
+    max_tokens: int = 256,
+    concurrency: int = 1,
+    prompt: str = DEFAULT_PROMPT,
+    price_per_hr: float | None = None,
+    timeout: float = 180.0,
+) -> BenchResult:
+    """Fire `n` requests `concurrency`-at-a-time and aggregate the stats."""
+    results: list[ReqResult] = []
+    wall_start = time.time()
+    with ThreadPoolExecutor(max_workers=concurrency) as pool:
+        futures = [
+            pool.submit(_one_request, base_url, model, api_key, prompt, max_tokens, timeout)
+            for _ in range(n)
+        ]
+        for f in futures:
+            results.append(f.result())
+    wall = time.time() - wall_start
+    return BenchResult(
+        n=n, concurrency=concurrency, wall_time=wall, results=results, price_per_hr=price_per_hr
+    )

aiod/bootstrap.py

@@ -0,0 +1,102 @@
+"""Builds the serving configuration that runs on the rented box: the Docker
+image, the OpenAI-API server command, environment, and the published port. The
+provider clients (vast/runpod) consume this when creating the instance.
+
+Two engines:
+  * vllm     — safetensors / AWQ / GPTQ / fp8 models (image vllm/vllm-openai)
+  * llamacpp — GGUF models (image ghcr.io/ggml-org/llama.cpp:server-cuda); loads
+               the GGUF straight from HF with -hf repo:QUANT (multi-part shards
+               auto-detected), spreads across all GPUs with -ngl 999, and enables
+               tool calling with --jinja. We pass --port 8000 so the published
+               port is identical to vLLM's.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+VLLM_IMAGE = "vllm/vllm-openai:latest"
+LLAMACPP_IMAGE = "ghcr.io/ggml-org/llama.cpp:server-cuda"
+
+# Port the server listens on inside the container (published to a public port).
+CONTAINER_PORT = 8000
+
+# Map our quant keys to vLLM's --quantization values. bf16 = no flag.
+VLLM_QUANT_FLAG = {
+    "bf16": None,
+    "fp16": None,
+    "fp8": "fp8",
+    "awq-int4": "awq",
+    "gptq-int4": "gptq",
+}
+
+
+@dataclass
+class ServerConfig:
+    repo_id: str
+    num_gpus: int
+    quant: str
+    api_key: str
+    engine: str = "vllm"  # "vllm" | "llamacpp"
+    image: str | None = None  # defaults per engine
+    port: int = CONTAINER_PORT
+    max_model_len: int | None = None
+    gpu_memory_utilization: float = 0.92
+    tool_call_parser: str = "hermes"
+    extra_args: list[str] = field(default_factory=list)
+    hf_token: str | None = None
+    gguf_quant: str | None = None  # llamacpp: the GGUF quant tag, e.g. "UD-IQ1_M"
+
+    def __post_init__(self) -> None:
+        if self.image is None:
+            self.image = LLAMACPP_IMAGE if self.engine == "llamacpp" else VLLM_IMAGE
+
+    def server_args(self) -> list[str]:
+        if self.engine == "llamacpp":
+            return self._llamacpp_args()
+        return self._vllm_args()
+
+    def _vllm_args(self) -> list[str]:
+        args = [
+            "--host", "0.0.0.0",
+            "--port", str(self.port),
+            "--model", self.repo_id,
+            "--served-model-name", self.repo_id,
+            "--api-key", self.api_key,
+            "--tensor-parallel-size", str(self.num_gpus),
+            "--gpu-memory-utilization", str(self.gpu_memory_utilization),
+            # Claude Code is extremely tool-call heavy; enable robust tool parsing.
+            "--enable-auto-tool-choice",
+            "--tool-call-parser", self.tool_call_parser,
+        ]
+        flag = VLLM_QUANT_FLAG.get(self.quant)
+        if flag:
+            args += ["--quantization", flag]
+        if self.max_model_len:
+            args += ["--max-model-len", str(self.max_model_len)]
+        args += self.extra_args
+        return args
+
+    def _llamacpp_args(self) -> list[str]:
+        # -hf repo[:quant]; the tag is a case-insensitive substring match against
+        # filenames (works across subfolders) and auto-downloads all shards.
+        ref = f"{self.repo_id}:{self.gguf_quant}" if self.gguf_quant else self.repo_id
+        args = [
+            "-hf", ref,
+            "--host", "0.0.0.0",
+            "--port", str(self.port),
+            "--api-key", self.api_key,
+            "-ngl", "999",          # all layers on GPU; default layer-split spreads them
+            "--jinja",              # enable OpenAI tool/function calling
+            "-c", str(self.max_model_len or 32768),
+        ]
+        args += self.extra_args
+        return args
+
+    def env(self) -> dict[str, str]:
+        e: dict[str, str] = {}
+        if self.hf_token:
+            # vLLM/huggingface_hub read both; llama.cpp reads HF_TOKEN.
+            e["HF_TOKEN"] = self.hf_token
+            e["HUGGING_FACE_HUB_TOKEN"] = self.hf_token
+        return e

aiod/branding.py

@@ -0,0 +1,37 @@
+"""Project links and the maintainer referral links.
+
+To earn referral credit, paste YOUR referral link below for each provider.
+  * vast.ai  -> Settings -> Referral Link   (3% of referred spend, for life)
+  * RunPod   -> Settings -> Referrals       (credits on referred spend)
+
+Leave an entry empty to fall back to the plain signup URL. These constants feed
+the `aiod init` wizard and the README signup links, so you only edit them here.
+"""
+
+from __future__ import annotations
+
+# >>> EDIT THESE: your referral links per provider <<<
+VAST_REFERRAL_URL = "https://cloud.vast.ai/?ref_id=25480"
+RUNPOD_REFERRAL_URL = "https://runpod.io?ref=p8hj7fq3"
+
+# Plain signup fallbacks (used if the matching referral link above is empty).
+SIGNUP_FALLBACK = {
+    "vast": "https://cloud.vast.ai/",
+    "runpod": "https://www.runpod.io/",
+}
+REFERRAL_URLS = {
+    "vast": VAST_REFERRAL_URL,
+    "runpod": RUNPOD_REFERRAL_URL,
+}
+
+# Other links.
+VAST_KEYS_URL = "https://cloud.vast.ai/manage-keys/"
+RUNPOD_KEYS_URL = "https://www.runpod.io/console/user/settings"
+HF_TOKENS_URL = "https://huggingface.co/settings/tokens"
+CCR_INSTALL_CMD = "npm install -g @musistudio/claude-code-router"
+
+
+def signup_url(provider: str = "vast") -> str:
+    """The link to send people who don't have an account on `provider` yet."""
+    ref = (REFERRAL_URLS.get(provider) or "").strip()
+    return ref or SIGNUP_FALLBACK.get(provider, SIGNUP_FALLBACK["vast"])

aiod/ccr.py

@@ -0,0 +1,90 @@
+"""Generate / merge the Claude Code Router config so `ccr code` routes Claude
+Code to our remote vLLM box.
+
+Verified schema notes (musistudio/claude-code-router, classic CLI):
+  * config path: ~/.claude-code-router/config.json
+  * Providers[].api_base_url must be the FULL URL ending in /chat/completions
+    (NOT just the base) for a custom OpenAI-compatible endpoint like vLLM.
+  * Providers[].api_key is sent upstream as `Authorization: Bearer`.
+  * Providers[].models lists names exactly as /v1/models reports them.
+  * transformer is omitted for a vanilla OpenAI-compatible endpoint.
+  * Router values are the string "providerName,modelName".
+
+We MERGE: existing providers and top-level settings are preserved; we only
+replace our own provider (matched by name) and repoint the Router.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+
+CONFIG_DIR = Path.home() / ".claude-code-router"
+CONFIG_FILE = CONFIG_DIR / "config.json"
+PROVIDER_NAME = "aiod-vllm"
+
+
+def build_provider(base_url: str, api_key: str, model: str) -> dict:
+    """base_url ends in /v1 (e.g. http://1.2.3.4:33526/v1)."""
+    return {
+        "name": PROVIDER_NAME,
+        "api_base_url": f"{base_url.rstrip('/')}/chat/completions",
+        "api_key": api_key or "dummy",
+        "models": [model],
+    }
+
+
+def build_router(model: str) -> dict:
+    ref = f"{PROVIDER_NAME},{model}"
+    return {
+        "default": ref,
+        "background": ref,
+        "think": ref,
+        "longContext": ref,
+        "longContextThreshold": 60000,
+        "webSearch": ref,
+    }
+
+
+def _load_existing() -> dict:
+    if CONFIG_FILE.exists():
+        try:
+            return json.loads(CONFIG_FILE.read_text())
+        except json.JSONDecodeError:
+            pass
+    return {}
+
+
+def write_config(base_url: str, api_key: str, model: str) -> Path:
+    """Merge our provider+router into the CCR config, preserving everything else.
+    Backs up an existing config to config.json.bak. Returns the config path."""
+    cfg = _load_existing()
+
+    if CONFIG_FILE.exists():
+        shutil.copy2(CONFIG_FILE, CONFIG_FILE.with_suffix(".json.bak"))
+
+    cfg.setdefault("LOG", False)
+    cfg.setdefault("HOST", "127.0.0.1")
+    cfg.setdefault("PORT", 3456)
+
+    providers = cfg.get("Providers")
+    if not isinstance(providers, list):
+        providers = []
+    providers = [p for p in providers if p.get("name") != PROVIDER_NAME]
+    providers.append(build_provider(base_url, api_key, model))
+    cfg["Providers"] = providers
+
+    cfg["Router"] = build_router(model)
+
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    CONFIG_FILE.write_text(json.dumps(cfg, indent=2))
+    return CONFIG_FILE
+
+
+def config_snippet(base_url: str, api_key: str, model: str) -> str:
+    """The exact provider+router block, for printing / dry-run."""
+    return json.dumps(
+        {"Providers": [build_provider(base_url, api_key, model)], "Router": build_router(model)},
+        indent=2,
+    )