rokroskar/sparkhpc

#!/usr/bin/env python
from __future__ import print_function
import os
from os.path import expanduser, exists
import subprocess
import sys
import re
import socket
from IPython.terminal.ipapp import launch_new_instance
import shutil
import click
import pkg_resources
 
import logging
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('hpcnotebook')
 
# try to figure out which scheduler we have
def which(program):
    import os
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)
 
    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file
 
    return None
 
if which('bjobs') is not None: 
    SCHEDULER = 'LSF'
    HOST_ENVIRON = 'LSB_HOSTS'
elif which('squeue') is not None: 
    SCHEDULER = 'SLURM'
    HOST_ENVIRON = 'SLURM_NODELIST'
else: 
    # no scheduler
    HOST_ENVIRON = 'NO_SCHEDULER_HERE'
 
#
# Initial inspiration for this script from https://github.com/felixcheung/vagrant-projects
#
 
notebook_config_template = """c = get_config()
c.NotebookApp.ip = '*'
c.NotebookApp.open_browser = False
c.NotebookApp.password = "{password}"
c.NotebookApp.certfile = "{certfile}"
c.NotebookApp.port = {port}
"""
 
templates = {'LSF': 'hpcnotebook.lsf.template'}
submit_command = {'LSF': 'bsub < %s'}
job_regex = {'LSF': 'Job <(\d+)>'}
 
class bc:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
 
# get home directory
home = expanduser("~")
 
# setup the path to the jupyter notebook configuration
jupyter_config_path = "{home}/.jupyter_notebook".format(home=home)
 
@click.group()
@click.option('--port', default=8889, help='Port for the notebook server')
@click.pass_context
def cli(ctx, port): 
    ctx.obj['port'] = port
 
@cli.command()
@click.option('--force', is_flag=True, help='Force a reset of the notebook, even if it has already been done')
@click.pass_context
def setup(ctx, force):
    """Setup the notebook"""
 
    port = ctx.obj['port']
 
    if force and exists(jupyter_config_path):
        shutil.rmtree(jupyter_config_path)
 
    # if the profile doesn't exist, create it -- otherwise we've probably
    # already done this step
    if not exists(jupyter_config_path):
        # get a password
        from notebook.auth import passwd
        
        logger.info("""
            This script will create a Jupyter notebook profile for working remotely.
            When it is finished, you can find the configuration in {conf_path}.
            First, please enter a *new* password for your Jupyter notebook:
            """.format(conf_path=bc.UNDERLINE + jupyter_config_path + bc.ENDC))
        
        new_pass = passwd()
 
        logger.info("Creating an SSL certificate to enable a secure connection; the certificate will be in your ~/.ssh directory")
        
        # make sure the .ssh directory is there before continuing
        sshdir = '{home}/.ssh'.format(home=home)
        if not os.path.exists(sshdir):
            os.makedirs(sshdir)
            os.fchmod(sshdir, 700)                        
 
        # make an ssl certificate
        certfile = "{home}/.ssh/notebook_cert.pem".format(home=home)
 
        out = subprocess.check_output(['openssl', 'req', '-x509', '-nodes', '-days', '365', '-newkey', 'rsa:1024', '-keyout', '%s'%certfile, '-out', '%s'%certfile, '-batch'], stderr=subprocess.STDOUT).decode()
 
        lines = out.split('\n')
        for l in lines : 
            logger.info(bc.OKGREEN + '[openssl] ' + bc.ENDC + l)
        
        # write the notebook config
 
        # create the directory
        os.makedirs(jupyter_config_path)
        with open("{profile_path}/jupyter_notebook_config.py".format(profile_path=jupyter_config_path), 'w') as f:
            f.write(notebook_config_template.format(
                password=new_pass, certfile=certfile, port=port))
        
        logger.info(bc.BOLD + 'Notebook setup complete' + bc.ENDC)
            
    else:
        logger.error(bc.FAIL + "The jupyter notebook already looks set up; if you want to force setup, use --force".format(dir=jupyter_config_path) + bc.ENDC)
 
@cli.command()
@click.pass_context
def launch(ctx):
    """Launch the notebook"""
    port = ctx.obj['port']
 
    argv = sys.argv[:1]
    argv.append('notebook')
 
    argv.append('--config={profile_path}/jupyter_notebook_config.py'.format(profile_path=jupyter_config_path))
 
    # check if we passed in a port that is different from the one in the
    # configuration
    with open("{profile_path}/jupyter_notebook_config.py".format(profile_path=jupyter_config_path), 'r') as conf:
        conf_port = int(re.findall('port = (\d+)', conf.read())[0])
 
    if conf_port != port:
        logger.warning(bc.WARNING + "Overriding the port found in the existing configuration" + bc.ENDC)
        argv.append('--port={port}'.format(port=port))
 
    # determine if we're running on a compute node
    if HOST_ENVIRON in os.environ:
        compute = True
    else:
        compute = False
 
    sys.argv = argv
 
    if compute:
        ip = socket.gethostbyname(socket.gethostname())
    else:
        ip = 'localhost'
 
    logger.info(bc.BOLD + "To access the notebook, inspect the output below for the port number, then point your browser to https://{ip}:<port_number>".format(ip=ip) + bc.ENDC)
    sys.stdout.flush()
    launch_new_instance()
 
@cli.command()
@click.option('--ncores', type=int, default=1, help='Requested number of cores')
@click.option('--walltime', default='01:00', help='Requested runtime in HH:MM')
@click.option('--memory', type=int, default=2000, help='Requested memory in MB')
@click.option('--template', default=None, help='Path to template to use for job submission')
@click.option('--jobname', default='hpcnotebook', help='Name for the batch job')
@click.pass_context
def submit(ctx, ncores, walltime, memory, template, jobname):
    """Submit a notebook job to the scheduler"""
    if SCHEDULER != 'LSF':
        raise RuntimeError('only the LSF scheduler is supported at the moment')
 
    if template is None: 
        template_file = templates[scheduler]
        template_str = pkg_resources.resource_string('sparkhpc', 'templates/%s'%template_file)
    else : 
        with open(self.template, 'r') as template_file: 
            template_str = template_file.read()
    
    job = template_str.format(walltime=walltime, 
                              ncores=ncores,
                              memory=memory,
                              jobname=jobname)
 
    with open('notebook_job', 'w') as jobfile: 
            jobfile.write(job)
 
    job_submit = subprocess.Popen(submit_command[scheduler]%'notebook_job', shell=True, 
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
    try: 
        jobid = re.findall(job_regex[scheduler], job_submit.stdout.read().decode())[0]
    except Exception as e: 
        logger.error('Job submission failed or jobid invalid')
        raise e
    return jobid
 
if __name__ == "__main__":
    cli(obj={})