File size: 2,591 Bytes
b85c683 |
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 |
# /// script
# requires-python = ">=3.10"
# dependencies = [
# "runpod>=1.6.0",
# "click>=8.0.0",
# ]
# ///
"""
Simple RunPod test script to verify API connection and GPU availability.
Usage:
export RUNPOD_API_KEY="your-api-key"
uv run scripts/runpod_simple_test.py
uv run scripts/runpod_simple_test.py --list-gpus
uv run scripts/runpod_simple_test.py --run-test
"""
import os
import click
import runpod
@click.command()
@click.option("--list-gpus", is_flag=True, help="List available GPU types")
@click.option("--run-test", is_flag=True, help="Run a quick test pod")
def main(list_gpus, run_test):
"""Test RunPod API connection and GPU availability."""
api_key = os.environ.get("RUNPOD_API_KEY")
if not api_key:
raise click.ClickException(
"Set RUNPOD_API_KEY environment variable.\n"
"Get your key at: https://runpod.io/console/user/settings"
)
runpod.api_key = api_key
# Test API connection
click.echo("Testing RunPod API connection...")
try:
pods = runpod.get_pods()
click.echo(f" Connected! Active pods: {len(pods)}")
except Exception as e:
raise click.ClickException(f"API connection failed: {e}")
# List GPUs
if list_gpus:
click.echo("\nAvailable GPU types:")
try:
gpus = runpod.get_gpus()
for gpu in gpus:
name = gpu.get("id", "Unknown")
mem = gpu.get("memoryInGb", "?")
click.echo(f" - {name} ({mem}GB)")
except Exception as e:
click.echo(f" Could not list GPUs: {e}")
# Run test pod
if run_test:
click.echo("\nLaunching test pod...")
test_script = "nvidia-smi && python3 -c 'import torch; print(f\"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}\")'"
pod = runpod.create_pod(
name="bamboo-1-test",
image_name="runpod/pytorch:2.1.0-py3.10-cuda11.8.0-devel-ubuntu22.04",
gpu_type_id="NVIDIA RTX A4000",
volume_in_gb=5,
docker_args=f"bash -c '{test_script}; sleep 60'",
)
click.echo(f" Pod ID: {pod['id']}")
click.echo(f" Monitor: https://runpod.io/console/pods")
click.echo(f"\n Terminate after checking:")
click.echo(f" uv run scripts/runpod_setup.py terminate {pod['id']}")
if not list_gpus and not run_test:
click.echo("\nUse --list-gpus to see available GPUs")
click.echo("Use --run-test to launch a quick test pod")
if __name__ == "__main__":
main()
|