#!/usr/bin/env python

import os
import sys
import getopt
import time
import shutil

verbose = False
script_file = 'spark.pbs'
spark_home = None
deploy_dir = None
num_nodes = '###'
account = 'xxx####'
walltime = '00:00:00'
deploy_timeout = '30'

# Default cluster settings (for Rhea)
worker_memory = '100g'
worker_cores = '30'
driver_memory = '64g'
executor_memory = '4g'

def printUsage():
    print 'Usage spark_setup.py [options]'
    print 'Options:'
    print '  -h,-?,--help           : Show help'
    print '  -s,--script <file>     : Generate PBS script <file>'
    print '  -a,--account <acc>     : Name of account to charge'
    print '  -n,--nodes <num>       : Number of nodes'
    print '  -w,--wall <time>       : Maximum walltime'
    print '  -d,--deploy-dir <path> : Spark deployment directory *'
    print '  -t,--timeout <sec>     : Set deployment timeout (default = 30 sec)'
    print '  --spark-home <path>    : Override/set SPARK_HOME env variable **\n'
    print '  --worker-mem <arg>     : Worker memory\n'
    print '  --worker-cores <arg>   : Worker cores\n'
    print '  --driver-mem <arg>     : Driver memory\n'
    print '  --exec-mem <arg>       : Executor memory\n'
    print '*  The deployment directory must be accessible from login AND compute'
    print '   nodes, and will also be used for scratch space.\n'
    print '** Overriding SPARK_HOME is not recommended, use with caution\n\n'
    print 'Note: It is not necessary to re-run spark_setup for subsequent spark'
    print 'deployments unless the SPARK_HOME and/or deployment directories are changed.'
    print 'The generated PBS script may be edited directly to change or add PBS'
    print 'options.'

def printOptions():
    print 'PBS script     :', script_file
    print 'Account        :', account
    print 'Num nodes      :', num_nodes
    print 'Wall time      :', walltime
    print 'Deployment dir :', deploy_dir
    print 'SPARK_HOME     :', spark_home
    print 'Worker memory  :', worker_memory
    print 'Worker cores   :', worker_cores
    print 'Driver memory  :', driver_memory
    print 'Executor memory:', executor_memory

def generateScript():
    outf = open( script_file, 'w' )
    outf.write( "#!/bin/bash\n\n" )
    outf.write( '#PBS -A ' + account + '\n' )
    outf.write( '#PBS -l walltime=' + walltime + ',nodes=' + num_nodes + '\n\n'  )

    outf.write( 'export SPARK_HOME=' + spark_home + '\n' )
    outf.write( 'export SPARK_DEPLOY_DIR=' + deploy_dir + '\n\n' )

    outf.write( 'source ${MODULESHOME}/init/bash\n' )
    outf.write( 'module load python\n' )
    outf.write( 'export PYTHONPATH=$SPARK_HOME/python:$PYTHONPATH\n\n' )

    outf.write( 'if [ -d "$SPARK_DEPLOY_DIR/nodes" ]; then\n' )
    outf.write( '  rm -rf $SPARK_DEPLOY_DIR/nodes/*\n' )
    outf.write( '  sync\n' )
    outf.write( 'fi\n\n' )

    outf.write( "echo =============== Running spark deployment ===============\n\n" )

    outf.write( 'mpirun -n ' + num_nodes + ' --npernode 1 $SPARK_HOME/sbin/spark_deploy.py $SPARK_HOME $SPARK_DEPLOY_DIR &\n\n' )
    outf.write( "MASTERFILE=$SPARK_DEPLOY_DIR/nodes/master\n" )
    outf.write( "timeout=" + deploy_timeout + "\n" )
    outf.write( "while [ ! -f $MASTERFILE ]\n" )
    outf.write( "do\n" )
    outf.write( "    sleep 1\n" )
    outf.write( "    let timeout=timeout-1\n" )
    outf.write( "    if [ $timeout -eq 0 ]; then\n" )
    outf.write( "       echo >&2 Spark deployment failed. Inspect per-node logfiles in deployment directory for diagnostics.\n" )
    outf.write( "       exit 1\n" )
    outf.write( "    fi\n" )
    outf.write( "done\n\n" )

    outf.write( "MASTERNODE=$(<$MASTERFILE)\n" )
    outf.write( "echo Master: $MASTERNODE\n\n" )

    outf.write( 'SPARK_SUBMIT="$SPARK_HOME/bin/spark-submit --driver-memory '+driver_memory+' --executor-memory '+executor_memory+' --master spark://$MASTERNODE:7077"\n\n' )

    outf.write( "echo =============== Running spark job ===============\n\n" )

    outf.write( "# Edit/add the following line to with your Spark program and arguments specified...\n" )
    outf.write( "# $SPARK_SUBMIT <application> [arguments...]\n\n")

    outf.write( "echo =============== Spark job finished ===============\n\n" )
    outf.close()

def setupDeployDir():
    try:
        if not os.path.exists( deploy_dir ):
            os.makedirs( deploy_dir )
        if not os.path.exists( deploy_dir + '/templates' ):
            os.makedirs( deploy_dir + '/templates')
        if not os.path.exists( deploy_dir + '/nodes' ):
            os.makedirs( deploy_dir + '/nodes')

        shutil.copyfile( spark_home + "/conf/spark-env.sh.template",
                         deploy_dir + "/templates/spark-env.sh.template" )

        # Append default (static) settings to spark-env.sh template
        conf = open( deploy_dir + "/templates/spark-env.sh.template", "a" )

        conf.write( '\n# Default (static) settings generated by spark_setup:\n' )

        conf.write( 'SPARK_HOME='+ spark_home + '\n' )
        conf.write( 'SPARK_DRIVER_MEMORY='+driver_memory+'\n' )
        conf.write( 'SPARK_WORKER_INSTANCES=1\n' )
        conf.write( 'SPARK_WORKER_MEMORY='+worker_memory+'\n' )
        conf.write( 'SPARK_WORKER_CORES='+worker_cores+'\n' )

        conf.write( '\n# Support for modules:\n' )
        conf.write( 'source /etc/profile\n' )
        conf.write( 'source $MODULESHOME/init/bash\n' )
        conf.close()

        # Append default (static) settings to spark-defaults.conf template
        shutil.copyfile( spark_home + "/conf/spark-defaults.conf.template",
                         deploy_dir + "/templates/spark-defaults.conf.template" )

        conf = open( deploy_dir + "/templates/spark-defaults.conf.template", "a" )
        conf.write( '\n# Default (static) settings generated by spark_setup:\n' )
        conf.write( 'spark.driver.memory  '+driver_memory+'\n' )
        conf.write( 'spark.driver.maxResultSize  0\n' )
        conf.write( 'spark.executor.memory  '+executor_memory+'\n' )
        #conf.write( 'spark.shuffle.consolidateFiles  true\n' )
        conf.write( '\n' )
        conf.close()

        shutil.copyfile( spark_home + "/conf/log4j.properties.template",
                         deploy_dir + "/templates/log4j.properties.template" )

    except Exception as e:
        print 'Failed creating deployment directories.'
        print 'Please ensure the deployment path is valid and accessible.\n'
        print e
        sys.exit(1)

try:
    opts, args = getopt.getopt( sys.argv[1:],
        '?hvs:d:a:n:w:',
        ['verbose','help','script=','deploy-dir=','account=','wall=','nodes=','spark-home='] )

    if len(args) > 0:
        print 'Unrecognized options: ', args
        printUsage()
        sys.exit(1)

except getopt.GetoptError as e:
    print e
    printUsage()
    sys.exit(1)

for opt, arg in opts:
    if opt in ( '-?', '-h', '--help' ):
        printUsage()
        sys.exit()
    elif opt in ( '-s', '--script' ):
        script_file = arg
    elif opt in ( '-d', '--deploy-dir' ):
        deploy_dir = arg
    elif opt in ( '-n', '--nodes' ):
        num_nodes = arg
    elif opt in ( '-a', '--account' ):
        account = arg
    elif opt in ( '-w', '--wall' ):
        walltime = arg
    elif opt in ( '--spark-home' ):
        spark_home = arg
    elif opt in ( '-t', '--timeout' ):
        deploy_timeout = arg
    elif opt in ( '-v', '--verbose' ):
        verbose = True
    elif opt in ( '--worker-mem' ):
        worker_memory = arg
    elif opt in ( '--worker-cores' ):
        worker_cores = arg
    elif opt in ( '--driver-mem' ):
        driver_memory = arg
    elif opt in ( '--exec-mem' ):
        executor_memory = arg

if spark_home == None:
    if "SPARK_HOME" in os.environ:
        spark_home = os.environ["SPARK_HOME"]
    else:
        print 'SPARK_HOME must either exported or specified with --spark-home option'
        sys.exit(1)

if verbose == True:
    printOptions()

if script_file != None:
    generateScript()

if deploy_dir != None:
    setupDeployDir()