File size: 4,645 Bytes
47ea602 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """Minimal RunPod GraphQL client (API key from ~/.runpod/config.toml).
python scripts/runpod.py balance
python scripts/runpod.py gpus
python scripts/runpod.py create NAME --gpu "NVIDIA H100 80GB HBM3" --count 1 [--cloud COMMUNITY] [--volume 300]
python scripts/runpod.py status POD_ID # prints ssh target when running
python scripts/runpod.py stop POD_ID | terminate POD_ID
"""
import argparse
import json
import re
import sys
import time
import urllib.request
from pathlib import Path
IMAGE = 'runpod/pytorch:1.4.1-cu1281-torch2130-ubuntu2404'
KEY_FILE = Path.home() / '.runpod' / 'ssh' / 'runpodctl-ssh-key'
def api_key():
text = (Path.home() / '.runpod' / 'config.toml').read_text()
return re.search(r"apikey\s*=\s*['\"]([^'\"]+)", text).group(1)
def gql(query, variables=None):
req = urllib.request.Request('https://api.runpod.io/graphql',
data=json.dumps({'query': query, 'variables': variables or {}}).encode(),
headers={'Content-Type': 'application/json', 'Authorization': f'Bearer {api_key()}',
'User-Agent': 'rapid-moss'})
with urllib.request.urlopen(req, timeout=60) as r:
out = json.load(r)
if out.get('errors'):
raise SystemExit(json.dumps(out['errors'], indent=1))
return out['data']
def pod(pod_id):
return gql('query($id: String!) { pod(input: {podId: $id}) { id name desiredStatus costPerHr gpuCount '
'machine { gpuDisplayName location } runtime { uptimeInSeconds ports { ip isIpPublic privatePort '
'publicPort type } } } }', {'id': pod_id})['pod']
def ssh_target(p):
for port in (p.get('runtime') or {}).get('ports') or []:
if port['privatePort'] == 22 and port['isIpPublic']:
return port['ip'], port['publicPort']
return None
def main():
ap = argparse.ArgumentParser()
sub = ap.add_subparsers(dest='cmd', required=True)
sub.add_parser('balance')
sub.add_parser('gpus')
c = sub.add_parser('create')
c.add_argument('name')
c.add_argument('--gpu', default='NVIDIA H100 80GB HBM3')
c.add_argument('--count', type=int, default=1)
c.add_argument('--cloud', default='COMMUNITY')
c.add_argument('--volume', type=int, default=300)
c.add_argument('--disk', type=int, default=60)
for name in ('status', 'stop', 'terminate', 'wait'):
sub.add_parser(name).add_argument('pod_id')
a = ap.parse_args()
if a.cmd == 'balance':
print(json.dumps(gql('query { myself { clientBalance currentSpendPerHr pods { id name desiredStatus costPerHr } } }')['myself'], indent=1))
elif a.cmd == 'gpus':
data = gql('query { gpuTypes { id memoryInGb securePrice communityPrice lowestPrice(input:{gpuCount:1}) { stockStatus } } }')
for g in sorted(data['gpuTypes'], key=lambda g: -(g['memoryInGb'] or 0)):
if (g['memoryInGb'] or 0) >= 80:
print(f"{g['id']:<60} {g['memoryInGb']:>4} GB secure {g['securePrice']} community {g['communityPrice']} stock {(g['lowestPrice'] or {}).get('stockStatus')}")
elif a.cmd == 'create':
pub = KEY_FILE.with_suffix('.pub').read_text().strip()
q = ('mutation($in: PodFindAndDeployOnDemandInput) { podFindAndDeployOnDemand(input: $in) '
'{ id costPerHr machine { gpuDisplayName location } } }')
inp = {'cloudType': a.cloud, 'gpuCount': a.count, 'volumeInGb': a.volume, 'containerDiskInGb': a.disk,
'gpuTypeId': a.gpu, 'name': a.name, 'imageName': IMAGE, 'ports': '22/tcp',
'volumeMountPath': '/workspace', 'startSsh': True, 'supportPublicIp': True,
'minCudaVersion': '12.8', 'env': [{'key': 'PUBLIC_KEY', 'value': pub}]}
print(json.dumps(gql(q, {'in': inp})['podFindAndDeployOnDemand'], indent=1))
elif a.cmd == 'status':
p = pod(a.pod_id)
print(json.dumps(p, indent=1))
t = ssh_target(p)
if t:
print(f'ssh -i {KEY_FILE} -p {t[1]} root@{t[0]}')
elif a.cmd == 'wait':
for _ in range(120):
t = ssh_target(pod(a.pod_id))
if t:
print(f'{t[0]} {t[1]}')
return
time.sleep(10)
sys.exit('pod did not expose ssh in 20 min')
elif a.cmd == 'stop':
print(gql('mutation($id: String!) { podStop(input: {podId: $id}) { id desiredStatus } }', {'id': a.pod_id}))
elif a.cmd == 'terminate':
print(gql('mutation($id: String!) { podTerminate(input: {podId: $id}) }', {'id': a.pod_id}))
if __name__ == '__main__':
main()
|