emcal_calib_year1/base/run.py

0001 #!/usr/bin/env python3
0002 """
0003 This module automates the EMCal calibration procedure.
0004 """
0005
0006 import argparse
0007 import os
0008 import sys
0009 import subprocess
0010 import datetime
0011 import time
0012 import shutil
0013 import logging
0014
0015 parser = argparse.ArgumentParser()
0016
0017 parser.add_argument('-i'
0018                     , '--input-list', type=str
0019                     , required=True
0020                     , help='Input DST List.')
0021
0022 parser.add_argument('-o'
0023                     , '--project-dir', type=str
0024                     , default='/gpfs02/sphenix/user/anarde/EMCal-Calib-Test'
0025                     , help='Project Directory. Default: .')
0026
0027 parser.add_argument('-it1'
0028                     , '--iter-start', type=int
0029                     , default=0
0030                     , help='Starting Iteration. Default: 0')
0031
0032 parser.add_argument('-it2'
0033                     , '--iter-end', type=int
0034                     , default=10
0035                     , help='Ending Iteration. Default: 10')
0036
0037 parser.add_argument('-s'
0038                     , '--memory', type=float
0039                     , default=2
0040                     , help='Memory (units of GB) to request per condor submission. Default: 2 GB.')
0041
0042 parser.add_argument('-l'
0043                     , '--log-file', type=str
0044                     , default='log.txt'
0045                     , help='Log File. Default: log.txt')
0046
0047 parser.add_argument('-l2'
0048                     , '--condor-log-dir', type=str
0049                     , default='/tmp/anarde/dump'
0050                     , help='Condor Log Directory. Default: /tmp/anarde/dump')
0051
0052 parser.add_argument('-p'
0053                     , '--files-per-hadd', type=int
0054                     , default=20
0055                     , help='Files Per hadd. Default: 20')
0056
0057 parser.add_argument('-s2'
0058                     , '--sleep-interval', type=int
0059                     , default=60
0060                     , help='Sleep Interval. Default: 60 seconds.')
0061
0062 parser.add_argument('-f'
0063                     , '--f4a-macro', type=str
0064                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/Fun4All_EMCal.C'
0065                     , help='Fun4All Macro. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/Fun4All_EMCal.C')
0066
0067 parser.add_argument('-f1'
0068                     , '--fit-calib-macro', type=str
0069                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/doFitAndCalibUpdate.C'
0070                     , help='Fit Calib Macro. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/doFitAndCalibUpdate.C')
0071
0072 parser.add_argument('-f2'
0073                     , '--tsc-fit-macro', type=str
0074                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/doTscFit.C'
0075                     , help='TSC Calib Macro. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/macros/doTscFit.C')
0076
0077 parser.add_argument('-f3'
0078                     , '--condor-script', type=str
0079                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/scripts/genCalib.sh'
0080                     , help='Condor Script. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/scripts/genCalib.sh')
0081
0082 parser.add_argument('-b'
0083                     , '--f4a-bin', type=str
0084                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/Fun4All_EMCal'
0085                     , help='Fun4All Bin. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/Fun4All_EMCal')
0086
0087 parser.add_argument('-b1'
0088                     , '--fit-calib-bin', type=str
0089                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/doFitAndCalibUpdate'
0090                     , help='Fit Calib Bin. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/doFitAndCalibUpdate')
0091
0092 parser.add_argument('-b2'
0093                     , '--tsc-fit-bin', type=str
0094                     , default='/gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/doTscFit'
0095                     , help='TSC Calib Bin. Default: /gpfs02/sphenix/user/anarde/sPHENIX/analysis-EMCal-Calibration/EMCal-Calibration-Study/bin/doTscFit')
0096
0097 parser.add_argument('-f4'
0098                     , '--calib-field', type=str
0099                     , default='CEMC_calib_ADC_to_ETower'
0100                     , help='Calib Field. Default: CEMC_calib_ADC_to_ETower')
0101
0102 args = parser.parse_args()
0103
0104 def setup_logging(log_file, log_level):
0105     """Configures the logging system to output to a file and console."""
0106
0107     # Create a logger instance
0108     logger = logging.getLogger(__name__) # Use __name__ to get a logger specific to this module
0109     logger.setLevel(log_level)
0110
0111     # Clear existing handlers to prevent duplicate output if run multiple times
0112     if logger.hasHandlers():
0113         logger.handlers.clear()
0114
0115     # Create a formatter for log messages
0116     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
0117
0118     # Create a FileHandler to save logs to a file
0119     file_handler = logging.FileHandler(log_file)
0120     file_handler.setLevel(log_level)
0121     file_handler.setFormatter(formatter)
0122     logger.addHandler(file_handler)
0123
0124     # Create a StreamHandler to also output logs to the console (optional)
0125     console_handler = logging.StreamHandler(sys.stdout)
0126     console_handler.setLevel(logging.INFO) # Console might only show INFO and above
0127     console_handler.setFormatter(formatter)
0128     logger.addHandler(console_handler)
0129
0130     return logger
0131
0132 def run_command_and_log(command, logger, current_dir = '.', description="Executing command"):
0133     """
0134     Runs an external command using subprocess and logs its stdout, stderr, and return code.
0135     """
0136     logger.info(f"{description}: '{command}'")
0137
0138     try:
0139         # subprocess.run is the recommended high-level API
0140         # capture_output=True: captures stdout and stderr
0141         # text=True: decodes output as text (usually UTF-8)
0142         # check=False: do NOT raise an exception for non-zero exit codes immediately.
0143         #              We want to log stderr even on failure before deciding to raise.
0144         result = subprocess.run(['bash','-c',command], cwd=current_dir, capture_output=True, text=True, check=False)
0145
0146         # Log stdout if any
0147         if result.stdout:
0148             # Using logger.debug allows capturing even verbose outputs
0149             logger.debug(f"  STDOUT from '{command}':\n{result.stdout.strip()}")
0150
0151         # Log stderr if any
0152         if result.stderr:
0153             # Using logger.error for stderr, as it often indicates problems
0154             logger.error(f"  STDERR from '{command}':\n{result.stderr.strip()}")
0155
0156         # Log the return code
0157         logger.info(f"  Command exited with code: {result.returncode}")
0158
0159         # You can choose to raise an exception here if the command failed
0160         if result.returncode != 0:
0161             logger.error(f"Command failed: '{command}' exited with non-zero code {result.returncode}")
0162             # Optionally, raise an error to stop execution
0163             # raise subprocess.CalledProcessError(result.returncode, command, output=result.stdout, stderr=result.stderr)
0164             return False
0165         return True
0166
0167     except FileNotFoundError:
0168         logger.critical(f"Error: Command '{command}' not found. Is it in your PATH?")
0169         return False
0170     except Exception as e:
0171         logger.critical(f"An unexpected error occurred while running '{command}': {e}")
0172         return False
0173
0174 def main():
0175     """
0176     Main Function
0177     """
0178     input_list = os.path.realpath(args.input_list)
0179     total_jobs = int(subprocess.run(['bash','-c',f'wc -l {input_list}'], capture_output=True, encoding='utf-8', check=False).stdout.strip().split()[0])
0180     project_dir = os.path.realpath(args.project_dir)
0181     iter_start = args.iter_start
0182     iter_end = args.iter_end
0183     memory = args.memory
0184     log_file  = os.path.realpath(args.log_file)
0185     log_dir   = os.path.dirname(log_file)
0186     condor_log_dir = os.path.realpath(args.condor_log_dir)
0187     files_per_hadd = args.files_per_hadd
0188     sleep_interval = args.sleep_interval # seconds
0189     f4a_macro = os.path.realpath(args.f4a_macro)
0190     fit_calib_macro = os.path.realpath(args.fit_calib_macro)
0191     tsc_fit_macro = os.path.realpath(args.tsc_fit_macro)
0192     f4a_bin = os.path.realpath(args.f4a_bin)
0193     fit_calib_bin = os.path.realpath(args.fit_calib_bin)
0194     tsc_fit_bin = os.path.realpath(args.tsc_fit_bin)
0195     condor_script = os.path.realpath(args.condor_script)
0196     calib_field = args.calib_field
0197
0198     # Check args
0199     if iter_start > iter_end or iter_start < 0:
0200         parser.error(f'ERROR: {iter_start} > {iter_end} or {iter_start} is negative.')
0201
0202     if sleep_interval <= 0:
0203         parser.error(f'ERROR: Negative sleep interval: {sleep_interval}.')
0204
0205     # Create Dirs
0206     os.makedirs(log_dir,exist_ok=True)
0207     os.makedirs(project_dir, exist_ok=True)
0208
0209     # Initialize the logger
0210     logger = setup_logging(log_file, logging.DEBUG)
0211
0212     if os.path.exists(condor_log_dir):
0213         shutil.rmtree(condor_log_dir)
0214         os.makedirs(condor_log_dir)
0215
0216     # Print Logs
0217     logger.info('#'*40)
0218     logger.info(f'LOGGING: {datetime.datetime.now()}')
0219     logger.info(f'Input DST List: {input_list}')
0220     logger.info(f'Total Jobs per Iteration: {total_jobs}')
0221     logger.info(f'Project Directory: {project_dir}')
0222     logger.info(f'Iteration Start: {iter_start}')
0223     logger.info(f'Iteration End: {iter_end}')
0224     logger.info(f'Condor Memory (per Job): {memory}')
0225     logger.info(f'Log File: {log_file}')
0226     logger.info(f'Files Per hadd: {files_per_hadd}')
0227     logger.info(f'Fun4All Macro: {f4a_macro}')
0228     logger.info(f'Fit Calib Macro: {fit_calib_macro}')
0229     logger.info(f'TSC Fit Macro: {tsc_fit_macro}')
0230     logger.info(f'Fun4All Bin: {f4a_bin}')
0231     logger.info(f'Fit Calib Bin: {fit_calib_bin}')
0232     logger.info(f'TSC Fit Bin: {tsc_fit_bin}')
0233     logger.info(f'Condor Script: {condor_script}')
0234     logger.info(f'Calib Field: {calib_field}')
0235
0236     # Iteration Loop
0237     for it in range(iter_start,iter_end+1):
0238         logger.info(f'Iteration: {it}')
0239
0240         if it == 0:
0241             # command = f'root -b -l -q \'{f4a_macro}(0,"{input_list}",{it})\''
0242             command = f'{f4a_bin} 0 {input_list} {it}'
0243             run_command_and_log(command, logger, project_dir)
0244             continue
0245
0246         iter_dir = os.path.join(project_dir, f'test-iter-{it}')
0247
0248         os.makedirs(iter_dir, exist_ok=True)
0249
0250         os.makedirs(f'{iter_dir}/output', exist_ok=True)
0251         os.makedirs(f'{iter_dir}/stdout', exist_ok=True)
0252         os.makedirs(f'{iter_dir}/error', exist_ok=True)
0253         if it == 3:
0254             os.makedirs(f'{iter_dir}/figures', exist_ok=True)
0255
0256         if it == 1:
0257             shutil.copy(f'{project_dir}/local_calib_copy.root', iter_dir)
0258         else:
0259             prev_iter_dir = os.path.join(project_dir, f'test-iter-{it-1}')
0260             shutil.copy(f'{prev_iter_dir}/local_calib_copy.root', iter_dir)
0261
0262         shutil.copy(f'{condor_script}', iter_dir)
0263         shutil.copy(f'{f4a_macro}', iter_dir)
0264         shutil.copy(f'{input_list}', iter_dir)
0265
0266         if it <= 3:
0267             shutil.copy(f'{tsc_fit_macro}', iter_dir)
0268         else:
0269             shutil.copy(f'{fit_calib_macro}', iter_dir)
0270
0271         with open(f'{iter_dir}/genCalib.sub', mode='w', encoding='utf-8') as file:
0272             file.write(f'executable    = {os.path.basename(condor_script)}\n')
0273             file.write(f'arguments     = {f4a_bin} 0 $(input_dst) {it} {iter_dir}/local_calib_copy.root {calib_field} {iter_dir}/output\n')
0274             file.write(f'log           = {condor_log_dir}/job-$(ClusterId)-$(Process).log\n')
0275             file.write('output         = stdout/job-$(ClusterId)-$(Process).out\n')
0276             file.write('error          = error/job-$(ClusterId)-$(Process).err\n')
0277             file.write(f'request_memory = {memory}GB\n')
0278
0279         logger.info(f'Iter: {it}, Condor Directory Generated: {datetime.datetime.now()}')
0280
0281         command = f'condor_submit genCalib.sub -queue "input_dst from {input_list}"'
0282         logger.info(command)
0283         run_command_and_log(command, logger, iter_dir)
0284
0285         while True:
0286             jobs = int(subprocess.run(['bash','-c','ls output | wc -l'], capture_output=True, encoding='utf-8', cwd=iter_dir, check=False).stdout.strip())
0287
0288             if jobs == total_jobs:
0289                 logger.info(f'Iter: {it}, All Jobs Finished: {datetime.datetime.now()}')
0290                 break
0291
0292             logger.info(f'Checking Condor Output Status: {datetime.datetime.now()}, Jobs: {jobs} out of {total_jobs}, {jobs * 100 / total_jobs:0.2f} %')
0293
0294             time.sleep(sleep_interval)
0295
0296
0297         logger.info(f'Iter: {it}, hadd start: {datetime.datetime.now()}')
0298
0299         command = f'hadd -n {files_per_hadd+1} test-iter{it}.root output/*'
0300         run_command_and_log(command, logger, iter_dir)
0301
0302         logger.info(f'Iter: {it}, hadd finish: {datetime.datetime.now()}')
0303
0304         shutil.copy(f'{iter_dir}/local_calib_copy.root', f'{iter_dir}/local_calib_copy_{it-1}.root')
0305
0306         # macro = fit_calib_macro if it >= 4 else tsc_fit_macro
0307         calib_bin = fit_calib_bin if it >= 4 else tsc_fit_bin
0308
0309         # command = f'root -b -l -q \'{macro}("test-iter{it}.root", "local_calib_copy.root", {it}, "{calib_field}")\''
0310         command = f'{calib_bin} test-iter{it}.root local_calib_copy.root {it} {calib_field}'
0311         run_command_and_log(command, logger, iter_dir)
0312
0313         logger.info(f'Iter: {it}, local_calib_copy.root updated: {datetime.datetime.now()}')
0314
0315 if __name__ == "__main__":
0316     main()