def neptune_ml_endpoint(args: argparse.Namespace, client: Client,
output: widgets.Output, params):
if args.which_sub == 'create':
if params is None or params == '' or params == {}:
params = {"id": args.job_id, 'instanceType': args.instance_type}
create_endpoint_res = client.endpoints_create(args.model_job_id,
**params)
create_endpoint_res.raise_for_status()
create_endpoint_job = create_endpoint_res.json()
if args.wait:
return wait_for_endpoint(create_endpoint_job['id'], client, output,
args.wait_interval, args.wait_timeout)
else:
return create_endpoint_job
elif args.which_sub == 'status':
if args.wait:
return wait_for_endpoint(args.job_id, client, output,
args.wait_interval, args.wait_timeout)
else:
endpoint_status = client.endpoints_status(args.job_id)
endpoint_status.raise_for_status()
return endpoint_status.json()
else:
return f'Sub parser "{args.which} {args.which_sub}" was not recognized'
def wait_for_endpoint(job_id: str,
client: Client,
output: widgets.Output,
wait_interval: int = DEFAULT_WAIT_INTERVAL,
wait_timeout: int = DEFAULT_WAIT_TIMEOUT):
job_id_output = widgets.Output()
update_status_output = widgets.Output()
with output:
display(job_id_output, update_status_output)
with job_id_output:
print(f'Wait called on endpoint creation job {job_id}')
with update_status_output:
beginning_time = datetime.datetime.utcnow()
while datetime.datetime.utcnow() - beginning_time < (
datetime.timedelta(seconds=wait_timeout)):
update_status_output.clear_output()
endpoint_status_res = client.endpoints_status(job_id)
endpoint_status_res.raise_for_status()
status = endpoint_status_res.json()
if status['status'] in ['InService', 'Failed']:
print('Endpoint creation is finished')
return status
else:
print(f'Status is {status["status"]}')
print(f'Waiting for {wait_interval} before checking again...')
time.sleep(wait_interval)
