Commit f70cb757 authored by John Tourtellott's avatar John Tourtellott

Add test scripts to access NERSC queue directly

parent 287be21a
"""
Log into nersc machine and get newt session id
"""
import argparse
import getpass
import requests
nersc_url = 'https://newt.nersc.gov/newt'
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Login to obtain newt session id')
parser.add_argument('--machine', '-m', default='cori', help='cori or edison')
parser.add_argument('--user', '-u', default='johnt', help='user name')
args = parser.parse_args()
# print(args)
print('Username: {}'.format(args.user))
print('Machine: {}'.format(args.machine))
password = getpass.getpass(prompt='Password: ')
# One time password for MFA
otp = input('Enter one time password (for MFA)')
# print('ready')
credentials = {
'username': args.user,
'password': password + str(otp)
}
url = '{}/login/'.format(nersc_url)
r = requests.post(url, data=credentials)
r.raise_for_status()
js = r.json()
if js.get('auth'):
session_id = js.get('newt_sessionid')
print('newt_sessionid: {}'.format(session_id))
else:
print('NOT authenticated, response:')
print(r.text)
"""
Continually logs user's queue on nersc machine
"""
import argparse
import datetime
import time
import requests
nersc_url = 'https://newt.nersc.gov/newt'
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Login to obtain newt session id')
parser.add_argument('--machine', '-m', default='cori', help='cori or edison')
parser.add_argument('--outfile', '-o', help='output logfile')
parser.add_argument('--user', '-u', default='johnt', help='user name')
parser.add_argument('--wait_sec', '-w', default=5, help='delay between each queue request')
# Add session id either as string or file to read
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('--sessionid', '-s', help='newt session id')
group.add_argument('--id_filename', '-i', help='filename containing newt session id')
args = parser.parse_args()
# print(args)
print('Username: {}'.format(args.user))
print('Machine: {}'.format(args.machine))
url = '{}/queue/{}?user={}'.format(nersc_url, args.machine, args.user)
newt_sessionid = None
if args.sessionid:
newt_sessionid
elif args.id_filename:
with open(args.id_filename) as f:
newt_sessionid = f.read().strip()
if not newt_sessionid:
raise RuntimeError('Not able to read session id file {}'.format(args.if_filename))
cookies = dict(newt_sessionid=newt_sessionid)
# Rest logfile
if args.outfile:
with open(args.outfile, 'w') as f:
pass
start_time = datetime.datetime.now()
td = None # time delta
print('Begin logging loop: use Control-C to break')
try:
while True:
if td is None:
td = datetime.timedelta()
else:
td = datetime.datetime.now() - start_time
r = requests.get(url, cookies=cookies)
r.raise_for_status()
js = r.json()
data = '{}: {}'.format(td, js)
print(data)
# Log to file
if args.outfile:
complete = False
with open(args.outfile, 'a') as f:
f.write(data)
f.write('\n')
complete = True
if not complete:
print('Unabled to write to output file {}'.format(args.outfile))
# Delay
time.sleep(args.wait_sec)
except KeyboardInterrupt:
print('done')
r = requests.post(url, data=credentials)
r.raise_for_status()
js = r.json()
if js.get('auth'):
session_id = js.get('newt_sessionid')
print('newt_sessionid: {}'.format(session_id))
else:
print('NOT authenticated, response:')
print(r.text)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment