75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
|
import requests
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
import sys
|
||
|
|
|
||
|
|
args = sys.argv
|
||
|
|
|
||
|
|
# TODO turn all this into config style options or @click-style flags/options
|
||
|
|
logging.basicConfig(level=logging.INFO)
|
||
|
|
|
||
|
|
pod_id = "7x0b7u16s6vyrc"
|
||
|
|
bearer_token = "EIWX9RO18PRXCD0RUSY26MSD062GUF6REQGGV6QB"
|
||
|
|
|
||
|
|
api = f"https://api.runpod.ai/v2/{pod_id}"
|
||
|
|
run_endpoint = f"{api}/run"
|
||
|
|
status_endpoint = f"{api}/status"
|
||
|
|
health_endpoint = f"{api}/health"
|
||
|
|
purge_endpoint = f"{api}/purge-queue"
|
||
|
|
|
||
|
|
headers = {
|
||
|
|
"Content-Type": "application/json",
|
||
|
|
"Authorization": f"Bearer {bearer_token}",
|
||
|
|
}
|
||
|
|
|
||
|
|
input_data = {"input": {"url": "https://0x0.st/H9PH.aac"}}
|
||
|
|
|
||
|
|
@click.command()
|
||
|
|
def cli():
|
||
|
|
...
|
||
|
|
|
||
|
|
def main(args: list[str]) -> None:
|
||
|
|
if len(args) <= 1:
|
||
|
|
sys.exit(0)
|
||
|
|
|
||
|
|
if args[1] == "health":
|
||
|
|
logging.info("requesting health status...")
|
||
|
|
response = requests.get(health_endpoint, headers=headers)
|
||
|
|
elif args[1] == "status":
|
||
|
|
if len(args) <= 2:
|
||
|
|
logging.error("No job id to get status from supplied.")
|
||
|
|
sys.exit(1)
|
||
|
|
logging.info(f"requesting job {args[2]} status...")
|
||
|
|
response = requests.get(f"{status_endpoint}/{args[2]}", headers=headers)
|
||
|
|
elif args[1] == "cancel":
|
||
|
|
if len(args) <= 2:
|
||
|
|
logging.error("No job id to cancel supplied.")
|
||
|
|
sys.exit(1)
|
||
|
|
logging.info(f"requesting job {args[2]} cancellation...")
|
||
|
|
response = requests.get(f"{status_endpoint}/{args[2]}", headers=headers)
|
||
|
|
elif args[1] == "purge":
|
||
|
|
logging.info("purging all jobs in queue...")
|
||
|
|
response = requests.post(purge_endpoint, headers=headers)
|
||
|
|
else:
|
||
|
|
logging.info("requesting new job...")
|
||
|
|
response = requests.post(run_endpoint, json=input_data, headers=headers)
|
||
|
|
|
||
|
|
json = response.json()
|
||
|
|
|
||
|
|
# the json will be similar to
|
||
|
|
# {'id': 'e3d2e250-ea81-4074-9838-1c52d006ddcf', 'status': 'IN_QUEUE'}
|
||
|
|
|
||
|
|
while "status" in json and (
|
||
|
|
json["status"] == "IN_QUEUE" or json["status"] == "IN_PROGRESS"
|
||
|
|
):
|
||
|
|
logging.info(f"{json['status']} for job {json['id']}, waiting...")
|
||
|
|
time.sleep(3)
|
||
|
|
response = requests.get(f"{status_endpoint}/{json['id']}", headers=headers)
|
||
|
|
json = response.json()
|
||
|
|
|
||
|
|
logging.info(json)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main(args)
|