| |
| |
| |
| |
| |
| |
| |
| """ |
| 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 |
|
|
| |
| 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}") |
|
|
| |
| 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}") |
|
|
| |
| 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() |
|
|
