251 lines
7.6 KiB
Python
251 lines
7.6 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
import requests
|
|
|
|
|
|
def parse_arguments():
|
|
"""Parse command line arguments"""
|
|
parser = argparse.ArgumentParser(description='Run tests on RunPod GPU')
|
|
parser.add_argument('--gpu-type', type=str, help='GPU type to use')
|
|
parser.add_argument('--gpu-count',
|
|
type=int,
|
|
help='Number of GPUs to use',
|
|
default=1)
|
|
parser.add_argument('--test-command', type=str, help='Test command to run')
|
|
parser.add_argument('--disk-size',
|
|
type=int,
|
|
default=20,
|
|
help='Container disk size in GB (default: 20)')
|
|
parser.add_argument('--volume-size',
|
|
type=int,
|
|
default=20,
|
|
help='Persistent volume size in GB (default: 20)')
|
|
parser.add_argument(
|
|
'--image',
|
|
type=str,
|
|
required=True,
|
|
help='Docker image to use')
|
|
return parser.parse_args()
|
|
|
|
|
|
args = parse_arguments()
|
|
API_KEY = os.environ['RUNPOD_API_KEY']
|
|
RUN_ID = os.environ['GITHUB_RUN_ID']
|
|
JOB_ID = os.environ['JOB_ID']
|
|
PODS_API = "https://rest.runpod.io/v1/pods"
|
|
HEADERS = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {API_KEY}"
|
|
}
|
|
|
|
|
|
def create_pod():
|
|
"""Create a RunPod instance"""
|
|
# Ensure image name is lowercase (Docker requirement)
|
|
image_name = args.image.lower()
|
|
print(f"Using specified image: {image_name}")
|
|
|
|
docker_start_cmd = [
|
|
"bash",
|
|
"-c",
|
|
"apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
|
|
]
|
|
|
|
print(f"Creating RunPod instance with GPU: {args.gpu_type}...")
|
|
payload = {
|
|
"name": f"fastvideo-{JOB_ID}-{RUN_ID}",
|
|
"containerDiskInGb": args.disk_size,
|
|
"volumeInGb": args.volume_size,
|
|
"gpuTypeIds": [args.gpu_type],
|
|
"gpuCount": args.gpu_count,
|
|
"imageName": image_name,
|
|
"allowedCudaVersions": ["12.4"],
|
|
"dockerStartCmd": docker_start_cmd
|
|
}
|
|
|
|
response = requests.post(PODS_API, headers=HEADERS, json=payload)
|
|
response_data = response.json()
|
|
print(f"Response: {json.dumps(response_data, indent=2)}")
|
|
|
|
return response_data["id"]
|
|
|
|
|
|
def wait_for_pod(pod_id):
|
|
"""Wait for pod to be in RUNNING state and fully ready with SSH access"""
|
|
print("Waiting for RunPod to be ready...")
|
|
|
|
# First wait for RUNNING status
|
|
max_attempts = 10
|
|
attempts = 0
|
|
while attempts < max_attempts:
|
|
response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
|
|
pod_data = response.json()
|
|
status = pod_data["desiredStatus"]
|
|
|
|
if status == "RUNNING":
|
|
print("RunPod is running! Now waiting for ports to be assigned...")
|
|
break
|
|
|
|
print(
|
|
f"Current status: {status}, waiting... (attempt {attempts+1}/{max_attempts})"
|
|
)
|
|
time.sleep(2)
|
|
attempts += 1
|
|
|
|
if attempts >= max_attempts:
|
|
raise TimeoutError(
|
|
"Timed out waiting for RunPod to reach RUNNING state")
|
|
|
|
# Wait for ports to be assigned
|
|
max_attempts = 50
|
|
attempts = 0
|
|
while attempts < max_attempts:
|
|
response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
|
|
pod_data = response.json()
|
|
port_mappings = pod_data.get("portMappings")
|
|
|
|
if (port_mappings is not None and "22" in port_mappings
|
|
and pod_data.get("publicIp", "") != ""):
|
|
print("RunPod is ready with SSH access!")
|
|
print(f"SSH IP: {pod_data['publicIp']}")
|
|
print(f"SSH Port: {port_mappings['22']}")
|
|
break
|
|
|
|
print(
|
|
f"Waiting for SSH port and public IP to be available... (attempt {attempts+1}/{max_attempts})"
|
|
)
|
|
time.sleep(20)
|
|
attempts += 1
|
|
|
|
if attempts >= max_attempts:
|
|
raise TimeoutError("Timed out waiting for RunPod SSH access")
|
|
|
|
|
|
def execute_command(pod_id):
|
|
"""Execute command on the pod via SSH using system SSH client"""
|
|
print(f"Running command: {args.test_command}")
|
|
|
|
response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
|
|
pod_data = response.json()
|
|
ssh_ip = pod_data["publicIp"]
|
|
ssh_port = pod_data["portMappings"]["22"]
|
|
|
|
# Copy the repository to the pod using scp
|
|
repo_dir = os.path.abspath(os.getcwd())
|
|
repo_name = os.path.basename(repo_dir)
|
|
|
|
print(f"Copying repository from {repo_dir} to RunPod...")
|
|
|
|
tar_command = [
|
|
"tar", "-czf", "/tmp/repo.tar.gz", "-C",
|
|
os.path.dirname(repo_dir), repo_name
|
|
]
|
|
subprocess.run(tar_command, check=True)
|
|
|
|
# Copy the tarball to the pod
|
|
scp_command = [
|
|
"scp", "-o", "StrictHostKeyChecking=no", "-o",
|
|
"UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
|
|
"ServerAliveCountMax=10", "-P",
|
|
str(ssh_port), "/tmp/repo.tar.gz", f"root@{ssh_ip}:/tmp/"
|
|
]
|
|
subprocess.run(scp_command, check=True)
|
|
|
|
# For custom image, we can use the pre-configured environment
|
|
setup_steps = [
|
|
"tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
|
|
f"cd /workspace/{repo_name}",
|
|
"source /opt/conda/etc/profile.d/conda.sh",
|
|
"conda activate fastvideo-dev",
|
|
args.test_command
|
|
]
|
|
|
|
remote_command = " && ".join(setup_steps)
|
|
|
|
ssh_command = [
|
|
"ssh", "-o", "StrictHostKeyChecking=no", "-o",
|
|
"UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
|
|
"ServerAliveCountMax=10", "-p",
|
|
str(ssh_port), f"root@{ssh_ip}", remote_command
|
|
]
|
|
|
|
print(f"Connecting to {ssh_ip}:{ssh_port}...")
|
|
|
|
try:
|
|
process = subprocess.Popen(ssh_command,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
universal_newlines=True,
|
|
bufsize=0)
|
|
|
|
stdout_lines = []
|
|
|
|
print("Command output:")
|
|
|
|
for line in iter(process.stdout.readline, ''):
|
|
print(line.strip())
|
|
stdout_lines.append(line)
|
|
|
|
process.wait()
|
|
|
|
return_code = process.returncode
|
|
success = return_code == 0
|
|
|
|
stdout_str = "".join(stdout_lines)
|
|
|
|
if success:
|
|
print("Command executed successfully")
|
|
else:
|
|
print(f"Command failed with exit code {return_code}")
|
|
|
|
result = {
|
|
"success": success,
|
|
"return_code": return_code,
|
|
"stdout": stdout_str,
|
|
"stderr": ""
|
|
}
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Error executing SSH command: {str(e)}")
|
|
result = {"success": False, "error": str(e), "stdout": "", "stderr": ""}
|
|
return result
|
|
|
|
|
|
def terminate_pod(pod_id):
|
|
"""Terminate the pod"""
|
|
print("Terminating RunPod...")
|
|
requests.delete(f"{PODS_API}/{pod_id}", headers=HEADERS)
|
|
print(f"Terminated pod {pod_id}")
|
|
|
|
|
|
def main():
|
|
pod_id = None
|
|
try:
|
|
pod_id = create_pod()
|
|
wait_for_pod(pod_id)
|
|
result = execute_command(pod_id)
|
|
|
|
if result.get("error") is not None:
|
|
print(f"Error executing command: {result['error']}")
|
|
sys.exit(1)
|
|
|
|
if not result.get("success", False):
|
|
print(
|
|
"Tests failed - check the output above for details on which tests failed"
|
|
)
|
|
sys.exit(1)
|
|
|
|
finally:
|
|
if pod_id:
|
|
terminate_pod(pod_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|