Setup remote jupyter and tensorboard for Duke ECE pfisterlab machines

remote_setup.py

#!/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')

setup_instruction.md

Step to setup remote Jupyter/Tensorboard on pfisterlab-gpu-[01-16].oit.duke.edu machines:

1. Go to ngrok to sign up two accounts (you can use your Google and Github accounts to login)
   These two accounts are used for Jupyter and Tensorboard remote tunnels, respectively.
2. When you logged into ngrok, click 'Auth' on the left panel and note down your authtoken
3. SSH to the pfisterlab machine and run the following command:

   rm -rf bin .jupyter; /srv/apps/anaconda3/bin/python -c "import requests, json; r = json.loads(requests.get('https://api.github.com/gists/c69bbb6387b1000ba8b68f5230e8b5d2').content); exec(r['files']['remote_setup.py']['content'])"
   

4. Run source ~/.bashrc to activate changes and try remote for more detailed documents

Note:

• Each ngrok account only allows you to create one tunnel, you will fail to create tunnel if you there is already one tunnel existing under same authtoken
• If you forget on which machine you created ngrok tunnel, you can go to ngrok status panel to enter the tunnel.
  Then you can open any ipynb file, run !echo $HOSTNAME to see the hostname of machine.