#!/usr/bin/env python3 """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())