qmcpack/labs/misc/jobrun_vesta

#! /usr/bin/env python

import os
import sys

#=======#
# usage # 
#=======#
#  jobrun_vesta exe_name exe_inputs/options


#=============#
# global data #
#=============#
# the account to charge time to 
#account = 'QMCPACK-training' # participant account!
account = 'QMCPACK'  # instructor account!

# executables jobrun_vesta knows how to handle
exe_name_list = 'rungms convert4qmc pw.x pw2qmcpack.x qmcpack qmcpack_comp'.split()
exe_names = set(exe_name_list)

# protect user against errors and print pedagogical error messages
guard = True
#guard = False

# examples displayed to user if no arguments are provided
examples = '''
jobrun_vesta rungms H2O.PBE

jobrun_vesta convert4qmc -gamessAscii H2O.CISD.orbs.output -ci H2O.CISD.orbs.output -threshold 0.01 -readInitialGuess 60  -add3BodyJ

jobrun_vesta pw.x scf.in

jobrun_vesta pw2qmcpack.x p2q.in

jobrun_vesta qmcpack opt.in.xml
'''

#==================#
# global functions #
#==================#
def error(msg):
    print 'jobrun error:\n  '+msg.replace('\n','\n  ')
    print 'exiting'
    sys.exit()
#end def error

def join_args(args):
    s = ''
    for a in args:
        s+=' '+a
    #end for
    return s.lstrip()
#end def join_args


#==================#
# actual execution #
#==================#
# get the tokens provided by the user
args = sys.argv[1:]

# require at least one argument, the exe name
if len(args)==0:
    error('no arguments provided to jobrun_vesta\nfirst argument to jobrun_vesta must be the exe name\nvalid options are: {0}\nan example for each executable is shown below\n{1}'.format(exe_name_list,examples))
#end if

# get the exe name and any arguments for the exe itself
exe_name = args[0]
exe_arg_list = args[1:]

# exe must be known
if exe_name not in exe_names:
    error('invalid exe name: {0}\nvalid options are: {1}'.format(exe_name,exe_name_list))
#end if

# generate submission command based on target executable
command = ''
if exe_name=='rungms':
    #========#
    # GAMESS #
    #========#
    if len(exe_arg_list)!=1:
        error('must provide a single argument to rungms, the input file prefix\narguments provided: {0}'.format(exe_arg_list))
    #end if
    infile_prefix = exe_arg_list[0]
    if guard:
        if infile_prefix.endswith('.inp'):
            infile = infile_prefix
            infile_prefix = infile_prefix.rsplit('.',1)[0]
            error('rungms expects the input file prefix, not the file itself\nfile provided: {0}\nprefix expected: {1}'.format(infile,infile_prefix))
        else:
            infile = infile_prefix+'.inp'
        #end if
        if not os.path.exists(infile):
            error('gamess input file does not exist: {0}'.format(infile))
        #end if
    #end if
    gms   = '/soft/applications/qmcpack/Binaries/gmsbgq'
    nodes = 32
    time  = 20
    mode  = 'c1'
    # rungms prefix nodes time mode account prefix
    command = '{0} {1} {2} {3} {4} {5} {1}'.format(gms,infile_prefix,nodes,time,mode,account,infile_prefix)
elif exe_name=='convert4qmc':
    #=============#
    # convert4qmc #
    #=============#
    convert  = '/soft/applications/qmcpack/Binaries/convert4qmc'
    nodes    = 1
    threads  = 1
    time     = 10
    queue    = 'default'
    output   = 'convert'
    mode     = 'c1'
    # qsub account queue nodes time output mode threads convert4qmc user_args
    command = 'qsub -A {0} -q {1} -n {2} -t {3} -O {4} --mode {5} --env BG_SHAREDMEMSIZE=32:OMP_NUM_THREADS={6} {7} {8}'.format(account,queue,nodes,time,output,mode,threads,convert,join_args(exe_arg_list))

elif exe_name=='pw.x':
    #======#
    # pw.x #
    #======#
    if len(exe_arg_list)!=1:
        error('must provide a single argument to pw.x, the input file\narguments provided: {0}\n(pw.x can accept other inputs outside the training environment)'.format(exe_arg_list))
    #end if
    infile = exe_arg_list[0]
    if guard:
        if not os.path.exists(infile):
            error('pw.x input file does not exist: {0}'.format(infile))
        #end if
    #end if
    if '.' in infile:
        prefix = infile.split('.')[0]
    else:
        prefix = infile
    #end if
    pw    = '/soft/applications/qmcpack/Binaries/pw.x'
    nodes = 9
    time  = 60
    queue = 'prod'
    script_text = '''#!/bin/sh

RANKS_PER_NODE=4
NODES=`cat $COBALT_NODEFILE | wc -l`
PROCS=$(($NODES * $RANKS_PER_NODE))

mpirun -f $COBALT_NODEFILE -n $PROCS {0} -npool 4 -inp {1}
    '''.format(pw,infile)
    script_file = prefix+'.submit.sh'
    open(script_file,'w').write(script_text)
    # qsub account nodes time submission_script
    command = 'qsub -A {0} -n {1} -t {2} {3}'.format(account,nodes,time,script_file)
elif exe_name=='pw2qmcpack.x':
    #==============#
    # pw2qmcpack.x #
    #==============#
    if len(exe_arg_list)!=1:
        error('must provide a single argument to pw2qmcpack.x, the input file\narguments provided: {0}\n(pw2qmcpack.x can accept other inputs outside the training environment)'.format(exe_arg_list))
    #end if
    infile = exe_arg_list[0]
    if guard:
        if not os.path.exists(infile):
            error('pw2qmcpack.x input file does not exist: {0}'.format(infile))
        #end if
    #end if
    if '.' in infile:
        prefix = infile.split('.')[0]
    else:
        prefix = infile
    #end if
    pw    = '/soft/applications/qmcpack/Binaries/pw2qmcpack.x'
    nodes = 9
    time  = 60
    queue = 'prod'
    script_text = '''#!/bin/sh

RANKS_PER_NODE=1
NODES=`cat $COBALT_NODEFILE | wc -l`
PROCS=$(($NODES * $RANKS_PER_NODE))

mpirun -f $COBALT_NODEFILE -n $PROCS {0} -npool 1 -inp {1}
    '''.format(pw,infile)
    script_file = prefix+'.submit.sh'
    open(script_file,'w').write(script_text)
    # qsub account nodes time submission_script
    command = 'qsub -A {0} -n {1} -t {2} {3}'.format(account,nodes,time,script_file)
elif exe_name=='qmcpack' or exe_name=='qmcpack_comp':
    #=========#
    # QMCPACK #
    #=========#
    if len(exe_arg_list)!=1:
        error('must provide a single argument to qmcpack, the input file\narguments provided: {0}\n(qmcpack can accept other inputs outside the training environment)'.format(exe_arg_list))
    #end if
    infile = exe_arg_list[0]
    if guard:
        if not os.path.exists(infile):
            error('qmcpack input file does not exist: {0}'.format(infile))
        #end if
    #end if
    if '.' in infile:
        output = infile.split('.')[0]
    else:
        output = infile
    #end if
    qmcpack = '/soft/applications/qmcpack/Binaries/'+exe_name
    nodes   = 512
    threads = 64
    time    = 60
    queue   = 'default'
    mode    = 'c1'
    command = 'qsub -A {0} -q {1} -n {2} -t {3} -O {4} --mode {5} --env BG_SHAREDMEMSIZE=32:OMP_NUM_THREADS={6} {7} {8}'.format(account,queue,nodes,time,output,mode,threads,qmcpack,infile)
#end if

# execute the submission command
#print command
os.system(command)