| """ |
| Launch a PC-SHO-DLM training pod on RunPod using the current REST API. |
| |
| Usage: |
| export RUNPOD_API_KEY=... |
| python3 deploy_runpod.py --list |
| python3 deploy_runpod.py --create --gpu "<exact RunPod GPU type ID>" --steps 10 |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import urllib.error |
| import urllib.parse |
| import urllib.request |
|
|
|
|
| RUNPOD_API = "https://rest.runpod.io/v1" |
| DEFAULT_IMAGE = "runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04" |
| DEFAULT_VOLUME_MOUNT = "/workspace" |
| DEFAULT_HF_CODE_REPO = "https://huggingface.co/zotowata/pc-sho-dlm-code" |
|
|
|
|
| def get_runpod_key() -> str: |
| key = os.environ.get("RUNPOD_API_KEY") or os.environ.get("RUNPOD_API_TOKEN") |
| if not key: |
| raise SystemExit( |
| "Missing RunPod API key. Export RUNPOD_API_KEY before using this script." |
| ) |
| return key |
|
|
|
|
| def request(method: str, path: str, payload: dict | None = None) -> dict | list: |
| key = get_runpod_key() |
| url = f"{RUNPOD_API}{path}" |
| body = None if payload is None else json.dumps(payload).encode() |
| req = urllib.request.Request( |
| url, |
| data=body, |
| method=method, |
| headers={ |
| "Authorization": f"Bearer {key}", |
| "Content-Type": "application/json", |
| }, |
| ) |
| try: |
| with urllib.request.urlopen(req, timeout=60) as resp: |
| data = resp.read().decode() |
| return json.loads(data) if data else {} |
| except urllib.error.HTTPError as e: |
| detail = e.read().decode(errors="ignore") |
| raise RuntimeError(f"RunPod API {e.code}: {detail}") from e |
|
|
|
|
| def list_pods() -> list[dict]: |
| data = request("GET", "/pods") |
| if not isinstance(data, list): |
| raise RuntimeError(f"Unexpected pods response: {data}") |
| return data |
|
|
|
|
| def list_gpu_types() -> list[dict]: |
| |
| |
| |
| return [] |
|
|
|
|
| def make_remote_bootstrap(args: argparse.Namespace) -> str: |
| lines = [ |
| "set -euo pipefail", |
| "cd /workspace", |
| "if ! command -v git >/dev/null 2>&1 || ! command -v git-lfs >/dev/null 2>&1; then apt-get update && apt-get install -y --no-install-recommends git git-lfs; fi", |
| "export GIT_LFS_SKIP_SMUDGE=1", |
| "git lfs install", |
| ] |
| if os.environ.get("HF_TOKEN"): |
| lines.append('git config --global credential.helper store') |
| lines.append('printf "https://user:%s@huggingface.co\\n" "$HF_TOKEN" > ~/.git-credentials') |
| lines += [ |
| f"if [ ! -d pc-sho-dlm ]; then git clone {sh_quote(args.code_repo)} pc-sho-dlm; fi", |
| "cd pc-sho-dlm", |
| "git fetch --all --tags || true", |
| ] |
| if args.code_revision: |
| lines.append(f"git checkout {sh_quote(args.code_revision)}") |
| lines += [ |
| "chmod +x ./train_runpod.sh", |
| "bash ./train_runpod.sh", |
| ] |
| return " && ".join(lines) |
|
|
|
|
| def sh_quote(value: str) -> str: |
| return "'" + value.replace("'", "'\"'\"'") + "'" |
|
|
|
|
| def create_pod(args: argparse.Namespace) -> dict: |
| env = { |
| "HF_TOKEN": os.environ.get("HF_TOKEN", ""), |
| "MODEL_PRESET": args.model_preset, |
| "MODE": args.mode, |
| "TOKENIZER": args.tokenizer, |
| "DATA": args.data, |
| "STEPS": str(args.steps), |
| "BATCH_SIZE": str(args.batch_size), |
| "LR": str(args.lr), |
| "LOG_INTERVAL": str(args.log_interval), |
| "SAVE_INTERVAL": str(args.save_interval), |
| "TENSORBOARD_DIR": args.tensorboard_dir, |
| "TENSORBOARD_PORT": str(args.tensorboard_port), |
| "PRETRAIN_SHUFFLE_BUFFER": str(args.pretrain_shuffle_buffer), |
| "PRETRAIN_SHUFFLE_SEED": str(args.pretrain_shuffle_seed), |
| "SEQ_LEN_OVERRIDE": str(args.seq_len_override or ""), |
| "EVAL_INTERVAL": str(args.eval_interval), |
| "EVAL_BATCHES": str(args.eval_batches), |
| "EVAL_BATCH_SIZE": str(args.eval_batch_size), |
| "EVAL_RECORD_SKIP": str(args.eval_record_skip), |
| "HF_REPO_ID": args.hf_repo_id, |
| "CODE_REPO_URL": args.code_repo, |
| "CODE_REPO_REVISION": args.code_revision or "", |
| } |
| docker_start_cmd = ( |
| ["bash", "-lc", "sleep infinity"] |
| if args.idle_start |
| else ["bash", "-lc", make_remote_bootstrap(args)] |
| ) |
| payload = { |
| "name": args.name, |
| "cloudType": args.cloud_type, |
| "computeType": "GPU", |
| "gpuCount": args.gpu_count, |
| "gpuTypeIds": [args.gpu], |
| "gpuTypePriority": "availability", |
| "containerDiskInGb": args.container_disk_gb, |
| "volumeInGb": args.volume_gb, |
| "volumeMountPath": DEFAULT_VOLUME_MOUNT, |
| "imageName": args.image, |
| "ports": ["22/tcp", "8888/http"], |
| "supportPublicIp": True, |
| "globalNetworking": True, |
| "interruptible": args.interruptible, |
| "env": env, |
| "dockerStartCmd": docker_start_cmd, |
| } |
| return request("POST", "/pods", payload) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="RunPod launcher for PC-SHO-DLM 2B") |
| parser.add_argument("--list", action="store_true", help="List existing pods") |
| parser.add_argument("--create", action="store_true", help="Create a new pod") |
| parser.add_argument("--name", default="pc-sho-dlm-2b") |
| parser.add_argument( |
| "--gpu", |
| default=None, |
| help="Exact RunPod GPU type ID, e.g. one of the values shown in the RunPod console/API", |
| ) |
| parser.add_argument("--gpu-count", type=int, default=1) |
| parser.add_argument("--cloud-type", choices=["SECURE", "COMMUNITY"], default="COMMUNITY") |
| parser.add_argument("--interruptible", action="store_true") |
| parser.add_argument( |
| "--idle-start", |
| action="store_true", |
| help="Start container with 'sleep infinity' instead of auto-running bootstrap", |
| ) |
| parser.add_argument("--image", default=DEFAULT_IMAGE) |
| parser.add_argument("--volume-gb", type=int, default=150) |
| parser.add_argument("--container-disk-gb", type=int, default=100) |
| parser.add_argument("--code-repo", default=DEFAULT_HF_CODE_REPO) |
| parser.add_argument("--code-revision", default=None) |
| parser.add_argument("--model-preset", default="2b") |
| parser.add_argument("--mode", default="unified") |
| parser.add_argument("--tokenizer", default="HuggingFaceTB/SmolLM2-1.7B") |
| parser.add_argument("--data", default="HuggingFaceFW/fineweb-edu") |
| parser.add_argument("--steps", type=int, default=50000) |
| parser.add_argument("--batch-size", type=int, default=1) |
| parser.add_argument("--lr", type=float, default=1e-4) |
| parser.add_argument("--log-interval", type=int, default=10) |
| parser.add_argument("--save-interval", type=int, default=500) |
| parser.add_argument("--tensorboard-dir", default="checkpoints/tensorboard") |
| parser.add_argument("--tensorboard-port", type=int, default=8888) |
| parser.add_argument("--pretrain-shuffle-buffer", type=int, default=4096) |
| parser.add_argument("--pretrain-shuffle-seed", type=int, default=17) |
| parser.add_argument("--seq-len-override", type=int, default=None) |
| parser.add_argument("--eval-interval", type=int, default=500) |
| parser.add_argument("--eval-batches", type=int, default=8) |
| parser.add_argument("--eval-batch-size", type=int, default=1) |
| parser.add_argument("--eval-record-skip", type=int, default=200000) |
| parser.add_argument("--hf-repo-id", default="zotowata/pc-sho-dlm-2b-pretrain") |
| args = parser.parse_args() |
|
|
| if args.list: |
| pods = list_pods() |
| print(json.dumps(pods, indent=2)) |
| return |
|
|
| if args.create: |
| if not args.gpu: |
| raise SystemExit("--gpu is required with --create") |
| pod = create_pod(args) |
| print(json.dumps(pod, indent=2)) |
| return |
|
|
| parser.print_help() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
