|
#!/srv/apps/anaconda3/bin/python |
|
import os |
|
import time |
|
|
|
HOME = os.environ.get('HOME') |
|
|
|
# edit bashrc |
|
with open(HOME + '/.bashrc', 'w') as f: |
|
ngrok_key1 = input('Enter the first ngrok authtoken > ') |
|
ngrok_key2 = input('Enter the second ngrok authtoken > ') |
|
bashrc_snippet = f'''# .bashrc |
|
|
|
# Source global definitions |
|
if [ -f /etc/bashrc ]; then |
|
. /etc/bashrc |
|
fi |
|
|
|
# Uncomment the following line if you don't like systemctl's auto-paging feature: |
|
# export SYSTEMD_PAGER= |
|
|
|
# User specific aliases and functions |
|
########################################################################################## |
|
# 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) |
|
|
|
# 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 /srv/apps/work_env/work_env.img 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, |
|
'jupyter_ngrok_binding_port': None, |
|
'tboard_port': None, |
|
'tboard_ngrok_port': None, |
|
'tboard_ngrok_binding_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] |
|
} |
|
|
|
|
|
def get_available_port(process): |
|
port = json.load(open(config_path))[f'{process}_port_default'] |
|
while True: |
|
if all(conn.laddr.port != port for conn in psutil.net_connections()): |
|
return port |
|
print(f'port {port} is not available, trying port {port + 1}') |
|
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 proc.username() == os.environ.get('USER') and \ |
|
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: |
|
for proc in psutil.process_iter(): |
|
if proc.username() == os.environ.get('USER') and \ |
|
any(kw in proc.name() for kw in keyword[process]) and \ |
|
any(conn.laddr.port == port for conn in proc.connections()): |
|
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 |
|
if process in ['jupyter_ngrok', 'tboard_ngrok']: |
|
cfg[f'{process}_binding_port'] = 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}') |
|
f.write('\\n') |
|
f.write(f'web_addr: localhost:{ngrok_port}') |
|
os.system(f'~/bin/ngrok http {port} -config={ngrok_config_path} > /dev/null &') |
|
cfg['jupyter_ngrok_binding_port'] = port |
|
with open(config_path, 'w') as f: |
|
json.dump(cfg, f) |
|
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}') |
|
f.write('\\n') |
|
f.write(f'web_addr: localhost:{ngrok_port}') |
|
os.system(f'~/bin/ngrok http {port} -config={ngrok_config_path} > /dev/null &') |
|
cfg['tboard_ngrok_binding_port'] = port |
|
with open(config_path, 'w') as f: |
|
json.dump(cfg, f) |
|
while find_process(process, False) is None: |
|
time.sleep(1) |
|
elif process in ['jupyter_ngrok', 'tboard_ngrok']: |
|
port, binding_port = cfg[f'{process[:-6]}_port'], cfg[f'{process}_binding_port'] |
|
if port != binding_port: |
|
print(f'{process}_binding port is different from {process[:-6]}_port, restarting ngrok') |
|
kill_process(process) |
|
start_process(process) |
|
|
|
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 not os.path.exists(config_path) or set(json.load(open(config_path, 'r')).keys()) != set(default_cfg.keys()): |
|
with open(config_path, 'w') as f: |
|
json.dump(default_cfg, f) |
|
|
|
help_doc = """Remote Jupyter/Tensorboard/ngrok configurator |
|
|
|
Manage remote Jupyter: |
|
remote start |
|
remote status |
|
remote clean |
|
remote kill [jupyter|jupyter_ngrok] |
|
|
|
Manage remote Tensorboard: |
|
remote tboard_start |
|
remote tboard_status |
|
remote tboard_clean |
|
remote tboard_kill [tboard|tboard_ngrok] |
|
|
|
Print config and help |
|
remote config |
|
remote help |
|
""" |
|
if len(sys.argv) == 1: |
|
print(help_doc) |
|
else: |
|
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))) |
|
elif sys.argv[1] == 'help': |
|
print(help_doc) |
|
|
|
''' |
|
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') |