From d24b5a47eff3d0831e2cd9fe8db93ede5d4eb00f Mon Sep 17 00:00:00 2001 From: Ashley Kleynhans Date: Thu, 13 Jul 2023 22:19:51 +0200 Subject: [PATCH] Updated test endpoint script --- test_endpoint.py | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/test_endpoint.py b/test_endpoint.py index e3053ba..956cc18 100755 --- a/test_endpoint.py +++ b/test_endpoint.py @@ -59,7 +59,10 @@ def save_result_image(resp_json): if 'output' in resp_json and 'image' in resp_json['output']: save_result_image(resp_json) else: - if 'status' in resp_json and resp_json['status'] == 'IN_QUEUE': + job_status = resp_json['status'] + print(f'Job status: {job_status}') + + if job_status == 'IN_QUEUE' or job_status == 'IN_PROGRESS': request_id = resp_json['id'] request_in_queue = True @@ -75,17 +78,27 @@ def save_result_image(resp_json): if r.status_code == 200: resp_json = r.json() + job_status = resp_json['status'] - if resp_json['status'] == 'IN_QUEUE' or resp_json['status'] == 'IN_PROGRESS': - print(f'RunPod inswapper request {request_id} is still in the queue, sleeping for 5 seconds...') + if job_status == 'IN_QUEUE' or job_status == 'IN_PROGRESS': + print(f'RunPod request {request_id} is {job_status}, sleeping for 5 seconds...') time.sleep(5) - elif resp_json['status'] == 'FAILED': - print(f'RunPod inswapper request {request_id} failed') - elif resp_json['status'] == 'COMPLETED': - print(f'RunPod inswapper request {request_id} completed') + elif job_status == 'FAILED': + request_in_queue = False + print(f'RunPod request {request_id} failed') + elif job_status == 'COMPLETED': + request_in_queue = False + print(f'RunPod request {request_id} completed') save_result_image(resp_json) + elif job_status == 'TIMED_OUT': + request_in_queue = False + print(f'ERROR: RunPod request {request_id} timed out') else: - print(f'Invalid status response from RunPod status endpoint') + request_in_queue = False + print(f'ERROR: Invalid status response from RunPod status endpoint') + print(json.dumps(resp_json, indent=4, default=str)) + elif job_status == 'COMPLETED' and resp_json['output']['status'] == 'error': + print(f'ERROR: {resp_json["output"]["message"]}') else: print(json.dumps(resp_json, indent=4, default=str)) else: