h2jupynb.py

#!/usr/bin/env python
from __future__ import print_function
"""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
import os
import platform
import random
from datetime import datetime
import sys, getopt

random.seed(datetime.now())

username = ''
username = str(username)
timeinhours = 2
timeinhours = int(timeinhours)
memoryingb = 1
memoryingb = int(memoryingb) 
parenv = 'shared'
numberofslots = 1
numberofslots = int(numberofslots)
port = random.randrange(6222, 10000, 2)
port = int(port)
directory = ''
directory = str(directory)
opsys = platform.system()
opsys = str(opsys)
pythonver = 'python/2.7.13'
pythonver = str(pythonver)
usegpu = 'no'
gpu = 'P4'
highp = 'no'
cudaver = '9.1'
exclusive = 'no'
arch = ''

#qrsh_error='no'

#print("ARGV      :", sys.argv[1:])


if sys.version_info<(2,6,0):
  sys.stderr.write("You need python 2.6 or later to run this script\n")
  exit(1)


if opsys == 'Windows':
    try:
        ret=call(['where','ssh.exe'])
	#mytest = 'MINGW' in os.environ['MSYSTEM']
    except:
        print("It appears that Git for Windows is not installed")
        print("or that you are not running this script from its")
        print("Git BASH terminal. Please install either Git for")
        print("Windows, availabe at: ")
        print("https://git-for-windows.github.io/")
        print("and run this script from its git BASH.")
        sys.exit(2)

myssh = 'ssh'


def main(argv):
    global  username, timeinhours, memoryingb, numberofslots, pythonver, port, directory, usegpu, gpu, cudaver, highp, exclusive, arch, parenv
    try:
        opts, args = getopt.getopt(sys.argv[1:],"hu:t:m:e:s:v:p:d:g:c:l:o:x:a:",["help","username=","timeinhours=","memoryingn=","parallel-enviromnment=","numberofslots=","pythonver=","port=","dir=","usegpu=","gpu=","cudaver=","highp=","exclusive=","arch="])
    except getopt.GetoptError:
        print("Usage:")
        print("\n\t h2jupynb [-u <Hoffman2 user name>] [-t <time, integer number of hours>] [-m <memory, integer number of GB per core>] [-e <parallel environment: 1 for shared, 2 for distributed >] [-s <number of slots>] [-v <python-version>] [-p <port>] [-g <yes/no>] [-c <gpu-card-type>] [-l <cuda-version>] [-o <run on group owned nodes: yes/no>] [-d <path to directory from which the jupyter NB/lab will start>] [-x <run on an exclusively reserved node: yes/no>] [-a <CPU-type>]\n")
        print("If no arguments are given to this script it is assumed that:")
        print("\t your Hoffman2 user name is the same as on your client machine")
        print("\t the time duration for your session is of 2 hours")
        print("\t the parallel environment is shared")
        print("\t the memory per slot for your session is of 1GB")
        print("\t the number of slots for your session is of 1")
        print("\t the python version for your notebook is 2.7.13")
        print("\t the port on which the server is started is 8789")
        print("\t the starting directory on Hoffman2 is your $HOME")
        print("\t use GPU? default is no")
        print("\t GPU type default is P4 (if -g yes)")
        print("\t CUDA version 9.1 (if -g yes)")
        print("\t not running on owned nodes")
        print("\t not running in exclusive mode")
        print("\t no specific CPU selected (see ARCH in qhost output)")
        print("\n\t python versions currently available are: 2.7.13, 2.7.16, 3.6.1, 3.7.0, 3.7.2")
        print("\t or: anaconda2, anaconda3 and anaconda/python3-4.2")
        print("\t python versions: 2.7.3 or 3.4 are available but deprecated.\n")
        print("\t cuda versions currently publicly available are: 7.0, 7.5, 8.0, 9.1 and 10.0 \n")
        print("\t GPU cards currently available: P4, RTX2080Ti and V100\n")  
        sys.exit(2)


    #print('OPTIONS   :', opts)


    for opt, arg in opts:
        #print(opt, arg)
        if opt == '-h':
            print("h2jupynb [-u <Hoffman2 user name>] [-t <time, integer number of hours>] [-m <memory, integer number of GB peer core>] [-e <parallel environment: 1 for shared, 2 for distributed>] [-s <number of slots>] [-v <python-version>] [-o <run on group owned nodes: yes/no>] [-x <run on an exclusively reserved node: yes/no>]  [-a <CPU-type>] [-d <path to directory from which the jupyter NB/lab will start>] [-p <port number>] [-g <run on a gpu node: yes/no>] [-c <gpu-card-type>] [-l <cuda-version>]\n")
            sys.exit()
        elif opt in ("-u", "--user"):
            username = arg
            username = str(username)
        elif opt in ("-t", "--time"):
            timeinhours = arg
            try:
                timeinhours=int(timeinhours)
            except exceptions.KeyError:
                print("Setting the time to 2 hours...")
                timeinhours = 2
        elif opt in ("-m", "--memory"):
            memoryingb = arg
            try:
                memoryingb=int(memoryingb)
            except exceptions.KeyError:
                print("Setting the memory to 1GB...")
                memoryingb = 1
        elif opt in ("-e", "--parallel-environment"):
            parenv = arg
            try:
                parenv=int(parenv)
            except exceptions.KeyError:
                print("Setting the parallel environmen to shared...")
                parenv = 1
            if (parenv == 1):
                parenv = 'shared'
            elif (parenv == 2):
                parenv = 'dc\*'
            else:
                print("Setting the parallel environmen to shared...")
                parenv = 1
        elif opt in ("-s", "--slots"):
            numberofslots = arg
            try:
                numberofslots=int(numberofslots)
            except exceptions.KeyError:
                print("Setting the number of slots to 1...")
                numberofslots = 1
        elif opt in ("-v", "--version"):
            pythonver = arg
            pythonver = str(pythonver)
            if (not pythonver == '2.7.3') and (not pythonver == '2.7.13') and (not pythonver == '2.7.16') and (not pythonver == '3.4') and (not pythonver == '3.6.1') and (not pythonver == 'anaconda2') and (not pythonver == '3.6.1') and (not pythonver == '3.7.0') and (not pythonver == '3.7.2') and (not pythonver == 'anaconda3') and (not pythonver == 'anaconda/python2-4.2') and (not pythonver == 'anaconda/python3-4.2'):
                print("version, ", pythonver,", of python not available")
                print("setting python version to 2.7.13")
                pythonver = 'python/2.7.13'
            if (pythonver == '2.7.3') or (pythonver == '2.7.13') or (pythonver == '2.7.16') or (pythonver == '3.4') or (pythonver == '3.6.1') or (pythonver == 'anaconda2') or (pythonver == '3.7.0') or (pythonver == '3.7.2') or (pythonver == 'anaconda3'):
                pythonver = "python/"+pythonver
        elif opt in ("-p", "--port"):
            port = arg
            try:
                port=int(port)
            except exceptions.KeyError:
                print("Setting the port to 2 8789...")
                port = 8789
        elif opt in ("-d", "--dir"):
            directory = arg
            directory = str(directory)
        elif opt in ("-g", "--usegpu"):
            usegpu = arg
            usegpu = str(usegpu)
            if ( usegpu == 'NO') or (usegpu == 'No') or (usegpu == 'no') or (usegpu == 'nO') or (usegpu == 'N') or (usegpu == 'n'):
                print("Jupyter notebook or lab will not run on a GPU node")
                usegpu = 'no'
            elif (usegpu == 'Yes') or (usegpu == 'yes') or (usegpu == 'YES') or (usegpu == 'Y') or (usegpu == 'y'):
                print("Jupyter notebook or lab will run on a GPU node")
                usegpu='yes'
            else:
                print("USEGPU=",usegpu)
        elif opt in ("-c", "--gpu_card"):
            gpu = arg
            gpu = str(gpu)
            if (not gpu == 'P4') and (not gpu == 'K40') and (not gpu == 'GTX1080Ti') and (not gpu == 'RTX2080Ti') and (not gpu == 'V100'):
                print("version, ",gpu,", of GPU not available")
                print("setting GPU version to P4")
                gpu = 'P4'
        elif opt in ("-l", "--cuda"):
            cudaver = arg
            cudaver = str(cudaver)
            if (not cudaver == '7.0') and (not cudaver == '7.5') and (not cudaver == '8.0') and (not cudaver == '9.1') and (not cudaver == '10.0'):
                print("version, ",cudaver,", not available")
                print("setting cuda version to 9.1")
                gpu = '9.1'            
        elif opt in ("-o", "--highp"):
            highp = arg
            highp = str(highp)
            print("HIGHP= ",highp)
            if ( highp == 'NO') or (highp == 'No') or (highp == 'no') or (highp == 'nO') or (highp == 'N') or (highp == 'n'): 
                print("HIGHP= ",highp," Jupyter notebook or lab will not run on owned node(s)")
                highp = 'no'
            elif ( highp == 'Yes') or ( highp == 'yes') or ( highp == 'YES') or ( highp == 'Y') or ( highp == 'y'):
                print("HIGHP= ",highp," Jupyter notebook or lab will run on owned node(s)")
                highp='yes'
        elif opt in ("-x", "--exclusive"):
            exclusive = arg
            #print("EXCL= ",exclusive)
            exclusive = str(exclusive)
            print("EXCL= ",exclusive)
            if ( exclusive == 'NO') or (exclusive == 'No') or (exclusive == 'no') or (exclusive == 'nO') or (exclusive == 'N') or (exclusive == 'n'):
                print("Jupyter notebook or lab will not run in exclusive mode")
                exclusive = 'no'
            elif (exclusive == 'Yes') or (exclusive == 'yes') or (exclusive == 'YES') or (exclusive == 'Y') or (exclusive == 'y'):
                print("Jupyter notebook or lab will run in exclusive mode")
                exclusive='yes'
            else:
                print("EXCLUSIVE=",exclusive)
        elif opt in ("-a", "--arch"):
            arch = arg
            arch = str(arch)
            if (not arch == 'amd-23540') and (not arch == 'amd-2376') and (not arch == 'amd-2380') and (not arch == 'amd-2382') and (not arch == 'amd-2384') and (not arch == 'amd-2431') and (not arch == 'amd-2435') and (not arch == 'amd-4176') and (not arch == 'amd-6136') and  (not arch == 'amd-6276') and (not arch == 'i-E5-2670v3') and  (not arch == 'intel-E5-2650') and (not arch == 'intel-E5-2670') and (not arch == 'intel-E5-2697') and (not arch == 'intel-E5-4620') and (not arch == 'intel-E5530') and (not arch == 'intel-E5540') and (not arch == 'intel-E7-8890') and (not arch == 'intel-gold-51') and (not arch == 'intel-gold-61') and (not arch == 'intel-X5550') and (not arch == 'intel-X5560') and (not arch == 'intel-X5650') and (not arch == 'lx-amd64'):
                print("version, ",arch,", of CPU not available")
                print("no CPU-type will be requested")
                arch = ''
        else:
          print("Usage:\n")
          print("h2jupynb [-u <Hoffman2 user name>] [-t <time, integer number of hours>] [-m <memory, integer number of GB per core>] [-e <parallel environment: 1 for shared, 2 for distributed>] [-s <number of slots>] [-v <python-version>] [-o <run on group owned nodes: yes/no>] [-x <run on an exclusively reserved node: yes/no>]  [-a <CPU-type>] [-d <path to directory from which the jupyter NB/lab will start>] [-p <port number>] [-g <run on a gpu node: yes/no>] [-c <gpu-card-type>] [-l <cuda-version>]\n")
          print("If no arguments are given to this script it is assumed that:")
          print("\t your Hoffman2 user name is the same as on your client machine")
          print("\t the time duration for your session is of 2 hours")
          print("\t the parallel environment is shared")
          print("\t the memory per slot for your session is of 1GB")
          print("\t the number of slots for your session is of 1")
          print("\t the python version for your notebook is 2.7.13")
          print("\t the port on which the server is started is 8789")
          print("\t the starting directory on Hoffman2 is your $HOME")
          print("\t use GPU? default is no")
          print("\t GPU type default is P4 (if -g yes)")
          print("\t CUDA version 9.1 (if -g yes)")
          print("\t not running on owned nodes")
          print("\t not running in exclusive mode")
          print("\t no specific CPU selected (see ARCH in qhost output)")
          print("\n\t python versions currently available are: 2.7.13, 2.7.16, 3.6.1, 3.7.0, 3.7.2")
          print("\t or: anaconda2, anaconda3 and anaconda/python3-4.2")
          print("\t python versions: 2.7.3 or 3.4 are available but deprecated.\n")
          print("\t cuda versions currently available are: 7.0, 7.5, 8.0, 9.1 and 10.0 \n")
          print("\t GPU cards currently publicly available: P4, RTX2080Ti and V100\n")
          sys.exit(2)
            
    if not username:
        if opsys == 'Windows':
           username = str(os.environ["USERNAME"])
        else:
           username = str(os.environ["USER"])


    print("Your Hoffman2 user name is", username )
    print("The time in hours is", timeinhours )
    print("The memory in GB per slots is", memoryingb)
    if (numberofslots > 1):
      print("The parallel environment is", parenv)
      print("The number of slots is", numberofslots)
    print("The version of python for the notebook is", pythonver)
    if (usegpu == 'no') and (highp == 'yes'):
       print("Your Jupyter NB or lab will run on owned resources")
    elif (usegpu == 'yes') and (highp == 'no'):
       print("Your Jupyter NB or lab will run on a GPU node")
       print("Your Jupyter NB or lab will load cuda version ",cudaver)
    elif (usegpu == 'yes') and (highp == 'yes'):
       print("Your Jupyter NB or lab will run on a GPU owned node")  
       print("Your Jupyter NB or lab will load cuda version ",cudaver)
    print("The port is" , port)
    if (exclusive == 'yes'):
       print("Your Jupyter NB or lab will reserve a node exclusively ")
    if not directory:
        print("The directory on Hoffman2 is $HOME")
    else:
        print("The directory on Hoffman2 is", directory)
    if ( arch != '' ):
        print("Your Jupyter NB or lab will run on ",arch," CPU")
        arch=',arch='+arch
        print("arch=",arch)

if __name__ == "__main__":
   main(sys.argv[1:])

print("usegpu= ",usegpu," highp= ",highp," exclusive= ",exclusive, "gputype= ",gpu, "arch= ",arch, "time= ",timeinhours, "mem= ",memoryingb, "PE= ", parenv, "NSLOTS= ",numberofslots)


if (usegpu == 'no') and (highp == 'no') and (exclusive == 'no'):
#   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l i,h_rt=%d:00:00,h_data=%dG%s -pe dc\* %d")
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l i,h_rt=%d:00:00,h_data=%dG%s -pe %s %d")
elif (usegpu == 'no') and (highp == 'yes') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,highp%s -pe %s %d -w e")
elif (usegpu == 'no') and (highp == 'no') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,exclusive%s -pe %s %d -w e")
elif (usegpu == 'no') and (highp == 'yes') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,exclusive,highp%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'no') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'yes') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s,highp%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'yes') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s,highp,exclusive%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'no') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s,exclusive%s -pe %s %d -w e")
MODULE_LOADCUDA_TEMP = "module load cuda/%s"
MODULE_LOAD_TEMP = "module load %s"


if (pythonver == 'python/2.7.3'):
  STARTNBCMD = "ipython notebook"
else:
  STARTNBCMD = "jupyter notebook"


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

pqsub=Popen(['ssh','-t','-t','-4','%s@hoffman2.idre.ucla.edu' % username],stdin=PIPE,stdout=PIPE,stderr=PIPE,universal_newlines=True)
pqsub.stdin.flush()
pqsub.stdout.flush()


#
#if (usegpu == 'no') and (highp == 'no'): 
#   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,arch,parenv,numberofslots))+"\n")
#elif (usegpu == 'no') and (highp == 'yes'):
#   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,arch,parenv,numberofslots))+"\n")
#elif (usegpu == 'yes') and (highp == 'no'):
#   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,arch,parenv,numberofslots))+"\n")
#elif (usegpu == 'yes') and (highp == 'yes'):
#   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,arch,parenv,numberofslots))+"\n")
#
if (usegpu == 'no'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,arch,parenv,numberofslots))+"\n")
elif (usegpu == 'yes'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,arch,parenv,numberofslots))+"\n") 


# check for error on qrsh --- WE CAN GET error: no suitable queues b/c of -w e: 
#def getqrsherror(line):
#    global qrsh_error
#    if (line.find('error: no suitable queues')>0):
#       qrsh_error = 'yes'
#       return True
#
#readwhile(pqsub.stdout, getqrsherror)
##readwhile(pqsub.stdout, lambda line: line.find('error: no suitable queues')>0)
#if (qrsh_error == 'yes'):
#    print('The combination of resources requested is ivalid please try again...')
#    pqsub.kill()
#    #raise SystemExit
#else:


if (usegpu == 'yes'):
   pqsub.stdin.write(MODULE_LOADCUDA_TEMP % (cudaver) +"\n")
pqsub.stdin.write(MODULE_LOAD_TEMP % (pythonver) +"\n")
pqsub.stdin.write('module li\n')
pqsub.stdin.write('echo HOSTNAME=`hostname`\n')
pqsub.stdin.flush()

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

readwhile(pqsub.stdout, gethostname)

if ('login' in hostname):
   print("We are on the wrong host. Exiting...")
   quit()

if directory: 
    pqsub.stdin.write('cd %s\n'%directory)
    pqsub.stdin.write('echo CD\n')
    readwhile(pqsub.stdout, lambda line: line.startswith('CD'))
 
def gethttpcommand(line):
    global httpcommand
    if (line.find('localhost')>0):
       httpcommand = line.split(' ')[-1].strip()
       return True


pqsub.stdin.write(STARTNBCMD + ' --port=%s\n'%port)
pqsub.stdin.flush()

if (STARTNBCMD == 'jupyter notebook'):
  #readwhile(pqsub.stdout, lambda line: line.find('localhost')>0)
  #print("\n\tUse the token above to log into your notebook as needed\n")
  readwhile(pqsub.stdout, gethttpcommand)
  print("\n\t", httpcommand, "\n")
else:
  readwhile(pqsub.stdout, lambda line: line.find('NotebookApp')>0)

if myssh == 'ssh':
    tunnel = ['ssh','-4', '-t', '-Y', '%s@hoffman2.idre.ucla.edu' % username, '-L', '%s:localhost:%s'%(port,port), 'ssh', '-t', '-Y', hostname, '-L', '%s:localhost:%s'%(port,port)]
elif myssh == 'plink':
    tunnel = ['plink','-4', '-t', '-X', '%s@hoffman2.idre.ucla.edu' % username, '-L', '%s:localhost:%s'%(port,port), 'ssh', '-t', '-X', hostname, '-L', '%s:localhost:%s'%(port,port)]

print(' '.join(tunnel))

ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE,universal_newlines=True)
ptunnel.stdin.write('echo TUNNEL\n')
ptunnel.stdin.flush()


readwhile(ptunnel.stdout,lambda line: line.startswith('TUNNEL'))

if ('CYGWIN' in opsys) or ('Cygwin' in opsys) or ('cygwin' in opsys):
   if (STARTNBCMD == 'jupyter notebook'):
     call(["cygstart",httpcommand])
   else:
     call(["cygstart",'http://localhost:%s'%(port)])
else:
   if (STARTNBCMD == 'jupyter notebook'):
     webbrowser.open(httpcommand)
   else:
     webbrowser.open('http://localhost:%s'%(port))

print("Succesfully opened notebook!")
print("Kill this process to end your notebook connection.")

time.sleep(timeinhours*3600)

pqsub.kill()
ptunnel.kill()


print("Succesfully cleaned up connections.")