rokroskar/sparkhpc

#!/bin/env python
#
#
# CLI for starting and running Spark standalone clusters on HPC resources
#
#

from __future__ import print_function
import click
import sparkhpc
from sparkhpc import sparkjob
import subprocess
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('sparkhpc')

home = os.path.expanduser('~')

SCHEDULER = sparkjob.get_scheduler()

@click.group()
def cli():
    pass

@cli.command()
@click.argument('ncores', type=int)
@click.option('--walltime', default="00:30", help="Walltime in HH:MM format")
@click.option('--jobname', default='sparkcluster', help='Name to use for the job')
@click.option('--template', default=None, help='Job template path')
@click.option('--memory-per-executor', default=2000, envvar='SPARK_EXECUTOR_MEMORY',
              help='Memory to reserve for each executor (i.e. the JVM) in MB')
@click.option('--memory-per-core', default=2000,
              help='Memory per core to request from scheduler in MB')
@click.option('--cores-per-executor', default=1,
              help='Cores per executor')
@click.option('--spark-home', default=os.path.join(home,'spark'), envvar='SPARK_HOME', 
              help='Location of the Spark distribution')
@click.option('--wait', default=False, is_flag=True, help='Wait until the job starts')
def start(ncores, 
          walltime, 
          jobname, 
          template, 
          memory_per_executor, 
          memory_per_core, 
          cores_per_executor,
          spark_home, 
          wait):
    """Start the spark cluster as a batch job"""
    
    sj = sparkjob.sparkjob(ncores=ncores, 
                           walltime=walltime, 
                           jobname=jobname, 
                           template=template, 
                           memory_per_core=memory_per_core,
                           memory_per_executor=memory_per_executor, 
                           cores_per_executor=cores_per_executor,
                           spark_home=spark_home)
    
    if wait: 
        logger.info(' Waiting for job to start - ctrl-c to stop')
        sj.wait_to_start()
    else:
        sj.submit()
    

@cli.command()
def info():
    """Get info about currently running clusters"""
    sparkhpc.show_clusters()


@cli.command()
@click.argument('clusterid')
def stop(clusterid):
    """Kill a currently running cluster ('all' to kill all clusters)"""
    sjs = sparkjob.sparkjob().current_clusters()

    if clusterid == 'all': 
        if len(sjs) == 0: 
            logger.info(' No clusters running')
        for sj in sjs: 
            sj.stop()
    elif int(clusterid) < len(sjs): 
        sjs[int(clusterid)].stop()
    else: 
        raise RuntimeError('Cluster %s does not exist'%clusterid)


@cli.command()
@click.option('--memory', default='2000M', help='Memory for each executor using a Java memory string')
@click.option('--timeout', default=30, help='Timeout for starting spark master')
@click.option('--cores-per-executor', default=1, help='Number of cores per executor')
def launch(memory, timeout, cores_per_executor):
    """Launch the Spark master and workers within a current job context"""
    sparkjob.start_cluster(memory, timeout=timeout, cores_per_executor=cores_per_executor)

if __name__ == "__main__":
    cli()