60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
import json
|
|
import logging
|
|
import time
|
|
|
|
import requests
|
|
|
|
from bots.data import config
|
|
|
|
|
|
class RunpodClient:
|
|
INSTANCE_UNAVAILABLE_MSG = (
|
|
"There are no longer any instances available with the requested specifications. "
|
|
"Please refresh and try again."
|
|
)
|
|
|
|
def __init__(self, logger=None):
|
|
self.url = f"https://api.runpod.io/graphql?api_key={config['runpod_token']}"
|
|
self.logger = logger or logging.getLogger(__name__)
|
|
|
|
# Return True if a pod with pod_name is currently running.
|
|
def is_running(self, pod_name):
|
|
fetch_pods_data = {
|
|
"query": "query Pods { myself { pods { id name } } }",
|
|
}
|
|
|
|
fetch_response = requests.post(
|
|
self.url,
|
|
headers={"Content-Type": "application/json"},
|
|
data=json.dumps(fetch_pods_data),
|
|
)
|
|
|
|
if not fetch_response.ok:
|
|
self.logger.error("Error fetching pods.")
|
|
return
|
|
|
|
pods = fetch_response.json()["data"]["myself"]["pods"]
|
|
return any(pod["name"] == pod_name for pod in pods)
|
|
|
|
def create_instance(self, json_file):
|
|
with open(json_file, "r") as file:
|
|
runpod_query = json.load(file)
|
|
|
|
while True:
|
|
response = requests.post(
|
|
self.url,
|
|
headers={"Content-Type": "application/json"},
|
|
json=runpod_query,
|
|
)
|
|
|
|
if "errors" in response.json():
|
|
error_message = response.json()["errors"][0]["message"]
|
|
|
|
if error_message == RunpodClient.INSTANCE_UNAVAILABLE_MSG:
|
|
logging.warn("No instances available, retrying in 1 second...")
|
|
time.sleep(1)
|
|
else:
|
|
logging.error(f"Unhandled error: {error_message}")
|
|
else:
|
|
break
|