1#! /usr/bin/env python3
2
3import os
4import sys
5
6#=======#
7# usage #
8#=======#
9#  jobrun_vesta exe_name exe_inputs/options
10
11
12#=============#
13# global data #
14#=============#
15# the account to charge time to
16account = 'QMCPACK-training' # participant account!
17#account = 'QMCPACK'  # instructor account!
18
19# executables jobrun_vesta knows how to handle
20exe_name_list = 'ppconvert rungms convert4qmc pw.x pw2qmcpack.x qmcpack qmcpack_comp'.split()
21exe_names = set(exe_name_list)
22
23# protect user against errors and print pedagogical error messages
24guard = True
25#guard = False
26
27# examples displayed to user if no arguments are provided
28examples = '''
29jobrun_vesta ppconvert --gamess_pot O.BFD.gamess --s_ref "1s(2)2p(4)" --p_ref "1s(2)2p(4)" --d_ref "1s(2)2p(4)" --log_grid --upf O.BFD.upf
30
31jobrun_vesta rungms H2O.PBE
32
33jobrun_vesta convert4qmc -gamess H2O.CISD.orbs.output -ci H2O.CISD.orbs.output -threshold 0.01 -readInitialGuess 60  -add3BodyJ
34
35jobrun_vesta pw.x scf.in
36
37jobrun_vesta pw2qmcpack.x p2q.in
38
39jobrun_vesta qmcpack opt.in.xml
40'''
41
42# portions of template qsub script
43cobalt_header = '''#!/bin/bash
44#COBALT -q qmcpack
45#COBALT -A {0}
46#COBALT -n {1}
47#COBALT -t {2}
48#COBALT -O {3}
49'''
50#cobalt_header = '''#!/bin/bash
51##COBALT -q default
52##COBALT -A {0}
53##COBALT -n {1}
54##COBALT -t {2}
55##COBALT -O {3}
56#'''
57
58
59cobalt_shape = '''LOCARGS="--block $COBALT_PARTNAME ${COBALT_CORNER:+--corner} $COBALT_CORNER ${COBALT_SHAPE:+--shape} $COBALT_SHAPE"
60echo "Cobalt location args: $LOCARGS" >&2
61'''
62runjob='runjob --np {0} -p {1} $LOCARGS --verbose=INFO --envs OMP_NUM_THREADS={2} : {3}{4}'
63
64
65#==================#
66# global functions #
67#==================#
68def error(msg):
69    print 'jobrun error:\n  '+msg.replace('\n','\n  ')
70    print 'exiting'
71    sys.exit()
72#end def error
73
74def join_args(args):
75    s = ''
76    for a in args:
77        s+=' '+a
78    #end for
79    return s.lstrip()
80#end def join_args
81
82
83#==================#
84# actual execution #
85#==================#
86# get the tokens provided by the user
87args = sys.argv[1:]
88
89# require at least one argument, the exe name
90if len(args)==0:
91    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))
92#end if
93
94# get the exe name and any arguments for the exe itself
95exe_name = args[0]
96exe_arg_list = args[1:]
97
98# exe must be known
99if exe_name not in exe_names:
100    error('invalid exe name: {0}\nvalid options are: {1}'.format(exe_name,exe_name_list))
101#end if
102
103# generate submission command based on target executable
104script = False
105runstr = None
106command = ''
107if exe_name=='ppconvert':
108    #===========#
109    # ppconvert #
110    #===========#
111    for n in range(len(exe_arg_list)-1): # put " back on arguments
112        arg1 = exe_arg_list[n]
113        arg2 = exe_arg_list[n+1]
114        if arg1.startswith('--') and arg1.endswith('ref'):
115            if not arg2.startswith('"'):
116                arg2 = '"'+arg2
117            #end if
118            if not arg2.endswith('"'):
119                arg2 = arg2+'"'
120            #end if
121            exe_arg_list[n+1] = arg2
122        #end if
123    #end for
124    script  = True
125    #exe     = '/soft/applications/qmcpack/Binaries/ppconvert'
126    exe     = '/home/yeluo/qmcdev/trunk/src/QMCTools/ppconvert/build/ppconvert'
127    prefix  = 'ppconv'
128    nodes   = 1
129    procs   = 1
130    threads = 1
131    time    = 60
132    argstr  = ' '+join_args(exe_arg_list)+'\n'
133    runstr  = exe+argstr
134
135elif exe_name=='rungms':
136    #========#
137    # GAMESS #
138    #========#
139    if len(exe_arg_list)!=1:
140        error('must provide a single argument to rungms, the input file prefix\narguments provided: {0}'.format(exe_arg_list))
141    #end if
142    infile_prefix = exe_arg_list[0]
143    if guard:
144        if infile_prefix.endswith('.inp'):
145            infile = infile_prefix
146            infile_prefix = infile_prefix.rsplit('.',1)[0]
147            error('rungms expects the input file prefix, not the file itself\nfile provided: {0}\nprefix expected: {1}'.format(infile,infile_prefix))
148        else:
149            infile = infile_prefix+'.inp'
150        #end if
151        if not os.path.exists(infile):
152            error('gamess input file does not exist: {0}'.format(infile))
153        #end if
154    #end if
155    #gms   = '/soft/applications/qmcpack/Binaries/gmsbgq'
156    gms   = '/home/yeluo/qmcdev/DFT_Binaries/gmsbgq.old' # 1 May 2012 (R2)
157    nodes = 32
158    time  = 20
159    mode  = 'c1'
160    # rungms prefix nodes time mode account prefix
161    command = '{0} {1} {2} {3} {4} {5} {1}'.format(gms,infile_prefix,nodes,time,mode,account,infile_prefix)
162elif exe_name=='convert4qmc':
163    #=============#
164    # convert4qmc #
165    #=============#
166    prefix   = 'c4q'
167    for n in range(len(exe_arg_list)):
168        if exe_arg_list[n]=='-prefix' and n+1<len(exe_arg_list):
169            prefix = exe_arg_list[n+1]
170        #end if
171    #end for
172    script  = True
173    exe     = '/soft/applications/qmcpack/Binaries/convert4qmc'
174    nodes   = 1
175    procs   = 1
176    threads = 1
177    time    = 10
178    argstr  = ' '+join_args(exe_arg_list)+'\n'
179
180elif exe_name=='pw.x':
181    #======#
182    # pw.x #
183    #======#
184    if len(exe_arg_list)!=1:
185        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))
186    #end if
187    infile = exe_arg_list[0]
188    if guard:
189        if not os.path.exists(infile):
190            error('pw.x input file does not exist: {0}'.format(infile))
191        #end if
192    #end if
193    if '.' in infile:
194        prefix = infile.rsplit('.',1)[0]
195    else:
196        prefix = infile
197    #end if
198    script  = True
199    exe     = '/soft/applications/qmcpack/Binaries/pw.x'
200    nodes   = 32
201    procs   = 16
202    threads = 1
203    time    = 60
204    argstr  =  ' -input '''+infile+'\n'
205
206elif exe_name=='pw2qmcpack.x':
207    #==============#
208    # pw2qmcpack.x #
209    #==============#
210    if len(exe_arg_list)!=1:
211        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))
212    #end if
213    infile = exe_arg_list[0]
214    if guard:
215        if not os.path.exists(infile):
216            error('pw2qmcpack.x input file does not exist: {0}'.format(infile))
217        #end if
218    #end if
219    if '.' in infile:
220        prefix = infile.rsplit('.',1)[0]
221    else:
222        prefix = infile
223    #end if
224    script  = True
225    exe     = '/soft/applications/qmcpack/Binaries/pw2qmcpack.x'
226    nodes   = 1
227    procs   = 1
228    threads = 1
229    time    = 60
230    argstr  = '<'+infile+'\n'
231
232elif exe_name=='qmcpack' or exe_name=='qmcpack_comp':
233    #=========#
234    # QMCPACK #
235    #=========#
236    if len(exe_arg_list)!=1:
237        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))
238    #end if
239    infile = exe_arg_list[0]
240    if guard:
241        if not os.path.exists(infile):
242            error('qmcpack input file does not exist: {0}'.format(infile))
243        #end if
244    #end if
245    if '.' in infile:
246        if infile.endswith('.in.xml'):
247            prefix = infile.rsplit('.',2)[0]
248        else:
249            prefix = infile.rsplit('.',1)[0]
250        #end if
251    else:
252        prefix = infile
253    #end if
254    script  = True
255    exe     = '/soft/applications/qmcpack/Binaries/'+exe_name
256    nodes   = 32
257    procs   = 1
258    threads = 16
259    #threads = 64
260    time    = 60
261    argstr  = ' '+infile+'\n'
262#end if
263
264if script:
265    script_text = cobalt_header.format(account,nodes,time,prefix)
266    script_text += cobalt_shape
267    if runstr is None:
268        script_text += runjob.format(nodes*procs,procs,threads,exe,argstr)
269    else:
270        script_text += runstr+'\n'
271    #end if
272    script_file = prefix+'.qsub.in'
273    open(script_file,'w').write(script_text)
274    os.system('chmod +x '+script_file)
275    command = 'qsub --mode script --env BG_SHAREDMEMSIZE=32 {0}'.format(script_file)
276#end if
277
278# execute the submission command
279print command
280os.system(command)
281