287 lines
8.3 KiB
Python
287 lines
8.3 KiB
Python
"""
|
|
RunPod Client - High-level API for pod management
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import requests
|
|
import shlex
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
|
|
class RunPodClient:
|
|
"""High-level client for RunPod pod management."""
|
|
|
|
def __init__(self, api_key: str, volume_id: str, registry_auth_id: Optional[str] = None):
|
|
"""
|
|
Initialize RunPod client.
|
|
|
|
Args:
|
|
api_key: RunPod API key
|
|
volume_id: RunPod network volume ID
|
|
registry_auth_id: Docker registry auth ID (optional)
|
|
"""
|
|
self.api_key = api_key
|
|
self.volume_id = volume_id
|
|
self.registry_auth_id = registry_auth_id
|
|
|
|
self.graphql_endpoint = "https://api.runpod.io/graphql"
|
|
self.rest_api_url = "https://rest.runpod.io/v1/pods"
|
|
|
|
# EUR-IS-1 datacenter (volume location)
|
|
self.datacenters = ['EUR-IS-1']
|
|
|
|
def get_available_gpus(self, min_vram: int = 16) -> List[Dict[str, Any]]:
|
|
"""
|
|
Query available GPU types with minimum VRAM.
|
|
|
|
Args:
|
|
min_vram: Minimum VRAM in GB (default: 16)
|
|
|
|
Returns:
|
|
List of GPU dictionaries sorted by price (cheapest first)
|
|
"""
|
|
query = """
|
|
{
|
|
gpuTypes {
|
|
id
|
|
displayName
|
|
memoryInGb
|
|
secureCloud
|
|
communityCloud
|
|
lowestPrice(input: {gpuCount: 1}) {
|
|
uninterruptablePrice
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
|
|
data = self._query_graphql(query)
|
|
if not data:
|
|
return []
|
|
|
|
gpu_types = data.get('data', {}).get('gpuTypes', [])
|
|
|
|
available_gpus = []
|
|
for gpu in gpu_types:
|
|
memory = gpu.get('memoryInGb', 0)
|
|
secure_count = gpu.get('secureCloud', 0)
|
|
lowest_price = gpu.get('lowestPrice', {})
|
|
price = lowest_price.get('uninterruptablePrice') if lowest_price else None
|
|
|
|
if memory >= min_vram and secure_count > 0 and price is not None:
|
|
available_gpus.append({
|
|
'id': gpu.get('id', ''),
|
|
'name': gpu.get('displayName', 'Unknown'),
|
|
'vram': memory,
|
|
'price': float(price),
|
|
'global_available': secure_count
|
|
})
|
|
|
|
# Sort by price (cheapest first)
|
|
available_gpus.sort(key=lambda x: x['price'])
|
|
|
|
return available_gpus
|
|
|
|
def deploy_pod(
|
|
self,
|
|
gpu_id: str,
|
|
image: str,
|
|
command: Optional[str] = None,
|
|
container_disk: int = 50,
|
|
ports: Optional[List[str]] = None,
|
|
env: Optional[Dict[str, str]] = None,
|
|
dry_run: bool = False
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Deploy a pod using REST API.
|
|
|
|
Args:
|
|
gpu_id: GPU type ID
|
|
image: Docker image name
|
|
command: Docker command (optional)
|
|
container_disk: Container disk size in GB
|
|
ports: Port mappings (e.g., ["8888/http", "22/tcp"])
|
|
env: Environment variables
|
|
dry_run: Show plan without deploying
|
|
|
|
Returns:
|
|
Pod data dictionary if successful, None otherwise
|
|
"""
|
|
pod_name = "foxhunt-training"
|
|
|
|
if ports is None:
|
|
ports = ["8888/http", "22/tcp"]
|
|
|
|
if env is None:
|
|
env = {}
|
|
|
|
deployment_payload = {
|
|
"cloudType": "SECURE",
|
|
"computeType": "GPU",
|
|
"dataCenterIds": self.datacenters,
|
|
"dataCenterPriority": "availability",
|
|
"gpuTypeIds": [gpu_id],
|
|
"gpuTypePriority": "availability",
|
|
"gpuCount": 1,
|
|
"name": pod_name,
|
|
"imageName": image,
|
|
"containerDiskInGb": container_disk,
|
|
"volumeInGb": 0,
|
|
"networkVolumeId": self.volume_id,
|
|
"volumeMountPath": "/runpod-volume",
|
|
"ports": ports,
|
|
"env": env,
|
|
"interruptible": False,
|
|
"minRAMPerGPU": 8,
|
|
"minVCPUPerGPU": 2
|
|
}
|
|
|
|
if command:
|
|
deployment_payload["dockerStartCmd"] = shlex.split(command)
|
|
|
|
if self.registry_auth_id:
|
|
deployment_payload["containerRegistryAuthId"] = self.registry_auth_id
|
|
|
|
if dry_run:
|
|
import json
|
|
print("\n" + "="*70)
|
|
print("DRY RUN: Deployment payload")
|
|
print("="*70)
|
|
print(json.dumps(deployment_payload, indent=2))
|
|
return None
|
|
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
|
|
try:
|
|
response = requests.post(
|
|
self.rest_api_url,
|
|
json=deployment_payload,
|
|
headers=headers,
|
|
timeout=60
|
|
)
|
|
|
|
if response.status_code in [200, 201]:
|
|
pod_data = response.json()
|
|
|
|
if not isinstance(pod_data, dict) or 'id' not in pod_data:
|
|
return None
|
|
|
|
return pod_data
|
|
|
|
elif response.status_code == 400:
|
|
try:
|
|
error_data = response.json()
|
|
error_msg = error_data.get('error', 'Unknown error')
|
|
print(f" ⚠️ Deployment failed: {error_msg}")
|
|
except ValueError:
|
|
print(f" ⚠️ Deployment failed: {response.text[:200]}")
|
|
|
|
return None
|
|
|
|
else:
|
|
print(f" ❌ Unexpected response (status {response.status_code})")
|
|
return None
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f" ⚠️ Deployment failed: {str(e)[:100]}")
|
|
return None
|
|
|
|
def stop_pod(self, pod_id: str) -> bool:
|
|
"""
|
|
Stop a running pod.
|
|
|
|
Args:
|
|
pod_id: Pod ID to stop
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
url = f"{self.rest_api_url}/{pod_id}/stop"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, timeout=30)
|
|
return response.status_code in [200, 204]
|
|
except requests.exceptions.RequestException:
|
|
return False
|
|
|
|
def terminate_pod(self, pod_id: str) -> bool:
|
|
"""
|
|
Terminate (delete) a pod.
|
|
|
|
Args:
|
|
pod_id: Pod ID to terminate
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
url = f"{self.rest_api_url}/{pod_id}/terminate"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
|
|
try:
|
|
response = requests.post(url, headers=headers, timeout=30)
|
|
return response.status_code in [200, 204]
|
|
except requests.exceptions.RequestException:
|
|
return False
|
|
|
|
def get_pod_status(self, pod_id: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get pod status.
|
|
|
|
Args:
|
|
pod_id: Pod ID
|
|
|
|
Returns:
|
|
Pod status dictionary if successful, None otherwise
|
|
"""
|
|
url = f"{self.rest_api_url}/{pod_id}"
|
|
headers = {
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
|
|
try:
|
|
response = requests.get(url, headers=headers, timeout=30)
|
|
if response.status_code == 200:
|
|
return response.json()
|
|
return None
|
|
except requests.exceptions.RequestException:
|
|
return None
|
|
|
|
def _query_graphql(self, query: str, variables: Optional[Dict] = None) -> Optional[Dict]:
|
|
"""Execute GraphQL query against RunPod API."""
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.api_key}"
|
|
}
|
|
|
|
payload = {"query": query}
|
|
if variables:
|
|
payload["variables"] = variables
|
|
|
|
try:
|
|
response = requests.post(
|
|
self.graphql_endpoint,
|
|
json=payload,
|
|
headers=headers,
|
|
timeout=30
|
|
)
|
|
response.raise_for_status()
|
|
result = response.json()
|
|
|
|
if 'errors' in result:
|
|
print(f"ERROR: GraphQL errors: {result['errors']}")
|
|
return None
|
|
|
|
return result
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"ERROR: Failed to query RunPod GraphQL API: {e}")
|
|
return None
|