Back to home page

Project CMSSW displayed by LXR

 
 

    


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         # dict to hold required variables. Can not be marked in argparse to allow
0024         # loading of options from config
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         # These variables are determined in the derived classes
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         # change to working directory
0039         os.chdir(self.options.working_dir)
0040 
0041     def check_missing_options(self, requirements_dict):
0042         missing_options = []
0043         # check if all required options exist
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         #check if all options to prepare the command are used
0065         self.check_missing_options(self.required_options_prepare_dict)
0066         self.prepare_workflow()
0067         # create output folder if they do not exist yet
0068         if not os.path.exists( self.local_path ):
0069             os.makedirs(self.local_path)
0070         # dump used options
0071         self.dump_options()
0072         #check if all options to run the command are used
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         # call chosen function
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         #if authPath: calibDB.DBParameters.authenticationPath = authPath
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