Moves reconstruction from a bash script to a python script.

client.py

@@ -2,22 +2,10 @@
 
 def compute(version_number, max_events, skip_events, event_file):
   ''' Runs reco.sh with the given parameters, using a tmp folder to clean up intermediary files. '''
-  import tempfile, subprocess, os, stat, shutil, socket, string, constants
-  log_file_handle, _ = tempfile.mkstemp(dir=constants.log_dir, prefix=(string.split(socket.gethostname(), sep='.')[0]+'_'), suffix='.log')
-  tmp_dir = tempfile.mkdtemp(dir=constants.tmp_dir)
-  arg = '{} {} {} {} {}'.format(constants.reco, version_number, max_events, skip_events, event_file)
-  process = subprocess.Popen(arg, executable='/bin/bash', cwd=tmp_dir, shell=True, stdout=log_file_handle, stderr=subprocess.STDOUT)
-  process.wait()
-  # move the aod file to the output directory, and make it immutable so that
-  # it is not accidentally deleted.
-  output_file = '{}.aod.pool.root'.format(version_number)
-  os.rename(tmp_dir + '/' + output_file, constants.aod_output_dir + '/' + output_file)
-  output_file = constants.aod_output_dir + '/' + output_file
-  st = os.stat(output_file)
-  not_writable = ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
-  os.chmod(output_file, st.st_mode & not_writable)
-  shutil.rmtree(tmp_dir)
-  return (socket.gethostname(), output_file)
+  import subprocess
+  arg = 'python {} -n {} -s {} {} {}'.format(constants.reco, max_events, skip_events, event_file, version_number)
+  subprocess.check_call(arg, executable='/bin/bash', shell=True)
+  return socket.gethostname()
 
 def dispatch_computations(batch_size, evnt_dir, test, local):
   import dispy, dispy.httpd, glob, math, string, constants
@@ -51,8 +39,8 @@ def dispatch_computations(batch_size, evnt_dir, test, local):
       job_id += 1
   if not test and not local:    
     for job in jobs:
-      host, aod_file = job()
-      print('{} executed job {} from {} to {} creating {}'.format(host, job.id, job.start_time, job.end_time, aod_file))
+      host = job()
+      print('{} executed job {} from {} to {}'.format(host, job.id, job.start_time, job.end_time))
     cluster.print_status()
     http_server.shutdown()
 

constants.py

@@ -1,5 +1,16 @@
+evnt_dir = '/phys/groups/tev/scratch3/users/WHHV/evnts'
 tmp_dir = '/phys/groups/tev/scratch3/users/WHHV/tmp'
 log_dir = '/phys/groups/tev/scratch3/users/WHHV/logs'
-reco = '/phys/users/schuya/whhv/Gen_Reco/reco.sh'
-aod_output_dir = '/phys/groups/tev/scratch3/users/WHHV/aod_output'
+aod_dir = '/phys/groups/tev/scratch3/users/WHHV/aod'
+reco = '/phys/users/schuya/whhv/Gen_Reco/reco.py'
+
+geometry_version = 'ATLAS-R2-2015-03-01-00_VALIDATION'
+conditions_tag = 'OFLCOND-RUN12-SDR-20'
+run_number = 304795
+gen_job_config = 'hss-runner.py'
+
+gen_release = '19.2.4.14'
+sim_release = '19.2.4.14'
+dig_release = '19.2.4.14'
+reco_release = '19.2.4.14'
 

generate.py

@@ -1,8 +1,13 @@
+import argparse
+import os
+import stat
+import shutil
+import tempfile
+import subprocess
+import time
+import constants
 
 def gen(batch_size, num_batches, evnt_dir, release, job_config, run_number):
-  import os, stat, tempfile, subprocess, time, constants, shutil
-  import pdb
-  pdb.set_trace()
   max_seed = 30081**2
   processes = []
   tmp_dirs = []
@@ -18,8 +23,10 @@ def gen(batch_size, num_batches, evnt_dir, release, job_config, run_number):
     arg = '. /phys/users/gwatts/bin/CommonScripts/configASetup.sh && . $AtlasSetup/scripts/asetup.sh here,{} && Generate_tf.py --jobConfig {} --maxEvents {} --runNumber {} --firstEvent {} --outputEVNTFile {} --ecmEnergy 13000 --randomSeed {}'.format(release, job_config, batch_size, run_number, first_event, evnt_file, seed)
     process = subprocess.Popen(arg, executable='/bin/bash', cwd=tmp_dir, shell=True, stdout=log_file_handle, stderr=subprocess.STDOUT)
     processes.append(process)
-  for process in processes:
-    process.wait()
+  for i, process in enumerate(processes):
+    return_code = process.wait()
+    if return_code != 0:
+      print ('ERROR: Generation of batch {} failed with code: {}'.format(i, return_code))
   not_writable = ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
   for evnt_file in evnt_files:
     st = os.stat(evnt_file)
@@ -28,14 +35,13 @@ def gen(batch_size, num_batches, evnt_dir, release, job_config, run_number):
     shutil.rmtree(tmp_dir)
 
 def main():
-  import argparse, os
   parser = argparse.ArgumentParser(description='Generates events according to a hidden valley theory.')
   parser.add_argument('-b', '--batch_size', type=int, default=5000, help='The number of events to generate in each batch.')
   parser.add_argument('-n', '--num_batches', type=int, default=20, help='The number of batches to generate.')
-  parser.add_argument('--evnt_dir', default='/phys/groups/tev/scratch3/users/WHHV/evnts', help='The directory in which to store events.')
-  parser.add_argument('-r', '--release', default='19.2.4.14', help='The athena release to use.')
-  parser.add_argument('--job_config', default='hss-runner.py', help='The job configuration script for generating the events.')
-  parser.add_argument('--run_number', default='304795', help='The run number.')
+  parser.add_argument('--evnt_dir', default=constants.evnt_dir, help='The directory in which to store events.')
+  parser.add_argument('-r', '--release', default=constants.gen_release, help='The athena release to use.')
+  parser.add_argument('--job_config', default=constants.gen_job_config, help='The job configuration script for generating the events.')
+  parser.add_argument('--run_number', default=constants.run_number, help='The run number.')
   args = parser.parse_args()
   gen(args.batch_size, args.num_batches, os.path.abspath(args.evnt_dir), args.release, os.path.abspath(args.job_config), args.run_number)
 

reco.py

@@ -0,0 +1,49 @@
+import constants
+import os
+import stat
+import tempfile
+import shutil
+import subprocess
+
+def reco(evnt_file, version, aod_dir, num_events, skip_events, geometry_version, conditions_tag, sim_release, dig_release, reco_release, log_dir, log_prefix):
+  hits_file = version + '.hits.pool.root'
+  rdo_file = version + '.rdo.pool.root'
+  esd_file = version + '.esd.pool.root'
+  aod_file = version + '.aod.pool.root'
+  log_file_handle = open(os.path.join(log_dir, log_prefix + '.log'), 'w+')
+  tmp_dir = tempfile.mkdtemp(dir=constants.tmp_dir)
+  source_arg = '. /phys/users/gwatts/bin/CommonScripts/configASetup.sh && . $AtlasSetup/scripts/asetup.sh here,'
+  sim_arg = source_arg + '{} && Sim_tf.py --inputEVNTFile {} --geometryVersion {} --conditionsTag {} --outputHITSFile {}  --physicsList "FTFP_BERT" --postInclude "PyJobTransforms/UseFrontier.py" --preInclude "EVNTtoHITS:SimulationJobOptions/preInclude.BeamPipeKill.py" --maxEvents {} --skipEvents {} --randomSeed "8" --simulator "MC12G4" --truthStrategy "MC12"'.format(sim_release, evnt_file, geometry_version, conditions_tag, hits_file, num_events, skip_events)
+  dig_arg = source_arg + '{} && Digi_tf.py --inputHitsFile {} --outputRDOFile {} --geometryVersion {} --conditionsTag {}'.format(dig_release, hits_file, rdo_file, geometry_version, conditions_tag)
+  reco_arg = source_arg + '{} && Reco_tf.py  --inputRDOFile {} --outputESDFile {} --DBRelease current  --autoConfiguration= "everything" && Reco_tf.py --inputESDFile {} --outputAODFile {} --DBRelease current --autoConfiguration="everything"'.format(reco_release, rdo_file, esd_file, esd_file, aod_file)
+  for arg in (sim_arg, dig_arg, reco_arg):
+    subprocess.check_call(arg, executable='/bin/bash', cwd=tmp_dir, shell=True, stdout=log_file_handle, stderr=subprocess.STDOUT)
+  # move the aod file to the output directory, and make it immutable so that
+  # it is not accidentally deleted.
+  os.rename(os.path.join(tmp_dir, aod_file), os.path.join(aod_dir, aod_file))
+  aod_file_path = os.path.join(aod_dir, aod_file)
+  st = os.stat(aod_file_path)
+  not_writable = ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
+  os.chmod(aod_file_path, st.st_mode & not_writable) 
+  shutil.rmtree(tmp_dir)
+
+def main():
+  import argparse
+  parser = argparse.ArgumentParser(description='Perform simulation, digitization and reconstruction on a given event file.')
+  parser.add_argument('evnt_file', help='The absolute path to the event file.')
+  parser.add_argument('version', help='The name (without file extensions) of the output AOD file.')
+  parser.add_argument('--aod_dir', default=constants.aod_dir, help='The path of the aod directory.')
+  parser.add_argument('-n', '--num_events', type=int, default=-1, help='The number of events.')
+  parser.add_argument('-s', '--skip_events', type=int, default=0, help='The number of events to skip.')
+  parser.add_argument('--geometry_version', default=constants.geometry_version, help='Geometry version for simulation and digitization.')
+  parser.add_argument('--conditions_tag', default=constants.conditions_tag, help='Conditions tag for simulation and digitization.')
+  parser.add_argument('--sim_release', default=constants.sim_release, help='The athena release to use for simulation.')
+  parser.add_argument('--dig_release', default=constants.dig_release, help='The athena release to use for digitization.')
+  parser.add_argument('--reco_release', default=constants.reco_release, help='The athena release to use for reconstruction.')
+  parser.add_argument('--log_dir', default=constants.log_dir, help='The log directory.')
+  parser.add_argument('--log_prefix', default='reco', help='The prefix for the log file.')
+  args = parser.parse_args()
+  reco(args.evnt_file, args.version, args.aod_dir, args.num_events, args.skip_events, args.geometry_version, args.conditions_tag, args.sim_release, args.dig_release, args.reco_release, args.log_dir, args.log_prefix)
+
+if __name__ == '__main__':
+  main()