129 lines
4.9 KiB
Python
129 lines
4.9 KiB
Python
"""Minimal RunPod API v2 client for the H3 single-node benchmark pod."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
API = "https://api.runpod.io/v2"
|
|
BLACKWELL_GPUS = (
|
|
"NVIDIA RTX PRO 6000 Blackwell Server Edition",
|
|
"NVIDIA RTX PRO 6000 Blackwell Workstation Edition",
|
|
"NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition",
|
|
)
|
|
DEFAULT_IMAGE = "runpod/pytorch:1.1.0-cu1300-torch291-ubuntu2404"
|
|
|
|
|
|
def request(method: str, path: str, body=None, query=None):
|
|
key = os.environ.get("RUNPOD_API_KEY")
|
|
if not key:
|
|
raise SystemExit("RUNPOD_API_KEY is required")
|
|
url = f"{API}{path}"
|
|
if query:
|
|
url += "?" + urllib.parse.urlencode(query)
|
|
data = None if body is None else json.dumps(body).encode("utf-8")
|
|
call = urllib.request.Request(url, data=data, method=method)
|
|
call.add_header("User-Agent", "h3-blackwell-runtime/0.1")
|
|
call.add_header("Authorization", f"Bearer {key}")
|
|
call.add_header("Accept", "application/json")
|
|
if data is not None:
|
|
call.add_header("Content-Type", "application/json")
|
|
try:
|
|
with urllib.request.urlopen(call, timeout=60) as response:
|
|
payload = response.read()
|
|
return json.loads(payload) if payload else {"status": response.status}
|
|
except urllib.error.HTTPError as error:
|
|
detail = error.read().decode("utf-8", errors="replace")
|
|
raise SystemExit(f"RunPod API returned HTTP {error.code}: {detail}") from error
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
commands = parser.add_subparsers(dest="command", required=True)
|
|
|
|
catalog_parser = commands.add_parser("catalog")
|
|
catalog_parser.add_argument("--count", type=int, default=8)
|
|
catalog_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")
|
|
|
|
create_parser = commands.add_parser("create")
|
|
create_parser.add_argument("--gpu", choices=BLACKWELL_GPUS, default=BLACKWELL_GPUS[0])
|
|
create_parser.add_argument("--count", type=int, default=8)
|
|
create_parser.add_argument("--cloud", choices=("SECURE", "COMMUNITY"), default="SECURE")
|
|
create_parser.add_argument("--datacenter")
|
|
create_parser.add_argument("--image", default=DEFAULT_IMAGE)
|
|
create_parser.add_argument("--disk", type=int, default=100)
|
|
create_parser.add_argument("--volume", type=int, default=100)
|
|
create_parser.add_argument("--network-volume")
|
|
create_parser.add_argument("--volume-mount-path", default="/runpod-volume")
|
|
create_parser.add_argument("--ssh-public-key", type=Path)
|
|
create_parser.add_argument("--yes", action="store_true")
|
|
|
|
get_parser = commands.add_parser("get")
|
|
get_parser.add_argument("pod_id")
|
|
|
|
terminate_parser = commands.add_parser("terminate")
|
|
terminate_parser.add_argument("pod_id")
|
|
terminate_parser.add_argument("--yes", action="store_true")
|
|
|
|
ssh_keys_parser = commands.add_parser("ssh-keys")
|
|
ssh_keys_parser.add_argument("--replace-with", type=Path)
|
|
ssh_keys_parser.add_argument("--yes", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
if args.command == "catalog":
|
|
response = request("GET", "/catalog/gpus", query={
|
|
"include": "AVAILABILITY",
|
|
"product": "POD",
|
|
"count": args.count,
|
|
"cloud": args.cloud,
|
|
"minCudaVersion": "12.8",
|
|
})
|
|
response["gpus"] = [gpu for gpu in response["gpus"] if gpu["id"] in BLACKWELL_GPUS]
|
|
elif args.command == "create":
|
|
if not args.yes:
|
|
raise SystemExit("create rents billable GPUs; repeat with --yes after checking catalog")
|
|
body = {
|
|
"name": "h3-blackwell-distributed",
|
|
"image": args.image,
|
|
"gpu": {"id": args.gpu, "count": args.count, "minCudaVersion": "12.8"},
|
|
"cloud": args.cloud,
|
|
"disk": args.disk,
|
|
"ports": ["22/tcp"],
|
|
"startSsh": True,
|
|
}
|
|
if args.network_volume:
|
|
body["mounts"] = {
|
|
"network": [{"volumeId": args.network_volume, "path": args.volume_mount_path}],
|
|
}
|
|
else:
|
|
body["mounts"] = {
|
|
"persistent": {"size": args.volume, "path": args.volume_mount_path},
|
|
}
|
|
if args.ssh_public_key:
|
|
body["env"] = {
|
|
"PUBLIC_KEY": args.ssh_public_key.read_text(encoding="utf-8").strip(),
|
|
}
|
|
if args.datacenter:
|
|
body["dataCenterIds"] = [args.datacenter]
|
|
response = request("POST", "/pods", body=body)
|
|
elif args.command == "get":
|
|
response = request("GET", f"/pods/{args.pod_id}")
|
|
elif args.command == "terminate":
|
|
if not args.yes:
|
|
raise SystemExit("termination is irreversible; repeat with --yes")
|
|
response = request("POST", f"/pods/{args.pod_id}/action", body={"action": "terminate"})
|
|
else:
|
|
if args.replace_with:
|
|
if not args.yes:
|
|
raise SystemExit("replacing account SSH keys requires --yes")
|
|
public_key = args.replace_with.read_text(encoding="utf-8").strip()
|
|
response = request("PUT", "/account/ssh-keys", body={"keys": [public_key]})
|
|
else:
|
|
response = request("GET", "/account/ssh-keys")
|
|
|
|
json.dump(response, sys.stdout, indent=2)
|
|
sys.stdout.write("\n")
|