6 # Run event generation and produce an EVNT file.
9 import os, sys, time, shutil
13 from PyJobTransforms.trfArgs
import addAthenaArguments
14 from PyJobTransforms.trfDecorators
import stdTrfExceptionHandler, sigUsrStackTrace
15 from EvgenJobTransforms.evgenTrfArgs
import addStdEvgenArgs
19 ListOfDefaultPositionalKeys=[
'--AMIConfig',
'--AMITag',
'--argJSON',
'--asetup',
'--athena',
20 '--athenaMPMergeTargetSize',
'--athenaopts',
'--attempt',
'--checkEventCount',
'--command',
21 '--dumpJSON',
'--dumpPickle',
'--ecmEnergy',
'--env',
'--eventAcceptanceEfficiency',
22 '--evgenJobOpts',
'--execOnly',
'--fileValidation',
'--firstEvent',
'--ignoreErrors',
23 '--ignoreFiles',
'--ignorePatterns',
'--imf',
'--inputEVNT_PreFile',
'--inputFileValidation',
24 '--inputGenConfFile',
'--inputGeneratorFile',
'--jobConfig',
'--jobid',
'--maxEvents',
'--orphanKiller',
'--outputEVNTFile',
'--outputEVNT_PreFile',
'--outputHEPMCFile',
'--outputFileValidation',
'--outputNTUP_TRUTHFile',
'--outputTXTFile',
'--parallelFileValidation',
'--postExec',
'--postInclude',
'--preExec',
'--preInclude',
'--wprintEvts',
'--randomSeed',
'--reportName',
'--reportType',
'--rivetAnas',
'--runNumber',
'--showGraph',
'--showPath',
'--showSteps',
'--skipEvents',
'--skipFileValidation',
'--skipInputFileValidation',
'--skipOutputFileValidation',
'--steering',
'--taskid',
'--tcmalloc',
'--valgrind',
'--valgrindbasicopts',
'--valgrindextraopts',
'--lheOnly',
'--localPath',
'--cleanOut',
'--saveList']
27 "Specialised trf executor class for event generation jobs"
28 def __init__(self, name="generate", skeleton=None, skeletonCA=None, substep=None, inData=set(), outData=
set()):
29 athenaExecutor.__init__(self, name=name, skeletonFile=skeleton, skeletonCA=skeletonCA,
30 substep=substep, inputEventTest=
False, tryDropAndReload=
False,
31 inData=inData, outData=outData)
34 "Get input tarball, unpack and set up env if an evgenJobOpts arg was provided."
37 super(EvgenExecutor, self).
preExecute(input, output)
39 def expand_if_archive(filename):
40 "Function to expand a file if it is a zip archive or tarball"
41 if ".tar" in filename:
43 with tarfile.open(filename)
as tf:
45 elif filename.endswith(
".zip"):
47 zf = zipfile.ZipFile(filename)
52 return [name
for name
in os.listdir(a_dir)
53 if os.path.isdir(os.path.join(a_dir, name))]
57 os.environ[
'LOCAL_INSTALL_DIR'] = (os.environ[
'JOBOPTSEARCHPATH']).
split(
":")[0]
58 if os.path.exists(
'/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common'):
59 CommonCvmfsDir =
'/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common'
60 os.environ[
"JOBOPTSEARCHPATH"] = CommonCvmfsDir+
":"+os.environ[
"JOBOPTSEARCHPATH"]
61 if os.path.exists(
'/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common/MadGraphControl/dat/'):
62 datCvmfsDir =
'/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common/MadGraphControl/dat/'
63 os.environ[
"DATAPATH"] = datCvmfsDir+
":"+os.environ[
"DATAPATH"]
65 dsidpar = (self._trf.argdict[
"jobConfig"].value).
values()
66 dsidparam =
list(dsidpar)[0][0]
68 if os.path.exists(
'/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/'):
69 BaseCvmfsPath =
"/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/"
71 if dsidparam.isdigit()
and (len(dsidparam)==6
or len(dsidparam)==7):
73 Jodir = dsidparam[:3]+
'xxx'
75 Jodir = dsidparam[0] +
'/' + dsidparam[:4]+
'xxx'
77 cwd_ful = os.path.join(cwdir, dsidparam)
78 if (os.path.isdir(cwd_ful)):
79 os.environ[
"JOBOPTSEARCHPATH"] = cwd_ful+
":"+os.environ[
"JOBOPTSEARCHPATH"]
80 os.environ[
"DATAPATH"] = cwd_ful+
":"+os.environ[
"DATAPATH"]
82 cwd_Jodir = os.path.join(cwdir,Jodir)
83 cwd_Jodir_ful = os.path.join(cwd_Jodir,dsidparam)
84 if (os.path.isdir(cwd_Jodir_ful)):
85 os.environ[
"JOBOPTSEARCHPATH"] = cwd_Jodir_ful+
":"+os.environ[
"JOBOPTSEARCHPATH"]
86 os.environ[
"DATAPATH"] = cwd_Jodir_ful+
":"+os.environ[
"DATAPATH"]
88 if (os.path.isdir(BaseCvmfsPath)):
89 JoCvmfsPath = os.path.join(BaseCvmfsPath, Jodir)
90 JoCvmfsPath_ful = os.path.join(JoCvmfsPath, dsidparam)
91 os.environ[
"JOBOPTSEARCHPATH"] = JoCvmfsPath_ful+
":"+os.environ[
"JOBOPTSEARCHPATH"]
92 os.environ[
"DATAPATH"] = JoCvmfsPath_ful+
":"+os.environ[
"DATAPATH"]
94 msg.error(
"No access to JOs cvmfs location. JOs should be placed in working directory or full path should be given.")
95 raise RuntimeError(
"JOs not found - no cvmfs access")
97 if dsidparam.startswith(
"Test"):
98 for token
in (os.environ[
'JOBOPTSEARCHPATH']).
split(
":"):
99 if "jobOptions" in token:
102 search_token +=
"/EvgenJobTransforms/EvgenTest/" + dsidparam.split(
"Test")[-1]
103 os.environ[
"JOBOPTSEARCHPATH"] = search_token+
":"+os.environ[
"JOBOPTSEARCHPATH"]
104 os.environ[
"DATAPATH"] = search_token+
":"+os.environ[
"DATAPATH"]
105 elif (os.path.isdir(dsidparam)):
106 os.environ[
"JOBOPTSEARCHPATH"] = dsidparam+
":"+os.environ[
"JOBOPTSEARCHPATH"]
107 os.environ[
"DATAPATH"] = dsidparam+
":"+os.environ[
"DATAPATH"]
109 msg.error(
"JOs not found, please check = '%s'" % dsidparam)
110 raise RuntimeError(
"JOs not found")
112 msg.info(
"Using JOBOPTSEARCHPATH = '%s'" % os.environ[
"JOBOPTSEARCHPATH"])
113 msg.info(
"Using DATAPATH = '%s'" % os.environ[
"DATAPATH"])
115 if "evgenJobOpts" in self._trf.argdict:
116 tarball = self._trf.argdict[
"evgenJobOpts"].value
119 if tarball.startswith(
"http"):
121 tarball = os.basename(tarball)
123 url =
"http://cern.ch/atlas-computing/links/kitsDirectory/EvgenJobOpts/" + tarball
125 if not os.path.exists(tarball):
126 from EvgenJobTransforms.download
import downloadUsingProxy
129 raise EnvironmentError(
'Error downloading tarball %s. Downloader reports: %s' % (tarball, output))
130 msg.info(
'Evgen tarball download success: %s' % output)
132 expand_if_archive(tarball)
137 FIRST_DIR = (os.environ[
'JOBOPTSEARCHPATH']).
split(
":")[0]
138 configFiles = [f
for f
in os.listdir(FIRST_DIR)
if (
"GRID" in f)]
140 if len(configFiles) == 1:
141 msg.info(
"gridpack for only one energy available ")
142 elif len(configFiles) >1:
143 msg.info(
"more then one gridpack ! ")
144 if len(configFiles) >=1:
145 if "--ecmEnergy" in str(sys.argv[1:]):
146 split_args=
str(sys.argv[1:]).
split(
"ecmEnergy",1)[1]
147 split_args=split_args.lstrip(
"\',=")
148 ener_GeV=split_args.split(
",")[0].strip(
" ,\']")
150 msg.info(
"Should be used gridpack for energy "+energy)
152 msg.info(
"no ecm energy given, assuming 13.6 TeV ")
154 for x
in configFiles:
155 gridS=
"mc_"+energy+
"TeV"
156 msg.info(
"Gridpack should start from "+gridS)
157 if x.startswith(gridS):
158 confFile = os.path.join(FIRST_DIR, x)
159 msg.info(
"using gridpack = "+confFile)
161 msg.error(
"No *GRID* config files, for requested energy = '%s' please check = '%s'" %(energy,dsidparam))
164 if confFile
is not None:
165 expand_if_archive(confFile)
167 msg.info(
"Configuration input gridpack found " + confFile)
170 loc_files = os.listdir(os.getcwd())
171 for loc_file
in loc_files:
172 if "GRID" not in loc_file:
173 expand_if_archive(loc_file)
176 if "inputGeneratorFile" in self._trf.argdict:
178 myinputfiles = self._trf.argdict[
"inputGeneratorFile"].value
179 genInputFiles = myinputfiles.split(
',')
180 for file
in genInputFiles:
181 expand_if_archive(file)
182 if "inputGenConfFile" in self._trf.argdict:
183 expand_if_archive(self._trf.argdict[
"inputGenConfFile"].value)
186 files = os.listdir(tmp_dir)
192 dest = main_dir+
"/"+f
199 msg.info(
"Transform arguments %s" % sys.argv[1:])
200 if "--outputEVNTFile" in str(sys.argv[1:]):
201 exeSet.add(
EvgenExecutor(name=
"generate", skeleton=
"EvgenJobTransforms/skel.GENtoEVGEN.py", skeletonCA=
"EvgenJobTransforms.GENtoEVGEN_Skeleton", inData=[
"inNULL"], outData=[
"YODA",
"EVNT",
"EVNT_Pre",
"TXT"]))
202 msg.info(
"Output EVNT file")
203 elif "--outputYODAFile" in str(sys.argv[1:]):
204 exeSet.add(
EvgenExecutor(name=
"generate", skeleton=
"EvgenJobTransforms/skel.GENtoEVGEN.py", inData=[
"inNULL"], outData=[
"YODA",
"TXT"]))
205 msg.info(
"Output EVNT file")
206 elif "--outputTXTFile" in str(sys.argv[1:]):
207 exeSet.add(
EvgenExecutor(name=
"generate", skeleton=
"EvgenJobTransforms/skel.GENtoTXT.py", inData=[
"inNULL"], outData=[
"TXT"]))
208 msg.info(
"Output TXT file")
209 elif "--outputHEPMCFile" not in str(sys.argv[1:]):
210 msg.error(
"Output cannot be recognised")
212 exeSet.add(
EvgenExecutor(name=
"afterburn", skeleton=
"EvgenJobTransforms/skel.ABtoEVGEN.py", inData=[
"EVNT_Pre"], outData=[
"EVNT"]))
213 exeSet.add(athenaExecutor(name =
"AODtoDPD", skeletonFile =
"PATJobTransforms/skeleton.AODtoDPD_tf.py",
214 substep =
"a2d", inData = [
"EVNT"], outData = [
"NTUP_TRUTH"], perfMonFile =
"ntuple_AODtoDPD.pmon.gz"))
215 exeSet.add(athenaExecutor(name =
'EVNTtoHEPMC', skeletonCA =
'EvgenJobTransforms.POOLtoHEPMC_Skeleton',
216 substep =
"a2d", perfMonFile =
'ntuple.pmon.gz', inData=[
'EVNT'], outData=[
'HEPMC']))
223 @stdTrfExceptionHandler
226 msg.info(
"This is %s" % sys.argv[0])
228 main_dir = os.getcwd()
230 trf.parseCmdLineArgs(sys.argv[1:])
231 if ((
"cleanOut" in trf.argdict)
and (trf.argdict[
"cleanOut"].value != 0)):
232 name_tmpdir =
"tmprun"
233 tmp_dir = os.path.join(main_dir, name_tmpdir)
234 if os.path.isdir(tmp_dir):
235 shutil.rmtree(tmp_dir, ignore_errors=
True)
238 tmp_dir = os.getcwd()
239 whitelist_in = [
'MC',
'group',
'TXT']
244 msg.info(
"%s stopped at %s, trf exit code %d" % (sys.argv[0], time.asctime(), trf.exitCode))
249 if ((
"cleanOut" in trf.argdict)
and (trf.argdict[
"cleanOut"].value!=0)):
250 whitelist_out = [
'log.generate',
'.root']
251 if "outputTXTFile" in trf.argdict:
252 whitelist_out.append(
'TXT')
253 if "saveList" in trf.argdict:
254 saveList_dic= trf.argdict[
"saveList"].value
255 saveList_str=
str(saveList_dic)
256 saveList_str=saveList_str[10:-3]
257 saveList= saveList_str.split(
",")
258 for item
in saveList:
259 test_ex = os.path.join(main_dir,
str(item))
260 if os.path.isdir(test_ex):
261 shutil.rmtree(test_ex, ignore_errors=
True)
262 elif os.path.isfile(test_ex):
264 if not saveList[0].isdigit():
265 whitelist_out=whitelist_out+saveList
269 if "saveList" not in trf.argdict:
270 shutil.rmtree(tmp_dir, ignore_errors=
True)
271 elif not saveList[0].isdigit():
272 shutil.rmtree(tmp_dir, ignore_errors=
True)
274 elif (
"inputGeneratorFile" in trf.argdict):
275 myinputfiles = trf.argdict[
"inputGeneratorFile"].value
276 genInputFiles = myinputfiles.split(
',')
277 numberOfFiles = len(genInputFiles)
278 merge_file =
'merged_lhef._0.events'
279 if((numberOfFiles>1)
and (os.path.exists(merge_file))):
280 os.remove(merge_file)
282 if ((
"lheOnly" in trf.argdict )
and (trf.argdict[
"lheOnly"].value == 1)):
283 outputName =
''.
join(trf.argdict[
"outputEVNTFile"].value)
284 os.remove(outputName)
285 sys.exit(trf.exitCode)
291 if __name__ ==
"__main__":