deploy/vast_node.py

"""Rent a dedicated GPU node on Vast.ai and run vLLM (OpenAI-compatible API).

  python deploy/vast_node.py offers                 # list the cheapest available machines
  python deploy/vast_node.py up [--offer ID] [--model lightning-nvfp4]
  python deploy/vast_node.py status                 # print base_url, ready to paste into config.json
  python deploy/vast_node.py down                   # destroy the instance and stop billing

Requires VAST_API_KEY. The node's vLLM API key is random and stored in deploy/.node.json (chmod 600, never committed).
"""
import argparse
import json
import os
import secrets
import sys
import time
import urllib.request
from pathlib import Path

API = "https://console.vast.ai/api/v0"
STATE = Path(__file__).parent / ".node.json"

MODELS = {
    # Blackwell (RTX PRO 6000) runs NVFP4 natively: ~18 GB of weights, leaving 70 GB+ for the KV cache
    "lightning-nvfp4": {"hf": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", "served": "kefuro-base",
                        "extra": ["--max-model-len", "32768", "--gpu-memory-utilization", "0.90"]},
    "lightning-bf16": {"hf": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", "served": "kefuro-base",
                       "extra": ["--max-model-len", "32768", "--gpu-memory-utilization", "0.92"]},
    # BF16 base + runtime LoRA loading (scp the adapter after boot, then POST /v1/load_lora_adapter)
    "lightning-bf16-lora": {"hf": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", "served": "kefuro-base",
                            "extra": ["--max-model-len", "16384", "--gpu-memory-utilization", "0.92", "--enable-lora", "--max-lora-rank", "16"]},
    # BF16 + tool calling + reasoning parser (per the NVIDIA model card), for tau-bench style agent evals
    "lightning-bf16-tools": {"hf": "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", "served": "kefuro-base",
                             "extra": ["--max-model-len", "65536", "--gpu-memory-utilization", "0.92", "--enable-prefix-caching",
                                       "--reasoning-parser", "nemotron_v3", "--enable-auto-tool-choice", "--tool-call-parser", "qwen3_coder",
                                       "--enable-lora", "--max-lora-rank", "16"]},
    "nano3-fp8": {"hf": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", "served": "kefuro-nano3",
                  "extra": ["--max-model-len", "32768", "--gpu-memory-utilization", "0.90"]},
}
GPU_PREFS = ["RTX PRO 6000 WS", "RTX PRO 6000 S", "H100 SXM", "H100 NVL"]


def _req(method, path, body=None):
    key = os.environ["VAST_API_KEY"]
    req = urllib.request.Request(API + path, method=method, data=json.dumps(body).encode() if body is not None else None,
                                 headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.load(r)


def offers(limit=8):
    out = []
    for g in GPU_PREFS:
        q = {"gpu_name": {"eq": g}, "num_gpus": {"eq": 1}, "rentable": {"eq": True}, "verified": {"eq": True},
             "reliability2": {"gte": 0.98}, "inet_down": {"gte": 500}, "disk_space": {"gte": 150},
             "order": [["dph_total", "asc"]], "type": "on-demand", "limit": 20}
        out += _req("POST", "/bundles/", q).get("offers", [])
    out.sort(key=lambda o: o["dph_total"])
    return out[:limit]


def cmd_offers(_):
    for o in offers():
        print(f"id={o['id']:<10} {o['gpu_name']:16s} {o['gpu_ram']/1024:.0f}GB ${o['dph_total']:.3f}/h "
              f"rel={o.get('reliability2', 0):.3f} down={o.get('inet_down', 0):.0f}Mbps {o.get('geolocation', '')}")


def cmd_up(a):
    if STATE.exists():
        sys.exit(f"A node record already exists at {STATE}; run down before up")
    m = MODELS[a.model]
    offer = a.offer or offers(1)[0]["id"]
    api_key = "sk-cs-" + secrets.token_hex(24)
    args = ["--model", m["hf"], "--served-model-name", m["served"], "--host", "0.0.0.0", "--port", "8000",
            "--api-key", api_key, "--trust-remote-code", *m["extra"]]
    body = {"client_id": "me", "image": "vllm/vllm-openai:latest", "disk": 150, "runtype": "args",
            "args": args,
            # without an HF token the Xet download path hangs on some hosts; force plain HTTP downloads
            "env": {"-p 8000:8000": "1", "HF_HUB_DISABLE_XET": "1", "HF_HUB_ENABLE_HF_TRANSFER": "0", "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True"}, "label": "kefuro"}
    if a.ssh:
        # SSH mode: Vast's init owns the entrypoint, vLLM starts in the background from onstart; SSH in to train
        body.update({"runtype": "ssh_direc ssh_proxy", "args": None,
                     "onstart": "nohup vllm serve " + " ".join(args[1:2]) + " " + " ".join(args[2:]) + " > /root/vllm.log 2>&1 &"})
        body.pop("args")
    r = _req("PUT", f"/asks/{offer}/", body)
    if not r.get("success"):
        sys.exit(f"create failed: {r}")
    if a.ssh:
        pub = Path(os.path.expanduser("~/.ssh/id_ed25519.pub")).read_text().strip()
        time.sleep(5)
        print("attach ssh key:", _req("POST", f"/instances/{r['new_contract']}/ssh/", {"ssh_key": pub}).get("success"))
    STATE.write_text(json.dumps({"instance_id": r["new_contract"], "offer": offer, "model": a.model, "ssh": bool(a.ssh),
                                 "served": m["served"], "api_key": api_key, "created": time.time()}))
    STATE.chmod(0o600)
    print(f"instance {r['new_contract']} created on offer {offer}; use status to see ports and readiness")


def _instance(st):
    return _req("GET", f"/instances/{st['instance_id']}/").get("instances") or {}


def cmd_status(_):
    st = json.loads(STATE.read_text())
    i = _instance(st)
    ports = (i.get("ports") or {}).get("8000/tcp") or []
    url = f"http://{i.get('public_ipaddr', '').strip()}:{ports[0]['HostPort']}/v1" if ports else None
    hours = (time.time() - st["created"]) / 3600
    ssh = None
    if st.get("ssh") and i.get("public_ipaddr"):
        p22 = ((i.get("ports") or {}).get("22/tcp") or [{}])[0].get("HostPort")
        ssh = f"ssh -p {p22} root@{i['public_ipaddr'].strip()}" if p22 else f"ssh -p {i.get('ssh_port')} root@{i.get('ssh_host')}"
    print(json.dumps({"instance": st["instance_id"], "state": i.get("actual_status"), "gpu": i.get("gpu_name"), "ssh": ssh,
                      "dph": i.get("dph_total"), "hours": round(hours, 2), "spent_est": round(hours * (i.get("dph_total") or 0), 2),
                      "base_url": url, "model": st["served"], "api_key_file": str(STATE)}, indent=2))


def cmd_down(_):
    st = json.loads(STATE.read_text())
    r = _req("DELETE", f"/instances/{st['instance_id']}/")
    print(r)
    if r.get("success"):
        STATE.unlink()


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("offers")
    u = sub.add_parser("up")
    u.add_argument("--offer", type=int)
    u.add_argument("--model", default="lightning-nvfp4", choices=list(MODELS))
    u.add_argument("--ssh", action="store_true", help="enable SSH (vLLM starts in the background from onstart), for training on the same machine")
    sub.add_parser("status")
    sub.add_parser("down")
    a = p.parse_args()
    {"offers": cmd_offers, "up": cmd_up, "status": cmd_status, "down": cmd_down}[a.cmd](a)