EIC code displayed by LXR

0001 #!/usr/local/bin/python
0002
0003 # same as make_dawn_views but stops at generating the prim file.
0004 # W. Armstrong (ANL), original bash script
0005 # C. Peng (ANL), translate to python and add flexible run time for simulation
0006
0007 import os
0008 import signal
0009 import subprocess
0010 import argparse
0011 import atexit
0012 import time
0013 from datetime import datetime
0014 import fcntl
0015 import psutil
0016
0017
0018 def readline_nonblocking(output):
0019     fd = output.fileno()
0020     fl = fcntl.fcntl(fd, fcntl.F_GETFL)
0021     fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
0022     try:
0023         return output.readline()
0024     except:
0025         return ''
0026
0027
0028 # arguments
0029 parser = argparse.ArgumentParser()
0030
0031 parser.add_argument('-c', '--compact-file', type=str, dest='compact',
0032         default=os.path.join(os.environ.get('DETECTOR_PATH', '.'), 'athena.xml'),
0033         help='Top level compact file for detectors')
0034
0035 parser.add_argument('-s', '--skip', type=int,
0036         default=0,
0037         help='Number of events number to skip in the input')
0038
0039 parser.add_argument('-i', '--input', type=str,
0040         default='scripts/input_data/few_events.hepmc',
0041         help='Input hepmc file')
0042
0043 parser.add_argument('-o', '--output-dir', type=str, dest='out_dir',
0044         default='sim_output',
0045         help='output directory')
0046
0047 parser.add_argument('-D', '--detector-only', action='store_true', dest='detector_only',
0048         help='only generate the prim files for the detector geometry')
0049
0050 parser.add_argument('-t', '--tag', type=str,dest='file_tag',
0051         default='view',
0052         help='Output file tag')
0053
0054 parser.add_argument('--timeout', type=int,
0055         default=3600,
0056         help='Timeout in seconds')
0057
0058 parser.add_argument('passthrough', nargs='*')
0059
0060 args = parser.parse_args()
0061
0062 macro = 'macro/dawn_picture.mac' if args.detector_only else 'macro/dawn_picture2.mac'
0063
0064 # raise error if cannot create a temporary working dir
0065 # os.makedirs('dawn_view_tmp', exist_ok=False)
0066 os.makedirs(args.out_dir, exist_ok=True)
0067
0068 # use absolute path so the chdir does not affect them
0069 args.input = os.path.abspath(args.input)
0070 args.out_dir = os.path.abspath(args.out_dir)
0071 args.compact = os.path.abspath(args.compact)
0072 macro = os.path.abspath(macro)
0073
0074 # adjust fiber radius to reduce the number of fibers
0075 compact_dir = os.path.dirname(os.path.realpath(args.compact))
0076 ci_ecal = os.path.join(compact_dir, 'compact', 'ci_ecal_scfi.xml')
0077 os.system('sed -i \'s/radius=\"EcalEndcapP_FiberRadius\"/radius=\"EcalEndcapP_FiberRadius*10\"/\' {}'.format(ci_ecal))
0078
0079 prim_file = 'g4_0000.prim'
0080 dawn_env = os.environ.copy()
0081 dawn_env['DAWN_BATCH'] = 'a'
0082 # sdir = os.path.dirname(os.path.realpath(__file__))
0083
0084 # Using a python warpper such as npsim introduces some problem in managing the subprocess.
0085 # The process1 managed by subprocess will generate another process with proc2_pid = proc1_pid + 1, which will not
0086 # be terminated by terminating or killing the process1.
0087 # In addition, running Geant4 with vis mode will never exit automatically (it waits for input).
0088 # Thus the created process 2 will occupy the system resources.
0089 sim_cmd = ['npsim', '--runType', 'vis',
0090         '--compact', args.compact,
0091         '--inputFiles', args.input,
0092         '--outputFile', 'derp.root',
0093         '--numberOfEvents', '1',
0094         '--skipNEvents', str(args.skip),
0095         '--macroFile', macro]
0096
0097 start = datetime.now()
0098 elapse = datetime.now() - start
0099 last_update = datetime.now()
0100 finished = False
0101
0102 # run simulation
0103 print(' '.join(sim_cmd))
0104 p = subprocess.Popen(args=sim_cmd, env=dawn_env,
0105                      stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
0106 __child_pid = p.pid
0107 while elapse.seconds < args.timeout:
0108     line = readline_nonblocking(p.stdout)
0109     elapse = datetime.now() - start
0110     time_left = args.timeout - elapse.seconds
0111     time_str = '[{:02d}:{:02d}]'.format(elapse.seconds // 60, elapse.seconds % 60)
0112
0113     if time_left < 10:
0114         print('{} === TIMEOUT ===: Terminating in {:d} seconds'.format(time_str, time_left))
0115
0116     if line:
0117         decoded_line = line.decode('utf-8').strip()
0118         print('{} {}'.format(time_str, decoded_line))
0119         # what we are looking for
0120         if decoded_line == 'File  {}  is generated.'.format(prim_file):
0121             print('{} === FINISHED ===: Got the prim file, terminating.'.format(time_str))
0122             finished = True
0123             break
0124         if decoded_line == 'Idle>':
0125             p.stdin.write(b'exit')
0126             break
0127         # do not sleep
0128         continue
0129
0130     # ended early before file
0131     if p.poll() is not None:
0132         print(p.poll())
0133         break
0134
0135     time.sleep(1)
0136
0137 p.kill()
0138 # use to kill the subprocess generated from the python wrapper
0139 # this is unsafe so maybe more checks required
0140 for proc in psutil.process_iter():
0141     pinfo = proc.as_dict(attrs=['pid', 'name', 'create_time'])
0142     if pinfo['pid'] == p.pid + 1 and pinfo['name'] == 'python':
0143         print('kill {}, generated from {}'.format(pinfo, p.pid))
0144         os.kill(pinfo['pid'], signal.SIGTERM)
0145
0146 # revert the change
0147 os.system('sed -i \'s/radius=\"EcalEndcapP_FiberRadius*10\"/radius=\"EcalEndcapP_FiberRadius\"/\' {}'.format(ci_ecal))
0148
0149 line = b'stderr outputs:\n'
0150 while line:
0151     print(line.decode('utf-8'), end='')
0152     line = readline_nonblocking(p.stderr)
0153
0154 if finished:
0155     print('Simulation finished')
0156 else:
0157     print('Simulation failed')
0158     exit(1)
0159
0160 # move the prim files (which can be quite large)
0161 # to the local pipeline storage path
0162 os.system('mv g4_0000.prim {}/{}.prim'.format(args.out_dir,args.file_tag))
0163
0164