| |
| """RunPod pod lifecycle over the REST API (classifier-plan §3.1, §3.4). Run on this host, never on the pod. |
| |
| pod.py create --image TAG [--volume ID] [--cloud COMMUNITY|SECURE] [--public-key FILE] [--yes] |
| pod.py status POD_ID read-only (no spend), prints state, cost per hour and the SSH host/port |
| pod.py ssh POD_ID prints "HOST PORT" of the exposed 22/tcp once the pod is up (exit 1 if not yet) |
| pod.py stop POD_ID [--yes] a stopped pod still bills its disk |
| pod.py terminate POD_ID [--yes] |
| |
| ``create``/``stop``/``terminate`` only print the request (method, URL, JSON body; never the key) unless ``--yes`` |
| is given. The key comes from ``RUNPOD_API_KEY`` in the environment and is never printed or logged; error messages |
| carry the HTTP status only. Pod spec: 1x H100 80 GB, 50 GB container disk, the network volume at /workspace, |
| SSH over exposed TCP 22 with a dedicated public key. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import httpx |
|
|
| API = "https://rest.runpod.io/v1" |
| KEY_ENV = "RUNPOD_API_KEY" |
| GPU_TYPES = ["NVIDIA H100 80GB HBM3", "NVIDIA H100 PCIe"] |
| DEFAULT_PUBKEY = Path.home() / ".ssh" / "naver_psy_runpod.pub" |
|
|
|
|
| class PodError(RuntimeError): |
| pass |
|
|
|
|
| def create_body(image: str, public_key: str, volume: str | None = None, cloud: str = "COMMUNITY", |
| name: str = "naver-psy-train", gpu_types: list[str] | None = None, container_gb: int = 50, |
| volume_gb: int = 100) -> dict[str, Any]: |
| if cloud not in ("COMMUNITY", "SECURE"): |
| raise PodError("cloud must be COMMUNITY or SECURE") |
| if not image or ":" not in image or image.endswith(":latest"): |
| raise PodError("image must be an exact tag (recorded in config.json), not latest") |
| if not public_key.startswith(("ssh-ed25519 ", "ssh-rsa ", "ecdsa-")): |
| raise PodError("public key file does not look like an SSH public key") |
| body: dict[str, Any] = { |
| "name": name, "imageName": image, "gpuTypeIds": gpu_types or GPU_TYPES, "gpuCount": 1, "cloudType": cloud, |
| "containerDiskInGb": container_gb, "ports": ["22/tcp"], "supportPublicIp": True, |
| "env": {"PUBLIC_KEY": public_key.strip()}, |
| } |
| if volume: |
| body["networkVolumeId"] = volume |
| else: |
| body.update(volumeInGb=volume_gb, volumeMountPath="/workspace") |
| return body |
|
|
|
|
| class Client: |
| def __init__(self, key: str | None = None, transport: httpx.BaseTransport | None = None, base: str = API) -> None: |
| key = key if key is not None else os.environ.get(KEY_ENV) |
| if not key: |
| raise PodError(f"{KEY_ENV} is not set") |
| self._headers = {"Authorization": f"Bearer {key}"} |
| self._http = httpx.Client(timeout=30, transport=transport) |
| self.base = base |
|
|
| def __repr__(self) -> str: |
| return "pod.Client()" |
|
|
| def call(self, method: str, path: str, body: Any = None) -> Any: |
| try: |
| r = self._http.request(method, self.base + path, headers=self._headers, json=body) |
| except httpx.HTTPError as e: |
| raise PodError(f"{method} {path}: transport error {type(e).__name__}") from None |
| if r.status_code >= 300: |
| raise PodError(f"{method} {path}: HTTP {r.status_code}") |
| return r.json() if r.content else {} |
|
|
|
|
| def ssh_address(pod: dict) -> tuple[str, int] | None: |
| """Public IP and port of the pod's exposed 22/tcp, if it is up.""" |
| ip = pod.get("publicIp") |
| ports = pod.get("portMappings") or {} |
| port = ports.get("22") if isinstance(ports, dict) else None |
| if ip and port: |
| return str(ip), int(port) |
| return None |
|
|
|
|
| def summary(pod: dict) -> dict: |
| addr = ssh_address(pod) |
| return {"id": pod.get("id"), "name": pod.get("name"), "desiredStatus": pod.get("desiredStatus"), |
| "costPerHr": pod.get("costPerHr"), "gpu": (pod.get("machine") or {}).get("gpuTypeId"), |
| "ssh": None if addr is None else {"host": addr[0], "port": addr[1]}} |
|
|
|
|
| def plan(method: str, path: str, body: Any = None) -> dict: |
| return {"dry_run": True, "method": method, "url": API + path, "body": body, |
| "note": "add --yes to send this request (the API key is read from RUNPOD_API_KEY and never shown)"} |
|
|
|
|
| def main(argv=None, transport: httpx.BaseTransport | None = None, out=sys.stdout) -> int: |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| sub = p.add_subparsers(dest="command", required=True) |
| c = sub.add_parser("create") |
| c.add_argument("--image", required=True, help="Exact RunPod PyTorch devel image tag") |
| c.add_argument("--volume", help="Network volume id (mounted at /workspace)") |
| c.add_argument("--cloud", default="COMMUNITY", choices=["COMMUNITY", "SECURE"]) |
| c.add_argument("--public-key", default=str(DEFAULT_PUBKEY)) |
| c.add_argument("--name", default="naver-psy-train") |
| c.add_argument("--gpu-type", action="append", help="Override the GPU type ids (repeatable)") |
| c.add_argument("--yes", action="store_true") |
| for name in ("status", "ssh", "stop", "terminate"): |
| s = sub.add_parser(name) |
| s.add_argument("pod_id") |
| if name in ("stop", "terminate"): |
| s.add_argument("--yes", action="store_true") |
| args = p.parse_args(argv) |
|
|
| def emit(value: Any) -> None: |
| out.write(json.dumps(value, ensure_ascii=False) + "\n") |
|
|
| try: |
| if args.command == "create": |
| body = create_body(args.image, Path(args.public_key).expanduser().read_text(), args.volume, args.cloud, |
| args.name, args.gpu_type) |
| if not args.yes: |
| emit(plan("POST", "/pods", body)) |
| return 0 |
| emit(summary(Client(transport=transport).call("POST", "/pods", body))) |
| return 0 |
| path = f"/pods/{args.pod_id}" |
| if args.command in ("stop", "terminate"): |
| method, target = ("POST", path + "/stop") if args.command == "stop" else ("DELETE", path) |
| if not args.yes: |
| emit(plan(method, target)) |
| return 0 |
| Client(transport=transport).call(method, target) |
| emit({"pod": args.pod_id, args.command: "requested"}) |
| return 0 |
| pod = Client(transport=transport).call("GET", path) |
| if args.command == "ssh": |
| addr = ssh_address(pod) |
| if addr is None: |
| return 1 |
| out.write(f"{addr[0]} {addr[1]}\n") |
| return 0 |
| emit(summary(pod)) |
| return 0 |
| except (PodError, OSError) as e: |
| print(f"pod.py: {e}", file=sys.stderr) |
| return 2 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
