255 lines
11 KiB
Python
255 lines
11 KiB
Python
|
"""Create and populate a regional RunPod network volume for H3 benchmarks."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import time
|
||
|
|
import urllib.error
|
||
|
|
import urllib.request
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
API_V1 = "https://rest.runpod.io/v1"
|
||
|
|
API_V2 = "https://api.runpod.io/v2"
|
||
|
|
CPU_IMAGE = "runpod/base:0.7.0-ubuntu2004"
|
||
|
|
MODEL_REPO = "coolthor/MiniMax-H3-pruned-NVFP4"
|
||
|
|
MODEL_REVISION = "fb06c2af47899f086a33bd599e91084b9f95bc54"
|
||
|
|
MODEL_REPO_PATH = "diffusion_models/minimax_h3_fl2va_pruned_nvfp4.safetensors"
|
||
|
|
MODEL_VOLUME_PATH = f"ComfyUI/models/{MODEL_REPO_PATH}"
|
||
|
|
MODEL_SHA256 = "9d49beb65ddc373a0df523b8ce715b61b73f88127bf3c8d3ffda76c5dda02bd4"
|
||
|
|
TEXT_REPO = "Comfy-Org/MiniMax-H3"
|
||
|
|
TEXT_REVISION = "0f7fb980293fcc4d55c1158cbda920806682ed5d"
|
||
|
|
TEXT_REPO_PATH = "text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors"
|
||
|
|
TEXT_VOLUME_PATH = f"ComfyUI/models/{TEXT_REPO_PATH}"
|
||
|
|
TEXT_SHA256 = "35a88d51044231fe332301d7a62aa81e3f2cba62febeb446e2c1e3e0ef76f2c6"
|
||
|
|
|
||
|
|
|
||
|
|
def request(api_key: str, method: str, url: str, body=None):
|
||
|
|
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 {api_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 RuntimeError(f"RunPod API returned HTTP {error.code}: {detail}") from error
|
||
|
|
|
||
|
|
|
||
|
|
def run(command: list[str], *, cwd: Path | None = None, input_text: str | None = None) -> None:
|
||
|
|
if input_text is None:
|
||
|
|
subprocess.run(command, cwd=cwd, text=True, check=True)
|
||
|
|
else:
|
||
|
|
# Preserve Unix LF endings when a Windows client streams a script to bash.
|
||
|
|
subprocess.run(command, cwd=cwd, input=input_text.encode("utf-8"), check=True)
|
||
|
|
|
||
|
|
|
||
|
|
def capture(command: list[str], *, cwd: Path) -> str:
|
||
|
|
return subprocess.check_output(command, cwd=cwd, text=True).strip()
|
||
|
|
|
||
|
|
|
||
|
|
def wait_for_ssh(api_key: str, pod_id: str, timeout_seconds: int) -> dict:
|
||
|
|
deadline = time.monotonic() + timeout_seconds
|
||
|
|
last_status = None
|
||
|
|
while time.monotonic() < deadline:
|
||
|
|
pod = request(api_key, "GET", f"{API_V2}/pods/{pod_id}")
|
||
|
|
status = pod.get("status")
|
||
|
|
if status != last_status:
|
||
|
|
print(f"pod {pod_id}: {status}", flush=True)
|
||
|
|
last_status = status
|
||
|
|
direct = (pod.get("ssh") or {}).get("direct")
|
||
|
|
if status == "RUNNING" and direct:
|
||
|
|
return direct
|
||
|
|
if status in {"ERROR", "TERMINATED"}:
|
||
|
|
raise RuntimeError(f"pod {pod_id} entered terminal status {status}")
|
||
|
|
time.sleep(5)
|
||
|
|
raise TimeoutError(f"pod {pod_id} did not expose direct SSH within {timeout_seconds}s")
|
||
|
|
|
||
|
|
|
||
|
|
def ssh_base(endpoint: dict, key: Path) -> list[str]:
|
||
|
|
return [
|
||
|
|
"ssh",
|
||
|
|
"-o", "BatchMode=yes",
|
||
|
|
"-o", "StrictHostKeyChecking=accept-new",
|
||
|
|
"-o", "ServerAliveInterval=15",
|
||
|
|
"-o", "ServerAliveCountMax=20",
|
||
|
|
"-p", str(endpoint["port"]),
|
||
|
|
"-i", str(key),
|
||
|
|
f'{endpoint["username"]}@{endpoint["host"]}',
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def remote_download_script(hf_token: str) -> str:
|
||
|
|
def download(repo: str, revision: str, repo_path: str, volume_path: str, checksum: str) -> str:
|
||
|
|
url = f"https://huggingface.co/{repo}/resolve/{revision}/{repo_path}?download=true"
|
||
|
|
target = f"/runpod-volume/{volume_path}"
|
||
|
|
directory, filename = target.rsplit("/", 1)
|
||
|
|
return f"""
|
||
|
|
mkdir -p '{directory}'
|
||
|
|
if echo '{checksum} {target}' | sha256sum --check --status; then
|
||
|
|
echo 'READY {volume_path}'
|
||
|
|
else
|
||
|
|
if [[ -f '{target}' && ! -f '{target}.aria2' ]]; then rm -f '{target}'; fi
|
||
|
|
echo 'DOWNLOAD {volume_path}'
|
||
|
|
aria2c --header="Authorization: Bearer $HF_TOKEN" \
|
||
|
|
--dir='{directory}' --out='{filename}' --continue=true \
|
||
|
|
--allow-overwrite=true --auto-file-renaming=false \
|
||
|
|
--max-connection-per-server=16 --split=16 --min-split-size=16M \
|
||
|
|
--file-allocation=none --max-tries=20 --retry-wait=5 --timeout=60 \
|
||
|
|
--console-log-level=warn --summary-interval=10 '{url}'
|
||
|
|
echo '{checksum} {target}' | sha256sum --check
|
||
|
|
fi
|
||
|
|
"""
|
||
|
|
|
||
|
|
return f"""#!/usr/bin/env bash
|
||
|
|
set -Eeuo pipefail
|
||
|
|
export HF_TOKEN={json.dumps(hf_token)}
|
||
|
|
echo "$(date -u +%FT%TZ) BOOTSTRAP_START"
|
||
|
|
apt-get update -qq
|
||
|
|
apt-get install -y -qq aria2 ca-certificates
|
||
|
|
{download(MODEL_REPO, MODEL_REVISION, MODEL_REPO_PATH, MODEL_VOLUME_PATH, MODEL_SHA256)}
|
||
|
|
{download(TEXT_REPO, TEXT_REVISION, TEXT_REPO_PATH, TEXT_VOLUME_PATH, TEXT_SHA256)}
|
||
|
|
mkdir -p /runpod-volume/h3-runtime
|
||
|
|
touch /runpod-volume/h3-runtime/artifacts-ready
|
||
|
|
echo "$(date -u +%FT%TZ) BOOTSTRAP_COMPLETE"
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--region", required=True, help="RunPod data center ID, for example US-MO-2")
|
||
|
|
parser.add_argument("--volume-id", help="Reuse and repair an existing volume instead of creating one")
|
||
|
|
parser.add_argument("--volume-name", default="h3-vortex-benchmark")
|
||
|
|
parser.add_argument("--volume-size", type=int, default=100)
|
||
|
|
parser.add_argument("--cpu-flavor", default="cpu3g")
|
||
|
|
parser.add_argument("--vcpu", type=int, default=2)
|
||
|
|
parser.add_argument("--ssh-key", type=Path, default=Path.home() / ".ssh" / "inceptal-dev-envs")
|
||
|
|
parser.add_argument("--ssh-timeout", type=int, default=300)
|
||
|
|
parser.add_argument("--keep-pod", action="store_true", help="Do not terminate the CPU pod after bootstrapping")
|
||
|
|
parser.add_argument("--dry-run", action="store_true", help="Print the plan without creating billable resources")
|
||
|
|
parser.add_argument("--yes", action="store_true", help="Confirm network-volume and CPU-pod charges")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
repo_root = Path(__file__).resolve().parents[1]
|
||
|
|
commit = capture(["git", "rev-parse", "HEAD"], cwd=repo_root)
|
||
|
|
dirty = bool(capture(["git", "status", "--porcelain"], cwd=repo_root))
|
||
|
|
archive_name = f"h3-blackwell-runtime-{commit[:12]}.tar.gz"
|
||
|
|
plan = {
|
||
|
|
"region": args.region,
|
||
|
|
"volume_id": args.volume_id,
|
||
|
|
"volume_name": args.volume_name,
|
||
|
|
"volume_size_gb": args.volume_size,
|
||
|
|
"cpu": {"id": args.cpu_flavor, "vcpu_count": args.vcpu},
|
||
|
|
"cpu_image": CPU_IMAGE,
|
||
|
|
"git_commit": commit,
|
||
|
|
"git_worktree_dirty": dirty,
|
||
|
|
"archive": f"h3-runtime/{archive_name}",
|
||
|
|
"artifacts": {
|
||
|
|
MODEL_VOLUME_PATH: MODEL_SHA256,
|
||
|
|
TEXT_VOLUME_PATH: TEXT_SHA256,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
if args.dry_run:
|
||
|
|
print(json.dumps(plan, indent=2))
|
||
|
|
return 0
|
||
|
|
if not args.yes:
|
||
|
|
parser.error("this creates billable storage and CPU compute; repeat with --yes after reviewing --dry-run")
|
||
|
|
|
||
|
|
api_key = os.environ.get("RUNPOD_API_KEY")
|
||
|
|
hf_token = os.environ.get("HF_TOKEN")
|
||
|
|
if not api_key:
|
||
|
|
parser.error("RUNPOD_API_KEY is required")
|
||
|
|
if not hf_token:
|
||
|
|
parser.error("HF_TOKEN is required for the gated model repository")
|
||
|
|
if not args.ssh_key.is_file():
|
||
|
|
parser.error(f"SSH private key not found: {args.ssh_key}")
|
||
|
|
|
||
|
|
volume_id = args.volume_id
|
||
|
|
if volume_id:
|
||
|
|
volume = request(api_key, "GET", f"{API_V1}/networkvolumes/{volume_id}")
|
||
|
|
if volume["dataCenterId"] != args.region:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"volume {volume_id} is in {volume['dataCenterId']}, not requested region {args.region}"
|
||
|
|
)
|
||
|
|
print(f"reusing volume {volume_id} in {args.region}", flush=True)
|
||
|
|
else:
|
||
|
|
volume = request(api_key, "POST", f"{API_V1}/networkvolumes", {
|
||
|
|
"name": args.volume_name,
|
||
|
|
"size": args.volume_size,
|
||
|
|
"dataCenterId": args.region,
|
||
|
|
})
|
||
|
|
volume_id = volume["id"]
|
||
|
|
print(f"created volume {volume_id} in {args.region}", flush=True)
|
||
|
|
|
||
|
|
pod_id = None
|
||
|
|
try:
|
||
|
|
pod = request(api_key, "POST", f"{API_V2}/pods", {
|
||
|
|
"name": f"h3-volume-bootstrap-{args.region.lower()}",
|
||
|
|
"image": CPU_IMAGE,
|
||
|
|
"cpu": {"id": args.cpu_flavor, "vcpuCount": args.vcpu},
|
||
|
|
"cloud": "SECURE",
|
||
|
|
"dataCenterIds": [args.region],
|
||
|
|
"disk": 5,
|
||
|
|
"ports": ["22/tcp"],
|
||
|
|
"mounts": {"network": [{"volumeId": volume_id, "path": "/runpod-volume"}]},
|
||
|
|
"startSsh": True,
|
||
|
|
})
|
||
|
|
pod_id = pod["id"]
|
||
|
|
print(f"created CPU pod {pod_id} at ${pod['cost']:.2f}/hour", flush=True)
|
||
|
|
endpoint = wait_for_ssh(api_key, pod_id, args.ssh_timeout)
|
||
|
|
ssh = ssh_base(endpoint, args.ssh_key)
|
||
|
|
|
||
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
||
|
|
temporary = Path(temporary_directory)
|
||
|
|
archive = temporary / archive_name
|
||
|
|
run(["git", "archive", "--format=tar.gz", f"--output={archive}", "HEAD"], cwd=repo_root)
|
||
|
|
run(ssh + ["mkdir", "-p", "/runpod-volume/h3-runtime"])
|
||
|
|
scp = [
|
||
|
|
"scp", "-q", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=accept-new",
|
||
|
|
"-P", str(endpoint["port"]), "-i", str(args.ssh_key), str(archive),
|
||
|
|
f'{endpoint["username"]}@{endpoint["host"]}:/runpod-volume/h3-runtime/{archive_name}',
|
||
|
|
]
|
||
|
|
run(scp)
|
||
|
|
print(f"uploaded application archive {archive_name}", flush=True)
|
||
|
|
|
||
|
|
print("starting pinned model downloads; progress follows", flush=True)
|
||
|
|
run(ssh + ["bash", "-s"], input_text=remote_download_script(hf_token))
|
||
|
|
|
||
|
|
manifest = dict(plan)
|
||
|
|
manifest.update({
|
||
|
|
"volume_id": volume_id,
|
||
|
|
"prepared_at": datetime.now(timezone.utc).isoformat(),
|
||
|
|
"cpu_pod_id": pod_id,
|
||
|
|
})
|
||
|
|
manifest_path = temporary / "regional-bootstrap-manifest.json"
|
||
|
|
manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
|
||
|
|
scp[-2] = str(manifest_path)
|
||
|
|
scp[-1] = (
|
||
|
|
f'{endpoint["username"]}@{endpoint["host"]}:'
|
||
|
|
"/runpod-volume/h3-runtime/regional-bootstrap-manifest.json"
|
||
|
|
)
|
||
|
|
run(scp)
|
||
|
|
|
||
|
|
print(json.dumps({"status": "ready", "region": args.region, "volume_id": volume_id}, indent=2))
|
||
|
|
return 0
|
||
|
|
finally:
|
||
|
|
if pod_id and not args.keep_pod:
|
||
|
|
try:
|
||
|
|
request(api_key, "POST", f"{API_V2}/pods/{pod_id}/action", {"action": "terminate"})
|
||
|
|
print(f"terminated CPU pod {pod_id}", flush=True)
|
||
|
|
except Exception as error:
|
||
|
|
print(f"WARNING: failed to terminate CPU pod {pod_id}: {error}", file=sys.stderr)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|