Start a Jupyter server on a remote server and tunnel to localhost

nbserver.py

#!/usr/bin/env python

from __future__ import print_function
import subprocess
import argparse
from contextlib import closing
import socket
import random

__author__ = 'Yao-Yuan Mao'


def check_port_open(port):
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
        try:
            s.bind(('', port)) # pylint: disable=E1101
        except socket.error:
            return False
    return True


def get_random_port():
    return random.randint(50000, 59999)


def get_random_open_port():
    while True:
        port = get_random_port()
        if check_port_open(port):
            return port


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('server')
    parser.add_argument('-p', '--remote-port', type=int)
    parser.add_argument('-l', '--local-port', type=int)
    parser.add_argument('-d', '--remote-notebook-dir', default='~')
    parser.add_argument('-e', '--use-existing-remote-instance', action='store_true')
    parser.add_argument('-n', '--notebook', action='store_true')
    parser.add_argument('-c', '--prepend-command')
    args = parser.parse_args()

    if args.local_port:
        if not check_port_open(args.local_port):
            raise RuntimeError('local port {} in use'.format(args.local_port))
        local_port = args.local_port
    else:
        local_port = get_random_open_port()

    if args.remote_port:
        remote_port = args.remote_port
    elif args.use_existing_remote_instance:
        raise parser.error('Must specify --remote-port [PORT] when --use-existing-remote-instance is set')
    elif args.local_port:
        remote_port = get_random_port()
    else:
        remote_port = local_port

    call = [
        'ssh',
        ('-N' if args.use_existing_remote_instance else '-t'),
        '-L',
        '{}:localhost:{}'.format(local_port, remote_port),
        args.server,
    ]

    if not args.use_existing_remote_instance:
        if args.prepend_command:
            prepend_command = args.prepend_command.strip()
            if prepend_command.endswith(';') or prepend_command.endswith('&&'):
                prepend_command += ' '
            else:
                prepend_command += ' && '
        else:
            prepend_command = ''

        call.append('{}jupyter {} --no-browser --notebook-dir={} --port-retries=0 --port={}'.format(
            prepend_command,
            ('notebook' if args.notebook else 'lab'),
            args.remote_notebook_dir,
            remote_port,
        ))

    print()
    print('Connecting to', args.server)
    print('Jupyter will be served at http://localhost:{}'.format(local_port))
    print()

    try:
        subprocess.call(call)
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()