Skip to content

Instantly share code, notes, and snippets.

@Joshua1989

Joshua1989/remote_setup

Created August 14, 2018 06:18
Show Gist options
  • Save Joshua1989/79c60786dd70331834b1f468b1b499a7 to your computer and use it in GitHub Desktop.
Save Joshua1989/79c60786dd70331834b1f468b1b499a7 to your computer and use it in GitHub Desktop.
Download ZIP
Setup remote jupyter and tensorboard for Duke ECE pfisterlab machines:
Raw
remote_setup
#!/srv/apps/anaconda3/bin/python
import os
import time
HOME = os.environ.get('HOME')
# edit bashrc
if 'Added by remote_setup.sh' not in open(HOME + '/.bashrc', 'r').read():
with open(HOME + '/.bashrc', 'a') as f:
ngrok_key1 = input('Enter the first ngrok authtoken > ')
ngrok_key2 = input('Enter the second ngrok authtoken > ')
bashrc_snippet = f'''
##########################################################################################
# Added by remote_setup.sh
# User specific aliases and functions
export ANACONDA_PATH = "/srv/apps/anaconda3/bin"
export ANACONDA_PACKAGES = "/srv/apps/anaconda3/lib/python3.6/site-packages"
export PATH = "$HOME/bin:$ANACONDA_PATH:$PATH"
# authorize two ngrok accounts, one for remote Jupyter one for tensorboard
export NGROK_KEY1 = "{ngrok_key1}"
export NGROK_KEY2 = "{ngrok_key2}"
# alias to start working environment
export VM_DRIVERS = "/srv/apps/work_env/vm_cuda_driver"
export VM_IMAGE = "/srv/apps/work_env/work_env.img"
export VM_CONFIG = "-B $VM_DRIVERS:/nvlib,$VM_DRIVERS:/nvbin,$ANACONDA_PACKAGES:/packages,$HOME/user_packages:/script"
alias load_anaconda = "singularity run $VM_CONFIG $VM_IMAGE"
alias utop = "top -U $USER"
alias remote = "$HOME/bin/remote_jupyter.py"
##########################################################################################
'''
f.write(bashrc_snippet)
os.system(f'source {HOME}/.bashrc')
# edit jupyter config
notebook_json = '''{
"load_extensions": {
"codefolding/main": true,
"execute_time/ExecuteTime": true,
"freeze/main": true,
"scratchpad/main": true,
"hide_input_all/main": true,
"varInspector/main": true,
"collapsible_headings/main": true,
"select_keymap/main": true
}
}
'''
if not os.path.exists(HOME + '/.jupyter'):
os.system(f'mkdir {HOME}/.jupyter')
if not os.path.exists(HOME + '/.jupyter/nbconfig'):
os.system(f'mkdir {HOME}/.jupyter/nbconfig')
with open(HOME + '/.jupyter/nbconfig/notebook.json', 'w') as f:
f.write(notebook_json)
if not os.path.exists(HOME + '/.jupyter/jupyter_notebook_config.json'):
print('set password for remote jupyter')
os.system('singularity exec $VM_IMAGE jupyter notebook password')
# set bin folder for user directory
remote_jupyter_py = '''#!/srv/apps/anaconda3/bin/python
import json
import os
import pprint
import psutil
import requests
import sys
import time
config_path = os.environ.get('HOME') + '/bin/remote_jupyter_config.json'
ngrok_config_path = os.environ.get('HOME') + '/bin/ngrok2.yml'
keyword = {
'jupyter': ['jupyter', 'ZMQbg'],
'jupyter_ngrok': ['ngrok'],
'tboard': ['tensorboard'],
'tboard_ngrok': ['ngrok']
}
default_cfg = {
'jupyter_port_default': 8898,
'jupyter_ngrok_port_default': 4000,
'tboard_port_default': 6006,
'tboard_ngrok_port_default': 4040,
'jupyter_port': None,
'jupyter_ngrok_port': None,
'tboard_port': None,
'tboard_ngrok_port': None,
'jupyter_pid': None,
'jupyter_ngrok_pid': None,
'tboard_pid': None,
'tboard_ngrok_pid': None,
'ngrok_keys': [v for k, v in os.environ.items() if 'NGROK_KEY' in k]
}
jupyter_cmd = 'singularity exec $VM_CONFIG $VM_IMAGE /usr/local/anaconda3/bin/jupyter notebook --no-browser --port={port} > /dev/null 2>&1 &'
if not os.path.exists(config_path):
with open(config_path, 'w') as f:
json.dump(default_cfg, f)
def get_available_port(process):
port = json.load(open(config_path))[f'{process}_port_default']
while True:
for conn in psutil.net_connections(kind='tcp'):
if not (conn.laddr[1] == port and conn.status == psutil.CONN_LISTEN):
return port
port += 1
def find_process(process, write_on_missing=True):
cfg = json.load(open(config_path))
pid, port = cfg[f'{process}_pid'], cfg[f'{process}_port']
if pid is not None and psutil.pid_exists(pid) and port is not None:
proc = psutil.Process(pid)
# both pid and port exists
if any(kw in proc.name() for kw in keyword[process]) and \
any(conn.laddr.port == port for conn in proc.connections()):
return proc
elif pid is None and port is not None:
# only port exists, need to find corresponding pid
procs = [proc for proc in psutil.process_iter() if
any(kw in proc.name() for kw in keyword[process]) and
any(conn.laddr.port == port for conn in proc.connections())]
if len(procs) > 0:
proc = procs[0]
cfg[f'{process}_pid'] = proc.pid
with open(config_path, 'w') as f:
json.dump(cfg, f)
return proc
if write_on_missing:
# Wrong pid and port, write back to None
cfg[f'{process}_pid'], cfg[f'{process}_port'] = None, None
with open(config_path, 'w') as f:
json.dump(cfg, f)
def start_process(process, **kwargs):
cfg = json.load(open(config_path))
proc = find_process(process)
if proc is None:
port = get_available_port(process)
cfg[f'{process}_port'] = port
with open(config_path, 'w') as f:
json.dump(cfg, f)
if process == 'jupyter':
print(f'start jupyter notebook at port {port}')
os.system(f'singularity exec $VM_CONFIG $VM_IMAGE /usr/local/anaconda3/bin/jupyter notebook --no-browser --port={port} > /dev/null 2>&1 &')
elif process == 'jupyter_ngrok':
ngrok_port = port
port, ngrok_key = cfg['jupyter_port'], cfg['ngrok_keys'][0]
print(f'start ngrok at port {ngrok_port} binding with jupyter port {port}')
with open(ngrok_config_path, 'w') as f:
f.write(f'authtoken: {ngrok_key}\\nweb_addr: localhost:{ngrok_port}')
os.system(f'~/bin/ngrok http {port} -config={ngrok_config_path} > /dev/null &')
elif process == 'tboard':
log_dir = kwargs.get('log_dir') or os.environ.get('HOME')
print(f'start tensorboard at port {port} with log dir {log_dir}')
os.system(f'tensorboard --logdir {log_dir} --host 0.0.0.0 --port {port} > /dev/null &')
elif process == 'tboard_ngrok':
ngrok_port = port
port, ngrok_key = cfg['tboard_port'], cfg['ngrok_keys'][1]
print(f'start ngrok at port {ngrok_port} binding with tensorboard port {port}')
with open(ngrok_config_path, 'w') as f:
f.write(f'authtoken: {ngrok_key}\\nweb_addr: localhost:{ngrok_port}')
os.system(f'~/bin/ngrok http {port} -config={ngrok_config_path} > /dev/null &')
while find_process(process, False) is None:
time.sleep(1)
return find_process(process)
def kill_process(process):
proc = find_process(process)
if proc is not None:
proc.kill()
cfg = json.load(open(config_path))
cfg[f'{process}_pid'], cfg[f'{process}_port'] = None, None
with open(config_path, 'w') as f:
json.dump(cfg, f)
def list_pid(mode):
proc = find_process(mode)
ngrok = find_process(f'{mode}_ngrok')
if None not in [proc, ngrok]:
cfg = json.load(open(config_path))
proc_pid, proc_port = cfg[f'{mode}_pid'], cfg[f'{mode}_port']
proc_ngrok_pid, proc_ngrok_port = cfg[f'{mode}_ngrok_pid'], cfg[f'{mode}_ngrok_port']
print(f'{mode} PID/Port = {proc_pid}/{proc_port}, ngrok PID/Port = {proc_ngrok_pid}/{proc_ngrok_port}')
return [proc and proc.pid, ngrok and ngrok.pid]
def ngrok_link(mode):
ngrok_port = json.load(open(config_path))[f'{mode}_ngrok_port']
retval = requests.get(f'http://localhost:{ngrok_port}/api/tunnels').json()['tunnels']
if len(retval) > 0:
print(f'{mode} ngrok link:', retval[0]['public_url'].strip())
return retval[0]['public_url'].strip()
else:
return ''
if __name__ == "__main__":
if len(sys.argv) >= 2:
if sys.argv[1] == 'start':
start_process('jupyter')
start_process('jupyter_ngrok')
wait = 0
while not ngrok_link('jupyter'):
time.sleep(1)
wait += 1
print(f'waited {wait} seconds', end='\\r')
elif sys.argv[1] == 'status':
if None not in list_pid('jupyter'):
ngrok_link('jupyter')
elif sys.argv[1] == 'clean':
kill_process('jupyter')
kill_process('jupyter_ngrok')
elif sys.argv[1] == 'kill':
if len(sys.argv) >= 3:
kill_process(sys.argv[2])
elif sys.argv[1] == 'tboard_start':
start_process('tboard', log_dir=sys.argv[2] if len(sys.argv) >= 3 else None)
start_process('tboard_ngrok')
wait = 0
while not ngrok_link('tboard'):
time.sleep(1)
wait += 1
print(f'waited {wait} seconds', end='\\r')
elif sys.argv[1] == 'tboard_status':
if None not in list_pid('tboard'):
ngrok_link('tboard')
elif sys.argv[1] == 'tboard_clean':
kill_process('tboard')
kill_process('tboard_ngrok')
elif sys.argv[1] == 'tboard_kill':
if len(sys.argv) >= 3:
kill_process(sys.argv[2])
elif sys.argv[1] == 'config':
pprint.pprint(json.load(open(config_path)))
'''
if not os.path.exists(HOME + '/bin'):
os.system(f'mkdir {HOME}/bin')
if not os.path.exists(HOME + '/bin/ngrok'):
os.system(f'wget -O {HOME}/bin/ngrok.zip https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip')
while not os.path.exists(HOME + '/bin/ngrok.zip'):
time.sleep(1)
os.system(f'unzip {HOME}/bin/ngrok.zip -d {HOME}/bin')
os.system(f'rm {HOME}/bin/ngrok.zip')
if not os.path.exists(HOME + '/bin/ngrok/remote_jupyter.py'):
with open(HOME + '/bin/remote_jupyter.py', 'w') as f:
f.write(remote_jupyter_py)
os.system(f'chmod 755 {HOME}/bin/remote_jupyter.py')
if not os.path.exists(HOME + '/user_packages'):
os.system(f'mkdir {HOME}/user_packages')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment