File indexing completed on 2025-04-17 02:41:48
0001 import os,sys
0002 import glob
0003 import logging
0004 import argparse
0005 import subprocess
0006 import time, datetime
0007 import urllib
0008 import json
0009
0010 from . import tools
0011 from .CLIHelper import CLIHelper
0012 from .CrabHelper import CrabHelper
0013 import FWCore.ParameterSet.Config as cms
0014 log = logging.getLogger(__name__)
0015
0016 class DTWorkflow(CLIHelper, CrabHelper):
0017 """ This is the base class for all DTWorkflows and contains some
0018 common tasks """
0019 def __init__(self, options):
0020 self.options = options
0021 super( DTWorkflow, self ).__init__()
0022 self.digilabel = "muonDTDigis"
0023
0024
0025 self.required_options_dict = {}
0026 self.required_options_prepare_dict = {}
0027 self.fill_required_options_dict()
0028 self.fill_required_options_prepare_dict()
0029
0030 self.pset_name = ""
0031 self.outpath_command_tag = ""
0032 self.output_files = []
0033 self.input_files = []
0034
0035 self.run_all_command = False
0036 self.files_reveived = False
0037 self._user = ""
0038
0039 os.chdir(self.options.working_dir)
0040
0041 def check_missing_options(self, requirements_dict):
0042 missing_options = []
0043
0044 if self.options.command in requirements_dict:
0045 for option in requirements_dict[self.options.command]:
0046 if not (hasattr(self.options, option)
0047 and ( (getattr(self.options,option))
0048 or isinstance(getattr(self.options,option), bool) )):
0049 missing_options.append(option)
0050 if len(missing_options) > 0:
0051 err = "The following CLI options are missing"
0052 err += " for command %s: " % self.options.command
0053 err += " ".join(missing_options)
0054 raise ValueError(err)
0055
0056 def run(self):
0057 """ Generalized function to run workflow command"""
0058 msg = "Preparing %s workflow" % self.options.workflow
0059 if hasattr(self.options, "command"):
0060 msg += " for command %s" % self.options.command
0061 log.info(msg)
0062 if self.options.config_path:
0063 self.load_options( self.options.config_path )
0064
0065 self.check_missing_options(self.required_options_prepare_dict)
0066 self.prepare_workflow()
0067
0068 if not os.path.exists( self.local_path ):
0069 os.makedirs(self.local_path)
0070
0071 self.dump_options()
0072
0073 self.check_missing_options(self.required_options_dict)
0074 try:
0075 run_function = getattr(self, self.options.command)
0076 except AttributeError:
0077 errmsg = "Class `{}` does not implement `{}` for workflow %s" % self.options.workflow
0078 if hasattr(self.options, "workflow_mode"):
0079 errmsg += "and workflow mode %s" % self.options.workflow_mode
0080 raise NotImplementedError( errmsg.format(self.__class__.__name__,
0081 self.options.command))
0082 log.debug("Running command %s" % self.options.command)
0083
0084 run_function()
0085
0086 def prepare_workflow(self):
0087 """ Abstract implementation of prepare workflow function"""
0088 errmsg = "Class `{}` does not implement `{}`"
0089 raise NotImplementedError( errmsg.format(self.__class__.__name__,
0090 "prepare_workflow"))
0091
0092 def all(self):
0093 """ generalized function to perform several workflow mode commands in chain.
0094 All commands mus be specified in self.all_commands list in workflow mode specific
0095 prepare function in child workflow objects.
0096 """
0097 self.run_all_command = True
0098 for command in self.all_commands:
0099 log.info(f"Will run command: {command}")
0100 self.options.command = command
0101 self.run()
0102
0103 def submit(self):
0104 self.submit_crab_task()
0105
0106 def check(self):
0107 """ Function to check status of submitted tasks """
0108 self.check_crabtask()
0109
0110 def write(self):
0111 self.runCMSSWtask()
0112
0113 def dump(self):
0114 self.runCMSSWtask()
0115
0116 def correction(self):
0117 self.runCMSSWtask()
0118
0119 def add_preselection(self):
0120 """ Add preselection to the process object stored in workflow_object"""
0121 if not hasattr(self, "process"):
0122 raise NameError("Process is not initalized in workflow object")
0123 pathsequence = self.options.preselection.split(':')[0]
0124 seqname = self.options.preselection.split(':')[1]
0125 self.process.load(pathsequence)
0126 tools.prependPaths(self.process, seqname)
0127
0128 def add_raw_option(self):
0129 getattr(self.process, self.digilabel).inputLabel = self.options.raw_data_label
0130 tools.prependPaths(self.process,self.digilabel)
0131
0132 def add_local_t0_db(self, local=False):
0133 """ Add a local t0 database as input. Use the option local is used
0134 if the pset is processed locally and not with crab.
0135 """
0136 if local:
0137 connect = os.path.abspath(self.options.inputT0DB)
0138 else:
0139 connect = os.path.basename(self.options.inputT0DB)
0140 self.addPoolDBESSource( process = self.process,
0141 moduleName = 't0DB',
0142 record = 'DTT0Rcd',
0143 tag = 't0',
0144 connect = 'sqlite_file:%s' % connect)
0145 self.input_files.append(os.path.abspath(self.options.inputT0DB))
0146
0147 def add_local_vdrift_db(self, local=False):
0148 """ Add a local vdrift database as input. Use the option local is used
0149 if the pset is processed locally and not with crab.
0150 """
0151 if local:
0152 connect = os.path.abspath(self.options.inputVDriftDB)
0153 else:
0154 connect = os.path.basename(self.options.inputVDriftDB)
0155 self.addPoolDBESSource( process = self.process,
0156 moduleName = 'vDriftDB',
0157 record = 'DTMtimeRcd',
0158 tag = 'vDrift',
0159 connect = 'sqlite_file:%s' % connect)
0160 self.input_files.append( os.path.abspath(self.options.inputVDriftDB) )
0161
0162 def add_local_calib_db(self, local=False):
0163 """ Add a local calib database as input. Use the option local is used
0164 if the pset is processed locally and not with crab.
0165 """
0166 label = ''
0167 if self.options.datasettype == "Cosmics":
0168 label = 'cosmics'
0169 if local:
0170 connect = os.path.abspath(self.options.inputCalibDB)
0171 else:
0172 connect = os.path.basename(self.options.inputCalibDB)
0173 self.addPoolDBESSource( process = self.process,
0174 moduleName = 'calibDB',
0175 record = 'DTTtrigRcd',
0176 tag = 'ttrig',
0177 connect = str("sqlite_file:%s" % connect),
0178 label = label
0179 )
0180 self.input_files.append( os.path.abspath(self.options.inputCalibDB) )
0181
0182 def add_local_custom_db(self):
0183 for option in ('inputDBRcd', 'connectStrDBTag'):
0184 if hasattr(self.options, option) and not getattr(self.options, option):
0185 raise ValueError("Option %s needed for custom input db" % option)
0186 self.addPoolDBESSource( process = self.process,
0187 record = self.options.inputDBRcd,
0188 tag = self.options.inputDBTag,
0189 connect = self.options.connectStrDBTag,
0190 moduleName = 'customDB%s' % self.options.inputDBRcd
0191 )
0192
0193 def prepare_common_submit(self):
0194 """ Common operations used in most prepare_[workflow_mode]_submit functions"""
0195 if not self.options.run:
0196 raise ValueError("Option run is required for submission!")
0197 if hasattr(self.options, "inputT0DB") and self.options.inputT0DB:
0198 self.add_local_t0_db()
0199
0200 if hasattr(self.options, "inputVDriftDB") and self.options.inputVDriftDB:
0201 self.add_local_vdrift_db()
0202
0203 if hasattr(self.options, "inputDBTag") and self.options.inputDBTag:
0204 self.add_local_custom_db()
0205
0206 if self.options.run_on_RAW:
0207 self.add_raw_option()
0208 if self.options.preselection:
0209 self.add_preselection()
0210
0211 def prepare_common_write(self, do_hadd=True):
0212 """ Common operations used in most prepare_[workflow_mode]_erite functions"""
0213 self.load_options_command("submit")
0214 output_path = os.path.join( self.local_path, "unmerged_results" )
0215 merged_file = os.path.join(self.result_path, self.output_file)
0216 crabtask = self.crabFunctions.CrabTask(crab_config = self.crab_config_filepath,
0217 initUpdate = False)
0218 if not (self.options.skip_stageout or self.files_reveived or self.options.no_exec):
0219 output_files = self.get_output_files(crabtask, output_path)
0220 if "xrootd" not in output_files.keys():
0221 raise RuntimeError("Could not get output files. No xrootd key found.")
0222 if len(output_files["xrootd"]) == 0:
0223 raise RuntimeError("Could not get output files. Output file list is empty.")
0224 log.info("Received files from storage element")
0225 log.info("Using hadd to merge output files")
0226 if not self.options.no_exec and do_hadd:
0227 returncode = tools.haddLocal(output_files["xrootd"], merged_file)
0228 if returncode != 0:
0229 raise RuntimeError("Failed to merge files with hadd")
0230 return crabtask.crabConfig.Data.outputDatasetTag
0231
0232 def prepare_common_dump(self, db_path):
0233 self.process = tools.loadCmsProcess(self.pset_template)
0234 self.process.calibDB.connect = 'sqlite_file:%s' % db_path
0235 try:
0236 path = self.result_path
0237 except:
0238 path = os.getcwd()
0239 print("path", path)
0240 out_path = os.path.abspath(os.path.join(path,
0241 os.path.splitext(db_path)[0] + ".txt"))
0242
0243 self.process.dumpToFile.outputFileName = out_path
0244
0245 @staticmethod
0246 def addPoolDBESSource( process,
0247 moduleName,
0248 record,
0249 tag,
0250 connect='sqlite_file:',
0251 label='',):
0252
0253 from CondCore.CondDB.CondDB_cfi import CondDB
0254
0255 calibDB = cms.ESSource("PoolDBESSource",
0256 CondDB,
0257 toGet = cms.VPSet(cms.PSet(
0258 record = cms.string(record),
0259 tag = cms.string(tag),
0260 label = cms.untracked.string(label)
0261 )),
0262 )
0263 calibDB.connect = cms.string( str(connect) )
0264
0265 if 'oracle:' in connect:
0266 calibDB.DBParameters.authenticationPath = '/afs/cern.ch/cms/DB/conddb'
0267 setattr(process,moduleName,calibDB)
0268 setattr(process,"es_prefer_" + moduleName,cms.ESPrefer('PoolDBESSource',
0269 moduleName)
0270 )
0271
0272 def get_output_files(self, crabtask, output_path):
0273 res = self.crab.callCrabCommand( ["getoutput",
0274 "--dump",
0275 "--xrootd",
0276 crabtask.crabFolder ] )
0277
0278 return res
0279
0280 def runCMSSWtask(self, pset_path=""):
0281 """ Run a cmsRun job locally. The member variable self.pset_path is used
0282 if pset_path argument is not given"""
0283 if self.options.no_exec:
0284 return 0
0285 process = subprocess.Popen( "cmsRun %s" % self.pset_path,
0286 stdout=subprocess.PIPE,
0287 stderr=subprocess.STDOUT,
0288 shell = True)
0289 stdout = process.communicate()[0]
0290 log.info(stdout)
0291 if process.returncode != 0:
0292 raise RuntimeError("Failed to use cmsRun for pset %s" % self.pset_name)
0293 return process.returncode
0294
0295 @property
0296 def remote_out_path(self):
0297 """ Output path on remote excluding user base path
0298 Returns a dict if crab is used due to crab path setting policy"""
0299 if self.options.command =="submit":
0300 return {
0301 "outLFNDirBase" : os.path.join( "/store",
0302 "user",
0303 self.user,
0304 'DTCalibration/',
0305 self.outpath_command_tag,
0306 self.outpath_workflow_mode_tag),
0307 "outputDatasetTag" : self.tag
0308 }
0309 else:
0310 return os.path.join( 'DTCalibration/',
0311 datasetstr,
0312 'Run' + str(self.options.run),
0313 self.outpath_command_tag,
0314 self.outpath_workflow_mode_tag,
0315 'v' + str(self.options.trial),
0316 )
0317 @property
0318 def outpath_workflow_mode_tag(self):
0319 if not self.options.workflow_mode in self.outpath_workflow_mode_dict:
0320 raise NotImplementedError("%s missing in outpath_workflow_mode_dict" % self.options.workflow_mode)
0321 return self.outpath_workflow_mode_dict[self.options.workflow_mode]
0322
0323 @property
0324 def tag(self):
0325 return 'Run' + str(self.options.run) + '_v' + str(self.options.trial)
0326
0327 @property
0328 def user(self):
0329 if self._user:
0330 return self._user
0331 if hasattr(self.options, "user") and self.options.user:
0332 self._user = self.options.user
0333 else:
0334 self._user = self.crab.checkusername()
0335 return self._user
0336
0337 @property
0338 def local_path(self):
0339 """ Output path on local machine """
0340 if self.options.run and self.options.label:
0341 prefix = "Run%d-%s_v%d" % ( self.options.run,
0342 self.options.label,
0343 self.options.trial)
0344 else:
0345 prefix = ""
0346 if self.outpath_workflow_mode_tag:
0347 path = os.path.join( self.options.working_dir,
0348 prefix,
0349 self.outpath_workflow_mode_tag)
0350 else:
0351 path = os.path.join( self.options.working_dir,
0352 prefix,
0353 self.outpath_command_tag )
0354 return path
0355
0356 @property
0357 def result_path(self):
0358 result_path = os.path.abspath(os.path.join(self.local_path,"results"))
0359 if not os.path.exists(result_path):
0360 os.makedirs(result_path)
0361 return result_path
0362
0363 @property
0364 def pset_template_base_bath(self):
0365 """ Base path to folder containing pset files for cmsRun"""
0366 return os.path.expandvars(os.path.join("$CMSSW_BASE",
0367 "src",
0368 "CalibMuon",
0369 "test",
0370 )
0371 )
0372
0373 @property
0374 def pset_path(self):
0375 """ full path to the pset file """
0376 basepath = os.path.join( self.local_path, "psets")
0377 if not os.path.exists( basepath ):
0378 os.makedirs( basepath )
0379 return os.path.join( basepath, self.pset_name )
0380
0381 def write_pset_file(self):
0382 if not hasattr(self, "process"):
0383 raise NameError("Process is not initalized in workflow object")
0384 if not os.path.exists(self.local_path):
0385 os.makedirs(self.local_path)
0386 with open( self.pset_path,'w') as pfile:
0387 pfile.write(self.process.dumpPython())
0388
0389 def get_config_name(self, command= ""):
0390 """ Create the name for the output json file which will be dumped"""
0391 if not command:
0392 command = self.options.command
0393 return "config_" + command + ".json"
0394
0395 def dump_options(self):
0396 with open(os.path.join(self.local_path, self.get_config_name()),"w") as out_file:
0397 json.dump(vars(self.options), out_file, indent=4)
0398
0399 def load_options(self, config_file_path):
0400 if not os.path.exists(config_file_path):
0401 raise IOError("File %s not found" % config_file_path)
0402 with open(config_file_path, "r") as input_file:
0403 config_json = json.load(input_file)
0404 for key, val in config_json.items():
0405 if not hasattr(self.options, key) or not getattr(self.options, key):
0406 setattr(self.options, key, val)
0407
0408 def load_options_command(self, command ):
0409 """Load options for previous command in workflow """
0410 if not self.options.config_path:
0411 if not self.options.run:
0412 raise RuntimeError("Option run is required if no config path specified")
0413 if not os.path.exists(self.local_path):
0414 raise IOError("Local path %s does not exist" % self.local_path)
0415 self.options.config_path = os.path.join(self.local_path,
0416 self.get_config_name(command))
0417 self.load_options( self.options.config_path )
0418
