ipyncarver.py

#!/usr/bin/env python
"""In order to use the script you need to copy your SSH key to the target server
and also copy the server SSH public key (usually .ssh/id_rsa.pub) to .ssh/authorized_keys, 
so that the computing node can ssh passwordless to the login node"""

from subprocess import Popen, PIPE, call
import sys
import webbrowser
from getopt import getopt
import time

args = dict(getopt(sys.argv[1:],[],['port=','dir='])[0])

port = args.get('--port',8789)
dir = args.get('--dir',None)


def readwhile(stream,func):
    while True:
        line = stream.readline()
        if line!='':
            print line[:-1]
            if func(line): break
        else:
            raise Exception("Disconnected unexpectedly.") 

pqsub=Popen(['ssh','-t','-t','-4','carver.nersc.gov'],stdin=PIPE,stdout=PIPE,stderr=PIPE)
pqsub.stdin.write('qsub -I -V -q usplanck -l nodes=1:ppn=1,pvmem=20gb -l walltime=12:00:00\n')
pqsub.stdin.write('echo HOSTNAME=`hostname`\n')

def gethostname(line):
    global hostname
    if line.startswith('HOSTNAME'):
        hostname = line.split('=')[1].strip()
        return True

readwhile(pqsub.stdout, gethostname)

if dir: 
    pqsub.stdin.write('cd %s\n'%dir)
    pqsub.stdin.write('echo CD\n')
    readwhile(pqsub.stdout, lambda line: line.startswith('CD'))
    
pqsub.stdin.write('ipython notebook --pylab=inline --port=%s\n'%port)
readwhile(pqsub.stdout, lambda line: line.startswith('[NotebookApp]'))


tunnel = ['ssh','-4', '-t', '-Y', 'carver.nersc.gov', '-L', '%s:localhost:%s'%(port,port), 'ssh', '-t', '-Y', hostname, '-L', '%s:localhost:%s'%(port,port)]
print ' '.join(tunnel)
ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE)
ptunnel.stdin.write('echo TUNNEL\n')
readwhile(ptunnel.stdout,lambda line: line.startswith('TUNNEL'))

webbrowser.open('http://localhost:%s'%port)

print "Succesfully opened notebook!"
print "Kill this process to end your notebook connection."
time.sleep(12*60**2)

pqsub.kill()
ptunnel.kill()

print "Succesfully cleanup up connections."