Skip to content

Instantly share code, notes, and snippets.

@yymao

yymao/nbserver.py

Last active Aug 8, 2019
Embed
What would you like to do?
Start a Jupyter server on a remote server and tunnel to localhost
#!/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'
server_dict = {
'ki-ls' : 'ki-ls.slac.stanford.edu',
'cori' : 'cori.nersc.gov',
'edision': 'edison.nersc.gov',
}
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', nargs='?', default='cori')
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
server = server_dict.get(args.server, args.server)
notebook_dir = args.remote_notebook_dir
call = [
'ssh',
('-N' if args.use_existing_remote_instance else '-t'),
'-L',
'{}:localhost:{}'.format(local_port, remote_port),
server,
]
if not args.use_existing_remote_instance:
if args.prepend_command:
prepend_command = args.prepend_command
if prepend_command.endswith(';') or prepend_command.endswith('&&'):
prepend_command += ' '
else:
prepend_command += ' && '
elif server.endswith('.nersc.gov'):
prepend_command = 'module unload python; module load python/3.6-anaconda-4.4 && '
else:
prepend_command = ''
call.append('{}jupyter {} --no-browser --notebook-dir={} --port-retries=0 --port={}'.format(
prepend_command,
('notebook' if args.notebook else 'lab'),
notebook_dir,
remote_port,
))
print()
print('Connecting to', server)
print('Jupyter will be served at http://localhost:{}'.format(local_port))
print()
try:
subprocess.call(call)
except KeyboardInterrupt:
pass
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment