91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import uuid
|
|
|
|
import requests
|
|
|
|
API_KEY = os.environ['RUNPOD_API_KEY']
|
|
RUN_ID = os.environ.get('GITHUB_RUN_ID', str(uuid.uuid4()))
|
|
PODS_API = "https://rest.runpod.io/v1/pods"
|
|
HEADERS = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {API_KEY}"
|
|
}
|
|
|
|
|
|
def get_job_ids():
|
|
"""Parse job IDs from environment variable"""
|
|
job_ids_str = os.environ.get('JOB_IDS')
|
|
try:
|
|
job_ids = json.loads(job_ids_str)
|
|
if not isinstance(job_ids, list):
|
|
print("Error: JOB_IDS is not a list.")
|
|
sys.exit(1)
|
|
return job_ids
|
|
except json.JSONDecodeError as e:
|
|
print(f"Error parsing JOB_IDS: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
def cleanup_pods():
|
|
"""Find and terminate RunPod instances"""
|
|
print(f"Run ID: {RUN_ID}")
|
|
|
|
single_job_id = os.environ.get('JOB_ID')
|
|
|
|
if single_job_id:
|
|
job_ids = [single_job_id]
|
|
print(f"Job ID: {single_job_id}")
|
|
else:
|
|
job_ids = get_job_ids()
|
|
print(f"Job IDs: {job_ids}")
|
|
|
|
# Get all pods associated with RunPod API_KEY
|
|
try:
|
|
response = requests.get(PODS_API, headers=HEADERS)
|
|
response.raise_for_status()
|
|
pods = response.json()
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error getting pods: {e}")
|
|
sys.exit(1)
|
|
|
|
# Find and terminate pods created by this workflow run
|
|
terminated_pods = []
|
|
for pod in pods:
|
|
pod_name = pod.get("name", "")
|
|
pod_id = pod.get("id")
|
|
|
|
# Check if this pod was created by one of our jobs
|
|
if any(f"{job_id}-{RUN_ID}" in pod_name for job_id in job_ids):
|
|
print(f"Found pod: {pod_id} ({pod_name})")
|
|
try:
|
|
print(f"Terminating pod {pod_id}...")
|
|
term_response = requests.delete(f"{PODS_API}/{pod_id}",
|
|
headers=HEADERS)
|
|
term_response.raise_for_status()
|
|
terminated_pods.append(pod_id)
|
|
print(f"Successfully terminated pod {pod_id}")
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error terminating pod {pod_id}: {e}")
|
|
sys.exit(1)
|
|
|
|
if terminated_pods:
|
|
if single_job_id:
|
|
print(f"Terminated pod: {terminated_pods[0]}")
|
|
else:
|
|
print(f"Terminated {len(terminated_pods)} pods: {terminated_pods}")
|
|
else:
|
|
if single_job_id:
|
|
print(f"No pod found matching pattern: {single_job_id}-{RUN_ID}")
|
|
else:
|
|
print("No pods found to terminate.")
|
|
|
|
|
|
def main():
|
|
cleanup_pods()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|